feat(training): decouple training method & task

models/docs/introduction/overview.rst

@@ -68,8 +68,6 @@ The package currently includes the following model architectures:
    base architecture
 -  **AnemoiEnsModelEncProcDec**: The CRPS-optimized ensemble version
    that injects noise in the processor, e.g. AIFS-CRPS
--  **AnemoiModelEncProcDecMultiOutInterpolator**: A specialized architecture for
-   time interpolation
 
 All models support flexible layer kernel configuration, allowing for
 customization of linear and normalization layers in different parts of

models/docs/modules/models.rst

@@ -69,19 +69,3 @@ the depth of the hierarchy, enabling efficient handling of data.
    :members:
    :no-undoc-members:
    :show-inheritance:
-
-*************************
- Time Interpolator Model
-*************************
-
-A specialized architecture for time interpolation tasks.
-
-Key features:
-
-   #. Ability to select time indices for forcing and predictions
-   #. Allows for provision of t0 and t6 and predictions of t1->5
-
-.. autoclass:: anemoi.models.models.interpolator.AnemoiModelEncProcDecMultiOutInterpolator
-   :members:
-   :no-undoc-members:
-   :show-inheritance:

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

@@ -59,6 +59,8 @@ def __init__(
         self,
         *,
         config: DotDict,
+        n_step_input: int,
+        n_step_output: int,
         graph_data: HeteroData,
         statistics: dict,
         data_indices: dict,
@@ -69,7 +71,8 @@ def __init__(
         super().__init__()
         self.config = config
         self.id = str(uuid.uuid4())
-        self.n_step_input = self.config.training.multistep_input
+        self.n_step_input = n_step_input
+        self.n_step_output = n_step_output
         self.graph_data = graph_data
         self.statistics = statistics
         self.statistics_tendencies = statistics_tendencies
@@ -141,8 +144,7 @@ def _build_tendency_processors(
             return self._build_processor_pair(processors_configs, data_indices, statistics_tendencies)
 
         lead_times = list(statistics_tendencies.get("lead_times") or [])
-        n_step_output = getattr(self.config.training, "multistep_output", None)
-        if n_step_output == 1:
+        if self.n_step_output == 1:
             step_stats = statistics_tendencies.get(lead_times[0]) if lead_times else None
             stats_for_tendencies = step_stats or statistics_tendencies
             return self._build_processor_pair(processors_configs, data_indices, stats_for_tendencies)
@@ -193,6 +195,8 @@ def _build_model(self) -> None:
             data_indices=self.data_indices,
             statistics=self.statistics,
             graph_data=self.graph_data,
+            n_step_input=self.n_step_input,
+            n_step_output=self.n_step_output,
             _recursive_=False,  # Disables recursive instantiation by Hydra
         )
 

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

@@ -11,7 +11,6 @@
 from .encoder_processor_decoder import AnemoiModelEncProcDec
 from .ens_encoder_processor_decoder import AnemoiEnsModelEncProcDec
 from .hierarchical import AnemoiModelEncProcDecHierarchical
-from .interpolator import AnemoiModelEncProcDecMultiOutInterpolator
 
 from .autoencoder import AnemoiModelAutoEncoder  # isort: skip
 from .hierarchical_autoencoder import AnemoiModelHierarchicalAutoEncoder  # isort: skip
@@ -28,5 +27,4 @@
     "AnemoiEnsModelEncProcDec",
     "AnemoiModelAutoEncoder",
     "AnemoiModelHierarchicalAutoEncoder",
-    "AnemoiModelEncProcDecMultiOutInterpolator",
 ]

models/src/anemoi/models/models/autoencoder.py

@@ -33,50 +33,6 @@ def _calculate_target_dim(self, dataset_name: str) -> int:
             + self.node_attributes.attr_ndims[dataset_name]
         )
 
