Corrected expected value of Gamma-Bernoulli distribution.

climax/core/predict.py

@@ -62,18 +62,18 @@ def predict_ERA5(net, ERA5_ds, predictand, batch_size=16, **kwargs):
         ds = {'tasmax': target[:, 0, ...].squeeze(),
               'tasmin': target[:, 1, ...].squeeze()}
     elif predictand == 'pr':
-        ds = {
-            # probability of precipitation
-            'prob': torch.sigmoid(
-                torch.as_tensor(target[:, 0, ...].squeeze(),
-                                dtype=torch.float32)).numpy(),
-
-            # amount of precipitation: expected value of Bernoulli-Gamma
-            # distribution
-            # pr = p * shape * scale
-            'precipitation': (ds['prob'] * np.exp(target[:, 1, ...].squeeze())
-                              * np.exp(target[:, 2, ...].squeeze()))
-            }
+
+        # probability of precipitation
+        prob = torch.sigmoid(torch.as_tensor(target[:, 0, ...].squeeze(),
+                                             dtype=torch.float32)).numpy()
+
+        # amount of precipitation: expected value of Bernoulli-Gamma distribution
+        # pr = p * shape * scale
+        pr = (prob * np.exp(target[:, 1, ...].squeeze()) *
+              np.exp(target[:, 2, ...].squeeze()))
+
+        ds = {'prob': prob, 'precipitation': pr}
+
     else:
         # single predictand
         ds = {predictand: target}