Coupled evaluator with prediction loader by elynnwu · Pull Request #856 · ai2cm/ace

elynnwu · 2026-02-19T23:29:56Z

prediction_loader is an already supported feature in ace evaluator, we add this feature for coupled evaluator so that we can easily compare with target data on wandb.

Resolves #805

elynnwu · 2026-02-19T23:32:37Z

fme/coupled/inference/evaluator.py

+            restrict_to_output_names=(
+                stepper.ocean.out_names,
+                stepper.atmosphere.out_names,
+            ),


We did not have to do this in ace evaluator because there was no merge operation, we need this here because we can have overlap in atmosphere and ocean in_names (e.g., land_fraction)

In the fme.ace version we log all variables including input-only forcings. I think this can be useful for examining the forcings and for post-hoc analyses.

See my other comment. I think we can instead add an @property called all_names to CoupledStepperConfig which returns a CoupledNames where:

atmosphere has its out names + atmosphere_forcing_exogenous_names + shared_forcing_exogenous_names

ocean has its out names + ocean_forcing_exogenous_names

Good point, added

elynnwu · 2026-02-20T01:06:13Z

Job here
wandb here

jpdunc23

Looks good! I have some suggestions for handling the names.

jpdunc23 · 2026-02-20T19:54:09Z

fme/coupled/inference/loop.py

+    deriver: CoupledDeriverABC,
+    writer: CoupledPairedDataWriter | NullDataWriter | None = None,
+    record_logs: Callable[[InferenceLogs], None] | None = None,
+    restrict_to_output_names: tuple[list[str], list[str]] | None = None,


Can you please add a new dataclass to replace tuple[list[str], list[str]] in fme/coupled/typing_.py:

@dataclasses.dataclass class CoupledNames: ocean: list[str] atmosphere: list[str]

~~You could also add an @property called out_names to CoupledStepperConfig which gets the names from the component steppers and returns a CoupledNames object.~~ See my other comment where I suggest instead adding all_names to CoupledStepperConfig.

jpdunc23 · 2026-02-20T19:59:10Z

fme/coupled/inference/test_evaluator.py

+        assert not os.path.exists(tmp_path / "atmosphere/restart.nc")
+        assert not os.path.exists(tmp_path / "atmosphere/initial_condition.nc")
+        assert not os.path.exists(tmp_path / "ocean/restart.nc")
+        assert not os.path.exists(tmp_path / "ocean/initial_condition.nc")


Would be good to add a check similar to the following in the fme.ace test:

# if these are off by something like 90% then probably the stepper # is being used instead of the prediction_data assert log[f"inference/mean/weighted_rmse/{var}"] == 0.0 assert log[f"inference/mean/weighted_bias/{var}"] == 0.0

jpdunc23 · 2026-02-20T20:02:01Z

fme/coupled/inference/loop.py

+                    initial_condition=CoupledPairedData(
+                        ocean_data=PairedData.from_batch_data(
+                            prediction=pred_ic.ocean_data,
+                            reference=target_ic.ocean_data,
+                        ),
+                        atmosphere_data=PairedData.from_batch_data(
+                            prediction=pred_ic.atmosphere_data,
+                            reference=target_ic.atmosphere_data,
+                        ),
+                    ),


Can you use CoupledPairedData.from_coupled_batch_data here?

Yes, updated

jpdunc23 · 2026-02-20T20:06:25Z

fme/coupled/inference/loop.py

+from fme.coupled.inference.data_writer import CoupledPairedDataWriter
+
+
+class CoupledDeriverABC(Protocol):


I think using Protocol rather than abc.ABC as in fme.ace is good, but we should call it CoupledDeriverProtocol or simply CoupledDeriver.

Changed to CoupledDeriver

jpdunc23 · 2026-02-20T20:27:25Z

fme/coupled/inference/evaluator.py

+            restrict_to_output_names=(
+                stepper.ocean.out_names,
+                stepper.atmosphere.out_names,
+            ),


In the fme.ace version we log all variables including input-only forcings. I think this can be useful for examining the forcings and for post-hoc analyses.

See my other comment. I think we can instead add an @property called all_names to CoupledStepperConfig which returns a CoupledNames where:

atmosphere has its out names + atmosphere_forcing_exogenous_names + shared_forcing_exogenous_names

ocean has its out names + ocean_forcing_exogenous_names

…-loader

jpdunc23

Couple of nits and a question, but LGTM.

jpdunc23 · 2026-02-23T22:06:20Z

fme/coupled/data_loading/test_data_loader.py

    timestep_size=1,
    timestep_start=0,
    nz=3,
+    masked_fill_value: float = float("nan"),


Do you know why this is needed for prediction_loader inference but not evaluator inference?

Okay this is because we didn't write mask_2d so variables like o_prog doesn't have the right mask, I reverted back to using nans as masked fill values and added the proper mask.

jpdunc23 · 2026-02-23T22:10:34Z

fme/coupled/stepper.py

+        atmosphere_names = (
+            self.atmosphere.stepper.output_names
+            + self._atmosphere_forcing_exogenous_names
+            + self._shared_forcing_exogenous_names
+        )


Shouldn't matter in practice, but agent review flagged that _shared_forcing_exogenous_names is a subset of _atmosphere_forcing_exogenous_names, so probably best to remove _shared_forcing_exogenous_names here.

jpdunc23 · 2026-02-23T22:12:41Z

fme/coupled/inference/loop.py

+    deriver: CoupledDeriver,
+    writer: CoupledPairedDataWriter | NullDataWriter | None = None,
+    record_logs: Callable[[InferenceLogs], None] | None = None,
+    restrict_to_output_names: CoupledNames | None = None,


Nit: Maybe call this all_names since it is no longer restricted to just output names.

elynnwu added 7 commits February 18, 2026 15:56

cpl pred loader

d0ea5ff

run exp

24513cd

fix type

f1d388f

fix time stamp

296b3d5

only report out_names

ccab097

cleanup

b41af48

add unit test

24a2882

elynnwu commented Feb 19, 2026

View reviewed changes

delete exp

05be474

jpdunc23 reviewed Feb 20, 2026

View reviewed changes

elynnwu and others added 4 commits February 20, 2026 14:45

address reviewer comments

b178d48

reduce domain size of flaky slow downscaling inline test

8d5a1a4

Merge branch 'main' into fix/slow_test_downscaling_inline

61ee35f

Merge branch 'fix/slow_test_downscaling_inline' into feature/cpl-pred…

c8fec2d

…-loader

jpdunc23 approved these changes Feb 23, 2026

View reviewed changes

elynnwu and others added 2 commits February 23, 2026 14:29

address PR comments

85c9a79

Merge branch 'main' into feature/cpl-pred-loader

a04512e

elynnwu enabled auto-merge (squash) February 23, 2026 22:29

elynnwu merged commit a40e75c into main Feb 23, 2026
7 checks passed

elynnwu deleted the feature/cpl-pred-loader branch February 23, 2026 22:42

		from fme.coupled.inference.data_writer import CoupledPairedDataWriter


		class CoupledDeriverABC(Protocol):

Comments

Conversation

elynnwu commented Feb 19, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

elynnwu commented Feb 20, 2026

Uh oh!

jpdunc23 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jpdunc23 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants