feat: metadata for multi datasets (ecmwf#762)

Update the metadata we store in checkpoints in order to provide more
information to inference

The checkpoint mentioned
[here](ecmwf#594 (comment))
was created based on this branch.

***As a contributor to the Anemoi framework, please ensure that your
changes include unit tests, updates to any affected dependencies and
documentation, and have been tested in a parallel setting (i.e., with
multiple GPUs). As a reviewer, you are also responsible for verifying
these aspects and requesting changes if they are not adequately
addressed. For guidelines about those please refer to
https://anemoi.readthedocs.io/en/latest/***

By opening this pull request, I affirm that all authors agree to the
[Contributor License
Agreement.](https://github.com/ecmwf/codex/blob/main/Legal/contributor_license_agreement.md)

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 2 people

models/src/anemoi/models/interface/__init__.py

@@ -77,6 +77,7 @@ def __init__(
         self.supporting_arrays = supporting_arrays if supporting_arrays is not None else {}
         self.data_indices = data_indices
         self._build_model()
+        self._update_metadata()
 
     def _build_processors_for_dataset(
         self, dataset_name: str, statistics: dict, data_indices: dict, statistics_tendencies: dict = None
@@ -207,3 +208,6 @@ def predict_step(
 
         # Delegate to the model's predict_step implementation with processors
         return self.model.predict_step(**predict_kwargs, **kwargs)
+
+    def _update_metadata(self) -> None:
+        self.model.fill_metadata(self.metadata)

models/src/anemoi/models/models/base.py

@@ -312,3 +312,8 @@ def predict_step(
                     y_hat[dataset_name] = gather_tensor(y_hat[dataset_name], -2, y_hat_shard_shapes, model_comm_group)
 
         return y_hat
+
+    @abstractmethod
+    def fill_metadata(self, md_dict) -> None:
+        """To be implemented in subclasses to fill model-specific metadata."""
+        pass

models/src/anemoi/models/models/diffusion_encoder_processor_decoder.py

@@ -589,6 +589,24 @@ def sample(
             grid_shard_shapes=grid_shard_shapes,
         )
 
+    def fill_metadata(self, md_dict) -> None:
+        for dataset in self.input_dim.keys():
+            shapes = {
+                "variables": self.input_dim[dataset],
+                "input_timesteps": self.multi_step,
+                "ensemble": 1,
+                "grid": None,  # grid size is dynamic
+            }
+            md_dict["metadata_inference"][dataset]["shapes"] = shapes
+
+            rel_date_indices = md_dict["metadata_inference"][dataset]["timesteps"]["relative_date_indices_training"]
+            input_rel_date_indices = rel_date_indices[:-1]
+            output_rel_date_indices = rel_date_indices[-1]
+            md_dict["metadata_inference"][dataset]["timesteps"]["input_relative_date_indices"] = input_rel_date_indices
+            md_dict["metadata_inference"][dataset]["timesteps"][
+                "output_relative_date_indices"
+            ] = output_rel_date_indices
+
 
 class AnemoiDiffusionTendModelEncProcDec(AnemoiDiffusionModelEncProcDec):
     """Diffusion model for tendency prediction."""

models/src/anemoi/models/models/encoder_processor_decoder.py

@@ -273,3 +273,21 @@ def forward(
             )
 
         return x_out_dict
+
+    def fill_metadata(self, md_dict) -> None:
+        for dataset in self.input_dim.keys():
+            shapes = {
+                "variables": self.input_dim[dataset],
+                "input_timesteps": self.multi_step,
+                "ensemble": 1,
+                "grid": None,  # grid size is dynamic
+            }
+            md_dict["metadata_inference"][dataset]["shapes"] = shapes
+
+            rel_date_indices = md_dict["metadata_inference"][dataset]["timesteps"]["relative_date_indices_training"]
+            input_rel_date_indices = rel_date_indices[:-1]
+            output_rel_date_indices = rel_date_indices[-1]
+            md_dict["metadata_inference"][dataset]["timesteps"]["input_relative_date_indices"] = input_rel_date_indices
+            md_dict["metadata_inference"][dataset]["timesteps"][
+                "output_relative_date_indices"
+            ] = output_rel_date_indices