Add separated channel module registry and step implementation

fme/ace/inference/inference.py

@@ -305,7 +305,7 @@ def run_inference_from_config(config: InferenceConfig):
         logging.info("Loading initial condition data")
         initial_condition = get_initial_condition(
             config.initial_condition.get_dataset(),
-            stepper_config.prognostic_names,
+            stepper_config.get_prognostic_names(),
             labels=config.labels,
             n_ensemble=config.n_ensemble_per_ic,
         )

fme/ace/step/fcn3.py

@@ -19,7 +19,6 @@
 from fme.core.distributed import Distributed
 from fme.core.normalizer import NetworkAndLossNormalizationConfig, StandardNormalizer
 from fme.core.ocean import Ocean, OceanConfig
-from fme.core.optimization import NullOptimization
 from fme.core.packer import Packer
 from fme.core.registry import CorrectorSelector
 from fme.core.step.args import StepArgs
@@ -229,7 +228,8 @@ def get_loss_normalizer(
             extra_residual_scaled_names = []
         return self.normalization.get_loss_normalizer(
             names=self._normalize_names + extra_names,
-            residual_scaled_names=self.prognostic_names + extra_residual_scaled_names,
+            residual_scaled_names=self.get_prognostic_names()
+            + extra_residual_scaled_names,
         )
 
     @classmethod
@@ -386,7 +386,6 @@ def __init__(
         self.module = dist.wrap_module(module)
         self._img_shape = dataset_info.img_shape
         self._config = config
-        self._no_optimization = NullOptimization()
 
         self._timestep = timestep
 
@@ -483,7 +482,7 @@ def network_call(input_norm: TensorDict) -> TensorDict:
             corrector=self._corrector,
             ocean=self.ocean,
             residual_prediction=self._config.residual_prediction,
-            prognostic_names=self.prognostic_names,
+            prognostic_names=self.get_prognostic_names(),
             prescribed_prognostic_names=self._config.prescribed_prognostic_names,
         )
 

fme/ace/stepper/single_module.py

@@ -525,7 +525,7 @@ def get_evaluation_window_data_requirements(
 
     def get_prognostic_state_data_requirements(self) -> PrognosticStateDataRequirements:
         return PrognosticStateDataRequirements(
-            names=self.prognostic_names,
+            names=self.get_prognostic_names(),
             n_timesteps=self.n_ic_timesteps,
         )
 
@@ -672,10 +672,9 @@ def next_step_forcing_names(self) -> list[str]:
         """
         return self.step.get_next_step_forcing_names()
 
-    @property
-    def prognostic_names(self) -> list[str]:
+    def get_prognostic_names(self) -> list[str]:
         """Names of variables which both inputs and outputs."""
-        return self.step.prognostic_names
+        return self.step.get_prognostic_names()
 
     @property
     def output_names(self) -> list[str]:
@@ -971,9 +970,8 @@ def get_base_weights(self) -> Weights | None:
         """
         return self._parameter_initializer.base_weights
 
-    @property
-    def prognostic_names(self) -> list[str]:
-        return self._step_obj.prognostic_names
+    def get_prognostic_names(self) -> list[str]:
+        return self._step_obj.get_prognostic_names()
 
     @property
     def out_names(self) -> list[str]:
@@ -1174,7 +1172,9 @@ def predict(
                     )
                     .remove_initial_condition(self.n_ic_timesteps)
                 )
-        prognostic_state = data.get_end(self.prognostic_names, self.n_ic_timesteps)
+        prognostic_state = data.get_end(
+            self.get_prognostic_names(), self.n_ic_timesteps
+        )
         data = BatchData.new_on_device(
             data=data.data,
             time=data.time,
@@ -1494,7 +1494,7 @@ def __init__(
 
         self._epoch: int | None = None  # to keep track of cached values
 
-        self._prognostic_names = self._stepper.prognostic_names
+        self._prognostic_names = self._stepper.get_prognostic_names()
         self._derive_func = self._stepper.derive_func
         self._loss_obj = self._stepper.build_loss(config.loss)
 

fme/core/generics/test_looper.py

@@ -173,7 +173,7 @@ def test_looper():
         data={n: spherical_data.data[n][:, :1] for n in spherical_data.data},
         time=time[:, 0:1],
     ).get_start(
-        prognostic_names=stepper.prognostic_names,
+        prognostic_names=stepper.get_prognostic_names(),
         n_ic_timesteps=1,
     )
     loader = MockLoader(shape, forcing_names, 3, time=time)
@@ -197,7 +197,7 @@ def test_looper_paired():
         data={n: spherical_data.data[n][:, :1] for n in spherical_data.data},
         time=time[:, 0:1],
     ).get_start(
-        prognostic_names=stepper.prognostic_names,
+        prognostic_names=stepper.get_prognostic_names(),
         n_ic_timesteps=1,
     )
     loader = MockLoader(shape, forcing_names, 3, time=time)
@@ -235,7 +235,7 @@ def test_looper_paired_with_derived_variables():
         data={n: spherical_data.data[n][:, :1] for n in spherical_data.data},
         time=time[:, 0:1],
     ).get_start(
-        prognostic_names=stepper.prognostic_names,
+        prognostic_names=stepper.get_prognostic_names(),
         n_ic_timesteps=1,
     )
     loader = MockLoader(shape, forcing_names, 2, time=time)
@@ -258,7 +258,7 @@ def test_looper_paired_with_target_data():
         data={n: spherical_data.data[n][:, :1] for n in spherical_data.data},
         time=time[:, 0:1],
     ).get_start(
-        prognostic_names=stepper.prognostic_names,
+        prognostic_names=stepper.get_prognostic_names(),
         n_ic_timesteps=1,
     )
     loader = MockLoader(shape, all_names, 2, time=time)
@@ -284,7 +284,7 @@ def test_looper_paired_with_target_data_and_derived_variables():
         data={n: spherical_data.data[n][:, :1] for n in spherical_data.data},
         time=time[:, 0:1],
     ).get_start(
-        prognostic_names=stepper.prognostic_names,
+        prognostic_names=stepper.get_prognostic_names(),
         n_ic_timesteps=1,
     )
     loader = MockLoader(shape, all_names, 2, time=time)

fme/core/registry/__init__.py

@@ -1,3 +1,4 @@
 from .corrector import CorrectorSelector
 from .module import ModuleSelector
 from .registry import Registry
+from .separated_module import SeparatedModuleSelector

fme/core/registry/separated_module.py

@@ -0,0 +1,204 @@
+import abc
+import dataclasses
+from collections.abc import Callable, Mapping
+from typing import Any, ClassVar
+
+import dacite
+import torch
+from torch import nn
+
+from fme.core.dataset_info import DatasetInfo
+from fme.core.labels import BatchLabels, LabelEncoding
+
+from .registry import Registry
+
+SeparatedModuleConfigType = type["SeparatedModuleConfig"]
+
+
+@dataclasses.dataclass
+class SeparatedModuleConfig(abc.ABC):
+    """
+    Builds an nn.Module that takes separate forcing and prognostic tensors
+    as input, and returns separate prognostic and diagnostic tensors as output.
+
+    The built nn.Module must have the forward signature::
+
+        forward(
+            forcing: Tensor,
+            prognostic: Tensor,
+            labels: Tensor | None = None,
+        ) -> tuple[Tensor, Tensor]
+
+    where the return value is (prognostic_out, diagnostic_out).
+    """
+
+    @abc.abstractmethod
+    def build(
+        self,
+        n_forcing_channels: int,
+        n_prognostic_channels: int,
+        n_diagnostic_channels: int,
+        dataset_info: DatasetInfo,
+    ) -> nn.Module:
+        """
+        Build a nn.Module with separated forcing/prognostic/diagnostic channels.
+
+        Args:
+            n_forcing_channels: number of input-only (forcing) channels
+            n_prognostic_channels: number of input-output (prognostic) channels
+            n_diagnostic_channels: number of output-only (diagnostic) channels
+            dataset_info: Information about the dataset, including img_shape,
+                horizontal coordinates, vertical coordinate, etc.
+
+        Returns:
+            An nn.Module whose forward method takes
+            (forcing, prognostic, labels=None) and returns
+            (prognostic_out, diagnostic_out) tensors.
+        """
+        ...
+
+    @classmethod
+    def from_state(cls, state: Mapping[str, Any]) -> "SeparatedModuleConfig":
+        return dacite.from_dict(
+            data_class=cls, data=state, config=dacite.Config(strict=True)
+        )
+
+
+class SeparatedModule:
+    """
+    Wrapper around an nn.Module with separated channel interface.
+
+    The wrapped module takes (forcing, prognostic) tensors and returns
+    (prognostic_out, diagnostic_out) tensors. This wrapper handles
+    optional label encoding for conditional models.
+    """
+
+    def __init__(self, module: nn.Module, label_encoding: LabelEncoding | None):
+        self._module = module
+        self._label_encoding = label_encoding
+
+    def __call__(
+        self,
+        forcing: torch.Tensor,
+        prognostic: torch.Tensor,
+        labels: BatchLabels | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if labels is not None and self._label_encoding is None:
+            raise ValueError(
+                "labels were provided but the module has no label encoding"
+            )
+        if labels is None and self._label_encoding is not None:
+            raise ValueError(
+                "labels were not provided but the module has a label encoding"
+            )
+        if labels is not None and self._label_encoding is not None:
+            encoded_labels = labels.conform_to_encoding(self._label_encoding)
+            return self._module(forcing, prognostic, labels=encoded_labels.tensor)
+        else:
+            return self._module(forcing, prognostic, labels=None)
+
+    @property
+    def torch_module(self) -> nn.Module:
+        return self._module
+
+    def get_state(self) -> dict[str, Any]:
+        if self._label_encoding is not None:
+            label_encoder_state = self._label_encoding.get_state()
+        else:
+            label_encoder_state = None
+        return {
+            **self._module.state_dict(),
+            "label_encoding": label_encoder_state,
+        }
+
+    def load_state(self, state: dict[str, Any]) -> None:
+        state = state.copy()
+        if state.get("label_encoding") is not None:
+            if self._label_encoding is None:
+                self._label_encoding = LabelEncoding.from_state(
+                    state.pop("label_encoding")
+                )
+            else:
+                self._label_encoding.conform_to_state(state.pop("label_encoding"))
+        state.pop("label_encoding", None)
+        self._module.load_state_dict(state)
+
+    def wrap_module(
+        self, callable: Callable[[nn.Module], nn.Module]
+    ) -> "SeparatedModule":
+        return SeparatedModule(callable(self._module), self._label_encoding)
+
+    def to(self, device: torch.device) -> "SeparatedModule":
+        return SeparatedModule(self._module.to(device), self._label_encoding)
+
+
+@dataclasses.dataclass
+class SeparatedModuleSelector:
+    """
+    A dataclass for building SeparatedModuleConfig instances from config dicts.
+
+    Mirrors ModuleSelector but for the separated channel interface.
+
+    Parameters:
+        type: the type of the SeparatedModuleConfig
+        config: data for a SeparatedModuleConfig instance of the indicated type
+    """
+
+    type: str
+    config: Mapping[str, Any]
+    registry: ClassVar[Registry[SeparatedModuleConfig]] = Registry[
+        SeparatedModuleConfig
+    ]()
+
+    def __post_init__(self):
+        if not isinstance(self.registry, Registry):
+            raise ValueError(
+                "SeparatedModuleSelector.registry should not be set manually"
+            )
+        self._instance = self.registry.get(self.type, self.config)
+
+    @property
+    def module_config(self) -> SeparatedModuleConfig:
+        return self._instance
+
+    @classmethod
+    def register(
+        cls, type_name: str
+    ) -> Callable[[SeparatedModuleConfigType], SeparatedModuleConfigType]:
+        return cls.registry.register(type_name)
+
+    def build(
+        self,
+        n_forcing_channels: int,
+        n_prognostic_channels: int,
+        n_diagnostic_channels: int,
+        dataset_info: DatasetInfo,
+    ) -> SeparatedModule:
+        """
+        Build a SeparatedModule with separated forcing/prognostic/diagnostic
+        channels.
+
+        Args:
+            n_forcing_channels: number of input-only (forcing) channels
+            n_prognostic_channels: number of input-output (prognostic) channels
+            n_diagnostic_channels: number of output-only (diagnostic) channels
+            dataset_info: Information about the dataset, including img_shape.
+
+        Returns:
+            a SeparatedModule object
+        """
+        if len(dataset_info.all_labels) > 0:
+            label_encoding = LabelEncoding(sorted(list(dataset_info.all_labels)))
+        else:
+            label_encoding = None
+        module = self._instance.build(
+            n_forcing_channels=n_forcing_channels,
+            n_prognostic_channels=n_prognostic_channels,
+            n_diagnostic_channels=n_diagnostic_channels,
+            dataset_info=dataset_info,
+        )
+        return SeparatedModule(module, label_encoding)
+
+    @classmethod
+    def get_available_types(cls):
+        return cls.registry._types.keys()