Clean implementation of model predictions.

climax/core/loss.py

@@ -4,9 +4,11 @@
 # -*- coding: utf-8 -*-
 
 # externals
+import numpy as np
 import torch
 import torch.nn.functional as F
 from torch.nn.modules.loss import _Loss
+from scipy.special import gamma
 
 
 class NaNLoss(_Loss):
@@ -46,7 +48,7 @@ class L1Loss(NaNLoss):
         return F.l1_loss(y_pred[mask], y_true[mask], reduction=self.reduction)
 
 
-class BernoulliGammaLoss(NaNLoss):
+class BernoulliLoss(NaNLoss):
 
     def __init__(self, size_average=None, reduce=None, reduction='mean',
                  min_amount=0):
@@ -55,6 +57,13 @@ class BernoulliGammaLoss(NaNLoss):
         # minimum amount of precipitation to be classified as precipitation
         self.min_amount = min_amount
 
+
+class BernoulliGammaLoss(BernoulliLoss):
+
+    def __init__(self, size_average=None, reduce=None, reduction='mean',
+                 min_amount=0):
+        super().__init__(size_average, reduce, reduction)
+
     def forward(self, y_pred, y_true):
 
         # convert to float32
@@ -96,16 +105,18 @@ class BernoulliGammaLoss(NaNLoss):
 
         return self.reduce(loss)
 
+    @staticmethod
+    def predict(p, shape, scale):
+        # pr = p * shape * scale
+        return p * shape * scale
+
 
-class BernoulliGenParetoLoss(NaNLoss):
+class BernoulliGenParetoLoss(BernoulliLoss):
 
     def __init__(self, size_average=None, reduce=None, reduction='mean',
                  min_amount=0):
         super().__init__(size_average, reduce, reduction)
 
-        # minimum amount of precipitation to be classified as precipitation
-        self.min_amount = min_amount
-
     def forward(self, y_pred, y_true):
 
         # convert to float32
@@ -146,15 +157,12 @@ class BernoulliGenParetoLoss(NaNLoss):
         return self.reduce(loss)
 
 
-class BernoulliWeibullLoss(NaNLoss):
+class BernoulliWeibullLoss(BernoulliLoss):
 
     def __init__(self, size_average=None, reduce=None, reduction='mean',
                  min_amount=0):
         super().__init__(size_average, reduce, reduction)
 
-        # minimum amount of precipitation to be classified as precipitation
-        self.min_amount = min_amount
-
     def forward(self, y_pred, y_true):
 
         # convert to float32
@@ -207,3 +215,8 @@ class BernoulliWeibullLoss(NaNLoss):
         # loss = torch.clamp(loss, max=100)
 
         return self.reduce(loss)
+
+    @staticmethod
+    def predict(p, shape, scale):
+        # pr = p * scale * gamma(1 + 1 / shape)
+        return p * scale *  gamma(1 + 1 / shape)

climax/core/predict.py

@@ -11,11 +11,9 @@ import torch
 import numpy as np
 import xarray as xr
 from torch.utils.data import DataLoader
-from scipy.special import gamma
 
 # locals
 from climax.core.dataset import EoDataset, NetCDFDataset
-from climax.core.loss import BernoulliGammaLoss, BernoulliWeibullLoss
 
 # module level logger
 LOGGER = logging.getLogger(__name__)
@@ -76,24 +74,14 @@ def predict_ERA5(net, ERA5_ds, predictand, loss, batch_size=16, **kwargs):
             prob = torch.sigmoid(torch.as_tensor(target[:, 0, ...].squeeze(),
                                                  dtype=torch.float32)).numpy()
 
-            # check which loss function is used
+            # shape and scale parameters
+            shape = np.exp(target[:, 1, ...].squeeze())
+            scale = np.exp(target[:, 2, ...].squeeze())
 
-            # Bernoulli-Gamma
-            if isinstance(loss, BernoulliGammaLoss):
-                # precipitation amount: expected value of Bernoulli-Gamma
-                # distribution
-                # pr = p * shape * scale
-                pr = (prob * np.exp(target[:, 1, ...].squeeze()) *
-                      np.exp(target[:, 2, ...].squeeze()))
-
-            # Bernoulli-Weibull
-            if isinstance(loss, BernoulliWeibullLoss):
-                # precipitation amount: expected value of Bernoulli-Weibull
-                # distribution
-                # pr = p * scale * tau(1 + 1 / shape)
-                pr = (prob * np.exp(target[:, 2, ...].squeeze()) *
-                      gamma(1 + 1 / np.exp(target[:, 1, ...])))
+            # precipitation amount
+            pr = loss.predict(prob, shape, scale)
 
+            # precipitation probability and amount
             ds = {'prob': prob, 'precipitation': pr}
 
     # add coordinates to arrays