feat(training): downscaling

models/src/anemoi/models/data_indices/collection.py

@@ -27,15 +27,15 @@
 class IndexCollection:
     """Collection of data and model indices."""
 
-    def __init__(self, config, name_to_index) -> None:
-        self.config = OmegaConf.to_container(config, resolve=True)
+    def __init__(self, data_config, name_to_index) -> None:
+        self.config = OmegaConf.to_container(data_config, resolve=True)
         self.name_to_index = dict(sorted(name_to_index.items(), key=operator.itemgetter(1)))
-        self.forcing = [] if config.data.forcing is None else OmegaConf.to_container(config.data.forcing, resolve=True)
+        self.forcing = [] if data_config.forcing is None else OmegaConf.to_container(data_config.forcing, resolve=True)
         self.diagnostic = (
-            [] if config.data.diagnostic is None else OmegaConf.to_container(config.data.diagnostic, resolve=True)
+            [] if data_config.diagnostic is None else OmegaConf.to_container(data_config.diagnostic, resolve=True)
         )
         self.target = (
-            [] if config.data.get("target", None) is None else OmegaConf.to_container(config.data.target, resolve=True)
+            [] if data_config.get("target", None) is None else OmegaConf.to_container(data_config.target, resolve=True)
         )
         defined_variables = set.union(set(self.forcing), set(self.diagnostic), set(self.target))
         self.prognostic = [v for v in self.name_to_index.keys() if v not in defined_variables]
@@ -70,7 +70,7 @@ def __init__(self, config, name_to_index) -> None:
         )
 
     def __repr__(self) -> str:
-        return f"IndexCollection(config={self.config}, name_to_index={self.name_to_index})"
+        return f"IndexCollection(data_config={self.config}, name_to_index={self.name_to_index})"
 
     def __eq__(self, other):
         if not isinstance(other, IndexCollection):

models/src/anemoi/models/data_indices/tensor.py

@@ -90,6 +90,7 @@ def todict(self):
             "forcing": self.forcing,
             "target": self.target,
             "full": self.full,
+            "name_to_index": self.name_to_index,
         }
 
     @staticmethod

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

@@ -16,6 +16,7 @@
 from torch_geometric.data import HeteroData
 
 from anemoi.models.preprocessing import Processors
+from anemoi.models.utils.config import get_multiple_datasets_config
 from anemoi.utils.config import DotDict
 
 
@@ -75,28 +76,77 @@ 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_model(self) -> None:
-        """Builds the model and pre- and post-processors."""
-        # Instantiate processors
+    def _build_processors_for_dataset(
+        self,
+        processors_configs: dict,
+        statistics: dict,
+        data_indices: dict,
+        statistics_tendencies: dict = None,
+    ):
+        """Build processors for a single dataset.
+
+        Parameters
+        ----------
+        processors_configs : dict
+            Configuration for the processors
+        statistics : dict
+            Statistics for the dataset
+        data_indices : dict
+            Data indices for the dataset
+        statistics_tendencies : dict, optional
+            Tendencies statistics for the dataset
+
+        Returns
+        -------
+        tuple
+            (pre_processors, post_processors, pre_processors_tendencies, post_processors_tendencies)
+        """
+        # Build processors for the dataset
         processors = [
-            [name, instantiate(processor, data_indices=self.data_indices, statistics=self.statistics)]
-            for name, processor in self.config.data.processors.items()
+            [name, instantiate(processor, data_indices=data_indices, statistics=statistics)]
+            for name, processor in processors_configs.items()
         ]
 
-        # Assign the processor list pre- and post-processors
-        self.pre_processors = Processors(processors)
-        self.post_processors = Processors(processors, inverse=True)
+        pre_processors = Processors(processors)
+        post_processors = Processors(processors, inverse=True)
 
