Asymmetric eval/inference support: coarse forcings can perform a high resolution rollout.

configs/fomo_om4/eval_cross_resolution.yaml

@@ -0,0 +1,56 @@
+# yaml-language-server: $schema=../schemas/EvalConfig.json
+#
+# Cross-resolution eval for the FOMO model (KR2, issue #615).
+#
+# The rollout starts at ¼° and consumes 1° boundary forcings at each step.
+# Prognostics stay at ¼° throughout the rollout — only the boundaries come
+# from a coarser source. Enabled by the dual-perceiver encoder (PR #702),
+# which produces the same latent grid for both streams when patch_extent is
+# in degrees.
+#
+# As in eval_multiscale.yaml, all three resolutions must appear in
+# `sources` so FOMOConfig.build computes `max_patch_size` correctly for the
+# Perceiver's Fourier frequency range. The first source is the one used for
+# inference (¼° here), and `boundary_source` specifies where to load
+# boundary forcings from (1° here).
+
+debug: false
+save_zarr: true
+disk_mode: true
+inference_time:
+  start: "2014-10-05"
+  end: "2015-10-05"
+num_model_steps_forward: 25
+
+experiment:
+  name: fomo_om4_eval_cross_resolution
+  rand_seed: 15
+  base_output_dir: .LOCAL
+  wandb:
+    mode: disabled
+    project: default
+  prognostic_vars_key: thermo_dynamic_all
+  boundary_vars_key: tau_hfds_hfds_anom
+
+data:
+  static_data_vars: []
+  concurrent_compute: true
+  sources:
+    # First source is used for inference prognostics.
+    - data_location: om4_quarterdeg_v2/OM4.zarr
+      data_means_location: om4_quarterdeg_v2/OM4_means.zarr
+      data_stds_location: om4_quarterdeg_v2/OM4_stds.zarr
+    - data_location: om4_halfdeg_v4/OM4.zarr
+      data_means_location: om4_halfdeg_v4/OM4_means.zarr
+      data_stds_location: om4_halfdeg_v4/OM4_stds.zarr
+    - data_location: om4_onedeg_v3/OM4.zarr
+      data_means_location: om4_onedeg_v3/OM4_means.zarr
+      data_stds_location: om4_onedeg_v3/OM4_stds.zarr
+  # Boundary forcings come from the 1° source, while prognostics come from
+  # `sources[0]` (¼°). This is the KR2 asymmetric-inference setup.
+  boundary_source:
+    data_location: om4_onedeg_v3/OM4.zarr
+    data_means_location: om4_onedeg_v3/OM4_means.zarr
+    data_stds_location: om4_onedeg_v3/OM4_stds.zarr
+
+model: !include model.yaml

src/ocean_emulators/config.py

@@ -194,6 +194,17 @@ class DataConfig(BaseConfig):
         ),
         min_length=1,
     )
+    boundary_source: DataSourceConfig | None = Field(
+        default=None,
+        description=(
+            "Optional separate boundary source for inference. When set, "
+            "InferenceDataset loads boundary forcings from this source "
+            "(typically at a coarser resolution than the prognostic source), "
+            "enabling cross-resolution rollouts such as ¼° prognostics + "
+            "1° boundaries. When unset, boundaries are loaded from the same "
+            "source as the prognostics (current default behavior)."
+        ),
+    )
     static_data_vars: list[str] | None = None
     loading: DataLoadingConfig = Field(default_factory=CpuDataLoadingConfig)
     hist: int = 1
@@ -271,12 +282,22 @@ def make_source(
             else None
         )
 
+        inference_boundary_source: DataSource | None = None
+        if self.boundary_source is not None:
+            inference_boundary_source, _ = make_source(
+                self.boundary_source.data_location,
+                self.boundary_source.data_means_location,
+                self.boundary_source.data_stds_location,
+                turn_on_dask=True,
+            )
+
         return DataContainer(
             sources,
             inference_source,
             loader_version,
             supports_fork,
             static_data,
+            inference_boundary_source=inference_boundary_source,
         )
 
 

src/ocean_emulators/datasets.py

@@ -60,6 +60,7 @@ def __init__(
         normalize_before_mask,
         masked_fill_value,
         long_rollout,
+        boundary_src: DataSource | None = None,
     ):
         super().__init__()
         # NOTE: Keep tensors on CPU during initialization. This allows the dataset
