Corrected expected value of Gamma-Bernoulli distribution.

climax/core/predict.py

@@ -68,10 +68,11 @@ def predict_ERA5(net, ERA5_ds, predictand, batch_size=16, **kwargs):
                 torch.as_tensor(target[:, 0, ...].squeeze(),
                                 dtype=torch.float32)).numpy(),
 
-            # amount of precipitation: expected value of gamma distribution
-            # pr = shape * scale
-            'precipitation': (np.exp(target[:, 1, ...]) *
-                              np.exp(target[:, 2, ...])).squeeze()
+            # 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()))
             }
     else:
         # single predictand