-        # If tendencies statistics are provided, instantiate the tendencies processors
-        if self.statistics_tendencies is not None:
-            processors = [
-                [name, instantiate(processor, data_indices=self.data_indices, statistics=self.statistics_tendencies)]
-                for name, processor in self.config.data.processors.items()
+        # Build tendencies processors if provided
+        pre_processors_tendencies = None
+        post_processors_tendencies = None
+        if statistics_tendencies is not None:
+            processors_tendencies = [
+                [name, instantiate(processor, data_indices=data_indices, statistics=statistics_tendencies)]
+                for name, processor in processors_configs.items()
             ]
-            # Assign the processor list pre- and post-processors
-            self.pre_processors_tendencies = Processors(processors)
-            self.post_processors_tendencies = Processors(processors, inverse=True)
+            pre_processors_tendencies = Processors(processors_tendencies)
+            post_processors_tendencies = Processors(processors_tendencies, inverse=True)
+
+        return pre_processors, post_processors, pre_processors_tendencies, post_processors_tendencies
+
+    def _build_model(self) -> None:
+        """Builds the model and pre- and post-processors."""
+        # Multi-dataset mode: create processors for each dataset
+        self.pre_processors = torch.nn.ModuleDict()
+        self.post_processors = torch.nn.ModuleDict()
+        self.pre_processors_tendencies = torch.nn.ModuleDict()
+        self.post_processors_tendencies = torch.nn.ModuleDict()
+
+        data_config = get_multiple_datasets_config(self.config.data)
+        for dataset_name in self.statistics.keys():
+            # Build processors for each dataset
+            pre, post, pre_tend, post_tend = self._build_processors_for_dataset(
+                data_config[dataset_name].processors,
+                self.statistics[dataset_name],
+                self.data_indices[dataset_name],
+                self.statistics_tendencies[dataset_name] if self.statistics_tendencies is not None else None,
+            )
+            self.pre_processors[dataset_name] = pre
+            self.post_processors[dataset_name] = post
+            if pre_tend is not None:
+                self.pre_processors_tendencies[dataset_name] = pre_tend
+                self.post_processors_tendencies[dataset_name] = post_tend
 
         # Instantiate the model
         # Only pass _target_ and _convert_ from model config to avoid passing diffusion as kwarg
@@ -117,13 +167,17 @@ def _build_model(self) -> None:
         self.forward = self.model.forward
 
     def predict_step(
-        self, batch: torch.Tensor, model_comm_group: Optional[ProcessGroup] = None, gather_out: bool = True, **kwargs
-    ) -> torch.Tensor:
+        self,
+        batch: dict[str, torch.Tensor],
+        model_comm_group: Optional[ProcessGroup] = None,
+        gather_out: bool = True,
+        **kwargs,
+    ) -> dict[str, torch.Tensor]:
         """Prediction step for the model.
 
         Parameters
         ----------
-        batch : torch.Tensor
+        batch : dict[str, torch.Tensor]
             Input batched data.
         model_comm_group : Optional[ProcessGroup], optional
             model communication group, specifies which GPUs work together
@@ -132,7 +186,7 @@ def predict_step(
 
         Returns
         -------
-        torch.Tensor
+        dict[str, torch.Tensor]
             Predicted data.
         """
         # Prepare kwargs for model's predict_step
@@ -152,3 +206,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/layers/bounding.py

@@ -306,7 +306,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
 