@@ -72,7 +73,22 @@ def __init__(
         data = src.data
         self.input_res = src.resolution
         self._prognostic_src = src.filter(prognostic_var_names, prefix="prognostic")
-        self._boundary_src = src.filter(boundary_var_names, prefix="boundary")
+        # When boundary_src is provided, boundaries come from a different source
+        # (typically a coarser resolution) — enabling cross-resolution rollouts.
+        # The rolling indices are built off the prognostic source's time axis and
+        # reused for both streams, so the two sources must share a time axis.
+        if boundary_src is not None:
+            if not boundary_src.data.time.equals(data.time):
+                raise ValueError(
+                    "Boundary source time axis does not match the prognostic "
+                    "source. Cross-resolution inference requires both sources "
+                    "to share a time axis."
+                )
+            self._boundary_src = boundary_src.filter(
+                boundary_var_names, prefix="boundary"
+            )
+        else:
+            self._boundary_src = src.filter(boundary_var_names, prefix="boundary")
         self._times = data.time
         self.normalize_before_mask = normalize_before_mask
         self.masked_fill_value = masked_fill_value
@@ -102,7 +118,12 @@ def __init__(
             )
 
         self.wet: PrognosticMask = src.masks.prognostic
-        self.wet_surface: GridMask = src.masks.boundary
+        # `wet_surface` is applied to the boundary tensor and must match its grid.
+        self.wet_surface: GridMask = (
+            boundary_src.masks.boundary
+            if boundary_src is not None
+            else src.masks.boundary
+        )
         self.wet_label = src.masks.prognostic_with_hist(self.hist)
         self.size = len(self.rolling_indices)
 

src/ocean_emulators/eval.py

@@ -97,6 +97,7 @@ def __init__(self, cfg: EvalConfig) -> None:
         )
 
         self.src = self.data_container.inference_source
+        self.boundary_src = self.data_container.inference_boundary_source
         self.data = self.src.data
         self.static_data = self.data_container.static_data
         self.metadata = construct_metadata(self.data)
@@ -165,6 +166,11 @@ def load_checkpoint(self, ckpt_path: str):
 
     def init_inference_store(self):
         sliced_src = self.src.slice(self.inference_time)
+        sliced_boundary_src = (
+            self.boundary_src.slice(self.inference_time)
+            if self.boundary_src is not None
+            else None
+        )
         self.num_time_steps = get_inference_steps(
             sliced_src,
             hist=self.hist,
@@ -177,6 +183,7 @@ def init_inference_store(self):
             normalize_before_mask=self.normalize_before_mask,
             masked_fill_value=self.masked_fill_value,
             long_rollout=True,
+            boundary_src=sliced_boundary_src,
         )
 
     def run(self) -> None:

src/ocean_emulators/utils/data.py

@@ -421,6 +421,11 @@ class DataContainer:
     # TODO(559): static_data should belong to the DataSource, since we now
     #  deal with multiple resolutions.
     static_data: xr.Dataset | None = None
+    # Optional separate source for boundary forcings at inference time.
+    # When set, InferenceDataset loads boundaries from this source (typically
+    # at a coarser resolution than `inference_source`), enabling cross-resolution
+    # rollouts such as ¼° prognostics + 1° boundaries.
+    inference_boundary_source: DataSource | None = None
 
     @property
     def primary_source(self) -> DataSource:

tests/conftest.py