-    def _assemble_input(self, x, batch_size, grid_shard_shapes=None, model_comm_group=None, dataset_name=None):
-        assert dataset_name is not None, "dataset_name must be provided when using multiple datasets."
-        node_attributes_data = self.node_attributes(dataset_name, batch_size=batch_size)
-        grid_shard_shapes = grid_shard_shapes[dataset_name] if grid_shard_shapes is not None else None
-        if grid_shard_shapes is not None:
-            shard_shapes_nodes = get_or_apply_shard_shapes(
-                node_attributes_data, 0, shard_shapes_dim=grid_shard_shapes, model_comm_group=model_comm_group
-            )
-            node_attributes_data = shard_tensor(node_attributes_data, 0, shard_shapes_nodes, model_comm_group)
-
-        x_input = x[:, : self.n_step_input, ...]
-        # normalize and add data positional info (lat/lon)
-        x_data_latent = torch.cat(
-            (
-                einops.rearrange(x_input, "batch time ensemble grid vars -> (batch ensemble grid) (time vars)"),
-                node_attributes_data,
-            ),
-            dim=-1,  # feature dimension
-        )
-        shard_shapes_data = get_or_apply_shard_shapes(
-            x_data_latent, 0, shard_shapes_dim=grid_shard_shapes, model_comm_group=model_comm_group
-        )
-
-        return x_data_latent, shard_shapes_data
-
-    def _assemble_output(self, x_out, batch_size, ensemble_size, dtype, dataset_name=None):
-        x_out = (
-            einops.rearrange(
-                x_out,
-                "(batch ensemble grid) (time vars) -> batch time ensemble grid vars",
-                batch=batch_size,
-                ensemble=ensemble_size,
-                time=self.n_step_output,
-            )
-            .to(dtype=dtype)
-            .clone()
-        )
-        assert dataset_name is not None, "dataset_name must be provided for multi-dataset case"
-
-        for bounding in self.boundings[dataset_name]:
-            # bounding performed in the order specified in the config file
-            x_out = bounding(x_out)
-        return x_out
-
     def _assemble_forcings(self, x, batch_size, grid_shard_shapes=None, model_comm_group=None, dataset_name=None):
         assert dataset_name is not None, "dataset_name must be provided when using multiple datasets."
         node_attributes_target = self.node_attributes(dataset_name, batch_size=batch_size)

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

@@ -40,6 +40,8 @@ def __init__(
         model_config: DotDict,
         data_indices: dict,
         statistics: dict,
+        n_step_input: int,
+        n_step_output: int,
         graph_data: HeteroData,
     ) -> None:
         """Initializes the graph neural network.
@@ -59,14 +61,14 @@ def __init__(
         self._graph_data = graph_data
         self.data_indices = data_indices
         self.statistics = statistics
+        self.n_step_input = n_step_input
+        self.n_step_output = n_step_output
 
         self.dataset_names = list(data_indices.keys())
 
         model_config = DotDict(model_config)
         self._graph_name_hidden = model_config.model.model.hidden_nodes_name
 
-        self.n_step_input = model_config.training.multistep_input
-        self.n_step_output = model_config.training.multistep_output
         self.num_channels = model_config.model.num_channels
         self.latent_skip = model_config.model.model.latent_skip
 

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

@@ -44,6 +44,8 @@ def __init__(
         model_config: DotDict,
         data_indices: dict,
         statistics: dict,
+        n_step_input: int,
+        n_step_output: int,
         graph_data: HeteroData,
     ) -> None:
 
@@ -62,6 +64,8 @@ def __init__(
             data_indices=data_indices,
             statistics=statistics,
             graph_data=graph_data,
+            n_step_input=n_step_input,
+            n_step_output=n_step_output,
         )
 
         self.noise_embedder = instantiate(diffusion_config.noise_embedder)
@@ -756,14 +760,6 @@ def fill_metadata(self, md_dict) -> None:
             }
             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[: self.n_step_input]
-            output_rel_date_indices = rel_date_indices[-self.n_step_output :]
-            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."""
@@ -774,6 +770,8 @@ def __init__(
         model_config: DotDict,
         data_indices: dict,
         statistics: dict,
+        n_step_input: int,
+        n_step_output: int,
         graph_data: HeteroData,
     ) -> None:
         model_config_local = DotDict(model_config)
@@ -783,6 +781,8 @@ def __init__(
             model_config=model_config,
             data_indices=data_indices,
             statistics=statistics,
+            n_step_input=n_step_input,
+            n_step_output=n_step_output,
             graph_data=graph_data,
         )
 

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

@@ -324,11 +324,3 @@ def fill_metadata(self, md_dict) -> None:
                 "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[: self.n_step_input]
-            output_rel_date_indices = rel_date_indices[-self.n_step_output :]
-            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

models/src/anemoi/models/models/ens_encoder_processor_decoder.py

@@ -36,13 +36,17 @@ def __init__(
         data_indices: dict,
         statistics: dict,
         graph_data: HeteroData,
+        n_step_input: int,
+        n_step_output: int,
     ) -> None:
         self.condition_on_residual = DotDict(model_config).model.condition_on_residual
         super().__init__(
             model_config=model_config,
             data_indices=data_indices,
             statistics=statistics,
             graph_data=graph_data,
+            n_step_input=n_step_input,
+            n_step_output=n_step_output,
         )
 
     def _build_networks(self, model_config):

models/src/anemoi/models/models/hierarchical_autoencoder.py

@@ -32,6 +32,8 @@ def __init__(
         data_indices: dict,
         statistics: dict,
         graph_data: HeteroData,
+        n_step_input: int,
+        n_step_output: int,
     ) -> None:
         """Initializes the graph neural network.
 
@@ -55,8 +57,8 @@ def __init__(
         model_config = DotDict(model_config)
         self._graph_name_hidden = model_config.model.model.hidden_nodes_name
 
-        self.n_step_input = model_config.training.multistep_input
-        self.n_step_output = model_config.training.multistep_output
+        self.n_step_input = n_step_input
+        self.n_step_output = n_step_output
 
         self.num_channels = model_config.model.num_channels