Allow for variable renaming in coupled modeling

fme/coupled/inference/evaluator.py

@@ -63,13 +63,16 @@ class StandaloneComponentCheckpointsConfig:
         atmosphere: The atmosphere component configuration. The stepper
             configuration must include 'ocean'.
         sst_name: Name of the sea surface temperature field in the ocean data.
+        atmosphere_output_rename: Optional mapping from ocean forcing names to
+            atmosphere output names. See ``CoupledStepperConfig``.
 
     """
 
     ocean: StandaloneComponentConfig
     atmosphere: StandaloneComponentConfig
     sst_name: str = "sst"
     ocean_fraction_prediction: CoupledOceanFractionConfig | None = None
+    atmosphere_output_rename: dict[str, str] | None = None
 
     def load_stepper_config(self) -> CoupledStepperConfig:
         return CoupledStepperConfig(
@@ -83,6 +86,7 @@ def load_stepper_config(self) -> CoupledStepperConfig:
             ),
             sst_name=self.sst_name,
             ocean_fraction_prediction=self.ocean_fraction_prediction,
+            atmosphere_output_rename=self.atmosphere_output_rename,
         )
 
     def load_stepper(self) -> CoupledStepper:

fme/coupled/stepper.py

@@ -238,15 +238,23 @@ class CoupledStepperConfig:
             ocean fraction to replace the ocean fraction variable specified in the
             atmosphere's OceanConfig. If the atmosphere uses the ocean fraction as
             an ML forcing, the generated ocean fraction is also passed as an input.
+        atmosphere_output_rename: Optional mapping from ocean forcing names to
+            atmosphere output names when the two components use different names
+            for the same quantity. Uses the same ``{target: source}`` convention
+            as ``CheckpointModelConfig.rename`` in downscaling, e.g.
+            ``eastward_surface_wind_stress: eastward_surface_stress`` couples the
+            ocean input to the atmosphere checkpoint output name.
 
     """
 
     ocean: ComponentConfig
     atmosphere: ComponentConfig
     sst_name: str = "sst"
     ocean_fraction_prediction: CoupledOceanFractionConfig | None = None
+    atmosphere_output_rename: dict[str, str] | None = None
 
     def __post_init__(self):
+        self._validate_atmosphere_output_rename()
         self._validate_component_configs()
 
         atmosphere_ocean_config = self.atmosphere.stepper.get_ocean()
@@ -264,9 +272,12 @@ def __post_init__(self):
         ).to_pytimedelta()
 
         # calculate forcing sets
+        self._atmosphere_to_ocean_forcing_names = (
+            self._compute_atmosphere_to_ocean_forcing_names()
+        )
         self._ocean_forcing_exogenous_names = list(
             set(self.ocean.stepper.input_only_names).difference(
-                self.atmosphere.stepper.output_names
+                self._atmosphere_to_ocean_forcing_names
             )
         )
         unfiltered_atmosphere_forcing_names = list(
@@ -290,11 +301,6 @@ def __post_init__(self):
                 self._atmosphere_forcing_exogenous_names
             )
         )
-        self._atmosphere_to_ocean_forcing_names = list(
-            set(self.ocean.stepper.input_only_names).intersection(
-                self.atmosphere.stepper.output_names
-            )
-        )
         extra_forcings_names = [self.sst_name]
         if self.ocean_fraction_prediction is not None:
             # NOTE: this is only necessary for the special case where the
@@ -415,6 +421,35 @@ def ocean_to_atmosphere_forcing_names(self) -> list[str]:
         """Atmosphere forcing variables that are outputs of the ocean."""
         return self._ocean_to_atmosphere_forcing_names
 
+    def atmosphere_output_name_for_ocean_forcing(self, ocean_forcing_name: str) -> str:
+        """Atmosphere output tensor name used to force the given ocean input."""
+        if self.atmosphere_output_rename is None:
+            return ocean_forcing_name
+        return self.atmosphere_output_rename.get(ocean_forcing_name, ocean_forcing_name)
+
+    def _compute_atmosphere_to_ocean_forcing_names(self) -> list[str]:
+        return [
+            name
+            for name in self.ocean.stepper.input_only_names
+            if self.atmosphere_output_name_for_ocean_forcing(name)
+            in self.atmosphere.stepper.output_names
+        ]
+
+    def _validate_atmosphere_output_rename(self) -> None:
+        if self.atmosphere_output_rename is None:
+            return
+        for ocean_name, atmos_name in self.atmosphere_output_rename.items():
+            if ocean_name not in self.ocean.stepper.input_only_names:
+                raise ValueError(
+                    f"atmosphere_output_rename key {ocean_name!r} is not an ocean "
+                    "input-only name."
+                )
+            if atmos_name not in self.atmosphere.stepper.output_names:
+                raise ValueError(
+                    f"atmosphere_output_rename[{ocean_name!r}]={atmos_name!r} is not "
+                    "an atmosphere output name."
+                )
+
     def _validate_component_configs(self):
         # validate atmosphere's OceanConfig
         atmosphere_ocean_config = self.atmosphere.stepper.get_ocean()
@@ -462,10 +497,8 @@ def _validate_component_configs(self):
 
         # all ocean inputs that are atmosphere outputs must be "next step"
         # forcings according to the ocean stepper config
-        atmosphere_to_ocean_forcing_names = list(
-            set(self.ocean.stepper.input_only_names).intersection(
-                self.atmosphere.stepper.output_names
-            )
+        atmosphere_to_ocean_forcing_names = (
+            self._compute_atmosphere_to_ocean_forcing_names()
         )
         missing_next_step_forcings = list(
             set(atmosphere_to_ocean_forcing_names).difference(
@@ -998,8 +1031,10 @@ def _get_ocean_forcings(
         # get time-averaged forcings from atmosphere
         forcings_from_atmosphere = {
             **{
-                k: atmos_gen[k].mean(time_dim, keepdim=True)
-                for k in self._atmosphere_to_ocean_forcing_names
+                ocean_name: atmos_gen[
+                    self._config.atmosphere_output_name_for_ocean_forcing(ocean_name)
+                ].mean(time_dim, keepdim=True)
+                for ocean_name in self._atmosphere_to_ocean_forcing_names
             },
             **{
                 k: atmos_forcings[k].mean(time_dim, keepdim=True)

fme/coupled/test_stepper.py

@@ -287,6 +287,90 @@ def test_config_names(inputs, expectations):
     )
 
 
+def test_atmosphere_output_rename_couples_mismatched_names():
+    atmosphere = ComponentConfig(
+        timedelta="6h",
+        stepper=StepperConfig(
+            step=StepSelector(
+                type="single_module",
+                config=dataclasses.asdict(
+                    SingleModuleStepConfig(
+                        builder=ModuleSelector(
+                            type="SphericalFourierNeuralOperatorNet",
+                            config={"scale_factor": 1, "embed_dim": 1, "num_layers": 1},
+                        ),
+                        in_names=["a", "frac", "eastward_surface_stress"],
+                        out_names=["a", "eastward_surface_stress"],
+                        normalization=NetworkAndLossNormalizationConfig(
+                            network=NormalizationConfig(
+                                means={
+                                    "a": 0.0,
+                                    "frac": 0.0,
+                                    "eastward_surface_stress": 0.0,
+                                },
+                                stds={
+                                    "a": 1.0,
+                                    "frac": 1.0,
+                                    "eastward_surface_stress": 1.0,
+                                },
+                            ),
+                        ),
+                        ocean=OceanConfig(
+                            surface_temperature_name="a",
+                            ocean_fraction_name="frac",
+                        ),
+                    ),
+                ),
+            ),
+        ),
+    )
+    ocean = ComponentConfig(
+        timedelta="12h",
+        stepper=StepperConfig(
+            step=StepSelector(
+                type="single_module",
+                config=dataclasses.asdict(
+                    SingleModuleStepConfig(
+                        builder=ModuleSelector(
+                            type="SphericalFourierNeuralOperatorNet",
+                            config={"scale_factor": 1, "embed_dim": 1, "num_layers": 1},
+                        ),
+                        in_names=["sst", "eastward_surface_wind_stress"],
+                        out_names=["sst"],
+                        next_step_forcing_names=["eastward_surface_wind_stress"],
+                        normalization=NetworkAndLossNormalizationConfig(
+                            network=NormalizationConfig(
+                                means={
+                                    "sst": 0.0,
+                                    "eastward_surface_wind_stress": 0.0,
+                                },
+                                stds={
+                                    "sst": 1.0,
+                                    "eastward_surface_wind_stress": 1.0,
+                                },
+                            ),
+                        ),
+                    ),
+                ),
+            ),
+        ),
+    )
+    config = CoupledStepperConfig(
+        atmosphere=atmosphere,
+        ocean=ocean,
+        sst_name="sst",
+        atmosphere_output_rename={
+            "eastward_surface_wind_stress": "eastward_surface_stress",
+        },
+    )
+    assert config.atmosphere_to_ocean_forcing_names == ["eastward_surface_wind_stress"]
+    assert "eastward_surface_wind_stress" not in config.ocean_forcing_exogenous_names
+    assert (
+        config.atmosphere_output_name_for_ocean_forcing("eastward_surface_wind_stress")
+        == "eastward_surface_stress"
+    )
+
+
 @pytest.mark.parametrize(
     "inputs, expectations",
     FORCING_TEST_PARAMS,