@@ -567,6 +567,47 @@ def trainer_pair(
         yield train_config, trainer
 
 
+def build_synthetic_source(
+    name: str,
+    h: int,
+    w: int,
+    n_times: int,
+    prognostic_var_names: list[str],
+    boundary_var_names: list[str],
+    seed: int | None = None,
+) -> DataSource:
+    """Build a synthetic DataSource at an arbitrary resolution, full-globe grid.
+
+    Lat/lon coordinates span the full globe so `patch_extent`-based slicing is
+    consistent across resolutions — required when combining two sources (e.g.
+    cross-resolution inference). Variable data is standard-normal; means/stds
+    are zero/one so `normalize_before_mask=True` is a no-op.
+    """
+    all_vars = list(prognostic_var_names) + list(boundary_var_names)
+    coords = {
+        "time": np.arange(n_times),
+        "lat": np.linspace(-90 + 180 / (2 * h), 90 - 180 / (2 * h), h),
+        "lon": np.linspace(360 / (2 * w), 360 - 360 / (2 * w), w),
+    }
+    generator = torch.Generator()
+    if seed is not None:
+        generator.manual_seed(seed)
+    arr = torch.randn(len(all_vars), n_times, h, w, generator=generator).numpy()
+    data = xr.Dataset(
+        {
+            v: xr.DataArray(arr[i], dims=["time", "lat", "lon"], coords=coords)
+            for i, v in enumerate(all_vars)
+        }
+    )
+    stats_coords = {"lat": [0], "lon": [0]}
+    mean_ds = xr.Dataset({v: 0.0 for v in all_vars}, coords=stats_coords)
+    std_ds = xr.Dataset({v: 1.0 for v in all_vars}, coords=stats_coords)
+    wet_surface = torch.ones(h, w)
+    wet = wet_surface.expand(len(prognostic_var_names), h, w)
+    masks = Masks(prognostic=wet, boundary=wet_surface)
+    return DataSource(name, data, mean_ds, std_ds, masks=masks)
+
+
 @pytest.fixture
 def dummy_src():
     h, w = 4, 8

tests/test_datasets.py

@@ -35,7 +35,13 @@
 from ocean_emulators.utils.multiton import MultitonScope
 from ocean_emulators.utils.samplers import EquivalenceGroupBatchSampler
 from ocean_emulators.utils.train import collate_raw_train_data
-from tests.conftest import DEFAULT_CONFIG, DataSourceDims, TrainPair, cache_dir
+from tests.conftest import (
+    DEFAULT_CONFIG,
+    DataSourceDims,
+    TrainPair,
+    build_synthetic_source,
+    cache_dir,
+)
 
 
 @pytest.fixture
@@ -719,6 +725,94 @@ def test_train_dataset_normalize_pre_fill(
         assert inf_prog[0, 0, 0, 0] == data
 
 
+def test_inference_dataset__cross_resolution():
+    """InferenceDataset yields mismatched prog/boundary shapes when given distinct sources."""
+    prognostic_var_names = ["prognostic1", "prognostic2"]
+    boundary_var_names = ["boundary1", "boundary2"]
+    prog_src = build_synthetic_source(
+        "high_res",
+        h=8,
+        w=16,
+        n_times=10,
+        prognostic_var_names=prognostic_var_names,
+        boundary_var_names=boundary_var_names,
+    )
+    boundary_src = build_synthetic_source(
+        "low_res",
+        h=2,
+        w=4,
+        n_times=10,
+        prognostic_var_names=prognostic_var_names,
+        boundary_var_names=boundary_var_names,
+    )
+
+    with MultitonScope():
+        Normalize.init_instance(
+            prog_src,
+            prognostic_var_names=prognostic_var_names,
+            boundary_var_names=boundary_var_names,
+        )
+        dataset = InferenceDataset(
+            src=prog_src,
+            prognostic_var_names=prognostic_var_names,
+            boundary_var_names=boundary_var_names,
+            hist=0,
+            normalize_before_mask=True,
+            masked_fill_value=0.0,
+            long_rollout=True,
+            boundary_src=boundary_src,
+        )
+
+        prog, boundary = dataset.get_initial_input()
+
+        # Prognostic sits on the high-res grid; boundary on the low-res grid.
+        assert prog.shape[-2:] == (8, 16)
+        assert boundary.shape[-2:] == (2, 4)
+        # Context tracks the prognostic resolution on both in and out.
+        assert dataset.ctx.input_resolution_cpu[0].shape == (8,)
+        assert dataset.ctx.output_resolution_cpu[0].shape == (8,)
+
+
+def test_inference_dataset__cross_resolution_time_mismatch_raises():
+    """Sources with different time axes must fail loudly."""
+    prognostic_var_names = ["prognostic1", "prognostic2"]
+    boundary_var_names = ["boundary1", "boundary2"]
+    prog_src = build_synthetic_source(
+        "prog",
+        h=8,
+        w=16,
+        n_times=10,
+        prognostic_var_names=prognostic_var_names,
+        boundary_var_names=boundary_var_names,
+    )
+    boundary_src = build_synthetic_source(
+        "boundary",
+        h=2,
+        w=4,
+        n_times=5,
+        prognostic_var_names=prognostic_var_names,
+        boundary_var_names=boundary_var_names,
+    )
+
+    with MultitonScope():
+        Normalize.init_instance(
+            prog_src,
+            prognostic_var_names=prognostic_var_names,
+            boundary_var_names=boundary_var_names,
+        )
+        with pytest.raises(ValueError, match="time axis"):
+            InferenceDataset(
+                src=prog_src,
+                prognostic_var_names=prognostic_var_names,
+                boundary_var_names=boundary_var_names,
+                hist=0,
+                normalize_before_mask=True,
+                masked_fill_value=0.0,
+                long_rollout=True,
+                boundary_src=boundary_src,
+            )
+
+
 @pytest.mark.manual
 @pytest.mark.parametrize(
     "data_source,config_name", [("mock", DEFAULT_CONFIG)], indirect=True