refactor losses

Author: liellnima

emulator/src/core/losses.py

@@ -1,14 +1,11 @@
+import abc
 import torch
 import torch.nn as nn
 import logging
 import gpytorch
 
-import xarray as xr
-
 from pytorch_lightning.utilities import rank_zero_only
 
-import numpy as np
-
 # import problems from utils
 def get_logger(name=__name__, level=logging.INFO) -> logging.Logger:
     """Initializes multi-GPU-friendly python logger."""
@@ -93,24 +90,64 @@ def forward(self, pred, y):
 
         return error
 
-class NRMSELoss_s_ClimateBench(nn.Module):
+# Parent class with latitude weights and some other functions
+class ClimateSetLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+    # weighting to account for decreasing grid-cell area towards poles
+    def get_latitude_weights(self, lat_size: int) -> torch.Tensor:
+        """ Returns latitude weights for a given number of data points along the latitude axis (y.shape[-2]).
+        The weights are 1 at the equator and decrease towards the poles. 
+
+        Parameters:
+            lat_size (int): How many latitude data points should be considered
+        Returns:
+            torch.Tensor: Latitude weights
+        """
+        # we are not using -90, 90, but a small offset to make sure none of the datapoints get weights of 0
+        lats = torch.linspace(-89.75, 89.75, lat_size)
+        weights = torch.cos((torch.pi * lats) / 180)
+        weights = weights.unsqueeze(-1)
+        return weights
+    
+    def weighted_global_mean(self, x: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
+        """ Get a weighted global mean of a tensor.
+        Parameters:
+            x (torch.Tensor): Values that are to be averaged
+            weights (torch.Tensor): Latitude weights for x
+        Returns:
+            torch.Tensor: Single number representing the global mean
+        """
+        # sum_lat(sum_lon(x * weights)) / N_lat * N_lon
+        # i.e.: sum(sum(x * weights)) / (96 * 144)
+        return torch.mean(x * weights, dim=(-2, -1)) # dims order does not matter
+    
+    def check_lat_lon(self, pred: torch.Tensor, y: torch.Tensor):
+        """ Functions that checks if latitude and longitude is behaving as expected.
+        Parameters:
+            pred (torch.Tensor): Predictions
+            y (torch.Tensor): Targets
+        """
+        # Expected shape: [4, 12, 96, 144] -> [batch, time, latitude, longitude]
+        if (pred.shape[-1] == 1) or (y[-1].shape == 1):
+            raise ValueError("Loss function: Last dimension (values/channels) must be squeezed away")
+        
+        if (pred.shape[-1] < pred.shape[-2]) or (y.shape[-1] < y.shape[-2]):
+            raise ValueError("There are more latitude than longitude grid cells. Check if you swapped longitude and latitude.")
+
+
+class NRMSELoss_s_ClimateBench(ClimateSetLoss):
     """
     Spatial normalized weighted RMSE taken from Climate Bench.
     Weighting to account for decreasing grid size towards the poles.
     """
-
     def __init__(self):
         super().__init__()
         self.mse = nn.MSELoss(reduction="none")
 
-    def forward(self, pred, y):
-        # weighting to account for decreasing grid-cell area towards poles
-        # latitude weights
-        lat_size = y.shape[-2]
-        lats = torch.linspace(-89.75, 89.75, lat_size)
-        weights = torch.cos((torch.pi * lats) / 180)
-
-        weights = weights.unsqueeze(-1)
+    def forward(self, pred: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        self.check_lat_lon(pred, y)
+        weights = self.get_latitude_weights(y.shape[-2])
         weights = weights.to(device)
 
         # nrmses = sqrt((weights * (pred_mean_t - y_mean_n_t)**2)_mean_s) / ((weights*y)_mean_s)_mean_t_n
@@ -127,28 +164,18 @@ def forward(self, pred, y):
 
         return nrmse_s
 
-    def weighted_global_mean(self, x, weights):
-        # sum_lat(sum_lon(x * weights)) / N_lat * N_lon
-        # i.e.: sum(sum(x * weights)) / (96 * 144)
-        return torch.mean(x * weights, dim=(-2, -1)) # dims order does not matter
-
-class NRMSELoss_g_ClimateBench(nn.Module):
+class NRMSELoss_g_ClimateBench(ClimateSetLoss):
     """
     Spatial normalized weighted RMSE taken from Climate Bench.
     Weighting to account for decreasing grid size towards the pole.
     """
-
     def __init__(self):
         super().__init__()
         self.mse = nn.MSELoss(reduction="none")
 
-    def forward(self, pred, y):
-        #latitude weighting to account for decreasing grid-cell area towards pole
-        lat_size = y.shape[-2]
-        lats = torch.linspace(-89.75, 89.75, lat_size)
-        # same like np.cos(np.deg2rad(lats))
-        weights = torch.cos((torch.pi * lats) / 180)
-        weights = weights.unsqueeze(-1)
+    def forward(self, pred: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        self.check_lat_lon(pred, y)
+        weights = self.get_latitude_weights(y.shape[-2])
         weights = weights.to(device)
 
         # nrmseg = sqrt(
@@ -164,83 +191,67 @@ def forward(self, pred, y):
             )
             / self.weighted_global_mean(y, weights).mean(dim=(0, 1))
         )
-        # TODO understand: the values are in the same range like nrmse_s - why do we need to adapt them?
         return nrmse_g
 
-    def weighted_global_mean(self, x, weights):
-        # sum_lat(sum_lon(x * weights)) / N_lat * N_lon
-        # i.e.: sum(sum(x * weights)) / (144 * 96)
-        return torch.mean(x * weights, dim=(-2, -1))
-
-class NRMSELoss_ClimateBench(nn.Module):
+class NRMSELoss_ClimateBench(ClimateSetLoss):
     """
     Combination of global weighted and spatially weighted nrmse.
-
     """
-
     def __init__(self, alpha: int = 5):
         super().__init__()
 
         self.nrmse_g = NRMSELoss_g_ClimateBench()
         self.nrmse_s = NRMSELoss_s_ClimateBench()
         self.alpha = alpha
 
-    def forward(self, pred, y):
+    def forward(self, pred: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        self.check_lat_lon(pred, y)
         nrmseg = self.nrmse_g(pred, y)
         nrmses = self.nrmse_s(pred, y)
         nrmse = nrmses + self.alpha * nrmseg
         return nrmse
 
-class LLWeighted_RMSELoss_WeatherBench(nn.Module):
-
+class LLWeighted_RMSELoss_WeatherBench(ClimateSetLoss):
     """
     Weigthed RMSE taken from Weather Bench.
     Weighting to account for decreasing grid sizes towards the pole.
 
     rmse = mean over forecasts and time of torch.sqrt( mean over lon lat L(lat_j)*)MSE(pred, y)
     weights = cos(latitude)/cos(latitude).mean()
     """
-
     def __init__(self):
         super().__init__()
 
         self.mse = nn.MSELoss(reduction="none")
 
-    def forward(self, pred, y):
-
-        lat_size = y.shape[-2]
-        lats = torch.linspace(-89.75, 89.75, lat_size)
-        weights = torch.cos((torch.pi * lats) / 180)
-        weights = weights.unsqueeze(-1)
+    def forward(self, pred: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        self.check_lat_lon(pred, y)
+        weights = self.get_latitude_weights(y.shape[-2])
+        weights = weights.to(device)
         weights = weights.to(device)
 
         #rmse_before = torch.sqrt(torch.mean(weights * self.mse(pred, y), dim=(-2, -1))).mean()
         rmse = torch.mean(torch.sqrt(torch.mean(weights * ((pred - y)**2), dim=([-2, -1]))))
 
         return rmse
 
-class LLweighted_MSELoss_Climax(nn.Module):
+class LLweighted_MSELoss_Climax(ClimateSetLoss):
     """
     Latitude weighted mean squared error taken from ClimaX.
     Allows to weight the loss by the cosine of the latitude to account for gridding differences at equator vs. poles.
     Applied per variable.
     If given a mask, normalized by sum of that.
-
     """
-
     def __init__(self, mask=None):
         super().__init__()
 
         self.mse = nn.MSELoss(reduction="none")
         self.mask = mask
 
-    def forward(self, pred, y):
+    def forward(self, pred: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        self.check_lat_lon(pred, y)
         mse = self.mse(pred, y)
-
-        lat_size = y.shape[-2]
-        lats = torch.linspace(-89.75, 89.75, lat_size)
-        weights = torch.cos((torch.pi * lats) / 180)
-        weights = weights.unsqueeze(-1)
+        weights = self.get_latitude_weights(y.shape[-2])
         weights = weights.to(device)
 
         # how they create the weights (does not work for us, results make no sense)
@@ -261,46 +272,32 @@ def forward(self, pred, y):
         return error
 
 
-class LLweighted_RMSELoss_Climax(nn.Module):
+class LLweighted_RMSELoss_Climax(ClimateSetLoss):
     """
     Latitude weighted root mean squared error taken from ClimaX.
     Allows to weight the loss by the cosine of the latitude to account for gridding differences at equator vs. poles.
     Applied per variable.
     If given a mask, normalized by sum of that.
     """
-
     def __init__(self, mask=None):
         super().__init__()
 
         self.mse = nn.MSELoss(reduction="none")
         self.mask = mask
 
-    def forward(self, pred, y):
+    def forward(self, pred: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         """ Latitude is expected to be on position -2
         """
-        lat_num_grid_cells = y.shape[-2]
-
-        # Expected shape: [4, 12, 96, 144] -> [batch, time, latitude, longitude]
-        if (pred.shape[-1] == 1) or (y[-1].shape == 1):
-            raise ValueError("Loss function: Last dimension (values/channels) must be squeezed away")
-        
-        if (pred.shape[-1] < pred.shape[-2]):
-            raise ValueError("There are more latitude than longitude grid cells. Check if you swapped longitude and latitude.")
-
+        self.check_lat_lon(pred, y)
         mse = self.mse(pred, y) # [batch, time, lat, lon]
 
-        latitudes = torch.linspace(-89.75, 89.75, lat_num_grid_cells)
-        # torch.abs: -90 and + 90 get -0.000X as weight -> make all weights positive
-        weights = torch.abs(torch.cos(torch.deg2rad(latitudes))) 
-        weights = weights.unsqueeze(-1)
+        weights = self.get_latitude_weights(y.shape[-2])
+        weights = weights.to(device)
 
         # ClimaX creates weird weights by dividing them by the mean 
         # this leads to the climax rmse and mse to be the exact same like the unweighted mse / rmse
         #weights = weights / weights.mean() # ignored in this code
 
-        # move weights to device
-        weights = weights.to(device)
-
         if self.mask is not None:
             raise NotImplementedError("Masking is not supported in the loss functions anymore.")
 
@@ -335,8 +332,7 @@ def forward(self, pred, y):
 
     llmse_cx = LLweighted_MSELoss_Climax()
     llrmse_cx = LLweighted_RMSELoss_Climax()
-
-    # MSE: CHECKED
+    
     loss = mse(dummy, targets)
     print("MSE loss", loss, loss.size())
     # np_dummy = dummy.cpu().detach().numpy()
@@ -363,17 +359,9 @@ def forward(self, pred, y):
     loss = llrmse_cx(dummy, targets)
     print("CX rmse loss", loss, loss.size())
 
-# TESTS for losses:
-# - with specific tensor of 1s + offset
-# - with specific random tensors
-# - make sure output size is only one number (except if several channels?)
-
-# - compare losses for ones: rmse == nrmse_g
-# - with channels for different variables (make sure it's not breaking) / doing whatever is needed
-# - make sure WB and CX rmse losses are the same
 
 # REFACTOR
-# - weight function should be one function (utils)
+# - weight function should be one function
 
 # Same tests needed for metrics
 

tests/test_scoring/test_losses.py

@@ -22,6 +22,10 @@
 
 PRECISION_VALUE = 0.0005
 
+# TODO add channel issue
+# TODO test weighting function
+# TODO test if numpy weights and torch weights are the same!
+
 @pytest.fixture 
 def rand_targets():
     return torch.rand(size=(batch_size, out_time, lat, lon))
@@ -87,7 +91,7 @@ def test_nrmse(rand_predics, rand_targets, ones_predics, ones_targets):
     expected_loss(loss_rand, 0.1142, 0.005)
 
 def test_wb_rmse(rand_predics, rand_targets, ones_predics, ones_targets):
-    error = LLWeighted_RMSELoss_WheatherBench()
+    error = LLWeighted_RMSELoss_WeatherBench()
     loss_ones = error(ones_predics, ones_targets) 
     loss_rand = error(rand_predics, rand_targets)
     expected_loss(loss_ones, 0.0795, PRECISION_VALUE)
@@ -112,7 +116,7 @@ def test_equality_ones(ones_predics, ones_targets):
     nrmse_s = NRMSELoss_s_ClimateBench()
     nrmse_g = NRMSELoss_g_ClimateBench()
     nrmse = NRMSELoss_ClimateBench()
-    wb_rmse = LLWeighted_RMSELoss_WheatherBench()
+    wb_rmse = LLWeighted_RMSELoss_WeatherBench()
     cx_rmse = LLweighted_RMSELoss_Climax()
 
     rmse_loss_ones = rmse(ones_predics, ones_targets)
@@ -127,7 +131,7 @@ def test_equality_ones(ones_predics, ones_targets):
     assert (nrmse_s_loss_ones.item() + 5 * nrmse_g_loss_ones.item()) == pytest.approx(nrmse_loss_ones.item(), abs=PRECISION_VALUE)
 
 def test_equality_rand(rand_predics, rand_targets):
-    wb_rmse = LLWeighted_RMSELoss_WheatherBench()
+    wb_rmse = LLWeighted_RMSELoss_WeatherBench()
     cx_rmse = LLweighted_RMSELoss_Climax()
     nrmse_s = NRMSELoss_s_ClimateBench()
     nrmse_g = NRMSELoss_g_ClimateBench()
@@ -143,4 +147,3 @@ def test_equality_rand(rand_predics, rand_targets):
     assert (nrmse_s_loss_rand.item() + 5 * nrmse_g_loss_rand.item()) == pytest.approx(nrmse_loss_rand.item(), abs=PRECISION_VALUE)
 
 
-# TODO add channel issue