Residual training in the downscaling_only mode + add dates to the Item class

py4cast/datasets/base.py

@@ -44,11 +44,13 @@ class Item:
     inputs has shape (timestep, lat, lon, features)
     outputs has shape (timestep, lat, lon, features)
     forcing has shape (timestep, lat, lon, features)
+    valdity_times has shape (timestep)
     """
 
     inputs: NamedTensor | None
     forcing: NamedTensor | None
     outputs: NamedTensor
+    validity_times: list[dt.datetime]
 
     def unsqueeze_(self, dim_name: str, dim_index: int):
         """
@@ -116,20 +118,21 @@ def __str__(self) -> str:
         """
         table = []
         for attr in (f.name for f in fields(self)):
-            nt: NamedTensor = getattr(self, attr)
-            if nt is not None:
-                for feature_name in nt.feature_names:
-                    tensor = nt[feature_name]
-                    table.append(
-                        [
-                            attr,
-                            nt.names,
-                            list(nt[feature_name].shape),
-                            feature_name,
-                            tensor.min(),
-                            tensor.max(),
-                        ]
-                    )
+            if attr != "validity_times":
+                nt: NamedTensor = getattr(self, attr)
+                if nt is not None:
+                    for feature_name in nt.feature_names:
+                        tensor = nt[feature_name]
+                        table.append(
+                            [
+                                attr,
+                                nt.names,
+                                list(nt[feature_name].shape),
+                                feature_name,
+                                tensor.min(),
+                                tensor.max(),
+                            ]
+                        )
         headers = [
             "Type",
             "Dimension Names",
@@ -175,16 +178,19 @@ def collate_fn(items: List[Item]) -> ItemBatch:
     """
     # Here we postpone that for each batch the same dimension should be present.
     batch_of_items = {}
-
     # Iterate over inputs, outputs and forcing fields
     for field_name in (f.name for f in fields(Item)):
-        batched_tensor = collate_tensor_fn(
-            [getattr(item, field_name).tensor for item in items]
-        ).type(torch.float32)
-
-        batch_of_items[field_name] = NamedTensor.expand_to_batch_like(
-            batched_tensor, getattr(items[0], field_name)
-        )
+        if field_name == "validity_times":
+            batched_valid_times = [getattr(item, field_name) for item in items]
+            batch_of_items[field_name] = batched_valid_times
+        else:
+            batched_tensor = collate_tensor_fn(
+                [getattr(item, field_name).tensor for item in items]
+            ).type(torch.float32)
+
+            batch_of_items[field_name] = NamedTensor.expand_to_batch_like(
+                batched_tensor, getattr(items[0], field_name)
+            )
 
     return ItemBatch(**batch_of_items)
 
@@ -517,6 +523,7 @@ def load(self, no_standardize: bool = False) -> Item:
             inputs=inputs,
             outputs=outputs,
             forcing=forcing,
+            validity_times=self.output_timestamps.validity_times,
         )
 
     def plot(self, item: Item, step: int, save_path: Path = None) -> None:

py4cast/lightning.py

@@ -318,6 +318,11 @@ def setup(self, stage=None):
             self.psd_plot_metric = MetricPSDK(self.save_path, pred_step=max_pred_step)
             self.acc_metric = MetricACC(self.dataset_info)
             self.configure_loggers()
