diff --git a/configs/fomo_om4/eval_cross_resolution.yaml b/configs/fomo_om4/eval_cross_resolution.yaml new file mode 100644 index 000000000..9f4ccbbcf --- /dev/null +++ b/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 diff --git a/src/ocean_emulators/config.py b/src/ocean_emulators/config.py index 0304d9219..14928feda 100644 --- a/src/ocean_emulators/config.py +++ b/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, ) diff --git a/src/ocean_emulators/datasets.py b/src/ocean_emulators/datasets.py index 45b1f9c43..367ede4af 100644 --- a/src/ocean_emulators/datasets.py +++ b/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) diff --git a/src/ocean_emulators/eval.py b/src/ocean_emulators/eval.py index 744a65e30..a6114b758 100644 --- a/src/ocean_emulators/eval.py +++ b/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: diff --git a/src/ocean_emulators/utils/data.py b/src/ocean_emulators/utils/data.py index 5e1f971a2..ded09d7d1 100644 --- a/src/ocean_emulators/utils/data.py +++ b/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: diff --git a/tests/conftest.py b/tests/conftest.py index 812a35bd6..295c65b2f 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 23ad642fe..350266566 100644 --- a/tests/test_datasets.py +++ b/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 diff --git a/tests/test_fomo_cross_resolution.py b/tests/test_fomo_cross_resolution.py index 244c62f97..bc26b9432 100644 --- a/tests/test_fomo_cross_resolution.py +++ b/tests/test_fomo_cross_resolution.py @@ -5,6 +5,7 @@ from perceiver_pytorch import Perceiver from perceiver_pytorch.perceiver_io import PerceiverIO +from ocean_emulators.datasets import InferenceDataset from ocean_emulators.models.fomo import FOMO from ocean_emulators.models.modules import ( PerceiverDecoder, @@ -18,6 +19,9 @@ CoreBlock, ) from ocean_emulators.utils.ctx import GridContext +from ocean_emulators.utils.data import Normalize +from ocean_emulators.utils.multiton import MultitonScope +from tests.conftest import build_synthetic_source LATENT_DIM = 4 EMBED_DIM = 8 @@ -234,3 +238,78 @@ def test_fomo_autoregressive_mix_schedule(add_3d_coordinates: bool): for out in outputs: assert out.shape == (1, out_channels, output_h, output_w) assert torch.isfinite(out).all(), "Output contains NaN or Inf." + + +def test_fomo_inference_cross_resolution_ar_rollout(): + """BaseModel.inference runs a multi-step AR rollout with asymmetric sources. + + Prognostics come from a high-res source (¼°-like grid, 8×16) and boundaries + from a low-res source (1°-like grid, 2×4), sharing a time axis. This is + the KR2 asymmetric-inference path: prognostics stay at the fine grid + throughout the rollout while boundary forcings come from a coarser source + at every step. + """ + prognostic_var_names = ["prognostic1", "prognostic2"] + boundary_var_names = ["boundary1", "boundary2"] + high_h, high_w = 8, 16 + low_h, low_w = 2, 4 + n_times = 10 + num_prog, num_boundary = ( + len(prognostic_var_names), + len(boundary_var_names), + ) + + prog_src = build_synthetic_source( + "high_res", + h=high_h, + w=high_w, + n_times=n_times, + prognostic_var_names=prognostic_var_names, + boundary_var_names=boundary_var_names, + ) + boundary_src = build_synthetic_source( + "low_res", + h=low_h, + w=low_w, + n_times=n_times, + prognostic_var_names=prognostic_var_names, + boundary_var_names=boundary_var_names, + ) + + model = _make_fomo( + prog_channels=num_prog, + boundary_channels=num_boundary, + out_channels=num_prog, + ) + model.eval() + + 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, + ) + initial_prog = dataset.initial_prognostic + + with torch.no_grad(): + output = model.inference( + dataset=dataset, + initial_prognostic=initial_prog, + steps_completed=0, + num_steps=7, + ) + + assert output.prediction.shape == (7, num_prog, high_h, high_w) + assert torch.isfinite(output.prediction).all(), ( + "AR rollout prediction contains NaN or Inf." + )