diff --git a/fme/coupled/inference/evaluator.py b/fme/coupled/inference/evaluator.py index 1ca99607f..dbbe4572d 100644 --- a/fme/coupled/inference/evaluator.py +++ b/fme/coupled/inference/evaluator.py @@ -63,6 +63,8 @@ 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``. """ @@ -70,6 +72,7 @@ class StandaloneComponentCheckpointsConfig: 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: diff --git a/fme/coupled/stepper.py b/fme/coupled/stepper.py index b4be42211..9f16f1479 100644 --- a/fme/coupled/stepper.py +++ b/fme/coupled/stepper.py @@ -238,6 +238,12 @@ 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. """ @@ -245,8 +251,10 @@ class CoupledStepperConfig: 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) diff --git a/fme/coupled/test_stepper.py b/fme/coupled/test_stepper.py index ff19d67fb..31dc9b096 100644 --- a/fme/coupled/test_stepper.py +++ b/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,