+            self.list_metrics = [
+                self.acc_metric,
+                self.psd_plot_metric,
+                self.rmse_psd_plot_metric,
+            ]
 
     def configure_loggers(self):
         layout = {
@@ -516,6 +521,8 @@ def _common_step(
 
         self.original_shape = None
 
+        ds = self.training_strategy == "downscaling_only"
+
         if self.model.model_type == ModelType.GRAPH:
             # Stack original shape to reshape later
             self.original_shape = batch.inputs.tensor.shape
@@ -531,11 +538,24 @@ def _common_step(
         # to check at inference time if the feature names are the same
         # also useful to build NamedTensor outputs with same feature and dim names
         # If model type is graph, flat the lon/lat dim before saving the dims
-        if batch_idx == 0 and phase == "train":
+        if batch_idx == 0:
             self.input_feature_names = batch.inputs.feature_names
             self.output_feature_names = batch.outputs.feature_names
             self.output_dim_names = batch.outputs.names
             self.output_dtype = batch.outputs.tensor.dtype
+            # save the indices of the features common to the forcings and outputs in downscaling_only mode
+            # useful for reconstructing the prediction from the residual
+            if ds:
+                forcing_feature_names = batch.forcing.feature_names
+                common_features_idx = []
+                for output_feature_name in self.output_feature_names:
+                    for i, feature_forcing_name in enumerate(forcing_feature_names):
+                        if (
+                            output_feature_name.split("_")[1:]
+                            == feature_forcing_name.split("_")[1:]
+                        ):
+                            common_features_idx.append(i)
+                self.common_features_idx = common_features_idx
 
         prev_states = batch.inputs
         prediction_list = []
@@ -575,8 +595,6 @@ def _common_step(
                 else:
                     y = self.model(x)
 
-                ds = self.training_strategy == "downscaling_only"
-
                 # select the last timestep
                 last_prev_state = prev_states.select_tensor_dim("timestep", -1).clone()
                 if self.mask_on_nan:
@@ -590,6 +608,17 @@ def _common_step(
                         + y * step_diff_std
                         + step_diff_mean
                     )
+                elif ds:
+                    # update the coarse forcings with our netwwork output
+                    last_coarse_step = batch.forcing.select_tensor_dim(
+                        "timestep", -1
+                    ).clone()
+                    if self.mask_on_nan:
+                        last_coarse_step = torch.nan_to_num(last_coarse_step, nan=0)
+                    # only add common features
+                    predicted_state = (
+                        last_coarse_step[:, :, :, self.common_features_idx] + y
+                    )
                 else:
                     predicted_state = last_prev_state * (1 - ds) + y
 
@@ -919,10 +948,9 @@ def on_validation_epoch_end(self):
 
         if self.logging_enabled:
             # Get dict of metrics' results
-            dict_metrics = dict()
-            dict_metrics.update(self.psd_plot_metric.compute())
-            dict_metrics.update(self.rmse_psd_plot_metric.compute())
-            dict_metrics.update(self.acc_metric.compute())
+            dict_metrics = {}
+            for metric in self.list_metrics:
+                dict_metrics.update(metric.compute())
             for name, elmnt in dict_metrics.items():
                 if isinstance(elmnt, matplotlib.figure.Figure):
                     # Tensorboard logger
@@ -1040,9 +1068,8 @@ def on_test_epoch_end(self):
         """
         if self.logging_enabled:
             dict_metrics = {}
-            dict_metrics.update(self.psd_plot_metric.compute(prefix="test"))
-            dict_metrics.update(self.rmse_psd_plot_metric.compute(prefix="test"))
-            dict_metrics.update(self.acc_metric.compute(prefix="test"))
+            for metric in self.list_metrics:
+                dict_metrics.update(metric.compute(prefix="test"))
 
             for name, elmnt in dict_metrics.items():
                 if isinstance(elmnt, matplotlib.figure.Figure):

tests/test_datasets.py

@@ -33,7 +33,11 @@ def test_item():
         names=["lat", "lon", "features"],
         feature_names=[f"forcing_{i}" for i in range(2)],
     )
-    item = Item(inputs=inputs, outputs=outputs, forcing=forcing)
+
+    validity_times = [datetime.datetime(year=2023, month=1, day=1, hour=18)]
+    item = Item(
+        inputs=inputs, outputs=outputs, forcing=forcing, validity_times=validity_times
+    )
     print(item)
 
     # test collate_fn
@@ -59,7 +63,12 @@ def test_item():
             feature_names=[f"feature_{i}" for i in range(4)],
         )
 
-        item = Item(inputs=inputs, outputs=outputs, forcing=forcing)
+        item = Item(
+            inputs=inputs,
+            outputs=outputs,
+            forcing=forcing,
+            validity_times=validity_times,
+        )
 
     # Input and Output must have the same feature names
     with pytest.raises(ValueError):
@@ -74,7 +83,12 @@ def test_item():
             feature_names=[f"f_{i}" for i in range(5)],
         )
 
-        item = Item(inputs=inputs, outputs=outputs, forcing=forcing)
+        item = Item(
+            inputs=inputs,
+            outputs=outputs,
+            forcing=forcing,
+            validity_times=validity_times,
+        )
 
     # Input and Output must have the same dim names
     with pytest.raises(ValueError):
@@ -89,7 +103,12 @@ def test_item():
             feature_names=[f"feature_{i}" for i in range(5)],
         )
 
-        item = Item(inputs=inputs, outputs=outputs, forcing=forcing)
+        item = Item(
+            inputs=inputs,
+            outputs=outputs,
+            forcing=forcing,
+            validity_times=validity_times,
+        )
 
 
 def test_date_forcing():