Upload folder using huggingface_hub

README.md

@@ -1,13 +1,40 @@
 ---
-title: Samudra2
-emoji: 🦀
-colorFrom: pink
-colorTo: indigo
+title: Samudra 2
+emoji: 🌊
+colorFrom: green
+colorTo: gray
 sdk: gradio
 sdk_version: 6.25.0
-python_version: '3.12'
 app_file: app.py
+short_description: Neural ocean emulator rollouts scored against GFDL OM4
+python_version: "3.12"
+startup_duration_timeout: 1h
 pinned: false
+license: apache-2.0
+models:
+  - M2LInES/Samudra2
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Samudra 2 — a neural ocean emulator
+
+[Samudra 2](https://huggingface.co/M2LInES/Samudra2) (M2LInES) is an autoregressive
+deep-learning emulator of the **GFDL OM4 global ocean model**. Given one ocean state plus
+atmospheric forcing it advances temperature, salinity and horizontal currents across 19
+depth levels together with sea surface height — 77 prognostic fields at once — 10 days per
+autoregressive step.
+
+This Space runs the released **1° checkpoint** (`onedeg/ema_ckpt.pt`, 84M parameters) on
+genuinely held-out years (2014–2022). The initial condition and the atmospheric forcing are
+streamed live from the public OM4 archive (`m2lines-pubs/Samudra` on the Open Storage
+Network), the emulator is rolled forward on ZeroGPU, and the emulated ocean is scored
+against the OM4 simulation it emulates — including a persistence baseline.
+
+## Notes
+
+- The vendored `samudra/` package is the upstream
+  [`m2lines/Samudra`](https://github.com/m2lines/Samudra) library (Apache-2.0), included
+  directly because its full dependency set (SkyPilot, xESMF) cannot be pip-resolved inside
+  a Space. Model configuration matches the upstream `samudra_om4_v2` config.
+- Only the 1° checkpoint is demoed: the half- and quarter-degree checkpoints require the
+  `hfds_anomalies` boundary field, which is not present in the corresponding public stores.
+- Each emulator step emits two 5-day snapshots, so *n* steps ⇒ 10·*n* days of forecast.

app.py

@@ -0,0 +1,400 @@
+"""Samudra 2 — an autoregressive neural ocean emulator (M2LInES).
+
+The app streams real GFDL-OM4 ocean state from public S3 storage, rolls the
+released 1-degree Samudra 2 checkpoint forward autoregressively on ZeroGPU,
+and compares the emulated ocean against the OM4 simulation it emulates.
+"""
+
+import spaces  # noqa: F401  (must precede any CUDA-touching import)
+
+import dataclasses
+import functools
+import os
+import tempfile
+import time
+
+import matplotlib
+
+matplotlib.use("Agg")
+
+import gradio as gr
+import imageio.v2 as imageio
+import matplotlib.pyplot as plt
+import numpy as np
+import s3fs
+import torch
+import xarray as xr
+from huggingface_hub import hf_hub_download
+
+import samudra.utils.device as sdev
+from samudra.config import BlockConfig, SamudraConfig, UNetBackboneConfig
+from samudra.constants import build_om4_spec
+from samudra.datasets import InferenceDataset
+from samudra.utils.data import DataSource, Normalize
+
+# Data pipeline work happens on CPU in the main process; the GPU worker flips
+# this to "cuda" for the forward passes only.
+sdev.set_device(torch.device("cpu"))
+
+# --------------------------------------------------------------------------
+# Data: public 1-degree OM4 store (Open Storage Network, anonymous access)
+# --------------------------------------------------------------------------
+ENDPOINT = "https://nyu1.osn.mghpcc.org"
+ROOT = "m2lines-pubs/Samudra"
+HIST = 1
+MAX_STEPS = 12
+
+_fs = s3fs.S3FileSystem(anon=True, client_kwargs={"endpoint_url": ENDPOINT})
+
+
+def _open(name: str) -> xr.Dataset:
+    return xr.open_zarr(
+        s3fs.S3Map(root=f"{ROOT}/{name}", s3=_fs, check=False),
+        consolidated=True,
+        chunks={},
+    )
+
+
+print("[samudra] opening public OM4 zarr stores ...", flush=True)
+_t0 = time.time()
+SPEC = build_om4_spec("thermo_dynamic_all", "tau_hfds_hfds_anom")
+SRC = DataSource.from_datasets(
+    _open("OM4"),
+    _open("OM4_means"),
+    _open("OM4_stds"),
+    dataset_spec=SPEC,
+    prognostic_var_names=SPEC.prognostic_var_names,
+    boundary_var_names=SPEC.boundary_var_names,
+    name="om4-1deg",
+)
+PROG_NAMES = list(SPEC.prognostic_var_names)
+BND_NAMES = list(SPEC.boundary_var_names)
+N_PROG = len(PROG_NAMES)
+_norm = Normalize(SRC, SPEC.prognostic_var_names, SPEC.boundary_var_names)
+PROG_MEAN = np.asarray(_norm._prognostic_mean_np, dtype=np.float64)
+PROG_STD = np.asarray(_norm._prognostic_std_np, dtype=np.float64)
+TIMES = [str(t)[:10] for t in SRC.data.time.values]
+print(f"[samudra] data source ready in {time.time() - _t0:.1f}s "
+      f"({len(TIMES)} snapshots, {N_PROG} prognostic vars)", flush=True)
+
+# Samudra 2 was trained on 1975-2013 and validated to 2014-10; everything from
+# 2014-10-10 onward is held out, so restrict the demo to genuine forecasts.
+FIRST_HELDOUT = TIMES.index("2014-10-10")
+LAST_START = len(TIMES) - (2 * MAX_STEPS + 2)
+START_DATES = TIMES[FIRST_HELDOUT : LAST_START + 1]
+
+# --------------------------------------------------------------------------
+# Model: the released 1-degree Samudra 2 checkpoint
+# --------------------------------------------------------------------------
+print("[samudra] building model ...", flush=True)
+_model_cfg = SamudraConfig(
+    checkpointing=None,
+    pred_residuals=False,
+    last_kernel_size=3,
+    pad="circular",
+    unet=UNetBackboneConfig(
+        ch_width=[280, 380, 480, 520],
+        dilation=[1, 2, 4, 8],
+        n_layers=[1, 1, 1, 1],
+        core_block=BlockConfig(
+            block_type="conv_next_block",
+            kernel_size=3,
+            activation="capped_gelu",
+            upscale_factor=2,
+            norm="batch",
+        ),
+        down_sampling_block="avg_pool",
+        up_sampling_block="zonally_periodic_upsample",
+    ),
+)
+_grid_src = dataclasses.replace(SRC, data=SRC.data.isel(time=slice(0, 4)))
+MODEL = _model_cfg.build(
+    prog_channels=(HIST + 1) * N_PROG,
+    boundary_channels=(HIST + 1) * len(BND_NAMES),
+    out_channels=(HIST + 1) * N_PROG,
+    hist=HIST,
+    srcs=[_grid_src],
+)
+_ckpt_path = hf_hub_download("M2LInES/Samudra2", "onedeg/ema_ckpt.pt")
+_ck = torch.load(_ckpt_path, map_location="cpu", weights_only=False)
+MODEL.load_state_dict(
+    {k.removeprefix("module."): v for k, v in _ck["model"].items()}, strict=True
+)
+del _ck
+DEVICE = os.environ.get("SAMUDRA_DEVICE", "cuda")
+MODEL.eval().to(DEVICE)
+N_PARAMS = sum(p.numel() for p in MODEL.parameters())
+print(f"[samudra] model ready ({N_PARAMS / 1e6:.1f}M params)", flush=True)
+
+# --------------------------------------------------------------------------
+# Variable menu
+# --------------------------------------------------------------------------
+LEVELS = [2.5, 10, 22.5, 40, 65, 105, 165, 250, 375, 550, 775, 1050, 1400,
+          1850, 2400, 3100, 4000, 5000, 6000]
+
+VARIABLES: dict[str, dict] = {
+    "Sea surface temperature (2.5 m)": dict(key="thetao_0", unit="°C", cmap="RdYlBu_r"),
+    "Ocean temperature at 105 m": dict(key="thetao_5", unit="°C", cmap="RdYlBu_r"),
+    "Ocean temperature at 550 m": dict(key="thetao_9", unit="°C", cmap="RdYlBu_r"),
+    "Sea surface salinity (2.5 m)": dict(key="so_0", unit="psu", cmap="viridis"),
+    "Sea surface height": dict(key="zos", unit="m", cmap="Spectral_r"),
+    "Zonal velocity (2.5 m)": dict(key="uo_0", unit="m/s", cmap="RdBu_r"),
+    "Meridional velocity (2.5 m)": dict(key="vo_0", unit="m/s", cmap="RdBu_r"),
+}
+for _meta in VARIABLES.values():
+    _meta["idx"] = PROG_NAMES.index(_meta["key"])
+
+
+# --------------------------------------------------------------------------
+# Rollout
+# --------------------------------------------------------------------------
+@functools.lru_cache(maxsize=4)
+def _prepare(i0: int, n_steps: int):
+    """Stream the initial ocean state and atmospheric forcing for a window."""
+    window = SRC.data.isel(time=slice(i0, i0 + 2 * n_steps + 2))
+    windowed_src = dataclasses.replace(SRC, data=window)
+    dataset = InferenceDataset(
+        src=windowed_src,
+        prognostic_var_names=SPEC.prognostic_var_names,
+        boundary_var_names=SPEC.boundary_var_names,
+        hist=HIST,
+        normalize_before_mask=True,
+        masked_fill_value=0.0,
+        long_rollout=False,
+    )
+    initial = dataset.initial_prognostic
+    boundaries = tuple(dataset.get_boundary(k) for k in range(n_steps))
+    target_times = dataset.get_target_time(0, n_steps)
+    return initial, boundaries, dataset.ctx, target_times, window
+
+
+def _gpu_duration(initial, boundaries, ctx, var_idx) -> int:
+    return int(12 + 2 * len(boundaries))
+
+
+@spaces.GPU(duration=_gpu_duration)
+def _rollout(initial, boundaries, ctx, var_idx: int):
+    """Autoregressively step the emulator; returns one normalised field on CPU."""
+    sdev.set_device(torch.device(DEVICE))
+    ctx = ctx.to(torch.device(DEVICE))
+    state = initial.to(DEVICE)
+    outputs = []
+    with torch.no_grad():
+        for boundary in boundaries:
+            state = MODEL.forward_once(state, boundary.to(DEVICE), ctx)
+            outputs.append(state)
+    stacked = torch.cat(outputs, dim=0)
+    stacked = stacked.reshape(-1, N_PROG, stacked.shape[-2], stacked.shape[-1])
+    return stacked[:, var_idx].float().cpu().numpy()
+
+
+# --------------------------------------------------------------------------
+# Rendering
+# --------------------------------------------------------------------------
+def _roll(field: np.ndarray) -> np.ndarray:
+    """Shift longitudes from 0..360 to -180..180 for a familiar world map."""
+    return np.roll(field, field.shape[-1] // 2, axis=-1)
+
+
+def _make_video(pred, truth, dates, label, unit, cmap, fps):
+    diff = pred - truth
+    finite = np.isfinite(truth)
+    vmin, vmax = np.nanpercentile(truth[finite], [1, 99])
+    dmax = float(np.nanpercentile(np.abs(diff[np.isfinite(diff)]), 99.5))
+    dmax = max(dmax, 1e-6)
+    extent = [-180, 180, -90, 90]
+
+    fig, axes = plt.subplots(3, 1, figsize=(7.2, 9.6), dpi=100)
+    artists = []
+    titles = [f"Samudra 2 emulator — {label}", f"GFDL OM4 simulation — {label}",
+              "Difference (emulator − OM4)"]
+    for ax, title, cm, lo, hi in zip(
+        axes, titles, [cmap, cmap, "RdBu_r"], [vmin, vmin, -dmax], [vmax, vmax, dmax]
+    ):
+        colormap = matplotlib.colormaps[cm].copy()
+        colormap.set_bad("0.85")
+        im = ax.imshow(
+            np.zeros_like(pred[0]), origin="lower", extent=extent,
+            cmap=colormap, vmin=lo, vmax=hi, interpolation="nearest",
+        )
+        ax.set_title(title, fontsize=10)
+        ax.set_xticks([-180, -90, 0, 90, 180])
+        ax.set_yticks([-90, -45, 0, 45, 90])
+        ax.tick_params(labelsize=7)
+        fig.colorbar(im, ax=ax, fraction=0.025, pad=0.02).ax.tick_params(labelsize=7)
+        artists.append(im)
+    suptitle = fig.suptitle("", fontsize=12)
+    fig.tight_layout(rect=[0, 0, 1, 0.96])
+
+    frames = []
+    for k, date in enumerate(dates):
+        artists[0].set_data(_roll(pred[k]))
+        artists[1].set_data(_roll(truth[k]))
+        artists[2].set_data(_roll(diff[k]))
+        suptitle.set_text(f"{date}   ·   lead time +{5 * (k + 1)} days   ·   [{unit}]")
+        fig.canvas.draw()
+        frames.append(np.asarray(fig.canvas.buffer_rgba())[..., :3].copy())
+    plt.close(fig)
+
+    frames = frames + [frames[-1]] * max(1, int(fps))
+    path = os.path.join(tempfile.mkdtemp(), "samudra2_rollout.mp4")
+    imageio.mimwrite(
+        path, frames, fps=fps, codec="libx264", quality=8, macro_block_size=16
+    )
+    return path
+
+
+def _make_skill_plot(rmse, persistence, unit, label):
+    leads = 5 * (np.arange(len(rmse)) + 1)
+    fig, ax = plt.subplots(figsize=(7.2, 3.6), dpi=100)
+    ax.plot(leads, persistence, "o--", color="0.55", label="Persistence baseline")
+    ax.plot(leads, rmse, "o-", color="#1f77b4", label="Samudra 2 emulator")
+    ax.set_xlabel("Lead time (days)")
+    ax.set_ylabel(f"RMSE vs OM4 [{unit}]")
+    ax.set_title(f"Emulator error growth — {label}", fontsize=11)
+    ax.grid(alpha=0.3)
+    ax.legend(fontsize=9)
+    ax.set_ylim(bottom=0)
+    fig.tight_layout()
+    path = os.path.join(tempfile.mkdtemp(), "samudra2_skill.png")
+    fig.savefig(path)
+    plt.close(fig)
+    return path
+
+
+# --------------------------------------------------------------------------
+# Handler
+# --------------------------------------------------------------------------
+def forecast(
+    variable: str,
+    start_date: str,
+    n_steps: int = 6,
+    fps: int = 3,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Roll Samudra 2 forward from a held-out ocean state and score it against OM4.
+
+    Args:
+        variable: Which ocean field to visualise (the emulator always advances
+            the full 77-channel ocean state).
+        start_date: Initial condition date, from the held-out 2014-2022 period.
+        n_steps: Number of autoregressive emulator steps; each step advances the
+            ocean by 10 days (two 5-day snapshots).
+        fps: Frames per second of the returned animation.
+
+    Returns:
+        An MP4 animation (emulator / OM4 / difference), an error-growth chart,
+        and a short text summary.
+    """
+    meta = VARIABLES[variable]
+    n_steps = int(n_steps)
+    i0 = TIMES.index(start_date)
+
+    progress(0.05, desc="Streaming ocean state from public OM4 archive…")
+    initial, boundaries, ctx, target_times, window = _prepare(i0, n_steps)
+
+    key = meta["key"]
+    idx = meta["idx"]
+
+    progress(0.45, desc="Running the emulator on GPU…")
+    t_gpu = time.time()
+    pred_norm = _rollout(initial, boundaries, ctx, idx)
+    gpu_seconds = time.time() - t_gpu
+
+    progress(0.7, desc="Fetching the OM4 reference…")
+    wet = SRC.masks.prognostic[idx].numpy().astype(bool)
+    pred = pred_norm.astype(np.float64) * PROG_STD[idx] + PROG_MEAN[idx]
+    pred = np.where(wet[None], pred, np.nan)
+    truth = np.asarray(window[key].sel(time=target_times).values, dtype=np.float32)
+    truth = np.where(wet[None], truth, np.nan)
+    initial_field = np.where(wet, np.asarray(window[key].isel(time=1).values), np.nan)
+
+    rmse = np.sqrt(np.nanmean((pred - truth) ** 2, axis=(1, 2)))
+    persistence = np.sqrt(np.nanmean((initial_field[None] - truth) ** 2, axis=(1, 2)))
+
+    dates = [str(t)[:10] for t in np.asarray(target_times)]
+    progress(0.8, desc="Rendering animation…")
+    video = _make_video(pred, truth, dates, variable, meta["unit"], meta["cmap"], int(fps))
+    plot = _make_skill_plot(rmse, persistence, meta["unit"], variable)
+
+    summary = (
+        f"**Initial condition** {TIMES[i0 + 1]} → **forecast** {dates[0]} … {dates[-1]}  \n"
+        f"**{n_steps} autoregressive steps** = {10 * n_steps} days "
+        f"({len(dates)} 5-day snapshots) in **{gpu_seconds:.1f} s** of GPU time.  \n"
+        f"**{variable}** RMSE vs OM4: {rmse[0]:.3f} {meta['unit']} at +5 days → "
+        f"{rmse[-1]:.3f} {meta['unit']} at +{5 * len(dates)} days "
+        f"(persistence: {persistence[-1]:.3f} {meta['unit']}).  \n"
+        f"Every step advances all {N_PROG} prognostic channels "
+        f"(temperature, salinity, currents over 19 depth levels, plus sea surface height)."
+    )
+    progress(1.0, desc="Done")
+    return video, plot, summary
+
+
+DESCRIPTION = """
+# 🌊 Samudra 2 — a neural ocean emulator
+
+[**Samudra 2**](https://huggingface.co/M2LInES/Samudra2) (M2LInES) is an autoregressive
+deep-learning emulator of the **GFDL OM4 global ocean model**. From one ocean state plus
+the atmospheric forcing, it advances *temperature, salinity, and horizontal currents across
+19 depth levels together with sea surface height* — 77 prognostic fields at once — 10 days
+per step, orders of magnitude faster than the ocean model it replaces.
+
+This Space runs the released **1° checkpoint** on genuinely held-out years (2014–2022):
+the initial condition and forcing are streamed live from the public OM4 archive, the
+emulator is rolled forward on ZeroGPU, and the result is scored against the OM4
+simulation itself.
+"""
+
+with gr.Blocks(title="Samudra 2 ocean emulator") as demo:
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            variable = gr.Dropdown(
+                choices=list(VARIABLES), value="Sea surface temperature (2.5 m)",
+                label="Field to visualise",
+                info="The emulator always advances the full ocean state.",
+            )
+            start_date = gr.Dropdown(
+                choices=START_DATES, value="2015-06-02",
+                label="Initial condition (held-out period)",
+            )
+            n_steps = gr.Slider(
+                1, MAX_STEPS, value=6, step=1,
+                label="Autoregressive steps (10 days each)",
+            )
+            run = gr.Button("Run the ocean emulator", variant="primary")
+            with gr.Accordion("Advanced", open=False):
+                fps = gr.Slider(1, 8, value=3, step=1, label="Animation frames per second")
+            summary = gr.Markdown()
+        with gr.Column(scale=2):
+            video = gr.Video(label="Emulator vs OM4", autoplay=True, loop=True)
+            plot = gr.Image(label="Error growth with lead time", type="filepath")
+
+    inputs = [variable, start_date, n_steps, fps]
+    outputs = [video, plot, summary]
+    run.click(forecast, inputs=inputs, outputs=outputs, api_name="forecast")
+
+    gr.Examples(
+        examples=[
+            ["Sea surface temperature (2.5 m)", "2015-06-02", 6],
+            ["Sea surface height", "2016-02-07", 8],
+            ["Sea surface salinity (2.5 m)", "2018-11-04", 6],
+            ["Ocean temperature at 105 m", "2020-03-18", 10],
+        ],
+        inputs=[variable, start_date, n_steps],
+        outputs=outputs,
+        fn=forecast,
+        cache_examples=True,
+        cache_mode="lazy",
+    )
+
+    gr.Markdown(
+        "Data: public 1° GFDL-OM4 archive (`m2lines-pubs/Samudra`, Open Storage Network). "
+        "Model: [M2LInES/Samudra2](https://huggingface.co/M2LInES/Samudra2) · "
+        "code: [m2lines/Samudra](https://github.com/m2lines/Samudra) (Apache-2.0)."
+    )
+
+if __name__ == "__main__":
+    demo.queue().launch(theme=gr.themes.Citrus(), mcp_server=True)

requirements.txt

@@ -0,0 +1,23 @@
+numpy>=1.26.4,<2
+pandas>=2.2
+scipy>=1.15.2
+matplotlib>=3.10.1
+imageio>=2.34
+imageio-ffmpeg>=0.5
+cftime>=1.6.4.post1
+dacite>=1.9.2
+dask>=2025.2,<2026
+einops>=0.7
+jaxtyping>=0.3
+microsoft-aurora>=1.8
+perceiver-pytorch>=0.8.8
+pydantic-settings>=2.8.1
+pyyaml>=6.0.2
+s3fs==2025.5.1
+torchinfo>=1.8
+tqdm>=4.67.1
+typing-extensions>=4.15
+wandb>=0.19.8
+xarray>=2025.1.2
+xarray-einstats>=0.8
+zarr<3

samudra/__init__.py

@@ -0,0 +1,12 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("samudra")
+except PackageNotFoundError:  # not installed (e.g. running from a bare source tree)
+    __version__ = "0.0.0+unknown"
+
+__all__ = ["__version__"]

samudra/aggregator/LICENSE-ACE.txt

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

samudra/aggregator/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from .inference import InferenceEvaluatorAggregator
+from .loss import (
+    get_channel_loss_dict,
+    get_channel_loss_scale_dict,
+    get_depth_loss_dict,
+    get_variable_loss_dict,
+)
+from .main import Aggregator
+from .validate import ValidateAggregator

samudra/aggregator/inference/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from .main import InferenceEvaluatorAggregator

samudra/aggregator/inference/main.py

@@ -0,0 +1,265 @@
+# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import wandb
+import xarray as xr
+
+from samudra.constants import TensorMap
+from samudra.utils.data import Normalize, get_aggregator_dicts
+from samudra.utils.output import ModelInferenceOutput
+from samudra.utils.wandb import Metrics, MetricsDict
+
+from ..validate.reduced import MeanAggregator as OneStepMeanAggregator
+from .reduced import MeanAggregator
+from .time_mean import TimeMeanEvaluatorAggregator
+
+
+class InferenceEvaluatorAggregator:
+    """
+    Aggregates statistics for inference comparing a generated and target series.
+    """
+
+    def __init__(
+        self,
+        n_timesteps: int,
+        metadata: dict[str, dict[str, str]],
+        hist: int,
+        area_weights: torch.Tensor,
+        wet: torch.Tensor,
+        num_prognostic_channels: int,
+        normalize: Normalize,
+        tensor_map: TensorMap,
+        record_step_20: bool = True,
+        log_global_mean_time_series: bool = True,
+        log_global_mean_norm_time_series: bool = True,
+        time_mean_reference_data: xr.Dataset | None = None,
+        channel_mean_names: list[str] | None = None,
+    ):
+        """
+        Args:
+            n_timesteps: Number of timesteps of inference that will be run.
+            metadata: Mapping of variable names their metadata that will
+                used in generating logged image captions.
+            hist: Number of timesteps of history.
+            area_weights: Area weights for the data.
+            wet: Wet mask for the data.
+            num_prognostic_channels: Number of prognostic channels in the data.
+            normalize: Normalization helper for prognostic channels.
+            tensor_map: Mapping from prognostic variables to tensor channels.
+            record_step_20: Whether to record the mean of the 20th steps.
+            log_global_mean_time_series: Whether to log global mean time series metrics.
+            log_global_mean_norm_time_series: Whether to log the normalized global mean
+                time series metrics.
+            time_mean_reference_data: Reference time means for computing bias stats.
+            channel_mean_names: List of channel names to compute the mean of.
+        """
+        self._aggregators: dict[
+            str, MeanAggregator | OneStepMeanAggregator | TimeMeanEvaluatorAggregator
+        ] = {}
+        self._time_dependent_aggregators: dict[str, TimeMeanEvaluatorAggregator] = {}
+        self._log_time_series = (
+            log_global_mean_time_series or log_global_mean_norm_time_series
+        )
+        if log_global_mean_time_series:
+            self._aggregators["mean"] = MeanAggregator(
+                target="denorm",
+                n_timesteps=n_timesteps,
+                metadata=metadata,
+                area_weights=area_weights,
+            )
+        if log_global_mean_norm_time_series:
+            self._aggregators["mean_norm"] = MeanAggregator(
+                target="norm",
+                n_timesteps=n_timesteps,
+                metadata=metadata,
+                area_weights=area_weights,
+            )
+        if record_step_20:
+            self._aggregators["mean_step_20"] = OneStepMeanAggregator(
+                target_time=20,
+                area_weights=area_weights,
+            )
+        self._aggregators["time_mean"] = TimeMeanEvaluatorAggregator(
+            metadata=metadata,
+            area_weights=area_weights,
+            reference_means=time_mean_reference_data,
+            channel_mean_names=channel_mean_names,
+        )
+        self._aggregators["time_mean_norm"] = TimeMeanEvaluatorAggregator(
+            metadata=metadata,
+            area_weights=area_weights,
+            target="norm",
+            reference_means=time_mean_reference_data,
+            channel_mean_names=channel_mean_names,
+        )
+
+        self._summary_aggregators = {
+            name: agg
+            for name, agg in list(self._aggregators.items())
+            + list(self._time_dependent_aggregators.items())
+            if name not in ["mean", "mean_norm"]
+        }
+        self._n_timesteps_seen = 0
+        self._normalize = normalize
+        self._tensor_map = tensor_map
+        self.num_prognostic_channels = num_prognostic_channels
+        self.hist = hist
+        self.wet = wet
+
+    @property
+    def log_time_series(self) -> bool:
+        return self._log_time_series
+
+    @torch.no_grad()
+    def record_batch(self, data: ModelInferenceOutput):
+        if len(data.prediction) == 0:
+            raise ValueError("No prediction values in data")
+        if len(data.target) == 0:
+            raise ValueError("No target values in data")
+        total_len = len(data.time)
+        assert data.prediction.shape[0] == total_len // (self.hist + 1)
+        target_norm_dict, target_unnorm_dict = get_aggregator_dicts(
+            data.target,
+            normalize=self._normalize,
+            tensor_map=self._tensor_map,
+            wet=self.wet,
+            long_rollout=True,
+            input_type="prognostic",
+            num_prognostic_channels=self.num_prognostic_channels,
+            hist=self.hist,
+        )
+        gen_norm_dict, gen_unnorm_dict = get_aggregator_dicts(
+            data.prediction,
+            normalize=self._normalize,
+            tensor_map=self._tensor_map,
+            wet=self.wet,
+            long_rollout=True,
+            input_type="prognostic",
+            num_prognostic_channels=self.num_prognostic_channels,
+            hist=self.hist,
+        )
+
+        for aggregator in self._aggregators.values():
+            aggregator.record_batch(
+                target_data=target_unnorm_dict,
+                gen_data=gen_unnorm_dict,
+                target_data_norm=target_norm_dict,
+                gen_data_norm=gen_norm_dict,
+                i_time_start=self._n_timesteps_seen,
+            )
+        # TODO: Add time-dependent aggregators
+        # for time_dependent_aggregator in self._time_dependent_aggregators.values():
+        #     time_dependent_aggregator.record_batch(
+        #         time=data.time,
+        #         target_data=target_unnorm_dict,
+        #         gen_data=gen_unnorm_dict,
+        #     )
+        key = list(target_unnorm_dict.keys())[0]
+        n_times = target_unnorm_dict[key].shape[1]
+        logs = self._get_inference_logs_slice(
+            step_slice=slice(self._n_timesteps_seen, self._n_timesteps_seen + n_times),
+        )
+        self._n_timesteps_seen += n_times
+        return logs
+
+    def record_initial_prognostic(
+        self,
+        initial_prognostic: torch.Tensor,
+    ):
+        if self._n_timesteps_seen != 0:
+            raise RuntimeError(
+                "record_initial_condition may only be called once, "
+                "before recording any batches"
+            )
+
+        data_norm_dict, data_unnorm_dict = get_aggregator_dicts(
+            initial_prognostic,
+            normalize=self._normalize,
+            tensor_map=self._tensor_map,
+            wet=self.wet,
+            long_rollout=True,
+            input_type="input",
+            num_prognostic_channels=self.num_prognostic_channels,
+            hist=self.hist,
+        )
+        for aggregator_name in ["mean", "mean_norm"]:
+            aggregator = self._aggregators.get(aggregator_name)
+            if aggregator is not None:
+                aggregator.record_batch(
+                    target_data=data_unnorm_dict,
+                    gen_data=data_unnorm_dict,
+                    target_data_norm=data_norm_dict,
+                    gen_data_norm=data_norm_dict,
+                    i_time_start=0,
+                )
+        key = list(data_unnorm_dict.keys())[0]
+        n_times = data_unnorm_dict[key].shape[1]
+        logs = self._get_inference_logs_slice(
+            step_slice=slice(0, n_times),
+        )
+        self._n_timesteps_seen = n_times
+        return logs
+
+    def get_summary_logs(self):
+        # These aggregators require a full timeseries of
+        # data and thus wandb logged at the end
+        logs: MetricsDict = {}
+        for name, aggregator in self._summary_aggregators.items():
+            logs.update(aggregator.get_logs(label=name))
+        return logs
+
+    @torch.no_grad()
+    def _get_logs(self) -> Metrics:
+        """
+        Returns logs as can be reported to WandB.
+        """
+        logs: MetricsDict = {}
+        for name, aggregator in self._aggregators.items():
+            logs.update(aggregator.get_logs(label=name))
+        for name, time_dependent_aggregator in self._time_dependent_aggregators.items():
+            logs.update(time_dependent_aggregator.get_logs(label=name))
+        return logs
+
+    @torch.no_grad()
+    def _get_inference_logs_slice(self, step_slice: slice):
+        """
+        Returns a subset of the time series for applicable metrics
+        for a specific slice of as can be reported to WandB.
+
+        Args:
+            step_slice: Timestep slice to determine the time series subset.
+
+        Returns:
+            Tuple of start index and list of logs.
+        """
+        logs = {}
+        for name, aggregator in self._aggregators.items():
+            if isinstance(aggregator, MeanAggregator):
+                logs.update(aggregator.get_logs(label=name, step_slice=step_slice))
+        return to_inference_logs(logs)
+
+
+def to_inference_logs(log):
+    # we have a dictionary which contains WandB tables
+    # which we will convert to a list of dictionaries, one for each
+    # row in the tables. Any scalar values will be reported in the last
+    # dictionary.
+    n_rows = 0
+    for val in log.values():
+        if isinstance(val, wandb.Table):
+            n_rows = max(n_rows, len(val.data))
+    logs: list[dict[str, float | int]] = []
+    for i in range(max(1, n_rows)):  # need at least one for non-series values
+        logs.append({})
+    for key, val in log.items():
+        if isinstance(val, wandb.Table):
+            for i, row in enumerate(val.data):
+                for j, col in enumerate(val.columns):
+                    key_without_table_name = key[: key.rfind("/")]
+                    logs[i][f"{key_without_table_name}/{col}"] = row[j]
+        else:
+            logs[-1][key] = val
+    return logs

samudra/aggregator/inference/reduced.py

@@ -0,0 +1,475 @@
+# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import dataclasses
+from collections import defaultdict
+from functools import partial
+from typing import Literal, Protocol
+
+import numpy as np
+import torch
+import wandb
+
+from samudra.aggregator.metrics import (
+    area_weighted_gradient_magnitude_percent_diff,
+    area_weighted_mean,
+    area_weighted_mean_bias,
+    area_weighted_rmse,
+    area_weighted_std,
+)
+from samudra.utils.device import get_device
+from samudra.utils.distributed import all_reduce_mean
+from samudra.utils.wandb import Metrics
+
+
+@dataclasses.dataclass
+class _SeriesData:
+    metric_name: str
+    var_name: str
+    data: np.ndarray
+
+    def get_wandb_key(self) -> str:
+        return f"{self.metric_name}/{self.var_name}"
+
+
+def get_gen_shape(gen_data: dict[str, torch.Tensor]):
+    for name in gen_data:
+        return gen_data[name].shape
+
+
+class MeanMetric(Protocol):
+    def record(self, target: torch.Tensor, gen: torch.Tensor, i_time_start: int):
+        """
+        Update metric for a batch of data.
+        """
+        ...
+
+    def get(self) -> torch.Tensor:
+        """
+        Get the total metric value, not divided by number of recorded batches.
+        """
+        ...
+
+
+class SingleTargetMeanMetric(Protocol):
+    def record(self, tensor: torch.Tensor, i_time_start: int):
+        """
+        Update metric for a batch of data.
+        """
+        ...
+
+    def get(self) -> torch.Tensor:
+        """
+        Get the total metric value, not divided by number of recorded batches.
+        """
+        ...
+
+
+class AreaWeightedFunction(Protocol):
+    """
+    A function that computes a metric on the true and predicted values,
+    weighted by area.
+    """
+
+    def __call__(
+        self,
+        target: torch.Tensor,
+        gen: torch.Tensor,
+    ) -> torch.Tensor: ...
+
+
+class AreaWeightedSingleTargetFunction(Protocol):
+    """
+    A function that computes a metric on a single value, weighted by area.
+    """
+
+    def __call__(
+        self,
+        tensor: torch.Tensor,
+    ) -> torch.Tensor: ...
+
+
+def compute_metric_on(
+    source: Literal["gen", "target"], metric: AreaWeightedSingleTargetFunction
+) -> AreaWeightedFunction:
+    """Turns a single-target metric function
+    (computed on only the generated or target data) into a function that takes in
+    both the generated and target data as arguments, as required for the APIs
+    which call generic metric functions.
+    """
+
+    def metric_wrapper(
+        target: torch.Tensor,
+        gen: torch.Tensor,
+    ) -> torch.Tensor:
+        if source == "gen":
+            return metric(gen)
+        elif source == "target":
+            return metric(target)
+
+    return metric_wrapper
+
+
+class AreaWeightedReducedMetric:
+    """
+    A wrapper around an area-weighted metric function.
+    """
+
+    def __init__(
+        self,
+        device: torch.device,
+        compute_metric: AreaWeightedFunction,
+        n_timesteps: int,
+    ):
+        self._compute_metric = compute_metric
+        self._total: torch.Tensor | None = None
+        self._n_batches = torch.zeros(
+            n_timesteps, dtype=torch.int32, device=get_device()
+        )
+        self._device = device
+        self._n_timesteps = n_timesteps
+
+    def record(self, target: torch.Tensor, gen: torch.Tensor, i_time_start: int):
+        """Add a batch of data to the metric.
+
+        Args:
+            target: Target data. Should have shape [batch, time, height, width].
+            gen: Generated data. Should have shape [batch, time, height, width].
+            i_time_start: The index of the first timestep in the batch.
+        """
+        time_dim = 1
+        if target.shape != gen.shape:
+            raise RuntimeError(
+                "target and gen must have the same shape, got "
+                f"{target.shape} and {gen.shape}"
+            )
+        new_value = (
+            self._compute_metric(target=target, gen=gen).mean(dim=0).to(self._device)
+        )
+        if self._total is None:
+            self._total = torch.zeros(
+                [self._n_timesteps], dtype=new_value.dtype, device=self._device
+            )
+        time_slice = slice(i_time_start, i_time_start + gen.shape[time_dim])
+        self._total[time_slice] += new_value
+        self._n_batches[time_slice] += 1
+
+    def get(self) -> torch.Tensor:
+        """Returns the mean metric across recorded batches."""
+        if self._total is None:
+            return torch.tensor(torch.nan)
+        return self._total / self._n_batches
+
+
+class MeanAggregator:
+    def __init__(
+        self,
+        target: Literal["norm", "denorm"],
+        n_timesteps: int,
+        area_weights: torch.Tensor,
+        metadata: dict[str, dict[str, str]] | None = None,
+    ):
+        self._variable_metrics: dict[str, dict[str, MeanMetric]] | None = None
+        self._shape_x = None
+        self._shape_y = None
+        self._target = target
+        self._n_timesteps = n_timesteps
+        self._area_weights = area_weights
+
+        if metadata is None:
+            self._metadata: dict[str, dict[str, str]] = {}
+        else:
+            self._metadata = metadata
+
+    def _get_variable_metrics(self, gen_data: dict[str, torch.Tensor]):
+        if self._variable_metrics is None:
+            self._variable_metrics = {}
+
+            device = get_device()
+            self._variable_metrics["weighted_rmse"] = defaultdict(
+                lambda: AreaWeightedReducedMetric(
+                    device=device,
+                    compute_metric=partial(
+                        area_weighted_rmse, area_weights=self._area_weights
+                    ),
+                    n_timesteps=self._n_timesteps,
+                )
+            )
+            if self._target == "denorm":
+                self._variable_metrics["weighted_grad_mag_percent_diff"] = defaultdict(
+                    lambda: AreaWeightedReducedMetric(
+                        device=device,
+                        compute_metric=partial(
+                            area_weighted_gradient_magnitude_percent_diff,
+                            area_weights=self._area_weights,
+                        ),  # noqa: E501
+                        n_timesteps=self._n_timesteps,
+                    )
+                )
+            self._variable_metrics["weighted_mean_gen"] = defaultdict(
+                lambda: AreaWeightedReducedMetric(
+                    device=device,
+                    compute_metric=compute_metric_on(
+                        source="gen",
+                        metric=lambda tensor: area_weighted_mean(
+                            tensor, area_weights=self._area_weights
+                        ),
+                    ),
+                    n_timesteps=self._n_timesteps,
+                )
+            )
+            self._variable_metrics["weighted_mean_target"] = defaultdict(
+                lambda: AreaWeightedReducedMetric(
+                    device=device,
+                    compute_metric=compute_metric_on(
+                        source="target",
+                        metric=lambda tensor: area_weighted_mean(
+                            tensor, area_weights=self._area_weights
+                        ),
+                    ),
+                    n_timesteps=self._n_timesteps,
+                )
+            )
+            self._variable_metrics["weighted_bias"] = defaultdict(
+                lambda: AreaWeightedReducedMetric(
+                    device=device,
+                    compute_metric=partial(
+                        area_weighted_mean_bias, area_weights=self._area_weights
+                    ),
+                    n_timesteps=self._n_timesteps,
+                )
+            )
+            self._variable_metrics["weighted_std_gen"] = defaultdict(
+                lambda: AreaWeightedReducedMetric(
+                    device=device,
+                    compute_metric=compute_metric_on(
+                        source="gen",
+                        metric=(
+                            lambda tensor: area_weighted_std(
+                                tensor, area_weights=self._area_weights
+                            )
+                        ),
+                    ),
+                    n_timesteps=self._n_timesteps,
+                )
+            )
+        return self._variable_metrics
+
+    @torch.no_grad()
+    def record_batch(
+        self,
+        target_data: dict[str, torch.Tensor],
+        gen_data: dict[str, torch.Tensor],
+        target_data_norm: dict[str, torch.Tensor],
+        gen_data_norm: dict[str, torch.Tensor],
+        i_time_start: int = 0,
+    ):
+        if self._target == "norm":
+            target_data = target_data_norm
+            gen_data = gen_data_norm
+        variable_metrics = self._get_variable_metrics(gen_data)
+        for name in gen_data.keys():
+            for metric in variable_metrics:
+                variable_metrics[metric][name].record(
+                    target=target_data[name],
+                    gen=gen_data[name],
+                    i_time_start=i_time_start,
+                )
+
+    def _get_series_data(self, step_slice: slice | None = None) -> list[_SeriesData]:
+        """Converts internally stored variable_metrics to a list."""
+        if self._variable_metrics is None:
+            raise ValueError("No batches have been recorded.")
+        data: list[_SeriesData] = []
+        for metric in self._variable_metrics:
+            sorted_keys = sorted(list(self._variable_metrics[metric].keys()))
+            for key in sorted_keys:
+                arr = self._variable_metrics[metric][key].get().detach()
+                if step_slice is not None:
+                    arr = arr[step_slice]
+                datum = _SeriesData(
+                    metric_name=metric,
+                    var_name=key,
+                    data=all_reduce_mean(arr).cpu().numpy(),
+                )
+                data.append(datum)
+        return data
+
+    @torch.no_grad()
+    def get_logs(self, label: str, step_slice: slice | None = None):
+        """
+        Returns logs as can be reported to WandB.
+
+        Args:
+            label: Label to prepend to all log keys.
+            step_slice: Slice of forecast steps to log.
+        """
+        logs = {}
+        series_data: dict[str, np.ndarray] = {
+            datum.get_wandb_key(): datum.data
+            for datum in self._get_series_data(step_slice)
+        }
+        init_step = 0 if step_slice is None else step_slice.start
+        table = data_to_table(series_data, init_step)
+        logs[f"{label}/series"] = table
+        return logs
+
+
+def data_to_table(data: dict[str, np.ndarray], init_step: int = 0) -> wandb.Table:
+    """
+    Convert a dictionary of 1-dimensional timeseries data to a wandb Table.
+
+    Args:
+        data: dictionary of timeseries data.
+        init_step: initial step corresponding to the first row's "forecast_step"
+    """
+    keys = sorted(list(data.keys()))
+    table = wandb.Table(columns=["forecast_step"] + keys)
+    if len(keys) > 0:
+        for i in range(len(data[keys[0]])):
+            row = [init_step + i]
+            for key in keys:
+                row.append(data[key][i])
+            table.add_data(*row)
+    return table
+
+
+class AreaWeightedSingleTargetReducedMetric:
+    """
+    A wrapper around an area-weighted metric function on a single data source.
+    """
+
+    def __init__(
+        self,
+        device: torch.device,
+        compute_metric: AreaWeightedSingleTargetFunction,
+        n_timesteps: int,
+    ):
+        self._compute_metric = compute_metric
+        self._total: torch.Tensor | None = None
+        self._n_batches = torch.zeros(
+            n_timesteps, dtype=torch.int32, device=get_device()
+        )
+        self._device = device
+        self._n_timesteps = n_timesteps
+
+    def record(self, tensor: torch.Tensor, i_time_start: int):
+        """Add a batch of data to the metric.
+
+        Args:
+            tensor: batch data. Should have shape [batch, time, height, width].
+            i_time_start: The index of the first timestep in the batch.
+        """
+        new_value = self._compute_metric(tensor).mean(dim=0).to(self._device)
+        if self._total is None:
+            self._total = torch.zeros(
+                [self._n_timesteps], dtype=new_value.dtype, device=self._device
+            )
+        time_slice = slice(i_time_start, i_time_start + tensor.shape[1])
+        self._total[time_slice] += new_value
+        self._n_batches[time_slice] += 1
+
+    def get(self) -> torch.Tensor:
+        """Returns the mean metric across recorded batches."""
+        if self._total is None:
+            return torch.tensor(torch.nan)
+        return self._total / self._n_batches
+
+
+class SingleTargetMeanAggregator:
+    def __init__(
+        self,
+        n_timesteps: int,
+        area_weights: torch.Tensor,
+        metadata: dict[str, dict[str, str]] | None = None,
+    ):
+        self._variable_metrics: dict[str, dict[str, SingleTargetMeanMetric]] | None = (
+            None
+        )
+        self._shape_x = None
+        self._shape_y = None
+        self._n_timesteps = n_timesteps
+        self._area_weights = area_weights
+        if metadata is None:
+            self._metadata: dict[str, dict[str, str]] = {}
+        else:
+            self._metadata = metadata
+
+    def _get_variable_metrics(self, gen_data: dict[str, torch.Tensor]):
+        if self._variable_metrics is None:
+            self._variable_metrics = {
+                "weighted_mean_gen": {},
+                "weighted_std_gen": {},
+            }
+
+            device = get_device()
+            self._variable_metrics["weighted_mean_gen"] = defaultdict(
+                lambda: AreaWeightedSingleTargetReducedMetric(
+                    device=device,
+                    compute_metric=lambda x: area_weighted_mean(x, self._area_weights),
+                    n_timesteps=self._n_timesteps,
+                )
+            )
+            self._variable_metrics["weighted_std_gen"] = defaultdict(
+                lambda: AreaWeightedSingleTargetReducedMetric(
+                    device=device,
+                    compute_metric=lambda x: area_weighted_std(x, self._area_weights),
+                    n_timesteps=self._n_timesteps,
+                )
+            )
+
+        return self._variable_metrics
+
+    @torch.no_grad()
+    def record_batch(
+        self,
+        data: dict[str, torch.Tensor],
+        i_time_start: int = 0,
+    ):
+        variable_metrics = self._get_variable_metrics(data)
+        for name in data.keys():
+            for metric in variable_metrics:
+                variable_metrics[metric][name].record(
+                    tensor=data[name],
+                    i_time_start=i_time_start,
+                )
+
+    def _get_series_data(self, step_slice: slice | None = None) -> list[_SeriesData]:
+        """Converts internally stored variable_metrics to a list."""
+        if self._variable_metrics is None:
+            raise ValueError("No batches have been recorded.")
+        data: list[_SeriesData] = []
+        for metric in self._variable_metrics:
+            sorted_keys = sorted(list(self._variable_metrics[metric].keys()))
+            for key in sorted_keys:
+                arr = self._variable_metrics[metric][key].get().detach()
+                if step_slice is not None:
+                    arr = arr[step_slice]
+                datum = _SeriesData(
+                    metric_name=metric,
+                    var_name=key,
+                    data=all_reduce_mean(arr).cpu().numpy(),
+                )
+                data.append(datum)
+        return data
+
+    @torch.no_grad()
+    def get_logs(self, label: str, step_slice: slice | None = None) -> Metrics:
+        """
+        Returns logs as can be reported to WandB.
+
+        Args:
+            label: Label to prepend to all log keys.
+            step_slice: Slice of forecast steps to log.
+        """
+        logs = {}
+        series_data: dict[str, np.ndarray] = {
+            datum.get_wandb_key(): datum.data
+            for datum in self._get_series_data(step_slice)
+        }
+        init_step = 0 if step_slice is None else step_slice.start
+        table = data_to_table(series_data, init_step)
+        logs[f"{label}/series"] = table
+        return logs

samudra/aggregator/inference/time_mean.py

@@ -0,0 +1,350 @@
+# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import dataclasses
+from typing import Literal
+
+import matplotlib.pyplot as plt
+import torch
+import wandb
+import xarray as xr
+
+from samudra.aggregator.metrics import area_weighted_mean_bias, area_weighted_rmse
+from samudra.aggregator.plotting import get_cmap_limits, plot_imshow
+from samudra.utils.distributed import all_reduce_mean
+from samudra.utils.wandb import Metrics, MetricsDict
+
+
+@dataclasses.dataclass
+class _TargetGenPair:
+    name: str
+    target: torch.Tensor
+    gen: torch.Tensor
+    area_weights: torch.Tensor
+
+    def bias(self):
+        return self.gen - self.target
+
+    def rmse(self) -> float:
+        ret = float(
+            area_weighted_rmse(
+                gen=self.gen,
+                target=self.target,
+                area_weights=self.area_weights,
+            )
+            .cpu()
+            .numpy()
+        )
+        return ret
+
+    def weighted_mean_bias(self) -> float:
+        return float(
+            area_weighted_mean_bias(
+                gen=self.gen,
+                target=self.target,
+                area_weights=self.area_weights,
+            )
+            .cpu()
+            .numpy()
+        )
+
+
+def get_gen_shape(gen_data: dict[str, torch.Tensor]):
+    for name in gen_data:
+        return gen_data[name].shape
+
+
+class TimeMeanAggregator:
+    _image_captions = {
+        "bias_map": "{name} time-mean bias (generated - reference) [{units}]",
+        "gen_map": "{name} time-mean generated [{units}]",
+    }
+
+    def __init__(
+        self,
+        area_weights: torch.Tensor,
+        target: Literal["norm", "denorm"] = "denorm",
+        metadata: dict[str, dict[str, str]] | None = None,
+        reference_means: xr.Dataset | None = None,
+    ):
+        """
+        Args:
+            area_weights: Tensor of shape [n_lat, n_lon] representing the area weights.
+            target: Whether to compute metrics on the normalized or denormalized data,
+                defaults to "denorm".
+            metadata: Mapping of variable names their metadata that will
+                used in generating logged image captions.
+            reference_means: Dataset containing reference time-mean values
+                for bias computation.
+        """
+        self._target = target
+        if metadata is None:
+            self._metadata: dict[str, dict[str, str]] = {}
+        else:
+            self._metadata = metadata
+        # Dictionaries of tensors of shape [n_lat, n_lon] represnting time means
+        self._data: dict[str, torch.Tensor] | None = None
+        self._n_timesteps = 0
+        self._n_samples: int | None = None
+        self._reference_means = reference_means
+        self._reference_validated = False
+        self._area_weights = area_weights
+
+    @staticmethod
+    def _add_or_initialize_time_mean(
+        maybe_dict: dict[str, torch.Tensor] | None,
+        new_data: dict[str, torch.Tensor],
+        ignore_initial: bool = False,
+    ) -> dict[str, torch.Tensor]:
+        sample_dim = 0
+        time_dim = 1
+        if ignore_initial:
+            time_slice = slice(1, None)
+        else:
+            time_slice = slice(0, None)
+        if maybe_dict is None:
+            d: dict[str, torch.Tensor] = {
+                name: tensor[:, time_slice].sum(dim=time_dim).sum(dim=sample_dim)
+                for name, tensor in new_data.items()
+            }
+        else:
+            d = dict(maybe_dict)
+            for name, tensor in new_data.items():
+                d[name] += tensor[:, time_slice].sum(dim=time_dim).sum(dim=sample_dim)
+        return d
+
+    @torch.no_grad()
+    def record_batch(
+        self,
+        data: dict[str, torch.Tensor],
+        i_time_start: int = 0,
+    ):
+        ignore_initial = i_time_start == 0
+        if ignore_initial:
+            raise Exception(
+                "This is wrong! Make sure this isnt called while recording "
+                "initial prognostic"
+            )
+        self._data = self._add_or_initialize_time_mean(self._data, data, ignore_initial)
+        if self._n_samples is None:
+            self._n_samples = data[list(data)[0]].size(0)
+        if ignore_initial:
+            self._n_timesteps = data[list(data)[0]].size(1) - 1
+        else:
+            self._n_timesteps += data[list(data)[0]].size(1)
+        if not self._reference_validated:
+            if self._reference_means is not None:
+                self.get_logs(label="")
+            self._reference_validated = True
+
+    def get_data(self) -> dict[str, torch.Tensor]:
+        if self._n_timesteps == 0 or self._data is None:
+            raise ValueError("No data recorded.")
+
+        assert self._n_samples is not None  # for type checker
+
+        ret = {}
+        names = sorted(list(self._data.keys()))  # sort for rank-consistent order
+        for name in names:
+            value = self._data[name]
+            gen = all_reduce_mean(value / self._n_timesteps / self._n_samples)
+            ret[name] = gen
+        return ret
+
+    @torch.no_grad()
+    def get_logs(self, label: str) -> Metrics:
+        logs: MetricsDict = {}
+        data = self.get_data()
+        gen_map_key = "gen_map"
+        for name, pred in data.items():
+            vmin_pred, vmax_pred = get_cmap_limits(pred.cpu().numpy())
+            prediction_image = wandb.Image(
+                plot_imshow(pred.cpu().numpy()),
+                caption=self._get_caption(gen_map_key, name, vmin_pred, vmax_pred),
+            )
+            plt.close("all")
+            logs.update(
+                {
+                    f"{gen_map_key}/{name}": prediction_image,
+                }
+            )
+            if self._reference_means is not None and name in self._reference_means:
+                pair = _TargetGenPair(
+                    name=name,
+                    target=torch.as_tensor(
+                        self._reference_means[name].values, device=pred.device
+                    ),
+                    gen=pred,
+                    area_weights=self._area_weights,
+                )
+                bias_map = pair.bias().cpu().numpy()
+                vmin_bias, vmax_bias = get_cmap_limits(bias_map, diverging=True)
+                cmap = plt.colormaps.get_cmap("RdBu_r")
+                cmap.set_bad(color=(0.7, 0.7, 0.7))
+
+                bias_fig: plt.Figure = plot_imshow(
+                    bias_map,
+                    vmin=vmin_bias,
+                    vmax=vmax_bias,
+                    cmap=cmap,
+                    nan_padding=False,
+                )
+                bias_image = wandb.Image(
+                    bias_fig,
+                    caption=self._get_caption("bias_map", name, vmin_bias, vmax_bias),
+                )
+                logs.update({f"ref_bias/{name}": pair.weighted_mean_bias()})
+                logs.update({f"ref_rmse/{name}": pair.rmse()})
+                logs.update({f"ref_bias_map/{name}": bias_image})
+
+        if len(label) != 0:
+            return {f"{label}/{key}": logs[key] for key in logs}
+        return logs
+
+    def _get_caption(self, key: str, name: str, vmin: float, vmax: float) -> str:
+        if name in self._metadata:
+            caption_name = self._metadata[name]["long_name"]
+            units = self._metadata[name]["units"]
+        else:
+            caption_name, units = name, "unknown_units"
+        caption = self._image_captions[key].format(name=caption_name, units=units)
+        caption += f" vmin={vmin:.4g}, vmax={vmax:.4g}."
+        return caption
+
+
+class TimeMeanEvaluatorAggregator:
+    """Statistics and images on the time-mean state.
+
+    This aggregator keeps track of the time-mean state, then computes
+    statistics and images on that time-mean state when logs are retrieved.
+    """
+
+    _image_captions = {
+        "bias_map": "{name} time-mean bias (generated - target) [{units}]",
+        "gen_map": "{name} time-mean generated [{units}]",
+    }
+
+    def __init__(
+        self,
+        area_weights: torch.Tensor,
+        target: Literal["norm", "denorm"] = "denorm",
+        metadata: dict[str, dict[str, str]] | None = None,
+        reference_means: xr.Dataset | None = None,
+        channel_mean_names: list[str] | None = None,
+    ):
+        """
+        Args:
+            area_weights: Tensor of shape [n_lat, n_lon] representing the area weights.
+            target: Whether to compute metrics on the normalized or denormalized data,
+                defaults to "denorm".
+            metadata: Mapping of variable names their metadata that will
+                used in generating logged image captions.
+            reference_means: Dataset containing reference time-mean values
+                for bias computation.
+            channel_mean_names: Names of variables whose RMSE will be averaged. If
+                not provided, all available variables will be used.
+        """
+        self._target = target
+        if metadata is None:
+            self._metadata: dict[str, dict[str, str]] = {}
+        else:
+            self._metadata = metadata
+        self._area_weights = area_weights
+        # Dictionaries of tensors of shape [n_lat, n_lon] represnting time means
+        self._target_agg = TimeMeanAggregator(
+            target=target, metadata=metadata, area_weights=self._area_weights
+        )
+        self._gen_agg = TimeMeanAggregator(
+            target=target,
+            metadata=metadata,
+            reference_means=reference_means,
+            area_weights=self._area_weights,
+        )
+        self._channel_mean_names = channel_mean_names
+
+    @torch.no_grad()
+    def record_batch(
+        self,
+        target_data: dict[str, torch.Tensor],
+        gen_data: dict[str, torch.Tensor],
+        target_data_norm: dict[str, torch.Tensor],
+        gen_data_norm: dict[str, torch.Tensor],
+        i_time_start: int = 0,
+    ):
+        if self._target == "norm":
+            target_data = target_data_norm
+            gen_data = gen_data_norm
+        self._target_agg.record_batch(target_data, i_time_start)
+        self._gen_agg.record_batch(gen_data, i_time_start)
+
+    def _get_target_gen_pairs(self) -> list[_TargetGenPair]:
+        target_data = self._target_agg.get_data()
+        gen_data = self._gen_agg.get_data()
+
+        ret = []
+        for name in gen_data.keys():
+            ret.append(
+                _TargetGenPair(
+                    gen=gen_data[name],
+                    target=target_data[name],
+                    name=name,
+                    area_weights=self._area_weights,
+                )
+            )
+        return ret
+
+    @torch.no_grad()
+    def get_logs(self, label: str) -> Metrics:
+        logs = dict(self._gen_agg.get_logs(""))
+        preds = self._get_target_gen_pairs()
+        bias_map_key = "bias_map"
+        rmse_all_channels = {}
+        for pred in preds:
+            bias_data = pred.bias().cpu().numpy()
+            vmin_bias, vmax_bias = get_cmap_limits(bias_data, diverging=True)
+            cmap = plt.colormaps.get_cmap("RdBu_r")
+            cmap.set_bad(color=(0.7, 0.7, 0.7))
+            bias_fig = plot_imshow(
+                bias_data, vmin=vmin_bias, vmax=vmax_bias, cmap=cmap, nan_padding=False
+            )
+            bias_image = wandb.Image(
+                bias_fig,
+                caption=self._get_caption(
+                    bias_map_key, pred.name, vmin_bias, vmax_bias
+                ),
+            )
+            plt.close("all")
+            rmse_all_channels[pred.name] = pred.rmse()
+            logs.update({f"rmse/{pred.name}": rmse_all_channels[pred.name]})
+            if self._target == "denorm":
+                logs.update(
+                    {
+                        f"{bias_map_key}/{pred.name}": bias_image,
+                        f"bias/{pred.name}": pred.weighted_mean_bias(),
+                    }
+                )
+        if self._target == "norm":
+            metric_name = "rmse/channel_mean"
+            if self._channel_mean_names is None:
+                values_to_average = list(rmse_all_channels.values())
+            else:
+                values_to_average = [
+                    rmse_all_channels[name] for name in self._channel_mean_names
+                ]
+            logs.update({metric_name: sum(values_to_average) / len(values_to_average)})
+
+        if len(label) != 0:
+            return {f"{label}/{key}": logs[key] for key in logs}
+        return logs
+
+    def _get_caption(self, key: str, name: str, vmin: float, vmax: float) -> str:
+        if name in self._metadata:
+            caption_name = self._metadata[name]["long_name"]
+            units = self._metadata[name]["units"]
+        else:
+            caption_name, units = name, "unknown_units"
+        caption = self._image_captions[key].format(name=caption_name, units=units)
+        caption += f" vmin={vmin:.4g}, vmax={vmax:.4g}."
+        return caption

samudra/aggregator/loss.py

@@ -0,0 +1,73 @@
+# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+from samudra.constants import TensorMap
+
+
+def get_depth_loss_dict(
+    label: str,
+    loss_per_channel: torch.Tensor,
+    *,
+    tensor_map: TensorMap,
+) -> dict[str, torch.Tensor]:
+    metrics = {}
+    for depth in tensor_map.DEPTH_SET:
+        metrics[f"{label}/loss/depth/depth_{depth}_loss"] = loss_per_channel[
+            tensor_map.DP_3D_IDX[depth]
+        ].mean()
+    return metrics
+
+
+def get_variable_loss_dict(
+    label: str,
+    loss_per_channel: torch.Tensor,
+    *,
+    tensor_map: TensorMap,
+) -> dict[str, torch.Tensor]:
+    metrics = {}
+    for variable in tensor_map.VAR_SET:
+        metrics[f"{label}/loss/variable/{variable}_loss"] = loss_per_channel[
+            tensor_map.VAR_3D_IDX[variable]
+        ].mean()
+    return metrics
+
+
+def get_channel_loss_dict(
+    label: str,
+    loss_per_channel: torch.Tensor,
+    *,
+    tensor_map: TensorMap,
+    loss_name: str = "loss",
+) -> dict[str, torch.Tensor]:
+    return get_channel_dict(label, loss_name, loss_per_channel, tensor_map=tensor_map)
+
+
+def get_channel_loss_scale_dict(
+    label: str,
+    loss_scale_per_channel: torch.Tensor,
+    *,
+    tensor_map: TensorMap,
+) -> dict[str, torch.Tensor]:
+    return get_channel_dict(
+        label,
+        "loss_scale",
+        loss_scale_per_channel,
+        tensor_map=tensor_map,
+    )
+
+
+def get_channel_dict(
+    prefix: str,
+    measure: str,
+    per_channel: torch.Tensor,
+    *,
+    tensor_map: TensorMap,
+) -> dict[str, torch.Tensor]:
+    metrics = {}
+    for i, channel in enumerate(tensor_map.prognostic_var_names):
+        metrics[f"{prefix}/{measure}/channel/{channel}_{measure}"] = per_channel[i]
+    return metrics

samudra/aggregator/main.py

@@ -0,0 +1,112 @@
+# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+The code in this directory is directly inspired by the ACE project by AI2.
+
+See the original repository at: https://github.com/ai2cm/ace/tree/39133c18524cda85965486ecdc8cb64aac06f4c3/fme/fme/ace/aggregator
+"""
+
+import torch
+
+from samudra.aggregator.inference import InferenceEvaluatorAggregator
+from samudra.aggregator.train import TrainAggregator
+from samudra.aggregator.validate import ValidateAggregator
+from samudra.aggregator.validate.map import MapAggregator
+from samudra.aggregator.validate.reduced import MeanAggregator
+from samudra.aggregator.validate.snapshot import SnapshotAggregator
+from samudra.aggregator.validate.sub_aggregator import ValidateSubAggregator
+from samudra.constants import TensorMap
+from samudra.utils.data import Normalize
+
+
+class Aggregator:
+    @staticmethod
+    def get_train_aggregator(tensor_map: TensorMap) -> TrainAggregator:
+        return TrainAggregator(tensor_map)
+
+    @staticmethod
+    def get_validation_aggregator(
+        metadata: dict[str, dict[str, str]],
+        hist: int,
+        area_weights: torch.Tensor,
+        num_prognostic_channels: int,
+        tensor_map: TensorMap,
+        normalize: Normalize,
+        *,
+        include_image_aggregators: bool = True,
+    ) -> ValidateAggregator:
+        val_aggregators: dict[str, ValidateSubAggregator] = {
+            "reduced": MeanAggregator(area_weights, hist),
+        }
+        if include_image_aggregators:
+            val_aggregators.update(
+                {
+                    "snapshot": SnapshotAggregator(metadata, hist),
+                    "mean_map": MapAggregator(metadata, hist),
+                }
+            )
+
+        return ValidateAggregator(
+            val_aggregators,
+            hist=hist,
+            num_prognostic_channels=num_prognostic_channels,
+            tensor_map=tensor_map,
+            normalize=normalize,
+        )
+
+    @staticmethod
+    def get_inline_inference_aggregator(
+        n_timesteps: int,
+        metadata: dict[str, dict[str, str]],
+        hist: int,
+        area_weights: torch.Tensor,
+        wet: torch.Tensor,
+        num_prognostic_channels: int,
+        tensor_map: TensorMap,
+        normalize: Normalize,
+        channel_mean_names: list[str] | None = None,
+    ) -> InferenceEvaluatorAggregator:
+        return InferenceEvaluatorAggregator(
+            n_timesteps=n_timesteps,
+            metadata=metadata,
+            hist=hist,
+            area_weights=area_weights,
+            wet=wet,
+            num_prognostic_channels=num_prognostic_channels,
+            normalize=normalize,
+            tensor_map=tensor_map,
+            record_step_20=(n_timesteps > 20),
+            log_global_mean_time_series=False,
+            log_global_mean_norm_time_series=False,
+            channel_mean_names=channel_mean_names,
+        )
+
+    @staticmethod
+    def get_standalone_inference_aggregator(
+        n_timesteps: int,
+        metadata: dict[str, dict[str, str]],
+        hist: int,
+        area_weights: torch.Tensor,
+        wet: torch.Tensor,
+        num_prognostic_channels: int,
+        tensor_map: TensorMap,
+        normalize: Normalize,
+        channel_mean_names: list[str] | None = None,
+    ) -> InferenceEvaluatorAggregator:
+        return InferenceEvaluatorAggregator(
+            n_timesteps=n_timesteps,
+            metadata=metadata,
+            hist=hist,
+            area_weights=area_weights,
+            wet=wet,
+            num_prognostic_channels=num_prognostic_channels,
+            normalize=normalize,
+            tensor_map=tensor_map,
+            record_step_20=(n_timesteps > 20),
+            log_global_mean_time_series=True,
+            log_global_mean_norm_time_series=True,
+            channel_mean_names=channel_mean_names,
+        )

samudra/aggregator/metrics.py

@@ -0,0 +1,142 @@
+# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+
+def weighted_sum(
+    tensor: torch.Tensor,
+    weights: torch.Tensor | None = None,
+    dim: tuple[int, ...] = (-2, -1),
+    keepdim: bool = False,
+) -> torch.Tensor:
+    """Computes the weighted sum across the specified list of dimensions.
+
+    Args:
+        tensor: torch.Tensor
+        weights: Weights to apply to the sum.
+        dim: Dimensions to compute the sum over.
+        keepdim: Whether the output tensor has `dim` retained or not.
+
+    Returns:
+        a tensor of the weighted sum averaged over the specified dimensions `dim`.
+    """
+    if weights is None:
+        return tensor.nansum(dim=dim, keepdim=keepdim)
+
+    weights = weights.to(tensor.device)
+    return (tensor * weights).nansum(dim=dim, keepdim=keepdim)
+
+
+def area_weighted_sum(
+    data: torch.Tensor,
+    area_weights: torch.Tensor,
+    dim: tuple[int, ...] = (-2, -1),
+    keepdim: bool = False,
+) -> torch.Tensor:
+    return weighted_sum(data, area_weights, dim=dim, keepdim=keepdim)
+
+
+def weighted_mean(
+    tensor: torch.Tensor,
+    weights: torch.Tensor | None = None,
+    dim: tuple[int, ...] = (-2, -1),
+    keepdim: bool = False,
+) -> torch.Tensor:
+    """Computes the weighted mean across the specified list of dimensions.
+
+    This function treats NaNs as land cells and does not include them in the
+    computation.
+
+    Args:
+        tensor: torch.Tensor
+        weights: Weights to apply to the mean.
+        dim: Dimensions to compute the mean over.
+        keepdim: Whether the output tensor has `dim` retained or not.
+
+    Returns:
+        a tensor of the weighted mean averaged over the specified dimensions `dim`.
+    """
+    if weights is None:
+        return tensor.nanmean(dim=dim, keepdim=keepdim)
+
+    weights = weights.to(tensor.device)
+    denom = torch.where(torch.isnan(tensor), 0.0, weights.expand(tensor.shape)).sum(
+        dim=dim, keepdim=keepdim
+    )
+    return (tensor * weights).nansum(dim=dim, keepdim=keepdim) / denom
+
+
+def area_weighted_mean(
+    data: torch.Tensor,
+    area_weights: torch.Tensor,
+    dim: tuple[int, ...] = (-2, -1),
+    keepdim: bool = False,
+) -> torch.Tensor:
+    return weighted_mean(data, area_weights, dim=dim, keepdim=keepdim)
+
+
+def area_weighted_std(
+    data: torch.Tensor,
+    area_weights: torch.Tensor,
+    keepdim: bool = False,
+):
+    return area_weighted_mean(
+        (data - area_weighted_mean(data, area_weights, keepdim=True)) ** 2,
+        area_weights,
+        keepdim=keepdim,
+    ).sqrt()
+
+
+def area_weighted_rmse(
+    target: torch.Tensor, gen: torch.Tensor, area_weights: torch.Tensor
+) -> torch.Tensor:
+    area_weights = area_weights.to(target.device)
+    return torch.sqrt(area_weighted_mean((gen - target) ** 2, area_weights))
+
+
+def area_weighted_mean_bias(
+    target: torch.Tensor, gen: torch.Tensor, area_weights: torch.Tensor
+) -> torch.Tensor:
+    area_weights = area_weights.to(target.device)
+    return area_weighted_mean(gen - target, area_weights)
+
+
+def gradient_magnitude(
+    tensor: torch.Tensor, dim: tuple[int, ...] = (-2, -1)
+) -> torch.Tensor:
+    """Compute the magnitude of gradient across the specified dimensions."""
+    gradients = torch.gradient(tensor, dim=dim)
+    return torch.sqrt(sum([g**2 for g in gradients]))  # type: ignore
+
+
+def weighted_mean_gradient_magnitude(
+    tensor: torch.Tensor,
+    weights: torch.Tensor | None = None,
+    dim: tuple[int, ...] = (-2, -1),
+) -> torch.Tensor:
+    """Compute weighted mean of gradient magnitude across the specified dimensions."""
+    return weighted_mean(gradient_magnitude(tensor, dim), weights=weights, dim=dim)
+
+
+def gradient_magnitude_percent_diff(
+    target: torch.Tensor,
+    gen: torch.Tensor,
+    weights: torch.Tensor | None = None,
+    dim: tuple[int, ...] = (-2, -1),
+) -> torch.Tensor:
+    """Compute the percent difference of the weighted mean gradient magnitude across
+    the specified dimensions.
+    """
+    target_grad_mag = weighted_mean_gradient_magnitude(target, weights, dim)
+    gen_grad_mag = weighted_mean_gradient_magnitude(gen, weights, dim)
+    return 100 * (gen_grad_mag - target_grad_mag) / target_grad_mag
+
+
+def area_weighted_gradient_magnitude_percent_diff(
+    target: torch.Tensor, gen: torch.Tensor, area_weights: torch.Tensor
+):
+    area_weights = area_weights.to(target.device)
+    return gradient_magnitude_percent_diff(target, gen, weights=area_weights)

samudra/aggregator/plotting.py

@@ -0,0 +1,140 @@
+# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import gc
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.colors import Colormap
+from matplotlib.figure import Figure
+
+from samudra.utils.wandb import WandBLogger
+
+
+def get_cmap_limits(data: np.ndarray, diverging=False) -> tuple[float, float]:
+    vmin = np.nanmin(data)
+    vmax = np.nanmax(data)
+    if diverging:
+        vmax = max(abs(vmin), abs(vmax))
+        vmin = -vmax
+    return vmin, vmax
+
+
+def plot_imshow(
+    data: np.ndarray,
+    vmin: float | None = None,
+    vmax: float | None = None,
+    cmap: Colormap | None = None,
+    flip_lat: bool = True,
+    use_colorbar: bool = True,
+    nan_padding: bool = True,
+) -> Figure:
+    """Plot a 2D array using imshow, ensuring figure size is same as array size."""
+    min_ = np.nanmin(data) if vmin is None else vmin
+    max_ = np.nanmax(data) if vmax is None else vmax
+
+    if flip_lat:
+        lat_dim = -2
+        data = np.flip(data, axis=lat_dim)
+
+    if use_colorbar:
+        height, width = data.shape
+        colorbar_width = max(1, int(0.025 * width))
+        range_ = np.linspace(min_, max_, height)
+        range_ = np.repeat(range_[:, np.newaxis], repeats=colorbar_width, axis=1)
+        range_ = np.flipud(range_)  # wandb images start from top (and left)
+        padding = np.zeros((height, colorbar_width))
+        if nan_padding:
+            padding = padding + np.nan  # Set when using non-diverging map
+        data = np.concatenate((data, padding, range_), axis=1)
+
+    # make figure size (in pixels) be the same as array size
+    figsize = np.array(data.T.shape) / plt.rcParams["figure.dpi"]
+    fig = Figure(figsize=figsize)  # create directly for cleanup when it leaves scope
+    ax = fig.add_axes((0, 0, 1, 1))
+    ax.imshow(data, cmap=cmap, vmin=vmin, vmax=vmax)
+    ax.set_axis_off()
+    return fig
+
+
+def plot_paneled_data(
+    data: list[list[np.ndarray]],
+    diverging: bool,
+    caption: str | None = None,
+):
+    """Plot a list of 2D data arrays in a paneled plot."""
+    if diverging:
+        cmap = plt.colormaps.get_cmap("RdBu_r")
+        cmap.set_bad(color=(0.7, 0.7, 0.7))
+    else:
+        cmap = plt.colormaps.get_cmap("viridis")
+        cmap.set_bad(color="white")
+    vmin = np.inf
+    vmax = -np.inf
+    for row in data:
+        for arr in row:
+            vmin = min(vmin, np.nanmin(arr))
+            vmax = max(vmax, np.nanmax(arr))
+    if diverging:
+        vmax = max(abs(vmin), abs(vmax))
+        vmin = -vmax
+    if caption is not None:
+        caption += " "
+    else:
+        caption = ""
+
+    caption += f"vmin={vmin:.4g}, vmax={vmax:.4g}."
+
+    if diverging:
+        fill_value = 0.5 * (vmin + vmax)
+    else:
+        fill_value = vmin
+    all_data = _stitch_data_panels(data, fill_value=fill_value)
+
+    fig = plot_imshow(
+        all_data, vmin=vmin, vmax=vmax, cmap=cmap, nan_padding=not diverging
+    )
+    wandb = WandBLogger.get_instance()
+    wandb_image = wandb.Image(fig, caption=caption)
+    plt.close(fig)
+
+    # necessary to avoid CUDA error in some contexts
+    # see https://github.com/ai2cm/full-model/issues/740#issuecomment-2086546187
+    gc.collect()
+
+    return wandb_image
+
+
+def _stitch_data_panels(data: list[list[np.ndarray]], fill_value) -> np.ndarray:
+    for row in data:
+        if len(row) != len(data[0]):
+            raise ValueError("All rows must have the same number of panels.")
+
+    n_rows = len(data)
+    n_cols = len(data[0])
+    for row in data:
+        for arr in row:
+            if arr.shape != data[0][0].shape:
+                raise ValueError("All panels must have the same shape.")
+
+    stitched_data = np.full(
+        (
+            n_rows * data[0][0].shape[0] + n_rows - 1,
+            n_cols * data[0][0].shape[1] + n_cols - 1,
+        ),
+        fill_value=fill_value,
+    )
+
+    # iterate over rows backwards, as the image starts in the bottom left
+    # and moves upwards
+    for i, row in enumerate(reversed(data)):
+        for j, arr in enumerate(row):
+            start_row = i * (arr.shape[0] + 1)
+            end_row = start_row + arr.shape[0]
+            start_col = j * (arr.shape[1] + 1)
+            end_col = start_col + arr.shape[1]
+            stitched_data[start_row:end_row, start_col:end_col] = arr
+
+    return stitched_data

samudra/aggregator/train.py

@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+from samudra.aggregator.loss import (
+    get_channel_loss_dict,
+    get_depth_loss_dict,
+    get_variable_loss_dict,
+)
+from samudra.constants import TensorMap
+from samudra.utils.distributed import all_reduce_mean
+from samudra.utils.output import TrainBatchOutput
+from samudra.utils.wandb import Metrics
+
+
+class TrainAggregator:
+    """Aggregates train statistics for an epoch."""
+
+    def __init__(self, tensor_map: TensorMap):
+        self.tensor_map = tensor_map
+        self._n_batches = 0
+        self._loss = torch.tensor(torch.nan)
+        self._loss_per_channel = torch.tensor(torch.nan)
+
+    @torch.no_grad()
+    def record_batch(self, batch: TrainBatchOutput):
+        if self._n_batches == 0:
+            self._loss = batch.loss
+            self._loss_per_channel = batch.loss_per_channel
+        else:
+            self._loss += batch.loss
+            self._loss_per_channel += batch.loss_per_channel
+        self._n_batches += 1
+
+    @torch.no_grad()
+    def get_logs(self, label: str = "train") -> Metrics:
+        loss = self._loss / self._n_batches
+
+        loss_per_channel = self._loss_per_channel / self._n_batches
+        depth_loss_dict = get_depth_loss_dict(
+            label, loss_per_channel, tensor_map=self.tensor_map
+        )
+        var_loss_dict = get_variable_loss_dict(
+            label, loss_per_channel, tensor_map=self.tensor_map
+        )
+        channel_loss_dict = get_channel_loss_dict(
+            label, loss_per_channel, tensor_map=self.tensor_map
+        )
+        logs = {
+            f"{label}/mean/loss": loss,
+            **depth_loss_dict,
+            **var_loss_dict,
+            **channel_loss_dict,
+        }
+        meaned_logs = {}
+        for key in sorted(logs.keys()):
+            meaned_logs[key] = float(all_reduce_mean(logs[key].detach()).cpu().numpy())
+
+        return meaned_logs

samudra/aggregator/validate/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from .main import ValidateAggregator

samudra/aggregator/validate/main.py

@@ -0,0 +1,117 @@
+# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+from samudra.aggregator.train import TrainAggregator
+from samudra.aggregator.validate.sub_aggregator import ValidateSubAggregator
+from samudra.constants import TensorMap
+from samudra.utils.data import Normalize, get_aggregator_dicts
+from samudra.utils.output import ValBatchOutput
+from samudra.utils.wandb import Metrics, MetricsDict
+
+
+class ValidateAggregator(TrainAggregator):
+    """Aggregates Validation Statistics."""
+
+    def __init__(
+        self,
+        aggregators: dict[str, ValidateSubAggregator],
+        hist: int,
+        num_prognostic_channels: int,
+        *,
+        tensor_map: TensorMap,
+        normalize: Normalize,
+    ):
+        super().__init__(tensor_map)
+        self._aggregators = aggregators
+        self.hist = hist
+        self.num_prognostic_channels = num_prognostic_channels
+        self.normalize = normalize
+
+    # TODO(jder): we could remove this by moving from inheritance
+    # to composition with the TrainAggregator functionality.
+    def record_batch(self, batch):
+        raise NotImplementedError(
+            "Call record_validation_batch instead of record_batch"
+        )
+
+    @torch.no_grad()
+    def record_validation_batch(self, batch: ValBatchOutput):
+        super().record_batch(batch)  # Record losses
+
+        # If there are no log aggregators, omit doing any extra work.
+        if not self._aggregators:
+            return
+
+        # Translate the GridContext mask by removing history.
+        target_data = batch.target_data  # [B, C*(hist+1), H, W]
+        wet = batch.ctx.label_mask  # [C*(hist+1), H, W]
+        assert wet.shape == target_data.shape[1:], (
+            "The wetmask must match the target data shape excluding batch."
+        )
+        assert wet.shape[0] % (self.hist + 1) == 0, (
+            "The wetmask channel count must be divisible by history size."
+        )
+        first_wetmask_chunk = wet.shape[0] // (self.hist + 1)
+        wet = wet[:first_wetmask_chunk]  # [C, H, W]
+
+        if len(target_data) == 0:
+            raise ValueError("No data in target_data")
+        if len(batch.gen_data) == 0:
+            raise ValueError("No data in gen_data")
+
+        assert target_data.shape[1] == self.num_prognostic_channels
+        target_data_dict, target_data_unnorm_dict = get_aggregator_dicts(
+            target_data,
+            normalize=self.normalize,
+            tensor_map=self.tensor_map,
+            wet=wet,
+            long_rollout=False,
+            input_type="prognostic",
+            num_prognostic_channels=self.num_prognostic_channels,
+            hist=self.hist,
+        )
+
+        gen_data_dict, gen_data_unnorm_dict = get_aggregator_dicts(
+            batch.gen_data,
+            normalize=self.normalize,
+            tensor_map=self.tensor_map,
+            wet=wet,
+            long_rollout=False,
+            input_type="prognostic",
+            num_prognostic_channels=self.num_prognostic_channels,
+            hist=self.hist,
+        )
+        input_data_dict, input_data_unnorm_dict = get_aggregator_dicts(
+            batch.input_data,
+            normalize=self.normalize,
+            tensor_map=self.tensor_map,
+            wet=wet,
+            long_rollout=False,
+            input_type="input",
+            num_prognostic_channels=self.num_prognostic_channels,
+            hist=self.hist,
+        )
+
+        for agg in self._aggregators.values():
+            agg.record_batch(
+                loss=batch.loss,
+                target_data=target_data_unnorm_dict,
+                gen_data=gen_data_unnorm_dict,
+                input_data=input_data_unnorm_dict,
+                target_data_norm=target_data_dict,
+                gen_data_norm=gen_data_dict,
+                input_data_norm=input_data_dict,
+            )
+
+    @torch.no_grad()
+    def get_logs(self, label: str = "train") -> Metrics:
+        logs: MetricsDict = dict(super().get_logs(label))
+        for agg_label in self._aggregators:
+            for k, v in self._aggregators[agg_label].get_logs(label=agg_label).items():
+                logs[f"{label}/{k}"] = v
+
+        return logs

samudra/aggregator/validate/map.py

@@ -0,0 +1,122 @@
+# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+from samudra.aggregator.plotting import plot_paneled_data
+from samudra.aggregator.validate.sub_aggregator import ValidateSubAggregator
+from samudra.utils.distributed import all_reduce_mean, is_main_process
+
+
+class MapAggregator(ValidateSubAggregator):
+    """
+    An aggregator that records the average over batches as function of lat and lon.
+    """
+
+    _captions = {
+        "full-field": (
+            "{name} one step mean full field; "
+            "(top) generated and (bottom) target [{units}]"
+        ),
+        "error": (
+            "{name} one step mean full field error (generated - target) [{units}]"
+        ),
+    }
+
+    def __init__(
+        self, metadata: dict[str, dict[str, str]] | None = None, hist: int = 0
+    ):
+        """
+        Args:
+            metadata: Mapping of variable names their metadata that will
+                used in generating logged image captions.
+            hist: Number of history steps to include in the snapshot.
+        """
+        if metadata is None:
+            metadata = {}
+        else:
+            self._metadata = metadata
+        self._n_batches = 0
+        self._target_data: dict[str, torch.Tensor] = {}
+        self._gen_data: dict[str, torch.Tensor] = {}
+        self.hist = hist
+
+    @torch.no_grad()
+    def record_batch(
+        self,
+        loss,
+        target_data,
+        gen_data,
+        input_data,
+        target_data_norm,
+        gen_data_norm,
+        input_data_norm,
+    ):
+        self._loss = loss
+        for name in target_data:
+            if name in self._target_data:
+                self._target_data[name] += target_data[name].mean(dim=0)
+            else:
+                self._target_data[name] = target_data[name].mean(dim=0)
+        for name in gen_data:
+            if name in self._gen_data:
+                self._gen_data[name] += gen_data[name].mean(dim=0)
+            else:
+                self._gen_data[name] = gen_data[name].mean(dim=0)
+        self._n_batches += 1
+
+    @torch.no_grad()
+    def get_logs(self, label: str):
+        """
+        Returns logs as can be reported to WandB.
+
+        Args:
+            label: Label to prepend to all log keys.
+        """
+        time_dim = 0
+        target_time = self.hist  # Use latest time step
+        image_logs = {}
+        render_images = is_main_process()
+        sorted_names = sorted(list(self._gen_data.keys()))
+        for name in sorted_names:
+            # use first sample in batch
+            gen = (
+                all_reduce_mean(
+                    self._gen_data[name].select(dim=time_dim, index=target_time)
+                )
+                / self._n_batches
+            )
+            target = (
+                all_reduce_mean(
+                    self._target_data[name].select(dim=time_dim, index=target_time)
+                )
+                / self._n_batches
+            )
+            if not render_images:
+                continue
+            image_logs[f"image-error/{name}"] = plot_paneled_data(
+                [[(gen - target).cpu().numpy()]],
+                diverging=True,
+                caption=self._get_caption("error", name),
+            )
+            image_logs[f"image-full-field/{name}"] = plot_paneled_data(
+                [
+                    [gen.cpu().numpy()],
+                    [target.cpu().numpy()],
+                ],
+                diverging=False,
+                caption=self._get_caption("full-field", name),
+            )
+        image_logs = {f"{label}/{key}": image_logs[key] for key in image_logs}
+        return image_logs
+
+    def _get_caption(self, key: str, name: str) -> str:
+        if name in self._metadata:
+            caption_name = self._metadata[name]["long_name"]
+            units = self._metadata[name]["units"]
+        else:
+            caption_name, units = name, "unknown_units"
+        caption = self._captions[key].format(name=caption_name, units=units)
+        return caption

samudra/aggregator/validate/reduced.py

@@ -0,0 +1,165 @@
+# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Callable
+from functools import partial
+
+import numpy as np
+import torch
+
+from samudra.aggregator.metrics import (
+    area_weighted_gradient_magnitude_percent_diff,
+    area_weighted_mean_bias,
+    area_weighted_rmse,
+)
+from samudra.aggregator.validate.sub_aggregator import ValidateSubAggregator
+from samudra.utils.device import get_device
+from samudra.utils.distributed import all_reduce_mean
+
+
+class AreaWeightedReducedMetric:
+    """
+    A wrapper around an area-weighted metric function.
+    """
+
+    def __init__(
+        self,
+        device: torch.device,
+        compute_metric: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+    ):
+        self._compute_metric = compute_metric
+        self._total: torch.Tensor | None = None
+        self._device = device
+
+    def record(self, target: torch.Tensor, gen: torch.Tensor):
+        """Add a batch of data to the metric.
+
+        Args:
+            target: Target data. Should have shape [batch, time, height, width].
+            gen: Generated data. Should have shape [batch, time, height, width].
+        """
+        new_value = self._compute_metric(target, gen).mean(dim=0)
+        if self._total is None:
+            self._total = torch.zeros_like(new_value, device=self._device)
+        self._total += new_value
+
+    def get(self) -> torch.Tensor:
+        """Returns the metric."""
+        assert self._total is not None
+        return self._total
+
+
+class MeanAggregator(ValidateSubAggregator):
+    """
+    Aggregator for mean-reduced metrics.
+
+    These are metrics such as means which reduce to a single float for each batch,
+    and then can be averaged across batches to get a single float for the
+    entire dataset. This is important because the aggregator uses the mean to combine
+    metrics across batches and processors.
+    """
+
+    def __init__(self, area_weights: torch.Tensor, target_time: int):
+        self._n_batches = 0
+        self._variable_metrics: dict | None = None
+        self._target_time = target_time
+        self._area_weights = area_weights
+
+    def _get_variable_metrics(self, gen_data):
+        if self._variable_metrics is None:
+            self._variable_metrics = {
+                "weighted_rmse": {},
+                "weighted_bias": {},
+                "weighted_grad_mag_percent_diff": {},
+            }
+            device = get_device()
+            for key in gen_data:
+                self._variable_metrics["weighted_rmse"][key] = (
+                    AreaWeightedReducedMetric(
+                        device=device,
+                        compute_metric=partial(
+                            area_weighted_rmse, area_weights=self._area_weights
+                        ),
+                    )
+                )
+                self._variable_metrics["weighted_bias"][key] = (
+                    AreaWeightedReducedMetric(
+                        device=device,
+                        compute_metric=partial(
+                            area_weighted_mean_bias, area_weights=self._area_weights
+                        ),
+                    )
+                )
+                self._variable_metrics["weighted_grad_mag_percent_diff"][key] = (
+                    AreaWeightedReducedMetric(
+                        device=device,
+                        compute_metric=partial(
+                            area_weighted_gradient_magnitude_percent_diff,
+                            area_weights=self._area_weights,
+                        ),  # noqa: E501
+                    )
+                )
+
+        return self._variable_metrics
+
+    @torch.no_grad()
+    def record_batch(
+        self,
+        *,
+        loss: torch.Tensor = torch.tensor(np.nan),
+        target_data,
+        gen_data,
+        target_data_norm,
+        gen_data_norm,
+        input_data: dict[str, torch.Tensor] | None = None,
+        input_data_norm: dict[str, torch.Tensor] | None = None,
+        i_time_start: int = 0,
+    ):
+        if input_data is None:
+            input_data = {}
+        if input_data_norm is None:
+            input_data_norm = {}
+
+        variable_metrics = self._get_variable_metrics(gen_data)
+        time_dim = 1
+        time_len = gen_data[list(gen_data.keys())[0]].shape[time_dim]
+        target_time = self._target_time - i_time_start
+        if target_time >= 0 and time_len > target_time:
+            for name in gen_data.keys():
+                target = target_data[name].select(dim=time_dim, index=target_time)
+                gen = gen_data[name].select(dim=time_dim, index=target_time)
+                for metric in variable_metrics:
+                    variable_metrics[metric][name].record(
+                        target=target,
+                        gen=gen,
+                    )
+            # only increment n_batches if we actually recorded a batch
+            self._n_batches += 1
+
+    def _get_data(self):
+        if self._variable_metrics is None or self._n_batches == 0:
+            raise ValueError("No batches have been recorded.")
+        data: dict[str, torch.Tensor] = {}
+        for metric in self._variable_metrics:
+            for key in self._variable_metrics[metric]:
+                data[f"{metric}/{key}"] = (
+                    self._variable_metrics[metric][key].get() / self._n_batches
+                )
+        meaned_data: dict[str, float] = {}
+        for key in sorted(data.keys()):
+            meaned_data[key] = float(all_reduce_mean(data[key].detach()).cpu().numpy())
+        return meaned_data
+
+    @torch.no_grad()
+    def get_logs(self, label: str):
+        """
+        Returns logs as can be reported to WandB.
+
+        Args:
+            label: Label to prepend to all log keys.
+        """
+        return {
+            f"{label}/{key}": data for key, data in sorted(self._get_data().items())
+        }

samudra/aggregator/validate/snapshot.py

@@ -0,0 +1,114 @@
+# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+from samudra.aggregator.plotting import plot_paneled_data
+from samudra.aggregator.validate.sub_aggregator import ValidateSubAggregator
+from samudra.utils.distributed import is_main_process
+from samudra.utils.wandb import Metrics
+
+
+class SnapshotAggregator(ValidateSubAggregator):
+    """
+    An aggregator that records the first sample of the last batch of data.
+    """
+
+    _captions = {
+        "full-field": (
+            "{name} one step full field for last sample; "
+            "(top) generated and (bottom) target [{units}]"
+        ),
+        "residual": (
+            "{name} one step residual (prediction - previous time) for last sample; "
+            "(top) generated and (bottom) target [{units}]"
+        ),
+        "error": (
+            "{name} one step full field error (generated - target) "
+            "for last sample [{units}]"
+        ),
+    }
+
+    def __init__(
+        self, metadata: dict[str, dict[str, str]] | None = None, hist: int = 0
+    ):
+        """
+        Args:
+            metadata: Mapping of variable names their metadata that will
+                used in generating logged image captions.
+            hist: Number of history steps to include in the snapshot.
+        """
+        if metadata is None:
+            metadata = {}
+        else:
+            self._metadata = metadata
+        self.hist = hist
+
+    @torch.no_grad()
+    def record_batch(
+        self,
+        loss: torch.Tensor,
+        target_data,
+        gen_data,
+        input_data,
+        target_data_norm,
+        gen_data_norm,
+        input_data_norm,
+    ):
+        self._loss = loss
+        self._target_data = target_data
+        self._gen_data = gen_data
+        self._target_data_norm = target_data_norm
+        self._gen_data_norm = gen_data_norm
+        self._input_data = input_data
+        self._input_data_norm = input_data_norm
+
+    @torch.no_grad()
+    def get_logs(self, label: str) -> Metrics:
+        """
+        Returns logs as can be reported to WandB.
+
+        Args:
+            label: Label to prepend to all log keys.
+        """
+        if not is_main_process():
+            return {}
+
+        time_dim = 1
+        target_time = 0  # first output time step
+        input_time = self.hist  # last input time step
+        image_logs = {}
+        for name in self._gen_data.keys():
+            # use first sample in batch
+            gen = self._gen_data[name].select(dim=time_dim, index=target_time)[0].cpu()
+            target = (
+                self._target_data[name].select(dim=time_dim, index=target_time)[0].cpu()
+            )
+            input = (
+                self._input_data[name].select(dim=time_dim, index=input_time)[0].cpu()
+            )
+            images = {}
+            images["error"] = [[(gen - target).numpy()]]
+            images["full-field"] = [[gen.numpy()], [target.numpy()]]
+            images["residual"] = [[(gen - input).numpy()], [(target - input).numpy()]]
+            for key, data in images.items():
+                if key == "error" or key == "residual":
+                    diverging = True
+                else:
+                    diverging = False
+                caption = self._get_caption(key, name)
+                wandb_image = plot_paneled_data(data, diverging, caption=caption)
+                image_logs[f"image-{key}/{name}"] = wandb_image
+        image_logs = {f"{label}/{key}": image_logs[key] for key in image_logs}
+        return image_logs
+
+    def _get_caption(self, key: str, name: str) -> str:
+        if name in self._metadata:
+            caption_name = self._metadata[name]["long_name"]
+            units = self._metadata[name]["units"]
+        else:
+            caption_name, units = name, "unknown_units"
+        caption = self._captions[key].format(name=caption_name, units=units)
+        return caption

samudra/aggregator/validate/sub_aggregator.py

@@ -0,0 +1,28 @@
+# SPDX-FileCopyrightText: 2024 Allen Institute for Artificial Intelligence
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+
+import torch
+
+from samudra.utils.wandb import Metrics
+
+
+class ValidateSubAggregator(ABC):
+    @abstractmethod
+    def get_logs(self, label: str) -> Metrics: ...
+
+    @abstractmethod
+    def record_batch(
+        self,
+        *,
+        loss: torch.Tensor,
+        target_data,
+        gen_data,
+        input_data,
+        target_data_norm,
+        gen_data_norm,
+        input_data_norm,
+    ): ...

samudra/backend.py

@@ -0,0 +1,107 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+import os
+
+from samudra.config import DistributedConfig, EvalBackendConfig, TrainBackendConfig
+from samudra.utils.device import set_device
+from samudra.utils.distributed import init_distributed_mode
+
+
+def _cuda_diagnostics() -> str:
+    """Return a short string that helps debug why CUDA is (not) available."""
+    visible = os.getenv("CUDA_VISIBLE_DEVICES", None)
+    return (
+        "CUDA diagnostics: "
+        f"torch={torch.__version__}, "
+        f"torch.version.cuda={torch.version.cuda}, "
+        f"torch.cuda.is_available()={torch.cuda.is_available()}, "
+        f"torch.cuda.device_count()={torch.cuda.device_count()}, "
+        f"CUDA_VISIBLE_DEVICES={visible!r}"
+    )
+
+
+def init_train_backend(
+    backend: TrainBackendConfig,
+) -> tuple[torch.device, DistributedConfig | None]:
+    """Given backend config, get the device and (if any) distributed configuration."""
+    match backend:
+        case "cpu":
+            device = torch.device("cpu")
+            dist_cfg = None
+        case "cuda":
+            if not torch.cuda.is_available():
+                raise RuntimeError(
+                    "Requested backend='cuda' but CUDA is not available. "
+                    + _cuda_diagnostics()
+                )
+            device = torch.device("cuda")
+            dist_cfg = None
+        case "nccl":
+            if not torch.cuda.is_available():
+                raise RuntimeError(
+                    "Requested backend='nccl' but CUDA is not available. "
+                    + _cuda_diagnostics()
+                )
+            device = torch.device("cuda")
+            dist_cfg = init_distributed_mode()
+        case "auto" if torch.cuda.is_available():
+            logger.info("auto backend detected CUDA")
+            device = torch.device("cuda")
+            try:
+                dist_cfg = init_distributed_mode()
+                logger.info("succeeded in initializing distributed mode")
+            except RuntimeError as e:
+                logger.info(
+                    f"Failed to initialize distributed mode, running on single GPU.",
+                    exc_info=e,
+                )
+                dist_cfg = None
+        case "auto":
+            logger.warning(
+                "auto backend: cuda not found, using CPU. " + _cuda_diagnostics()
+            )
+            device = torch.device("cpu")
+            dist_cfg = None
+        case _:
+            raise ValueError(f"Invalid backend: {backend}")
+
+    # We set this globally so we don't need to hand the device around.
+    # See https://github.com/m2lines/Samudra/issues/87.
+    set_device(device)
+
+    return device, dist_cfg
+
+
+def init_eval_backend(backend: EvalBackendConfig) -> torch.device:
+    """Given evaluation backend config, get the device."""
+    match backend:
+        case "cpu":
+            device = torch.device("cpu")
+        case "cuda":
+            if not torch.cuda.is_available():
+                raise RuntimeError(
+                    "Requested backend='cuda' but CUDA is not available. "
+                    + _cuda_diagnostics()
+                )
+            device = torch.device("cuda")
+        case "auto" if torch.cuda.is_available():
+            logger.info("auto backend detected CUDA")
+            device = torch.device("cuda")
+        case "auto":
+            logger.warning(
+                "auto backend: cuda not found, using CPU. " + _cuda_diagnostics()
+            )
+            device = torch.device("cpu")
+        case _:
+            raise ValueError(f"Invalid backend: {backend}")
+
+    set_device(device)
+    return device

samudra/cli.py

@@ -0,0 +1,77 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Console-script entry point: ``samudra <train|eval|viz> CONFIG [OVERRIDES...]``.
+
+Installed as the ``samudra`` command (see ``[project.scripts]`` in
+pyproject.toml), so a user who ``pip install samudra`` can run
+``samudra train path/to/config.yaml`` without cloning the repo. Each
+subcommand forwards to the same ``main`` the module entry points use
+(``python -m samudra.train`` etc.), so behaviour is identical.
+"""
+
+import sys
+
+_COMMANDS = ("train", "eval", "viz")
+
+_HELP = """\
+samudra — train and evaluate emulators of ocean physics
+
+Samudra is a PyTorch package for models that auto-regressively predict future
+ocean states — temperature, salinity, horizontal currents, sea-surface height,
+and surface heat flux — learned from the OM4 ocean model at 1°, 1/2°, and 1/4°
+resolution. See https://arxiv.org/abs/2412.03795 for the method.
+
+Usage:
+  samudra <command> CONFIG [--KEY VALUE ...]
+
+Commands:
+  train   Train a model from a config (checkpointing, W&B logging, multi-GPU).
+  eval    Roll a trained model out autoregressively and collect metrics.
+  viz     Render maps, time series, and PDFs from evaluation outputs.
+
+CONFIG is a path to a YAML file or the name of a bundled preset such as
+`samudra_om4/train.yaml`. Any config key can be overridden inline (e.g.
+`--epochs 100`); run `samudra <command> --help` for a command's full options.
+
+Examples:
+  samudra train samudra_om4/train.yaml --experiment.data_root ./data
+  samudra eval  samudra_om4/eval.yaml  --ckpt_path ./checkpoint.pt
+
+Docs: https://m2lines.github.io/Samudra/docs/
+"""
+
+
+def main() -> None:
+    """Dispatch ``samudra <command> ...`` to the matching task entry point."""
+    argv = sys.argv[1:]
+    if not argv or argv[0] in ("-h", "--help"):
+        sys.stdout.write(_HELP)
+        raise SystemExit(0)
+    if argv[0] not in _COMMANDS:
+        sys.stderr.write(f"samudra: unknown command {argv[0]!r}\n\n{_HELP}")
+        raise SystemExit(2)
+
+    command, rest = argv[0], argv[1:]
+    # The config loaders parse sys.argv directly, so strip the subcommand token
+    # and leave them an argv of `CONFIG [OVERRIDES...]`.
+    sys.argv = [f"samudra {command}", *rest]
+
+    if command == "train":
+        from samudra.train import main as train_main
+
+        train_main()
+    elif command == "eval":
+        from samudra.eval import main as eval_main
+
+        eval_main()
+    else:  # viz
+        from samudra.viz.config import VizConfig
+        from samudra.viz.config import main as viz_main
+
+        viz_main(VizConfig.from_yaml_and_cli())
+
+
+if __name__ == "__main__":
+    main()

samudra/config.py

@@ -0,0 +1,1339 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import abc
+import datetime
+from functools import cached_property
+from pathlib import Path
+from typing import Annotated, Literal, Self, assert_never
+
+import cftime
+import numpy as np
+import pandas as pd
+import pydantic
+import torch
+import xarray as xr
+from perceiver_pytorch import Perceiver as NaivePerceiver
+from pydantic import (
+    Field,
+    PlainSerializer,
+    PlainValidator,
+    WithJsonSchema,
+    model_validator,
+)
+from torch import nn
+from torch.nn import GELU
+
+from samudra.config_base import BaseConfig, TopLevelConfig
+from samudra.constants import (
+    DatasetSpec,
+    GridType,
+    LoaderVersion,
+    build_llc_spec,
+    build_om4_spec,
+)
+from samudra.models import Samudra, SamudraMini, SamudraMulti
+from samudra.models.base import BaseModel
+from samudra.models.modules import (
+    AvgPool,
+    BilinearUpsample,
+    CappedGELU,
+    ConvBlock,
+    ConvNeXtBlock,
+    CoreBlock,
+    CoreBlockBuilder,
+    MaxPool,
+    PerceiverDecoder,
+    PerceiverEncoder,
+    ReLU,
+    TransposedConvUpsample,
+    UNetBackbone,
+)
+from samudra.models.modules.augment_input import (
+    Concat3dCoordinates,
+    FourierFeatures2D,
+    fourier_features_2d_dim,
+)
+from samudra.models.modules.blocks import ZonallyPeriodicBilinearUpsample
+from samudra.models.modules.encoder import patch_from
+from samudra.utils.data import DataContainer, DataSource, DataSourceSplits
+from samudra.utils.llc import canonicalize_llc_datasets
+from samudra.utils.location import (
+    LocalLocation,
+    Location,
+    ResolvedLocation,
+    UnresolvedLocation,
+)
+from samudra.utils.loss import (
+    DynamicLoss,
+    GradientLoss,
+    LossFnWithContext,
+    LossMetric,
+    loss_fn_from_metric,
+)
+from samudra.utils.profiler import Profiler
+from samudra.utils.schedule import SchedulerConfig
+
+
+class WandBConfig(BaseConfig):
+    mode: Literal["online", "disabled"] = "disabled"
+    project: str = "default"
+    entity: str = "ocean_emulators"
+    group: str | None = None
+    tags: list[str] | None = None
+    notes: str | None = None
+
+
+class JulianDate:
+    """Represents a Julian date as a cftime.datetime at noon on the relevant day.
+
+    This is the format the OM4 data uses, so we match that here.
+    TODO(jder): probably worth asserting the date format when opening the data.
+    """
+
+    datetime: cftime.datetime
+
+    def __init__(self, value: str):
+        parsed = cftime.datetime.strptime(value, "%Y-%m-%d", calendar="julian")
+        self.datetime = parsed.replace(hour=12)
+
+    def __str__(self) -> str:
+        return self.datetime.strftime("%Y-%m-%d")
+
+
+def _julian_date_validator(value: str | JulianDate) -> JulianDate:
+    return JulianDate(value) if isinstance(value, str) else value
+
+
+JulianDateConfig = Annotated[
+    JulianDate,
+    PlainValidator(_julian_date_validator),
+    PlainSerializer(JulianDate.__str__),
+    WithJsonSchema({"type": "string", "format": "date"}),
+]
+
+
+# We reuse pydantic's AwareDatetime to match JSONSchema's expected RFC3339 format
+_aware_datetime_adapter = pydantic.TypeAdapter(pydantic.AwareDatetime)
+
+
+def _datetime64_validator(value: str | np.datetime64) -> np.datetime64:
+    if isinstance(value, np.datetime64):
+        return value.astype("datetime64[ns]")
+    parsed = _aware_datetime_adapter.validate_python(value)
+    utc = parsed.astimezone(datetime.UTC).replace(tzinfo=None)
+    return np.datetime64(utc, "ns")
+
+
+def _serialize_datetime64(value: np.datetime64) -> str:
+    return str(np.datetime_as_string(value, unit="ns", timezone="UTC"))
+
+
+LlcDatetimeConfig = Annotated[
+    np.datetime64,
+    PlainValidator(_datetime64_validator),
+    PlainSerializer(_serialize_datetime64),
+    WithJsonSchema({"type": "string", "format": "date-time"}),
+]
+
+
+class Om4TimeConfig(BaseConfig):
+    model_config = pydantic.ConfigDict(frozen=True)
+
+    start: JulianDateConfig
+    end: JulianDateConfig
+
+    @property
+    def time_slice(self) -> slice:
+        return slice(self.start.datetime, self.end.datetime)
+
+    def overlaps(self, other: Self) -> bool:
+        return (
+            self.start.datetime < other.end.datetime
+            and self.end.datetime > other.start.datetime
+        )
+
+    def __str__(self) -> str:
+        return f"{self.start} to {self.end}"
+
+
+class LlcTimeConfig(BaseConfig):
+    model_config = pydantic.ConfigDict(frozen=True)
+
+    start: LlcDatetimeConfig
+    end: LlcDatetimeConfig
+
+    @property
+    def time_slice(self) -> slice:
+        return slice(self.start, self.end)
+
+    def overlaps(self, other: Self) -> bool:
+        return bool(self.start < other.end and self.end > other.start)
+
+    def __str__(self) -> str:
+        return f"{self.start} to {self.end}"
+
+
+TimeConfig = Om4TimeConfig | LlcTimeConfig
+
+
+LOCATION_DOCS = (
+    "Use a string relative to the `data_root` or use a structured location "
+    "see location.py for possible types."
+)
+
+
+class BaseDataSourceConfig[SourceTimeConfigT: TimeConfig](BaseConfig, abc.ABC):
+    train_time: SourceTimeConfigT = Field(frozen=True)
+    val_time: SourceTimeConfigT = Field(frozen=True)
+    inference_times: tuple[SourceTimeConfigT, ...] = Field(default=(), frozen=True)
+    data_location: Location = Field(
+        description="Location of the data; " + LOCATION_DOCS
+    )
+    data_means_location: Location = Field(
+        description="Location of the data means; " + LOCATION_DOCS
+    )
+    data_stds_location: Location = Field(
+        description="Location of the data standard deviations; " + LOCATION_DOCS
+    )
+
+    @property
+    @abc.abstractmethod
+    def dataset_spec(self) -> DatasetSpec:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def canonicalize_datasets(
+        self,
+        data: xr.Dataset,
+        means: xr.Dataset,
+        stds: xr.Dataset,
+    ) -> tuple[xr.Dataset, xr.Dataset, xr.Dataset]:
+        raise NotImplementedError
+
+    def build(
+        self,
+        data_root: ResolvedLocation,
+        *,
+        use_dask: bool,
+        is_primary: bool,
+    ) -> DataSourceSplits:
+        source = self._build_source(
+            data_root,
+            turn_on_dask=use_dask,
+        )
+        inference_source = None
+        if is_primary and self.inference_times:
+            if use_dask:
+                full_inference_source = source
+            else:
+                full_inference_source = self._build_source(
+                    data_root,
+                    turn_on_dask=True,
+                )
+            # TODO: remove multiple inference time ranges altogether (see #813)
+            assert len(self.inference_times) == 1, (
+                "multiple inference time ranges have been deprecated"
+            )
+            inference_source = full_inference_source.slice(self.inference_times[0])
+
+        return DataSourceSplits(
+            train=source.slice(self.train_time),
+            val=source.slice(self.val_time),
+            inference=inference_source,
+        )
+
+    def _build_source(
+        self,
+        data_root: ResolvedLocation,
+        *,
+        turn_on_dask: bool,
+    ) -> DataSource:
+        resolved_data_location = data_root.resolve(self.data_location)
+        resolved_means_location = data_root.resolve(self.data_means_location)
+        resolved_stds_location = data_root.resolve(self.data_stds_location)
+
+        chunks: dict[str, int] | None = {} if turn_on_dask else None
+        data = resolved_data_location.open(chunks)
+        means = resolved_means_location.open(chunks)
+        stds = resolved_stds_location.open(chunks)
+        data, means, stds = self.canonicalize_datasets(
+            data,
+            means,
+            stds,
+        )
+        dataset_spec = self.dataset_spec
+
+        source = DataSource.from_datasets(
+            data,
+            means,
+            stds,
+            dataset_spec=dataset_spec,
+            prognostic_var_names=dataset_spec.prognostic_var_names,
+            boundary_var_names=dataset_spec.boundary_var_names,
+            name=f"{resolved_data_location}-{turn_on_dask}",
+        )
+        return source
+
+
+class BaseDataLoadingConfig(BaseConfig):
+    def num_pytorch_workers(self) -> int:
+        raise NotImplementedError
+
+    def persistent_pytorch_workers(self) -> bool:
+        raise NotImplementedError
+
+
+class CpuDataLoadingConfig(BaseDataLoadingConfig):
+    type: Literal["cpu"] = "cpu"
+    num_workers: int = Field(default=4, ge=0)
+    persistent_workers: bool = True
+
+    def num_pytorch_workers(self) -> int:
+        return self.num_workers
+
+    def persistent_pytorch_workers(self) -> bool:
+        return self.persistent_workers
+
+
+class GpuDataLoadingConfig(BaseDataLoadingConfig):
+    type: Literal["gpu"] = "gpu"
+    kvikio_task_size: int = Field(default=64 * 1024 * 1024, gt=0)
+    kvikio_num_threads: int = Field(default=8, gt=0)
+
+    def num_pytorch_workers(self) -> int:
+        # When loading data direct to GPU, we don't want worker processes.
+        # 0 means "load in the main process"
+        return 0
+
+    def persistent_pytorch_workers(self) -> bool:
+        return False
+
+
+DataLoadingConfig = Annotated[
+    CpuDataLoadingConfig | GpuDataLoadingConfig,
+    Field(discriminator="type"),
+]
+
+
+class Om4DataSourceConfig(BaseDataSourceConfig[Om4TimeConfig]):
+    type: Literal["om4"] = "om4"
+    prognostic_vars_key: str = "thermo_dynamic_all"
+    boundary_vars_key: str = "tau_hfds"
+    grid_type: GridType = "gaussian"
+
+    @property
+    def dataset_spec(self) -> DatasetSpec:
+        return build_om4_spec(
+            self.prognostic_vars_key,
+            self.boundary_vars_key,
+            grid_type=self.grid_type,
+        )
+
+    @pydantic.model_validator(mode="after")
+    def validate_time_splits(self) -> Self:
+        if self.train_time.overlaps(self.val_time):
+            raise ValueError(
+                f"Training time range {self.train_time} overlaps "
+                f"with validation time range {self.val_time}"
+            )
+        return self
+
+    def canonicalize_datasets(
+        self,
+        data: xr.Dataset,
+        means: xr.Dataset,
+        stds: xr.Dataset,
+    ) -> tuple[xr.Dataset, xr.Dataset, xr.Dataset]:
+        # TODO: move OM4 canonicalization in here instead of in validation afte.
+        return data, means, stds
+
+
+class LlcDataSourceConfig(BaseDataSourceConfig[LlcTimeConfig]):
+    type: Literal["llc"] = "llc"
+    prognostic_vars_key: str = "single_1"
+    boundary_vars_key: str = "single_1"
+    face: int = Field(default=1, ge=0)
+    i_start: int = Field(default=0, ge=0)
+    i_end: int = Field(default=720, gt=0)
+    j_start: int = Field(default=0, ge=0)
+    j_end: int = Field(default=720, gt=0)
+
+    @property
+    def dataset_spec(self) -> DatasetSpec:
+        return build_llc_spec(
+            self.prognostic_vars_key,
+            self.boundary_vars_key,
+        )
+
+    @pydantic.model_validator(mode="after")
+    def validate_time_splits(self) -> Self:
+        if self.train_time.overlaps(self.val_time):
+            raise ValueError(
+                f"Training time range {self.train_time} overlaps "
+                f"with validation time range {self.val_time}"
+            )
+        return self
+
+    @pydantic.model_validator(mode="after")
+    def validate_crop_bounds(self) -> Self:
+        if self.i_end <= self.i_start:
+            raise ValueError("LLC crop bounds must satisfy i_start < i_end")
+        if self.j_end <= self.j_start:
+            raise ValueError("LLC crop bounds must satisfy j_start < j_end")
+        return self
+
+    def canonicalize_datasets(
+        self,
+        data: xr.Dataset,
+        means: xr.Dataset,
+        stds: xr.Dataset,
+    ) -> tuple[xr.Dataset, xr.Dataset, xr.Dataset]:
+        return canonicalize_llc_datasets(
+            data,
+            means,
+            stds,
+            face=self.face,
+            i_start=self.i_start,
+            i_end=self.i_end,
+            j_start=self.j_start,
+            j_end=self.j_end,
+            dataset_spec=self.dataset_spec,
+        )
+
+
+DataSourceConfig = Annotated[
+    Om4DataSourceConfig | LlcDataSourceConfig,
+    Field(discriminator="type"),
+]
+
+
+class DataConfig(BaseConfig):
+    sources: list[DataSourceConfig] = Field(
+        description=(
+            "Data sources to include, each with explicit data/means/stds "
+            "locations. These are resolved relative to data_root."
+        ),
+        min_length=1,
+    )
+    loading: DataLoadingConfig = Field(default_factory=CpuDataLoadingConfig)
+    hist: int = 1
+    loader_version: str = str(LoaderVersion.OM4_TORCH.value)
+    normalize_before_mask: bool = True
+    masked_fill_value: float = 0.0
+    concurrent_compute: bool = False
+
+    def build(
+        self,
+        data_root: ResolvedLocation,
+    ) -> DataContainer:
+        loader_version = LoaderVersion(self.loader_version)
+        use_dask = loader_version != LoaderVersion.OM4_TORCH
+        dataset_spec = self.sources[0].dataset_spec
+        assert all(source.dataset_spec == dataset_spec for source in self.sources), (
+            "All data sources must use the same dataset spec"
+        )
+
+        source_splits = [
+            source_cfg.build(
+                data_root,
+                use_dask=use_dask,
+                is_primary=index == 0,
+            )
+            for index, source_cfg in enumerate(self.sources)
+        ]
+        train_sources = [splits.train for splits in source_splits]
+        val_sources = [splits.val for splits in source_splits]
+
+        return DataContainer(
+            train_sources=train_sources,
+            val_sources=val_sources,
+            inference_source=source_splits[0].inference,
+            loader_version=loader_version,
+            dataset_spec=dataset_spec,
+        )
+
+
+BlockType = Literal["conv_next_block", "conv_block"]
+ActivationType = Literal["relu", "gelu", "capped_gelu"]
+NormType = Literal["batch", "instance", "layer"]
+
+
+class BlockConfig(BaseConfig):
+    block_type: BlockType = "conv_next_block"
+    kernel_size: int = 3
+    activation: ActivationType = "capped_gelu"
+    upscale_factor: int = 4
+    norm: NormType = "batch"
+    pointwise_linear: bool = False
+
+    def build(self) -> CoreBlockBuilder:
+        match self.activation:
+            case "relu":
+                activation: type[nn.Module] = ReLU
+            case "capped_gelu":
+                activation = CappedGELU
+            case "gelu":
+                activation = GELU
+            case _:
+                assert_never(self.activation)
+
+        def create_block(
+            in_channels: int,
+            out_channels: int,
+            dilation: int,
+            n_layers: int,
+            pad: str,
+            checkpoint_simple: bool,
+        ) -> CoreBlock:
+            match self.block_type:
+                case "conv_block":
+                    return ConvBlock(
+                        in_channels=in_channels,
+                        out_channels=out_channels,
+                        dilation=dilation,
+                        n_layers=n_layers,
+                        pad=pad,
+                        checkpoint_simple=checkpoint_simple,
+                        kernel_size=self.kernel_size,
+                        activation=activation,
+                    )
+                case "conv_next_block":
+                    return ConvNeXtBlock(
+                        in_channels=in_channels,
+                        out_channels=out_channels,
+                        dilation=dilation,
+                        n_layers=n_layers,
+                        pad=pad,
+                        checkpoint_simple=checkpoint_simple,
+                        kernel_size=self.kernel_size,
+                        upscale_factor=self.upscale_factor,
+                        norm=self.norm,
+                        activation=activation,
+                        pointwise_linear=self.pointwise_linear,
+                    )
+                case _:
+                    assert_never(self.block_type)
+
+        return create_block
+
+
+PerceiverImpl = Literal["auto", "naive", "flash"]
+
+
+class PerceiverConfig(BaseConfig):
+    """A standard config interface to various perceiver implementations.
+
+    Builds either a regular Perceiver (for the encoder, via ``build``) or a
+    PerceiverIO (for the decoder, via ``build_io``).  Both respect the shared
+    ``implementation`` setting from ``SamudraMultiConfig.perceiver_implementation``.
+    """
+
+    depth: int = 6
+    latent_dim: int = Field(
+        default=128,
+        description="The small, latent dimension of the Perceiver. This is the `N` dimension for the Perceiver's `O(M*N)` complexity",
+    )
+    num_latents: int = Field(
+        default=512,
+        description="The number of latent vectors in the Perceiver. This is the `M` dimension for the Perceiver's `O(M*N)` complexity",
+    )
+
+    def build(
+        self,
+        in_channels: int,
+        out_channels: int,
+        max_patch_size: tuple[int, int],
+        implementation: PerceiverImpl,
+    ) -> nn.Module:
+        """Build a regular Perceiver (used by the encoder)."""
+        # This is not really a "frequency" but a maximum of the width appears to be reasonable from looking at the code.
+        max_freq = max(*max_patch_size)
+
+        num_freq_bands = 4
+        fourier_dim = fourier_features_2d_dim(num_freq_bands)
+        # Use the same explicit 2D Fourier features in both implementations so
+        # intra-patch positions are encoded equivalently.
+        if _use_flash(implementation):
+            try:
+                from flash_perceiver import Perceiver as FlashPerceiver  # type: ignore
+            except ImportError as e:
+                raise _flash_import_error() from e
+            from einops.layers.torch import Rearrange
+
+            perceiver: nn.Module = nn.Sequential(
+                FourierFeatures2D(num_freq_bands=num_freq_bands, max_freq=max_freq),
+                Rearrange("b ph pw v -> b (ph pw) v"),
+                FlashPerceiver(
+                    depth=self.depth,
+                    input_dim=in_channels + fourier_dim,
+                    output_dim=out_channels,
+                    output_mode="average",
+                    latent_dim=self.latent_dim,
+                    num_latents=self.num_latents,
+                    use_flash_attn=True,
+                    weight_tie_layers=True,
+                    self_per_cross_attn=2,
+                ),
+            )
+        elif _use_naive(implementation):
+            perceiver = nn.Sequential(
+                FourierFeatures2D(num_freq_bands=num_freq_bands, max_freq=max_freq),
+                NaivePerceiver(
+                    # Required by perceiver-pytorch even when its internal
+                    # Fourier encoding is disabled below.
+                    num_freq_bands=num_freq_bands,
+                    max_freq=max_freq,
+                    depth=self.depth,
+                    input_axis=2,
+                    input_channels=in_channels + fourier_dim,
+                    num_classes=out_channels,
+                    latent_dim=self.latent_dim,
+                    num_latents=self.num_latents,
+                    weight_tie_layers=True,
+                    fourier_encode_data=False,
+                    self_per_cross_attn=2,
+                ),
+            )
+        else:
+            raise ValueError(f"Unknown perceiver implementation: {implementation}.")
+
+        return perceiver
+
+    def build_io(
+        self,
+        in_channels: int,
+        queries_dim: int,
+        out_channels: int,
+        implementation: PerceiverImpl,
+    ) -> nn.Module:
+        """Build a PerceiverIO (used by the decoder)."""
+        if _use_flash(implementation):
+            try:
+                from flash_perceiver.perceiver import (  # type: ignore
+                    PerceiverIO as FlashPerceiverIO,  # type: ignore
+                )
+            except ImportError as e:
+                raise _flash_import_error() from e
+            perceiver_io: nn.Module = FlashPerceiverIO(
+                depth=self.depth,
+                input_dim=in_channels,
+                query_dim=queries_dim,
+                proj_dim=out_channels,
+                num_latents=self.num_latents,
+                latent_dim=self.latent_dim,
+                use_flash_attn=True,
+                weight_tie_layers=True,
+            )
+        elif _use_naive(implementation):
+            from perceiver_pytorch.perceiver_io import PerceiverIO as NaivePerceiverIO
+
+            perceiver_io = NaivePerceiverIO(
+                depth=self.depth,
+                dim=in_channels,
+                queries_dim=queries_dim,
+                logits_dim=out_channels,
+                num_latents=self.num_latents,
+                latent_dim=self.latent_dim,
+                weight_tie_layers=True,
+                decoder_ff=True,
+            )
+        else:
+            raise ValueError(f"Unknown perceiver implementation: {implementation}.")
+
+        return perceiver_io
+
+
+def _use_flash(implementation: PerceiverImpl) -> bool:
+    return (
+        implementation == "auto" and torch.cuda.is_available()
+    ) or implementation == "flash"
+
+
+def _use_naive(implementation: PerceiverImpl) -> bool:
+    return (
+        implementation == "auto" and not torch.cuda.is_available()
+    ) or implementation == "naive"
+
+
+def _flash_import_error() -> ValueError:
+    return ValueError(
+        "`implementation==flash` or flash was automatically chosen for `implementation==auto`, "
+        "but the flash attention dependencies could not be imported. "
+        "Please run `uv sync --extra cuda` or specify the `naive` attention implementation."
+    )
+
+
+class EncoderConfig(BaseConfig):
+    perceiver: PerceiverConfig = PerceiverConfig()
+
+    def build(
+        self,
+        in_channels: int,
+        out_channels: int,
+        patch_extent: tuple[float, float],
+        max_lat_size: int,
+        max_lon_size: int,
+        implementation: PerceiverImpl,
+    ) -> PerceiverEncoder:
+        max_patch_size = patch_from(patch_extent, max_lat_size, max_lon_size)
+        return PerceiverEncoder(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            patch_extent=patch_extent,
+            perceiver=self.perceiver.build(
+                in_channels, out_channels, max_patch_size, implementation
+            ),
+        )
+
+
+class DecoderConfig(BaseConfig):
+    """A PerceiverIO-based decoder configuration.
+
+    Uses PerceiverIO (with an explicit query mechanism) rather than a regular
+    Perceiver.  Output pixel positions are encoded as queries, so the output
+    size is determined by the query count — not by ``num_latents``.
+
+    When ``window_patches`` is set, the decoder tiles the output grid into
+    spatial blocks of that many patches per side.  Each block's PerceiverIO
+    call receives only the overlapping latent tokens plus ``context_patches``
+    extra rings of neighbors, keeping cost bounded even when the latent grid
+    is large (i.e. fine ``patch_extent``).
+    """
+
+    perceiver: PerceiverConfig = PerceiverConfig()
+    queries_dim: int = Field(
+        default=64,
+        description="Embedding dimension for pixel-position queries in the PerceiverIO decoder head.",
+    )
+    window_patches: int | None = Field(
+        default=4096,
+        description="Side length (in patches) of each spatial decode window. "
+        "None = decode all patches at once (global attention). "
+        "E.g. window_patches=8 means each PerceiverIO call covers an 8x8 block of patches.",
+    )
+    context_patches: int | None = Field(
+        default=1,
+        description="Number of extra patch rings around each window to include as data context. "
+        "Only used when window_patches is set. None = full context (every window sees all latent tokens).",
+    )
+
+    def build(
+        self,
+        in_channels: int,
+        out_channels: int,
+        patch_extent: tuple[float, float],
+        implementation: PerceiverImpl,
+    ) -> PerceiverDecoder:
+        return PerceiverDecoder(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            patch_extent=patch_extent,
+            queries_dim=self.queries_dim,
+            perceiver_io=self.perceiver.build_io(
+                in_channels, self.queries_dim, out_channels, implementation
+            ),
+            window_patches=self.window_patches,
+            context_patches=self.context_patches,
+        )
+
+
+DownSamplingBlocks = Literal["avg_pool", "max_pool"]
+UpSamplingBlocks = Literal[
+    "bilinear_upsample", "transposed_conv", "zonally_periodic_upsample"
+]
+Checkpointing = Literal["all", "simple"]
+
+
+class UNetBackboneConfig(BaseConfig):
+    ch_width: list[int] = [200, 250, 300, 400]
+    dilation: list[int] = [1, 2, 4, 8]
+    n_layers: list[int] = [1, 1, 1, 1]
+    core_block: BlockConfig = BlockConfig()
+    down_sampling_block: DownSamplingBlocks = "avg_pool"
+    up_sampling_block: UpSamplingBlocks = "zonally_periodic_upsample"
+    drop_path_rate: float = Field(
+        default=0.0,
+        description="Shortcut dropout rate. The chance we turn off skip connections in the UNet. Reasonable values are 0.1-0.3. Use 0.0 to disable.",
+    )
+
+    def build(
+        self,
+        in_channels: int,
+        pad: str,
+        checkpointing: Checkpointing | None,
+    ) -> UNetBackbone:
+        assert len(self.ch_width) == len(self.dilation) == len(self.n_layers), (
+            "`ch_width`, `dilation`, and `n_layers` must have the same length."
+        )
+
+        def create_upsampling_block(in_channels: int, out_channels: int):
+            match self.up_sampling_block:
+                case "bilinear_upsample":
+                    return BilinearUpsample(
+                        in_channels=in_channels, out_channels=out_channels
+                    )
+                case "transposed_conv":
+                    return TransposedConvUpsample(
+                        in_channels=in_channels, out_channels=out_channels
+                    )
+                case "zonally_periodic_upsample":
+                    return ZonallyPeriodicBilinearUpsample()
+                case _:
+                    assert_never(self.up_sampling_block)
+
+        match self.down_sampling_block:
+            case "avg_pool":
+                downsampling_block: nn.Module = AvgPool()
+            case "max_pool":
+                downsampling_block = MaxPool()
+            case _:
+                assert_never(self.down_sampling_block)
+
+        return UNetBackbone(
+            in_channels=in_channels,
+            ch_width=self.ch_width,
+            dilation=self.dilation,
+            n_layers=self.n_layers,
+            pad=pad,
+            create_block=self.core_block.build(),
+            downsampling_block=downsampling_block,
+            create_upsampling_block=create_upsampling_block,
+            checkpointing=checkpointing,
+            drop_path_rate=self.drop_path_rate,
+        )
+
+
+class BaseModelConfig(BaseConfig, abc.ABC):
+    pred_residuals: bool = False
+    last_kernel_size: int = 3
+    pad: str = "circular"
+
+    checkpointing: Checkpointing | None = Field(
+        default=None,
+        description="""Strategy for storing activations for the model for use in
+        the backward pass. If not set, the model will store all activations in memory
+        (fast but lots of memory). If set to 'all', the model will recompute each
+        top-level layer (CoreBlocks, scaling layers, etc.) in the backward pass.
+        If set to 'simple', the model will recompute only cheap layers like scales
+        and nonlinearities.""",
+    )
+
+    gradient_detach_interval: int = Field(
+        default=0,
+        description="""Interval for detaching gradients in autoregressive training. `0` means no detaching.""",
+    )
+
+    add_3d_coordinates: bool = Field(
+        default=False,
+        description="Add 3d coordinates representing position on the Earth (cartesian coordinates on a unit sphere) to the input channels.",
+    )
+
+    @abc.abstractmethod
+    def build(
+        self,
+        prog_channels: int,
+        boundary_channels: int,
+        out_channels: int,
+        hist: int,
+        srcs: list[DataSource],
+    ) -> BaseModel:
+        pass
+
+
+class SamudraConfig(BaseModelConfig):
+    unet: UNetBackboneConfig = UNetBackboneConfig()
+    pos_channels: int = Field(
+        default=0,
+        description="""Number of channels used for a learned positional embedding""",
+    )
+    use_bfloat16: bool = Field(
+        default=False,
+        description="Use bfloat16 for most layers rather than float32.",
+    )
+
+    def build(
+        self,
+        prog_channels: int,
+        boundary_channels: int,
+        out_channels: int,
+        hist: int,
+        srcs: list[DataSource],
+    ) -> Samudra:
+        if len(srcs) != 1:
+            raise ValueError(
+                "Samudra only supports training at a single scale! Please configure exactly one data source."
+            )
+        src = srcs[0]
+        in_channels = prog_channels + boundary_channels
+        total_in_channels = (
+            in_channels + self.pos_channels + (3 if self.add_3d_coordinates else 0)
+        )
+        add_3d_coordinates = Concat3dCoordinates() if self.add_3d_coordinates else None
+        return Samudra(
+            in_channels=total_in_channels,
+            out_channels=out_channels,
+            pred_residuals=self.pred_residuals,
+            last_kernel_size=self.last_kernel_size,
+            pad=self.pad,
+            unet=self.unet.build(
+                in_channels=total_in_channels,
+                pad=self.pad,
+                checkpointing=self.checkpointing,
+            ),
+            pos_channels=self.pos_channels,
+            add_3d_coordinates=add_3d_coordinates,
+            hist=hist,
+            grid_size=src.grid_size,
+            gradient_detach_interval=self.gradient_detach_interval,
+            use_bfloat16=self.use_bfloat16,
+        )
+
+
+class SamudraMultiConfig(BaseModelConfig):
+    encoder: EncoderConfig = EncoderConfig()
+    processor: UNetBackboneConfig = UNetBackboneConfig()
+    decoder: DecoderConfig = DecoderConfig()
+    perceiver_implementation: PerceiverImpl = Field(
+        default="auto",
+        description="Perceiver attention implementation shared by the encoder and decoder. "
+        "'auto' selects flash attention when CUDA is available, otherwise naive.",
+    )
+    patch_extent: list[float] = Field(
+        default=[6.0, 10.0],
+        description="Target physical extent of each patch in degrees [height_deg, width_deg]. "
+        "Shared by the encoder and decoder for consistent spatial semantics.",
+    )
+    embedding_dim: int = 128
+    use_bfloat16: bool = Field(
+        default=True,
+        description="Use bfloat16 for most layers rather than float32. Required for flash attention.",
+    )
+
+    def build(
+        self,
+        prog_channels: int,
+        boundary_channels: int,
+        out_channels: int,
+        hist: int,
+        srcs: list[DataSource],
+    ) -> SamudraMulti:
+        assert len(self.patch_extent) == 2, "patch_extent must be a pair of floats."
+        extent = self.patch_extent[0], self.patch_extent[1]
+
+        all_grid_sizes = [s.grid_size for s in srcs]
+        max_lat_size, max_lon_size = (
+            max(g[0] for g in all_grid_sizes),
+            max(g[1] for g in all_grid_sizes),
+        )
+
+        impl = self.perceiver_implementation
+        if _use_flash(impl) and not self.use_bfloat16:
+            raise ValueError(
+                "Perceiver implementation resolves to flash attention. "
+                "Please set `use_bfloat16=True` or `perceiver_implementation='naive'`."
+            )
+
+        in_channels = prog_channels + boundary_channels
+        total_in_channels = in_channels + (3 if self.add_3d_coordinates else 0)
+
+        encoder = self.encoder.build(
+            total_in_channels,
+            self.embedding_dim,
+            extent,
+            max_lat_size,
+            max_lon_size,
+            impl,
+        )
+        processor = self.processor.build(
+            self.embedding_dim,
+            self.pad,
+            self.checkpointing,
+        )
+        decoder = self.decoder.build(
+            processor.out_channels,
+            out_channels,
+            extent,
+            impl,
+        )
+
+        add_3d_coordinates = Concat3dCoordinates() if self.add_3d_coordinates else None
+        return SamudraMulti(
+            in_channels=total_in_channels,
+            out_channels=out_channels,
+            pred_residuals=self.pred_residuals,
+            last_kernel_size=self.last_kernel_size,
+            pad=self.pad,
+            encoder=encoder,
+            processor=processor,
+            decoder=decoder,
+            add_3d_coordinates=add_3d_coordinates,
+            hist=hist,
+            checkpointing=self.checkpointing,
+            gradient_detach_interval=self.gradient_detach_interval,
+            use_bfloat16=self.use_bfloat16,
+        )
+
+
+class SamudraMiniConfig(BaseModelConfig):
+    perceiver: PerceiverConfig = PerceiverConfig()
+    perceiver_implementation: PerceiverImpl = Field(
+        default="auto",
+        description="Perceiver attention implementation for the single PerceiverIO model. "
+        "'auto' selects flash attention when CUDA is available, otherwise naive.",
+    )
+    embedding_dim: int = Field(
+        default=128,
+        description="Dimension of data-token embeddings before PerceiverIO.",
+    )
+    queries_dim: int = Field(
+        default=128,
+        description="Dimension of PerceiverIO output queries.",
+    )
+    coordinate_embedding_dim: int = Field(
+        default=64,
+        description="Hidden dimension used by learned 3D Cartesian coordinate embeddings.",
+    )
+    query_chunk_size: int | None = Field(
+        default=None,
+        description="Optional chunk size for query decoding. If set, PerceiverIO is called "
+        "over query chunks to reduce memory use.",
+    )
+    use_bfloat16: bool = Field(
+        default=True,
+        description="Use bfloat16 for most layers rather than float32. Required for flash attention.",
+    )
+
+    def build(
+        self,
+        prog_channels: int,
+        boundary_channels: int,
+        out_channels: int,
+        hist: int,
+        srcs: list[DataSource],
+    ) -> SamudraMini:
+        if self.add_3d_coordinates:
+            raise ValueError(
+                "SamudraMini always uses learned Cartesian coordinate embeddings. "
+                "Please set `add_3d_coordinates=False`."
+            )
+
+        impl = self.perceiver_implementation
+        if _use_flash(impl) and not self.use_bfloat16:
+            raise ValueError(
+                "Perceiver implementation resolves to flash attention. "
+                "Please set `use_bfloat16=True` or `perceiver_implementation='naive'`."
+            )
+
+        in_channels = prog_channels + boundary_channels
+        perceiver_io = self.perceiver.build_io(
+            self.embedding_dim,
+            self.queries_dim,
+            out_channels,
+            impl,
+        )
+        return SamudraMini(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            pred_residuals=self.pred_residuals,
+            last_kernel_size=self.last_kernel_size,
+            pad=self.pad,
+            input_embedding_dim=self.embedding_dim,
+            coordinate_embedding_dim=self.coordinate_embedding_dim,
+            queries_dim=self.queries_dim,
+            query_chunk_size=self.query_chunk_size,
+            perceiver_io=perceiver_io,
+            hist=hist,
+            checkpointing=self.checkpointing,
+            gradient_detach_interval=self.gradient_detach_interval,
+            use_bfloat16=self.use_bfloat16,
+        )
+
+
+AnyModelConfig = SamudraConfig | SamudraMultiConfig | SamudraMiniConfig
+
+
+class DistributedConfig(BaseConfig):
+    dist_url: str | None = None
+    world_size: int | None = None
+    rank: int | None = None
+    gpu: int | None = None
+    dist_backend: str | None = None
+
+
+class ExperimentConfig(BaseConfig):
+    name: str = "cm4_samudra"
+    rand_seed: int = 1
+    base_output_dir: str = "train"
+    # we require this to be set by the user but have optional here
+    # so we can leave it out of config files
+    data_root: Location | None = None
+    wandb: WandBConfig
+
+    @cached_property
+    def output_dir(self) -> Path:
+        return Path(self.base_output_dir) / f"{self.name}"
+
+    @cached_property
+    def nets_dir(self) -> Path:
+        return self.output_dir / "saved_nets"
+
+    @cached_property
+    def resolved_data_root(self) -> ResolvedLocation:
+        # Default to the current directory when no data_root is given. Absolute
+        # data locations (e.g. an s3:// demo source) ignore the root entirely, so
+        # they run flagless; relative locations resolve against cwd and still
+        # fail loudly at open() if the data isn't there.
+        default_root = LocalLocation(path=Path.cwd())
+        if self.data_root is None:
+            return default_root
+        return default_root.resolve(self.data_root)
+
+
+class ProfilerConfig(BaseConfig):
+    # How often (in batches processed) to take a snapshot of the CUDA memory
+    # (None = no snapshots)
+    cuda_snapshot_frequency: int | None = None
+
+    def build(self, output_dir: Path, device: torch.device) -> Profiler:
+        if self.cuda_snapshot_frequency is not None and device.type != "cuda":
+            raise ValueError(
+                "cuda_snapshot_frequency is only supported on CUDA devices, got "
+                f"{device.type}"
+            )
+        return Profiler(output_dir, self.cuda_snapshot_frequency)
+
+
+# See backend.py for how these are turned into concrete devices
+TrainBackendConfig = Literal["cpu", "cuda", "nccl", "auto"]
+
+
+class DynamicLossConfig(pydantic.BaseModel):
+    type: Literal["dynamic"] = "dynamic"
+    metric: LossMetric = "mse"
+    limit: float | None = Field(
+        description="The ratio of the largest weight to the smallest weight across all channels which we'll allow. Default of None means no limit.",
+        default=None,
+        ge=1.0,
+    )
+
+
+class GradientLossConfig(pydantic.BaseModel):
+    type: Literal["gradient"] = "gradient"
+    # at the moment this metric is only used for the non-gradient loss
+    # (and would take a bit of refactoring to make it work for the gradient loss too)
+    # so we fix it to MAE for now until it's clear we what flexibility is needed here.
+    # TODO(#497): support other metrics for the gradient loss
+    metric: Literal["mae"] = "mae"
+    alpha: float = Field(
+        description="Scaling factor for the gradient penalty term (alpha in the gradient-weighted loss).",
+        default=0.1,
+        ge=0.0,
+    )
+
+
+Loss = LossMetric | DynamicLossConfig | GradientLossConfig
+
+
+def build_loss_fn(
+    loss_cfg: Loss,
+    device: torch.device,
+    num_channels: int,
+    pad_mode: str,
+) -> LossFnWithContext:
+    match loss_cfg:
+        case str():
+            return loss_fn_from_metric(loss_cfg)
+        case DynamicLossConfig(metric=metric, limit=limit):
+            loss_fn = loss_fn_from_metric(metric)
+            return DynamicLoss(
+                loss_fn=loss_fn,
+                limit=limit,
+                device=device,
+                num_channels=num_channels,
+            )
+        case GradientLossConfig(metric=metric, alpha=alpha):
+            loss_fn = loss_fn_from_metric(metric)
+            return GradientLoss(
+                loss_fn=loss_fn,
+                gradient_weight=alpha,
+                pad_mode=pad_mode,
+            )
+        case _:
+            assert_never(loss_cfg)
+
+
+class TrainConfig(TopLevelConfig):
+    # Training parameters
+    disk_mode: bool = True
+    pin_mem: bool = True
+    save_freq: int = 5
+    validation_image_log_freq: int = Field(
+        default=10,
+        ge=1,
+        description=(
+            "How often to log expensive validation images. Epochs are 1-based, so "
+            "a value of 10 logs on epochs 1, 11, 21, ..."
+        ),
+    )
+    epochs: int = 120
+    preemptible: bool = True
+    batch_size: int = 2
+    learning_rate: float = 2e-4
+    gradient_accumulation_steps: int = 1
+    scheduler: SchedulerConfig | None = None
+    loss: Loss = "mse"
+    finetune: bool = False
+    resume_ckpt_path: str | None = None
+    debug: bool = False
+    test_using_ema: bool = True
+    ema_decay: float = 0.999
+    faster_decay_at_start: bool = True
+    delayed_loss_estimate: bool = False
+    backend: TrainBackendConfig = "auto"
+
+    # Profiling parameters
+    profiler: ProfilerConfig = ProfilerConfig()
+
+    # Data parameters at root level
+    data_stride: list[int] = [1]
+    steps: list[int] = [4]
+    step_transition: list[int] = []
+    inference_epochs: list[int] = [-1]
+
+    # Config components
+    experiment: ExperimentConfig
+    data: DataConfig
+    model: AnyModelConfig
+
+    def prepare_output_dirs(self) -> None:
+        self.experiment.nets_dir.mkdir(parents=True, exist_ok=True)
+        self.experiment.output_dir.mkdir(parents=True, exist_ok=True)
+
+
+class ObsMetricsConfig(BaseConfig):
+    """Where the observation products live, and over what period to score.
+
+    Supplying this block is what turns observation metrics on: an eval job then
+    compares its rollout against DUACS, OISST and ARGO-IAP once the rollout is
+    on disk. Omit it (or pass `--observations=null`) to skip that phase.
+    """
+
+    duacs_location: Location = Field(
+        default=UnresolvedLocation(path="obs/duacs.zarr"),
+        description="DUACS surface geostrophic velocity, on its native grid.",
+    )
+    oisst_location: Location = Field(
+        default=UnresolvedLocation(path="obs/oisst.zarr"),
+        description="OISST sea-surface temperature, on its native grid.",
+    )
+    argo_iap_location: Location = Field(
+        default=UnresolvedLocation(path="obs/argo-iap.zarr"),
+        description="ARGO-IAP gridded temperature, on its native grid.",
+    )
+    rmse_start: str = Field(
+        default="2015-01-01",
+        description=(
+            "Start of the primary scoring window. Must begin a complete calendar "
+            "year: the score and its bootstrap both use equal-year blocks."
+        ),
+    )
+    rmse_end: str = Field(
+        default="2022-12-31",
+        description="End of the primary scoring window; must end a complete calendar year.",
+    )
+    bootstrap_samples: int = Field(
+        default=10_000,
+        ge=0,
+        description="Calendar-year block-bootstrap draws for the 95% CI; 0 disables.",
+    )
+    velocity_kind: Literal["absolute", "anomaly"] = Field(
+        default="absolute",
+        description=(
+            "Which DUACS geostrophic velocity to compare against. 'absolute' "
+            "matches the model velocity derived from zos."
+        ),
+    )
+    baselines: list[str] = Field(
+        default_factory=lambda: ["om4"],
+        description=(
+            "Reference rollouts scored alongside the model. 'om4' scores the "
+            "ground-truth OM4 data already staged for the eval, which is what "
+            "makes the model's own numbers interpretable."
+        ),
+    )
+
+    @model_validator(mode="after")
+    def _check(self) -> "ObsMetricsConfig":
+        unknown = set(self.baselines) - {"om4"}
+        if unknown:
+            raise ValueError(f"Unknown baselines {sorted(unknown)}; supported: ['om4']")
+        if self.window[0] > self.window[1]:
+            raise ValueError(
+                f"rmse_start ({self.rmse_start}) must not be after rmse_end "
+                f"({self.rmse_end})"
+            )
+        return self
+
+    @property
+    def window(self) -> tuple[pd.Timestamp, pd.Timestamp]:
+        """The scoring window, with a date-only end bound taken to end-of-day.
+
+        OM4 samples are stamped at 12:00, so a bare `2022-12-31` end bound would
+        exclude that day's sample. Currently masked because the observation
+        products end earlier and the window gets trimmed to the shared span, but
+        it becomes a silent one-sample loss the moment coverage extends past
+        `rmse_end` -- and the 7.5-day tolerance in the complete-calendar-year
+        check is far too loose to notice.
+        """
+        end = pd.Timestamp(self.rmse_end)
+        if end == end.normalize():
+            end = end + pd.Timedelta(days=1) - pd.Timedelta(1, "ns")
+        return pd.Timestamp(self.rmse_start), end
+
+
+# See backend.py for how these are turned into concrete devices
+EvalBackendConfig = Literal["cpu", "cuda", "auto"]
+
+
+class EvalConfig(TopLevelConfig):
+    # Basic parameters
+    debug: bool = False
+    save_zarr: bool = False
+    disk_mode: bool = True
+    # we require this to be set by the user but have optional here
+    # so we can leave it out of config files
+    ckpt_path: str | None = None
+    num_model_steps_forward: int = 200
+    backend: EvalBackendConfig = "auto"
+
+    # Config components
+    experiment: ExperimentConfig
+    data: DataConfig
+    model: AnyModelConfig
+    observations: ObsMetricsConfig | None = Field(
+        default=None,
+        description=(
+            "Observation-based metrics computed after the rollout finishes. "
+            "Present means enabled; omit to skip the phase."
+        ),
+    )
+
+    @pydantic.model_validator(mode="after")
+    def _observations_need_a_saved_rollout(self) -> Self:
+        # Caught here rather than after the rollout: observation metrics score a
+        # rollout read back from disk, and discovering the misconfiguration at
+        # the end would waste the whole job.
+        if self.observations is not None and not self.save_zarr:
+            raise ValueError(
+                "observations requires save_zarr=true: the metrics are computed "
+                "from the predictions.zarr the rollout writes."
+            )
+        return self
+
+    def prepare_output_dirs(self) -> None:
+        self.experiment.output_dir.mkdir(parents=True, exist_ok=True)
+
+
+AnyTopLevelConfig = TrainConfig | EvalConfig

samudra/config_base.py

@@ -0,0 +1,166 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import importlib.resources
+import os
+import textwrap
+from pathlib import Path
+from typing import Any, Self
+
+import yaml
+from pydantic import BaseModel, ConfigDict
+from pydantic.fields import FieldInfo
+from pydantic_settings import (
+    BaseSettings,
+    CliSettingsSource,
+    PydanticBaseSettingsSource,
+    SettingsConfigDict,
+    YamlConfigSettingsSource,
+)
+
+
+def register_include_constructor():
+    """Set up yaml.safe_load to include other yaml files via !include."""
+
+    def include_constructor(loader: yaml.Loader, node: yaml.Node) -> Any:
+        if hasattr(loader.stream, "name"):
+            name = loader.stream.name  # type: ignore
+        else:
+            raise ValueError(
+                "To support includes, you must load a file object, not a string"
+            )
+        filename = os.path.normpath(os.path.join(os.path.dirname(name), node.value))
+        with open(filename) as f:
+            return yaml.safe_load(f)
+
+    # This is arguably unsafe, but we don't parse untrusted YAML
+    yaml.loader.SafeLoader.add_constructor("!include", include_constructor)
+
+
+register_include_constructor()
+
+
+def resolve_config_path(config: str) -> Path:
+    """Resolve a config argument to a file on disk.
+
+    A real filesystem path wins, so a checkout behaves exactly as before. If it
+    doesn't exist, fall back to a preset bundled in the installed package (e.g.
+    ``samudra_om4/train.yaml``), so ``samudra train samudra_om4/train.yaml``
+    works with no checkout. The bundled configs install as real co-located
+    files, so their relative ``!include``s resolve unchanged.
+    """
+    path = Path(config).expanduser()
+    if not path.exists():
+        bundled = importlib.resources.files("samudra") / "configs" / config
+        if bundled.is_file():
+            path = Path(str(bundled))
+    return path.resolve()
+
+
+class BaseConfig(BaseModel):
+    """
+    Base class for all configs.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class TopLevelConfig(BaseSettings):
+    """
+    Base class for top-level configs (ie tasks like train or eval).
+    """
+
+    model_config = SettingsConfigDict(nested_model_default_partial_update=True)
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @classmethod
+    def settings_customise_sources(
+        cls,
+        settings_cls: type[BaseSettings],
+        init_settings: PydanticBaseSettingsSource,
+        env_settings: PydanticBaseSettingsSource,
+        dotenv_settings: PydanticBaseSettingsSource,
+        file_secret_settings: PydanticBaseSettingsSource,
+    ) -> tuple[PydanticBaseSettingsSource, ...]:
+        # We don't need env/dotenv/secrets, yaml & CLI are injected in from_yaml_and_cli
+        return (init_settings,)
+
+    @classmethod
+    def from_yaml_and_cli(
+        cls,
+        args_to_parse: list[str] | None = None,
+    ) -> Self:
+        """Load config from YAML & CLI with validation."""
+        parser = argparse.ArgumentParser(
+            description=cls.__doc__,
+            epilog=textwrap.dedent(
+                """
+                CONFIG may be a path to a YAML file or, when samudra is installed,
+                the name of a bundled preset such as `samudra_om4/train.yaml`.
+                YAML files can include other YAML files using the !include tag,
+                as in `data: !include ../data/om4.yaml` (resolved relative to the
+                including file). You can also replace any JSON argument listed above
+                with a YAML file by specifying it with an @ symbol,
+                eg `--some_param=@../data/om4.yaml`.
+                """
+            ),
+        )
+        parser.add_argument(
+            "config",
+            type=str,
+            help="Path to a config YAML file, or a bundled preset name",
+        )
+
+        cli_source = IncludeYamlCliSettingsSource(
+            cls,
+            root_parser=parser,
+            # If args_to_parse is None, we parse argv, which is what `True` does
+            cli_parse_args=args_to_parse if args_to_parse is not None else True,
+        )
+
+        # We do this after creating CliSettingsSource (which populates the parser)
+        # so the help is complete on error.
+        args = parser.parse_args(args_to_parse)
+
+        # Then we read the YAML file specified in the CLI. A bare path resolves
+        # against the filesystem; if it's missing we try a preset bundled in the
+        # installed package before giving up.
+        # Note that by default, YamlConfigSettingsSource will ignore missing files
+        config_path = resolve_config_path(args.config)
+        if not config_path.exists():
+            raise FileNotFoundError(
+                f"Config file `{args.config}` (full path: {config_path}) not found"
+            )
+        yaml_values = YamlConfigSettingsSource(cls, yaml_file=config_path)()
+
+        return cls(
+            _cli_settings_source=cli_source,
+            **yaml_values,
+        )
+
+    def save_yaml(self, save_path: Path) -> None:
+        """Save config to YAML file."""
+        with open(save_path, "w") as f:
+            yaml.dump(self.model_dump(), f)
+
+
+class IncludeYamlCliSettingsSource(CliSettingsSource):
+    """CliSettingsSource which permits @filename.yaml for JSON arguments."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def decode_complex_value(
+        self, field_name: str, field: FieldInfo, value: Any
+    ) -> Any:
+        if isinstance(value, str) and value.startswith("@"):
+            # Resolve like the positional config: a real path wins, otherwise
+            # fall back to a bundled preset (so `--data @data/om4_demo.yaml`
+            # works without a checkout).
+            with open(resolve_config_path(value[1:])) as f:
+                return yaml.safe_load(f)
+        return super().decode_complex_value(field_name, field, value)

samudra/config_schema.py

@@ -0,0 +1,177 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import json
+import re
+import sys
+from pathlib import Path
+from typing import get_args, get_origin
+
+import pydantic
+import yaml
+
+from samudra.config import EvalConfig, TrainConfig
+from samudra.viz import VizConfig
+
+
+def get_pydantic_models(
+    model: type[pydantic.BaseModel],
+    models: dict[str, type[pydantic.BaseModel]] | None = None,
+    seen: set[int] | None = None,
+) -> dict[str, type[pydantic.BaseModel]]:
+    """Recursively find all Pydantic models in a model's fields.
+
+    Args:
+        model: The Pydantic model to start from
+        models: Collected models keyed by model name
+        seen: Visited annotation objects keyed by identity
+
+    Returns:
+        Dictionary of model names to model classes
+    """
+    if models is None:
+        models = {}
+    if seen is None:
+        seen = set()
+
+    _collect_pydantic_models(model, models, seen)
+    return models
+
+
+def _collect_pydantic_models(
+    annotation: object,
+    models: dict[str, type[pydantic.BaseModel]],
+    seen: set[int],
+) -> None:
+    annotation_id = id(annotation)
+    if annotation_id in seen:
+        return
+    seen.add(annotation_id)
+
+    if isinstance(annotation, type) and issubclass(annotation, pydantic.BaseModel):
+        models[annotation.__name__] = annotation
+        for field in annotation.model_fields.values():
+            _collect_pydantic_models(field.annotation, models, seen)
+        return
+
+    # We check if this is a parameterized type and recurse if so (eg Union[int, str])
+    # (Union is the origin in that case.)
+    origin = get_origin(annotation)
+    if origin is None:
+        return
+
+    for arg in get_args(annotation):
+        _collect_pydantic_models(arg, models, seen)
+
+
+def generate_schemas(
+    output_dir: Path,
+    models: dict[str, type[pydantic.BaseModel]],
+) -> None:
+    """Generate JSON schemas for all Pydantic models and save them to output_dir.
+
+    Args:
+        output_dir: Directory to save the generated schemas
+        models: Dictionary of model names to model classes
+    """
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    for model_name, model in models.items():
+        schema = model.model_json_schema()
+        output_path = output_dir / f"{model_name}.json"
+
+        # Check if file exists and content is different
+        if output_path.exists():
+            with open(output_path) as f:
+                existing_schema = json.load(f)
+                if existing_schema == schema:
+                    continue
+
+        # Write only if content differs or file doesn't exist
+        with open(output_path, "w") as f:
+            json.dump(schema, f, indent=2)
+            print(f"🆕 Updated schema for {model_name} at {output_path}")
+
+
+def validate_config_files(
+    config_dir: Path,
+    models: dict[str, type[pydantic.BaseModel]],
+) -> bool:
+    """Validate YAML configuration files against their Pydantic models.
+
+    Args:
+        config_dir: Directory containing YAML configuration files to validate
+        models: Dictionary of model names to model classes
+    """
+    valid = True
+
+    # Find all YAML files in the config directory
+    yaml_files = list(config_dir.rglob("*.yaml"))
+
+    for yaml_file in yaml_files:
+        # Read the YAML file
+        with open(yaml_file) as f:
+            yaml_content = f.read()
+
+        # Extract the schema path from the YAML file's first line
+        schema_match = re.search(
+            r"\s*#\s+yaml-language-server:\s+\$schema=(.*\.json)", yaml_content
+        )
+        if not schema_match:
+            print(f"⚠️ No schema specified in {yaml_file}")
+            continue
+
+        schema_name = Path(schema_match.group(1)).stem
+        if schema_name not in models:
+            print(f"✗ Model {schema_name} not found for {yaml_file}")
+            valid = False
+            continue
+
+        # Load and validate the YAML content
+        try:
+            with open(yaml_file) as f:
+                config = yaml.safe_load(f)
+                # Validate using Pydantic model
+                models[schema_name].model_validate(config)
+                print(f"✅ {yaml_file} is a valid {schema_name}")
+        except Exception as e:
+            print(f"❌ {yaml_file} is an invalid {schema_name}: {str(e)}")
+            valid = False
+
+    return valid
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Generate JSON schemas and validate config files"
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        default=Path("src/samudra/configs/schemas"),
+        help="Directory to save the generated schemas",
+    )
+    parser.add_argument(
+        "--validate-dir",
+        type=Path,
+        default=Path("src/samudra/configs"),
+        help="Directory of config files to validate",
+    )
+
+    args = parser.parse_args()
+
+    # Get all available models
+    models = get_pydantic_models(TrainConfig)
+    models.update(get_pydantic_models(EvalConfig))
+    models.update(get_pydantic_models(VizConfig))
+
+    generate_schemas(args.output_dir, models)
+
+    if not validate_config_files(args.validate_dir, models):
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

samudra/configs/data/llc.yaml

@@ -0,0 +1,26 @@
+# SPDX-FileCopyrightText: 2026 Ocean Emulator Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/DataConfig.json
+
+sources:
+  - type: llc
+    face: 1
+    i_start: 0
+    i_end: 720
+    j_start: 0
+    j_end: 720
+    train_time:
+      start: "2011-09-10T12:00:00Z"
+      end: "2012-09-01T12:00:00Z"
+    val_time:
+      start: "2012-09-01T12:00:00Z"
+      end: "2012-11-15T12:00:00Z"
+    data_location: LLC.zarr
+    data_means_location: LLC_means.zarr
+    data_stds_location: LLC_stds.zarr
+concurrent_compute: false
+loading:
+  type: cpu
+  num_workers: 4

samudra/configs/data/obs.yaml

@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# Observation-based evaluation metrics, scored once the rollout is on disk.
+# Including this file is what enables them; omit it to skip the phase.
+#
+# The three products are published anonymously readable, so an eval run needs
+# neither credentials nor a staging step to produce these metrics. They were
+# built from the providers by `python -m ocean_preprocessing.obs_preprocessing`
+# (see data/README.md), and each store records its own source, attribution and
+# provenance in its zarr attrs.
+#
+# To read a locally staged copy instead -- worth it on repeat runs, since DUACS
+# alone is 46 GB -- set OSN_OBS_PATH in the skypilot eval job and override the
+# locations to the relative paths `obs/duacs.zarr` and so on, which resolve
+# against experiment.data_root.
+
+duacs_location: &obs_s3
+  type: s3
+  bucket: m2lines-pubs
+  path: Samudra/v2026-07/obs/duacs.zarr
+  endpoint_url: https://nyu1.osn.mghpcc.org
+  anon: true
+oisst_location:
+  <<: *obs_s3
+  path: Samudra/v2026-07/obs/oisst.zarr
+argo_iap_location:
+  <<: *obs_s3
+  path: Samudra/v2026-07/obs/argo-iap.zarr
+
+# Eight complete calendar years, so the score and its block bootstrap share
+# equal-weight blocks.
+rmse_start: "2015-01-01"
+rmse_end: "2022-12-31"

samudra/configs/data/om4.yaml

@@ -0,0 +1,25 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/DataConfig.json
+
+sources:
+  - type: om4
+    prognostic_vars_key: thermo_dynamic_all
+    boundary_vars_key: tau_hfds_hfds_anom
+    train_time:
+      start: "1975-01-03"
+      end: "2013-10-05"
+    val_time:
+      start: "2013-10-05"
+      end: "2014-10-05"
+    inference_times:
+      - start: "2014-10-10"
+        end: "2022-12-24"
+    data_location: OM4.zarr
+    data_means_location: OM4_means.zarr
+    data_stds_location: OM4_stds.zarr
+concurrent_compute: true
+loading:
+  type: cpu

samudra/configs/data/om4_demo.yaml

@@ -0,0 +1,51 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/DataConfig.json
+
+# Two-degree (2°) OM4 demo data source that streams straight from the public
+# m2lines-pubs bucket on OSN over S3 — no local download or credentials needed
+# (anon: true reads the public bucket anonymously). The locations below are
+# absolute (type: s3), so they are read as-is and `experiment.data_root` is
+# unused — you can omit it entirely.
+
+sources:
+  - type: om4
+    prognostic_vars_key: thermo_dynamic_all
+    # tau_hfds (tauuo, tauvo, hfds) rather than the samudra_om4 presets'
+    # tau_hfds_hfds_anom: the public 2deg data has no `hfds_anomalies` variable.
+    # This is a self-contained train -> eval -> viz demo (train a fresh model on
+    # this data); its boundary channels differ from the released checkpoints, so
+    # don't eval a preset-trained checkpoint against this source.
+    boundary_vars_key: tau_hfds
+    train_time:
+      start: "1975-01-03"
+      end: "2013-10-05"
+    val_time:
+      start: "2013-10-05"
+      end: "2014-10-05"
+    inference_times:
+      - start: "2014-10-10"
+        end: "2022-12-24"
+    data_location:
+      type: s3
+      endpoint_url: "https://nyu1.osn.mghpcc.org"
+      anon: true
+      bucket: m2lines-pubs
+      path: Samudra/v2026-07/om4_twodeg/OM4.zarr
+    data_means_location:
+      type: s3
+      endpoint_url: "https://nyu1.osn.mghpcc.org"
+      anon: true
+      bucket: m2lines-pubs
+      path: Samudra/v2026-07/om4_twodeg/OM4_means.zarr
+    data_stds_location:
+      type: s3
+      endpoint_url: "https://nyu1.osn.mghpcc.org"
+      anon: true
+      bucket: m2lines-pubs
+      path: Samudra/v2026-07/om4_twodeg/OM4_stds.zarr
+concurrent_compute: true
+loading:
+  type: cpu

samudra/configs/samudra_mini_om4/README.md

@@ -0,0 +1,11 @@
+<!--
+SPDX-FileCopyrightText: 2026 Samudra Authors
+
+SPDX-License-Identifier: CC-BY-4.0
+-->
+
+# SamudraMini OM4
+
+These configurations are for SamudraMini, a single PerceiverIO model that treats
+each `(lat, lon)` point as a token and uses learned embeddings of 3D Cartesian
+coordinates for both data tokens and output queries.

samudra/configs/samudra_mini_om4/model.yaml

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/SamudraMiniConfig.json
+
+pred_residuals: false
+last_kernel_size: 3
+pad: "circular"
+checkpointing: "all"
+perceiver_implementation: "auto"
+embedding_dim: 128
+queries_dim: 128
+coordinate_embedding_dim: 64
+query_chunk_size: 8192
+
+perceiver:
+  depth: 4
+  latent_dim: 96
+  num_latents: 256

samudra/configs/samudra_mini_om4/train.yaml

@@ -0,0 +1,36 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/TrainConfig.json
+
+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_stride: [1]
+steps: [4]
+step_transition: []
+preemptible: false
+
+backend: auto
+
+experiment:
+  name: samudra_mini_om4
+  rand_seed: 15
+  base_output_dir: .LOCAL
+  wandb:
+    mode: disabled
+    project: default
+data: !include ../data/om4.yaml
+model: !include model.yaml

samudra/configs/samudra_multi_om4/README.md

@@ -0,0 +1,11 @@
+<!--
+SPDX-FileCopyrightText: 2026 Samudra Authors
+
+SPDX-License-Identifier: CC-BY-4.0
+-->
+
+# samudra-multi OM4
+
+These configurations are for the samudra-multi (encoder-processor-decoder) model,
+trained on OM4 data. The intent is for this to support heterogeneous
+resolutions and modalities.

samudra/configs/samudra_multi_om4/eval.yaml

@@ -0,0 +1,71 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/EvalConfig.json
+#
+# Eval config for the multi-scale samudra-multi model.
+#
+# The eval code runs inference on the *first* data source only.
+# To evaluate all three scales, run three separate eval jobs and
+# reorder the sources list so the target resolution is first.
+# All three sources must always be present so the model builds
+# with the correct max grid size for the encoder patch grid.
+# The max grid size affects the rotary positional embedding
+# dimension for the encoder's Perceiver. If it is set wrong,
+# it would likely cause hard-to-debug model level errors.
+# Setting multiple data sources lets us use checkpoints for
+# samudra-multi models that use multiple scales.
+
+debug: false
+save_zarr: true
+disk_mode: true
+num_model_steps_forward: 25
+
+experiment:
+  name: samudra_multi_om4_eval
+  rand_seed: 15
+  base_output_dir: .LOCAL
+  wandb:
+    mode: disabled
+    project: default
+
+data:
+  concurrent_compute: true
+  sources:
+    # First source is used for inference. Reorder to eval a different scale.
+    - type: om4
+      prognostic_vars_key: thermo_dynamic_all
+      boundary_vars_key: tau_hfds
+      train_time: &train_time
+        start: "1975-01-03"
+        end: "2013-10-05"
+      val_time: &val_time
+        start: "2013-10-05"
+        end: "2014-10-05"
+      inference_times: &inference_times
+        - start: "2014-10-05"
+          end: "2015-10-05"
+      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
+    - type: om4
+      prognostic_vars_key: thermo_dynamic_all
+      boundary_vars_key: tau_hfds
+      train_time: *train_time
+      val_time: *val_time
+      inference_times: *inference_times
+      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
+    - type: om4
+      prognostic_vars_key: thermo_dynamic_all
+      boundary_vars_key: tau_hfds
+      train_time: *train_time
+      val_time: *val_time
+      inference_times: *inference_times
+      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
+
+model: !include model.yaml

samudra/configs/samudra_multi_om4/model.yaml

@@ -0,0 +1,45 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/SamudraMultiConfig.json
+
+pred_residuals: false
+last_kernel_size: 3
+pad: "circular"
+checkpointing: "all"
+embedding_dim: 128  # Reduced from 280
+perceiver_implementation: "auto"
+# TODO(alxmrs): How fine of a resolution should each patch be? This produces a 60x72 latent grid.
+patch_extent: [3.0, 5.0]
+
+encoder:
+  perceiver:
+    depth: 2
+    latent_dim: 64
+    num_latents: 256
+
+decoder:
+  perceiver:
+    depth: 2
+    latent_dim: 64
+    num_latents: 256
+  queries_dim: 64  # on the small end, but typical for PercieverIO models (usually in the ~100s-1000s).
+  window_patches: 12  # Must evenly divide latent grid (60x72). Enables real spatial windowing.
+  context_patches: 1  # Look at this many window_patches around the target patch.
+
+processor:
+  ch_width: [380, 480, 520]
+  dilation: [1, 1, 1]
+  n_layers: [1, 1, 1]
+
+  core_block:
+    block_type: "conv_next_block"
+    kernel_size: 3
+    activation: "capped_gelu"
+    upscale_factor: 2
+    norm: "instance"
+    pointwise_linear: true
+
+  down_sampling_block: "avg_pool"
+  up_sampling_block: "zonally_periodic_upsample"

samudra/configs/samudra_multi_om4/train.yaml

@@ -0,0 +1,75 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/TrainConfig.json
+
+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
+  limit: 20
+finetune: false
+resume_ckpt_path: null
+inference_epochs: []
+data_stride: [1]
+steps: [4]
+step_transition: []
+preemptible: false
+
+backend: auto
+
+experiment:
+  name: samudra_multi_om4
+  rand_seed: 15
+  base_output_dir: .LOCAL
+  wandb:
+    mode: disabled
+    project: default
+data:
+  concurrent_compute: true
+  loading:
+    type: cpu
+    num_workers: 4
+  sources:
+    - type: om4
+      prognostic_vars_key: thermo_dynamic_all
+      boundary_vars_key: tau_hfds
+      train_time: &train_time
+        start: "1975-01-03"
+        end: "2013-10-05"
+      val_time: &val_time
+        start: "2013-10-05"
+        end: "2014-10-05"
+      inference_times: &inference_times
+        - start: "2014-10-05"
+          end: "2015-10-05"
+      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
+    - type: om4
+      prognostic_vars_key: thermo_dynamic_all
+      boundary_vars_key: tau_hfds
+      train_time: *train_time
+      val_time: *val_time
+      inference_times: *inference_times
+      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
+    - type: om4
+      prognostic_vars_key: thermo_dynamic_all
+      boundary_vars_key: tau_hfds
+      train_time: *train_time
+      val_time: *val_time
+      inference_times: *inference_times
+      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

samudra/configs/samudra_om4/README.md

@@ -0,0 +1,10 @@
+<!--
+SPDX-FileCopyrightText: 2026 Samudra Authors
+
+SPDX-License-Identifier: CC-BY-4.0
+-->
+
+# Samudra OM4
+
+These configurations are our current WIP Samudra-style (UNet) model,
+trained on OM4 data.

samudra/configs/samudra_om4/eval.yaml

@@ -0,0 +1,21 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/EvalConfig.json
+
+debug: false
+save_zarr: true
+disk_mode: true
+num_model_steps_forward: 25
+
+experiment:
+  name: samudra_om4_eval
+  rand_seed: 15
+  base_output_dir: .LOCAL
+  wandb:
+    mode: disabled
+    project: default
+data: !include ../data/om4.yaml
+model: !include model.yaml
+observations: !include ../data/obs.yaml

samudra/configs/samudra_om4/model.yaml

@@ -0,0 +1,25 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/SamudraConfig.json
+
+checkpointing: all
+pred_residuals: false
+last_kernel_size: 3
+pad: "circular"
+
+unet:
+  ch_width: [280, 380, 480, 520]
+  dilation: [1, 2, 4, 8]
+  n_layers: [1, 1, 1, 1]
+
+  core_block:
+    block_type: "conv_next_block"
+    kernel_size: 3
+    activation: "capped_gelu"
+    upscale_factor: 2
+    norm: "batch"
+
+  down_sampling_block: "avg_pool"
+  up_sampling_block: "zonally_periodic_upsample"

samudra/configs/samudra_om4/train.yaml

@@ -0,0 +1,37 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/TrainConfig.json
+
+debug: false
+disk_mode: true
+pin_mem: true
+save_freq: 5
+epochs: 70
+batch_size: 4
+learning_rate: 0.0006
+gradient_accumulation_steps: 1
+scheduler: { type: cosine }
+loss:
+  type: dynamic
+  metric: mse
+finetune: false
+resume_ckpt_path: null
+inference_epochs: []
+data_stride: [1]
+steps: [4]
+step_transition: []
+preemptible: false
+
+backend: auto
+
+experiment:
+  name: samudra_om4
+  rand_seed: 15
+  base_output_dir: .LOCAL
+  wandb:
+    mode: disabled
+    project: default
+data: !include ../data/om4.yaml
+model: !include model.yaml

samudra/configs/samudra_om4/viz.yaml

@@ -0,0 +1,25 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/VizConfig.json
+
+base_output_dir: ".LOCAL"
+name: "samudra_om4_viz"
+dataset_name: "OM4"
+# Ground truth comes from this data source (same shape as train/eval), so
+# `--data @data/om4_demo.yaml` points viz at the demo data just like the others.
+data: !include ../data/om4.yaml
+basins_location:
+  type: s3
+  anon: true
+  bucket: "m2lines-pubs"
+  path: "Samudra/basins/basin_masks_original.zarr"
+  endpoint_url: "https://nyu1.osn.mghpcc.org"
+groundtruth_time_range:
+  start: "2014-10-20"
+  end: "2022-12-24"
+runs:
+  - name: "samudra_om4_eval"
+    location: ".LOCAL/samudra_om4_eval/predictions.zarr"
+observations: !include ../data/obs.yaml

samudra/configs/samudra_om4_v1/README.md

@@ -0,0 +1,10 @@
+<!--
+SPDX-FileCopyrightText: 2026 Samudra Authors
+
+SPDX-License-Identifier: CC-BY-4.0
+-->
+
+# Samudra OM4 V1
+
+These configurations are our replication of the [original Samudra model](https://arxiv.org/abs/2412.03795)
+trained on OM4 data.

samudra/configs/samudra_om4_v1/eval.yaml

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/EvalConfig.json
+
+debug: false
+save_zarr: true
+disk_mode: true
+num_model_steps_forward: 25
+
+experiment:
+  name: samudra_om4_v1_eval
+  rand_seed: 15
+  base_output_dir: .LOCAL
+  wandb:
+    mode: disabled
+    project: default
+data: !include ../data/om4.yaml
+model: !include model.yaml

samudra/configs/samudra_om4_v1/model.yaml

@@ -0,0 +1,25 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/SamudraConfig.json
+
+pred_residuals: false
+last_kernel_size: 3
+pad: "circular"
+checkpointing: "all"
+
+unet:
+  ch_width: [200, 250, 300, 400]
+  dilation: [1, 2, 4, 8]
+  n_layers: [1, 1, 1, 1]
+
+  core_block:
+    block_type: "conv_next_block"
+    kernel_size: 3
+    activation: "capped_gelu"
+    upscale_factor: 4
+    norm: "batch"
+
+  down_sampling_block: "avg_pool"
+  up_sampling_block: "bilinear_upsample"

samudra/configs/samudra_om4_v1/train.yaml

@@ -0,0 +1,34 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/TrainConfig.json
+
+debug: false
+disk_mode: true
+pin_mem: true
+save_freq: 5
+epochs: 70
+batch_size: 4
+learning_rate: 0.0006
+scheduler: { type: cosine }
+loss: mse
+finetune: false
+resume_ckpt_path: null
+inference_epochs: []
+data_stride: [1]
+steps: [4]
+step_transition: []
+preemptible: false
+
+backend: auto
+
+experiment:
+  name: samudra_om4_v1
+  rand_seed: 15
+  base_output_dir: .LOCAL
+  wandb:
+    mode: disabled
+    project: default
+data: !include ../data/om4.yaml
+model: !include model.yaml

samudra/configs/samudra_om4_v1/viz.yaml

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/VizConfig.json
+
+base_output_dir: ".LOCAL"
+name: "samudra_om4_v1_viz"
+dataset_name: "OM4"
+groundtruth_location: "OM4.zarr"
+basins_location:
+  type: s3
+  anon: true
+  bucket: "m2lines-pubs"
+  path: "Samudra/basins/basin_masks_original.zarr"
+  endpoint_url: "https://nyu1.osn.mghpcc.org"
+groundtruth_time_range:
+  start: "2014-10-20"
+  end: "2022-12-24"
+runs:
+  - name: "samudra_om4_v1_eval"
+    location: ".LOCAL/samudra_om4_v1_eval/predictions.zarr"

samudra/configs/samudra_om4_v2/README.md

@@ -0,0 +1,11 @@
+<!--
+SPDX-FileCopyrightText: 2026 Samudra Authors
+
+SPDX-License-Identifier: CC-BY-4.0
+-->
+
+# Samudra OM4 V2
+
+These configurations are the v2 version of Samudra-style (UNet) model,
+trained on OM4 data. Major changes from the v1 model are that it is
+wider and uses a dynamically-weighted loss.

samudra/configs/samudra_om4_v2/eval.yaml

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/EvalConfig.json
+
+debug: false
+save_zarr: true
+disk_mode: true
+num_model_steps_forward: 25
+
+experiment:
+  name: samudra_om4_v2_eval
+  rand_seed: 15
+  base_output_dir: .LOCAL
+  wandb:
+    mode: disabled
+    project: default
+data: !include ../data/om4.yaml
+model: !include model.yaml

samudra/configs/samudra_om4_v2/model.yaml

@@ -0,0 +1,25 @@
+# SPDX-FileCopyrightText: 2026 Samudra Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# yaml-language-server: $schema=../schemas/SamudraConfig.json
+
+checkpointing: all
+pred_residuals: false
+last_kernel_size: 3
+pad: "circular"
+
+unet:
+  ch_width: [280, 380, 480, 520]
+  dilation: [1, 2, 4, 8]
+  n_layers: [1, 1, 1, 1]
+
+  core_block:
+    block_type: "conv_next_block"
+    kernel_size: 3
+    activation: "capped_gelu"
+    upscale_factor: 2
+    norm: "batch"
+
+  down_sampling_block: "avg_pool"
+  up_sampling_block: "zonally_periodic_upsample"