Add CorrectorConfigABC (ai2cm#965)

Refactors corrector config typing with `CorrectorConfigABC` with a
`get_corrector` abstract method which accepts `DatasetInfo` and returns
a `CorrectorABC`. This allows the corrector configs to build their
corresponding correctors, removing that responsibility from the vertical
coordinates.

Changes:
- `CorrectorSelector.registry` was previously unused but is now once
again active via `CorrectorSelector.__post_init__`.
- Adds properties `atmosphere_vertical_coordinate` and
`ocean_vertical_coordinate` to `DatasetInfo`.

- [x] Tests added

Author: jpdunc23

fme/ace/step/fcn3.py

@@ -314,11 +314,7 @@ def get_step(
         init_weights: Callable[[list[nn.Module]], None],
     ) -> "FCN3Step":
         logging.info("Initializing stepper from provided config")
-        corrector = dataset_info.vertical_coordinate.build_corrector(
-            config=self.corrector,
-            gridded_operations=dataset_info.gridded_operations,
-            timestep=dataset_info.timestep,
-        )
+        corrector = self.corrector.get_corrector(dataset_info)
         normalizer = self.normalization.get_network_normalizer(self._normalize_names)
         return FCN3Step(
             config=self,

fme/core/atmosphere_data.py

@@ -1,5 +1,5 @@
 from collections.abc import Mapping
-from typing import Protocol
+from typing import Protocol, runtime_checkable
 
 import torch
 
@@ -40,6 +40,7 @@
 }
 
 
+@runtime_checkable
 class HasAtmosphereVerticalIntegral(Protocol):
     def vertical_integral(
         self,

fme/core/coordinates.py

@@ -11,17 +11,12 @@
 
 from fme.core import metrics
 from fme.core.constants import EARTH_RADIUS, GRAVITY
-from fme.core.corrector.atmosphere import AtmosphereCorrector, AtmosphereCorrectorConfig
-from fme.core.corrector.ice import IceCorrector, IceCorrectorConfig
-from fme.core.corrector.ocean import OceanCorrector, OceanCorrectorConfig
-from fme.core.corrector.registry import CorrectorABC
 from fme.core.derived_variables import compute_derived_quantities
 from fme.core.device import get_device
 from fme.core.distributed import Distributed
 from fme.core.gridded_ops import GriddedOperations, HEALPixOperations, LatLonOperations
 from fme.core.mask_provider import MaskProvider, MaskProviderABC, NullMaskProvider
 from fme.core.ocean_derived_variables import compute_ocean_derived_quantities
-from fme.core.registry.corrector import CorrectorSelector
 from fme.core.typing_ import TensorDict, TensorMapping
 from fme.core.winds import lon_lat_to_xyz
 
@@ -136,15 +131,6 @@ def __repr__(self) -> str:
     def __eq__(self, other) -> bool:
         pass
 
-    @abc.abstractmethod
-    def build_corrector(
-        self,
-        config: AtmosphereCorrectorConfig | CorrectorSelector,
-        gridded_operations: GriddedOperations,
-        timestep: timedelta,
-    ) -> CorrectorABC:
-        pass
-
     @abc.abstractmethod
     def build_derive_function(
         self,
@@ -208,35 +194,6 @@ def __len__(self):
         """The number of vertical layer interfaces."""
         return len(self.ak)
 
-    def build_corrector(
-        self,
-        config: AtmosphereCorrectorConfig | CorrectorSelector,
-        gridded_operations: GriddedOperations,
-        timestep: timedelta,
-    ) -> AtmosphereCorrector:
-        if (
-            isinstance(config, CorrectorSelector)
-            and config.type != "atmosphere_corrector"
-        ):
-            raise ValueError(
-                f"Cannot build corrector for vertical coordinate {self} with "
-                f"corrector selector {config}."
-            )
-        if isinstance(config, CorrectorSelector):
-            config_instance = dacite.from_dict(
-                data_class=AtmosphereCorrectorConfig,
-                data=config.config,
-                config=dacite.Config(strict=True),
-            )
-        else:
-            config_instance = config
-        return AtmosphereCorrector(
-            config=config_instance,
-            gridded_operations=gridded_operations,
-            vertical_coordinate=self,
-            timestep=timestep,
-        )
-
     def build_derive_function(
         self,
         timestep: timedelta,
@@ -392,30 +349,6 @@ def __len__(self):
         """The number of vertical layer interfaces."""
         return len(self.idepth)
 
-    def build_corrector(
-        self,
-        config: AtmosphereCorrectorConfig | CorrectorSelector,
-        gridded_operations: GriddedOperations,
-        timestep: timedelta,
-    ) -> OceanCorrector:
-        if isinstance(config, AtmosphereCorrectorConfig):
-            raise ValueError(
-                "Cannot build corrector for depth coordinate with an "
-                "AtmosphereCorrectorConfig."
-            )
-        elif config.type != "ocean_corrector":
-            raise ValueError(
-                f"Cannot build corrector for vertical coordinate {self} with "
-                f"corrector selector {config}."
-            )
-        config_instance = OceanCorrectorConfig.from_state(config.config)
-        return OceanCorrector(
-            config=config_instance,
-            gridded_operations=gridded_operations,
-            vertical_coordinate=self,
-            timestep=timestep,
-        )
-
     def build_derive_function(
         self,
         timestep: timedelta,
@@ -521,56 +454,6 @@ def __repr__(self) -> str:
     def __len__(self) -> int:
         return 0
 
-    def build_corrector(
-        self,
-        config: AtmosphereCorrectorConfig | CorrectorSelector,
-        gridded_operations: GriddedOperations,
-        timestep: timedelta,
-    ) -> CorrectorABC:
-        if isinstance(config, AtmosphereCorrectorConfig):
-            return AtmosphereCorrector(
-                config=config,
-                gridded_operations=gridded_operations,
-                vertical_coordinate=None,
-                timestep=timestep,
-            )
-        if config.type == "atmosphere_corrector":
-            config_instance = dacite.from_dict(
-                data_class=AtmosphereCorrectorConfig,
-                data=config.config,
-                config=dacite.Config(strict=True),
-            )
-            return AtmosphereCorrector(
-                config=config_instance,
-                gridded_operations=gridded_operations,
-                vertical_coordinate=None,
-                timestep=timestep,
-            )
-        elif config.type == "ocean_corrector":
-            config_instance = OceanCorrectorConfig.from_state(config.config)
-            return OceanCorrector(
-                config=config_instance,
-                gridded_operations=gridded_operations,
-                vertical_coordinate=None,
-                timestep=timestep,
-            )
-        elif config.type == "ice_corrector":
-            config_instance = dacite.from_dict(
-                data_class=IceCorrectorConfig,
-                data=config.config,
-                config=dacite.Config(strict=True),
-            )
-            return IceCorrector(
-                config=config_instance,
-                gridded_operations=gridded_operations,
-                timestep=timestep,
-            )
-        else:
-            raise ValueError(
-                f"Invalid corrector type: {config.type}. "
-                "Must be either 'atmosphere_corrector' or 'ocean_corrector'."
-            )
-
     def build_derive_function(
         self,
         timestep: timedelta,

fme/core/corrector/atmosphere.py

@@ -13,8 +13,9 @@
     compute_layer_thickness,
 )
 from fme.core.constants import GRAVITY, SPECIFIC_HEAT_OF_DRY_AIR_CONST_VOLUME
-from fme.core.corrector.registry import CorrectorABC
+from fme.core.corrector.registry import CorrectorABC, CorrectorConfigABC
 from fme.core.corrector.utils import force_positive
+from fme.core.dataset_info import DatasetInfo
 from fme.core.gridded_ops import GriddedOperations
 from fme.core.registry.corrector import CorrectorSelector
 from fme.core.typing_ import TensorDict, TensorMapping
@@ -40,7 +41,7 @@ class EnergyBudgetConfig:
 
 @CorrectorSelector.register("atmosphere_corrector")
 @dataclasses.dataclass
-class AtmosphereCorrectorConfig:
+class AtmosphereCorrectorConfig(CorrectorConfigABC):
     r"""
     Configuration for the post-step state corrector.
 
@@ -138,6 +139,17 @@ def from_state(cls, state: Mapping[str, Any]) -> "AtmosphereCorrectorConfig":
             data_class=cls, data=state, config=dacite.Config(strict=True)
         )
 
+    def get_corrector(
+        self,
+        dataset_info: DatasetInfo,
+    ) -> "AtmosphereCorrector":
+        return AtmosphereCorrector(
+            self,
+            dataset_info.gridded_operations,
+            dataset_info.atmosphere_vertical_coordinate,
+            dataset_info.timestep,
+        )
+
 
 class AtmosphereCorrector(CorrectorABC):
     def __init__(

fme/core/corrector/ice.py

@@ -6,7 +6,8 @@
 import dacite
 import torch
 
-from fme.core.corrector.registry import CorrectorABC
+from fme.core.corrector.registry import CorrectorABC, CorrectorConfigABC
+from fme.core.dataset_info import DatasetInfo
 from fme.core.gridded_ops import GriddedOperations
 from fme.core.registry.corrector import CorrectorSelector
 from fme.core.typing_ import TensorDict, TensorMapping
@@ -185,8 +186,7 @@ def __call__(
 
 @CorrectorSelector.register("ice_corrector")
 @dataclasses.dataclass
-class IceCorrectorConfig:
-    # Correctors here. Can add more as needed
+class IceCorrectorConfig(CorrectorConfigABC):
     budget_correction: IceBudgetCorrectionConfig | None = None
 
     @classmethod
@@ -195,6 +195,16 @@ def from_state(cls, state: Mapping[str, Any]) -> "IceCorrectorConfig":
             data_class=cls, data=state, config=dacite.Config(strict=True)
         )
 
+    def get_corrector(
+        self,
+        dataset_info: DatasetInfo,
+    ) -> "IceCorrector":
+        return IceCorrector(
+            self,
+            dataset_info.gridded_operations,
+            dataset_info.timestep,
+        )
+
 
 class IceCorrector(CorrectorABC):
     """

fme/core/corrector/ocean.py

@@ -12,8 +12,9 @@
     LATENT_HEAT_OF_VAPORIZATION,
     SPECIFIC_HEAT_OF_SEA_WATER_CM4,
 )
-from fme.core.corrector.registry import CorrectorABC
+from fme.core.corrector.registry import CorrectorABC, CorrectorConfigABC
 from fme.core.corrector.utils import force_positive
+from fme.core.dataset_info import DatasetInfo
 from fme.core.gridded_ops import GriddedOperations
 from fme.core.ocean_data import HasOceanDepthIntegral, OceanData
 from fme.core.registry.corrector import CorrectorSelector
@@ -110,7 +111,7 @@ class SurfaceEnergyFluxCorrectionConfig:
 
 @CorrectorSelector.register("ocean_corrector")
 @dataclasses.dataclass
-class OceanCorrectorConfig:
+class OceanCorrectorConfig(CorrectorConfigABC):
     force_positive_names: list[str] = dataclasses.field(default_factory=list)
     sea_ice_fraction_correction: SeaIceFractionConfig | None = None
     surface_energy_flux_correction: SurfaceEnergyFluxCorrectionConfig | None = None
@@ -147,6 +148,17 @@ def remove_deprecated_keys(cls, state: Mapping[str, Any]) -> dict[str, Any]:
                     )
         return state_copy
 
+    def get_corrector(
+        self,
+        dataset_info: DatasetInfo,
+    ) -> "OceanCorrector":
+        return OceanCorrector(
+            self,
+            dataset_info.gridded_operations,
+            dataset_info.ocean_vertical_coordinate,
+            dataset_info.timestep,
+        )
+
 
 class OceanCorrector(CorrectorABC):
     def __init__(

fme/core/corrector/registry.py

@@ -1,8 +1,17 @@
 import abc
 
+from fme.core.dataset_info import DatasetInfo
 from fme.core.typing_ import TensorDict, TensorMapping
 
 
+class CorrectorConfigABC(abc.ABC):
+    @abc.abstractmethod
+    def get_corrector(
+        self,
+        dataset_info: DatasetInfo,
+    ) -> "CorrectorABC": ...
+
+
 class CorrectorABC(abc.ABC):
     @abc.abstractmethod
     def __call__(

fme/core/dataset_info.py

@@ -3,6 +3,7 @@
 from collections.abc import Mapping
 from typing import Any
 
+from fme.core.atmosphere_data import HasAtmosphereVerticalIntegral
 from fme.core.coordinates import (
     HorizontalCoordinates,
     NullVerticalCoordinate,
@@ -14,6 +15,7 @@
 from fme.core.dataset.utils import decode_timestep, encode_timestep
 from fme.core.gridded_ops import GriddedOperations
 from fme.core.mask_provider import MaskProvider, MaskProviderABC, NullMaskProvider
+from fme.core.ocean_data import HasOceanDepthIntegral
 
 
 class MissingDatasetInfo(ValueError):
@@ -181,6 +183,28 @@ def vertical_coordinate(self) -> VerticalCoordinate:
             raise MissingDatasetInfo("vertical_coordinate")
         return self._vertical_coordinate
 
+    @property
+    def atmosphere_vertical_coordinate(self) -> HasAtmosphereVerticalIntegral | None:
+        if isinstance(self._vertical_coordinate, HasAtmosphereVerticalIntegral):
+            return self._vertical_coordinate
+        elif isinstance(self._vertical_coordinate, NullVerticalCoordinate):
+            return None
+        raise RuntimeError(
+            f"{self._vertical_coordinate} cannot be used as an atmosphere vertical "
+            "coordinate."
+        )
+
+    @property
+    def ocean_vertical_coordinate(self) -> HasOceanDepthIntegral | None:
+        if isinstance(self._vertical_coordinate, HasOceanDepthIntegral):
+            return self._vertical_coordinate
+        elif isinstance(self._vertical_coordinate, NullVerticalCoordinate):
+            return None
+        raise RuntimeError(
+            f"{self._vertical_coordinate} cannot be used as an ocean vertical "
+            "coordinate."
+        )
+
     @property
     def mask_provider(self) -> MaskProvider:
         if self._mask_provider is None:

fme/core/ocean_data.py

@@ -1,7 +1,7 @@
 import math
 from collections.abc import Mapping
 from types import MappingProxyType
-from typing import Protocol
+from typing import Protocol, runtime_checkable
 
 import torch
 
@@ -30,6 +30,7 @@
 )
 
 
+@runtime_checkable
 class HasOceanDepthIntegral(Protocol):
     def depth_integral(
         self,