-def build_boundings(
+def _build_dataset_boundings(
     model_config: Any,
     data_indices: Any,
     statistics: dict | None,
@@ -355,3 +355,40 @@ def build_boundings(
             for cfg in bounding_cfgs
         ]
     )
+
+
+def build_boundings(
+    model_config: Any,
+    data_indices: Any,
+    statistics: dict | None,
+) -> nn.ModuleDict:
+    """Build the model-output bounding modules from configuration.
+
+    This is a thin factory that creates a ``nn.ModuleDict`` of bounding
+    modules by invoking ``_build_dataset_boundings`` for each dataset
+    specified in ``data_indices``.
+
+    Parameters
+    ----------
+    model_config : Any
+        Object with a ``model`` attribute containing an iterable ``bounding``
+        (e.g. a list of Hydra configs). If absent or empty, an empty
+        ``nn.ModuleDict`` is returned.
+    data_indices : Any
+        Dictionary mapping dataset names to data indices objects. Each
+        data indices object must provide the mappings:
+        ``data_indices.model.output.name_to_index`` and
+        ``data_indices.data.input.name_to_index``. These are forwarded to each
+        instantiated bounding module as ``name_to_index`` and
+        ``name_to_index_stats`` respectively.
+    statistics : dict | None
+        Dictionary mapping dataset names to optional dataset/model statistics
+        passed to each bounding module. Use ``None`` if not required by the
+        configured classes.
+    """
+    bounding_modules = nn.ModuleDict()
+    for dataset_name, dataset_indices in data_indices.items():
+        bounding_modules[dataset_name] = _build_dataset_boundings(
+            model_config, dataset_indices, statistics[dataset_name]
+        )
+    return bounding_modules

models/src/anemoi/models/migrations/scripts/1767108147_move_to_multiple_datasets.py

@@ -0,0 +1,93 @@
+# (C) Copyright 2025 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+from anemoi.models.migrations import CkptType
+from anemoi.models.migrations import MigrationMetadata
+
+# DO NOT CHANGE -->
+metadata = MigrationMetadata(
+    versions={
+        "migration": "1.0.0",
+        "anemoi-models": "%NEXT_ANEMOI_MODELS_VERSION%",
+    },
+)
+# <-- END DO NOT CHANGE
+
+
+def migrate(ckpt: CkptType) -> CkptType:
+    """Migrate the checkpoint.
+
+    Parameters
+    ----------
+    ckpt : CkptType
+        The checkpoint dict.
+
+    Returns
+    -------
+    CkptType
+        The migrated checkpoint dict.
+    """
+    dummy_dataset_name = "data"
+
+    updates = {}
+    for key in list(ckpt["state_dict"].keys()):
+        # Update pre-processors
+        if key.startswith("model.pre_processors."):
+            new_key = key.replace("model.pre_processors.", f"model.pre_processors.{dummy_dataset_name}.")
+            updates[new_key] = ckpt["state_dict"][key]
+            del ckpt["state_dict"][key]
+
+        # Update post-processors
+        if key.startswith("model.post_processors."):
+            new_key = key.replace("model.post_processors.", f"model.post_processors.{dummy_dataset_name}.")
+            updates[new_key] = ckpt["state_dict"][key]
+            del ckpt["state_dict"][key]
+
+        # Update node attributes
+        if key.startswith("model.model.node_attributes."):
+            new_key = key.replace("model.model.node_attributes.", f"model.model.node_attributes.{dummy_dataset_name}.")
+            updates[new_key] = ckpt["state_dict"][key]
+            del ckpt["state_dict"][key]
+
+        # Adjust model components
+        for model_component in ["encoder", "encoder_graph_provider", "decoder", "decoder_graph_provider"]:
+            prefix = f"model.model.{model_component}."
+
+            if key.startswith(prefix):
+                new_key = key.replace(prefix, f"{prefix}{dummy_dataset_name}.")
+                updates[new_key] = ckpt["state_dict"][key]
+                del ckpt["state_dict"][key]
+
+    ckpt["state_dict"].update(updates)
+
+    ckpt["hyper_parameters"]["data_indices"] = {dummy_dataset_name: ckpt["hyper_parameters"].pop("data_indices")}
+    ckpt["hyper_parameters"]["statistics"] = {dummy_dataset_name: ckpt["hyper_parameters"].pop("statistics")}
+    ckpt["hyper_parameters"]["statistics_tendencies"] = {
+        dummy_dataset_name: ckpt["hyper_parameters"].pop("statistics_tendencies")
+    }
+    ckpt["hyper_parameters"]["supporting_arrays"] = {
+        dummy_dataset_name: ckpt["hyper_parameters"].pop("supporting_arrays")
+    }
+    return ckpt
+
+
+def rollback(ckpt: CkptType) -> CkptType:
+    """Rollback the checkpoint.
+
+    Parameters
+    ----------
+    ckpt : CkptType
+        The checkpoint dict.
+
+    Returns
+    -------
+    CkptType
+        The rollbacked checkpoint dict.
+    """
+    return ckpt