Implement first function for latitude weighting (ecmwf#705)

* Changed logging level for some messages.

* Refactored loss computation to improve performance.

* Working around ruff issue

* - Refactored code to improve structure and readability
- Fixed problem with incomplete normalization over loss functions
- Solved problem with mse_weighted as loss function when mse is specified

* Fixed problems with multi-worker training

* add location weights, first commit

* assertion on mask and len(location_weights)

* restructuring of location weights and fixes in mse_channel_location_weighted function

* fix coords_raw dependency on offset and fstep

* ruff

* addressing review commits and fixing bug

* rm location_weight from default stream config

---------

Co-authored-by: Christian Lessig <christian.lessig@ecmwf.int>
Co-authored-by: Julian Kuehnert <julian.kuehnert@ecwmf.int>

Author: 3 people

packages/evaluate/src/weathergen/evaluate/plotter.py

@@ -8,10 +8,12 @@
 import xarray as xr
 
 from weathergen.utils.config import _load_private_conf
-work_dir = Path( _load_private_conf(None)['path_shared_working_dir']) / 'assets/cartopy'
+
+work_dir = Path(_load_private_conf(None)["path_shared_working_dir"]) / "assets/cartopy"
 import cartopy
-cartopy.config['data_dir'] = str(work_dir)
-cartopy.config['pre_existing_data_dir'] = str(work_dir)
+
+cartopy.config["data_dir"] = str(work_dir)
+cartopy.config["pre_existing_data_dir"] = str(work_dir)
 os.environ["CARTOPY_DATA_DIR"] = str(work_dir)
 
 np.seterr(divide="ignore", invalid="ignore")

src/weathergen/train/loss.py

@@ -135,9 +135,15 @@ def mse_channel_location_weighted(
     mask_nan = ~torch.isnan(target)
     pred = pred[0] if pred.shape[0] == 0 else pred.mean(0)
 
-    diff2 = torch.square(torch.where(mask_nan, target, 0) - torch.where(mask_nan, pred, 0)).mean(0)
-    wl = weights_points
-    loss_chs = ((diff2.transpose(1, 0) * wl).transpose(1, 0) if wl else diff2).mean(0)
+    diff2 = torch.square(torch.where(mask_nan, target, 0) - torch.where(mask_nan, pred, 0))
+    if weights_points is not None:
+        diff2 = (diff2.transpose(1, 0) * weights_points).transpose(1, 0)
+    loss_chs = diff2.mean(0)
     loss = torch.mean(loss_chs * weights_channels if weights_channels else loss_chs)
 
     return loss, loss_chs
+
+
+def cosine_latitude(stream_data, forecast_offset, fstep, min_value=1e-3, max_value=1.0):
+    latitudes_radian = stream_data.target_coords_raw[forecast_offset + fstep][:, 0] * np.pi / 180
+    return (max_value - min_value) * np.cos(latitudes_radian) + min_value

src/weathergen/train/loss_calculator.py

@@ -105,6 +105,16 @@ def _get_weights(self, stream_info):
 
         return stream_info_loss_weight, weights_channels
 
+    def _get_location_weights(self, stream_info, stream_data, forecast_offset, fstep):
+        location_weight_type = stream_info.get("location_weight", None)
+        if location_weight_type is None:
+            return None
+        weights_locations_fct = getattr(losses, location_weight_type)
+        weights_locations = weights_locations_fct(stream_data, forecast_offset, fstep)
+        weights_locations = weights_locations.to(device=self.device, non_blocking=True)
+
+        return weights_locations
+
     def _get_substep_masks(self, stream_info, fstep, stream_data):
         """
         Find substeps and create corresponding masks (reused across loss functions)
@@ -140,7 +150,7 @@ def _loss_per_loss_function(
 
         ctr_substeps = 0
         for mask_t in substep_masks:
-            assert mask_t.sum() == len(weights_locations) if weights_locations else True
+            assert mask_t.sum() == len(weights_locations) if weights_locations is not None else True
 
             loss, loss_chs = loss_fct(
                 target[mask_t], pred[:, mask_t], weights_channels, weights_locations
@@ -220,9 +230,6 @@ def compute_loss(
 
             stream_data = streams_data[i_batch][i_stream_info]
 
-            # TODO: set from stream info
-            weights_locations = None
-
             loss_fsteps = torch.tensor(0.0, device=self.device, requires_grad=True)
             ctr_fsteps = 0
             for fstep, target in enumerate(targets):
@@ -240,6 +247,11 @@ def compute_loss(
                 # get weigths for current streams
                 stream_loss_weight, weights_channels = self._get_weights(stream_info)
 
+                # get weights for locations
+                weights_locations = self._get_location_weights(
+                    stream_info, stream_data, self.cf.forecast_offset, fstep
+                )
+
                 # get masks for sub-time steps
                 substep_masks = self._get_substep_masks(stream_info, fstep, stream_data)