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

configs/fomo_om4/train_cross_resolution.yaml

@@ -0,0 +1,75 @@
+# yaml-language-server: $schema=../schemas/TrainConfig.json
+#
+# Cross-resolution training for FOMO (KR2, issue #615).
+#
+# Prognostic inputs and labels come from the first source (0.25 degree here),
+# while boundary forcings come from the explicit 1 degree boundary_source.
+# The additional sources are listed so FOMOConfig.build can compute
+# max_patch_size across the supported scales.
+
+debug: false
+disk_mode: true
+pin_mem: true
+save_freq: 5
+epochs: 70
+batch_size: 2
+learning_rate: 0.0006
+scheduler: { type: cosine }
+loss:
+  type: dynamic
+  metric: mse
+finetune: false
+resume_ckpt_path: null
+inference_epochs: []
+data_percent: 1.0
+train_time:
+  start: "1975-01-03"
+  end: "2013-10-05"
+val_time:
+  start: "2013-10-05"
+  end: "2014-10-05"
+inference_times:
+  - start: "1975-01-03"
+    end: "1976-01-03"
+  - start: "1985-01-03"
+    end: "1986-01-03"
+  - start: "1995-01-03"
+    end: "1996-01-03"
+  - start: "2004-10-05"
+    end: "2005-10-05"
+data_stride: [1]
+steps: [4]
+step_transition: []
+preemptible: false
+
+backend: auto
+
+experiment:
+  name: fomo_om4_train_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:
+    - 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_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. When set, training, validation, "
+            "and inference load boundary forcings from this source (typically at "
+            "a coarser resolution than the prognostic source), enabling "
+            "cross-resolution runs 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
@@ -249,8 +260,6 @@ def make_source(
             )
             sources.append(src)
             supports_forks.append(fork)
-        supports_fork = all(supports_forks)
-
         primary_source = sources[0]
         if use_dask:
             # If we're already using dask, we don't need a second source
@@ -271,12 +280,35 @@ def make_source(
             else None
         )
 
+        boundary_source: DataSource | None = None
+        inference_boundary_source: DataSource | None = None
+        if self.boundary_source is not None:
+            boundary_source, boundary_fork = make_source(
+                self.boundary_source.data_location,
+                self.boundary_source.data_means_location,
+                self.boundary_source.data_stds_location,
+            )
+            supports_forks.append(boundary_fork)
+            if use_dask:
+                inference_boundary_source = boundary_source
+            else:
+                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,
+                )
+
+        supports_fork = all(supports_forks)
+
         return DataContainer(
             sources,
             inference_source,
             loader_version,
             supports_fork,
             static_data,
+            boundary_source=boundary_source,
+            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)
 
@@ -469,11 +490,13 @@ def __init__(
         masked_fill_value: float,
         stride: int = 1,
         concurrent_compute_: bool = False,
+        boundary_src: DataSource | None = None,
     ):
         super().__init__()
         self.id = f"{self.__class__.__name__}_{str(id(self))}"
         # If the src and dst DataSource are the same, we can do a lot less work.
         srcs = [src, dst] if dst else [src]
+        boundary_data_src = boundary_src if boundary_src is not None else src
 
         self.hist: int = hist
         self.steps: int = steps
@@ -487,11 +510,18 @@ def __init__(
         assert np.array_equal(srcs[0].data.time, srcs[-1].data.time), (
             "src and dst DataSource have different time slices!"
         )
+        if not boundary_data_src.data.time.equals(src.data.time):
+            raise ValueError(
+                "Boundary source time axis does not match the prognostic source. "
+                "Cross-resolution training requires both sources to share a time axis."
+            )
         time_ = src.data.time
         self.prognostic_srcs = [
             src.filter(prognostic_var_names, prefix="prog") for src in srcs
         ]
-        self.boundary_src = src.filter(boundary_var_names, prefix="boundary")
+        self.boundary_src = boundary_data_src.filter(
+            boundary_var_names, prefix="boundary"
+        )
 
         # This class will be used only for training and validation
         total_steps: int = 2 * self.hist + 2
@@ -515,7 +545,7 @@ def __init__(
         self.wet_prognostic: list[PrognosticMask] = [
             src.masks.prognostic for src in srcs
         ]
-        self.wet_surface: GridMask = src.masks.boundary
+        self.wet_surface: GridMask = boundary_data_src.masks.boundary
 
         self.ctx = GridContext(
             label_mask=self.prognostic_srcs[-1].masks.prognostic_with_hist(self.hist),

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/train.py

@@ -205,6 +205,7 @@ def __init__(self, cfg: TrainConfig) -> None:
         # TODO(jder): Could rewrite inference dataset like we did for TorchTrainDataset
         # see https://github.com/suryadheeshjith/Ocean_Emulator/issues/208
         self.inference_src = self.data_container.inference_source
+        self.boundary_src = self.data_container.boundary_source
 
         self.loader_version = self.data_container.loader_version
 
@@ -793,10 +794,22 @@ def init_data_loaders(self, cur_step: int) -> None:
             case _:
                 assert_never(self.train_schedule)
 
+        train_boundary_src = (
+            self.boundary_src.slice(self.train_time)
+            if self.boundary_src is not None
+            else None
+        )
+        val_boundary_src = (
+            self.boundary_src.slice(self.val_time)
+            if self.boundary_src is not None
+            else None
+        )
+
         train_datasets = [
             TorchTrainDataset(
                 src=src.slice(self.train_time),
                 dst=dst.slice(self.train_time) if dst else None,
+                boundary_src=train_boundary_src,
                 prognostic_var_names=self.prognostic_var_names,
                 boundary_var_names=self.boundary_var_names,
                 hist=self.hist,
@@ -814,6 +827,7 @@ def init_data_loaders(self, cur_step: int) -> None:
             TorchTrainDataset(
                 src=src.slice(self.val_time),
                 dst=dst.slice(self.val_time) if dst else None,
+                boundary_src=val_boundary_src,
                 prognostic_var_names=self.prognostic_var_names,
                 boundary_var_names=self.boundary_var_names,
                 hist=self.hist,
@@ -855,9 +869,12 @@ def init_data_loaders(self, cur_step: int) -> None:
                 )
 
         # Create batch samplers - branch on distributed vs non-distributed
-        # Group by input AND label resolution to handle all training schedules
+        # Group by input, boundary, and label resolution so stacked batches have
+        # matching tensor shapes across all training schedules.
         def group_key(ds):
-            return tuple(prog.grid_size for prog in ds.prognostic_srcs)
+            return tuple(prog.grid_size for prog in ds.prognostic_srcs) + (
+                ds.boundary_src.grid_size,
+            )
 
         if self.distributed is not None:
             # Distributed training

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 during training/validation.
+    # When unset, training boundaries come from each prognostic input source.
+    boundary_source: DataSource | None = None
+    # Optional dask-backed boundary source for inference/eval rollouts.
+    inference_boundary_source: DataSource | None = None
 
     @property
     def primary_source(self) -> DataSource: