Jk/develop/gamma decay (ecmwf#998)

* Update to develop, prepare for new experiment series

* gamma decay over fsteps first commit

* add gamma decay factor to config

* working gamma decay weighting

* rm breakpoint

* rm eval and plot configs

* reverting default config

---------

Co-authored-by: Matthias Karlbauer <matthias.karlbauer@ecmwf.int>
Co-authored-by: Julian Kuehnert <julian.kuehnert@ecwmf.int>

Author: 3 people

src/weathergen/train/loss.py

@@ -147,3 +147,9 @@ def mse_channel_location_weighted(
 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
+
+
+def gamma_decay(forecast_steps, gamma):
+    fsteps = np.arange(forecast_steps)
+    weights = gamma**fsteps
+    return weights * (len(fsteps) / np.sum(weights))

src/weathergen/train/loss_calculator.py

@@ -107,6 +107,13 @@ def _get_weights(self, stream_info):
 
         return stream_info_loss_weight, weights_channels
 
+    def _get_fstep_weights(self, forecast_steps):
+        timestep_weight_config = self.cf.get("timestep_weight")
+        if timestep_weight_config is None:
+            return [1.0 for _ in range(forecast_steps)]
+        weights_timestep_fct = getattr(losses, timestep_weight_config[0])
+        return weights_timestep_fct(forecast_steps, timestep_weight_config[1])
+
     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:
@@ -232,14 +239,18 @@ def compute_loss(
 
             stream_data = streams_data[i_batch][i_stream_info]
 
+            fstep_loss_weights = self._get_fstep_weights(len(targets))
+
             loss_fsteps = torch.tensor(0.0, device=self.device, requires_grad=True)
             ctr_fsteps = 0
             stream_is_spoof = streams_data[i_batch][i_stream_info].is_spoof
             if stream_is_spoof:
                 mask = torch.tensor(0.0, device=self.device, requires_grad=False)
             else:
                 mask = torch.tensor(1.0, device=self.device, requires_grad=False)
-            for fstep, target in enumerate(targets):
+            for fstep, (target, fstep_weight) in enumerate(
+                zip(targets, fstep_loss_weights, strict=False)
+            ):
                 # skip if either target or prediction has no data points
                 pred = preds[fstep][i_stream_info]
                 if not (target.shape[0] > 0 and pred.shape[0] > 0):
@@ -280,7 +291,9 @@ def compute_loss(
 
                     # Add the weighted and normalized loss from this loss function to the total
                     # batch loss
-                    loss_fstep = loss_fstep + (loss_fct_weight * loss_lfct * stream_loss_weight)
+                    loss_fstep = loss_fstep + (
+                        loss_fct_weight * loss_lfct * stream_loss_weight * fstep_weight
+                    )
                     ctr_loss_fcts += 1 if loss_lfct > 0.0 else 0
 
                 loss_fsteps = loss_fsteps + (loss_fstep / ctr_loss_fcts if ctr_loss_fcts > 0 else 0)