diff --git a/earth2studio/nvcoupler/README.md b/earth2studio/nvcoupler/README.md
new file mode 100644
index 000000000..477b07498
--- /dev/null
+++ b/earth2studio/nvcoupler/README.md
@@ -0,0 +1,273 @@
+# nvcoupler
+
+A NUOPC/ESMF-inspired, Python-native coupling framework for AI Earth-system
+inference — and, by design, for future coupled fine-tuning.
+
+```python
+import earth2studio.nvcoupler as nvc
+```
+
+## Why
+
+AI Earth-system models are coupled today in three ad-hoc ways: baked into a
+datapipe (PhysicsNeMo `ConstantCoupler` / `TrailingAverageCoupler`),
+hard-coded inside one model's forward pass (DLESyM's atmos→ocean exchange),
+or left entirely to the caller (StormScope `call_with_conditioning`).
+Changing how DLESyM couples means forking `dlesym.py`; swapping its ocean for
+observed SST means a rewrite.
+
+nvcoupler factors coupling out of the models. Independent **Components**
+exchange **Fields** by standard name through **Connectors** (regridding,
+masking, time policies, windowed reductions) and **Mediators** (multi-source
+windowed reductions), scheduled by a **Driver** executing a NUOPC-style
+**run sequence** on a shared **Clock**.
+Swapping a modeled ocean for prescribed observations becomes a one-line
+change; a coupling-order experiment becomes a one-line DSL edit.
+
+## Concepts (NUOPC → nvcoupler)
+
+| NUOPC / ESMF | nvcoupler | Module |
+|---------------------|--------------------------|-----------------|
+| ESMF_Field / State | `Field`, `State` | `field.py` |
+| Field dictionary | `FieldDictionary` | `dictionary.py` |
+| ESMF Clock | `Clock` | `clock.py` |
+| NUOPC_Model | `Component` subclasses | `component.py` |
+| NUOPC_Connector | `Connector` | `connector.py` |
+| NUOPC_Mediator | `Mediator` subclasses | `mediator.py` |
+| NUOPC_Driver, runSeq| `Driver`, `RunSequence` | `driver.py`, `sequence.py` |
+
+### Field & State
+
+A `Field` is one physical quantity: a **torch tensor** plus its earth2studio
+`CoordSystem`, a canonical identity (`standard_name`, `units`), a
+`valid_time`, and optional `mask` (True = valid, e.g. ocean points for SST)
+and `vertical` metadata. A `State` is a bag of Fields keyed by standard name;
+every component owns an import State and an export State.
+
+Field data never round-trips through numpy inside the framework, so autograd
+graphs survive every exchange.
+
+### FieldDictionary
+
+Connectors match fields by **standard name**, never raw model variable
+strings. Aliases map model vocabularies onto standard names
+(`z1000 → geopotential_at_1000hpa`); units are checked (not converted) on
+every match. Derived fields carry a machine-readable `CellMethod`:
+
+```python
+nvc.FieldEntry(
+ "total_precipitation_48h_sum", "kg m-2",
+ cell_method=nvc.CellMethod("total_precipitation_6h", "sum", np.timedelta64(48, "h")),
+)
+```
+
+which is how mediators (and future auto-wiring) know that a 48 h precip sum
+derives from the 6 h precip field — no suffix string-parsing.
+
+### Components
+
+All components implement the NUOPC phases
+`advertise → realize(clock) → initialize(x, coords) → run(time)* → finalize`
+and declare their imports/exports as standard names.
+
+- **`PrognosticComponent`** wraps any earth2studio `PrognosticModel`.
+ Timestep and exports are inferred from `input_coords()`/`output_coords()`.
+- **`CallableComponent`** wraps a plain `fn(x, coords) -> (x, coords)` — the
+ entry point for synthetic components **and non-ML models** (a process-based
+ hydrology or crop model can join the coupled system through it).
+
+**ImportAdapters** — the critical seam. Real models receive coupled fields in
+different call shapes, so the adapter owns the model invocation:
+
+| Adapter | Call shape | Real-world pattern |
+|---|---|---|
+| `VariableOverwriteAdapter` (default) | overwrite state-variable slices, `model(x, coords)` | prescribed forcing as a state channel |
+| `ConditioningKwargAdapter` | `model.call_with_conditioning(x, coords, conditioning=...)` | StormScope |
+| `ExtraTensorAdapter` | `model(x, coords, coupling)` | DLESyM / PhysicsNeMo 4-tensor |
+| `PullAdapter` | installs a `StateDataSource` shim on `model.conditioning_data_source`, then `model(x, coords)` | StormCast |
+
+The pull pattern covers models that *fetch* their forcing internally rather
+than accepting it as an argument — the shim answers the model's own
+`fetch_data` calls from the import State; because that path crosses xarray/
+numpy, pull-coupled components are inference-only (no autograd through the
+exchange).
+
+Models that manage a sliding input window internally take a
+`next_input(prev_x, prev_coords, out, out_coords)` hook; the default handles
+single-window models.
+
+### Connector
+
+`Connector(src, dst)` transfers the intersection of src's exports and dst's
+imports (or an explicit `fields=[...]`). Each transfer runs a pipeline:
+
+1. **Time policy** — `"constant"` holds the latest export (the PhysicsNeMo
+ `ConstantCoupler` behavior); `"linear"` extrapolates from the two most
+ recent exports toward the current time.
+2. **Vertical** — when the destination declares `import_vertical` and it
+ differs from the field's coordinate, hybrid→pressure interpolation runs
+ (see below).
+3. **Mask fill** — `"zero"`, or `"nearest"` (each invalid point takes its
+ nearest valid neighbor via a cached KDTree — the principled version of
+ DLESyM's SST NaN-interpolation hack). Always applied *before* regridding
+ so invalid values can't bleed into the interpolation.
+4. **Spatial regrid** — lazily built and cached per grid pair: identity when
+ grids match, else bilinear via `earth2studio.utils.interp.
+ latlon_interpolation_regular` (regular 1D lat/lon sources). HEALPix
+ (`face` dim) and curvilinear sources need a user `regridder=` callable
+ (build one with `earth2grid` as `models/px/dlesym.py` does).
+
+Setting `window=` and `reduce=` together makes a **windowed connector** — the
+preferred path for simple fast→slow windowed coupling: each execute folds the
+source export into a running reduction (`mean`/`sum`/`max`/`min`), and on
+window boundaries the reduced field is delivered under the *derived* standard
+name the destination's dictionary declares via a matching
+`CellMethod(base, method, window)` entry. No mediator, no extra slot actions.
+
+### Mediator
+
+The multi-source, general form of the same accumulator core windowed
+connectors use. `AccumulationMediator("med",
+["geopotential_at_1000hpa_48h_mean"])` reads the CellMethod off each derived
+field: it imports the base field, accumulates a **running** reduction
+(mean/sum/max/min — O(1) memory in window length) on every connector
+delivery, and exports the reduced field when its `med.compute` action runs in
+the slow slot. `TrailingAverageMediator` is the mean-only restriction
+matching DLESyM's ocean forcing exactly. Duplicate deliveries (same
+`valid_time`) are ignored. For one source feeding one destination, prefer
+`Connector(..., window=, reduce=)`.
+
+### Vertical coordinates
+
+For components with an explicit `level` dimension (chiefly chemistry
+emulators on hybrid sigma-pressure levels):
+
+```python
+hybrid = nvc.HybridLevels(a=(30000., 20000., 0.), b=(0., 0.5, 1.0)) # p = a + b·ps
+pressure = nvc.PressureLevels((500., 850.)) # hPa
+
+met = CallableComponent(..., export_vertical={"ozone_mixing_ratio": hybrid})
+chem = CallableComponent(..., import_vertical={"ozone_mixing_ratio": pressure})
+```
+
+The connector interpolates linearly in log-pressure (differentiable), pulling
+`surface_pressure` from the source's exports automatically and raising
+`VerticalMismatchError` with a concrete fix when it can't. Models that encode
+levels in variable names (`z500`, `t850`) never touch this machinery.
+
+### Driver & run sequence
+
+A system is declared as **components plus connections**; the Driver derives
+the canonical (lagged) run sequence from the coupling graph:
+
+```python
+driver = nvc.Driver(
+ {"atmos": atmos, "ocean": ocean},
+ clock=nvc.Clock("2024-01-01", "2024-03-01", "6h"),
+ connectors=[
+ ("ocean", "atmos"), # default Connector
+ nvc.Connector(atmos, ocean, window="48h", reduce="mean"),
+ ],
+)
+driver.initialize({"atmos": (x_a, coords_a), "ocean": (x_o, coords_o)})
+
+datasets = driver.run() # dict[str, xr.Dataset]
+for time, states in driver.steps(): # or notebook-style iteration
+ ...
+driver.probe("ocean->atmos") # last exchanged fields
+```
+
+When the *ordering* is the experiment, pass an explicit run sequence — a DSL
+mirroring NUOPC's runSeq. **Coupling semantics are pure ordering**: a connect
+placed before the source's run in the same slot is lagged (NUOPC-explicit)
+coupling; placed after, sequential.
+
+```
+@6h
+ atmos -> med # accumulate atmos fields into the mediator
+ ocean -> atmos # lagged: atmos sees the ocean's previous state
+ atmos
+@48h
+ med.compute
+ med -> ocean # 48h-averaged forcing
+ ocean
+@
+```
+
+`nvc.couple(*components, start=, stop=)` goes one step further and discovers
+the connections themselves by matching standard names (windowed connectors
+included, from CellMethod entries).
+
+Validation is front-loaded: unknown names (with did-you-mean suggestions),
+cadence misalignment, unmatched imports, unit mismatches, and unconsumed
+exports are all reported at `initialize`, not mid-rollout.
+
+## Training / coupled fine-tuning
+
+The exchange path is autograd-clean end to end (regrid gathers, mediator
+reductions, functional import injection). `driver.rollout(n_steps)` keeps
+the graph when gradients are enabled:
+
+```python
+with torch.enable_grad():
+ states = driver.rollout(16)
+loss = criterion(states["atmos"]["geopotential_at_1000hpa"].data, target)
+loss.backward() # gradients reach parameters of BOTH components
+```
+
+This enables jointly fine-tuning coupled emulators so each learns to
+tolerate the other's imperfect output — the standard remedy for coupled
+drift. Optimizer loops, truncated BPTT, and per-component GPU placement are
+intentionally out of scope for v1. `run()`/`steps()` execute under
+`torch.inference_mode()`.
+
+## Errors
+
+All configuration errors derive from `CouplingError` and name the components,
+the field, and the concrete fix: `UnknownFieldError`, `UnmatchedImportError`,
+`UnitsMismatchError`, `IncompatibleFieldError`, `VerticalMismatchError`,
+`CadenceError`, `SequenceError`, `AmbiguousCouplingError`.
+
+## Current limitations (v1)
+
+- HEALPix / curvilinear source grids need a user-supplied `regridder=`.
+- Units are checked, not converted (no pint); convert in a Mediator.
+- No checkpoint/restart, coupled ensembles, or concurrent slot execution.
+- `Driver` IO is in-memory xarray (`collect=True`) plus per-component
+ `IOBackend` streaming (`io=`); `couple()` does not yet configure IO.
+
+## Documentation
+
+Detailed docs live in [`docs/`](docs/) (every code snippet in them is executed
+against the toy components before inclusion):
+
+- [Concepts](docs/concepts.md) — Field/State semantics, exchange contracts,
+ adapters, the connector pipeline, coupling semantics
+- [User guide](docs/user_guide.md) — task recipes from quickstart to coupled
+ fine-tuning, each with its most likely gotcha
+- [DSL & YAML reference](docs/dsl_and_yaml_reference.md) — run-sequence grammar,
+ validation rules, full YAML schema
+- [API reference](docs/api_reference.md) — exact signatures for every public export
+- [Errors & troubleshooting](docs/errors_and_troubleshooting.md) — every
+ exception with greppable messages, causes, and fixes; silent-failure guide
+- [Design & roadmap](docs/design_and_roadmap.md) — decisions with alternatives,
+ the verification story, honest limitations, v2 plans
+
+## Examples
+
+See `examples/09_nvcoupler/`:
+
+1. `01_coupled_toy_workflow.py` — the full atmos⇄ocean loop, declared as a
+ coupling graph with a windowed connector
+2. `02_lagged_vs_sequential.py` — coupling order as a one-line experiment
+3. `03_impact_chain.py` — windowed precip-sum connector + t2m-max mediator
+ feeding an impact index
+4. `04_vertical_chemistry.py` — hybrid→pressure coupling with auto ps dependency
+5. `05_coupled_finetuning.py` — gradients across the exchange + a training step
+6. `06_pull_conditioning.py` — pull-pattern conditioning (StormCast-style):
+ a `PullAdapter` serving live coupled forcing to a model's internal fetch
+6. `06_pull_conditioning.py` — pull-pattern coupling (StormCast-style) via
+ `PullAdapter`
+
+Tests (`test/nvcoupler/`) double as executable specification, including a
+fully hand-computed 96 h coupled run in `test_driver.py`.
diff --git a/earth2studio/nvcoupler/__init__.py b/earth2studio/nvcoupler/__init__.py
new file mode 100644
index 000000000..62ce96c75
--- /dev/null
+++ b/earth2studio/nvcoupler/__init__.py
@@ -0,0 +1,132 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""nvcoupler: NUOPC/ESMF-inspired coupling framework for ML inference.
+
+Couples independent AI Earth-system components (prognostic models, data
+sources, mediators) through Fields exchanged by standard name, with
+per-component cadences, connector regridding, and a configurable run
+sequence. See earth2studio/nvcoupler/driver.py for the entry point.
+"""
+
+from .api import couple, coupled, describe, describe_html
+from .clock import Clock
+from .component import (
+ CallableComponent,
+ Component,
+ ConditioningKwargAdapter,
+ DataComponent,
+ DiagnosticComponent,
+ Exchange,
+ ExtraTensorAdapter,
+ ImportAdapter,
+ PrognosticComponent,
+ VariableOverwriteAdapter,
+)
+from .config import from_yaml, to_yaml
+from .connector import Connector
+from .dictionary import (
+ DEFAULT_DICTIONARY,
+ CellMethod,
+ FieldDictionary,
+ FieldEntry,
+)
+from .dlesym_split import (
+ DLESYM_DICTIONARY,
+ build_dlesym_driver,
+ split_dlesym,
+)
+from .driver import Driver
+from .errors import (
+ AmbiguousCouplingError,
+ CadenceError,
+ CouplingError,
+ IncompatibleFieldError,
+ SequenceError,
+ UnitsMismatchError,
+ UnknownFieldError,
+ UnmatchedImportError,
+ VerticalMismatchError,
+)
+from .field import Field, State
+from .mediator import AccumulationMediator, Mediator, TrailingAverageMediator
+from .points import PointSet
+from .pull import PullAdapter, StateDataSource
+from .sequence import (
+ ConnectAction,
+ MediateAction,
+ RunAction,
+ RunSequence,
+ Slot,
+ derive_sequence,
+ parse_run_sequence,
+)
+from .vertical import HybridLevels, PressureLevels
+
+__all__ = [
+ "Clock",
+ "CallableComponent",
+ "DataComponent",
+ "DiagnosticComponent",
+ "couple",
+ "coupled",
+ "describe",
+ "describe_html",
+ "from_yaml",
+ "to_yaml",
+ "split_dlesym",
+ "build_dlesym_driver",
+ "DLESYM_DICTIONARY",
+ "Component",
+ "ConditioningKwargAdapter",
+ "Connector",
+ "Driver",
+ "Exchange",
+ "ExtraTensorAdapter",
+ "HybridLevels",
+ "ImportAdapter",
+ "AccumulationMediator",
+ "Mediator",
+ "PointSet",
+ "PressureLevels",
+ "PrognosticComponent",
+ "PullAdapter",
+ "StateDataSource",
+ "TrailingAverageMediator",
+ "VariableOverwriteAdapter",
+ "ConnectAction",
+ "MediateAction",
+ "RunAction",
+ "RunSequence",
+ "Slot",
+ "derive_sequence",
+ "parse_run_sequence",
+ "DEFAULT_DICTIONARY",
+ "CellMethod",
+ "FieldDictionary",
+ "FieldEntry",
+ "Field",
+ "State",
+ "AmbiguousCouplingError",
+ "CadenceError",
+ "CouplingError",
+ "IncompatibleFieldError",
+ "SequenceError",
+ "UnitsMismatchError",
+ "UnknownFieldError",
+ "UnmatchedImportError",
+ "VerticalMismatchError",
+]
diff --git a/earth2studio/nvcoupler/api.py b/earth2studio/nvcoupler/api.py
new file mode 100644
index 000000000..cc6a31dcf
--- /dev/null
+++ b/earth2studio/nvcoupler/api.py
@@ -0,0 +1,386 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""User-facing API: auto-wiring, one-call runs, and plan inspection.
+
+The NUOPC analogy stops at the door here — NUOPC makes you write the runSeq;
+:func:`couple` declares the coupling graph and lets the Driver derive it.
+Matching is by field-dictionary standard name: every advertised import is
+paired with its unique exporter, derived fields (entries carrying a
+:class:`~.dictionary.CellMethod`) become windowed connectors reducing the
+base field across the cadence gap (an :class:`~.mediator.AccumulationMediator`
+is synthesized only when the pair already carries a plain transfer), and the
+run sequence follows from the graph (:func:`~.sequence.derive_sequence`).
+:func:`describe` renders the resulting plan (terraform-plan style) before a
+single tensor moves, and :func:`coupled` is the notebook one-liner from
+initial conditions to xarray Datasets.
+"""
+
+import math
+from collections.abc import Sequence as SequenceABC
+
+import numpy as np
+
+from .clock import Clock, DeltaLike, TimeLike, as_timedelta, fmt_timedelta
+from .component import Component
+from .connector import Connector
+from .driver import Driver
+from .errors import AmbiguousCouplingError, UnmatchedImportError
+from .mediator import AccumulationMediator
+from .sequence import ConnectAction, MediateAction, RunAction
+
+__all__ = ["couple", "coupled", "describe", "describe_html"]
+
+
+# ---------------------------------------------------------------------------
+# couple(): auto-wiring
+# ---------------------------------------------------------------------------
+def _short_name(component: Component, standard_name: str) -> str:
+ """Shortest registered alias of a standard name (else the name itself),
+ sanitized to the run-sequence name grammar."""
+ entry = component.dictionary.resolve(standard_name)
+ candidates = [entry.standard_name, *entry.aliases]
+ short = min(candidates, key=len)
+ return "".join(c if (c.isalnum() or c in "_-") else "_" for c in short)
+
+
+def _gcd_timestep(components: SequenceABC[Component]) -> np.timedelta64:
+ ns = [c.timestep.astype("timedelta64[ns]").astype(np.int64) for c in components]
+ return np.timedelta64(int(math.gcd(*(int(n) for n in ns))), "ns")
+
+
+def couple(
+ *components: Component,
+ start: TimeLike,
+ stop: TimeLike,
+ dt: DeltaLike | None = None,
+ connectors: list[Connector] | None = None,
+ collect: bool = True,
+) -> Driver:
+ """Auto-wire components into a ready-to-initialize :class:`Driver`.
+
+ Every advertised import is matched to its exporter by standard name.
+ Derived imports (dictionary entries with a CellMethod) whose base field
+ is exported become windowed connectors — the base exporter's fields are
+ reduced across the cadence gap on the connector itself. Only when the
+ (src, dst) pair already carries a plain transfer is an
+ AccumulationMediator synthesized instead. The run sequence is derived
+ from this graph in the canonical lagged (NUOPC-explicit) shape: connects
+ precede the runs of each slot.
+
+ Parameters
+ ----------
+ *components : Component
+ Participants (mediators may be included explicitly, or synthesized).
+ start, stop : TimeLike
+ Clock span.
+ dt : DeltaLike, optional
+ Coupling interval; defaults to the GCD of all component timesteps.
+ connectors : list[Connector], optional
+ Pre-built connectors overriding the defaults for their (src, dst).
+ collect : bool
+ Keep per-ring exports in memory for ``to_xarray()``.
+
+ Returns
+ -------
+ Driver
+ Not yet initialized — call ``driver.initialize(ics)``.
+ """
+ comps = list(components)
+ prebuilt = {(c.src.name, c.dst.name): c for c in connectors or []}
+ exports_by_comp = {c.name: list(c.export_names) for c in comps}
+
+ def exporter_of(field: str, importer: Component) -> Component | None:
+ found = [
+ c for c in comps if c.name != importer.name and field in c.export_names
+ ]
+ if len(found) > 1:
+ raise AmbiguousCouplingError(field, importer.name, [c.name for c in found])
+ return found[0] if found else None
+
+ # pass 1: direct (unique-exporter) connections
+ wired: dict[tuple[str, str], Connector] = {}
+ derived: list[tuple[Component, str]] = [] # (importer, derived std name)
+ for comp in comps:
+ for imp in comp.import_names:
+ src = exporter_of(imp, comp)
+ if src is None:
+ derived.append((comp, imp))
+ elif (src.name, comp.name) not in wired:
+ key = (src.name, comp.name)
+ wired[key] = prebuilt.pop(key, None) or Connector(src, comp)
+
+ # pass 2: derived imports — reduce a base field across the cadence gap
+ for comp, imp in derived:
+ cm = comp.dictionary.resolve(imp).cell_method
+ src = exporter_of(cm.base, comp) if cm is not None else None
+ if src is None:
+ raise UnmatchedImportError(comp.name, imp, exports_by_comp)
+ key = (src.name, comp.name)
+ pre = prebuilt.pop(key, None)
+ if pre is not None and pre.window is not None:
+ wired[key] = pre # user already declared the windowed transfer
+ elif key in wired or pre is not None:
+ # the pair already carries a plain transfer, which a windowed
+ # connector cannot share — a mediator is genuinely needed
+ if pre is not None:
+ wired[key] = pre
+ med = AccumulationMediator(
+ f"med_{_short_name(comp, imp)}", [imp], dictionary=comp.dictionary
+ )
+ comps.append(med)
+ wired[(src.name, med.name)] = Connector(src, med)
+ wired[(med.name, comp.name)] = Connector(med, comp)
+ else:
+ wired[key] = Connector(
+ src, comp, fields=[cm.base], window=cm.window, reduce=cm.method
+ )
+
+ # extra user wiring for pairs the import matching did not discover
+ for key, conn in prebuilt.items():
+ wired.setdefault(key, conn)
+
+ if dt is None:
+ dt = _gcd_timestep(comps)
+ return Driver(
+ {c.name: c for c in comps},
+ clock=Clock(start, stop, dt),
+ connectors=list(wired.values()),
+ collect=collect,
+ )
+
+
+# ---------------------------------------------------------------------------
+# coupled(): one-call convenience entry point
+# ---------------------------------------------------------------------------
+def coupled(
+ time: TimeLike,
+ stop_or_nsteps: TimeLike | int,
+ components: "SequenceABC[Component] | dict[str, Component]",
+ ics: dict[str, tuple],
+ dt: DeltaLike | None = None,
+ collect: bool = True,
+ verbose: bool = True,
+) -> dict:
+ """Build, initialize, and run a coupled system in one call.
+
+ Parameters
+ ----------
+ time : TimeLike
+ Start time.
+ stop_or_nsteps : TimeLike | int
+ Stop time, or a number of driver (dt) steps.
+ components : list[Component] | dict[str, Component]
+ Participants; a dict's values are used (keys are cosmetic).
+ ics : dict[str, tuple[torch.Tensor, CoordSystem]]
+ Initial condition per non-mediator component name.
+ dt : DeltaLike, optional
+ Coupling interval; defaults to the GCD of the component timesteps.
+ verbose : bool
+ Show a tqdm progress bar over driver steps.
+
+ Returns
+ -------
+ dict[str, xarray.Dataset]
+ One Dataset of collected exports per component.
+ """
+ from tqdm import tqdm
+
+ comps = (
+ list(components.values()) if isinstance(components, dict) else list(components)
+ )
+ dt_td = as_timedelta(dt) if dt is not None else _gcd_timestep(comps)
+ if isinstance(stop_or_nsteps, (int, np.integer)):
+ stop = np.datetime64(time) + int(stop_or_nsteps) * dt_td
+ else:
+ stop = stop_or_nsteps
+ driver = couple(*comps, start=time, stop=stop, dt=dt_td, collect=collect)
+ driver.initialize(ics)
+
+ counts = {n: c.run_count for n, c in driver.components.items()}
+ with tqdm(
+ total=driver.clock.n_steps,
+ desc="Running coupled inference",
+ disable=(not verbose),
+ ) as pbar:
+ for step_time, _ in driver.steps():
+ ran = [n for n, c in driver.components.items() if c.run_count != counts[n]]
+ counts = {n: c.run_count for n, c in driver.components.items()}
+ pbar.set_postfix_str(
+ f"{np.datetime_as_string(step_time, unit='h')} ran {'+'.join(ran)}"
+ )
+ pbar.update(1)
+ return driver.to_xarray()
+
+
+# ---------------------------------------------------------------------------
+# describe(): the coupling plan, before anything runs
+# ---------------------------------------------------------------------------
+def _table(headers: list[str], rows: list[list[str]]) -> list[str]:
+ widths = [
+ max(len(headers[i]), *(len(r[i]) for r in rows)) if rows else len(headers[i])
+ for i in range(len(headers))
+ ]
+ fmt = " ".join(f"{{:<{w}}}" for w in widths)
+ lines = [fmt.format(*headers), fmt.format(*("-" * w for w in widths))]
+ lines.extend(fmt.format(*row) for row in rows)
+ return lines
+
+
+def _connector_rows(driver: Driver) -> list[dict]:
+ """One row per ConnectAction: fields, policies, and coupling mode.
+
+ Mode is per exchange: does the destination consume state the source
+ produced in this same slot iteration (sequential — the connect follows
+ the source's run/compute in its slot) or earlier (lagged)?
+ """
+ rows = []
+ prebuilt = {(c.src.name, c.dst.name): c for c in driver._connectors.values()}
+ for slot in driver.sequence.slots:
+ produced: set[str] = set() # components that ran earlier in this slot
+ for action in slot.actions:
+ if isinstance(action, RunAction):
+ produced.add(action.component)
+ elif isinstance(action, MediateAction):
+ produced.add(action.mediator)
+ elif isinstance(action, ConnectAction):
+ conn = prebuilt.get((action.src, action.dst))
+ if conn is not None:
+ fields = conn.match()
+ time_policy, fill = conn.time_policy, conn.fill
+ else:
+ src = driver.components[action.src]
+ dst = driver.components[action.dst]
+ _, exports = src.advertise()
+ imports, _ = dst.advertise()
+ fields = [n for n in imports if n in exports]
+ time_policy, fill = "constant", "none"
+ mode = "sequential" if action.src in produced else "lagged"
+ rows.append(
+ {
+ "name": f"{action.src} -> {action.dst}",
+ "fields": fields,
+ "time_policy": time_policy,
+ "fill": fill,
+ "mode": mode,
+ "slot": fmt_timedelta(slot.interval),
+ }
+ )
+ return rows
+
+
+def describe(driver: Driver) -> str:
+ """Terraform-plan-style text summary of a coupled system.
+
+ Works before ``initialize()``: only advertised imports/exports, the run
+ sequence, and the clock are consulted (grids are unknown until realize).
+ """
+ lines = [f"Coupled system: {driver.clock!r}", "", "Components:"]
+ comp_rows = []
+ for name, comp in driver.components.items():
+ imports, exports = comp.advertise()
+ comp_rows.append(
+ [
+ name,
+ type(comp).__name__,
+ fmt_timedelta(comp.timestep),
+ ", ".join(imports) or "-",
+ ", ".join(exports) or "-",
+ ]
+ )
+ lines.extend(
+ " " + row
+ for row in _table(["name", "type", "cadence", "imports", "exports"], comp_rows)
+ )
+ lines.extend(["", "Connectors:"])
+ conn_rows = [
+ [
+ r["name"],
+ ", ".join(r["fields"]) or "-",
+ r["time_policy"],
+ r["fill"],
+ r["mode"],
+ r["slot"],
+ ]
+ for r in _connector_rows(driver)
+ ]
+ if conn_rows:
+ lines.extend(
+ " " + row
+ for row in _table(
+ ["connector", "fields", "time_policy", "fill", "mode", "slot"],
+ conn_rows,
+ )
+ )
+ else:
+ lines.append(" (none)")
+ lines.extend(["", "Run sequence:"])
+ lines.extend(" " + line for line in str(driver.sequence).splitlines())
+ return "\n".join(lines)
+
+
+_HTML_STYLE = """
+.nvc-plan { font-family: -apple-system, Segoe UI, sans-serif; color: #1a1a1a; }
+.nvc-plan h4 { margin: 0.6em 0 0.3em; }
+.nvc-boxes { display: flex; flex-wrap: wrap; gap: 10px; }
+.nvc-box { border: 1px solid #888; border-radius: 6px; padding: 8px 12px;
+ background: #f7f7f7; min-width: 180px; }
+.nvc-box .nvc-name { font-weight: 600; }
+.nvc-box .nvc-meta { font-size: 0.85em; color: #444; }
+.nvc-arrows { list-style: none; padding-left: 0; }
+.nvc-arrows li { padding: 2px 0; font-size: 0.9em; }
+.nvc-arrow { color: #0a6; font-weight: 600; }
+.nvc-seq { background: #f0f0f0; border-radius: 6px; padding: 8px 12px;
+ font-size: 0.9em; }
+"""
+
+
+def describe_html(driver: Driver) -> str:
+ """Self-contained HTML rendering of the coupling plan (for Jupyter)."""
+ import html
+
+ def esc(s: str) -> str:
+ return html.escape(str(s))
+
+ boxes = []
+ for name, comp in driver.components.items():
+ imports, exports = comp.advertise()
+ boxes.append(
+ f'
{esc(name)}
'
+ f'
{esc(type(comp).__name__)} @ '
+ f"{esc(fmt_timedelta(comp.timestep))}
"
+ f"imports: {esc(', '.join(imports) or '-')}
"
+ f"exports: {esc(', '.join(exports) or '-')}
"
+ )
+ arrows = []
+ for r in _connector_rows(driver):
+ src, dst = r["name"].split(" -> ")
+ arrows.append(
+ f'{esc(src)} → {esc(dst)}: '
+ f"{esc(', '.join(r['fields']) or '-')} "
+ f"[{esc(r['time_policy'])}, fill={esc(r['fill'])}, {esc(r['mode'])}, "
+ f"@{esc(r['slot'])}]"
+ )
+ seq = "
".join(esc(line) for line in str(driver.sequence).splitlines())
+ return (
+ f""
+ ''
+ f"
Coupled system
{esc(repr(driver.clock))}
"
+ f'
Components
{"".join(boxes)}
'
+ f'
Connectors
'
+ f'
Run sequence
{seq}
'
+ "
"
+ )
diff --git a/earth2studio/nvcoupler/clock.py b/earth2studio/nvcoupler/clock.py
new file mode 100644
index 000000000..76758fefa
--- /dev/null
+++ b/earth2studio/nvcoupler/clock.py
@@ -0,0 +1,147 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Clock: coupled-system time management (ESMF analog).
+
+The Driver owns one Clock stepping at the coupling interval dt; each
+Component declares its own timestep, and the Driver runs a component only
+on driver steps aligned with that cadence, which is how a 6 h atmosphere
+and a 48 h ocean coexist in one loop.
+
+All times are np.datetime64 / np.timedelta64 to match earth2studio's
+TimeArray / LeadTimeArray conventions.
+"""
+
+from collections.abc import Iterator
+
+import numpy as np
+
+from .errors import CadenceError
+
+TimeLike = str | np.datetime64
+DeltaLike = str | np.timedelta64
+
+
+def as_datetime(t: TimeLike) -> np.datetime64:
+ """Coerce to nanosecond datetime64 (accepts ISO strings)."""
+ return np.datetime64(t).astype("datetime64[ns]")
+
+
+def as_timedelta(d: DeltaLike) -> np.timedelta64:
+ """Coerce to nanosecond timedelta64. Strings use '' (e.g. '6h')."""
+ if isinstance(d, np.timedelta64):
+ return d.astype("timedelta64[ns]")
+ if isinstance(d, (bool, int, np.integer)):
+ raise ValueError(
+ f"Bare number {d!r} is ambiguous as a timedelta (hours? steps?) — "
+ "pass a string like '6h' or '2D', or a np.timedelta64"
+ )
+ if isinstance(d, str):
+ s = d.strip()
+ i = 0
+ while i < len(s) and (s[i].isdigit() or s[i] == "-"):
+ i += 1
+ if i == 0 or i == len(s):
+ raise ValueError(f"Cannot parse timedelta {d!r}; expected e.g. '6h', '2D'")
+ value, unit = int(s[:i]), s[i:]
+ unit = {"d": "D", "H": "h", "min": "m", "S": "s"}.get(unit, unit)
+ return np.timedelta64(value, unit).astype("timedelta64[ns]")
+ raise TypeError(f"Cannot interpret {d!r} as a timedelta")
+
+
+def fmt_timedelta(d: np.timedelta64) -> str:
+ """Human-readable timedelta: whole hours/days where possible."""
+ ns = d.astype("timedelta64[ns]").astype(np.int64)
+ hour = 3_600_000_000_000
+ if ns % (24 * hour) == 0:
+ return f"{ns // (24 * hour)}D"
+ if ns % hour == 0:
+ return f"{ns // hour}h"
+ if ns % 60_000_000_000 == 0:
+ return f"{ns // 60_000_000_000}m"
+ return str(d)
+
+
+def is_multiple(interval: np.timedelta64, dt: np.timedelta64) -> bool:
+ interval_ns = interval.astype("timedelta64[ns]").astype(np.int64)
+ dt_ns = dt.astype("timedelta64[ns]").astype(np.int64)
+ return dt_ns > 0 and interval_ns > 0 and interval_ns % dt_ns == 0
+
+
+class Clock:
+ """Driver clock stepping from start to stop (inclusive) at dt.
+
+ Iterating yields each time after the start: the initial condition is at
+ ``start`` and is the caller's step 0; the first yielded time is
+ ``start + dt``. This mirrors run.py where the iterator's 0th output is
+ the IC and each subsequent step advances one dt.
+ """
+
+ def __init__(self, start: TimeLike, stop: TimeLike, dt: DeltaLike):
+ self.start = as_datetime(start)
+ self.stop = as_datetime(stop)
+ self.dt = as_timedelta(dt)
+ if self.dt <= np.timedelta64(0, "ns"):
+ raise ValueError(f"Clock dt must be positive, got {dt!r}")
+ if self.stop <= self.start:
+ raise ValueError(f"Clock stop {stop!r} must be after start {start!r}")
+ span = (self.stop - self.start).astype("timedelta64[ns]")
+ if not is_multiple(span, self.dt):
+ raise CadenceError("Clock span (stop - start)", str(span), str(self.dt))
+ self._step = 0
+
+ @property
+ def current(self) -> np.datetime64:
+ return self.start + self._step * self.dt
+
+ @property
+ def step_index(self) -> int:
+ return self._step
+
+ @property
+ def n_steps(self) -> int:
+ span_ns = (self.stop - self.start).astype("timedelta64[ns]").astype(np.int64)
+ return int(span_ns // self.dt.astype(np.int64))
+
+ def elapsed(self) -> np.timedelta64:
+ return self.current - self.start
+
+ def done(self) -> bool:
+ return self.current >= self.stop
+
+ def advance(self) -> np.datetime64:
+ if self.done():
+ raise StopIteration(f"Clock already at stop time {self.stop}")
+ self._step += 1
+ return self.current
+
+ def reset(self) -> None:
+ self._step = 0
+
+ def times(self) -> np.ndarray:
+ """All times from start to stop inclusive (length n_steps + 1)."""
+ return self.start + np.arange(self.n_steps + 1) * self.dt
+
+ def __iter__(self) -> Iterator[np.datetime64]:
+ while not self.done():
+ yield self.advance()
+
+ def __repr__(self) -> str:
+ return (
+ f"Clock({np.datetime_as_string(self.start, unit='m')} -> "
+ f"{np.datetime_as_string(self.stop, unit='m')}, "
+ f"dt={fmt_timedelta(self.dt)}, step={self._step})"
+ )
diff --git a/earth2studio/nvcoupler/component.py b/earth2studio/nvcoupler/component.py
new file mode 100644
index 000000000..1ed134c8c
--- /dev/null
+++ b/earth2studio/nvcoupler/component.py
@@ -0,0 +1,797 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Components: the NUOPC_Model analog wrapping steppable things.
+
+A Component owns an internal model state, an import State (fields other
+components provide) and an export State (fields it offers), and a timestep
+defining its cadence. The Driver calls the NUOPC-style phases:
+advertise -> realize(clock) -> initialize(x, coords) -> run(time)* -> finalize.
+
+The critical seam is the :class:`ImportAdapter`: real models receive coupled
+fields in different call shapes (state-variable overwrite, conditioning
+kwarg, extra input tensor), so the adapter — not the component — owns the
+model invocation. The adapter receives everything about the step bundled in
+an :class:`Exchange`.
+"""
+
+import abc
+from collections import OrderedDict
+from collections.abc import Callable, Iterable, Mapping
+from dataclasses import dataclass
+from dataclasses import field as dc_field
+from typing import Any, Protocol, runtime_checkable
+
+import numpy as np
+import torch
+
+from earth2studio.utils.type import CoordSystem
+
+from .clock import Clock, DeltaLike, as_datetime, as_timedelta, is_multiple
+from .dictionary import DEFAULT_DICTIONARY, FieldDictionary
+from .errors import CadenceError, CouplingError
+from .field import State
+from .points import PointSet
+
+StepFn = Callable[[torch.Tensor, CoordSystem], tuple[torch.Tensor, CoordSystem]]
+NextInputFn = Callable[
+ [torch.Tensor, CoordSystem, torch.Tensor, CoordSystem],
+ tuple[torch.Tensor, CoordSystem],
+]
+
+
+# ---------------------------------------------------------------------------
+# Exchange and import adapters
+# ---------------------------------------------------------------------------
+def _broadcast_to_slice(data: torch.Tensor, slice_shape: torch.Size) -> torch.Tensor:
+ """Left-pad `data` with singleton dims and expand to `slice_shape`."""
+ if data.ndim > len(slice_shape):
+ raise CouplingError(
+ f"Imported field with shape {tuple(data.shape)} has more dims than "
+ f"the model state slice {tuple(slice_shape)}"
+ )
+ data = data.reshape((1,) * (len(slice_shape) - data.ndim) + tuple(data.shape))
+ return data.expand(slice_shape)
+
+
+def _stacking_order(
+ field_order: list[str] | None, imports: State, who: str
+) -> list[str]:
+ """Resolve the channel order for stacking imported fields.
+
+ Models are channel-order-sensitive; silently stacking in alphabetical
+ order would feed permuted inputs that run fine and predict garbage. With
+ more than one import the order must therefore be explicit.
+ """
+ if field_order is not None:
+ missing = [n for n in field_order if n not in imports]
+ if missing:
+ raise CouplingError(
+ f"{who}: field_order names {missing} are not in the "
+ f"import state (present: {sorted(imports)})"
+ )
+ return list(field_order)
+ if len(imports) > 1:
+ raise CouplingError(
+ f"{who}: {len(imports)} imported fields but no field_order= "
+ "given — the model's channel order cannot be inferred. Pass "
+ f"field_order=[...] with an explicit ordering of {sorted(imports)}"
+ )
+ return list(imports)
+
+
+@dataclass(frozen=True)
+class Exchange:
+ """Everything an :class:`ImportAdapter` needs for one coupled model step.
+
+ Attributes
+ ----------
+ x : torch.Tensor
+ The component's current model state tensor.
+ coords : CoordSystem
+ Coordinates of ``x``.
+ imports : State
+ The imported Fields available this step (already subset to the
+ component's advertised imports).
+ std_to_raw : Mapping[str, str]
+ Standard field name -> this model's raw variable name.
+ time : np.datetime64 | None
+ The valid time the step advances to.
+
+ The two accessors cover the common delivery shapes: :meth:`inject` for
+ state-variable overwrite and :meth:`stacked` for channel-stacked forcing
+ tensors. Both are pure torch and autograd-safe.
+ """
+
+ x: torch.Tensor
+ coords: CoordSystem
+ imports: State
+ std_to_raw: Mapping[str, str] = dc_field(default_factory=dict)
+ time: np.datetime64 | None = None
+
+ def inject(self) -> torch.Tensor:
+ """Return ``x`` with each imported field overwritten into its matching
+ variable slice (resolved through ``std_to_raw``).
+
+ The overwrite composes a cloned tensor via index_copy_ on the clone,
+ which keeps the autograd graph intact.
+ """
+ x, imports = self.x, self.imports
+ if not imports:
+ return x
+ if "variable" not in self.coords:
+ raise CouplingError(
+ "Exchange.inject requires a 'variable' dim in the model state "
+ "coords; use ConditioningKwargAdapter or ExtraTensorAdapter "
+ "for models without one"
+ )
+ var_axis = list(self.coords).index("variable")
+ variables = self.coords["variable"]
+ x = x.clone()
+ for std_name, field in imports.items():
+ raw = self.std_to_raw.get(std_name, std_name)
+ pos = np.flatnonzero(variables == raw)
+ if pos.size == 0:
+ raise CouplingError(
+ f"Imported field {std_name!r} (model name {raw!r}) is not a "
+ f"state variable of the model (variables: {list(variables)}). "
+ "If the model takes forcing as a conditioning kwarg or an "
+ "extra tensor, pass the matching ImportAdapter."
+ )
+ slice_shape = x.select(var_axis, int(pos[0])).shape
+ data = _broadcast_to_slice(field.data, slice_shape).unsqueeze(var_axis)
+ index = torch.tensor([int(pos[0])], device=x.device)
+ x.index_copy_(var_axis, index, data.to(dtype=x.dtype, device=x.device))
+ return x
+
+ def stacked(
+ self, field_order: list[str] | None = None, *, who: str = "Exchange.stacked"
+ ) -> tuple[torch.Tensor, CoordSystem]:
+ """Stack the imported fields into one tensor with a leading
+ 'variable' dim.
+
+ ``field_order`` fixes the channel order; with more than one import it
+ is required (see :func:`_stacking_order`). ``who`` names the caller in
+ error messages.
+ """
+ names = _stacking_order(field_order, self.imports, who)
+ return self.imports.as_tensor(names)
+
+
+@runtime_checkable
+class ImportAdapter(Protocol):
+ """Runs one model step with the import State injected.
+
+ The adapter owns the model call because coupled models disagree on how
+ forcing arrives: DLESyM/PhysicsNeMo expect an extra input tensor,
+ StormScope expects a conditioning kwarg, and prescribed-forcing setups
+ overwrite state variables. Implementations must be autograd-safe (no
+ in-place mutation of tensors that may carry grad).
+ """
+
+ def __call__(
+ self, model: Any, exchange: Exchange
+ ) -> tuple[torch.Tensor, CoordSystem]: ...
+
+
+class VariableOverwriteAdapter:
+ """Default adapter: overwrite matching variable slices in x, then call
+ ``model(x, coords)``.
+
+ Suits toys and prescribed-forcing-as-state setups where imported fields
+ are also state variables of the model (e.g. an atmosphere carrying sst
+ as an input channel). See :meth:`Exchange.inject`.
+ """
+
+ def __call__(
+ self, model: Any, exchange: Exchange
+ ) -> tuple[torch.Tensor, CoordSystem]:
+ return model(exchange.inject(), exchange.coords)
+
+
+class ConditioningKwargAdapter:
+ """Pass imports as a conditioning tensor kwarg (StormScope pattern).
+
+ Calls ``model.call_with_conditioning(x, coords, conditioning=...,
+ conditioning_coords=...)`` with the stacked import fields.
+ """
+
+ def __init__(
+ self,
+ field_order: list[str] | None = None,
+ method: str = "call_with_conditioning",
+ ):
+ self.field_order = field_order
+ self.method = method
+
+ def __call__(
+ self, model: Any, exchange: Exchange
+ ) -> tuple[torch.Tensor, CoordSystem]:
+ conditioning, conditioning_coords = exchange.stacked(
+ self.field_order, who=type(self).__name__
+ )
+ fn = getattr(model, self.method)
+ return fn(
+ exchange.x,
+ exchange.coords,
+ conditioning=conditioning,
+ conditioning_coords=conditioning_coords,
+ )
+
+
+class ExtraTensorAdapter:
+ """Pass imports as an extra positional tensor (DLESyM / PhysicsNeMo
+ 4-tensor pattern): ``model(x, coords, coupling)`` by default.
+ """
+
+ def __init__(self, field_order: list[str] | None = None, kwarg: str | None = None):
+ self.field_order = field_order
+ self.kwarg = kwarg
+
+ def __call__(
+ self, model: Any, exchange: Exchange
+ ) -> tuple[torch.Tensor, CoordSystem]:
+ coupling, _ = exchange.stacked(self.field_order, who=type(self).__name__)
+ if self.kwarg is not None:
+ return model(exchange.x, exchange.coords, **{self.kwarg: coupling})
+ return model(exchange.x, exchange.coords, coupling)
+
+
+# ---------------------------------------------------------------------------
+# Component base
+# ---------------------------------------------------------------------------
+class Component(abc.ABC):
+ """NUOPC_Model analog: a steppable participant in the coupled system."""
+
+ # Whether initialize() needs an (x, coords) initial condition. Subclasses
+ # whose initialize() is safely callable with no arguments (mediators,
+ # data components, diagnostics) set this False so the Driver can
+ # initialize them without an ics entry.
+ requires_ic: bool = True
+
+ def __init__(
+ self,
+ name: str,
+ timestep: DeltaLike,
+ imports: Iterable[str] = (),
+ exports: Iterable[str] = (),
+ dictionary: FieldDictionary | None = None,
+ variable_aliases: Mapping[str, str] | None = None,
+ export_masks: Mapping[str, torch.Tensor] | None = None,
+ import_vertical: Mapping[str, Any] | None = None,
+ export_vertical: Mapping[str, Any] | None = None,
+ points: PointSet | None = None,
+ ):
+ self.name = name
+ # Scattered sample-location grid (stations, sites, query points),
+ # the "point" analog of a lat/lon mesh. When set, grid_coords()
+ # reports this instead of any lat/lon coords the component happens
+ # to carry internally, and a Connector delivering to this component
+ # samples onto these locations (see connector.py's `sample=`).
+ self.points = points
+ self.timestep = as_timedelta(timestep)
+ self.dictionary = FieldDictionary(dictionary or DEFAULT_DICTIONARY)
+ # variable_aliases: raw model variable name -> standard name
+ self._raw_to_std: dict[str, str] = dict(variable_aliases or {})
+ for raw, std in self._raw_to_std.items():
+ if raw not in self.dictionary:
+ self.dictionary.add_alias(std, raw)
+ self._std_to_raw = {std: raw for raw, std in self._raw_to_std.items()}
+ self.import_names = [self.dictionary.standard_name(n) for n in imports]
+ self.export_names = [self.dictionary.standard_name(n) for n in exports]
+ self.export_masks = dict(export_masks or {})
+ # std name -> VerticalCoordinate this component expects imports on /
+ # publishes exports on (only fields with a "level" dim need these)
+ self.import_vertical: dict[str, Any] = dict(import_vertical or {})
+ self.export_vertical: dict[str, Any] = dict(export_vertical or {})
+ self.import_state = State(f"{name}.imports")
+ self.export_state = State(f"{name}.exports")
+ self.clock: Clock | None = None
+ self.run_count = 0
+
+ # -- NUOPC phases ---------------------------------------------------------
+ def advertise(self) -> tuple[list[str], list[str]]:
+ return list(self.import_names), list(self.export_names)
+
+ def realize(self, clock: Clock) -> None:
+ if not is_multiple(self.timestep, clock.dt):
+ raise CadenceError(
+ f"Component {self.name!r} timestep", str(self.timestep), str(clock.dt)
+ )
+ self.clock = clock
+
+ @abc.abstractmethod
+ def initialize(self, x: torch.Tensor, coords: CoordSystem) -> None:
+ """Set internal state from an initial condition and seed export_state
+ (so lagged coupling has data at t0)."""
+
+ @abc.abstractmethod
+ def run(self, time: np.datetime64) -> None:
+ """Advance one component timestep; exports become valid at `time`."""
+
+ def finalize(self) -> None:
+ pass
+
+ # -- helpers ---------------------------------------------------------------
+ def _exchange(
+ self, x: torch.Tensor, coords: CoordSystem, time: np.datetime64
+ ) -> Exchange:
+ """Bundle the current step's state and imports for an ImportAdapter."""
+ imports = self.import_state.subset(
+ [n for n in self.import_names if n in self.import_state]
+ )
+ return Exchange(x, coords, imports, self.resolve_std_to_raw(coords), time)
+
+ def grid_coords(self) -> CoordSystem | None:
+ """Spatial coordinates of this component's grid (None = no grid of
+ its own, e.g. mediators — connectors then pass fields through).
+
+ A component with ``points`` set (a scattered sample-location target)
+ reports that instead of any mesh coords it happens to carry
+ internally — the point set is the authoritative spatial target for
+ everything a Connector delivers to it.
+ """
+ if self.points is not None:
+ return self.points.grid_coords()
+ coords = getattr(self, "_coords", None)
+ if coords is None:
+ return None
+ from .field import _SPATIAL_DIMS # local import to avoid cycle at module load
+
+ spatial = OrderedDict((k, v) for k, v in coords.items() if k in _SPATIAL_DIMS)
+ return spatial or None
+
+ def resolve_std_to_raw(self, coords: CoordSystem) -> dict[str, str]:
+ """Map standard names to this model's raw variable names, derived from
+ the actual variable coordinate plus any explicit variable_aliases."""
+ mapping: dict[str, str] = {}
+ for raw in coords.get("variable", ()): # type: ignore[union-attr]
+ raw = str(raw)
+ if raw in self.dictionary:
+ mapping[self.dictionary.standard_name(raw)] = raw
+ mapping.update(self._std_to_raw)
+ return mapping
+
+ def publish(
+ self, x: torch.Tensor, coords: CoordSystem, valid_time: np.datetime64
+ ) -> None:
+ """Populate export_state from a model output tensor."""
+ state = State.from_tensor(
+ f"{self.name}.exports",
+ x,
+ coords,
+ self.dictionary,
+ valid_time=valid_time,
+ source=self.name,
+ strict=False,
+ )
+ for std_name in self.export_names:
+ if std_name not in state:
+ raise CouplingError(
+ f"Component {self.name!r} advertises export {std_name!r} but "
+ f"its output variables are {list(coords.get('variable', []))}"
+ )
+ field = state[std_name]
+ if std_name in self.export_masks:
+ field.mask = self.export_masks[std_name]
+ if std_name in self.export_vertical:
+ field.vertical = self.export_vertical[std_name]
+ self.export_state.add(field)
+
+ def __repr__(self) -> str:
+ from .clock import fmt_timedelta
+
+ return (
+ f"{type(self).__name__}({self.name!r}, dt={fmt_timedelta(self.timestep)}, "
+ f"imports={self.import_names}, exports={self.export_names})"
+ )
+
+
+# ---------------------------------------------------------------------------
+# Concrete components
+# ---------------------------------------------------------------------------
+class CallableComponent(Component):
+ """Wraps a plain ``fn(x, coords) -> (x, coords)`` step function.
+
+ The workhorse for synthetic components and non-ML models (any Python
+ process model can join the coupled system through this class).
+ """
+
+ def __init__(
+ self,
+ name: str,
+ fn: StepFn,
+ timestep: DeltaLike,
+ imports: Iterable[str] = (),
+ exports: Iterable[str] = (),
+ import_adapter: ImportAdapter | None = None,
+ **kwargs: Any,
+ ):
+ super().__init__(name, timestep, imports, exports, **kwargs)
+ self.fn = fn
+ self.import_adapter: ImportAdapter = (
+ import_adapter or VariableOverwriteAdapter()
+ )
+ self._x: torch.Tensor | None = None
+ self._coords: CoordSystem | None = None
+
+ def initialize(self, x: torch.Tensor, coords: CoordSystem) -> None:
+ self._x, self._coords = x, OrderedDict(coords)
+ start = self.clock.start if self.clock is not None else None
+ self.publish(x, self._coords, valid_time=start)
+
+ def run(self, time: np.datetime64) -> None:
+ if self._x is None:
+ raise CouplingError(f"Component {self.name!r} not initialized")
+ y, ycoords = self.import_adapter(
+ self.fn, self._exchange(self._x, self._coords, time)
+ )
+ self._x, self._coords = y, OrderedDict(ycoords)
+ self.publish(y, ycoords, valid_time=time)
+ self.run_count += 1
+
+ @property
+ def state(self) -> tuple[torch.Tensor, CoordSystem]:
+ return self._x, self._coords
+
+
+class PrognosticComponent(Component):
+ """Wraps an earth2studio PrognosticModel (``models/px/base.py``).
+
+ Owns the model state tensor and steps by calling the model directly
+ (rather than ``create_iterator``) so imports can be injected between
+ steps. Models that manage multi-window inputs internally need a
+ ``next_input`` hook mapping (prev_x, prev_coords, out, out_coords) to the
+ next step's input; the default handles single-window models by reusing
+ the output with the model's input lead_time coordinates.
+ """
+
+ def __init__(
+ self,
+ name: str,
+ model: Any,
+ timestep: DeltaLike | None = None,
+ imports: Iterable[str] = (),
+ exports: Iterable[str] | None = None,
+ import_adapter: ImportAdapter | None = None,
+ next_input: NextInputFn | None = None,
+ **kwargs: Any,
+ ):
+ self.model = model
+ in_coords = model.input_coords()
+ out_coords = model.output_coords(in_coords)
+ if timestep is None:
+ timestep = (
+ out_coords["lead_time"][-1] - in_coords["lead_time"][-1]
+ ).astype("timedelta64[ns]")
+ if exports is None:
+ exports = []
+ dictionary = kwargs.get("dictionary") or DEFAULT_DICTIONARY
+ aliases = kwargs.get("variable_aliases") or {}
+ for raw in out_coords["variable"]:
+ raw = str(raw)
+ if raw in aliases:
+ exports.append(aliases[raw])
+ elif raw in dictionary:
+ exports.append(dictionary.standard_name(raw))
+ super().__init__(name, timestep, imports, exports, **kwargs)
+ self.import_adapter: ImportAdapter = (
+ import_adapter or VariableOverwriteAdapter()
+ )
+ self.next_input = next_input or self._default_next_input
+ self._x: torch.Tensor | None = None
+ self._coords: CoordSystem | None = None
+
+ def _default_next_input(
+ self,
+ prev_x: torch.Tensor,
+ prev_coords: CoordSystem,
+ out: torch.Tensor,
+ out_coords: CoordSystem,
+ ) -> tuple[torch.Tensor, CoordSystem]:
+ in_lead = self.model.input_coords()["lead_time"]
+ out_lead = out_coords.get("lead_time", np.empty(0))
+ if len(out_lead) != len(in_lead):
+ raise CouplingError(
+ f"Component {self.name!r}: model outputs {len(out_lead)} lead "
+ f"times but takes {len(in_lead)} as input — supply a "
+ "next_input hook to manage the sliding input window"
+ )
+ coords = OrderedDict(out_coords)
+ coords["lead_time"] = in_lead.copy()
+ return out, coords
+
+ def initialize(self, x: torch.Tensor, coords: CoordSystem) -> None:
+ self._x, self._coords = x, OrderedDict(coords)
+ start = self.clock.start if self.clock is not None else None
+ self.publish(x, self._coords, valid_time=start)
+
+ def run(self, time: np.datetime64) -> None:
+ if self._x is None:
+ raise CouplingError(f"Component {self.name!r} not initialized")
+ y, ycoords = self.import_adapter(
+ self.model, self._exchange(self._x, self._coords, time)
+ )
+ self._x, self._coords = self.next_input(self._x, self._coords, y, ycoords)
+ self.publish(y, ycoords, valid_time=time)
+ self.run_count += 1
+
+ def publish(
+ self, x: torch.Tensor, coords: CoordSystem, valid_time: np.datetime64
+ ) -> None:
+ """Publish exchange-shaped exports: singleton batch/time/lead_time
+ dims (and their stale size-1 coord values, e.g. the model's 'time')
+ are squeezed away so exported Fields carry plain spatial coords like
+ every other component's. The internal model state (`self._x`,
+ `self._coords`) keeps the full model dims."""
+ x, coords = _squeeze_singletons(x, coords)
+ super().publish(x, coords, valid_time)
+
+ def to(self, device: Any) -> "PrognosticComponent":
+ self.model = self.model.to(device)
+ if self._x is not None:
+ self._x = self._x.to(device)
+ return self
+
+ @property
+ def state(self) -> tuple[torch.Tensor, CoordSystem]:
+ return self._x, self._coords
+
+
+def _squeeze_singletons(
+ x: torch.Tensor,
+ coords: CoordSystem,
+ dims: tuple[str, ...] = ("batch", "time", "lead_time"),
+) -> tuple[torch.Tensor, CoordSystem]:
+ """Drop size-1 batch/time/lead_time dims so published Fields carry the
+ plain spatial coords (lat, lon) the rest of the coupler expects."""
+ out = OrderedDict(coords)
+ keys = list(out)
+ for axis in range(len(keys) - 1, -1, -1):
+ key = keys[axis]
+ if key in dims and x.shape[axis] == 1:
+ x = x.squeeze(axis)
+ del out[key]
+ return x, out
+
+
+class DataComponent(Component):
+ """Prescribed-forcing component wrapping an earth2studio DataSource.
+
+ The NUOPC "data component" analog: instead of stepping a model it
+ fetches fields from a data source (ERA5, GFS, an OISST archive, ...) at
+ its own cadence and publishes them as exports. Swapping a modeled ocean
+ for ``DataComponent("ocean", source=wb2, exports=
+ ["sea_surface_temperature"], timestep="24h")`` turns a two-way coupled
+ system into a prescribed-SST run with no other changes — the connectors,
+ mediators, and run sequence are untouched.
+
+ Parameters
+ ----------
+ name : str
+ source : DataSource
+ earth2studio data source; called through
+ :func:`earth2studio.data.utils.fetch_data`.
+ exports : Iterable[str]
+ Standard names (or aliases) to fetch and export.
+ timestep : DeltaLike
+ Fetch cadence (e.g. "24h" for daily analysis fields).
+ variable_map : Mapping[str, str], optional
+ standard name -> raw source variable name, for sources whose
+ vocabulary is not in the field dictionary (the raw name is also
+ registered as an alias so exports resolve). Exports without an entry
+ fall back to the dictionary's aliases.
+ interp_to : CoordSystem, optional
+ Forwarded to fetch_data for source-side regridding; usually left
+ None so the Connector regrids onto each destination grid instead.
+ device : torch.device | str
+ Device fetched tensors are loaded to, by default "cpu".
+ """
+
+ # initialize() with no arguments fetches at clock.start — no IC needed
+ requires_ic = False
+
+ def __init__(
+ self,
+ name: str,
+ source: Any,
+ exports: Iterable[str],
+ timestep: DeltaLike,
+ variable_map: Mapping[str, str] | None = None,
+ interp_to: CoordSystem | None = None,
+ device: Any = "cpu",
+ **kwargs: Any,
+ ):
+ super().__init__(name, timestep, imports=(), exports=exports, **kwargs)
+ self.source = source
+ self.interp_to = interp_to
+ self.device = device
+ self._variable_map: dict[str, str] = {}
+ for std, raw in (variable_map or {}).items():
+ std = self.dictionary.standard_name(std)
+ self._variable_map[std] = raw
+ if raw not in self.dictionary:
+ self.dictionary.add_alias(std, raw)
+ self._coords: CoordSystem | None = None
+
+ def _raw_name(self, std_name: str) -> str:
+ """Source variable name for a standard name: variable_map first, then
+ explicit variable_aliases, then the dictionary's aliases."""
+ if std_name in self._variable_map:
+ return self._variable_map[std_name]
+ if std_name in self._std_to_raw:
+ return self._std_to_raw[std_name]
+ entry = self.dictionary.resolve(std_name)
+ return min(entry.aliases) if entry.aliases else std_name
+
+ def _fetch(self, time: np.datetime64) -> tuple[torch.Tensor, CoordSystem]:
+ # Local import: keeps nvcoupler importable without pulling the whole
+ # data-source dependency stack until a DataComponent actually runs.
+ from earth2studio.data.utils import fetch_data
+
+ raw_names = [self._raw_name(n) for n in self.export_names]
+ x, coords = fetch_data(
+ self.source,
+ time=np.array([as_datetime(time)]),
+ variable=np.array(raw_names),
+ device=self.device,
+ interp_to=self.interp_to,
+ )
+ return _squeeze_singletons(x, coords)
+
+ def initialize(
+ self, x: torch.Tensor | None = None, coords: CoordSystem | None = None
+ ) -> None:
+ """Seed exports at the clock start (lagged coupling needs t0 data).
+
+ Needs no initial condition: with ``x``/``coords`` omitted the source
+ is queried at ``clock.start``. An explicit (x, coords) pair — with a
+ "variable" dim — is published as-is instead (e.g. to avoid a fetch
+ in tests or restarts).
+ """
+ if x is not None and coords is not None:
+ data, dcoords = x, OrderedDict(coords)
+ else:
+ if self.clock is None:
+ raise CouplingError(
+ f"DataComponent {self.name!r} cannot fetch initial data "
+ "before realize(clock) — the driver calls realize first, "
+ "or pass an explicit (x, coords) initial condition"
+ )
+ data, dcoords = self._fetch(self.clock.start)
+ self._coords = dcoords
+ start = self.clock.start if self.clock is not None else None
+ self.publish(data, dcoords, valid_time=start)
+
+ def run(self, time: np.datetime64) -> None:
+ x, coords = self._fetch(time)
+ self._coords = coords
+ self.publish(x, coords, valid_time=as_datetime(time))
+ self.run_count += 1
+
+
+class DiagnosticComponent(Component):
+ """Wraps an earth2studio DiagnosticModel (``models/dx/base.py``).
+
+ A stateless single-step transform: each run it stacks its imported
+ Fields into the model's expected variable order, calls
+ ``model(x, coords)``, and publishes the outputs — no internal time
+ state. Import/export lists default to the model's own
+ ``input_coords()``/``output_coords()`` variables resolved through the
+ field dictionary, so registered diagnostics wire up with just a name,
+ the model, and a cadence.
+ """
+
+ # stateless transform: no-arg initialize() derives the grid from the
+ # model's input_coords(), so no IC tensor is needed
+ requires_ic = False
+
+ def __init__(
+ self,
+ name: str,
+ model: Any,
+ timestep: DeltaLike,
+ imports: Iterable[str] | None = None,
+ exports: Iterable[str] | None = None,
+ **kwargs: Any,
+ ):
+ self.model = model
+ in_coords = model.input_coords()
+ out_coords = model.output_coords(in_coords)
+ self._input_raw = [str(v) for v in in_coords["variable"]]
+ dictionary = kwargs.get("dictionary") or DEFAULT_DICTIONARY
+ aliases = dict(kwargs.get("variable_aliases") or {})
+ if imports is None:
+ # every model input must resolve — a missing entry means the
+ # coupler cannot know what to wire in, so resolve() raises with
+ # suggestions rather than silently dropping the variable
+ imports = [
+ aliases[raw] if raw in aliases else dictionary.standard_name(raw)
+ for raw in self._input_raw
+ ]
+ if exports is None:
+ exports = []
+ for raw in out_coords["variable"]:
+ raw = str(raw)
+ if raw in aliases:
+ exports.append(aliases[raw])
+ elif raw in dictionary:
+ exports.append(dictionary.standard_name(raw))
+ super().__init__(name, timestep, imports, exports, **kwargs)
+ self._coords: CoordSystem | None = OrderedDict(in_coords)
+
+ def initialize(
+ self, x: torch.Tensor | None = None, coords: CoordSystem | None = None
+ ) -> None:
+ """No state tensor to set; records the model grid for grid_coords().
+
+ An optional (x, coords) input — in the model's own vocabulary — is
+ pushed through the model once to seed exports at t0 (lagged chains).
+ """
+ self._coords = OrderedDict(self.model.input_coords())
+ if x is not None and coords is not None:
+ y, ycoords = self.model(x, OrderedDict(coords))
+ y, ycoords = _squeeze_singletons(y, ycoords)
+ start = self.clock.start if self.clock is not None else None
+ self.publish(y, ycoords, valid_time=start)
+
+ def _conform_to_input(
+ self, x: torch.Tensor, coords: CoordSystem, time: np.datetime64
+ ) -> tuple[torch.Tensor, CoordSystem]:
+ """Add singleton dims (batch, time, lead_time, ...) the model's
+ input_coords declare but the stacked import Fields lack."""
+ in_coords = self.model.input_coords()
+ out: CoordSystem = OrderedDict()
+ for key, value in in_coords.items():
+ if key in coords:
+ out[key] = coords[key]
+ continue
+ x = x.unsqueeze(len(out))
+ if key == "time":
+ out[key] = np.array([as_datetime(time)])
+ elif key == "lead_time":
+ out[key] = np.array([np.timedelta64(0, "h")], dtype="timedelta64[ns]")
+ else:
+ out[key] = np.asarray(value) # e.g. batch: np.empty(0)
+ for key, value in coords.items():
+ if key not in out:
+ out[key] = value
+ return x, out
+
+ def run(self, time: np.datetime64) -> None:
+ missing = [n for n in self.import_names if n not in self.import_state]
+ if missing:
+ raise CouplingError(
+ f"DiagnosticComponent {self.name!r} is missing imports "
+ f"{missing} at {time} (present: {sorted(self.import_state)}) — "
+ "check the run sequence connects its source before this "
+ "component runs"
+ )
+ # stack imports in the model's raw variable order, then relabel the
+ # variable coordinate back to the model's vocabulary
+ std_order = [self.dictionary.standard_name(raw) for raw in self._input_raw]
+ x, coords = self.import_state.as_tensor(std_order)
+ coords = OrderedDict(coords)
+ coords["variable"] = np.array(self._input_raw)
+ x, coords = self._conform_to_input(x, coords, time)
+ y, ycoords = self.model(x, coords)
+ y, ycoords = _squeeze_singletons(y, ycoords)
+ self.publish(y, ycoords, valid_time=as_datetime(time))
+ self.run_count += 1
+
+ def to(self, device: Any) -> "DiagnosticComponent":
+ self.model = self.model.to(device)
+ return self
diff --git a/earth2studio/nvcoupler/config.py b/earth2studio/nvcoupler/config.py
new file mode 100644
index 000000000..0cefbd04f
--- /dev/null
+++ b/earth2studio/nvcoupler/config.py
@@ -0,0 +1,406 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""YAML configuration: serialize/reconstruct a coupled system (NUOPC config analog).
+
+Where NUOPC drivers read run sequences and field dictionaries from config
+files, :func:`to_yaml` / :func:`from_yaml` round-trip an nvcoupler Driver
+through a small YAML schema::
+
+ clock: {start, stop, dt}
+ sequence: | # hand-written run-sequence DSL, verbatim
+ @6h
+ ...
+ dictionary: [...] # only non-default FieldEntry items
+ aliases: {alias: standard_name} # add_alias() additions vs the default
+ components:
+ : {class: , kwargs: {...}}
+ connectors: [{src, dst, fields, time_policy, fill, window, reduce}]
+
+Systems whose sequence was derived from the coupling graph (``Driver``
+built without a sequence) serialize it as ``sequence: {derived: true,
+text: |...}`` — the text is informational; :func:`from_yaml` re-derives the
+sequence from components + connectors, which reproduces it exactly.
+
+Only import-path-constructible components round-trip in v1: the ``class``
+key must name a module-level class or factory callable that rebuilds the
+component from ``kwargs`` alone. Components wrapping closures (a bare
+CallableComponent) cannot be serialized unless they carry a ``yaml_spec``
+attribute — a ``{"class": ..., "kwargs": ...}`` dict declaring how to
+rebuild them. Model checkpoints referenced by load paths (e.g.
+``{load: 'earth2studio.models.px.Persistence'}``) are out of scope for v1.
+"""
+
+import importlib
+import inspect
+import os
+from collections import OrderedDict
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import yaml
+
+from .clock import Clock, as_timedelta, fmt_timedelta
+from .connector import Connector
+from .dictionary import DEFAULT_DICTIONARY, CellMethod, FieldDictionary, FieldEntry
+from .driver import Driver
+from .errors import CouplingError
+from .mediator import AccumulationMediator
+
+if TYPE_CHECKING:
+ from .component import Component
+
+
+class _LiteralDumper(yaml.SafeDumper):
+ """SafeDumper rendering multi-line strings as literal blocks (|)."""
+
+
+def _str_representer(dumper: yaml.Dumper, data: str) -> yaml.Node:
+ style = "|" if "\n" in data else None
+ return dumper.represent_scalar("tag:yaml.org,2002:str", data, style=style)
+
+
+_LiteralDumper.add_representer(str, _str_representer)
+
+
+def _sanitize(value: Any) -> Any:
+ """Coerce numpy scalars/arrays and time types to YAML-safe primitives."""
+ if isinstance(value, np.timedelta64):
+ return fmt_timedelta(value)
+ if isinstance(value, np.datetime64):
+ return str(np.datetime_as_string(value, unit="s"))
+ if isinstance(value, np.generic):
+ return value.item()
+ if isinstance(value, np.ndarray):
+ return [_sanitize(v) for v in value.tolist()]
+ if isinstance(value, dict):
+ return {k: _sanitize(v) for k, v in value.items()}
+ if isinstance(value, (list, tuple, set, frozenset)):
+ return [_sanitize(v) for v in value]
+ return value
+
+
+# ---------------------------------------------------------------------------
+# Serialization (Driver -> YAML)
+# ---------------------------------------------------------------------------
+def _entry_to_dict(entry: FieldEntry, aliases: list[str]) -> dict:
+ d: dict[str, Any] = {
+ "standard_name": entry.standard_name,
+ "canonical_units": entry.canonical_units,
+ "description": entry.description,
+ "aliases": sorted(aliases),
+ }
+ if entry.cell_method is not None:
+ cm = entry.cell_method
+ d["cell_method"] = {
+ "base": cm.base,
+ "method": cm.method,
+ "window": fmt_timedelta(cm.window),
+ }
+ return d
+
+
+def _custom_entries(components: dict[str, "Component"]) -> list[dict]:
+ """FieldEntry items present in any component dictionary but not in (or
+ differing from) the default dictionary."""
+ out: dict[str, dict] = {}
+ for comp in components.values():
+ d = comp.dictionary
+ for std in d.standard_names():
+ entry = d.resolve(std)
+ default = (
+ DEFAULT_DICTIONARY.resolve(std) if std in DEFAULT_DICTIONARY else None
+ )
+ if entry == default:
+ continue
+ aliases = [a for a, s in d._aliases.items() if s == std and a != std]
+ out[std] = _entry_to_dict(entry, aliases)
+ return [out[k] for k in sorted(out)]
+
+
+def _custom_aliases(components: dict[str, "Component"]) -> dict[str, str]:
+ """Alias -> standard-name additions relative to the default dictionary.
+
+ Aliases added via :meth:`FieldDictionary.add_alias` after registration
+ leave the FieldEntry itself equal to the default, so they are invisible
+ to :func:`_custom_entries`; serialize the alias-map delta explicitly.
+ Aliases originating from a component's ``variable_aliases`` kwarg are
+ skipped — they are rebuilt from the component spec itself.
+ """
+ out: dict[str, str] = {}
+ for comp in components.values():
+ raw_to_std = getattr(comp, "_raw_to_std", {})
+ for alias, std in comp.dictionary._aliases.items():
+ if alias == std:
+ continue
+ if DEFAULT_DICTIONARY._aliases.get(alias) == std:
+ continue
+ if raw_to_std.get(alias) == std:
+ continue
+ existing = out.get(alias)
+ if existing is not None and existing != std:
+ raise CouplingError(
+ f"Alias {alias!r} maps to {existing!r} in one component "
+ f"dictionary and to {std!r} in another; make the alias "
+ "consistent across components before serializing to YAML"
+ )
+ out[alias] = std
+ return {a: out[a] for a in sorted(out)}
+
+
+def _component_spec(name: str, comp: "Component") -> dict:
+ spec = getattr(comp, "yaml_spec", None)
+ if spec is not None:
+ if not isinstance(spec, dict) or "class" not in spec:
+ raise CouplingError(
+ f"Component {name!r}: yaml_spec must be a dict with a 'class' "
+ f"import path (and optional 'kwargs'), got {spec!r}"
+ )
+ return {
+ "class": spec["class"],
+ "kwargs": _sanitize(spec.get("kwargs", {})),
+ }
+ if isinstance(comp, AccumulationMediator):
+ cls = type(comp)
+ return {
+ "class": f"{cls.__module__}.{cls.__qualname__}",
+ "kwargs": {
+ "name": comp.name,
+ "fields": list(comp.methods),
+ "window": fmt_timedelta(comp.timestep),
+ },
+ }
+ raise CouplingError(
+ f"Component {name!r} ({type(comp).__name__}) is not serializable: it "
+ "wraps Python state (a closure or model object) that YAML cannot "
+ "reconstruct. Set a yaml_spec attribute on the component — a dict "
+ "{'class': '', 'kwargs': {...}} that "
+ "rebuilds it — or construct this system in Python. Model components "
+ "referenced by load paths are out of scope for YAML round-trips in v1."
+ )
+
+
+def to_yaml(driver: Driver, path: str | os.PathLike | None = None) -> str:
+ """Serialize a Driver to YAML text (optionally also written to `path`).
+
+ Raises
+ ------
+ CouplingError
+ If any component is neither an AccumulationMediator nor carries a
+ ``yaml_spec`` attribute describing how to rebuild it.
+ """
+ doc: dict[str, Any] = OrderedDict()
+ doc["clock"] = {
+ "start": str(np.datetime_as_string(driver.clock.start, unit="s")),
+ "stop": str(np.datetime_as_string(driver.clock.stop, unit="s")),
+ "dt": fmt_timedelta(driver.clock.dt),
+ }
+ doc["sequence"] = (
+ {"derived": True, "text": str(driver.sequence)}
+ if driver.sequence_derived
+ else str(driver.sequence)
+ )
+ entries = _custom_entries(driver.components)
+ if entries:
+ doc["dictionary"] = entries
+ aliases = _custom_aliases(driver.components)
+ if aliases:
+ doc["aliases"] = aliases
+ doc["components"] = {
+ name: _component_spec(name, comp) for name, comp in driver.components.items()
+ }
+ connectors = []
+ for conn in driver._connectors.values():
+ item: dict[str, Any] = {
+ "src": conn.src.name,
+ "dst": conn.dst.name,
+ "time_policy": conn.time_policy,
+ "fill": conn.fill,
+ }
+ if conn._fields is not None:
+ item["fields"] = list(conn._fields)
+ if conn.sample is not None:
+ item["sample"] = conn.sample
+ if conn.window is not None:
+ item["window"] = fmt_timedelta(conn.window)
+ item["reduce"] = conn.reduce
+ connectors.append(item)
+ if connectors:
+ doc["connectors"] = connectors
+ text = yaml.dump(
+ dict(doc), Dumper=_LiteralDumper, sort_keys=False, default_flow_style=False
+ )
+ if path is not None:
+ with open(path, "w") as f:
+ f.write(text)
+ return text
+
+
+# ---------------------------------------------------------------------------
+# Deserialization (YAML -> Driver)
+# ---------------------------------------------------------------------------
+def _resolve_import(path: str) -> Any:
+ module_path, _, attr = path.rpartition(".")
+ if not module_path:
+ raise CouplingError(
+ f"Component class {path!r} is not a dotted import path "
+ "(expected e.g. 'earth2studio.nvcoupler.mediator.TrailingAverageMediator')"
+ )
+ try:
+ module = importlib.import_module(module_path)
+ except ImportError as e:
+ raise CouplingError(
+ f"Cannot import module {module_path!r} for component class {path!r}: {e}"
+ ) from e
+ try:
+ return getattr(module, attr)
+ except AttributeError:
+ raise CouplingError(
+ f"Module {module_path!r} has no attribute {attr!r} "
+ f"(from component class {path!r})"
+ ) from None
+
+
+def _accepts_kwarg(fn: Any, name: str) -> bool:
+ try:
+ params = inspect.signature(fn).parameters
+ except (TypeError, ValueError):
+ return False
+ if name in params:
+ return True
+ return any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
+
+
+def _build_dictionary(items: list[dict]) -> FieldDictionary:
+ dictionary = FieldDictionary(DEFAULT_DICTIONARY)
+ for item in items:
+ cm = None
+ if item.get("cell_method"):
+ raw = item["cell_method"]
+ cm = CellMethod(
+ base=raw["base"],
+ method=raw["method"],
+ window=as_timedelta(raw["window"]),
+ )
+ dictionary.register(
+ FieldEntry(
+ standard_name=item["standard_name"],
+ canonical_units=item["canonical_units"],
+ description=item.get("description", ""),
+ aliases=frozenset(item.get("aliases", ())),
+ cell_method=cm,
+ )
+ )
+ return dictionary
+
+
+def from_yaml(path_or_str: str | os.PathLike) -> Driver:
+ """Build an (uninitialized) Driver from YAML text or a YAML file path.
+
+ Components are reconstructed by importing each ``class`` path and calling
+ it with ``kwargs``; call ``driver.initialize(ics)`` afterwards as usual.
+ """
+ source = str(path_or_str)
+ if "\n" not in source and os.path.exists(source):
+ with open(source) as f:
+ source = f.read()
+ doc = yaml.safe_load(source)
+ if not isinstance(doc, dict):
+ raise CouplingError(
+ "YAML config must be a mapping with 'clock', 'sequence' and "
+ f"'components' keys, got {type(doc).__name__}"
+ )
+ missing = [k for k in ("clock", "sequence", "components") if k not in doc]
+ if missing:
+ raise CouplingError(f"YAML config is missing required keys: {missing}")
+
+ clock_cfg = doc["clock"]
+ clock = Clock(clock_cfg["start"], clock_cfg["stop"], clock_cfg["dt"])
+
+ seq_cfg = doc["sequence"]
+ if isinstance(seq_cfg, dict):
+ if not seq_cfg.get("derived"):
+ raise CouplingError(
+ "YAML 'sequence' must be run-sequence DSL text, or "
+ "{derived: true, text: ...} for a graph-derived sequence; "
+ f"got {seq_cfg!r}"
+ )
+ sequence = None # re-derived from components + connectors
+ else:
+ sequence = seq_cfg
+
+ dictionary = None
+ if doc.get("dictionary"):
+ dictionary = _build_dictionary(doc["dictionary"])
+ if doc.get("aliases"):
+ if dictionary is None:
+ dictionary = FieldDictionary(DEFAULT_DICTIONARY)
+ for alias, std in doc["aliases"].items():
+ dictionary.add_alias(std, alias)
+
+ components: dict[str, Component] = {}
+ for name, spec in doc["components"].items():
+ if not isinstance(spec, dict) or "class" not in spec:
+ raise CouplingError(
+ f"Component {name!r}: expected {{class: , "
+ f"kwargs: {{...}}}}, got {spec!r}"
+ )
+ factory = _resolve_import(spec["class"])
+ kwargs = dict(spec.get("kwargs") or {})
+ if (
+ dictionary is not None
+ and "dictionary" not in kwargs
+ and _accepts_kwarg(factory, "dictionary")
+ ):
+ kwargs["dictionary"] = dictionary
+ try:
+ components[name] = factory(**kwargs)
+ except CouplingError:
+ raise
+ except Exception as e:
+ raise CouplingError(
+ f"Component {name!r}: {spec['class']}(**{spec.get('kwargs', {})}) "
+ f"failed: {e}"
+ ) from e
+
+ connectors: list[Connector] = []
+ for item in doc.get("connectors") or []:
+ src, dst = item["src"], item["dst"]
+ for endpoint in (src, dst):
+ if endpoint not in components:
+ raise CouplingError(
+ f"Connector {src}->{dst}: {endpoint!r} is not a configured "
+ f"component; have {sorted(components)}"
+ )
+ connectors.append(
+ Connector(
+ components[src],
+ components[dst],
+ fields=item.get("fields"),
+ time_policy=item.get("time_policy", "constant"),
+ fill=item.get("fill", "none"),
+ sample=item.get("sample"),
+ window=item.get("window"),
+ reduce=item.get("reduce"),
+ )
+ )
+
+ return Driver(
+ components,
+ sequence,
+ clock,
+ connectors=connectors or None,
+ )
diff --git a/earth2studio/nvcoupler/connector.py b/earth2studio/nvcoupler/connector.py
new file mode 100644
index 000000000..2096b02b7
--- /dev/null
+++ b/earth2studio/nvcoupler/connector.py
@@ -0,0 +1,688 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Connector: moves Fields from one component's exports to another's imports.
+
+The NUOPC_Connector analog. Matching is by standard name; each transfer runs
+a pipeline of (1) time policy, (2) vertical interpolation when source and
+destination vertical coordinates differ, (3) mask fill, (4) spatial regrid
+onto the destination grid. Regridders and mask fillers are built lazily and
+cached per grid signature (the dxwrapper.py pattern). All tensor math is
+torch, so autograd graphs survive the exchange.
+"""
+
+from collections import OrderedDict
+from collections.abc import Callable
+from dataclasses import replace
+from typing import Literal
+
+import numpy as np
+import torch
+from loguru import logger
+
+from earth2studio.utils.interp import latlon_interpolation_regular
+from earth2studio.utils.type import CoordSystem
+
+from .clock import DeltaLike, as_datetime, as_timedelta, fmt_timedelta
+from .component import Component
+from .errors import (
+ CouplingError,
+ IncompatibleFieldError,
+ VerticalMismatchError,
+)
+from .field import _SPATIAL_DIMS, Field
+from .mediator import _RunningReduction
+from .points import PointSet
+from .vertical import HybridLevels, PressureLevels, interp_to_pressure
+
+Regridder = Callable[[torch.Tensor], torch.Tensor]
+
+
+def _is_regular(v: np.ndarray) -> bool:
+ return v.ndim == 1 and len(v) > 1 and np.allclose(np.diff(v), v[1] - v[0])
+
+
+def _build_latlon_regridder(src: CoordSystem, dst: CoordSystem) -> Regridder:
+ """Bilinear lat/lon regridder using earth2studio's regular-grid kernel.
+
+ Requires 1D equally-spaced source lat/lon (the common case for global
+ models); index clamping at the grid edge stands in for extrapolation.
+ """
+ src_lat, src_lon = np.asarray(src["lat"]), np.asarray(src["lon"])
+ if not (_is_regular(src_lat) and _is_regular(src_lon)):
+ raise IncompatibleFieldError(
+ "Auto regrid requires a regular 1D source lat/lon grid; pass a "
+ "custom regridder=... to the Connector for curvilinear or "
+ "unstructured source grids"
+ )
+ flip_lat = src_lat[0] > src_lat[-1]
+ lat0 = torch.as_tensor(src_lat[::-1].copy() if flip_lat else src_lat)
+ lon0 = torch.as_tensor(src_lon)
+ lat1g, lon1g = np.meshgrid(
+ np.asarray(dst["lat"]), np.asarray(dst["lon"]), indexing="ij"
+ )
+ lat1 = torch.as_tensor(lat1g)
+ lon1 = torch.as_tensor(lon1g)
+
+ def regrid(data: torch.Tensor) -> torch.Tensor:
+ if flip_lat:
+ data = torch.flip(data, dims=(-2,))
+ return latlon_interpolation_regular(
+ data,
+ lat0.to(device=data.device, dtype=data.dtype),
+ lon0.to(device=data.device, dtype=data.dtype),
+ lat1.to(device=data.device, dtype=data.dtype),
+ lon1.to(device=data.device, dtype=data.dtype),
+ )
+
+ return regrid
+
+
+def _build_point_sampler(
+ src: CoordSystem, points: PointSet, method: Literal["nearest", "bilinear"]
+) -> Regridder:
+ """Grid-to-point sampler: a regular lat/lon field -> PointSet locations.
+
+ "bilinear" reuses the same regular-grid kernel as the mesh regridder,
+ just evaluated at scattered (lat, lon) pairs instead of a destination
+ mesh (an [N, 1] "mesh" degenerates to per-point bilinear interpolation).
+ "nearest" is a great-circle nearest-neighbor gather, matching the mask
+ filler's KDTree approach.
+ """
+ src_lat, src_lon = np.asarray(src["lat"]), np.asarray(src["lon"])
+ if not (_is_regular(src_lat) and _is_regular(src_lon)):
+ raise IncompatibleFieldError(
+ "Auto sample requires a regular 1D source lat/lon grid; pass a "
+ "custom regridder=... for curvilinear or unstructured source "
+ "grids"
+ )
+ if method == "bilinear":
+ flip_lat = src_lat[0] > src_lat[-1]
+ lat0 = torch.as_tensor(src_lat[::-1].copy() if flip_lat else src_lat)
+ lon0 = torch.as_tensor(src_lon)
+ # [N] point locations reshaped as an [N, 1] "mesh" so the existing
+ # regular-grid kernel evaluates one interpolated value per point.
+ lat1 = torch.as_tensor(points.lat).unsqueeze(-1)
+ lon1 = torch.as_tensor(points.lon).unsqueeze(-1)
+
+ def sample(data: torch.Tensor) -> torch.Tensor:
+ if flip_lat:
+ data = torch.flip(data, dims=(-2,))
+ out = latlon_interpolation_regular(
+ data,
+ lat0.to(device=data.device, dtype=data.dtype),
+ lon0.to(device=data.device, dtype=data.dtype),
+ lat1.to(device=data.device, dtype=data.dtype),
+ lon1.to(device=data.device, dtype=data.dtype),
+ )
+ return out.squeeze(-1)
+
+ return sample
+
+ if method != "nearest":
+ raise CouplingError(
+ f"Unsupported sample={method!r}; choose 'nearest' or 'bilinear'"
+ )
+
+ from scipy.spatial import cKDTree
+
+ lat2d, lon2d = np.meshgrid(src_lat, src_lon, indexing="ij")
+ phi, lam = np.deg2rad(lat2d).ravel(), np.deg2rad(lon2d).ravel()
+ src_xyz = np.stack(
+ [np.cos(phi) * np.cos(lam), np.cos(phi) * np.sin(lam), np.sin(phi)], axis=1
+ )
+ tree = cKDTree(src_xyz)
+ dst_phi, dst_lam = np.deg2rad(points.lat), np.deg2rad(points.lon)
+ dst_xyz = np.stack(
+ [
+ np.cos(dst_phi) * np.cos(dst_lam),
+ np.cos(dst_phi) * np.sin(dst_lam),
+ np.sin(dst_phi),
+ ],
+ axis=1,
+ )
+ _, nearest_flat = tree.query(dst_xyz, k=1)
+ index = torch.as_tensor(nearest_flat, dtype=torch.long)
+
+ def sample(data: torch.Tensor) -> torch.Tensor:
+ flat = data.reshape(*data.shape[:-2], -1)
+ return torch.index_select(flat, -1, index.to(data.device))
+
+ return sample
+
+
+def _build_mask_filler(coords: CoordSystem, mask: torch.Tensor) -> Regridder:
+ """Nearest-valid fill on the source grid: every invalid point takes the
+ value of its nearest valid neighbor (great-circle metric via unit-sphere
+ KDTree). Differentiable (pure gather)."""
+ from scipy.spatial import cKDTree
+
+ lat = np.asarray(coords["lat"], dtype=np.float64)
+ lon = np.asarray(coords["lon"], dtype=np.float64)
+ lat2d, lon2d = np.meshgrid(lat, lon, indexing="ij")
+ phi, lam = np.deg2rad(lat2d).ravel(), np.deg2rad(lon2d).ravel()
+ xyz = np.stack(
+ [np.cos(phi) * np.cos(lam), np.cos(phi) * np.sin(lam), np.sin(phi)], axis=1
+ )
+ valid = mask.reshape(-1).cpu().numpy().astype(bool)
+ if not valid.any():
+ raise IncompatibleFieldError("Mask fill impossible: no valid source points")
+ tree = cKDTree(xyz[valid])
+ _, nearest = tree.query(xyz, k=1)
+ valid_index = np.flatnonzero(valid)[nearest]
+ index = torch.as_tensor(valid_index, dtype=torch.long)
+
+ def fill(data: torch.Tensor) -> torch.Tensor:
+ flat = data.reshape(*data.shape[:-2], -1)
+ filled = torch.index_select(flat, -1, index.to(data.device))
+ return filled.reshape(data.shape)
+
+ return fill
+
+
+class Connector:
+ """Moves matched Fields src.exports -> dst.imports each time it executes.
+
+ Parameters
+ ----------
+ src, dst : Component
+ fields : list[str], optional
+ Standard names to transfer; defaults to the intersection of src's
+ advertised exports and dst's advertised imports.
+ time_policy : "constant" | "linear"
+ What the destination sees between source updates: hold the latest
+ export (constant, the PhysicsNeMo ConstantCoupler behavior), or
+ linearly extrapolate from the two most recent exports.
+ fill : "none" | "zero" | "nearest"
+ Treatment of masked (invalid) source points before regridding.
+ regridder : callable, optional
+ Override the spatial regrid for all fields of this connector
+ (signature: tensor[..., H, W] -> tensor[..., H', W']). Required when
+ the grids differ and the auto path cannot handle them (HEALPix
+ 'face' dims, curvilinear grids); identical grids — including
+ identical face grids — pass through as identity without one.
+ sample : "nearest" | "bilinear", optional
+ Grid-to-point sampling: set when the destination is a scattered
+ sample-location target (``dst.points`` is a :class:`.points.PointSet`
+ — stations, sites, arbitrary query coordinates) rather than a mesh.
+ "bilinear" reuses the mesh regridder's kernel per point; "nearest"
+ is a great-circle nearest-neighbor lookup. Mutually exclusive with
+ `regridder=`; a point-target destination with neither set raises at
+ `execute()` time rather than guessing.
+ window, reduce : optional
+ Set both to make this a *windowed* connector: each execute folds the
+ source fields into a running reduction ("mean" | "sum" | "max" |
+ "min"), and delivery happens only at execute times aligned to
+ `window`. The delivered Field carries the DERIVED standard name — the
+ destination must import a dictionary entry whose CellMethod is
+ (base=source export, method=`reduce`, window=`window`). Between
+ window boundaries the destination's previous import is untouched.
+ The window origin is the valid_time of the first execute's source
+ field (the clock start under lagged coupling, where the connector
+ runs before the source in its slot), so no driver hook is needed.
+ This replaces a single-source AccumulationMediator; `time_policy`
+ does not apply on the windowed path.
+ """
+
+ def __init__(
+ self,
+ src: Component,
+ dst: Component,
+ fields: list[str] | None = None,
+ time_policy: Literal["constant", "linear"] = "constant",
+ fill: Literal["none", "zero", "nearest"] = "none",
+ regridder: Regridder | None = None,
+ sample: Literal["nearest", "bilinear"] | None = None,
+ window: DeltaLike | None = None,
+ reduce: Literal["mean", "sum", "max", "min"] | None = None,
+ ):
+ self.src = src
+ self.dst = dst
+ self.time_policy = time_policy
+ self.fill = fill
+ if sample is not None and regridder is not None:
+ raise CouplingError(
+ f"Connector {src.name}->{dst.name}: sample= and regridder= "
+ "are mutually exclusive — pass one or the other"
+ )
+ if sample is not None and sample not in ("nearest", "bilinear"):
+ raise CouplingError(
+ f"Connector {src.name}->{dst.name}: unsupported sample="
+ f"{sample!r}; choose 'nearest' or 'bilinear'"
+ )
+ self.sample = sample
+ if (window is None) != (reduce is None):
+ raise CouplingError(
+ f"Connector {src.name}->{dst.name}: window= and reduce= must "
+ "be set together — a windowed reduction needs both the window "
+ "length and the reduction method"
+ )
+ if reduce is not None and reduce not in ("mean", "sum", "max", "min"):
+ raise CouplingError(
+ f"Connector {src.name}->{dst.name}: unsupported reduce="
+ f"{reduce!r}; choose 'mean', 'sum', 'max' or 'min'"
+ )
+ self.window = as_timedelta(window) if window is not None else None
+ self.reduce = reduce
+ self._reduction = _RunningReduction()
+ self._derived: dict[str, str] = {} # src export name -> derived dst name
+ self._origin: np.datetime64 | None = None # window alignment origin
+ self._user_regridder = regridder
+ self._fields = list(fields) if fields is not None else None
+ self._matched: list[str] | None = None
+ self._regridders: dict[tuple, Regridder] = {}
+ self._fillers: dict[tuple, Regridder] = {}
+ self._samplers: dict[tuple, Regridder] = {}
+ # 2-deep export history per field: (previous, latest), rotated only
+ # when a genuinely new export (different valid_time) arrives
+ self._history: dict[str, tuple[Field | None, Field]] = {}
+ self._linear_warned: set[str] = set()
+ self.last_transfer: dict[str, Field] = {}
+
+ @property
+ def name(self) -> str:
+ return f"{self.src.name}->{self.dst.name}"
+
+ # -- matching --------------------------------------------------------------
+ def match(self) -> list[str]:
+ """Resolve matched standard names (cached).
+
+ Plain connectors match by name intersection. Windowed connectors
+ match each source export against a destination import whose
+ CellMethod derives from it; the returned list then contains both the
+ consumed source names and the delivered derived names, so driver-side
+ bookkeeping (fed imports, consumed exports) sees the full mapping.
+ """
+ if self._matched is not None:
+ return self._matched
+ _, src_exports = self.src.advertise()
+ dst_imports, _ = self.dst.advertise()
+ if self.window is not None:
+ return self._match_windowed(src_exports, dst_imports)
+ if self._fields is not None:
+ missing = [
+ f for f in self._fields if f not in src_exports or f not in dst_imports
+ ]
+ if missing:
+ raise IncompatibleFieldError(
+ f"Connector {self.name}: fields {missing} are not in both "
+ f"{self.src.name!r} exports ({src_exports}) and "
+ f"{self.dst.name!r} imports ({dst_imports})"
+ )
+ self._matched = list(self._fields)
+ else:
+ self._matched = [n for n in dst_imports if n in src_exports]
+ if not self._matched:
+ raise IncompatibleFieldError(
+ f"Connector {self.name}: no fields match — {self.src.name!r} "
+ f"exports {src_exports}, {self.dst.name!r} imports {dst_imports}"
+ )
+ # units validation against the (shared) dictionary
+ for n in self._matched:
+ entry_src = self.src.dictionary.resolve(n)
+ entry_dst = self.dst.dictionary.resolve(n)
+ self.dst.dictionary.check_units(
+ n, entry_src.canonical_units, src=self.src.name, dst=self.dst.name
+ )
+ del entry_dst
+ return self._matched
+
+ def _match_windowed(
+ self, src_exports: list[str], dst_imports: list[str]
+ ) -> list[str]:
+ """Pair source exports with the destination's derived imports.
+
+ A source export `base` maps to a destination import whose dictionary
+ entry carries CellMethod(base, self.reduce, self.window); the field
+ is delivered under that derived name. No matching derived import is
+ an error — the coupler never invents names.
+ """
+ wanted = self._fields if self._fields is not None else src_exports
+ for name in dst_imports:
+ cm = self.dst.dictionary.resolve(name).cell_method
+ if (
+ cm is not None
+ and cm.method == self.reduce
+ and as_timedelta(cm.window) == self.window
+ and cm.base in wanted
+ and cm.base in src_exports
+ ):
+ self._derived[cm.base] = name
+ unmatched = [f for f in wanted if f not in self._derived]
+ if not self._derived or (self._fields is not None and unmatched):
+ w, r = fmt_timedelta(self.window), self.reduce
+ raise CouplingError(
+ f"Connector {self.name}: window={w!r}/reduce={r!r} is set but "
+ f"{self.dst.name!r} imports no derived field for "
+ f"{unmatched or src_exports} — register a "
+ f"FieldEntry(cell_method=CellMethod(base, {r!r}, window={w!r})) "
+ f"in the destination's dictionary and add its standard name to "
+ f"{self.dst.name!r}'s imports (destination imports: "
+ f"{dst_imports})"
+ )
+ for base, derived in self._derived.items():
+ self.dst.dictionary.check_units(
+ derived,
+ self.src.dictionary.resolve(base).canonical_units,
+ src=self.src.name,
+ dst=self.dst.name,
+ )
+ self._matched = list(self._derived) + list(self._derived.values())
+ return self._matched
+
+ # -- pipeline stages ---------------------------------------------------------
+ def _apply_time_policy(self, field: Field, time: np.datetime64) -> Field:
+ prev, latest = self._history.get(field.standard_name, (None, None))
+ # Rotate the (prev, latest) history only when the incoming export is
+ # genuinely new (different valid_time); re-seeing the same export on
+ # subsequent executes must not collapse the extrapolation baseline.
+ is_new = (
+ latest is None
+ or field.valid_time is None
+ or latest.valid_time is None
+ or as_datetime(field.valid_time) != as_datetime(latest.valid_time)
+ )
+ if is_new:
+ prev = latest
+ self._history[field.standard_name] = (prev, field)
+ if self.time_policy == "constant" or prev is None:
+ return field
+ if "lead_time" in field.coords or "window" in field.coords:
+ # a lead-time-resolved field carries many valid times; a single
+ # valid_time extrapolation is ill-defined for it
+ if field.standard_name not in self._linear_warned:
+ self._linear_warned.add(field.standard_name)
+ logger.warning(
+ "Connector {}: time_policy='linear' is undefined for "
+ "field {!r} with a lead_time/window dimension — falling "
+ "back to 'constant' for it",
+ self.name,
+ field.standard_name,
+ )
+ return field
+ if prev.valid_time is None or field.valid_time is None:
+ return field
+ dt_hist = (
+ (as_datetime(field.valid_time) - as_datetime(prev.valid_time))
+ .astype("timedelta64[ns]")
+ .astype(np.int64)
+ )
+ if dt_hist <= 0:
+ return field
+ dt_ahead = (
+ (as_datetime(time) - as_datetime(field.valid_time))
+ .astype("timedelta64[ns]")
+ .astype(np.int64)
+ )
+ if dt_ahead == 0:
+ return field
+ w = dt_ahead / dt_hist
+ data = field.data + (field.data - prev.data) * w
+ return replace(field, data=data, valid_time=as_datetime(time))
+
+ def _apply_vertical(self, field: Field) -> Field:
+ want = self.dst.import_vertical.get(field.standard_name)
+ if want is None:
+ return field
+ have = field.vertical
+ if have == want:
+ return field
+ if have is None:
+ raise VerticalMismatchError(
+ f"Connector {self.name}: {self.dst.name!r} expects "
+ f"{field.standard_name!r} on {want}, but the source field has "
+ "no vertical coordinate"
+ )
+ if not isinstance(want, PressureLevels):
+ raise VerticalMismatchError(
+ f"Connector {self.name}: only interpolation onto PressureLevels "
+ f"is supported in v1 (destination wants {type(want).__name__})"
+ )
+ ps = None
+ if isinstance(have, HybridLevels):
+ ps_std = self.src.dictionary.standard_name(have.ps_field)
+ if ps_std not in self.src.export_state:
+ raise VerticalMismatchError(
+ f"Connector {self.name}: hybrid->pressure interpolation of "
+ f"{field.standard_name!r} needs {ps_std!r} in "
+ f"{self.src.name!r} exports — add it to the source's "
+ "export list"
+ )
+ ps = self.src.export_state[ps_std].data
+ data, coords = interp_to_pressure(field.data, field.coords, have, want, ps)
+ return replace(field, data=data, coords=coords, vertical=want)
+
+ def _apply_fill(self, field: Field) -> Field:
+ if field.mask is None or self.fill == "none":
+ return field
+ if self.fill == "zero":
+ data = torch.where(field.mask.to(field.data.device), field.data, 0.0)
+ return replace(field, data=data, mask=None)
+ key = (field.grid_signature(), field.mask.cpu().numpy().tobytes())
+ if key not in self._fillers:
+ self._fillers[key] = _build_mask_filler(field.coords, field.mask)
+ return replace(field, data=self._fillers[key](field.data), mask=None)
+
+ def _apply_regrid(self, field: Field) -> Field:
+ dst_grid = self.dst.grid_coords()
+ if dst_grid is None:
+ return field # destination has no grid of its own (e.g. mediator)
+ src_spatial = OrderedDict(
+ (k, v) for k, v in field.coords.items() if k in _SPATIAL_DIMS
+ )
+ # Identity fast path: every spatial dim of the field (lat/lon, but
+ # also HEALPix-style face/height/width) exists in the destination
+ # grid with an equal coordinate array — nothing to regrid.
+ same = bool(src_spatial) and all(
+ k in dst_grid and np.array_equal(np.asarray(v), np.asarray(dst_grid[k]))
+ for k, v in src_spatial.items()
+ )
+ if same and self._user_regridder is None:
+ return field
+ if "point" in dst_grid:
+ return self._apply_sample(field, src_spatial, dst_grid)
+ if ("face" in field.coords or "face" in dst_grid) and (
+ self._user_regridder is None
+ ):
+ raise IncompatibleFieldError(
+ f"Connector {self.name}: source and destination HEALPix "
+ "'face' grids differ — pass a custom regridder= (e.g. built "
+ "with earth2grid, see models/px/dlesym.py)"
+ )
+ if self._user_regridder is not None:
+ # Explicit override: apply the user regridder to the trailing
+ # spatial dims of ANY layout (lat/lon, HEALPix face/height/width,
+ # curvilinear y/x) and rebuild coords from the destination grid.
+ data = self._user_regridder(field.data)
+ coords = OrderedDict(
+ (k, v) for k, v in field.coords.items() if k not in _SPATIAL_DIMS
+ )
+ for k, v in dst_grid.items():
+ coords[k] = np.asarray(v).copy()
+ return replace(field, data=data, coords=coords)
+ key = (field.grid_signature(),)
+ if not (
+ "lat" in src_spatial
+ and "lon" in src_spatial
+ and "lat" in dst_grid
+ and "lon" in dst_grid
+ ):
+ raise IncompatibleFieldError(
+ f"Connector {self.name}: auto regrid needs lat/lon on both "
+ f"grids (source dims {list(src_spatial)}, destination dims "
+ f"{list(dst_grid)}) — pass a custom regridder="
+ )
+ if key not in self._regridders:
+ self._regridders[key] = _build_latlon_regridder(src_spatial, dst_grid)
+ regrid = self._regridders[key]
+ # regrid acts on the trailing two (lat, lon) dims
+ spatial_last = list(field.coords)[-2:] == ["lat", "lon"]
+ if not spatial_last:
+ raise IncompatibleFieldError(
+ f"Connector {self.name}: field {field.standard_name!r} must "
+ f"have (lat, lon) as trailing dims, got {list(field.coords)}"
+ )
+ data = regrid(field.data)
+ # preserve original dim order: everything except lat/lon, then dst grid
+ coords = OrderedDict(
+ (k, v) for k, v in field.coords.items() if k not in ("lat", "lon")
+ )
+ coords["lat"] = np.asarray(dst_grid["lat"]).copy()
+ coords["lon"] = np.asarray(dst_grid["lon"]).copy()
+ return replace(field, data=data, coords=coords)
+
+ def _apply_sample(
+ self, field: Field, src_spatial: CoordSystem, dst_grid: CoordSystem
+ ) -> Field:
+ """Grid-to-point delivery: destination advertises a "point" dim.
+
+ Mirrors `_apply_regrid`'s user-override / auto-build split, but the
+ destination's actual (lat, lon) locations live on `self.dst.points`
+ (`dst_grid["point"]` is only the dim's own labels, same as any other
+ CoordSystem entry — it does not carry coordinates by itself).
+ """
+ if self._user_regridder is not None:
+ data = self._user_regridder(field.data)
+ coords = OrderedDict(
+ (k, v) for k, v in field.coords.items() if k not in _SPATIAL_DIMS
+ )
+ coords["point"] = np.asarray(dst_grid["point"]).copy()
+ return replace(field, data=data, coords=coords)
+ if self.sample is None:
+ raise CouplingError(
+ f"Connector {self.name}: destination {self.dst.name!r} is a "
+ "point target (a scattered sample-location grid) but this "
+ "connector has neither sample= nor regridder= set — pass "
+ "sample='nearest' or sample='bilinear', or a custom "
+ "regridder= for non-lat/lon sources"
+ )
+ points: PointSet | None = self.dst.points
+ if points is None:
+ raise CouplingError(
+ f"Connector {self.name}: destination {self.dst.name!r} "
+ "advertises a 'point' dim but has no points= location "
+ "metadata set — construct it with points=PointSet(lat=..., "
+ "lon=...)"
+ )
+ if not ("lat" in src_spatial and "lon" in src_spatial):
+ raise IncompatibleFieldError(
+ f"Connector {self.name}: auto sample needs lat/lon on the "
+ f"source grid (source dims {list(src_spatial)}) — pass a "
+ "custom regridder= for non-lat/lon sources"
+ )
+ spatial_last = list(field.coords)[-2:] == ["lat", "lon"]
+ if not spatial_last:
+ raise IncompatibleFieldError(
+ f"Connector {self.name}: field {field.standard_name!r} must "
+ f"have (lat, lon) as trailing dims, got {list(field.coords)}"
+ )
+ key = (field.grid_signature(), points.signature(), self.sample)
+ if key not in self._samplers:
+ self._samplers[key] = _build_point_sampler(src_spatial, points, self.sample)
+ data = self._samplers[key](field.data)
+ coords = OrderedDict(
+ (k, v) for k, v in field.coords.items() if k not in ("lat", "lon")
+ )
+ coords["point"] = points.labels().copy()
+ return replace(field, data=data, coords=coords)
+
+ # -- execution ----------------------------------------------------------------
+ def execute(self, time: np.datetime64) -> None:
+ self.match()
+ if self.window is not None:
+ self._execute_windowed(as_datetime(time))
+ return
+ for name in self._matched:
+ field = self._apply_time_policy(self._source_field(name), time)
+ self._deliver(field)
+
+ def _source_field(self, name: str) -> Field:
+ if name not in self.src.export_state:
+ raise CouplingError(
+ f"Connector {self.name}: {self.src.name!r} has not produced "
+ f"{name!r} yet — check the run sequence ordering"
+ )
+ return self.src.export_state[name]
+
+ def _deliver(self, field: Field) -> None:
+ """Run the spatial pipeline and hand the field to the destination."""
+ field = self._apply_vertical(field)
+ field = self._apply_fill(field)
+ field = self._apply_regrid(field)
+ self.dst.import_state.add(field)
+ self.last_transfer[field.standard_name] = field
+ logger.debug(
+ "exchange {}: {} (valid {})",
+ self.name,
+ field.standard_name,
+ field.valid_time,
+ )
+
+ def _execute_windowed(self, time: np.datetime64) -> None:
+ """Fold sources into the running reduction; deliver on window boundaries.
+
+ The origin for boundary alignment is the valid_time of the first
+ execute's source field — under lagged coupling (connector before the
+ source's RunAction in the slot) that is the clock start, so the first
+ delivery lands exactly one window after t0 with no driver hook.
+ """
+ for base in self._derived:
+ field = self._source_field(base)
+ if self._origin is None:
+ self._origin = (
+ as_datetime(field.valid_time)
+ if field.valid_time is not None
+ else time
+ )
+ self._reduction.add(base, field, self.reduce)
+ elapsed_ns = (time - self._origin).astype("timedelta64[ns]").astype(np.int64)
+ if elapsed_ns <= 0 or elapsed_ns % self.window.astype(np.int64) != 0:
+ return # mid-window: accumulate only, previous import stands
+ for base, derived in self._derived.items():
+ data, coords = self._reduction.emit(base, self.reduce)
+ entry = self.dst.dictionary.resolve(derived)
+ self._deliver(
+ Field(
+ data=data,
+ coords=coords,
+ standard_name=derived,
+ units=entry.canonical_units,
+ valid_time=time,
+ source=self.src.name,
+ )
+ )
+ self._reduction.reset()
+
+ def reset(self) -> None:
+ """Clear per-run exchange state (history, running reduction, probes)."""
+ self._history.clear()
+ self.last_transfer.clear()
+ self._reduction.reset()
+ self._origin = None
+
+ def __repr__(self) -> str:
+ fields = self._matched or self._fields or "auto"
+ windowed = (
+ f", window={fmt_timedelta(self.window)!r}, reduce={self.reduce!r}"
+ if self.window is not None
+ else ""
+ )
+ sampled = f", sample={self.sample!r}" if self.sample is not None else ""
+ return (
+ f"Connector({self.name}, fields={fields}, "
+ f"time_policy={self.time_policy!r}, fill={self.fill!r}"
+ f"{sampled}{windowed})"
+ )
diff --git a/earth2studio/nvcoupler/dictionary.py b/earth2studio/nvcoupler/dictionary.py
new file mode 100644
index 000000000..57c36ddb2
--- /dev/null
+++ b/earth2studio/nvcoupler/dictionary.py
@@ -0,0 +1,257 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""CF-style field dictionary: canonical standard names, units, and aliases.
+
+The NUOPC field-dictionary analog. Connectors match exported to imported
+fields by *standard name*, never by raw model variable strings; aliases map
+model vocabularies (``z1000``, ``ws10m``) onto standard names. Derived
+fields (e.g. a 48-hour mean) are first-class entries carrying a
+:class:`CellMethod`, which lets ``couple()`` synthesize the right mediator
+instead of string-parsing suffixes.
+"""
+
+from dataclasses import dataclass, field
+from typing import Literal
+
+import numpy as np
+
+from .errors import UnitsMismatchError, UnknownFieldError
+
+# Small normalization table so cosmetically-different unit strings compare
+# equal. v1 checks equality only; it never converts values.
+#
+# Kept behaviorally aligned with the lexicon normalizer in
+# earth2studio/lexicon/earthmover.py (normalize_units), which lowercases and
+# strips '**', '^', and spaces before matching — any two spellings that
+# normalizer treats as equal (e.g. 'm s**-1', 'm s^-1', 'M/S'; '(0-1)',
+# 'fraction', 'dimensionless'; 'degC', 'celsius'; '%', 'percent') compare
+# equal here too. Not imported from the lexicon so nvcoupler stays importable
+# standalone. Keys are in collapsed form (lowercase, no '**'/'^'/spaces).
+_UNIT_SYNONYMS = {
+ "m2/s2": "m2 s-2",
+ "m2s-2": "m2 s-2",
+ "m/s": "m s-1",
+ "ms-1": "m s-1",
+ "kelvin": "K",
+ "k": "K",
+ "pa": "Pa",
+ "hpa": "hPa",
+ "kg/m2": "kg m-2",
+ "kgm-2": "kg m-2",
+ "mm": "kg m-2", # precipitation depth-equivalence
+ "1": "",
+ "(0-1)": "",
+ "0-1": "",
+ "fraction": "",
+ "dimensionless": "",
+ "degree_celsius": "degC",
+ "degreec": "degC",
+ "degc": "degC",
+ "celsius": "degC",
+ "degreescelsius": "degC",
+ "%": "percent",
+ "percent": "percent",
+ "mofwaterequivalent": "m",
+ "mwe": "m",
+}
+
+
+def normalize_units(units: str) -> str:
+ """Normalize a unit string for comparison (no value conversion).
+
+ Collapses case, '**'/'^' exponent markers, and whitespace exactly like
+ ``earth2studio.lexicon.earthmover.normalize_units`` before applying the
+ synonym table, so the two normalizers agree on overlapping inputs.
+ """
+ u = units.strip().lower().replace("**", "").replace("^", "").replace(" ", "")
+ return _UNIT_SYNONYMS.get(u, u)
+
+
+@dataclass(frozen=True)
+class CellMethod:
+ """CF-style cell method describing a derived (time-reduced) field.
+
+ A field entry with a cell method declares "I am `method` of `base` over
+ `window`" — e.g. the 48 h mean of geopotential_at_1000hpa. This is the
+ machine-readable convention that lets auto-wiring insert the right
+ AccumulationMediator between components of different cadence.
+ """
+
+ base: str
+ method: Literal["mean", "sum", "max", "min"]
+ window: np.timedelta64
+
+ def __post_init__(self) -> None:
+ if self.method not in ("mean", "sum", "max", "min"):
+ raise ValueError(f"Unsupported cell method {self.method!r}")
+
+
+@dataclass(frozen=True)
+class FieldEntry:
+ """One entry in the field dictionary."""
+
+ standard_name: str
+ canonical_units: str
+ description: str = ""
+ aliases: frozenset[str] = field(default_factory=frozenset)
+ cell_method: CellMethod | None = None
+
+ def __post_init__(self) -> None:
+ if not isinstance(self.aliases, frozenset):
+ object.__setattr__(self, "aliases", frozenset(self.aliases))
+
+
+class FieldDictionary:
+ """Registry resolving standard names and aliases to :class:`FieldEntry`.
+
+ Lookup is case-sensitive on standard names and aliases. An alias may map
+ to exactly one standard name; a standard name may have many aliases.
+ """
+
+ def __init__(self, entries: "FieldDictionary | list[FieldEntry] | None" = None):
+ self._entries: dict[str, FieldEntry] = {}
+ self._aliases: dict[str, str] = {}
+ if isinstance(entries, FieldDictionary):
+ self._entries = dict(entries._entries)
+ self._aliases = dict(entries._aliases)
+ elif entries:
+ for entry in entries:
+ self.register(entry)
+
+ def register(self, entry: FieldEntry) -> None:
+ """Register an entry; re-registering a standard name replaces it."""
+ if entry.standard_name in self._aliases:
+ raise ValueError(
+ f"{entry.standard_name!r} is already an alias for "
+ f"{self._aliases[entry.standard_name]!r}"
+ )
+ self._entries[entry.standard_name] = entry
+ for alias in entry.aliases:
+ self.add_alias(entry.standard_name, alias)
+
+ def add_alias(self, standard_name: str, alias: str) -> None:
+ if standard_name not in self._entries:
+ raise UnknownFieldError(standard_name, self._entries.keys())
+ existing = self._aliases.get(alias)
+ if existing is not None and existing != standard_name:
+ raise ValueError(
+ f"Alias {alias!r} already maps to {existing!r}; cannot remap "
+ f"to {standard_name!r}"
+ )
+ if alias in self._entries and alias != standard_name:
+ raise ValueError(f"Alias {alias!r} collides with a standard name")
+ self._aliases[alias] = standard_name
+
+ def resolve(self, name: str) -> FieldEntry:
+ """Resolve a standard name or alias to its entry."""
+ if name in self._entries:
+ return self._entries[name]
+ if name in self._aliases:
+ return self._entries[self._aliases[name]]
+ raise UnknownFieldError(
+ name, list(self._entries.keys()) + list(self._aliases.keys())
+ )
+
+ def standard_name(self, name: str) -> str:
+ return self.resolve(name).standard_name
+
+ def __contains__(self, name: str) -> bool:
+ return name in self._entries or name in self._aliases
+
+ def standard_names(self) -> list[str]:
+ return list(self._entries.keys())
+
+ def check_units(
+ self, standard_name: str, units: str, *, src: str, dst: str
+ ) -> None:
+ """Raise :class:`UnitsMismatchError` if units disagree with canonical."""
+ canonical = self.resolve(standard_name).canonical_units
+ if normalize_units(units) != normalize_units(canonical):
+ raise UnitsMismatchError(standard_name, src, units, dst, canonical)
+
+ def derived_from(self, standard_name: str) -> CellMethod | None:
+ """Return the cell method if `standard_name` is a derived field."""
+ return self.resolve(standard_name).cell_method
+
+
+def _default_entries() -> list[FieldEntry]:
+ """Curated v1 vocabulary covering the earth2studio surface variables and
+ pressure-level fields used by the coupled models in this repo, plus the
+ accumulation/impact fields the mediators produce."""
+ e = FieldEntry
+ hours = lambda h: np.timedelta64(h, "h") # noqa: E731
+ return [
+ # surface / single-level
+ e("sea_surface_temperature", "K", "SST", frozenset({"sst"})),
+ e("air_temperature_2m", "K", "2 m air temperature", frozenset({"t2m"})),
+ e("dewpoint_temperature_2m", "K", "2 m dewpoint", frozenset({"d2m"})),
+ e("wind_speed_10m", "m s-1", "10 m wind speed", frozenset({"ws10m", "ws10"})),
+ e("eastward_wind_10m", "m s-1", "10 m u-wind", frozenset({"u10m", "u10"})),
+ e("northward_wind_10m", "m s-1", "10 m v-wind", frozenset({"v10m", "v10"})),
+ e("surface_pressure", "Pa", "surface pressure", frozenset({"sp"})),
+ e("mean_sea_level_pressure", "Pa", "MSLP", frozenset({"msl", "mslp"})),
+ e("total_column_water_vapour", "kg m-2", "TCWV", frozenset({"tcwv"})),
+ e(
+ "total_precipitation_6h",
+ "kg m-2",
+ "6 h accumulated precip",
+ frozenset({"tp06"}),
+ ),
+ # pressure-level (levels encoded in the name, earth2studio convention)
+ e("geopotential_at_1000hpa", "m2 s-2", "z at 1000 hPa", frozenset({"z1000"})),
+ e("geopotential_at_500hpa", "m2 s-2", "z at 500 hPa", frozenset({"z500"})),
+ e("geopotential_at_250hpa", "m2 s-2", "z at 250 hPa", frozenset({"z250"})),
+ e("air_temperature_at_850hpa", "K", "t at 850 hPa", frozenset({"t850"})),
+ e(
+ "geopotential_thickness_300_700hpa",
+ "m2 s-2",
+ "z300 - z700 thickness",
+ frozenset({"tau300-700"}),
+ ),
+ # derived / windowed fields (CellMethod-carrying entries)
+ e(
+ "geopotential_at_1000hpa_48h_mean",
+ "m2 s-2",
+ "trailing 48 h mean of z1000",
+ frozenset({"z1000-48H"}),
+ CellMethod("geopotential_at_1000hpa", "mean", hours(48)),
+ ),
+ e(
+ "wind_speed_10m_48h_mean",
+ "m s-1",
+ "trailing 48 h mean of 10 m wind speed",
+ frozenset({"ws10-48H"}),
+ CellMethod("wind_speed_10m", "mean", hours(48)),
+ ),
+ e(
+ "total_precipitation_48h_sum",
+ "kg m-2",
+ "48 h accumulated precipitation",
+ frozenset(),
+ CellMethod("total_precipitation_6h", "sum", hours(48)),
+ ),
+ e(
+ "air_temperature_2m_24h_max",
+ "K",
+ "24 h maximum 2 m temperature",
+ frozenset(),
+ CellMethod("air_temperature_2m", "max", hours(24)),
+ ),
+ ]
+
+
+DEFAULT_DICTIONARY = FieldDictionary(_default_entries())
diff --git a/earth2studio/nvcoupler/dlesym_split.py b/earth2studio/nvcoupler/dlesym_split.py
new file mode 100644
index 000000000..8c6c7bf62
--- /dev/null
+++ b/earth2studio/nvcoupler/dlesym_split.py
@@ -0,0 +1,648 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""DLESyM split adapter: one coupled model, two nvcoupler components.
+
+The NUOPC "split a monolithic executable into gridded components" move.
+:class:`earth2studio.models.px.DLESyM` internally couples an atmosphere and
+an ocean HEALPix U-Net inside a single ``__call__``; :func:`split_dlesym`
+re-exposes those sub-models as two :class:`~earth2studio.nvcoupler.component.Component`
+instances so the exchange (SST down, windowed z1000/ws10m up) runs through
+explicit Connectors on the coupler's clock instead of being hidden inside
+``DLESyM._forward``.
+
+Both components step the full parent cadence (96 h, one native ``__call__``)
+and close over the parent DLESyM instance for normalization (``center`` /
+``scale`` buffers), insolation (``_make_insolation_tensor``), the coupling
+tensor construction (``_make_atmos_coupling`` / ``_make_ocean_coupling`` are
+called directly, not re-implemented), constants and the precomputed
+lead-time / variable index tables:
+
+- **Atmos** runs first. It consumes the ``sea_surface_temperature`` import
+ with DLESyM's persisted-at-t0 semantics (``_make_atmos_coupling``: the
+ lead-0 SST repeated for every internal sub-step), produces the 16
+ 6-hourly atmos outputs, and exports its prognostic fields plus the
+ ocean-coupling variables already chunk-averaged into the two 48 h ocean
+ windows (the exact ``_make_ocean_coupling`` math) as derived
+ window-mean fields.
+- **Ocean** consumes those window means, rebuilds the DLESyM ocean coupling
+ tensor (lead, batch, window-major variables, face, height, width), runs
+ the ocean model to produce SST at 48 h / 96 h, and exports the 96 h SST —
+ which the run sequence feeds back to the atmosphere *lagged* across steps,
+ exactly like ``_next_step_inputs`` carrying the ocean output into the next
+ window.
+
+Exchange happens on the shared HEALPix (face, height, width) grid, so the
+connectors are identity transfers.
+
+All exchange-path tensor math is pure torch (normalize / chunk-mean /
+re-normalize round-trips through physical units are exact up to float
+rounding), so autograd survives the split.
+
+Honesty note
+------------
+Execution against real DLESyM weights is **untested** — this module has only
+been exercised against structural mocks. A weights-equivalence test (needs
+the ``dlesym`` optional dependencies plus the NGC package) would assert, for
+an n-step rollout from the same initial condition:
+
+1. the atmos component's prognostic exports at each 96 h ring equal the
+ final-lead slice of ``DLESyM.retrieve_valid_atmos_outputs`` from the
+ native iterator;
+2. the ocean component's SST export equals the 96 h slice of
+ ``DLESyM.retrieve_valid_ocean_outputs``;
+3. the coupling tensor delivered to ``ocean_model`` equals
+ ``DLESyM._make_ocean_coupling(atmos_outputs)`` bit-for-bit modulo the
+ denormalize/renormalize round trip;
+4. the lagged SST the atmosphere sees at step k equals the native step k-1
+ ocean output at 96 h.
+"""
+
+from collections import OrderedDict
+from collections.abc import Mapping
+from typing import Any
+
+import numpy as np
+import torch
+
+from earth2studio.utils.type import CoordSystem
+
+from .clock import as_datetime, as_timedelta
+from .component import Component, _broadcast_to_slice
+from .dictionary import DEFAULT_DICTIONARY, CellMethod, FieldDictionary, FieldEntry
+from .driver import Driver
+from .errors import CouplingError
+from .field import Field, State
+from .sequence import ConnectAction, RunAction, RunSequence, Slot
+
+# Module-level dictionary copy: the extension point for DLESyM vocabulary.
+# The standard DLESyM-V1-ERA5 variables (z1000, ws10m, sst, tau300-700, ...)
+# and the 48 h window-mean entries (geopotential_at_1000hpa_48h_mean,
+# wind_speed_10m_48h_mean) are already in DEFAULT_DICTIONARY; split_dlesym
+# auto-registers anything a non-default DLESyM config adds.
+DLESYM_DICTIONARY = FieldDictionary(DEFAULT_DICTIONARY)
+
+_HOUR_NS = 3_600_000_000_000
+
+
+class _DLESyMSplitComponent(Component):
+ """Shared machinery for the two halves of a split DLESyM."""
+
+ def __init__(
+ self,
+ name: str,
+ parent: Any,
+ dictionary: FieldDictionary,
+ imports: list[str],
+ exports: list[str],
+ ):
+ step = as_timedelta(parent.atmos_output_times[-1])
+ super().__init__(name, step, imports, exports, dictionary=dictionary)
+ self.parent = parent
+ self._x: torch.Tensor | None = None
+ self._coords: CoordSystem | None = None
+ self._times: np.ndarray | None = None
+ self._batch: int = 1
+
+ # -- shared helpers ------------------------------------------------------
+ def _validate_ic(self, x: torch.Tensor, coords: CoordSystem) -> None:
+ dims = list(coords)
+ expected = ["batch", "time", "lead_time", "variable", "face", "height", "width"]
+ if x.ndim != 7 or dims != expected:
+ raise CouplingError(
+ f"Component {self.name!r} expects the DLESyM input layout "
+ f"{expected}, got dims {dims} for tensor of shape "
+ f"{tuple(x.shape)}"
+ )
+ lead = coords["lead_time"]
+ rel = lead - lead[-1]
+ want = self.parent.full_input_times - self.parent.full_input_times[-1]
+ if len(rel) != len(want) or not np.array_equal(
+ rel.astype("timedelta64[ns]"), want.astype("timedelta64[ns]")
+ ):
+ raise CouplingError(
+ f"Component {self.name!r}: initial condition lead_time window "
+ f"{list(lead)} does not match DLESyM full_input_times "
+ f"{list(self.parent.full_input_times)}"
+ )
+
+ def _anchor_times(self, time: np.datetime64) -> np.ndarray:
+ """Absolute window-end times for insolation, one per (batch, time)
+ element, matching DLESyM's anchor + lead_time[-1] arithmetic."""
+ if self.clock is None or self._times is None:
+ raise CouplingError(f"Component {self.name!r} not realized/initialized")
+ offset = (as_datetime(time) - self.clock.start) - self.timestep
+ anchor = self._times + offset
+ if self._batch > 1:
+ anchor = np.concatenate([anchor] * self._batch)
+ return anchor
+
+ def _publish_from_tensor(
+ self,
+ x: torch.Tensor,
+ variables: np.ndarray,
+ names: list[str],
+ valid_time: np.datetime64 | None,
+ ) -> None:
+ """Publish instantaneous fields from a (batch, time, variable, face,
+ height, width) tensor onto export_state."""
+ if self._coords is None:
+ raise CouplingError(f"Component {self.name!r} not initialized")
+ coords: CoordSystem = OrderedDict(
+ (k, v) for k, v in self._coords.items() if k != "lead_time"
+ )
+ coords["variable"] = np.asarray(variables)
+ coords.move_to_end("variable", last=False)
+ coords.move_to_end("time", last=False)
+ coords.move_to_end("batch", last=False)
+ state = State.from_tensor(
+ f"{self.name}.publish",
+ x,
+ coords,
+ self.dictionary,
+ valid_time=valid_time,
+ source=self.name,
+ strict=False,
+ )
+ for std in names:
+ if std not in state:
+ raise CouplingError(
+ f"Component {self.name!r} advertises export {std!r} but the "
+ f"model output variables are {list(variables)}"
+ )
+ self.export_state.add(state[std])
+
+ def _center_scale(
+ self, var_idx: list[int], shape: tuple[int, ...]
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Per-variable center/scale views aligned to a var axis of `shape`.
+
+ Mirrors the normalization constants of ``DLESyM._normalize_input`` /
+ ``_denormalize_output`` restricted to a variable subset; the parent
+ methods cannot be called directly because they broadcast over the
+ FULL variable dimension while the split components carry tensors
+ holding only their own (atmos or ocean) variables.
+ """
+ c = self.parent.center.reshape(-1)[var_idx].view(shape)
+ s = self.parent.scale.reshape(-1)[var_idx].view(shape)
+ return c, s
+
+ def run(self, time: np.datetime64) -> None: # pragma: no cover - abstract-ish
+ raise NotImplementedError
+
+
+class DLESyMAtmosComponent(_DLESyMSplitComponent):
+ """The atmosphere half of a split DLESyM.
+
+ Imports ``sea_surface_temperature`` (persisted at the window end, the
+ ``_make_atmos_coupling`` semantics), exports its prognostic fields at the
+ 96 h window end plus the ocean-coupling variables chunk-averaged into the
+ ocean output windows (carrying a leading ``window`` dimension of length
+ ``len(ocean_output_times)``).
+ """
+
+ def __init__(
+ self,
+ parent: Any,
+ dictionary: FieldDictionary,
+ derived_names: Mapping[str, str],
+ name: str = "atmos",
+ ):
+ d = FieldDictionary(dictionary)
+ imports = [d.standard_name(v) for v in parent.atmos_coupling_variables]
+ self._prognostic_exports = [d.standard_name(v) for v in parent.atmos_variables]
+ # raw ocean-coupling variable -> derived window-mean standard name
+ self._derived = dict(derived_names)
+ exports = self._prognostic_exports + list(self._derived.values())
+ super().__init__(name, parent, d, imports, exports)
+ # window-lead index of "now" (lead 0) in full_input_times
+ self._zero_idx = parent.atmos_input_lt_idx[-1]
+ # next-step window: index of (t + step) in atmos_output_times per
+ # full_input_times entry (DLESyM's _next_step_inputs slice)
+ out_t = parent.atmos_output_times.astype("timedelta64[ns]")
+ step = as_timedelta(parent.atmos_output_times[-1])
+ self._next_idx = []
+ for t in parent.full_input_times.astype("timedelta64[ns]"):
+ hits = np.flatnonzero(out_t == t + step)
+ if hits.size == 0:
+ raise CouplingError(
+ f"DLESyM atmos window time {t} + {step} is not an atmos "
+ f"output time {list(parent.atmos_output_times)} — the "
+ "input window cannot be rebuilt from one step's outputs"
+ )
+ self._next_idx.append(int(hits[0]))
+ # position of each full-variable index within the atmos output var dim
+ self._atmos_pos = {vi: k for k, vi in enumerate(parent.atmos_var_idx)}
+
+ def initialize(self, x: torch.Tensor, coords: CoordSystem) -> None:
+ self._validate_ic(x, coords)
+ self._x = x
+ self._coords = OrderedDict(coords)
+ self._times = np.asarray(coords["time"], dtype="datetime64[ns]")
+ self._batch = x.shape[0]
+ start = self.clock.start if self.clock is not None else None
+ inst = x[:, :, self._zero_idx]
+ self._publish_from_tensor(
+ inst, coords["variable"], self._prognostic_exports, start
+ )
+
+ def _inject_imports(self, x: torch.Tensor) -> torch.Tensor:
+ """Overwrite the coupling variables at the persisted (lead-0)
+ coupling indices with imported fields, non-destructively."""
+ p = self.parent
+ updates = []
+ for raw, vi in zip(p.atmos_coupling_variables, p.atmos_coupling_var_idx):
+ std = self.dictionary.standard_name(raw)
+ if std in self.import_state:
+ updates.append((vi, self.import_state[std]))
+ if not updates:
+ return x
+ x = x.clone()
+ for lead_idx in sorted(set(p.atmos_coupled_input_lt_idx)):
+ for vi, field in updates:
+ slot = x[:, :, lead_idx, vi]
+ data = _broadcast_to_slice(
+ field.data.to(device=x.device, dtype=x.dtype), slot.shape
+ )
+ x[:, :, lead_idx, vi] = data
+ return x
+
+ def run(self, time: np.datetime64) -> None:
+ if self._x is None or self._coords is None:
+ raise CouplingError(f"Component {self.name!r} not initialized")
+ ic_coords = self._coords
+ p = self.parent
+ x = self._inject_imports(self._x)
+ b, t = x.shape[0], x.shape[1]
+
+ # Mirrors DLESyM._normalize_input; not called directly because the
+ # parent method assumes its center/scale buffers already live on the
+ # input's device/dtype, while the split component moves them here.
+ xn = (x - p.center.to(device=x.device, dtype=x.dtype)) / p.scale.to(
+ device=x.device, dtype=x.dtype
+ )
+ xf = xn.reshape(-1, *xn.shape[2:]) # (B, lead, var, face, h, w)
+
+ atmos_state = xf[:, p.atmos_input_lt_idx][
+ ..., p.atmos_var_idx, :, :, :
+ ].permute(0, 3, 1, 2, 4, 5)
+ insolation = p._make_insolation_tensor(
+ anchor_times=self._anchor_times(time), timedeltas=p.atmos_sol_times
+ )
+ # the parent's persisted-at-t0 coupling selection, verbatim
+ coupling = p._make_atmos_coupling(xf, ic_coords)
+ inputs = [
+ y.to(device=xf.device, dtype=xf.dtype)
+ for y in [atmos_state, insolation, p.atmos_constants, coupling]
+ ]
+ out = p.atmos_model(inputs) # (B, face, n_lead, n_atmos_var, h, w), normalized
+
+ # -- derived exports: the parent's own chunk-mean math, denormalized --
+ # _make_ocean_coupling returns (lead=1, B, window-major variables,
+ # face, h, w): window w, coupling var k lives at index w * C + k.
+ n_windows = len(p.ocean_output_times)
+ n_coupling = len(p.ocean_coupling_var_idx)
+ mc = p._make_ocean_coupling(out, ic_coords)[0]
+ window_coords: CoordSystem = OrderedDict(
+ {
+ "batch": np.arange(out.shape[0]),
+ "window": np.arange(n_windows),
+ "face": np.asarray(ic_coords["face"]).copy(),
+ "height": np.asarray(ic_coords["height"]).copy(),
+ "width": np.asarray(ic_coords["width"]).copy(),
+ }
+ )
+ for k, (raw, vi) in enumerate(
+ zip(p.ocean_coupling_variables, p.ocean_coupling_var_idx)
+ ):
+ c, s = self._center_scale([vi], (1,))
+ mean_k = torch.stack(
+ [mc[:, w * n_coupling + k] for w in range(n_windows)], dim=1
+ ) # (B, window, face, h, w)
+ data = mean_k * s.to(device=out.device, dtype=out.dtype) + c.to(
+ device=out.device, dtype=out.dtype
+ )
+ derived = self._derived[raw]
+ entry = self.dictionary.resolve(derived)
+ self.export_state.add(
+ Field(
+ data=data,
+ coords=OrderedDict(
+ (k2, v.copy()) for k2, v in window_coords.items()
+ ),
+ standard_name=derived,
+ units=entry.canonical_units,
+ valid_time=as_datetime(time),
+ source=self.name,
+ )
+ )
+
+ # -- denormalize outputs and rebuild the sliding input window --------
+ # mirrors DLESyM._denormalize_output restricted to atmos_var_idx
+ # (the atmos output tensor lacks the ocean variables)
+ c_a, s_a = self._center_scale(
+ p.atmos_var_idx, (1, 1, 1, len(p.atmos_var_idx), 1, 1)
+ )
+ out_phys = out * s_a.to(device=out.device, dtype=out.dtype) + c_a.to(
+ device=out.device, dtype=out.dtype
+ )
+ out_bt = out_phys.permute(0, 2, 3, 1, 4, 5).reshape(
+ b, t, out.shape[2], out.shape[3], out.shape[1], *out.shape[-2:]
+ ) # (b, t, lead, var_atmos, face, h, w)
+ window = out_bt[:, :, self._next_idx] # (b, t, n_window_lead, A, f, h, w)
+ n_lead = window.shape[2]
+ pieces = []
+ for vi in range(x.shape[3]):
+ if vi in self._atmos_pos:
+ pieces.append(window[:, :, :, self._atmos_pos[vi]])
+ else:
+ # non-atmos variables (sst): persist the current lead-0 value;
+ # it is refreshed from the SST import before the next step
+ carry = x[:, :, self._zero_idx, vi].unsqueeze(2)
+ pieces.append(carry.expand(b, t, n_lead, *carry.shape[-3:]))
+ self._x = torch.stack(pieces, dim=3)
+
+ # -- prognostic exports at the window end (96 h) -----------------------
+ inst = out_bt[:, :, -1] # (b, t, var_atmos, face, h, w)
+ self._publish_from_tensor(
+ inst,
+ np.array(p.atmos_variables),
+ self._prognostic_exports,
+ as_datetime(time),
+ )
+ self.run_count += 1
+
+ @property
+ def state(self) -> tuple[torch.Tensor, CoordSystem]:
+ return self._x, self._coords
+
+
+class DLESyMOceanComponent(_DLESyMSplitComponent):
+ """The ocean half of a split DLESyM.
+
+ Imports the atmosphere's window-mean coupling fields, rebuilds the
+ DLESyM ocean coupling tensor (lead=1, batch, window-major variables,
+ face, height, width), and exports ``sea_surface_temperature`` valid at
+ the 96 h window end.
+ """
+
+ def __init__(
+ self,
+ parent: Any,
+ dictionary: FieldDictionary,
+ derived_names: Mapping[str, str],
+ name: str = "ocean",
+ ):
+ d = FieldDictionary(dictionary)
+ self._derived = dict(derived_names)
+ imports = [self._derived[v] for v in parent.ocean_coupling_variables]
+ self._sst_exports = [d.standard_name(v) for v in parent.ocean_variables]
+ super().__init__(name, parent, d, imports, self._sst_exports)
+ step = as_timedelta(parent.atmos_output_times[-1])
+ out_t = parent.ocean_output_times.astype("timedelta64[ns]")
+ self._next_idx = []
+ for t in parent.ocean_input_times.astype("timedelta64[ns]"):
+ hits = np.flatnonzero(out_t == t + step)
+ if hits.size == 0:
+ raise CouplingError(
+ f"DLESyM ocean window time {t} + {step} is not an ocean "
+ f"output time {list(parent.ocean_output_times)} — the "
+ "input window cannot be rebuilt from one step's outputs"
+ )
+ self._next_idx.append(int(hits[0]))
+
+ def initialize(self, x: torch.Tensor, coords: CoordSystem) -> None:
+ self._validate_ic(x, coords)
+ p = self.parent
+ window = x[:, :, p.ocean_input_lt_idx][:, :, :, p.ocean_var_idx]
+ self._x = window # (b, t, n_ocean_lead, n_ocean_var, face, h, w), physical
+ self._coords = OrderedDict(coords)
+ self._times = np.asarray(coords["time"], dtype="datetime64[ns]")
+ self._batch = x.shape[0]
+ start = self.clock.start if self.clock is not None else None
+ inst = window[:, :, -1]
+ self._publish_from_tensor(
+ inst, np.array(p.ocean_variables), self._sst_exports, start
+ )
+
+ def _build_coupling(self, ref: torch.Tensor) -> torch.Tensor:
+ """Reassemble _make_ocean_coupling's tensor from imported window-mean
+ fields: (lead=1, batch, window-major variables, face, h, w).
+
+ Mirrors ``DLESyM._make_ocean_coupling``'s window-major layout but
+ cannot call it: the parent method chunk-averages the atmos model's
+ normalized output tensor, which never crosses the coupling seam —
+ here only the already-averaged, physical-unit import Fields exist,
+ so the tensor is rebuilt (renormalized) from those instead.
+ """
+ p = self.parent
+ n_windows = len(p.ocean_output_times)
+ per_var: list[torch.Tensor] = []
+ for raw, vi in zip(p.ocean_coupling_variables, p.ocean_coupling_var_idx):
+ derived = self._derived[raw]
+ if derived not in self.import_state:
+ raise CouplingError(
+ f"Component {self.name!r} needs import {derived!r} before it "
+ "can run — schedule the atmos component and the "
+ "atmos -> ocean connector earlier in the same slot"
+ )
+ field = self.import_state[derived]
+ data = field.data.to(device=ref.device, dtype=ref.dtype)
+ if data.ndim != 5 or data.shape[1] != n_windows:
+ raise CouplingError(
+ f"Component {self.name!r}: import {derived!r} must have "
+ f"shape (batch, window={n_windows}, face, height, width), "
+ f"got {tuple(data.shape)}"
+ )
+ c, s = self._center_scale([vi], (1,))
+ per_var.append(
+ (data - c.to(device=ref.device, dtype=ref.dtype))
+ / s.to(device=ref.device, dtype=ref.dtype)
+ )
+ blocks = [
+ torch.stack([v[:, w] for v in per_var], dim=1) for w in range(n_windows)
+ ] # each (B, C, face, h, w)
+ return torch.cat(blocks, dim=1).unsqueeze(0)
+
+ def run(self, time: np.datetime64) -> None:
+ if self._x is None:
+ raise CouplingError(f"Component {self.name!r} not initialized")
+ p = self.parent
+ window = self._x
+ b, t = window.shape[0], window.shape[1]
+
+ # mirrors DLESyM._normalize_input restricted to ocean_var_idx (the
+ # ocean window tensor lacks the atmos variables)
+ c_o, s_o = self._center_scale(
+ p.ocean_var_idx, (1, 1, 1, len(p.ocean_var_idx), 1, 1, 1)
+ )
+ xn = (window - c_o.to(device=window.device, dtype=window.dtype)) / s_o.to(
+ device=window.device, dtype=window.dtype
+ )
+ xf = xn.reshape(-1, *xn.shape[2:]) # (B, lead, var, face, h, w)
+ ocean_state = xf.permute(0, 3, 1, 2, 4, 5)
+ insolation = p._make_insolation_tensor(
+ anchor_times=self._anchor_times(time), timedeltas=p.ocean_sol_times
+ )
+ coupling = self._build_coupling(xf)
+ inputs = [
+ y.to(device=xf.device, dtype=xf.dtype)
+ for y in [ocean_state, insolation, p.ocean_constants, coupling]
+ ]
+ out = p.ocean_model(inputs) # (B, face, n_ocean_lead, n_ocean_var, h, w)
+
+ # mirrors DLESyM._denormalize_output restricted to ocean_var_idx
+ c2, s2 = self._center_scale(
+ p.ocean_var_idx, (1, 1, 1, len(p.ocean_var_idx), 1, 1)
+ )
+ out_phys = out * s2.to(device=out.device, dtype=out.dtype) + c2.to(
+ device=out.device, dtype=out.dtype
+ )
+ out_bt = out_phys.permute(0, 2, 3, 1, 4, 5).reshape(
+ b, t, out.shape[2], out.shape[3], out.shape[1], *out.shape[-2:]
+ ) # (b, t, lead, var, face, h, w)
+ self._x = out_bt[:, :, self._next_idx]
+
+ inst = out_bt[:, :, -1] # SST at the 96 h window end
+ self._publish_from_tensor(
+ inst, np.array(p.ocean_variables), self._sst_exports, as_datetime(time)
+ )
+ self.run_count += 1
+
+ @property
+ def state(self) -> tuple[torch.Tensor, CoordSystem]:
+ return self._x, self._coords
+
+
+def split_dlesym(
+ dlesym: Any, dictionary: FieldDictionary | None = None
+) -> tuple[DLESyMAtmosComponent, DLESyMOceanComponent]:
+ """Expose a DLESyM's atmos/ocean sub-models as two nvcoupler components.
+
+ Parameters
+ ----------
+ dlesym : DLESyM
+ A constructed :class:`earth2studio.models.px.DLESyM` (or any object
+ exposing the same attributes: ``atmos_model`` / ``ocean_model``,
+ ``center`` / ``scale`` / ``*_constants`` buffers, the ``*_variables``
+ and ``*_input_times`` / ``*_output_times`` config, the precomputed
+ ``*_lt_idx`` / ``*_var_idx`` index tables, ``_make_insolation_tensor``
+ and the ``_make_atmos_coupling`` / ``_make_ocean_coupling`` methods).
+ dictionary : FieldDictionary, optional
+ Vocabulary to extend; defaults to a copy of
+ :data:`DLESYM_DICTIONARY`. Unknown raw variables and missing
+ window-mean entries are auto-registered on a private copy.
+
+ Returns
+ -------
+ tuple[DLESyMAtmosComponent, DLESyMOceanComponent]
+ """
+ d = FieldDictionary(dictionary or DLESYM_DICTIONARY)
+ for raw in list(dlesym.atmos_variables) + list(dlesym.ocean_variables):
+ if raw not in d:
+ d.register(FieldEntry(raw, "", f"DLESyM variable {raw!r}"))
+
+ n_windows = len(dlesym.ocean_output_times)
+ n_lead = len(dlesym.atmos_output_times)
+ if n_windows == 0 or n_lead % n_windows != 0:
+ raise CouplingError(
+ f"DLESyM atmos output times ({n_lead}) do not chunk evenly into "
+ f"{n_windows} ocean windows — cannot replicate _make_ocean_coupling"
+ )
+ step = as_timedelta(dlesym.atmos_output_times[-1])
+ window_ns = step.astype(np.int64) // n_windows
+ if window_ns % _HOUR_NS != 0:
+ raise CouplingError(
+ f"DLESyM ocean coupling window {window_ns} ns is not a whole "
+ "number of hours; cannot name the derived window-mean fields"
+ )
+ window_h = int(window_ns // _HOUR_NS)
+ window = np.timedelta64(window_h, "h")
+
+ derived: dict[str, str] = {}
+ for raw in dlesym.ocean_coupling_variables:
+ std = d.standard_name(raw)
+ name = f"{std}_{window_h}h_mean"
+ if name not in d:
+ base = d.resolve(std)
+ d.register(
+ FieldEntry(
+ name,
+ base.canonical_units,
+ f"trailing {window_h} h mean of {std} (DLESyM ocean coupling)",
+ frozenset(),
+ CellMethod(std, "mean", window),
+ )
+ )
+ derived[raw] = name
+
+ atmos = DLESyMAtmosComponent(dlesym, d, derived)
+ ocean = DLESyMOceanComponent(dlesym, d, derived)
+ return atmos, ocean
+
+
+def build_dlesym_driver(
+ dlesym: Any,
+ start: Any,
+ stop: Any,
+ dictionary: FieldDictionary | None = None,
+ collect: bool = True,
+) -> Driver:
+ """Wire a split DLESyM into a Driver matching the native internal loop.
+
+ The run sequence reproduces ``DLESyM._forward`` + ``_next_step_inputs``
+ ordering: within a step the atmosphere runs first and its window means
+ flow to the ocean (sequential coupling); the ocean's SST flows to the
+ atmosphere *before* it runs, i.e. lagged across steps::
+
+ @96h
+ ocean -> atmos # lagged SST (previous step's 96 h output)
+ atmos
+ atmos -> ocean # window-mean coupling, same step
+ ocean
+ @
+
+ Initialize with the same DLESyM-layout initial condition for both halves::
+
+ driver = build_dlesym_driver(model, "2024-01-01", "2024-01-09")
+ driver.initialize({"atmos": (x, coords), "ocean": (x, coords)})
+ driver.run()
+
+ ``stop - start`` must be a multiple of the 96 h step.
+ """
+ atmos, ocean = split_dlesym(dlesym, dictionary)
+ from .clock import Clock
+
+ clock = Clock(start, stop, dt=atmos.timestep)
+ sequence = RunSequence(
+ [
+ Slot(
+ atmos.timestep,
+ [
+ ConnectAction(ocean.name, atmos.name),
+ RunAction(atmos.name),
+ ConnectAction(atmos.name, ocean.name),
+ RunAction(ocean.name),
+ ],
+ )
+ ]
+ )
+ return Driver(
+ {atmos.name: atmos, ocean.name: ocean},
+ sequence,
+ clock,
+ collect=collect,
+ )
diff --git a/earth2studio/nvcoupler/docs/api_reference.md b/earth2studio/nvcoupler/docs/api_reference.md
new file mode 100644
index 000000000..7c110206a
--- /dev/null
+++ b/earth2studio/nvcoupler/docs/api_reference.md
@@ -0,0 +1,883 @@
+# API reference
+
+Lookup page for everything `earth2studio.nvcoupler` exports (`import
+earth2studio.nvcoupler as nvc`). Signatures are copied from the source;
+behavior notes are one to three lines — for usage and rationale see
+[user_guide.md](user_guide.md) and [concepts.md](concepts.md). Common
+argument types: `TimeLike = str | np.datetime64`,
+`DeltaLike = str | np.timedelta64` (string forms in
+[dsl_and_yaml_reference.md](dsl_and_yaml_reference.md#interval-strings-as_timedelta)),
+`CoordSystem` = earth2studio's ordered coordinate dict.
+
+## field
+
+### `Field` (dataclass)
+
+```python
+Field(
+ data: torch.Tensor,
+ coords: CoordSystem,
+ standard_name: str,
+ units: str,
+ valid_time: np.datetime64 | None = None,
+ source: str | None = None,
+ mask: torch.Tensor | None = None,
+ vertical: VerticalCoordinate | None = None,
+)
+```
+
+One exchanged quantity: a torch tensor plus its coordinates and canonical
+identity. Raises `CouplingError` at construction if `coords` contains a
+`"variable"` dimension (a Field is one variable) or if `data.ndim` disagrees
+with `len(coords)`. `mask` is boolean, True = valid. Methods:
+`to(device) -> Field`, `clone() -> Field`, `grid_signature() -> tuple`
+(hashable spatial-grid key used for regridder caching).
+
+### `State`
+
+```python
+State(name: str, fields: Iterable[Field] = ())
+```
+
+A `MutableMapping[str, Field]` keyed by standard name; every component owns an
+import and an export State. `state[key] = field` enforces
+`key == field.standard_name` (`CouplingError` otherwise); missing keys raise a
+`KeyError` listing the present fields.
+
+- `add(field: Field, replace: bool = True) -> None` — `replace=False` raises
+ `CouplingError` on duplicates.
+- `subset(names: Iterable[str]) -> State`
+- `to(device: Any) -> State`
+- `as_tensor(names: list[str] | None = None) -> tuple[torch.Tensor, CoordSystem]`
+ — stacks fields along a new `"variable"` axis inserted before the first
+ spatial dim; all selected fields must share identical coords. Default order
+ is `sorted(fields)`.
+- `State.from_tensor(name, x, coords, dictionary, valid_time=None, source=None, strict=True) -> State`
+ (classmethod) — splits a multi-variable tensor into Fields, resolving raw
+ variable names through the dictionary; unknown names raise
+ `UnknownFieldError` unless `strict=False` (then they are skipped).
+
+## dictionary
+
+### `CellMethod` (frozen dataclass)
+
+```python
+CellMethod(base: str, method: Literal["mean", "sum", "max", "min"], window: np.timedelta64)
+```
+
+Machine-readable "I am `method` of `base` over `window`" tag for derived
+fields; drives windowed `Connector`s, `AccumulationMediator`, and `couple()`'s
+windowed-connector synthesis. Unsupported methods raise `ValueError`.
+
+### `FieldEntry` (frozen dataclass)
+
+```python
+FieldEntry(
+ standard_name: str,
+ canonical_units: str,
+ description: str = "",
+ aliases: frozenset[str] = frozenset(),
+ cell_method: CellMethod | None = None,
+)
+```
+
+One dictionary entry. `aliases` accepts any iterable (coerced to frozenset).
+
+### `FieldDictionary`
+
+```python
+FieldDictionary(entries: FieldDictionary | list[FieldEntry] | None = None)
+```
+
+Registry resolving standard names and aliases to entries; constructing from
+another `FieldDictionary` copies it (the standard way to extend the default).
+
+- `register(entry: FieldEntry) -> None` — re-registering a standard name
+ replaces it; a name that is already an alias raises `ValueError`.
+- `add_alias(standard_name: str, alias: str) -> None` — raises
+ `UnknownFieldError` for unknown standard names, `ValueError` when the alias
+ is taken or collides with a standard name.
+- `resolve(name: str) -> FieldEntry` — raises `UnknownFieldError` (with
+ did-you-mean suggestions) for unknown names.
+- `standard_name(name: str) -> str`, `standard_names() -> list[str]`,
+ `__contains__(name: str) -> bool`
+- `check_units(standard_name: str, units: str, *, src: str, dst: str) -> None`
+ — raises `UnitsMismatchError` when normalized units disagree with canonical
+ (checked, not converted).
+- `derived_from(standard_name: str) -> CellMethod | None`
+
+### `DEFAULT_DICTIONARY`
+
+Module-level `FieldDictionary` with the curated v1 vocabulary (earth2studio
+surface/pressure-level variables plus the built-in derived window fields).
+
+## clock
+
+### `Clock`
+
+```python
+Clock(start: TimeLike, stop: TimeLike, dt: DeltaLike)
+```
+
+Driver clock stepping `start → stop` inclusive at `dt`. Raises `ValueError`
+for non-positive `dt` or `stop <= start`, and `CadenceError` when the span is
+not a multiple of `dt`. Iterating yields each time **after** `start` (the IC
+lives at `start`, step 0). Properties `current`, `step_index`, `n_steps`;
+methods `elapsed()`, `done()`, `advance()` (raises `StopIteration` past
+`stop`), `reset()`, `times()` (all `n_steps + 1` times including `start`).
+
+The module also provides (not re-exported at package level, import from
+`earth2studio.nvcoupler.clock`): `as_datetime(t)`, `as_timedelta(d)`,
+`fmt_timedelta(d)`, `is_multiple(interval, dt)`.
+
+## component
+
+### `Exchange` (frozen dataclass) — one step's coupling bundle
+
+```python
+Exchange(
+ x: torch.Tensor,
+ coords: CoordSystem,
+ imports: State,
+ std_to_raw: Mapping[str, str] = {},
+ time: np.datetime64 | None = None,
+)
+```
+
+Everything an `ImportAdapter` needs for one coupled model step: the
+component's current state tensor and coords, the import `State` (already
+subset to the fields actually delivered), the standard-name →
+raw-model-variable map, and the valid time the step advances to. Two
+accessors cover the common delivery shapes, both pure torch and
+autograd-safe:
+
+- `inject() -> torch.Tensor` — returns `x` with each imported field
+ overwritten into its matching variable slice (clone + `index_copy_`,
+ autograd-intact). Requires a `"variable"` dim in the state coords and that
+ every imported field is a state variable (`CouplingError` otherwise).
+- `stacked(field_order=None, *, who="Exchange.stacked") -> tuple[torch.Tensor, CoordSystem]`
+ — stacks the imports into one tensor with a leading `"variable"` dim. With
+ more than one import, `field_order=` is mandatory (`CouplingError`).
+
+### `ImportAdapter` (runtime-checkable `Protocol`) — the model-invocation contract
+
+```python
+class ImportAdapter(Protocol):
+ def __call__(
+ self, model: Any, exchange: Exchange
+ ) -> tuple[torch.Tensor, CoordSystem]: ...
+```
+
+An adapter **owns the model call**: it receives the model/step-fn and the
+step's `Exchange` and must return the stepped `(x, coords)`. Implementations
+must be autograd-safe (no in-place mutation of tensors that may carry grad).
+Any callable matching this signature can be passed as `import_adapter=`.
+
+### `VariableOverwriteAdapter`
+
+```python
+VariableOverwriteAdapter()
+```
+
+Default adapter: `model(exchange.inject(), exchange.coords)` — overwrite
+matching variable slices, then step (see `Exchange.inject` for the
+requirements and errors).
+
+### `ConditioningKwargAdapter`
+
+```python
+ConditioningKwargAdapter(field_order: list[str] | None = None, method: str = "call_with_conditioning")
+```
+
+Stacks imports (`exchange.stacked(field_order)`) and calls
+`model.(x, coords, conditioning=..., conditioning_coords=...)`
+(StormScope pattern). With more than one import, `field_order=` is mandatory
+— channel order cannot be inferred (`CouplingError`).
+
+### `ExtraTensorAdapter`
+
+```python
+ExtraTensorAdapter(field_order: list[str] | None = None, kwarg: str | None = None)
+```
+
+Stacks imports and calls `model(x, coords, coupling)` positionally, or
+`model(x, coords, **{kwarg: coupling})` when `kwarg` is given (DLESyM /
+PhysicsNeMo 4-tensor pattern). Same `field_order` rule as above.
+
+### `PullAdapter` (from `pull.py` — see the [pull section](#pull) below)
+
+The fourth built-in adapter, for models that fetch forcing internally
+(StormCast-style) instead of accepting it as an argument.
+
+### `Component` (abstract base)
+
+```python
+Component(
+ name: str,
+ timestep: DeltaLike,
+ imports: Iterable[str] = (),
+ exports: Iterable[str] = (),
+ dictionary: FieldDictionary | None = None,
+ variable_aliases: Mapping[str, str] | None = None,
+ export_masks: Mapping[str, torch.Tensor] | None = None,
+ import_vertical: Mapping[str, Any] | None = None,
+ export_vertical: Mapping[str, Any] | None = None,
+ points: PointSet | None = None,
+)
+```
+
+NUOPC_Model analog. `imports`/`exports` may be aliases (resolved to standard
+names via the dictionary); `variable_aliases` maps raw model variable names to
+standard names and registers them as aliases. Class attribute
+`requires_ic: bool = True` (subclasses that can initialize without an
+`(x, coords)` pair set it False). `points` declares a scattered
+sample-location grid (see [`PointSet`](#points)) instead of a mesh; when set,
+`grid_coords()` reports it in place of whatever the component's own
+`_coords` carry, and a Connector delivering to this component samples onto
+those locations (`Connector(sample=...)`, see below). Phases:
+`advertise() -> tuple[list[str], list[str]]`,
+`realize(clock: Clock) -> None` (raises `CadenceError` when `timestep` is not
+a multiple of `clock.dt`), abstract `initialize(x, coords)` and
+`run(time)`, `finalize()`. Helpers: `should_run(time) -> bool` (a temporary
+shim — cadence gating lives in the Driver's slot alignment; slated for
+deletion), `grid_coords() -> CoordSystem | None`,
+`resolve_std_to_raw(coords) -> dict[str, str]`,
+`publish(x, coords, valid_time) -> None` (splits a model output tensor into
+export Fields, attaching declared masks/verticals; raises `CouplingError` if
+an advertised export is missing from the output variables).
+
+### `CallableComponent`
+
+```python
+CallableComponent(
+ name: str,
+ fn: StepFn, # (x, coords) -> (x, coords)
+ timestep: DeltaLike,
+ imports: Iterable[str] = (),
+ exports: Iterable[str] = (),
+ import_adapter: ImportAdapter | None = None,
+ **kwargs, # forwarded to Component
+)
+```
+
+Wraps a plain step function — the entry point for synthetic components and
+non-ML models. `initialize` publishes the IC at `clock.start` so lagged
+coupling has t0 data. Property `state -> tuple[torch.Tensor, CoordSystem]`.
+
+### `PrognosticComponent`
+
+```python
+PrognosticComponent(
+ name: str,
+ model: Any, # earth2studio PrognosticModel
+ timestep: DeltaLike | None = None,
+ imports: Iterable[str] = (),
+ exports: Iterable[str] | None = None,
+ import_adapter: ImportAdapter | None = None,
+ next_input: NextInputFn | None = None, # (prev_x, prev_coords, out, out_coords) -> (x, coords)
+ **kwargs,
+)
+```
+
+Wraps an earth2studio `PrognosticModel`, calling it directly (not via
+`create_iterator`) so imports can be injected between steps. `timestep`
+defaults to the model's output/input lead-time difference; `exports` default
+to the output variables resolvable through the dictionary. Published exports
+have singleton batch/time/lead_time dims squeezed away. The default
+`next_input` handles single-window models and raises `CouplingError` for
+multi-window ones (supply the hook). Also: `to(device)`, property `state`.
+
+### `DataComponent`
+
+```python
+DataComponent(
+ name: str,
+ source: Any, # earth2studio DataSource
+ exports: Iterable[str],
+ timestep: DeltaLike,
+ variable_map: Mapping[str, str] | None = None, # std name -> raw source name
+ interp_to: CoordSystem | None = None,
+ device: Any = "cpu",
+ **kwargs,
+)
+```
+
+Prescribed forcing from a data source via `fetch_data` at its own cadence
+(`requires_ic = False`). `initialize()` with no arguments fetches at
+`clock.start`; an explicit `(x, coords)` pair (with a `"variable"` dim) is
+published as-is instead. Raw source names resolve through `variable_map`,
+then `variable_aliases`, then the dictionary's aliases.
+
+### `DiagnosticComponent`
+
+```python
+DiagnosticComponent(
+ name: str,
+ model: Any, # earth2studio DiagnosticModel
+ timestep: DeltaLike,
+ imports: Iterable[str] | None = None,
+ exports: Iterable[str] | None = None,
+ **kwargs,
+)
+```
+
+Stateless single-step transform (`requires_ic = False`): stacks its imports in
+the model's raw variable order, adds singleton dims the model expects, calls
+`model(x, coords)`, publishes the outputs. Imports/exports default to the
+model's `input_coords()`/`output_coords()` variables resolved through the
+dictionary (every input must resolve, or `UnknownFieldError`). Missing
+imports at run time raise `CouplingError` pointing at the run sequence.
+Also `to(device)`.
+
+## connector
+
+### `Regridder` — the spatial-regrid callable contract
+
+```python
+Regridder = Callable[[torch.Tensor], torch.Tensor]
+```
+
+A regridder maps a tensor whose **trailing dims are the source spatial dims**
+to the same tensor on the destination grid (leading batch/window dims are
+preserved): `tensor[..., H, W] -> tensor[..., H', W']`, or for HEALPix
+`[..., face, h, w]` layouts the trailing three. It must be pure torch (no
+numpy round-trip) to keep autograd intact. When passed as
+`Connector(regridder=...)` it is applied to every field of that connector,
+and the output coords are rebuilt from the destination's `grid_coords()`.
+Build HEALPix ones with `earth2grid`, as `models/px/dlesym.py` does.
+
+### `Connector`
+
+```python
+Connector(
+ src: Component,
+ dst: Component,
+ fields: list[str] | None = None,
+ time_policy: Literal["constant", "linear"] = "constant",
+ fill: Literal["none", "zero", "nearest"] = "none",
+ regridder: Regridder | None = None,
+ sample: Literal["nearest", "bilinear"] | None = None,
+ window: DeltaLike | None = None,
+ reduce: Literal["mean", "sum", "max", "min"] | None = None,
+)
+```
+
+Moves matched fields `src.exports -> dst.imports` through the pipeline
+time policy → vertical → mask fill → spatial regrid. `fields=None` matches
+the intersection of advertised exports/imports; an explicit list must be in
+both (`IncompatibleFieldError` otherwise, also raised when the intersection
+is empty). `match() -> list[str]` also unit-checks every field
+(`UnitsMismatchError`). `execute(time: np.datetime64) -> None` performs the
+transfer (raises `CouplingError` if the source has not produced a field yet);
+`last_transfer: dict[str, Field]` holds the most recent delivery (see
+`Driver.probe`); `reset()` clears per-run exchange state (history, running
+reduction, probes). `time_policy="linear"` extrapolates from the two most
+recent exports and falls back to constant (with one warning) for fields
+carrying a `lead_time`/`window` dim. Auto-regrid requires regular 1D lat/lon
+source grids; identical grids pass through as identity; differing HEALPix
+`face` grids require `regridder=` (`IncompatibleFieldError`).
+
+`sample` targets a destination whose `grid_coords()` advertises a `"point"`
+dim (`dst.points` is a [`PointSet`](#points) — stations, sites, arbitrary
+query coordinates) rather than a mesh: `"bilinear"` reuses the mesh
+regridder's kernel evaluated per point; `"nearest"` is a great-circle
+nearest-neighbor lookup. Mutually exclusive with `regridder=`
+(`CouplingError`); a point-target destination with neither set also raises
+`CouplingError` at `execute()` time rather than guessing, as does a
+`"point"`-dim destination with no `points=` set on it, or a source without a
+regular 1D lat/lon grid (`IncompatibleFieldError`). A custom `regridder=`
+still works against a point destination — `sample=` is a convenience over
+the same generic-override path, not the only way in.
+
+`window`/`reduce` must be set together (`CouplingError` otherwise) and make
+this a **windowed connector**: each `execute` folds the source exports into a
+trailing running reduction, and delivery happens only at execute times
+aligned to `window`. Matching pairs each source export `base` with a
+destination import whose dictionary entry carries
+`CellMethod(base, reduce, window)` — the delivered Field carries that
+*derived* standard name; no matching derived import raises `CouplingError`
+(the coupler never invents names), and `match()` returns base names plus
+derived names. The window origin is the `valid_time` of the first execute's
+source field (the clock start under lagged coupling). Mid-window the
+destination's previous import is untouched; `time_policy` does not apply on
+the windowed path. This is the preferred replacement for a single-source
+`AccumulationMediator`.
+
+## pull
+
+Pull-pattern coupling for models that fetch their own forcing via
+`fetch_data(self.conditioning_data_source, ...)` inside `__call__`
+(StormCast is the canonical case). The pull path crosses `fetch_data`'s
+xarray/numpy boundary, so pull-coupled components are **inference-only** —
+no autograd through the exchange.
+
+### `StateDataSource`
+
+```python
+StateDataSource(
+ state: State,
+ raw_to_std: Mapping[str, str] | None = None,
+ strict_time: bool = False,
+ dictionary: FieldDictionary | None = None, # defaults to DEFAULT_DICTIONARY
+)
+```
+
+An in-memory object satisfying the earth2studio DataSource protocol:
+`__call__(time, variable) -> xr.DataArray` with dims
+`(time, variable, lat, lon)`, built from the State's Fields (which must be
+exchange-shaped `(lat, lon)`; other dims raise `CouplingError`). Requested
+names resolve in order: **(1)** a standard name already in the State,
+**(2)** the `raw_to_std` map (built by `PullAdapter` from the Exchange's
+`std_to_raw`, which only covers state variables), **(3)** dictionary
+fallback — a raw/alias name (e.g. `u10m`, `t2m`) is resolved to its standard
+name through `dictionary` and looked up in the State. No hit raises
+`CouplingError` naming the held fields. The source is a **snapshot view**:
+every requested time receives whatever the connector last delivered;
+`strict_time=True` raises `CouplingError` when a requested time differs from
+a served field's `valid_time`.
+
+### `PullAdapter`
+
+```python
+PullAdapter(
+ attribute: str = "conditioning_data_source",
+ strict_time: bool = False,
+ dictionary: FieldDictionary | None = None,
+)
+```
+
+`ImportAdapter` for pull-pattern models. Each call sets
+`model.` to a fresh `StateDataSource` over `exchange.imports`
+(with `raw_to_std` inverted from `exchange.std_to_raw`, and
+`strict_time`/`dictionary` forwarded), then returns
+`model(exchange.x, exchange.coords)` unchanged — the model's own fetch
+receives this step's coupled forcing. A model without the attribute raises
+`CouplingError` pointing at `ConditioningKwargAdapter` for
+argument-style conditioning. `strict_time=False` (the default) serves the
+snapshot and lets the run sequence own cadence alignment — a sequential
+connect before the pulling component's run guarantees fresh forcing.
+
+## vertical
+
+### `PressureLevels` (frozen dataclass)
+
+```python
+PressureLevels(levels: tuple[float, ...])
+```
+
+Constant pressure levels in hPa, ordered top to bottom (increasing, else
+`ValueError`). Method `pressure_pa() -> np.ndarray`.
+
+### `HybridLevels` (frozen dataclass)
+
+```python
+HybridLevels(a: tuple[float, ...], b: tuple[float, ...], ps_field: str = "surface_pressure")
+```
+
+Hybrid sigma-pressure levels `p_k = a_k + b_k * p_s` (`a` in Pa, `b`
+dimensionless, top to bottom). Validates equal lengths and strict pressure
+monotonicity across the plausible surface-pressure range [500, 1100] hPa
+(`ValueError`). `ps_field` names the surface-pressure export a connector
+pulls from the source automatically; `len(hybrid)` is the level count.
+
+The interpolation kernel `interp_to_pressure(x, coords, src, dst, ps=None)`
+lives in `earth2studio.nvcoupler.vertical` (not re-exported); connectors call
+it for you. Linear in log-pressure, clamped at column ends, differentiable.
+
+## points
+
+### `PointSet` (frozen dataclass)
+
+```python
+PointSet(lat: np.ndarray, lon: np.ndarray, names: tuple[str, ...] | None = None)
+```
+
+A fixed set of N scattered sample locations (stations, sites, arbitrary
+query points) — the "point" analog of a lat/lon mesh. `lat`/`lon` are 1-D,
+equal length, at least one point (`CouplingError` otherwise); `names`, if
+given, must match that length. `labels() -> np.ndarray` is `names` when set,
+else an integer index `0..N-1` — this is what the `"point"` dim's own
+coordinate array carries, matching every other CoordSystem dim.
+`grid_coords() -> CoordSystem` returns `{"point": labels()}`, what
+`Component.grid_coords()` reports when the component's `points=` is set.
+`signature()` is a hashable cache key (mirrors `Field.grid_signature`), used
+to cache built samplers per destination point set.
+
+Construct a point-target component by passing `points=PointSet(...)` to any
+`Component` subclass, then deliver to it with `Connector(..., sample=...)`
+(see the [`connector`](#connector) section above).
+
+## mediator
+
+### `Mediator` (base class)
+
+```python
+Mediator(name: str, timestep: Any, imports=(), exports=(), **kwargs)
+```
+
+`Component` whose import `State` forwards every delivered field to
+`accumulate(field)`; when scheduled by a `MediateAction`, `run(time)` calls
+`compute(time)`, which must populate `export_state`. Subclass and implement
+both to build custom reductions (including unit conversions —
+[the v1 remedy for mismatched units](errors_and_troubleshooting.md)).
+`requires_ic = False`.
+
+### `AccumulationMediator`
+
+```python
+AccumulationMediator(name: str, fields: list[str], window: Any = None, **kwargs)
+```
+
+Windowed running reduction (O(1) memory in window length). Each entry of
+`fields` must be a *derived* dictionary entry carrying a `CellMethod`
+(`CouplingError` otherwise); the cell method supplies the base import, the
+reduction (mean/sum/max/min), and the window, which becomes the mediator's
+timestep unless `window=` overrides it (`CouplingError` when fields disagree
+on windows and no override is given). Duplicate deliveries with the same
+`valid_time` are ignored; `compute` with zero samples raises `CouplingError`.
+After each compute, `samples_last_window: dict[str, int]` reports the counts
+and the accumulators reset. For one source feeding one destination, prefer
+the windowed connector (`Connector(..., window=, reduce=)`), which shares the
+same accumulator core; mediators are the multi-source / custom-reduction
+generalization.
+
+### `TrailingAverageMediator`
+
+```python
+TrailingAverageMediator(name: str, fields: list[str], window: Any = None, **kwargs)
+```
+
+`AccumulationMediator` restricted to mean reductions (`CouplingError` for
+non-mean fields) — the exact semantics of DLESyM's ocean coupling and
+PhysicsNeMo's `TrailingAverageCoupler`.
+
+## sequence
+
+### Action dataclasses (frozen)
+
+```python
+RunAction(component: str)
+ConnectAction(src: str, dst: str)
+MediateAction(mediator: str, phase: str = "compute")
+```
+
+`Action = RunAction | ConnectAction | MediateAction`; `str()` of each emits
+its DSL line.
+
+### `Slot` / `RunSequence`
+
+```python
+Slot(interval: np.timedelta64, actions: list[Action] = [])
+RunSequence(slots: list[Slot])
+```
+
+`Slot.interval` is coerced through `as_timedelta` (interval strings accepted).
+`RunSequence` methods: `components_run() -> set[str]`,
+`connections() -> list[ConnectAction]`,
+`validate(components: dict, dt: DeltaLike) -> None` (raises `SequenceError` /
+`CadenceError`; full rule list in
+[dsl_and_yaml_reference.md](dsl_and_yaml_reference.md#validation)), and
+`__str__` emitting round-trippable DSL.
+
+### `parse_run_sequence`
+
+```python
+parse_run_sequence(text: str) -> RunSequence
+```
+
+Parses the DSL (grammar in
+[dsl_and_yaml_reference.md](dsl_and_yaml_reference.md#grammar)); raises
+`SequenceError` with the offending line number.
+
+### `derive_sequence`
+
+```python
+derive_sequence(
+ components: dict[str, Component],
+ connectors: Iterable[Connector | tuple[str, str]] | None = None,
+ lagged: set[tuple[str, str]] | Literal["all"] = "all",
+) -> RunSequence
+```
+
+Derives the canonical run sequence from the coupling graph: one slot per
+distinct component cadence, fast to slow. Within a slot: lagged connects
+delivered at this cadence (sorted by component declaration order), then each
+mediator's compute followed by its outgoing connects, then component runs
+topologically ordered over the sequential (non-lagged) edges, each run
+followed by its outgoing sequential connects. `lagged="all"` (default) is the
+NUOPC-explicit shape; edges not in the `lagged` set are sequential, and a
+cycle of sequential edges among same-cadence components raises
+`SequenceError` telling you to mark one edge lagged. Unknown endpoint names
+raise `SequenceError` with did-you-mean suggestions. This is what
+`Driver(sequence=None)` and `couple()` call.
+
+## driver
+
+### `Driver`
+
+```python
+Driver(
+ components: dict[str, Component],
+ sequence: RunSequence | str | None = None,
+ clock: Clock | None = None,
+ connectors: list[Connector | tuple[str, str]] | None = None,
+ collect: bool = True,
+ io: dict[str, IOBackend] | None = None,
+ allow_unfed_imports: bool = False,
+)
+```
+
+Executes a coupled system declared as components + connections. With
+`sequence=None` (the default, declarative form) the schedule is derived from
+the coupling graph via `derive_sequence(components, connectors)` at
+construction; an explicit `RunSequence` or DSL string overrides it (the
+escape hatch for sequential coupling or hand-tuned action order). The
+attribute `sequence_derived: bool` records which path was taken. `clock`
+keeps its positional slot but omitting it raises an actionable
+`CouplingError`. `connectors` accepts `Connector` instances or bare
+`(src, dst)` name tuples (auto-built into default Connectors; unknown names
+raise `CouplingError`); with an explicit sequence, any `ConnectAction`
+without a prebuilt connector still gets a default `Connector(src, dst)`.
+`io=` streams each component's exports to an earth2studio `IOBackend` (one
+array per export field, leading `time` axis over the component's ring times
+including t0; NaN-initialized), independent of `collect`. Unknown `io` keys
+raise `CouplingError`.
+
+- `initialize(ics: dict[str, tuple[torch.Tensor, CoordSystem]] | None = None) -> None`
+ — validates the sequence, matches connectors (units checks), rejects unfed
+ imports (`UnmatchedImportError`, downgraded to a warning by
+ `allow_unfed_imports=True`), warns on unconsumed exports and oversized
+ in-memory collection, then realizes and initializes every component.
+ Components with `requires_ic=False` (mediators, data components,
+ diagnostics) need no `ics` entry; missing ICs for the rest raise
+ `CouplingError`.
+- `run() -> dict[str, xr.Dataset]` — runs to `clock.stop` under
+ `torch.inference_mode()`; returns `to_xarray()` when `collect=True`, else
+ `{}`.
+- `steps() -> Iterator[tuple[np.datetime64, dict[str, State]]]` — yields
+ `(time, {component: export State})` after every driver step (inference
+ mode); the notebook-inspection path.
+- `rollout(n_steps: int) -> dict[str, State]` — advances `n_steps` keeping
+ the autograd graph when grad is enabled (the coupled fine-tuning entry
+ point); raises `CouplingError` when fewer steps remain.
+- `reset() -> None` — rewinds the clock, clears records, and calls
+ `Connector.reset()` on every connector (time-policy history, probes,
+ windowed running reductions and window origins); `initialize(ics)` must be
+ called again before running. Running an exhausted clock raises
+ `CouplingError`.
+- `probe(connector: str) -> dict[str, Field]` — last exchanged fields on a
+ connector addressed as `"src->dst"` (`KeyError` listing known names).
+- `to_xarray() -> dict[str, xr.Dataset]` — collected exports, one Dataset per
+ component with a leading `time` axis of that component's ring times.
+- `describe() -> str` / `_repr_html_()` — delegate to `api.describe` /
+ `api.describe_html`.
+
+## api
+
+### `couple`
+
+```python
+couple(
+ *components: Component,
+ start: TimeLike,
+ stop: TimeLike,
+ dt: DeltaLike | None = None,
+ connectors: list[Connector] | None = None,
+ collect: bool = True,
+) -> Driver
+```
+
+Auto-wires components into a ready-to-initialize `Driver`: every import is
+matched to its unique exporter by standard name (`AmbiguousCouplingError` for
+several, `UnmatchedImportError` for none). A derived import whose base field
+someone exports becomes a windowed
+`Connector(src, dst, fields=[base], window=cm.window, reduce=cm.method)`; an
+`AccumulationMediator` (named `med_`) is synthesized only
+when the `(src, dst)` pair already carries a plain transfer, which a windowed
+connector cannot share. A user-prebuilt windowed connector for the pair is
+honored. The run sequence is derived from the graph (`sequence_derived=True`)
+in the canonical lagged layout, one slot per cadence. `dt` defaults to the
+GCD of the component timesteps. Returns an *uninitialized* driver.
+
+### `coupled`
+
+```python
+coupled(
+ time: TimeLike,
+ stop_or_nsteps: TimeLike | int,
+ components: Sequence[Component] | dict[str, Component],
+ ics: dict[str, tuple],
+ dt: DeltaLike | None = None,
+ collect: bool = True,
+ verbose: bool = True,
+) -> dict
+```
+
+One call from initial conditions to `dict[str, xarray.Dataset]`:
+`couple(...)` + `initialize(ics)` + a tqdm-wrapped run. An integer
+`stop_or_nsteps` means that many driver (`dt`) steps.
+
+### `describe` / `describe_html`
+
+```python
+describe(driver: Driver) -> str
+describe_html(driver: Driver) -> str
+```
+
+Terraform-plan-style preview (text / self-contained Jupyter HTML) of
+components, connectors (fields, policies, lagged/sequential mode, slot), and
+the run sequence. The mode column is per exchange: `sequential` iff the
+source ran (or the mediator computed) earlier in the same slot — the
+destination consumes state produced this iteration — else `lagged`. Works
+before `initialize()` — only advertised names, the sequence, and the clock
+are consulted.
+
+## config
+
+### `to_yaml`
+
+```python
+to_yaml(driver: Driver, path: str | os.PathLike | None = None) -> str
+```
+
+Serializes a `Driver` to YAML text (also written to `path` when given).
+Hand-written sequences serialize as plain DSL text; derived sequences as
+`sequence: {derived: true, text: }` (the text is informational —
+`from_yaml` re-derives). Windowed connectors carry `window`/`reduce` keys.
+Raises `CouplingError` for any component that is neither an
+`AccumulationMediator` nor carries a `yaml_spec` attribute. Schema and rules
+in [dsl_and_yaml_reference.md](dsl_and_yaml_reference.md#the-yaml-schema).
+
+### `from_yaml`
+
+```python
+from_yaml(path_or_str: str | os.PathLike) -> Driver
+```
+
+Builds an **uninitialized** `Driver` from YAML text or a file path (a
+newline-free string naming an existing file is read as a path). Components
+are rebuilt by importing each `class` path and calling it with `kwargs`;
+failures raise `CouplingError` naming the path. A `sequence` mapping with
+`derived: true` re-derives the schedule from components + connectors
+(deterministic, so round-trips reproduce identical runs); a mapping without
+it raises `CouplingError`. Windowed connectors are rebuilt from their
+`window`/`reduce` keys. Call `driver.initialize(ics)` afterwards.
+
+## errors
+
+All configuration errors derive from `CouplingError`; messages name the
+components/fields involved and the concrete fix. See
+[errors_and_troubleshooting.md](errors_and_troubleshooting.md) for triggers
+and remedies.
+
+| Exception | Bases | Raised when |
+|---|---|---|
+| `CouplingError` | `Exception` | Base class; also raised directly for generic misconfiguration (missing ICs, unproduced exports, bad `yaml_spec`, ...) |
+| `UnknownFieldError(name, candidates)` | `CouplingError` | A name is not a registered standard name or alias (did-you-mean suggestions) |
+| `UnmatchedImportError(component, field, available_exports)` | `CouplingError` | An advertised import that nothing exports/delivers |
+| `UnitsMismatchError(field, src, src_units, dst, dst_units)` | `CouplingError` | Matched fields disagree on (normalized) units |
+| `IncompatibleFieldError` | `CouplingError` | A connector cannot reconcile matched fields (grid layout, mask fill, missing regridder) |
+| `VerticalMismatchError` | `CouplingError` | Vertical coordinates cannot be reconciled (missing `vertical`, missing surface pressure, non-monotone hybrid levels) |
+| `CadenceError(what, interval, dt)` | `CouplingError` | An interval is not a positive multiple of the reference dt |
+| `AmbiguousCouplingError(field, importer, exporters)` | `CouplingError` | `couple()` found several exporters for one import |
+| `SequenceError` | `CouplingError` | Run-sequence parse or validation failure |
+
+## dlesym_split
+
+### `split_dlesym`
+
+```python
+split_dlesym(dlesym: Any, dictionary: FieldDictionary | None = None)
+ -> tuple[DLESyMAtmosComponent, DLESyMOceanComponent]
+```
+
+Re-exposes a constructed `earth2studio.models.px.DLESyM`'s internal atmos and
+ocean sub-models as two components exchanging SST and 48 h window-mean
+coupling fields through explicit connectors (identity transfers on the shared
+HEALPix grid). Both components step the full 96 h parent cadence and call the
+parent's own coupling/insolation methods. Unknown DLESyM variables and the
+window-mean entries are auto-registered on a private copy of
+`DLESYM_DICTIONARY`. **Honest limitation:** this module has only been
+exercised against structural mocks — the real-weights equivalence gate
+(`test/nvcoupler/test_dlesym_weights_equivalence.py`) has not been run.
+
+### `build_dlesym_driver`
+
+```python
+build_dlesym_driver(
+ dlesym: Any,
+ start: Any,
+ stop: Any,
+ dictionary: FieldDictionary | None = None,
+ collect: bool = True,
+) -> Driver
+```
+
+`split_dlesym` plus a `Driver` whose 96 h run sequence reproduces the native
+`DLESyM._forward` ordering (SST lagged across steps, window means sequential
+within a step). `stop - start` must be a multiple of 96 h. Initialize both
+halves with the same DLESyM-layout IC:
+`driver.initialize({"atmos": (x, coords), "ocean": (x, coords)})`.
+
+### `DLESYM_DICTIONARY`
+
+Module-level `FieldDictionary` copy of the default — the extension point for
+non-default DLESyM vocabularies.
+
+## testing
+
+Deterministic toy components — public API for downstream tests and the docs'
+own snippets. A two-component system with the DLESyM cadence structure whose
+values are hand-computable from spatially constant ICs (the executable spec in
+`test/nvcoupler/test_driver.py` relies on this).
+
+```python
+ATMOS_GRID = (32, 64)
+OCEAN_GRID = (16, 32)
+
+grid_coords(nlat: int, nlon: int) -> CoordSystem
+```
+
+Regular lat/lon coords, 90 → −90, 0 → 360 (endpoint excluded).
+
+```python
+fake_atmos(gain: torch.Tensor | float = 1.0, timestep: str = "6h") -> CallableComponent
+```
+
+`"atmos"`: imports `sea_surface_temperature`, exports
+`geopotential_at_1000hpa`; update `z ← z + 1 + gain·0.1·sst`. Pass
+`gain=torch.tensor(1.0, requires_grad=True)` to test gradient flow across the
+exchange.
+
+```python
+fake_ocean(gain: torch.Tensor | float = 1.0, timestep: str = "48h", with_mask: bool = False) -> CallableComponent
+```
+
+`"ocean"`: imports `geopotential_at_1000hpa_48h_mean`, exports
+`sea_surface_temperature`; update `sst ← sst + gain·0.01·z48m`.
+`with_mask=True` attaches a land mask (northern half invalid) to the SST
+export.
+
+```python
+atmos_ic(z0: float = 0.0, sst0: float = 2.0) -> tuple[torch.Tensor, CoordSystem]
+ocean_ic(sst0: float = 2.0, z48m0: float = 0.0) -> tuple[torch.Tensor, CoordSystem]
+```
+
+Spatially constant initial conditions on the matching grids, with `variable`
+coords `["z1000", "sst"]` / `["sst", "z48m"]`.
+
+Minimal end-to-end use (executed):
+
+```python
+import earth2studio.nvcoupler as nvc
+from earth2studio.nvcoupler.testing import atmos_ic, fake_atmos, fake_ocean, ocean_ic
+
+driver = nvc.couple(fake_atmos(), fake_ocean(), start="2024-01-01", stop="2024-01-05")
+print(driver.describe()) # plan preview, incl. the synthesized windowed connector
+driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+datasets = driver.run() # dict[str, xr.Dataset]
+```
+
+See `examples/09_nvcoupler/`
+([README](../../../examples/09_nvcoupler/README.rst)) for the five worked
+examples built on these toys.
diff --git a/earth2studio/nvcoupler/docs/concepts.md b/earth2studio/nvcoupler/docs/concepts.md
new file mode 100644
index 000000000..36a966000
--- /dev/null
+++ b/earth2studio/nvcoupler/docs/concepts.md
@@ -0,0 +1,603 @@
+# Concepts
+
+The deep conceptual reference for nvcoupler. The [package README](../README.md)
+is the overview; this page pins down the exact semantics of each abstraction.
+For task-oriented recipes see the [user guide](user_guide.md), for the DSL and
+YAML grammar the [DSL and YAML reference](dsl_and_yaml_reference.md), for
+signatures the [API reference](api_reference.md), and for every error class the
+[errors and troubleshooting page](errors_and_troubleshooting.md). Rationale for
+these designs lives in [design and roadmap](design_and_roadmap.md).
+
+## Field and State
+
+A `Field` is exactly one physical quantity: a torch tensor, its earth2studio
+`CoordSystem`, a canonical identity (`standard_name`, `units`), an optional
+`valid_time`, `source` (provenance), `mask`, and `vertical` metadata. Two rules
+are enforced at construction:
+
+- **One variable per Field.** A `"variable"` dimension in `coords` raises
+ `CouplingError` immediately — multi-variable tensors are split with
+ `State.from_tensor`.
+- **Coords insertion order is dimension order.** `data.ndim` must equal
+ `len(coords)`, and the *i*-th key of the coords dict describes the *i*-th
+ tensor axis. There is no name-based reordering anywhere in the framework.
+
+```python
+from collections import OrderedDict
+import numpy as np
+import torch
+import earth2studio.nvcoupler as nvc
+
+coords = OrderedDict(
+ {"lat": np.linspace(90.0, -90.0, 8), "lon": np.linspace(0.0, 360.0, 16, endpoint=False)}
+)
+sst = nvc.Field(
+ data=torch.full((8, 16), 290.0),
+ coords=coords,
+ standard_name="sea_surface_temperature",
+ units="K",
+ valid_time=np.datetime64("2024-01-01"),
+)
+# Field('sea_surface_temperature' [K], lat: 8, lon: 16, valid_time=2024-01-01)
+```
+
+A `State` is a mutable mapping of Fields keyed by standard name (the key must
+equal the field's `standard_name`, enforced on `__setitem__`). Every component
+owns an import State and an export State; Connectors move Fields between them.
+`State.from_tensor` splits a tensor along `"variable"`, resolving raw model
+names through a `FieldDictionary`; `State.as_tensor(names)` stacks Fields back
+along a new `"variable"` axis inserted immediately before the first spatial
+dimension (earth2studio's `batch, time, lead_time, variable, spatial...`
+convention). All stacked fields must share identical coords — regrid through a
+Connector first.
+
+```python
+x = torch.zeros(2, 8, 16)
+tensor_coords = OrderedDict({"variable": np.array(["z1000", "sst"]), **coords})
+state = nvc.State.from_tensor(
+ "demo", x, tensor_coords, nvc.DEFAULT_DICTIONARY,
+ valid_time=np.datetime64("2024-01-01"),
+)
+sorted(state) # ['geopotential_at_1000hpa', 'sea_surface_temperature']
+y, ycoords = state.as_tensor(["sea_surface_temperature", "geopotential_at_1000hpa"])
+```
+
+Field data is torch end to end — never round-tripped through numpy inside the
+framework — so autograd graphs survive every exchange.
+
+### The physical-units exchange contract
+
+Fields cross the coupling seam in **physical units**. A model that trains on
+normalized inputs normalizes internally: the DLESyM split adapter
+([`dlesym_split.py`](../dlesym_split.py)) is the reference implementation —
+each half denormalizes its outputs (`data * scale + center`) before publishing
+and renormalizes imports (`(data - center) / scale`) before calling its U-Net.
+Because both directions use the same per-variable constants, the round trip is
+exact up to float rounding, and `test/nvcoupler/test_dlesym_split.py` asserts
+the reconstructed coupling tensor matches the parent's own math. The payoff:
+any component can consume any other's exports without knowing its
+normalization statistics.
+
+Units are **checked, not converted** (see the dictionary section below); no
+scaling ever happens silently in a Connector.
+
+### valid_time semantics — and an honest caveat
+
+`valid_time` is the single instant the data is valid for. Components stamp it
+when publishing (`initialize` publishes at `clock.start`, so lagged coupling
+has data at t0; each `run(time)` publishes at `time`), and Connector time
+policies reason about it.
+
+The caveat: some exported Fields are not instantaneous. A field carrying a
+`lead_time` or `window` dimension (e.g. the DLESyM atmos component's
+window-mean exports, which have a leading `window` axis over the ocean's 48 h
+chunks) represents *many* times but carries one `valid_time` stamp — the end
+of the producing step's window. That is fine for `"constant"` transfers, but a
+single-timestamp extrapolation is ill-defined for such fields, so the
+Connector's `"linear"` policy detects `lead_time`/`window` dims and **falls
+back to `"constant"` for that field**, logging a one-time warning per field.
+
+### Masks
+
+`mask` is an optional boolean tensor broadcastable to `data`, with
+**True = valid** (e.g. ocean points for SST; the toy ocean in
+[`testing.py`](../testing.py) marks the northern half of its grid as land with
+`False`). Masks are declared per export via a component's `export_masks=` and
+consumed by the Connector's fill stage; after filling, the transferred field's
+mask is cleared (`None`) since every point is then valid.
+
+### Grid signatures
+
+`Field.grid_signature()` returns a hashable tuple of
+`(dim, shape, values.tobytes())` for every spatial dim
+(`level, face, lat, lon, hpx, height, width, y, x`). Connectors key their
+lazily built regridders and mask fillers on it, so the KDTree/interpolation
+setup cost is paid once per distinct grid (and, for fillers, per distinct
+mask), not per step — the `dxwrapper.py` caching pattern.
+
+## FieldDictionary
+
+Connectors match exports to imports by **standard name**, never by raw model
+variable strings. A `FieldDictionary` maps both standard names and aliases to
+`FieldEntry(standard_name, canonical_units, description, aliases, cell_method)`.
+Lookup is case-sensitive; an alias maps to exactly one standard name (remapping
+raises), an alias may not collide with a standard name, and re-registering a
+standard name replaces its entry. `DEFAULT_DICTIONARY` ships a curated v1
+vocabulary (earth2studio surface variables, the pressure-level fields used by
+the coupled models in this repo, and the mediator-produced derived fields);
+components each get a private *copy*, extended by their `variable_aliases=`
+(raw model name → standard name).
+
+### Units: checked, not converted
+
+`normalize_units` collapses cosmetic differences before comparing: lowercase,
+strip `**`/`^` exponent markers and spaces (behaviorally aligned with
+`earth2studio.lexicon.earthmover.normalize_units`), then apply a synonym table
+(`m/s` ≡ `m s-1`, `kelvin` ≡ `K`, `mm` ≡ `kg m-2` for precipitation depth,
+the dimensionless family `1`/`(0-1)`/`fraction` ≡ `""`, `celsius` ≡ `degC`,
+...). Values are **never converted** — a genuine disagreement raises
+`UnitsMismatchError` at connector match time, naming both components and
+suggesting the fix (convert in a Mediator or align the dictionary entries).
+
+```python
+from earth2studio.nvcoupler.dictionary import normalize_units
+
+normalize_units("m s**-1") == normalize_units("m/s") == "m s-1" # True
+
+nvc.DEFAULT_DICTIONARY.check_units(
+ "sea_surface_temperature", "degC", src="ocean", dst="atmos"
+) # raises UnitsMismatchError: 'ocean' exports 'degC' but 'atmos' expects 'K'
+```
+
+### CellMethod: derived fields as first-class dictionary entries
+
+A time-reduced field (a 48 h precipitation sum, a 24 h temperature max) is a
+dictionary entry carrying a machine-readable `CellMethod(base, method, window)`
+with `method` in `mean | sum | max | min` — the declaration "I am *method* of
+*base* over *window*". No suffix string-parsing anywhere.
+
+Three consumers read it:
+
+- **Windowed `Connector`s** — `Connector(src, dst, window=..., reduce=...)`
+ pairs each source export `base` with a destination import whose entry
+ carries `CellMethod(base, reduce, window)` and delivers under that derived
+ name (see [the connector pipeline](#the-connector-pipeline)).
+- **`AccumulationMediator`** — constructed with the *derived* names, it
+ resolves each entry, takes the base field as its import, the method as its
+ running reduction, and the (common) window as its timestep unless
+ `window=` overrides it. A derived name without a `cell_method` raises.
+- **`couple()`** — when a component imports a derived field nobody exports,
+ auto-wiring checks the entry's cell method: if some unique component exports
+ the *base* field, it wires a windowed
+ `Connector(base-exporter, importer, fields=[base], window=cm.window,
+ reduce=cm.method)`. Only when that `(src, dst)` pair already carries a
+ plain transfer (the importer also imports the base directly, or a second
+ derived field rides the same pair) is an `AccumulationMediator` synthesized
+ and wired `base-exporter -> mediator -> importer` instead; a user-prebuilt
+ windowed connector for the pair is honored. No cell method or no base
+ exporter raises `UnmatchedImportError`; multiple base exporters raise
+ `AmbiguousCouplingError`.
+
+```python
+entry = nvc.DEFAULT_DICTIONARY.resolve("total_precipitation_48h_sum")
+entry.cell_method
+# CellMethod(base='total_precipitation_6h', method='sum', window=numpy.timedelta64(48,'h'))
+
+med = nvc.AccumulationMediator("med", ["total_precipitation_48h_sum"])
+med.import_names # ['total_precipitation_6h']
+med.export_names # ['total_precipitation_48h_sum']
+med.timestep # 48 hours (from the cell method's window)
+```
+
+## Clock and the cadence model
+
+Three intervals coexist and must nest:
+
+1. **Driver dt** — the `Clock`'s step, the finest granularity at which
+ anything can happen.
+2. **Component timesteps** — each component declares its own cadence as a
+ plain `timestep`; the Driver runs it only in run-sequence slots whose
+ interval equals that timestep, and those slots execute only on clock
+ steps aligned with it (slot alignment). There is no per-component alarm
+ object and no offset mechanism — cadence is purely timestep alignment.
+3. **Slot intervals** — each run-sequence `@interval` slot executes on driver
+ steps aligned with it.
+
+The validation rules, all raising `CadenceError` at construction or
+`initialize` time (never mid-rollout):
+
+- `Clock(start, stop, dt)`: `stop - start` must be a positive whole multiple
+ of `dt`.
+- `Component.realize(clock)`: the component's timestep must be a positive
+ whole multiple of `dt`.
+- `RunSequence.validate`: every slot interval must be a multiple of `dt`, and
+ every component (or mediator) scheduled in a slot must have a timestep
+ **equal** to that slot's interval. Components never scheduled anywhere raise
+ `SequenceError`.
+
+```python
+import numpy as np
+
+clock = nvc.Clock("2024-01-01", "2024-01-03", "6h")
+clock.n_steps # 8
+timestep = np.timedelta64(48, "h") # a 48 h component on this 6 h clock
+[(t - clock.start) % timestep == np.timedelta64(0) for t in clock.times()]
+# [True, False, False, False, False, False, False, False, True]
+
+nvc.Clock("2024-01-01", "2024-01-02", "7h") # raises CadenceError (24 h span, 7 h dt)
+```
+
+Iterating a Clock yields times *after* start: the initial condition lives at
+`start` (the caller's step 0), and the first yielded time is `start + dt` —
+mirroring `earth2studio.run` where the iterator's 0th output is the IC. Every
+timestep is trivially aligned at `start` itself (elapsed time zero), but no
+slot executes there — the driver only runs actions at times strictly after
+`start`; that is why `initialize` seeds export states at t0.
+
+Clocks are **one-shot**. Running past `stop` raises, and calling
+`run()`/`steps()`/`rollout()` on an exhausted driver raises a `CouplingError`
+telling you to call `driver.reset()` — which rewinds the clock, clears
+collected records, connector history, and IO state — and then
+`driver.initialize(ics)` again before rerunning. A reset-and-reinitialized run
+reproduces the original exactly (asserted in `test_driver.py`).
+
+## Components and the NUOPC phase lifecycle
+
+Every component walks the NUOPC phases, driven by the Driver:
+
+1. **`advertise()`** — return the import and export standard-name lists
+ (declared at construction; names are resolved through the dictionary, so
+ aliases work).
+2. **`realize(clock)`** — validate the component's timestep against the
+ driver dt and attach the shared clock.
+3. **`initialize(x, coords)`** — set internal state from an initial condition
+ and *seed the export state at `clock.start`*, so lagged coupling has data
+ at t0.
+4. **`run(time)`** (repeated) — advance one component timestep; exports become
+ valid at `time`.
+5. **`finalize()`** — cleanup hook (no-op by default).
+
+### The requires_ic contract
+
+`requires_ic` (class attribute, default `True`) tells the Driver whether
+`initialize` needs an `(x, coords)` pair. During `Driver.initialize(ics)`:
+components present in `ics` get their entry; components with
+`requires_ic = False` (mediators, `DataComponent`, `DiagnosticComponent`) are
+initialized with no arguments — a DataComponent fetches at `clock.start`, a
+DiagnosticComponent just records its grid, a Mediator does nothing; a
+`requires_ic = True` component missing from `ics` raises `CouplingError`
+naming it. `requires_ic = False` components still *accept* an explicit IC
+(a DataComponent publishes it instead of fetching; a DiagnosticComponent
+pushes it through the model once to seed t0 exports for lagged chains).
+
+### The four component kinds
+
+- **`CallableComponent`** — wraps a plain `fn(x, coords) -> (x, coords)` step
+ function. The entry point for synthetic components and **non-ML models**
+ (process-based hydrology, crop models, anything Python-callable). Owns an
+ `(x, coords)` state; imports are injected through its ImportAdapter each run.
+- **`PrognosticComponent`** — wraps an earth2studio `PrognosticModel`.
+ Timestep defaults to the model's output-minus-input lead time; exports
+ default to the output variables resolvable through the dictionary. It steps
+ the model directly (not via `create_iterator`) so imports can be injected
+ between steps. Models with multi-window sliding inputs need a
+ `next_input(prev_x, prev_coords, out, out_coords)` hook; the default handles
+ single-window models and raises with instructions otherwise. Publishing is
+ **exchange-shaped**: singleton `batch`/`time`/`lead_time` dims are squeezed
+ so exported Fields carry plain spatial coords like everyone else's (the seam
+ tests in `test/nvcoupler/test_seams.py` exist precisely for this).
+- **`DataComponent`** — prescribed forcing from an earth2studio `DataSource`
+ (`requires_ic = False`). Instead of stepping a model it fetches its export
+ variables at its own cadence via `fetch_data`. Swapping a modeled ocean for
+ observed SST is a one-line component substitution; connectors, mediators,
+ and the run sequence are untouched.
+- **`DiagnosticComponent`** — wraps an earth2studio `DiagnosticModel`
+ (`requires_ic = False`): a stateless per-run transform that stacks its
+ imported Fields in the model's raw variable order, conforms singleton dims
+ the model expects, calls it, and publishes the outputs. Imports/exports
+ default to the model's own `input_coords()`/`output_coords()` variables.
+
+`Mediator` subclasses (below) are the fifth participant type; they run in the
+sequence like components but compute reductions rather than stepping models.
+
+## ImportAdapters: who owns the model call
+
+Real models disagree on how coupled forcing arrives, so the **adapter — not
+the component — owns the model invocation**. Each of the four built-ins maps
+to a real-world coupling pattern:
+
+| Adapter | Call shape | Real-world pattern |
+|---|---|---|
+| `VariableOverwriteAdapter` (default) | overwrite matching variable slices of `x`, then `model(x, coords)` | prescribed forcing as a state channel |
+| `ConditioningKwargAdapter` | `model.call_with_conditioning(x, coords, conditioning=..., conditioning_coords=...)` | StormScope |
+| `ExtraTensorAdapter` | `model(x, coords, coupling)` (or a named kwarg) | DLESyM / PhysicsNeMo 4-tensor |
+| `PullAdapter` | install a `StateDataSource` on `model.conditioning_data_source`, then `model(x, coords)` | StormCast |
+
+An adapter receives the model and an `Exchange` — a frozen bundle of the
+step's state tensor and coords, the delivered import `State`, the
+standard-name → raw-variable map, and the step time — and returns the stepped
+`(x, coords)`. The two `Exchange` accessors cover the common delivery shapes:
+`exchange.inject()` (variable-slice overwrite) and `exchange.stacked()`
+(channel-stacked forcing tensor).
+
+`VariableOverwriteAdapter` requires the imported field to be a state variable
+of the model (resolved through `std_to_raw` aliases); `Exchange.inject`
+clones the state tensor and uses `index_copy_` on the clone, keeping autograd
+intact. It raises with a pointer to the other adapters when the model has no
+`"variable"` dim or the import is not a state channel.
+
+The stacking adapters (`ConditioningKwargAdapter`, `ExtraTensorAdapter`)
+require an **explicit `field_order=[...]`** whenever more than one field is
+imported. Silently stacking in alphabetical order is forbidden by design:
+models are channel-order-sensitive, and a permuted conditioning tensor *runs
+without error and predicts garbage* — the worst failure mode in ML systems.
+With a single import the order is trivially inferable and `field_order` may be
+omitted.
+
+### The pull pattern: masquerade, not argument
+
+The fourth delivery shape exists because some models (StormCast is the
+canonical case) offer *no* argument to deliver forcing through: the model
+calls `fetch_data(self.conditioning_data_source, time, variables, ...)`
+inside its own `__call__`, and the only injection point it exposes is that
+settable data-source attribute. So `PullAdapter`'s job is a **masquerade
+rather than an argument**: before each step it sets the attribute to a
+`StateDataSource` — a tiny in-memory object satisfying the DataSource
+protocol that answers the model's fetches from the component's import State —
+then calls `model(x, coords)` unchanged. The model runs its unmodified
+production fetch path (fetch → interpolate → concatenate) believing it is
+reading GFS; it is reading the coupler. There is precedent for exactly this
+masquerade in earth2studio's serve workflows: `stormcast_conus_workflow.py`
+stages a full conditioning forecast to temp files and replays it through an
+`InferenceOutputSource`; the shim is the same trick minus the staging — the
+"source" is this step's live exchange.
+
+The cadence-alignment contract: a `StateDataSource` is a snapshot view —
+whatever the connector last delivered is what *every* requested time
+receives. Alignment is the run sequence's job: a **sequential** connect
+placed before the pulling component's run in the same slot
+(`global`, then `global -> stormcast`, then `stormcast`) guarantees the
+served fields are fresh at the pulled time. `strict_time=True` turns that
+contract into a check, raising when the model pulls times that do not match
+the served fields' `valid_time`.
+
+Honest limitation: the pull path runs through the model's own
+`fetch_data`/xarray machinery, so field data crosses a numpy boundary and the
+autograd graph is severed there — pull-coupled components are
+**inference-only** (no gradients through the exchange). The push-pattern
+adapters above keep autograd intact.
+
+Adapters (other than `PullAdapter`, per the boundary just described) must be
+autograd-safe (no in-place mutation of tensors that may carry
+grad); any object matching the `ImportAdapter` protocol can be passed as
+`import_adapter=`.
+
+## The Connector pipeline
+
+`Connector(src, dst)` transfers the intersection of src's advertised exports
+and dst's advertised imports, or an explicit `fields=[...]` (each of which
+must appear in *both* lists, else `IncompatibleFieldError`). An empty match
+raises. Units are checked against the dictionary at match time. Every
+`execute(time)` runs each matched field through four stages, in this order:
+
+### 1. Time policy
+
+The connector keeps a **2-deep history** per field, `(previous, latest)`,
+rotated only when a genuinely *new* export arrives (different `valid_time`) —
+re-seeing the same export on repeated executes (a slow source polled by a fast
+slot) must not collapse the extrapolation baseline.
+
+- `"constant"` (default): deliver the latest export as-is — the destination
+ holds the source's last state between updates (PhysicsNeMo
+ `ConstantCoupler` behavior).
+- `"linear"`: extrapolate from the two most recent exports toward the current
+ time: `data + (data - prev) * dt_ahead / dt_hist`, restamping `valid_time`
+ to `time`. Falls back to constant when there is no previous export, when
+ either `valid_time` is missing, when history is non-increasing, when
+ `dt_ahead == 0` — and, permanently with a one-time warning, for fields
+ carrying a `lead_time`/`window` dimension (see the valid_time caveat above).
+
+### 2. Vertical interpolation
+
+Triggered only when the **destination** declares `import_vertical` for this
+field *and* it differs from the field's `vertical` metadata. v1 supports
+interpolation onto `PressureLevels` only (destination wanting anything else
+raises `VerticalMismatchError`, as does a source field with no vertical
+metadata). For `HybridLevels` sources (`p_k = a_k + b_k * p_s`) the connector
+pulls the surface-pressure field named by `HybridLevels.ps_field` from the
+source's exports automatically, raising with a concrete fix ("add it to the
+source's export list") if absent. Interpolation is linear in log-pressure via
+`torch.searchsorted` + gathers (differentiable in values), clamped at the
+column ends. Models that encode levels in variable names (`z500`, `t850`)
+never touch this stage.
+
+### 3. Mask fill
+
+Applies only when the field carries a mask and `fill != "none"`:
+
+- `"zero"`: invalid points become 0.
+- `"nearest"`: each invalid point takes its nearest valid neighbor
+ (great-circle metric via a unit-sphere KDTree; pure gather, so
+ differentiable) — the principled version of DLESyM's SST NaN-interpolation
+ hack. The filler is cached per (grid signature, mask bytes). A mask with no
+ valid points raises.
+
+Fill runs **before** regridding by design: interpolating first would bleed
+invalid (e.g. land) values into valid ocean points near the coast.
+
+### 4. Spatial regrid
+
+The selection ladder, evaluated per field:
+
+1. **Identity fast path** — every spatial dim of the field exists in the
+ destination grid with an equal coordinate array (works for lat/lon and for
+ identical HEALPix `face/height/width` grids): pass through untouched. Only
+ taken when no user regridder is set.
+2. **Point target** — the destination advertises a `"point"` dim (its
+ `points=` is a `PointSet`, a scattered set of sample locations rather than
+ a mesh — stations, sites, arbitrary query coordinates; see
+ [`points`](#points-a-scattered-sample-location-grid) below). Handled
+ before the HEALPix guard since a point destination has no `face`/`lat`/
+ `lon` mesh of its own to compare against.
+3. **HEALPix guard** — a `face` dim on either side with *differing* grids and
+ no user regridder raises `IncompatibleFieldError` pointing at
+ `regridder=` (build one with `earth2grid`, as `models/px/dlesym.py` does).
+4. **User regridder** — a `regridder=` callable on the connector overrides
+ everything, including a point destination: it is applied to the trailing
+ spatial dims of any layout, and the output coords are rebuilt from the
+ destination grid. The contract:
+ `tensor[..., *src_spatial] -> tensor[..., *dst_spatial]`, operating on the
+ trailing spatial axes and preserving leading (batch/window/...) axes.
+5. **Auto bilinear** — both grids must expose 1D `lat`/`lon`, the source must
+ be *regular* (equally spaced), and the field's trailing two dims must be
+ `(lat, lon)`; otherwise `IncompatibleFieldError` with the `regridder=`
+ escape hatch. Uses `earth2studio.utils.interp.latlon_interpolation_regular`
+ (edge clamping stands in for extrapolation), built lazily and cached per
+ source grid signature.
+
+A destination with no grid of its own (`grid_coords()` is `None` — mediators)
+skips regridding entirely: fields pass through on the source grid, and
+reduction happens there.
+
+### `points`: a scattered sample-location grid
+
+A component constructed with `points=PointSet(lat=..., lon=...)` targets N
+arbitrary locations instead of a mesh — `grid_coords()` reports
+`{"point": labels}` (`labels` is `PointSet.names` if given, else an integer
+index) in place of whatever the component's own state coords happen to be.
+Delivering to it needs `Connector(..., sample="nearest" | "bilinear")`:
+
+- `"bilinear"` reuses the mesh regridder's kernel, reshaping the N
+ destination points as a degenerate `[N, 1]` mesh — one bilinearly
+ interpolated value per point. Same regularity requirement as auto bilinear
+ above.
+- `"nearest"` is a great-circle nearest-neighbor lookup via a unit-sphere
+ KDTree (the same construction as the mask filler's nearest-fill), gathering
+ one source grid cell per point.
+
+Both are built lazily and cached per `(source grid signature, point set
+signature, method)`. `sample=` and `regridder=` are mutually exclusive
+(`CouplingError`); a point destination with neither set raises `CouplingError`
+at `execute()` rather than silently picking one — as does a `"point"`-dim
+destination with no `points=` registered on it (reachable if a component
+hand-builds coords carrying a `"point"` key without going through `points=`),
+and a source without a regular 1D lat/lon grid (`IncompatibleFieldError`).
+This is the primitive downscaling-style applications need to sample a dense
+forecast (or a static context raster) down to station or site coordinates.
+
+After the pipeline the field lands in `dst.import_state`, and a copy of the
+reference is kept in `connector.last_transfer` for `driver.probe("src->dst")`.
+
+### Windowed connectors: window= and reduce=
+
+Setting `window=` and `reduce=` **together** (either alone raises
+`CouplingError`) turns the connector into a windowed reduction — the
+preferred path for simple fast→slow coupling that needs "the trailing 48 h
+mean" rather than the instantaneous field. The semantics:
+
+- **Matching is by CellMethod, not name intersection.** Each source export
+ `base` is paired with a destination import whose dictionary entry carries
+ `CellMethod(base, reduce, window)`, and the delivered Field carries that
+ *derived* standard name (`geopotential_at_1000hpa` in,
+ `geopotential_at_1000hpa_48h_mean` out). A missing derived import raises —
+ the coupler never invents names. `match()` returns both the consumed base
+ names and the delivered derived names.
+- **Every `execute` folds the source export into a running reduction**
+ (`"mean" | "sum" | "max" | "min"`; one accumulator per field, duplicate
+ `valid_time`s ignored) — the same accumulator core the mediators use.
+- **Delivery happens only at execute times aligned to `window`**; mid-window
+ the destination's previous import stands untouched. The alignment origin is
+ the `valid_time` of the first execute's source field — under lagged
+ coupling (connector before the source's run in the slot) that is the clock
+ start, so the first delivery lands exactly one window after t0, with no
+ driver hook.
+- `time_policy` does not apply on the windowed path; the spatial pipeline
+ (vertical → fill → regrid) still runs on each delivery.
+
+```python
+from earth2studio.nvcoupler.testing import fake_atmos, fake_ocean
+
+conn = nvc.Connector(fake_atmos(), fake_ocean(), window="48h", reduce="mean")
+conn.match()
+# ['geopotential_at_1000hpa', 'geopotential_at_1000hpa_48h_mean']
+```
+
+A windowed connector replaces a single-source `AccumulationMediator` plus its
+two connects and its `med.compute` slot action. Reach for a Mediator (below)
+when several sources feed one reduction or the reduction needs custom code.
+
+## Mediators
+
+A `Mediator` is the multi-source, generalized form of the windowed-reduction
+machinery — it shares the same running-accumulator core as windowed
+connectors. It sits between cadences as its own participant: its import state
+forwards every arriving field to `accumulate(field)`, and when its slot runs
+the driver calls `compute(time)` (the `med.compute` action), which must
+populate the export state.
+
+`AccumulationMediator(name, [derived_names])` implements windowed reductions:
+
+- **Running, O(1) memory.** Reductions are running torch ops — `add` for
+ mean/sum, `torch.maximum`/`torch.minimum` — so memory is one accumulator
+ per derived field regardless of window length. Mean divides by the sample
+ count at compute time. Gradients flow through mean/sum; max/min propagate to
+ the extremal sample.
+- **One base, many derived.** Several derived fields may reduce the same base
+ import (the 24 h max *and* 24 h mean of t2m accumulate from each delivered
+ t2m field); the mediator imports each base once and fans deliveries out.
+- **Duplicate dedup.** A field arriving with the same `valid_time` as the last
+ accumulated one (two connectors feeding the mediator, or a re-executed slot)
+ is ignored rather than double-counted.
+- Compute with zero samples raises `CouplingError` ("is a connector feeding
+ this mediator in a faster slot?"); after compute the accumulators clear, so
+ each window is independent (trailing, non-overlapping).
+
+`TrailingAverageMediator` is the mean-only restriction — the exact semantics
+of DLESyM's ocean forcing and PhysicsNeMo's `TrailingAverageCoupler`; non-mean
+fields raise at construction.
+
+When one source feeds one destination, prefer the
+[windowed connector](#windowed-connectors-window-and-reduce) — it produces
+the same numbers from the same accumulator core with fewer moving parts.
+
+## Coupling semantics: ordering is the coupling mode
+
+There is no `lagged=True` flag on connectors. Whether coupling is lagged
+(NUOPC-explicit) or sequential is **purely the position of the connect action
+relative to the source's run in the same slot**: a connect executed before
+the source runs delivers its previous export; after, the export just
+produced. The driver executes a slot's actions strictly in order and adds no
+hidden exchanges. (`derive_sequence`'s `lagged=` parameter is not a flag on
+the exchange — it just selects where the connect is placed.)
+
+A minimal pair, hand-checkable: `src` increments its `t2m` state by 1 each
+step; `dst` copies its imported `t2m` into its `d2m` export.
+
+```python
+lagged = """
+@6h
+ src -> dst # before dst runs: dst sees src's PREVIOUS export
+ src
+ dst
+@
+"""
+sequential = """
+@6h
+ src
+ src -> dst # after src runs: dst sees the export just produced
+ dst
+@
+"""
+```
+
+After one 6 h step from `t2m = 0`: the lagged system's `d2m` is **0.0** (the
+t0 export), the sequential system's is **1.0** (the export produced in the
+same step). Derived sequences (`Driver(sequence=None)`, `couple()`,
+`derive_sequence(lagged="all")`) always generate the lagged shape — connects
+precede runs — with one exception: mediator deliveries follow their
+`med.compute` in the same slot and are therefore sequential. `describe()`
+labels each connect `lagged` or `sequential` in its plan table (sequential
+iff the source ran or computed earlier in the same slot). Run
+[`examples/09_nvcoupler/02_lagged_vs_sequential.py`](../../../examples/09_nvcoupler/02_lagged_vs_sequential.py)
+to see the divergence over a real rollout, and see the
+[DSL reference](dsl_and_yaml_reference.md) for the full grammar.
diff --git a/earth2studio/nvcoupler/docs/design_and_roadmap.md b/earth2studio/nvcoupler/docs/design_and_roadmap.md
new file mode 100644
index 000000000..1f43716f0
--- /dev/null
+++ b/earth2studio/nvcoupler/docs/design_and_roadmap.md
@@ -0,0 +1,241 @@
+# Design and roadmap
+
+Why nvcoupler is shaped the way it is: the NUOPC inheritance, the decisions
+that diverge from it, how the implementation was verified, and what is
+honestly not done yet. For what each abstraction *means*, see
+[concepts](concepts.md); for how to use them, the [user guide](user_guide.md)
+and [API reference](api_reference.md).
+
+## Why NUOPC concepts for ML inference
+
+AI Earth-system models are coupled today in three ad-hoc ways (see prior art
+below), each of which welds the coupling decision to code that should not own
+it. The Earth-system-modeling community solved this problem once already:
+NUOPC/ESMF's component/connector/mediator/driver decomposition is thirty years
+of institutional knowledge about which seams matter. nvcoupler ports the
+concepts, not the code.
+
+**What transfers directly:**
+
+- Components with advertise/realize/initialize/run/finalize phases, so a
+ misconfigured system fails at initialize with a named, actionable error
+ rather than mid-rollout.
+- Connectors as the sole path between components, matching by a field
+ dictionary's standard names rather than model vocabularies.
+- Mediators for cadence-bridging reductions.
+- A driver executing an ordered run sequence on a shared clock, each
+ component gated by its own declared timestep (slot alignment) — which is
+ exactly what a 6 h atmosphere and a 48 h ocean need to coexist.
+- The "data component" move (prescribed forcing is just another component) and
+ the "split a monolithic executable into gridded components" move
+ ([`dlesym_split.py`](../dlesym_split.py) is the latter, applied to DLESyM).
+
+**What was deliberately dropped for v1:**
+
+- **Concurrency.** NUOPC runs components on disjoint PE layouts; nvcoupler
+ executes slot actions strictly sequentially in one process. ML inference
+ steps are GPU-bound and fast; the ordering-as-semantics model (below) also
+ *requires* deterministic sequencing. Multi-GPU concurrency is roadmap, not
+ regret.
+- **Conservative regridding and flux exchange.** ESMF couplers conserve energy
+ and mass across grids because the physics demands it. ML emulators are not
+ conservation-constrained, exchange states rather than fluxes, and v1 ships
+ bilinear interpolation only (plus a custom-`regridder=` escape hatch).
+- **Fortran-era config machinery** in favor of a small Python API, a
+ NUOPC-flavored [runSeq DSL, and YAML round-tripping](dsl_and_yaml_reference.md).
+
+**The fundamental caveat:** a physical model tolerates any dynamically
+consistent forcing; an ML model only tolerates forcing that looks like its
+training data. You can re-plumb *how* DLESyM's halves exchange SST, but you
+cannot make its atmosphere accept hourly SST, a new variable, or an
+out-of-distribution ocean and expect skill. nvcoupler makes coupling
+structure explicit and swappable; it cannot make models coupleable in ways
+they were not trained for. (Coupled fine-tuning, below, is the remedy the
+architecture is built to enable.)
+
+## Key decisions
+
+### The adapter owns the model call
+
+Alternatives considered: (a) require every model to implement a common
+coupled-model interface — rejected, because the point is to couple *existing*
+models unmodified; (b) have components translate imports into each model's
+call shape — rejected, because it multiplies component subclasses by call
+shapes. Instead a small `ImportAdapter` protocol owns the invocation, and the
+three built-ins are transcriptions of the three call shapes observed in the
+wild: state-channel overwrite (prescribed forcing), StormScope's
+`call_with_conditioning` kwarg, and the DLESyM/PhysicsNeMo extra coupling
+tensor. A new call shape is a ~20-line adapter, not a framework change.
+Corollary decision: stacking adapters demand an explicit `field_order=` for
+multiple imports, because alphabetical stacking would feed channel-permuted
+inputs that run fine and predict garbage.
+
+### Physical-units exchange
+
+Alternative: exchange in each model's normalized space, avoiding a
+denormalize/renormalize round trip per step. Rejected because normalization
+statistics are private per model — normalized exchange couples every
+component to every other's training pipeline and makes a `DataComponent`
+(observations, in physical units) a special case. The round trip through the
+same per-variable constants is exact up to float rounding (asserted for the
+DLESyM split in `test/nvcoupler/test_dlesym_split.py`). Units are checked at
+match time, not converted — a wrong-units pairing should be a loud
+configuration error, not a silent multiply (see
+[errors and troubleshooting](errors_and_troubleshooting.md)).
+
+### Slot ordering as coupling semantics
+
+Alternative: a `mode="lagged"` flag on connectors. Rejected in favor of the
+NUOPC convention that a runSeq *is* the coupling semantics: a connect before
+the source's run delivers the source's previous state (lagged), after it
+the fresh state (sequential). Derived sequences make the canonical lagged
+shape the default without giving up the mechanism — `derive_sequence`'s
+`lagged=` parameter only chooses where each connect is *placed*, never adds
+a second mechanism. One mechanism, zero redundant configuration to
+disagree with itself, and a coupling-order experiment is a one-line DSL edit
+([example 02](../../../examples/09_nvcoupler/02_lagged_vs_sequential.py)).
+`describe()` derives and displays the mode per connect so the ordering is
+never implicit knowledge.
+
+### Dictionary CellMethods over string parsing
+
+Alternative: infer "48 h mean of z1000" by parsing the suffix of
+`geopotential_at_1000hpa_48h_mean`. Rejected — name-grammar coupling is how
+lexicons rot. A derived field is a first-class `FieldEntry` carrying
+`CellMethod(base, method, window)`, which is what lets windowed
+`Connector`s, `AccumulationMediator`, and `couple()`'s windowed-connector
+synthesis operate on data, not regexes
+([concepts](concepts.md#cellmethod-derived-fields-as-first-class-dictionary-entries)).
+
+### Pure-torch exchange for training readiness
+
+Every stage of the exchange path — regrid gathers, mask-fill gathers,
+log-pressure vertical interpolation, mediator running reductions, functional
+import injection (clone + `index_copy_`) — is differentiable torch. That is a
+tax during inference (numpy would sometimes be simpler) paid for one payoff:
+`driver.rollout(n_steps)` keeps the autograd graph across the whole coupled
+system, while `run()`/`steps()` execute under `torch.inference_mode()` and
+record/IO paths detach so collection never pins graphs.
+
+```python
+import torch
+import earth2studio.nvcoupler as nvc
+from earth2studio.nvcoupler.testing import atmos_ic, fake_atmos, fake_ocean, ocean_ic
+
+gain = torch.tensor(1.0, requires_grad=True) # a shared "parameter"
+driver = nvc.couple(
+ fake_atmos(gain), fake_ocean(gain), start="2024-01-01", stop="2024-01-05"
+)
+driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+with torch.enable_grad():
+ states = driver.rollout(16) # 16 x 6 h, graph intact
+loss = states["atmos"]["geopotential_at_1000hpa"].data.mean()
+loss.backward()
+gain.grad # non-zero: crossed the exchange
+```
+
+What coupled fine-tuning enables: jointly training coupled emulators so each
+learns to tolerate the other's imperfect output — the standard remedy for
+coupled drift, previously unavailable because the exchange lived in numpy
+datapipes. Optimizer loops, truncated BPTT, and per-component device placement
+stay out of scope for v1 ([example 05](../../../examples/09_nvcoupler/05_coupled_finetuning.py)
+shows a complete training step).
+
+### Framework coupling, not datapipe or model-internal coupling
+
+The three prior-art patterns this package factors out:
+
+- **Datapipe-level** — PhysicsNeMo's `ConstantCoupler` /
+ `TrailingAverageCoupler` bake exchange policy into data loading; changing
+ the coupling means changing the datapipe, and two-way interaction is out of
+ reach. nvcoupler keeps both as one-line configurations (`time_policy=
+ "constant"`; `Connector(window=, reduce=)` or `TrailingAverageMediator`)
+ on a two-way-capable substrate.
+- **Model-internal** — `earth2studio/models/px/dlesym.py` hard-codes the
+ atmos↔ocean exchange inside `__call__`; swapping the ocean for observations
+ means forking the model. `split_dlesym` re-exposes the halves as components.
+- **Caller-owned** — StormScope's `call_with_conditioning` leaves coupling
+ entirely to user scripts; correct, but unshareable and unvalidated.
+ `ConditioningKwargAdapter` gives that pattern the same validation, cadence,
+ and regridding machinery as everything else.
+
+## Verification story
+
+The code ships with 176 passing tests (`test/nvcoupler/`; a 177th — the
+real-weights gate below — is collected but skipped), built on a deliberate
+strategy:
+
+- **Hand-computed toys.** The `testing.py` components are linear maps with
+ spatially constant ICs, so every intermediate value of a 96 h coupled run is
+ computable on paper; `test_driver.py::test_cadence_and_hand_computed_values`
+ asserts the full trajectory, and the lagged/sequential, gradient-flow, and
+ reset-reproducibility tests reuse the same closed-form system. Tests are the
+ executable specification — mediator dedup, connector history rotation, and
+ every error path in [errors_and_troubleshooting.md](errors_and_troubleshooting.md)
+ are pinned there.
+- **Seam tests** (`test_seams.py`). The historically bug-rich boundary is
+ shape conventions between component kinds: a `PrognosticComponent`'s exports
+ (model coords carry singleton batch/time/lead_time) feeding a
+ `CallableComponent`'s overwrite adapter, stacking with a `DataComponent`'s
+ fields in one import state, and driving a `DiagnosticComponent` end to end.
+- **Adversarial reviews.** The design underwent a design review and the
+ implementation an execution-verified code review (findings reproduced by
+ running code, then fixed and regression-tested). Representative bug classes
+ caught this way, each now defended in code and tests: axis renumbering after
+ `tensor.select` silently slicing the wrong dimension (see the pointed
+ helper comment in `test_dlesym_weights_equivalence.py`); connector history
+ collapsing when a slow source's export is re-seen by a fast slot (the
+ "rotate only on new valid_time" rule); in-memory records and IO writes
+ pinning autograd graphs during `rollout` (detached clones, off the exchange
+ path); zarr's default 0.0 fill letting never-written rows masquerade as
+ physical values (NaN initialization); and double-counted mediator samples on
+ duplicate deliveries.
+- **The DLESyM weights-equivalence gate — currently UNRUN.** All
+ `dlesym_split` tests run against a mock authored from a *reading* of
+ `dlesym.py`, which makes them structurally circular: a misreading
+ (normalization order, insolation anchors, window chunking) would pass every
+ mock test and fail on real weights.
+ `test_dlesym_weights_equivalence.py` is the actual proof — it drives the
+ split components through nvcoupler and asserts equality with native
+ `DLESyM.__call__` on real checkpoints — but it is gated behind
+ `NVCOUPLER_DLESYM_WEIGHTS=1` (needs physicsnemo plus the multi-GB
+ `hf://nvidia/dlesym-v1-era5` package) and **has not yet been executed
+ anywhere**. Until it passes, treat "nvcoupler can host DLESyM" as
+ structurally validated but numerically unverified.
+
+## Honest limitations and v2 roadmap
+
+Current, stated plainly:
+
+- **Real-weights validation is pending** (the gate above). Highest-priority
+ item; it either closes the loop or finds the misreading the mocks cannot.
+- **Units are checked, not converted.** No pint dependency in v1; a genuine
+ mismatch requires a converting Mediator or aligned dictionary entries.
+- **HEALPix / curvilinear sources need a user regridder.** The auto path
+ handles regular 1D lat/lon only; differing `face` grids raise with an
+ `earth2grid` pointer rather than guessing.
+- **No checkpoint/restart.** A crashed 3-month coupled run restarts from t0.
+ Restart needs serializing component internal state (model windows, mediator
+ accumulators, connector history), which the YAML layer deliberately does not
+ attempt yet.
+- **No coupled ensembles.** Batch dims flow through the exchange, but there is
+ no ensemble-aware driver API (perturbations, per-member IO).
+- **Single-process, single-device execution.** No concurrent slot execution,
+ no per-component GPU placement; components that could run in parallel
+ (independent branches of an impact chain) do not.
+- **`DiagnosticComponent` is shallow** — single-grid, all-imports-every-step;
+ windowed or multi-resolution diagnostics need mediator help.
+- **The field dictionary is package-local.** `DEFAULT_DICTIONARY` is a curated
+ v1 vocabulary; reconciling it with (or upstreaming it into)
+ `earth2studio.lexicon` — whose `normalize_units` it already mirrors — is
+ open, and drift between the two is a real risk.
+- **Long gradient rollouts hold every step's graph.** `rollout(n)` has no
+ gradient checkpointing or truncated-BPTT support; memory grows linearly in
+ n, which caps practical fine-tuning horizons.
+- **Driver IO** is in-memory xarray plus streaming `IOBackend` writes; the
+ `couple()` auto-wiring layer does not yet configure IO.
+
+The v2 order of attack roughly follows that list: run the weights gate, then
+checkpoint/restart and gradient checkpointing (they share the
+state-serialization work), then ensembles and multi-GPU placement, with unit
+conversion and lexicon reconciliation as the dictionary matures.
diff --git a/earth2studio/nvcoupler/docs/dsl_and_yaml_reference.md b/earth2studio/nvcoupler/docs/dsl_and_yaml_reference.md
new file mode 100644
index 000000000..0de9c74d1
--- /dev/null
+++ b/earth2studio/nvcoupler/docs/dsl_and_yaml_reference.md
@@ -0,0 +1,386 @@
+# Run-sequence DSL and YAML reference
+
+This page is the normative reference for the two text formats nvcoupler
+understands: the run-sequence DSL (`parse_run_sequence`, `Driver(sequence=...)`)
+and the YAML system schema (`to_yaml` / `from_yaml`). For the concepts behind
+them see [concepts.md](concepts.md); for task-oriented walkthroughs see
+[user_guide.md](user_guide.md). Every snippet below has been executed against
+the toy components in `earth2studio.nvcoupler.testing`.
+
+## The run-sequence DSL
+
+Most systems never write a sequence: `Driver(sequence=None)` (the default)
+and `couple()` derive the canonical lagged sequence from the coupling graph
+via `derive_sequence`. The DSL is the explicit override for hand-tuned
+ordering — it mirrors NUOPC's runSeq: slots opened by `@` headers,
+each containing an ordered list of actions.
+
+```
+@6h
+ atmos -> med # accumulate atmos exports into the mediator
+ ocean -> atmos # lagged: before atmos runs
+ atmos
+@48h
+ med.compute
+ med -> ocean
+ ocean
+@
+```
+
+### Grammar
+
+Line-oriented. On every line, everything from the first `#` to the end of the
+line is a comment; blank (or comment-only) lines are ignored. The remaining
+lines are:
+
+| Line form | Meaning | Regex (after comment strip) |
+|---|---|---|
+| `@` | Open a new slot with that interval | `^@(\S+)?$` |
+| `@` | Close the current slot (terminator) | same |
+| `src -> dst` | `ConnectAction(src, dst)` — transfer matched fields | `^(\w[\w.-]*)\s*->\s*(\w[\w.-]*)$` |
+| `name.phase` | `MediateAction(name, phase)` — run a mediator | `^(\w[\w-]*)\.(\w+)$` |
+| `name` | `RunAction(name)` — run a component | `^(\w[\w-]*)$` |
+
+Notes on the grammar, all verified against `parse_run_sequence`:
+
+- Actions are matched in the order connect, mediate, run — so a bare name
+ containing a dot always parses as a `MediateAction`. Component names must
+ therefore not contain dots (allowed characters: word characters and `-`).
+- The mediate phase is free-form (`med.compute` by convention; `med.finalize`
+ parses equally). The `Driver` ignores the phase at execution time and calls
+ the mediator's `run(time)` (which calls `compute`) for every `MediateAction`.
+- The trailing bare `@` terminator is what `str(RunSequence)` emits, but it is
+ optional on input: end-of-input also closes the last slot. An action line
+ appearing after a bare `@` (or before any `@`) raises
+ `SequenceError("... outside any @interval slot")` with the line number.
+- An empty sequence (no slots at all) raises `SequenceError("Run sequence is
+ empty")`.
+- An unparseable action line raises `SequenceError` naming the line and the
+ three accepted forms.
+
+### Interval strings (`as_timedelta`)
+
+Slot headers — and every timestep/window/`dt` argument in the package — are
+coerced through `earth2studio.nvcoupler.clock.as_timedelta`, which accepts:
+
+- **`np.timedelta64`** values (converted to nanosecond precision).
+- **Strings of the form ``**: an integer (optionally negative)
+ followed by a numpy timedelta64 unit code, with four convenience spellings
+ mapped first: `d → D`, `H → h`, `min → m`, `S → s`. Verified examples:
+ `"6h"`, `"12H"`, `"2D"`, `"2d"`, `"1W"`, `"90m"`, `"30min"`, `"45s"`,
+ `"500ms"`. Note `m` is minutes; `M` is calendar months, which numpy converts
+ using an average-month length — avoid `M`/`Y` for coupling intervals.
+
+Rejected:
+
+- **Bare numbers** — `as_timedelta(6)` raises
+ `ValueError: Bare number 6 is ambiguous as a timedelta (hours? steps?) —
+ pass a string like '6h' or '2D', or a np.timedelta64`. There is no implicit
+ "hours" or "steps" unit anywhere in nvcoupler.
+- Strings with no leading digits (`"h6"`) or no unit (`"6"`) raise
+ `ValueError: Cannot parse timedelta ...; expected e.g. '6h', '2D'`. Inside
+ `parse_run_sequence` this is re-raised as `SequenceError("Line N: ...")`.
+- A caveat, honestly: an unknown unit *letter* (e.g. `"6x"`) propagates
+ numpy's own `TypeError: Invalid datetime unit "x" in metadata` rather than a
+ `SequenceError` — only `ValueError`s are wrapped with the line number.
+- Negative intervals parse (`"-6h"`) but fail validation: `is_multiple`
+ requires a positive interval, so `validate()` raises `CadenceError`, and
+ `Clock` rejects a non-positive `dt` at construction.
+
+### Formal semantics
+
+Let the driver `Clock` run from `start` to `stop` in steps of `dt`. The
+execution rule, exactly as implemented by `Driver._execute_time`:
+
+1. **Slot alignment.** At each clock time `t` (the first is `start + dt`;
+ actions never execute at `start` itself — the initial conditions seed the
+ export states there), a slot with interval `I` is *aligned* iff
+ `(t - start) > 0` and `(t - start) % I == 0`. A 6 h slot on a 6 h clock is
+ aligned at every step; a 48 h slot at every 8th.
+2. **Execution order.** All aligned slots execute at `t` in the order they
+ appear in the sequence, and within a slot the actions execute strictly in
+ listed order. When a 6 h and a 48 h slot are both aligned (every 48 h), the
+ 6 h slot's actions run first because it is listed first — the order is
+ textual, never cadence-derived.
+3. **Action effects.**
+ - `RunAction(c)`: calls `c.run(t)` — the component consumes whatever is
+ currently in its import `State` and republishes its export `State` with
+ `valid_time = t`.
+ - `ConnectAction(src, dst)`: the connector copies whatever is *currently*
+ in `src.export_state` through its pipeline (time policy → vertical →
+ mask fill → regrid) into `dst.import_state`. If `src` has not yet
+ produced a matched field, this raises
+ `CouplingError("... has not produced yet — check the run
+ sequence ordering")`. A *windowed* connector (`window=`/`reduce=`)
+ behaves differently: each execute folds the source export into its
+ running reduction, and it delivers the derived field only at times
+ aligned to its window — mid-window executes leave the destination's
+ import untouched.
+ - `MediateAction(m, phase)`: calls `m.run(t)`, i.e. the mediator's
+ `compute(t)`, which turns its accumulated samples into exported derived
+ fields. (Accumulation itself is not an action — it happens as a side
+ effect of every connector delivery into the mediator.)
+
+**Lagged vs sequential is purely positional.** A connect transfers the source
+state *as of the moment the connect executes*:
+
+- Connect placed **before** the source's run at the same time (or in a slot
+ where the source does not run at all this time) delivers the source's
+ *previous* export — **lagged** (NUOPC-explicit) coupling. In the example
+ above, `ocean -> atmos` before `atmos` delivers the ocean state from its
+ last 48 h ring.
+- Connect placed **after** the source's run (or mediator's compute) at the
+ same time delivers the export *just produced* — **sequential** coupling.
+ `med -> ocean` after `med.compute` hands the ocean the freshly reduced
+ 48 h mean.
+
+Swapping the two flavors is a one-line reorder; see
+`examples/09_nvcoupler/02_lagged_vs_sequential.py`
+([examples README](../../../examples/09_nvcoupler/README.rst)). `describe()`
+labels a connect's mode with exactly this source-relative rule: `sequential`
+iff the source ran (or the mediator computed) earlier in the same slot, else
+`lagged`.
+
+### Validation
+
+`RunSequence.validate(components, dt)` runs inside `Driver.initialize()`
+(after parsing, before anything executes) and performs, in order:
+
+1. **Slot cadence**: every slot interval must be a positive whole multiple of
+ the driver `dt`, else `CadenceError` (`@7h` on a 6 h clock fails).
+2. Per action:
+ - `RunAction`: the component name must be a key of `components` (unknown
+ names raise `SequenceError` with a did-you-mean suggestion, e.g.
+ `atmoss` → `Did you mean: 'atmos'?`), **and** the component's `timestep`
+ must equal the slot interval exactly — a 48 h ocean listed in a `@6h`
+ slot raises `CadenceError`.
+ - `ConnectAction`: both endpoint names must resolve (same suggestion
+ machinery). No cadence constraint — connects may sit in any slot.
+ - `MediateAction`: the mediator name must resolve, and its `timestep`
+ (its window) must equal the slot interval, exactly like a `RunAction`
+ (`med.compute` for a 48 h mediator in a `@6h` slot raises
+ `CadenceError`). The phase string is not validated.
+3. **Completeness**: every component in `components` must appear in some
+ `RunAction` or `MediateAction`; idle components raise
+ `SequenceError("Components never run by the sequence: [...] — add a
+ RunAction (bare component name) to a slot matching their timestep")`.
+
+`Driver.initialize()` layers further checks on top (connector field matching
+and units, unfed-import detection, unconsumed-export warnings) — see
+[errors_and_troubleshooting.md](errors_and_troubleshooting.md).
+
+### Round-trip guarantees
+
+`str(RunSequence)` emits valid DSL and `parse_run_sequence(str(seq)) == seq`
+(actions are frozen dataclasses, so equality is structural). Formatting rules:
+
+- Intervals that are a whole number of hours print NUOPC-style as hours —
+ including multi-day ones: a slot built with `np.timedelta64(2, "D")` prints
+ as `@48h`. The parsed interval is identical either way.
+- Sub-hourly intervals fall back to `fmt_timedelta`, so `@90m` and `@30m`
+ survive round-trips exactly (this was a real regression: an earlier
+ formatter truncated `@90m` to `@1h` through YAML round-trips; the fix is
+ pinned by `test_str_preserves_subhour_intervals`).
+- Output is always two-space-indented actions and a final bare `@`; comments
+ are not preserved (they are stripped at parse time).
+
+For hand-written sequences the YAML `sequence` key stores
+`str(driver.sequence)` verbatim, so these guarantees are exactly what makes
+those YAML round-trips faithful. Derived sequences are stored as
+`{derived: true, text: ...}` and re-derived on load instead (see below).
+
+## The YAML schema
+
+`to_yaml(driver)` serializes a `Driver` to a small YAML document;
+`from_yaml(text_or_path)` rebuilds an *uninitialized* driver from it (call
+`driver.initialize(ics)` afterwards as usual — initial conditions are tensors
+and are never serialized).
+
+### Top-level keys
+
+| Key | Required | Type | Meaning |
+|---|---|---|---|
+| `clock` | yes | mapping `{start, stop, dt}` | ISO-8601 `start`/`stop` strings, `dt` an interval string. Rebuilt as `Clock(start, stop, dt)`. |
+| `sequence` | yes | string (literal block) or mapping | Hand-written sequences: the run-sequence DSL, verbatim (`str(driver.sequence)`). Derived sequences (`driver.sequence_derived`): `{derived: true, text: }` — the `text` is informational; `from_yaml` re-derives the schedule from components + connectors (deterministic, so round-trips reproduce identical runs). A mapping without `derived: true` raises `CouplingError`. |
+| `components` | yes | mapping `name -> {class, kwargs}` | `class` is a dotted import path to a module-level class or factory; it is imported and called as `factory(**kwargs)`. |
+| `dictionary` | no | list of entry mappings | Only `FieldEntry` items **absent from or differing from** `DEFAULT_DICTIONARY`. Each has `standard_name`, `canonical_units`, `description`, `aliases` (list), and optional `cell_method: {base, method, window}`. |
+| `aliases` | no | mapping `alias -> standard_name` | Alias additions relative to the default dictionary (see below). |
+| `connectors` | no | list of `{src, dst, time_policy, fill, fields?, sample?, window?, reduce?}` | Connector settings. `fields` appears only when the connector was built with an explicit list; `time_policy` defaults to `"constant"` and `fill` to `"none"` on load. `sample` (`"nearest"` or `"bilinear"`) appears only when set — a plain string, it round-trips like `time_policy`/`fill`. `window`/`reduce` appear (together) for windowed connectors and rebuild them on load. |
+
+`from_yaml` raises `CouplingError` when the document is not a mapping, when
+any of `clock`/`sequence`/`components` is missing, when a component spec lacks
+`class`, when the import path cannot be resolved (the error names the module
+and attribute), when the factory call itself fails (wrapped with the class
+path and kwargs), or when a connector endpoint is not a configured component.
+A string argument containing no newline that names an existing file is read as
+a file path; anything else is parsed as YAML text.
+
+If a `dictionary` and/or `aliases` section is present, `from_yaml` builds one
+extended `FieldDictionary` (a copy of the default plus the entries and
+aliases) and passes it as the `dictionary` kwarg to every component factory
+that accepts one (inspected via its signature, `**kwargs` counts) and whose
+`kwargs` don't already set it.
+
+### Annotated example
+
+This exact document is `to_yaml` output for the toy atmos–ocean–mediator
+system, and executing `from_yaml(to_yaml(driver))` reproduces the original
+run bit-for-bit (same xarray outputs; the assertion is in the scratch check
+for this page and in `test/nvcoupler/test_config.py::test_round_trip_identical_outputs`):
+
+```yaml
+clock:
+ start: '2024-01-01T00:00:00' # ISO-8601, second precision
+ stop: '2024-01-05T00:00:00'
+ dt: 6h # interval string (as_timedelta form)
+sequence: |- # the run-sequence DSL, verbatim
+ @6h
+ atmos -> med
+ ocean -> atmos
+ atmos
+ @48h
+ med.compute
+ med -> ocean
+ ocean
+ @
+components:
+ atmos:
+ class: earth2studio.nvcoupler.testing.fake_atmos # module-level factory
+ kwargs:
+ gain: 1.0
+ timestep: 6h
+ ocean:
+ class: earth2studio.nvcoupler.testing.fake_ocean
+ kwargs:
+ gain: 1.0
+ timestep: 48h
+ med: # AccumulationMediator: auto-serialized
+ class: earth2studio.nvcoupler.mediator.TrailingAverageMediator
+ kwargs:
+ name: med
+ fields:
+ - geopotential_at_1000hpa_48h_mean
+ window: 2D # 48 h, printed in fmt_timedelta's day form
+```
+
+The Python side, executed to verify:
+
+```python
+import earth2studio.nvcoupler as nvc
+from earth2studio.nvcoupler.testing import atmos_ic, fake_atmos, fake_ocean, ocean_ic
+
+atmos = fake_atmos(gain=1.0)
+atmos.yaml_spec = {
+ "class": "earth2studio.nvcoupler.testing.fake_atmos",
+ "kwargs": {"gain": 1.0, "timestep": "6h"},
+}
+ocean = fake_ocean(gain=1.0)
+ocean.yaml_spec = {
+ "class": "earth2studio.nvcoupler.testing.fake_ocean",
+ "kwargs": {"gain": 1.0, "timestep": "48h"},
+}
+med = nvc.TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"])
+driver = nvc.Driver(
+ {"atmos": atmos, "ocean": ocean, "med": med},
+ "@6h\n atmos -> med\n ocean -> atmos\n atmos\n@48h\n med.compute\n med -> ocean\n ocean\n@",
+ nvc.Clock("2024-01-01", "2024-01-05", "6h"),
+)
+rebuilt = nvc.from_yaml(nvc.to_yaml(driver)) # round-trip
+rebuilt.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+datasets = rebuilt.run() # identical to driver.run()
+```
+
+### Derived sequences and windowed connectors, executed
+
+A driver built declaratively (no `sequence=`) serializes its schedule as the
+mapping form, and windowed connectors carry `window`/`reduce`; `from_yaml`
+re-derives the sequence and rebuilds the windowed connector, so the
+round-trip reproduces the run exactly (pinned in
+`test/nvcoupler/test_config.py`):
+
+```python
+atmos2, ocean2 = fake_atmos(gain=1.0), fake_ocean(gain=1.0)
+atmos2.yaml_spec, ocean2.yaml_spec = atmos.yaml_spec, ocean.yaml_spec
+declared = nvc.Driver(
+ {"atmos": atmos2, "ocean": ocean2},
+ clock=nvc.Clock("2024-01-01", "2024-01-05", "6h"),
+ connectors=[
+ ("ocean", "atmos"),
+ nvc.Connector(atmos2, ocean2, window="48h", reduce="mean"),
+ ],
+)
+text = nvc.to_yaml(declared)
+assert "derived: true" in text and "window: 2D" in text and "reduce: mean" in text
+rebuilt2 = nvc.from_yaml(text)
+assert rebuilt2.sequence_derived
+assert str(rebuilt2.sequence) == str(declared.sequence)
+```
+
+### Serialization rules
+
+`to_yaml` decides per component, in order:
+
+1. **`yaml_spec` attribute wins.** If the component carries a `yaml_spec`
+ attribute, it must be a dict with a `"class"` key (dotted import path) and
+ optional `"kwargs"`; anything else raises `CouplingError`. The kwargs are
+ sanitized: `np.timedelta64` → interval string, `np.datetime64` → ISO
+ string, numpy scalars → Python scalars, arrays/sets/tuples → lists, dicts
+ recursively. This is the escape hatch for `CallableComponent`s and anything
+ wrapping Python state — the contract is simply *"calling
+ `(**kwargs)` rebuilds an equivalent component"*. Note `from_yaml`
+ does **not** re-attach `yaml_spec` to rebuilt components; re-tag them if
+ you intend to serialize again.
+2. **`AccumulationMediator` subclasses auto-serialize** (this includes
+ `TrailingAverageMediator`): class path plus
+ `{name, fields: , window: }` are recovered
+ from the instance itself — no `yaml_spec` needed.
+3. **Everything else raises `CouplingError`**, with the honest explanation:
+ the component wraps a closure or model object YAML cannot reconstruct; set
+ `yaml_spec` or build the system in Python. Model components referenced by
+ load paths (e.g. `{load: 'earth2studio.models.px.Persistence'}`) are
+ explicitly out of scope for v1.
+
+Not serialized, honestly stated: connector `regridder=` callables (a rebuilt
+connector falls back to the auto lat/lon path — HEALPix/curvilinear systems
+need Python construction); `io=` backends; `collect`; and
+`allow_unfed_imports`. Custom `Connector` subclasses lose their type: only
+`src`/`dst`/`fields`/`time_policy`/`fill`/`sample`/`window`/`reduce`
+round-trip. A destination's `points=PointSet(...)` is a *component* kwarg,
+not a connector setting — it round-trips the same way `model=` does: only
+through that component's own `yaml_spec` (`PointSet` isn't itself YAML-safe
+out of the box, being a dataclass of numpy arrays; a component author wanting
+one reconstructed needs a `yaml_spec["kwargs"]` shaped for it, same as any
+other non-primitive component kwarg).
+
+### The aliases delta mechanism
+
+The `dictionary` section only contains entries that differ from the built-in
+default — but an alias added with `FieldDictionary.add_alias("std", "alias")`
+*after* registration leaves the `FieldEntry` itself equal to the default and
+would be silently lost. The `aliases` key exists to carry exactly that delta:
+`to_yaml` walks every component dictionary and emits each
+`alias -> standard_name` pair that is neither in the default dictionary nor
+explained by the component's own `variable_aliases` kwarg (those are rebuilt
+from the component spec). If two components map the same alias to different
+standard names, `to_yaml` raises `CouplingError` asking you to make the alias
+consistent before serializing. On load, the pairs are re-applied with
+`add_alias` to the shared rebuilt dictionary.
+
+### Window and time string forms
+
+All interval-valued fields in the document — `clock.dt`, mediator `window`
+kwargs, `cell_method.window` — are emitted by `fmt_timedelta`: whole days as
+`"2D"`, else whole hours as `"48h"`, else whole minutes as `"90m"`, else the
+raw numpy repr. All are read back by `as_timedelta`, so the day/hour spelling
+is cosmetic (`window: 2D` ≡ `window: 48h`; verified equal after parsing).
+Clock `start`/`stop` are ISO-8601 strings at second precision, parsed by
+`np.datetime64`. The `sequence` block follows the DSL's own round-trip rules
+above, including the sub-hourly `@90m` guarantee.
+
+## See also
+
+- [concepts.md](concepts.md) — why sequences and slots look the way they do
+- [user_guide.md](user_guide.md) — building systems in Python vs YAML
+- [api_reference.md](api_reference.md) — exact signatures for `parse_run_sequence`, `to_yaml`, `from_yaml`
+- [errors_and_troubleshooting.md](errors_and_troubleshooting.md) — the full error catalogue
+- [design_and_roadmap.md](design_and_roadmap.md) — planned schema extensions (checkpointed models, IO)
diff --git a/earth2studio/nvcoupler/docs/errors_and_troubleshooting.md b/earth2studio/nvcoupler/docs/errors_and_troubleshooting.md
new file mode 100644
index 000000000..64bf30be6
--- /dev/null
+++ b/earth2studio/nvcoupler/docs/errors_and_troubleshooting.md
@@ -0,0 +1,498 @@
+# Errors and troubleshooting
+
+Every configuration error in nvcoupler derives from `CouplingError`
+(`earth2studio/nvcoupler/errors.py`) and is designed to fire at
+`Driver.initialize()` — not mid-rollout — naming the components, the field,
+and the concrete fix. This page catalogs every exception class with its
+actual raise sites, a representative message (quoted fragments are exact, so
+you can grep for them), root causes, and fixes; a
+[troubleshooting section](#troubleshooting-non-error-failure-modes) covers
+the failure modes that are warnings or silent by design. Recipes for doing
+things right the first time are in the [user guide](user_guide.md); the
+class-by-class API is in the [API reference](api_reference.md).
+
+The hierarchy: `UnknownFieldError`,
+`UnmatchedImportError`, `UnitsMismatchError`, `IncompatibleFieldError`,
+`VerticalMismatchError`, `CadenceError`, `AmbiguousCouplingError`, and
+`SequenceError` all subclass `CouplingError`; a number of checks raise plain
+`CouplingError` directly.
+
+## UnknownFieldError
+
+A name could not be resolved in the field dictionary.
+
+Raised from:
+
+- `FieldDictionary.resolve()` / `standard_name()` — any lookup of an
+ unregistered name, including component `imports=`/`exports=` lists at
+ construction time and `DiagnosticComponent`'s default import resolution.
+- `FieldDictionary.add_alias()` — aliasing to a standard name that does not
+ exist.
+- `State.from_tensor(strict=True)` — a `variable` coordinate value the
+ dictionary does not know.
+
+Message:
+
+```text
+Field name 'sea_surface_temp' is not a registered standard name or alias.
+Did you mean: 'sea_surface_temperature', 'surface_pressure'? Register it with
+FieldDictionary.register(FieldEntry(...)) or add an alias with
+FieldDictionary.add_alias(...).
+```
+
+Root causes: a typo (the did-you-mean list usually nails it), or a genuinely
+new field. Fix: correct the spelling, or register a
+`FieldEntry(standard_name, canonical_units, ...)` — for model-vocabulary
+names, prefer `variable_aliases={"raw": "standard_name"}` on the component,
+which registers the alias for you. Lookups are case-sensitive.
+
+## UnmatchedImportError
+
+A component advertises an import that nothing delivers. Two distinct raise
+sites with the same message shape:
+
+1. **`couple()` auto-wiring** (`api.couple`): no component exports the
+ imported field, and it is not a derived (CellMethod) entry whose base
+ field someone exports.
+2. **`Driver._check_unfed_imports` at `initialize()`**: components *do*
+ export the field, but no connector in *your* run sequence delivers it —
+ you forgot a `src -> dst` line. Without this check the component would
+ silently run the whole simulation on its stale initial-condition forcing.
+
+Message (case 2 — note the exporter is listed, which tells you the connect
+line to add):
+
+```text
+Component 'atmos' imports 'sea_surface_temperature' but no component exports
+it. Did you mean: 'sea_surface_temperature'? Available exports: atmos exports
+geopotential_at_1000hpa; ocean exports sea_surface_temperature; med exports
+geopotential_at_1000hpa_48h_mean. Add an alias, a Mediator producing the
+derived field, or a DataComponent supplying it from a data source.
+```
+
+Fixes, in order of likelihood: add the missing `(src, dst)` connection (or
+`src -> dst` sequence line); add an alias so the exporter's name resolves to
+the same standard name; register a CellMethod entry (so `couple()` can
+synthesize a windowed connector) or build a windowed
+`Connector(window=, reduce=)`/mediator explicitly; supply the field from a
+`DataComponent`.
+Deliberately unfed imports are opt-in via
+`Driver(..., allow_unfed_imports=True)`, which downgrades case 2 to a
+warning (grep `no connector in the run sequence delivers it`).
+
+## UnitsMismatchError
+
+Raised from `FieldDictionary.check_units()`, called by `Connector.match()`
+for every matched field — so it fires at `initialize()`.
+
+```text
+Field 'sea_surface_temperature': 'ocean' exports units 'degC' but 'atmos'
+expects 'K'. Unit conversion is not performed in v1 — convert in a Mediator
+or align the FieldDictionary entries.
+```
+
+Root cause: the two components' dictionaries carry different
+`canonical_units` for the same standard name (units are normalized before
+comparison — `m s**-1`, `m s^-1`, and `M/S` compare equal, as do
+`degC`/`celsius` and `(0-1)`/`dimensionless`). nvcoupler checks units, never
+converts values. Fix: make the entries agree, or insert a converting
+Mediator. This is an honest v1 limitation (no pint).
+
+## IncompatibleFieldError
+
+A connector could not reconcile matched fields. Raise sites in
+`connector.py`:
+
+- `match()` — the explicit `fields=[...]` list contains names not in both
+ endpoints: `fields ['air_temperature_2m'] are not in both 'ocean' exports
+ (...) and 'atmos' imports (...)`.
+- `match()` — nothing matches at all: `Connector atmos->atmos: no fields
+ match — 'atmos' exports [...], 'atmos' imports [...]`. Usually an alias or
+ advertisement problem, or a connector between the wrong pair.
+- `_build_latlon_regridder` — `Auto regrid requires a regular 1D source
+ lat/lon grid; pass a custom regridder=...` for curvilinear/unstructured
+ sources.
+- `_apply_regrid` — `source and destination HEALPix 'face' grids differ —
+ pass a custom regridder=` (identical face grids pass through as identity;
+ differing ones need e.g. an `earth2grid`-built callable, see
+ `models/px/dlesym.py`).
+- `_apply_regrid` — `auto regrid needs lat/lon on both grids` when either
+ side lacks lat/lon spatial dims, and `must have (lat, lon) as trailing
+ dims` when the field's dim order puts something after lon.
+- `_build_mask_filler` — `Mask fill impossible: no valid source points`
+ (`fill="nearest"` with an all-False mask).
+- `_apply_sample` / `_build_point_sampler` — `auto sample needs lat/lon on
+ the source grid` when the source lacks lat/lon spatial dims, and `must
+ have (lat, lon) as trailing dims` (same as the mesh path); `Auto sample
+ requires a regular 1D source lat/lon grid; pass a custom regridder=...`
+ for curvilinear/unstructured sources — see the "Point sampling" group
+ under `CouplingError (direct raises)` below for the rest of this path's
+ errors.
+
+Fixes: pass `regridder=` on the Connector for anything the bilinear
+regular-lat/lon kernel cannot handle; reorder dims so lat/lon trail; check
+the fields list against `src.advertise()` / `dst.advertise()`.
+
+## VerticalMismatchError
+
+Source and destination vertical coordinates cannot be reconciled. Only
+components declaring `import_vertical`/`export_vertical` (fields with a
+`level` dim) ever see these. Raise sites:
+
+In `connector._apply_vertical`:
+
+- destination declares `import_vertical` but the incoming field has no
+ `vertical` metadata: `'chem' expects 'ozone_mixing_ratio' on
+ PressureLevels(...), but the source field has no vertical coordinate` —
+ add `export_vertical={...}` on the source.
+- destination wants something other than pressure levels: `only
+ interpolation onto PressureLevels is supported in v1`.
+- hybrid source without surface pressure in the source's exports:
+ `hybrid->pressure interpolation of 'ozone_mixing_ratio' needs
+ 'surface_pressure' in 'met' exports — add it to the source's export list`.
+
+In `vertical.py` (`interp_to_pressure` / `_log_source_pressure`):
+
+- `coords have no 'level' dim` — you called the interpolation on a field
+ without a level axis.
+- data `level` coordinate does not match the declared `PressureLevels`
+ source: `does not match the declared PressureLevels source ... Reorder the
+ data so levels increase top to bottom, or fix the source component's
+ export_vertical declaration`.
+- `'level' coord length N != tensor level size M`.
+- `Non-positive pressure from hybrid coefficients`, and `Hybrid levels
+ a + b * ps are not strictly increasing along the level axis` — unphysical
+ ps values or misordered coefficients at interpolation time.
+
+Related `ValueError`s at construction: `PressureLevels` rejects
+non-increasing levels; `HybridLevels` rejects `a`/`b` of unequal length and
+coefficients that produce `non-increasing pressures` anywhere in the
+plausible surface-pressure range [50000, 110000] Pa. Order everything top to
+bottom (increasing pressure).
+
+## CadenceError
+
+A time interval does not divide cleanly. All three raise sites share the
+suffix `is not a positive multiple of the driver clock dt ... Choose a
+driver dt that divides every component timestep (typically their GCD).`
+
+- `Component.realize()` — component timestep vs clock dt (a 5 h component on
+ a 6 h clock).
+- `RunSequence.validate()` — a slot interval vs clock dt, and a component or
+ mediator scheduled in a slot that does not equal its own timestep:
+ `Component 'atmos' (timestep 21600000000000 nanoseconds) scheduled in a
+ slot of interval 43200000000000 nanoseconds ...`.
+- `Clock.__init__` — `Clock span (stop - start)` not a multiple of dt.
+
+Fix: pick a driver dt that divides every component timestep (their GCD —
+what `couple()` does by default), put each component's `RunAction` in the
+slot matching its exact timestep, and make `stop - start` a whole number of
+dt steps. Related `ValueError`s: `Clock dt must be positive`, `Clock stop ...
+must be after start`, and — from
+`as_timedelta` — `Bare number 6 is ambiguous as a timedelta (hours? steps?) —
+pass a string like '6h' or '2D', or a np.timedelta64`.
+
+## AmbiguousCouplingError
+
+Only `couple()` raises this (`api.couple`), in two spots: an imported field
+is exported by more than one component, or a derived field's *base* has
+multiple exporters.
+
+```text
+Import 'sea_surface_temperature' of component 'atmos' is exported by
+multiple components: ocean, ocean2. Auto-wiring cannot choose — build the
+Driver explicitly with Connector(src, dst, fields=[...]) or a run-sequence
+DSL.
+```
+
+Fix: exactly what it says — auto-wiring refuses to guess, so hand-build the
+Driver with explicit connectors and a DSL when several components export the
+same field.
+
+## SequenceError
+
+The run sequence references unknown names or is malformed. Raise sites:
+
+`derive_sequence()` (fires at `Driver` construction when no sequence is
+passed, and inside `couple()`):
+
+- `Coupling graph references unknown connection source 'atmso'. Did you
+ mean: 'atmos'? ...` — a connector endpoint is not a component name.
+- `Sequential coupling cycle among components [...] at cadence 6h:
+ connections [...] require each destination to see state its source
+ produces in the same step, so no run order exists. Mark one edge lagged
+ (e.g. lagged={...}) or pass an explicit run sequence` — only reachable
+ with `derive_sequence(lagged=)`; the default `lagged="all"` cannot
+ cycle.
+- `Cannot derive a run sequence from an empty component dict`.
+
+`RunSequence.validate()` (fires inside `Driver.initialize`):
+
+- `Run sequence references unknown component 'atmso'. Did you mean:
+ 'atmos'? Known components: [...]` — the DSL name does not match a key of
+ the `components` dict (also raised for connector endpoints and mediators).
+- `Components never run by the sequence: ['ocean'] — add a RunAction (bare
+ component name) to a slot matching their timestep` — every component,
+ mediators included, must run somewhere.
+
+`parse_run_sequence()`:
+
+- `Line 1: action 'atmos' outside any @interval slot` — actions must follow
+ an `@6h`-style header.
+- `Line 2: cannot parse 'atmos ->' — expected 'name', 'src -> dst', or
+ 'mediator.compute'`.
+- `Line N: Cannot parse timedelta ...` — a bad `@interval` header.
+- `Run sequence is empty` — no slots at all (comments only, or a stray
+ string).
+
+## CouplingError (direct raises)
+
+The base class is also raised directly for lifecycle and shape problems.
+The complete set, grouped:
+
+**Driver lifecycle** (`driver.py`):
+
+- `Driver needs a clock — Driver(components, sequence, clock) or
+ Driver(components, clock=Clock(start, stop, dt), connectors=[...])` —
+ `clock` is keyword-optional in the signature but always required.
+- `Connection ('atmso', 'ocean'): ['atmso'] are not component names; known
+ components: [...]` — a bare `(src, dst)` connector tuple names an unknown
+ component.
+- `Component 'atmos' needs an initial condition — pass ics={'atmos': (x,
+ coords)}` — every `requires_ic` component (Prognostic, Callable) needs an
+ ics entry; mediators, DataComponents, and DiagnosticComponents do not.
+- `Driver.initialize(ics) must be called before running` — also after
+ `reset()`, which deliberately invalidates initialization.
+- `Driver clock exhausted: already at stop time ... Call driver.reset() and
+ driver.initialize(ics) to run again` — `run()`/`steps()`/`rollout()` on a
+ finished driver.
+- `rollout(5) ran past the clock stop time ...: only 4 steps remained. Use
+ n_steps <= clock.n_steps (4) ...` — mid-rollout exhaustion.
+- `io= keys ['atmso'] are not component names; known components: [...]`.
+- `Component 'x' export '...' carries a 'time' dimension of size 2; the
+ driver records one snapshot per ring and can only absorb a size-1 'time'
+ dim — publish a single time-slice per run` — recording/IO cannot absorb
+ multi-time exports.
+
+**Connector ordering** (`connector.py`):
+
+- `Connector med->ocean: 'med' has not produced
+ 'geopotential_at_1000hpa_48h_mean' yet — check the run sequence ordering`
+ — a connect scheduled before its source ever ran/computed (classic:
+ `med -> ocean` placed before `med.compute`).
+
+**Windowed connectors** (`connector.py`):
+
+- `Connector atmos->ocean: window= and reduce= must be set together — a
+ windowed reduction needs both the window length and the reduction method`
+ — one of the two was passed alone.
+- `... unsupported reduce='median'; choose 'mean', 'sum', 'max' or 'min'`.
+- `Connector ocean->atmos: window='2D'/reduce='mean' is set but 'atmos'
+ imports no derived field for [...] — register a
+ FieldEntry(cell_method=CellMethod(base, 'mean', window='2D')) in the
+ destination's dictionary and add its standard name to 'atmos''s imports`
+ — the destination has no dictionary entry deriving from the source export
+ with that exact method and window (a window mismatch, e.g. 24h vs the
+ entry's 48h, fails the same way); the coupler never invents derived names.
+
+**Point sampling** (`connector.py`):
+
+- `Connector atmos->stations: sample= and regridder= are mutually exclusive
+ — pass one or the other`.
+- `... unsupported sample='linear'; choose 'nearest' or 'bilinear'`.
+- `Connector atmos->stations: destination 'stations' is a point target (a
+ scattered sample-location grid) but this connector has neither sample=
+ nor regridder= set — pass sample='nearest' or sample='bilinear', or a
+ custom regridder= for non-lat/lon sources` — the destination advertises a
+ `"point"` dim (`grid_coords()` has a `"point"` key) and neither delivery
+ path was configured; the coupler does not guess a default.
+- `Connector atmos->stations: destination 'stations' advertises a 'point'
+ dim but has no points= location metadata set — construct it with
+ points=PointSet(lat=..., lon=...)` — a `"point"` dim showed up in
+ `grid_coords()` without the component's `points=` being set (only
+ reachable by hand-building a component's coords with a `"point"` key
+ directly, bypassing `points=`).
+
+**Import adapters / Exchange** (`component.py`):
+
+- `Exchange.inject requires a 'variable' dim in the model state
+ coords; use ConditioningKwargAdapter or ExtraTensorAdapter ...`.
+- `Imported field 'air_temperature_2m' (model name '...') is not a state
+ variable of the model (variables: [...]). If the model takes forcing as a
+ conditioning kwarg or an extra tensor, pass the matching ImportAdapter.`
+- `ExtraTensorAdapter: 2 imported fields but no field_order= given — the
+ model's channel order cannot be inferred. Pass field_order=[...]` — see
+ [troubleshooting](#forgotten-field_order) below.
+- `... field_order names [...] are not in the import state (present: ...)`.
+- `Imported field with shape ... has more dims than the model state slice`.
+
+**Pull-pattern coupling** (`pull.py`):
+
+- `Pull-coupled model requested variable 'msl', but the import state holds
+ ['air_temperature_2m', 'eastward_wind_10m'] (raw-name map: {}). Add the
+ field to the component's imports and wire a connector delivering it before
+ the model runs.` — the model's internal `fetch_data` asked the
+ `StateDataSource` for a variable resolving to nothing in the import State
+ (not a held standard name, not in the raw-name map, and no dictionary alias
+ naming a held field). Fix as it says: add the standard name to `imports=`
+ and wire a connector delivering it — plus a `variable_aliases`/dictionary
+ entry when the model's raw name is unknown.
+- `Pull for 'eastward_wind_10m' at ['2024-01-01T06:00:00.000000000'] but the
+ served field is valid at 2024-01-01 — check the run-sequence ordering (the
+ connector must run before the pulling component in the same slot)` — only
+ with `strict_time=True`: the model pulled a time that does not match the
+ served field's `valid_time`, i.e. stale forcing. Classic cause: lagged
+ ordering (the connect placed before the source's run); make it sequential —
+ source run, then connect, then the pulling component, in one slot.
+- `PullAdapter: model 'NoPull' has no attribute 'conditioning_data_source' —
+ this adapter is for models that fetch forcing from a settable data source
+ (e.g. StormCast's conditioning_data_source). For models taking conditioning
+ as an argument use ConditioningKwargAdapter.` — no settable data-source
+ attribute under the configured name. Pass `PullAdapter(attribute="...")`
+ if it lives elsewhere, or use `ConditioningKwargAdapter` /
+ `ExtraTensorAdapter` for argument-style conditioning.
+- `StateDataSource serves (lat, lon) fields; got dims [...]` — the delivered
+ fields are not exchange-shaped `(lat, lon)`.
+
+**Component phases** (`component.py`):
+
+- `Component 'atmos' not initialized` — `run()` before `initialize()`;
+ `Component 'atmos' not realized` — `should_run()` before `realize()` (the
+ Driver sequences these for you; you only hit them driving components by
+ hand).
+- `Component 'c' advertises export 'sea_surface_temperature' but its output
+ variables are [...]` — the model output has no variable resolving to an
+ advertised export; check exports/aliases against the model's actual output
+ coords.
+- `Component 'prog': model outputs 1 lead times but takes 2 as input —
+ supply a next_input hook to manage the sliding input window`.
+- `DataComponent 'ocean' cannot fetch initial data before realize(clock)`.
+- `DiagnosticComponent 'diag' is missing imports [...] at ... — check the
+ run sequence connects its source before this component runs`.
+
+**Mediators** (`mediator.py`):
+
+- `AccumulationMediator 'med': 'sea_surface_temperature' has no cell_method
+ in the field dictionary — register a
+ FieldEntry(cell_method=CellMethod(base, method, window)) ...`.
+- `... fields have differing windows [...]; split them across mediators or
+ pass window= explicitly`.
+- `TrailingAverageMediator 'med': fields [...] are not mean reductions — use
+ AccumulationMediator`.
+- `Mediator 'med': no samples of 'geopotential_at_1000hpa' accumulated
+ before compute at ... — is a connector feeding this mediator in a faster
+ slot?` — the `atmos -> med` connect must live in a slot faster than
+ `med.compute`.
+
+**Field/State invariants** (`field.py`): a `Field` must not carry a
+`variable` dim (`use State.from_tensor to split`), data dims must match its
+coords, `State` keys must equal the field's standard name, and
+`State.as_tensor` with no fields raises. A missing key raises `KeyError:
+State '...' has no field '...'; present: [...]`; `Driver.probe` with an
+unknown name raises `KeyError: No connector '...'; have [...]`.
+
+**YAML** (`config.py`): `to_yaml` raises `Component 'atmos'
+(CallableComponent) is not serializable ... Set a yaml_spec attribute on the
+component — a dict {'class': ..., 'kwargs': {...}} ...`; `from_yaml` raises
+for missing top-level keys (`YAML config is missing required keys`), bad
+import paths (`Cannot import module ...`, `... has no attribute ...`),
+malformed component specs, connector endpoints that are not configured
+components, a `sequence` mapping without `derived: true` (`YAML 'sequence'
+must be run-sequence DSL text, or {derived: true, text: ...} for a
+graph-derived sequence`), and wraps any constructor failure as
+`Component 'x': (**kwargs) failed: ...`.
+
+**DLESyM split** (`dlesym_split.py`): layout/window checks on the real-model
+split — `expects the DLESyM input layout [...]`, `initial condition
+lead_time window ... does not match DLESyM full_input_times`, `atmos output
+times (N) do not chunk evenly into M ocean windows`, and the ocean run
+needing its imports first (`schedule the atmos component and the atmos ->
+ocean connector earlier in the same slot`). An honest caveat: the DLESyM
+real-weights equivalence gate (`test_dlesym_weights_equivalence.py`) is
+skipped unless `NVCOUPLER_DLESYM_WEIGHTS=1` is set with physicsnemo and the
+checkpoints available, and has not been run here — the split is verified
+against mock-model tests only.
+
+## Troubleshooting: non-error failure modes
+
+Things that used to fail silently and are now loud, plus the ones that remain
+warnings or intentional behavior.
+
+### Forgotten field_order
+
+Multi-import `ConditioningKwargAdapter`/`ExtraTensorAdapter` **now raise**
+instead of stacking alphabetically: silently permuted channels run fine and
+predict garbage, which is the worst failure mode in ML coupling. With one
+import no order is needed; with more, pass
+`field_order=["field_a", "field_b", ...]` matching the model's channel order.
+
+### Stale-IC forcing
+
+Forgetting a connect line used to mean the destination ran the whole
+simulation on its t0 import — plausible-looking, subtly wrong output. This
+**now raises `UnmatchedImportError` at initialize** (see above). If a
+free-running component is what you want, pass
+`Driver(..., allow_unfed_imports=True)` and you get the explicit warning
+`... it will run on stale initial-condition forcing` instead.
+
+### Linear time policy that looks constant
+
+`time_policy="linear"` extrapolates from the two most recent *distinct*
+exports. A historical bug where repeated executes between source updates
+collapsed the (prev, latest) baseline — degrading linear to constant after
+the first step — is fixed: the history rotates only when a genuinely new
+export (different `valid_time`) arrives, so the slope holds across every
+intermediate step (see
+`test_connector.py::test_time_policy_linear_holds_slope_across_repeated_executes`).
+Two semantics still surprise people:
+
+- The **first** transfer has one export in history and falls back to
+ constant — extrapolation needs two points.
+- Fields carrying a `lead_time` or `window` dimension have no single valid
+ time, so linear is undefined for them; the connector warns once (grep
+ `time_policy='linear' is undefined for field`) and holds constant.
+
+### NaN in outputs
+
+Two unrelated causes, both intentional:
+
+1. **Unwritten IO rows.** Streamed backend arrays are NaN-initialized so
+ never-written rows cannot masquerade as physical values (zarr's default
+ fill reads back as 0.0). A mediator's t0 row is always NaN (it exports
+ nothing before its first compute), as is everything after the point where
+ a run crashed. NaN at exact ring boundaries of a slow component = normal;
+ NaN spreading through a field mid-run = look upstream.
+2. **Mask fill not set.** If a source's data is NaN over masked-out points
+ (real SST products are NaN over land) and the connector has the default
+ `fill="none"`, those NaNs bleed through the bilinear regrid into coastal
+ destination cells. Set `fill="nearest"` (or `"zero"`) on every connector
+ leaving a masked source.
+
+### Exhausted clock
+
+A Driver runs its clock exactly once. A second `run()`, `steps()`, or
+`rollout()` raises `Driver clock exhausted` rather than silently yielding
+nothing (the old generator behavior). The rerun recipe is `driver.reset()`
+then `driver.initialize(ics)` — reset clears records, connector history, and
+IO-ready state, and deliberately invalidates initialization so you cannot
+resume from stale component state by accident.
+
+### Duplicate mediator deliveries
+
+Wiring the same field into a mediator twice per step (two connectors, or a
+re-executed slot) does **not** double-count: `AccumulationMediator` ignores
+arrivals whose `valid_time` it has already accumulated for that derived
+field. This is silent by design. Corollary: if your source republishes with
+an *unchanged* `valid_time`, the new values are ignored too — publish with
+the current ring time.
+
+### Warnings worth not ignoring
+
+- `Component 'x' exports [...] but no connector consumes them` — harmless if
+ the export is only for IO/records; a wiring bug otherwise.
+- `In-memory collection will hold ~N GB of export fields; pass collect=False
+ and a real IO backend (e.g. ZarrBackend) for runs of this size` — the
+ memory guard at initialize; see [IO and outputs](user_guide.md#io-and-outputs).
+- HEALPix/curvilinear sources: auto-regrid does not support them and raises
+ (`pass a custom regridder=`) — this is a v1 limitation, not a bug; see
+ [design and roadmap](design_and_roadmap.md).
diff --git a/earth2studio/nvcoupler/docs/user_guide.md b/earth2studio/nvcoupler/docs/user_guide.md
new file mode 100644
index 000000000..4c974cf7e
--- /dev/null
+++ b/earth2studio/nvcoupler/docs/user_guide.md
@@ -0,0 +1,881 @@
+# User guide
+
+Task-oriented recipes for building coupled systems with nvcoupler. Every
+snippet on this page runs as-is against the toy components in
+`earth2studio.nvcoupler.testing` — no weights, no network. For the mental
+model behind Fields, Connectors, and the run sequence see
+[concepts](concepts.md); for the full DSL and YAML grammar see the
+[DSL and YAML reference](dsl_and_yaml_reference.md); for symbol-by-symbol
+docs see the [API reference](api_reference.md). When something fails, start
+at [errors and troubleshooting](errors_and_troubleshooting.md).
+
+## Quickstart: couple two components in ten lines
+
+`couple()` auto-wires components by standard name, bridges any cadence gap a
+derived import needs (a windowed connector, or a mediator when the pair also
+carries a plain transfer), derives the run sequence from the coupling graph,
+and returns a ready-to-initialize `Driver`:
+
+```python
+import earth2studio.nvcoupler as nvc
+from earth2studio.nvcoupler.testing import atmos_ic, fake_atmos, fake_ocean, ocean_ic
+
+driver = nvc.couple(
+ fake_atmos(), fake_ocean(), start="2024-01-01", stop="2024-01-05"
+)
+print(driver.describe()) # the coupling plan, before anything runs
+driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+datasets = driver.run() # dict[str, xarray.Dataset]
+print(datasets["atmos"]["geopotential_at_1000hpa"].shape) # (17, 32, 64)
+```
+
+`describe()` prints a terraform-plan-style summary: one table row per
+component (type, cadence, imports, exports), one per connector (fields, time
+policy, fill, lagged/sequential mode, slot), then the derived run sequence.
+Here the ocean imports `geopotential_at_1000hpa_48h_mean` — a derived field
+nobody exports — so `couple()` synthesized a **windowed connector**
+(`Connector(atmos, ocean, window="48h", reduce="mean")`) that reduces the
+atmosphere's z1000 across the 6 h → 48 h cadence gap; its `match()` lists
+both the consumed base name and the delivered derived name. `run()` executes
+to the clock's stop time under `torch.inference_mode()` and returns one
+`xarray.Dataset` per component; the time axis is each component's own ring
+times including t0, so atmos has 17 rows and the 48 h ocean has 3.
+
+**Gotcha:** `couple()` uses lagged (NUOPC-explicit) coupling — every connect
+precedes the runs in its slot, so each destination sees the source's
+*previous* export. That is the reproducible default, not the only choice; see
+the next section. The full walkthrough is
+[example 01](../../../examples/09_nvcoupler/01_coupled_toy_workflow.py).
+
+## Hand-built systems: declarative Driver, DSL for ordering control
+
+`couple()` is a thin layer over the `Driver`'s own declarative form —
+components plus connections, no run sequence. Build it yourself when you
+need connector options (regridder, fill, window) or explicit control over
+which components participate; the schedule is still derived from the graph:
+
+```python
+from earth2studio.nvcoupler import Clock, Connector, Driver
+from earth2studio.nvcoupler.testing import atmos_ic, fake_atmos, fake_ocean, ocean_ic
+
+atmos, ocean = fake_atmos(), fake_ocean()
+driver = Driver(
+ {"atmos": atmos, "ocean": ocean},
+ clock=Clock("2024-01-01", "2024-01-05", "6h"),
+ connectors=[
+ ("ocean", "atmos"), # bare tuple -> default Connector
+ Connector(atmos, ocean, window="48h", reduce="mean"),
+ ],
+)
+assert driver.sequence_derived
+```
+
+When the coupling *order* is the experiment, pass the run sequence yourself
+(`sequence=` as DSL text or a `RunSequence`). Coupling semantics are pure
+ordering: a `src -> dst` line before `src`'s run in the same slot is lagged;
+after it, sequential.
+
+```python
+import numpy as np
+from earth2studio.nvcoupler import Clock, Driver, TrailingAverageMediator
+from earth2studio.nvcoupler.testing import atmos_ic, fake_atmos, fake_ocean, ocean_ic
+
+LAGGED = """
+@6h
+ atmos -> med
+ atmos
+@48h
+ med.compute
+ ocean -> atmos
+ med -> ocean
+ ocean
+@
+"""
+
+SEQUENTIAL = """
+@6h
+ atmos -> med
+ atmos
+@48h
+ med.compute
+ med -> ocean
+ ocean
+ ocean -> atmos
+@
+"""
+
+def build(dsl):
+ components = {
+ "atmos": fake_atmos(),
+ "ocean": fake_ocean(),
+ "med": TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"]),
+ }
+ driver = Driver(components, dsl, Clock("2024-01-01", "2024-01-05", "6h"))
+ driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+ return driver
+
+z_lag = build(LAGGED).run()["atmos"]["geopotential_at_1000hpa"].values[-1]
+z_seq = build(SEQUENTIAL).run()["atmos"]["geopotential_at_1000hpa"].values[-1]
+assert np.allclose(z_lag, 19.2, atol=1e-4) # forced by the previous SST
+assert np.allclose(z_seq, 19.2336, atol=1e-4) # forced by the fresh SST
+```
+
+The only difference between the two DSLs is where `ocean -> atmos` sits
+relative to `ocean`'s run — a one-line experiment
+([example 02](../../../examples/09_nvcoupler/02_lagged_vs_sequential.py)).
+
+**Gotcha:** a mediator's exports only exist after `med.compute`, so
+`med -> ocean` must come after the `MediateAction` in the slot, or the
+connector raises `CouplingError: ... has not produced ... yet`. Validation of
+names, cadences, and unfed imports all happens at `initialize()`, not
+mid-rollout.
+
+## Wrapping your model as a PrognosticComponent
+
+`PrognosticComponent` wraps any earth2studio `PrognosticModel`
+(`models/px/base.py` interface). The contract it reads off the model:
+
+- `input_coords()` / `output_coords(input_coords)` must return earth2studio
+ `CoordSystem`s with a `"variable"` coordinate.
+- **Timestep inference:** `timestep` defaults to
+ `output_coords["lead_time"][-1] - input_coords["lead_time"][-1]`. Pass
+ `timestep=` explicitly if your model has no `lead_time` axis.
+- **Export inference:** with `exports=None`, every output variable whose raw
+ name resolves through the field dictionary (or through your
+ `variable_aliases`) becomes an export; unknown variables are silently
+ skipped, so add aliases for anything you want exchanged.
+
+```python
+from collections import OrderedDict
+
+import numpy as np
+
+from earth2studio.nvcoupler import PrognosticComponent
+from earth2studio.nvcoupler.testing import grid_coords
+
+
+class MyModel:
+ """Stands in for any earth2studio PrognosticModel."""
+
+ def input_coords(self):
+ return OrderedDict(
+ {
+ "batch": np.array([0]),
+ "time": np.array([np.datetime64("2024-01-01")]),
+ "lead_time": np.array([np.timedelta64(0, "h")]),
+ "variable": np.array(["z1000", "sst"]),
+ **grid_coords(8, 16),
+ }
+ )
+
+ def output_coords(self, input_coords):
+ out = OrderedDict({k: v.copy() for k, v in input_coords.items()})
+ out["lead_time"] = input_coords["lead_time"] + np.timedelta64(6, "h")
+ return out
+
+ def __call__(self, x, coords):
+ return x + 1.0, self.output_coords(coords)
+
+
+atmos = PrognosticComponent(
+ "atmos",
+ MyModel(),
+ imports=["sea_surface_temperature"], # sst is also a state channel
+)
+# timestep inferred as 6h; exports inferred as both dictionary-known variables
+imports, exports = atmos.advertise()
+assert exports == ["geopotential_at_1000hpa", "sea_surface_temperature"]
+```
+
+`z1000` and `sst` resolve because they are registered aliases in the default
+dictionary. For model vocabularies the dictionary does not know, map raw
+names to standard names with `variable_aliases={"raw_name": "standard_name"}`
+— the alias is also registered so exports resolve.
+
+Published exports are *exchange-shaped*: size-1 `batch`/`time`/`lead_time`
+dims are squeezed off the exported Fields (the internal model state keeps
+them), so a `(1, 1, 1, var, lat, lon)` model couples cleanly to a plain
+`(var, lat, lon)` one.
+
+### Choosing an ImportAdapter
+
+The adapter owns the model call, because real models disagree on how coupled
+forcing arrives:
+
+| Your model receives forcing as... | Adapter | Call shape |
+|---|---|---|
+| a state variable you overwrite before stepping | `VariableOverwriteAdapter` (default) | `model(x, coords)` after injecting import slices |
+| a conditioning kwarg (StormScope) | `ConditioningKwargAdapter` | `model.call_with_conditioning(x, coords, conditioning=..., conditioning_coords=...)` |
+| an extra positional/keyword tensor (DLESyM, PhysicsNeMo 4-tensor) | `ExtraTensorAdapter` | `model(x, coords, coupling)` or `model(x, coords, **{kwarg: coupling})` |
+| its own internal `fetch_data` from a settable data source (StormCast) | `PullAdapter` | installs a `StateDataSource` on `model.conditioning_data_source`, then `model(x, coords)` — inference-only |
+
+The default only works when every imported field is *also* a variable of the
+model state (`"variable"` must be in the state coords); otherwise it raises
+with a pointer to the other adapters. **Gotcha:** with more than one
+import, `ConditioningKwargAdapter` and `ExtraTensorAdapter` require
+`field_order=[...]` — channel order cannot be inferred, and stacking
+alphabetically would run fine and predict garbage, so the framework refuses
+to guess.
+
+### Multi-window models: the next_input hook
+
+The default next-step input reuses the model output with the input
+`lead_time` coordinates — correct only when the model outputs as many lead
+times as it takes in. A model that consumes a sliding window (2 inputs, 1
+output) raises
+`CouplingError: model outputs 1 lead times but takes 2 as input — supply a
+next_input hook`. Provide
+`next_input=lambda prev_x, prev_coords, out, out_coords: (next_x, next_coords)`
+to roll the window yourself; `earth2studio/nvcoupler/dlesym_split.py` is the
+worked real-model example.
+
+### requires_ic
+
+`PrognosticComponent` and `CallableComponent` have `requires_ic = True`:
+`Driver.initialize(ics)` demands an `(x, coords)` entry for each, and fails
+with the exact `ics={...}` line to add. Mediators, `DataComponent`, and
+`DiagnosticComponent` set `requires_ic = False` and can be initialized with
+no arguments.
+
+## Coupling a pull-pattern model (StormCast-style)
+
+Some models take no forcing argument at all: they *pull* it, calling
+`fetch_data(self.conditioning_data_source, ...)` inside their own
+`__call__`. `PullAdapter` couples them without wrapping or modifying the
+model — before each step it installs a `StateDataSource` (an in-memory
+DataSource serving the component's import State) on the model's
+`conditioning_data_source` attribute, so the model's unmodified production
+fetch path receives this step's coupled forcing. Because the data crosses
+`fetch_data`'s xarray/numpy boundary, pull-coupled components are
+**inference-only**.
+
+```python
+from collections import OrderedDict
+
+import numpy as np
+import torch
+
+from earth2studio.data.utils import fetch_data
+from earth2studio.nvcoupler import (
+ DEFAULT_DICTIONARY,
+ CallableComponent,
+ Clock,
+ Connector,
+ Driver,
+ FieldDictionary,
+ FieldEntry,
+ PrognosticComponent,
+ PullAdapter,
+)
+from earth2studio.nvcoupler.testing import grid_coords
+
+GRID = (8, 16)
+T0 = np.datetime64("2024-01-01")
+
+
+class PullModel:
+ """StormCast-shaped: no forcing argument — pulls u10m/t2m through the
+ REAL fetch_data inside __call__, from a settable data-source attribute."""
+
+ conditioning_data_source = None # the injection point
+
+ def input_coords(self):
+ return OrderedDict(
+ {
+ "time": np.empty(0),
+ "lead_time": np.array([np.timedelta64(0, "h")]),
+ "variable": np.array(["refc"]),
+ **grid_coords(*GRID),
+ }
+ )
+
+ def output_coords(self, input_coords):
+ out = OrderedDict({k: v.copy() for k, v in input_coords.items()})
+ out["lead_time"] = input_coords["lead_time"] + np.timedelta64(1, "h")
+ return out
+
+ def __call__(self, x, coords):
+ cond, _ = fetch_data( # the production fetch path
+ self.conditioning_data_source,
+ time=np.atleast_1d(coords["time"]),
+ variable=np.array(["u10m", "t2m"]),
+ )
+ return (
+ x + 1.0 + cond[0, 0, 0].mean() + 0.1 * cond[0, 0, 1].mean(),
+ self.output_coords(coords),
+ )
+
+
+d = FieldDictionary(DEFAULT_DICTIONARY)
+d.register(FieldEntry("radar_reflectivity", "dBZ", aliases=frozenset({"refc"})))
+
+def global_step(x, coords): # u10m grows 1 m/s per step; t2m constant
+ return torch.stack([x[0] + 1.0, x[1]]), coords
+
+glob = CallableComponent(
+ "global", global_step, timestep="1h",
+ exports=["eastward_wind_10m", "air_temperature_2m"],
+)
+stormcast = PrognosticComponent(
+ "stormcast", PullModel(),
+ imports=["eastward_wind_10m", "air_temperature_2m"],
+ exports=["radar_reflectivity"],
+ import_adapter=PullAdapter(),
+ variable_aliases={"refc": "radar_reflectivity"},
+ dictionary=d,
+)
+driver = Driver(
+ {"global": glob, "stormcast": stormcast},
+ sequence="""
+ @1h
+ global
+ global -> stormcast # sequential: stormcast pulls FRESH forcing
+ stormcast
+ @
+ """,
+ clock=Clock(T0, "2024-01-01T03:00", "1h"),
+ connectors=[Connector(glob, stormcast)],
+)
+ic_sc = PullModel().input_coords()
+ic_sc["time"] = np.array([T0])
+driver.initialize(
+ {
+ "global": (
+ torch.stack([torch.full(GRID, 2.0), torch.full(GRID, 280.0)]),
+ OrderedDict({"variable": np.array(["u10m", "t2m"]), **grid_coords(*GRID)}),
+ ),
+ "stormcast": (torch.zeros(1, 1, 1, *GRID), ic_sc),
+ }
+)
+driver.run()
+# hour k pulls the CURRENT u = 2 + k: refc increments 32, 33, 34
+refc = stormcast.export_state["radar_reflectivity"]
+assert torch.allclose(refc.data, torch.full(GRID, 32.0 + 33.0 + 34.0))
+```
+
+The model's raw conditioning names (`u10m`, `t2m`) resolve against the import
+State through the component's aliases and the dictionary, so the shim answers
+them without any extra mapping. This is exactly the shape of earth2studio's
+`serve/server/example_workflows/stormcast_conus_workflow.py`, which today
+stages the full conditioning forecast to temp files and replays it through an
+`InferenceOutputSource` — with `PullAdapter` the same masquerade happens live,
+per step, with no staging. See
+[example 06](../../../examples/09_nvcoupler/06_pull_conditioning.py)
+(pull-pattern conditioning, StormCast-style) in `examples/09_nvcoupler/` for
+the full walkthrough.
+
+**Gotcha:** the served source is a snapshot of whatever the connector last
+delivered — under the default *lagged* ordering (connect before the source's
+run) the model pulls **stale** forcing from the previous step. Use an
+explicit sequential DSL as above (`global`, then `global -> stormcast`, then
+`stormcast` in one slot) so each pull sees fresh forcing;
+`PullAdapter(strict_time=True)` turns a misalignment into a
+`CouplingError` instead of a silent stale read.
+
+## Prescribed forcing: DataComponent
+
+Swapping a modeled ocean for observed SST is the classic AMIP-style
+experiment. A `DataComponent` fetches from any earth2studio `DataSource` at
+its own cadence and publishes the results as exports — the atmos, connectors,
+and sequence are untouched:
+
+```python
+import numpy as np
+import xarray as xr
+
+from earth2studio.nvcoupler import Clock, DataComponent, Driver
+from earth2studio.nvcoupler.testing import atmos_ic, fake_atmos, grid_coords
+
+
+class ConstantSST:
+ """Minimal DataSource: __call__(time, variable) -> xr.DataArray with
+ dims (time, variable, lat, lon)."""
+
+ grid = grid_coords(16, 32)
+
+ def __call__(self, time, variable) -> xr.DataArray:
+ time = np.atleast_1d(np.asarray(time, dtype="datetime64[ns]"))
+ variable = np.atleast_1d(np.asarray(variable))
+ lat, lon = self.grid["lat"], self.grid["lon"]
+ data = np.full((len(time), len(variable), len(lat), len(lon)), 3.0)
+ return xr.DataArray(
+ data,
+ dims=["time", "variable", "lat", "lon"],
+ coords={"time": time, "variable": variable, "lat": lat, "lon": lon},
+ )
+
+
+ocean = DataComponent(
+ "ocean",
+ source=ConstantSST(),
+ exports=["sea_surface_temperature"],
+ timestep="24h",
+)
+dsl = """
+@6h
+ ocean -> atmos
+ atmos
+@24h
+ ocean
+@
+"""
+driver = Driver(
+ {"atmos": fake_atmos(), "ocean": ocean},
+ dsl,
+ Clock("2024-01-01", "2024-01-03", "6h"),
+)
+driver.initialize({"atmos": atmos_ic()}) # no ocean entry: requires_ic=False
+ds = driver.run()
+```
+
+The mock above is the whole DataSource contract for testing: a callable
+`(time, variable) -> xr.DataArray` with dims `(time, variable, lat, lon)`.
+`DataComponent` fetches through `earth2studio.data.utils.fetch_data`, so any
+real source (WB2, GFS, an OISST archive) drops in. Standard names are mapped
+to source vocabulary through `variable_map={"standard_name": "raw_name"}`
+when the source's names are not dictionary aliases. On `initialize()` with no
+IC it fetches at `clock.start`, so lagged coupling has t0 data; the connector
+regrids the source grid onto each destination.
+
+**Gotcha:** the DataComponent must be *realized* before a no-arg
+`initialize()` (the Driver does this for you); standalone use raises
+`CouplingError: ... cannot fetch initial data before realize(clock)`.
+
+## Derived fields and impact chains
+
+Derived (time-reduced) fields are first-class dictionary entries carrying a
+`CellMethod` — machine-readable "I am `method` of `base` over `window`".
+Two consumers turn that declaration into a running windowed reduction: the
+**windowed connector** (the short form, preferred for one source feeding one
+destination) and the **AccumulationMediator** (the general form for multiple
+sources or custom reductions). Both share the same accumulator core and
+produce identical numbers.
+
+### The short form: a windowed connector
+
+Set `window=` and `reduce=` on the connector itself; it accumulates the
+source export every step and delivers the *derived* field (matched through
+the destination's `CellMethod` entry) on each window boundary — no mediator,
+no `med.compute` action, no extra slot lines:
+
+```python
+from earth2studio.nvcoupler import Clock, Connector, Driver
+from earth2studio.nvcoupler.testing import atmos_ic, fake_atmos, fake_ocean, ocean_ic
+
+atmos, ocean = fake_atmos(), fake_ocean() # ocean imports the 48h mean
+conn = Connector(atmos, ocean, window="48h", reduce="mean")
+
+DSL_WINDOWED = """
+@6h
+ atmos -> ocean # windowed: accumulates every step, delivers each 48h
+ ocean -> atmos
+ atmos
+@48h
+ ocean
+@
+"""
+driver = Driver(
+ {"atmos": atmos, "ocean": ocean},
+ DSL_WINDOWED,
+ Clock("2024-01-01", "2024-01-05", "6h"),
+ connectors=[conn],
+)
+driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+driver.run()
+# probes carry the DERIVED name, delivered on the ocean grid
+z48 = conn.last_transfer["geopotential_at_1000hpa_48h_mean"]
+assert z48.data.shape == (16, 32)
+```
+
+The destination must import a dictionary entry whose `CellMethod` is
+`(base=the source export, method=reduce, window=window)`; a plain import
+raises `CouplingError: ... register a FieldEntry(...)` — the coupler never
+invents names. Mid-window the destination's previous import stands; the
+first delivery lands one window after t0 (the window origin is the
+`valid_time` of the first execute's source field, i.e. the clock start under
+lagged coupling). `time_policy` does not apply on the windowed path.
+
+### The general form: an AccumulationMediator
+
+Register the entry, and an `AccumulationMediator` knows what to import,
+which reduction to run, and its cadence:
+
+```python
+import numpy as np
+
+from earth2studio.nvcoupler import (
+ AccumulationMediator,
+ CellMethod,
+ DEFAULT_DICTIONARY,
+ FieldDictionary,
+ FieldEntry,
+)
+
+dictionary = FieldDictionary(DEFAULT_DICTIONARY)
+dictionary.register(
+ FieldEntry(
+ "geopotential_at_1000hpa_24h_max",
+ "m2 s-2",
+ "24 h max of z1000",
+ cell_method=CellMethod(
+ "geopotential_at_1000hpa", "max", np.timedelta64(24, "h")
+ ),
+ )
+)
+med = AccumulationMediator(
+ "med", ["geopotential_at_1000hpa_24h_max"], dictionary=dictionary
+)
+assert med.import_names == ["geopotential_at_1000hpa"]
+```
+
+Wire it like any component: `atmos -> med` in the fast slot accumulates every
+delivery (running mean/sum/max/min, O(1) memory in window length), and
+`med.compute` in the slow slot exports the reduced field. Reach for the
+mediator rather than a windowed connector when several sources feed one
+reduction, when one reduced field fans out to several destinations, or when
+the reduction needs custom code (subclass `Mediator` — also the v1 home for
+unit conversions). Several derived fields can reduce the same base import in
+one mediator, but with mixed windows you must pass `window=` explicitly or
+split them across mediators. Deliveries carrying an already-seen `valid_time`
+are ignored, never double-counted — in mediators and windowed connectors
+alike.
+
+The impact-chain pattern stacks a `DiagnosticComponent` downstream: mediators
+turn fast fields into a 48 h precip sum and a 24 h t2m max, a diagnostic
+model consumes both and exports an impact index. `DiagnosticComponent` is a
+stateless single-step transform — imports/exports default to the model's own
+`input_coords()`/`output_coords()` variables resolved through the dictionary,
+so registered diagnostics wire up with just a name, the model, and a cadence.
+See [example 03](../../../examples/09_nvcoupler/03_impact_chain.py) for the
+full chain.
+
+**Gotcha:** if the mediator's `med.compute` runs before any sample arrived
+you get `CouplingError: no samples of ... accumulated before compute` — the
+connector feeding the mediator must sit in a *faster* slot than the compute.
+
+## Vertical coupling
+
+Only components with an explicit `level` dimension (chiefly chemistry
+emulators on hybrid sigma-pressure levels) touch this machinery. Models that
+encode levels in variable names (`z500`, `t850`) never do — those are just
+distinct fields.
+
+Declare what you publish and what you expect; the connector interpolates
+(linearly in log-pressure, differentiably) when they differ:
+
+```python
+from collections import OrderedDict
+
+import numpy as np
+import torch
+
+from earth2studio.nvcoupler import (
+ CallableComponent,
+ Clock,
+ Connector,
+ DEFAULT_DICTIONARY,
+ FieldDictionary,
+ FieldEntry,
+ HybridLevels,
+ PressureLevels,
+)
+from earth2studio.nvcoupler.field import Field
+from earth2studio.nvcoupler.testing import grid_coords
+
+d = FieldDictionary(DEFAULT_DICTIONARY)
+d.register(FieldEntry("ozone_mixing_ratio", "kg kg-1", aliases=frozenset({"o3"})))
+
+hybrid = HybridLevels(a=(30000.0, 20000.0, 0.0), b=(0.0, 0.5, 1.0)) # p = a + b*ps
+pressure = PressureLevels((500.0, 850.0)) # hPa
+
+identity = lambda x, coords: (x, coords)
+met = CallableComponent(
+ "met", identity, "6h",
+ exports=["ozone_mixing_ratio"],
+ dictionary=d,
+ export_vertical={"ozone_mixing_ratio": hybrid},
+)
+chem = CallableComponent(
+ "chem", identity, "6h",
+ imports=["ozone_mixing_ratio"],
+ dictionary=d,
+ import_vertical={"ozone_mixing_ratio": pressure},
+)
+clock = Clock("2024-01-01", "2024-01-02", "6h")
+met.realize(clock)
+chem.realize(clock)
+
+grid = grid_coords(4, 8)
+o3 = torch.arange(3.0).view(1, 3, 1, 1).expand(1, 3, 4, 8).clone()
+met.initialize(
+ o3, OrderedDict({"variable": np.array(["o3"]), "level": np.arange(3.0), **grid})
+)
+# surface pressure has no level dim, so it cannot ride in the same state
+# tensor as o3; add it to the export state directly
+met.export_state.add(
+ Field(
+ torch.full((4, 8), 100000.0), OrderedDict(grid),
+ "surface_pressure", "Pa",
+ valid_time=np.datetime64("2024-01-01"), source="met",
+ )
+)
+chem.initialize(
+ torch.zeros(1, 2, 4, 8),
+ OrderedDict(
+ {"variable": np.array(["o3"]), "level": np.array([500.0, 850.0]), **grid}
+ ),
+)
+Connector(met, chem, fields=["ozone_mixing_ratio"]).execute(np.datetime64("2024-01-01"))
+assert chem.import_state["ozone_mixing_ratio"].vertical == pressure
+```
+
+Hybrid sources depend on surface pressure: the connector pulls
+`surface_pressure` (or whatever `HybridLevels.ps_field` names) from the
+*source's* export state automatically, and raises `VerticalMismatchError`
+with the fix (`add it to the source's export list`) when it is absent.
+Interpolation clamps to the source column ends — no extrapolation beyond the
+top/bottom levels — and v1 only interpolates *onto* `PressureLevels`.
+[Example 04](../../../examples/09_nvcoupler/04_vertical_chemistry.py) runs
+the met/chem pair end to end in a Driver.
+
+**Gotcha:** for `PressureLevels` sources, the data's `level` coordinate (hPa)
+must equal the declared levels exactly, in top-to-bottom (increasing)
+order — a mismatch raises rather than silently pairing slices with wrong
+pressures.
+
+## Masked fields
+
+Components exporting fields that are only valid on part of the grid (SST on
+ocean points) declare `export_masks={"standard_name": bool_tensor}` (True =
+valid). The mask travels on the Field; the *connector's* `fill=` option
+decides what invalid points become, always before regridding so garbage
+cannot bleed into the interpolation:
+
+```python
+from earth2studio.nvcoupler import Clock, Connector, Driver, TrailingAverageMediator
+from earth2studio.nvcoupler.testing import atmos_ic, fake_atmos, fake_ocean, ocean_ic
+
+DSL = """
+@6h
+ atmos -> med
+ ocean -> atmos
+ atmos
+@48h
+ med.compute
+ med -> ocean
+ ocean
+@
+"""
+components = {
+ "atmos": fake_atmos(),
+ "ocean": fake_ocean(with_mask=True), # northern half is land
+ "med": TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"]),
+}
+driver = Driver(
+ components,
+ DSL,
+ Clock("2024-01-01", "2024-01-05", "6h"),
+ connectors=[
+ Connector(components["ocean"], components["atmos"], fill="nearest")
+ ],
+)
+driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+driver.run()
+assert driver.probe("ocean->atmos")["sea_surface_temperature"].mask is None
+```
+
+- `fill="nearest"`: every invalid point takes its nearest valid neighbor
+ (great-circle KDTree, differentiable gather) — the principled version of
+ DLESyM's SST NaN-interpolation. The filler is cached per (grid, mask).
+- `fill="zero"`: invalid points become 0.0.
+- Both consume the mask (`field.mask is None` downstream).
+
+**Without fill** (`fill="none"`, the default) nothing happens: invalid values
+pass straight into the regrid, bleeding into neighboring destination cells,
+and the mask stays attached but describes the *source* grid while the data is
+now on the destination grid. Import adapters ignore masks entirely. If a
+source declares `export_masks`, set `fill=` on every connector that leaves
+it — this is also why `couple()`'s describe output shows the fill column.
+
+## IO and outputs
+
+Two independent output paths:
+
+- **In-memory collection** (`collect=True`, the default): every ring's export
+ fields are recorded and `run()` / `to_xarray()` returns one
+ `xarray.Dataset` per component. Records are detached clones, off the
+ exchange path.
+- **Streaming IO** (`io={"name": backend}`): each component gets its own
+ `IOBackend`; arrays are allocated with a leading `time` axis covering the
+ component's ring times including t0, and a row is written after the IC seed
+ and after every run.
+
+```python
+import numpy as np
+
+from earth2studio.io import ZarrBackend
+from earth2studio.nvcoupler import Clock, Driver, TrailingAverageMediator
+from earth2studio.nvcoupler.testing import atmos_ic, fake_atmos, fake_ocean, ocean_ic
+
+io = {"atmos": ZarrBackend(), "med": ZarrBackend()}
+components = {
+ "atmos": fake_atmos(),
+ "ocean": fake_ocean(),
+ "med": TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"]),
+}
+driver = Driver(
+ components, DSL, Clock("2024-01-01", "2024-01-05", "6h"),
+ io=io,
+ collect=False, # nothing kept in memory
+)
+driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+result = driver.run() # {} when collect=False
+z = io["atmos"]["geopotential_at_1000hpa"][:] # (17, 32, 64)
+zm = io["med"]["geopotential_at_1000hpa_48h_mean"][:]
+assert np.all(np.isnan(zm[0])) # t0 row: NaN, by design
+```
+
+(`DSL` as in the previous section.)
+
+**NaN semantics:** backend arrays are NaN-initialized deliberately. Zarr's
+default fill reads back as 0.0, which would let never-written rows — a
+mediator's t0 row (mediators export nothing until their first compute), or
+the tail of a crashed run — masquerade as physical values. NaN rows in your
+output mean "this ring was never written", not a numerical blow-up.
+
+**Memory guard:** with `collect=True`, `initialize()` estimates the total
+collected size and logs a warning above ~4 GB telling you to pass
+`collect=False` plus a real IO backend. IO and collection are independent —
+you can have both, either, or neither.
+
+## Gradients and coupled fine-tuning
+
+`run()` and `steps()` execute under `torch.inference_mode()` — fast, no
+graphs. `rollout(n_steps)` is the training entry point: under
+`torch.enable_grad()` the autograd graph survives the entire exchange path
+(regrid gathers, mediator reductions, functional import injection):
+
+```python
+import torch
+
+from earth2studio.nvcoupler import Clock, Driver, TrailingAverageMediator
+from earth2studio.nvcoupler.testing import atmos_ic, fake_atmos, fake_ocean, ocean_ic
+
+gain_atmos = torch.tensor(1.0, requires_grad=True)
+gain_ocean = torch.tensor(1.0, requires_grad=True)
+components = {
+ "atmos": fake_atmos(gain=gain_atmos),
+ "ocean": fake_ocean(gain=gain_ocean),
+ "med": TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"]),
+}
+driver = Driver(components, DSL, Clock("2024-01-01", "2024-01-05", "6h"))
+driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+
+with torch.enable_grad():
+ states = driver.rollout(16) # keeps the graph
+loss = states["atmos"]["geopotential_at_1000hpa"].data.sum()
+loss.backward() # reaches BOTH components' parameters
+assert gain_atmos.grad is not None and gain_ocean.grad is not None
+```
+
+Collected records and IO writes are detached — they only feed
+`to_xarray()`/backends and never pin a step's graph, so collection stays on
+during training without a memory explosion. Optimizer loops, truncated BPTT,
+and per-component device placement are out of scope for v1 (see
+[design and roadmap](design_and_roadmap.md));
+[example 05](../../../examples/09_nvcoupler/05_coupled_finetuning.py) runs a
+real training step.
+
+**Gotcha:** `rollout(n)` with more steps than remain on the clock raises a
+`CouplingError` naming `clock.n_steps` — size your rollout to the clock, or
+`reset()` + `initialize()` between epochs.
+
+## YAML round-trip
+
+`to_yaml(driver)` / `from_yaml(text_or_path)` serialize the clock, the run
+sequence (hand-written sequences verbatim as DSL text; derived sequences as
+`sequence: {derived: true, text: ...}`, re-derived deterministically on
+load), non-default dictionary entries and aliases, component specs, and
+connector settings (`src`, `dst`, `fields`, `time_policy`, `fill`, and
+`window`/`reduce` for windowed connectors). What round-trips:
+
+- `AccumulationMediator` / `TrailingAverageMediator`: automatically (name,
+ fields, window).
+- Anything else: only if it carries a `yaml_spec` attribute — a
+ `{"class": "", "kwargs": {...}}` dict naming a module-level
+ class or factory that rebuilds the component from kwargs alone.
+
+```python
+import earth2studio.nvcoupler as nvc
+from earth2studio.nvcoupler import Clock, Driver, TrailingAverageMediator
+from earth2studio.nvcoupler.testing import atmos_ic, fake_atmos, fake_ocean, ocean_ic
+
+atmos = fake_atmos()
+atmos.yaml_spec = {
+ "class": "earth2studio.nvcoupler.testing.fake_atmos",
+ "kwargs": {"gain": 1.0, "timestep": "6h"},
+}
+ocean = fake_ocean()
+ocean.yaml_spec = {
+ "class": "earth2studio.nvcoupler.testing.fake_ocean",
+ "kwargs": {"gain": 1.0, "timestep": "48h"},
+}
+med = TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"])
+driver = Driver(
+ {"atmos": atmos, "ocean": ocean, "med": med},
+ DSL,
+ Clock("2024-01-01", "2024-01-05", "6h"),
+)
+text = nvc.to_yaml(driver) # or to_yaml(driver, path="system.yaml")
+rebuilt = nvc.from_yaml(text) # uninitialized; call initialize(ics) as usual
+rebuilt.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+```
+
+What cannot round-trip: components wrapping closures or live model objects
+without a `yaml_spec` (`to_yaml` raises with the fix), custom `regridder=`
+callables and custom ImportAdapter instances (connector `fields`/policies
+serialize; callables do not), and model checkpoints referenced by load paths
+(out of scope for v1). Initial conditions are never serialized — a config
+describes the system, not its state. Full schema in the
+[DSL and YAML reference](dsl_and_yaml_reference.md).
+
+## Inspection and debugging
+
+- `driver.describe()` (or `nvc.describe(driver)`) — the plan: components,
+ connectors with time policy/fill/mode/slot, and the run sequence. Works
+ *before* `initialize()`; in Jupyter, a bare `driver` renders the HTML
+ version.
+- `driver.probe("src->dst")` — the last Fields exchanged on a connector
+ (post-pipeline: time policy, vertical, fill, regrid already applied).
+ Spaces are tolerated: `probe("ocean -> atmos")`.
+- `driver.steps()` — iterate `(time, {component: export State})` per driver
+ step, notebook-style, instead of one opaque `run()`.
+- loguru levels: warnings (unconsumed exports, memory guard, linear-policy
+ fallback) are on by default; every exchange also logs at DEBUG:
+
+```python
+import sys
+
+from loguru import logger
+
+logger.remove()
+logger.add(sys.stderr, level="DEBUG") # show per-exchange lines
+
+import earth2studio.nvcoupler as nvc
+from earth2studio.nvcoupler.testing import atmos_ic, fake_atmos, fake_ocean, ocean_ic
+
+driver = nvc.couple(fake_atmos(), fake_ocean(), start="2024-01-01", stop="2024-01-03")
+driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+for time, states in driver.steps():
+ z = states["atmos"]["geopotential_at_1000hpa"]
+ print(time, float(z.data.mean()))
+print(driver.probe("ocean -> atmos"))
+```
+
+Each DEBUG line reads
+`exchange ocean->atmos: sea_surface_temperature (valid 2024-01-01T00:00...)`
+— the `valid` timestamp is the fastest way to see lagged coupling in action
+(the delivered field is older than the current step). When something raises
+instead, every message names the components, the field, and the fix; the
+complete catalog is in
+[errors and troubleshooting](errors_and_troubleshooting.md).
diff --git a/earth2studio/nvcoupler/driver.py b/earth2studio/nvcoupler/driver.py
new file mode 100644
index 000000000..b72500c2d
--- /dev/null
+++ b/earth2studio/nvcoupler/driver.py
@@ -0,0 +1,518 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Driver: owns the clock and executes the run sequence (NUOPC_Driver analog).
+
+Declare the coupling graph — components plus connections — and the schedule
+follows: with no sequence given the Driver derives the canonical run
+sequence from the graph (:func:`~.sequence.derive_sequence`). A hand-written
+sequence (RunSequence object or DSL text) is the override for schedules the
+graph cannot express. The lifecycle mirrors NUOPC: construct, then
+``initialize(ics)`` (advertise -> connector matching -> realize -> component
+initialize), then ``run()`` / ``steps()`` for inference or ``rollout(n)``
+for a gradient-carrying advance. Coupling order is entirely the run
+sequence's action order; the driver adds no hidden exchanges.
+"""
+
+from collections import OrderedDict
+from collections.abc import Iterator
+from dataclasses import replace
+from typing import TYPE_CHECKING
+
+import numpy as np
+import torch
+from loguru import logger
+
+if TYPE_CHECKING:
+ from earth2studio.io import IOBackend
+
+from earth2studio.utils.type import CoordSystem
+
+from .clock import Clock, as_timedelta
+from .component import Component
+from .connector import Connector
+from .errors import CouplingError, UnmatchedImportError
+from .field import Field, State
+from .mediator import Mediator
+from .sequence import (
+ ConnectAction,
+ MediateAction,
+ RunAction,
+ RunSequence,
+ derive_sequence,
+ parse_run_sequence,
+)
+
+MEMORY_WARN_BYTES = 4e9
+
+
+class Driver:
+ """Executes a coupled system declared as components + connections.
+
+ The declarative form needs no run sequence — the schedule is derived
+ from the coupling graph::
+
+ Driver(
+ {"atmos": atmos, "ocean": ocean},
+ clock=Clock("2024-01-01", "2024-01-05", "6h"),
+ connectors=[("ocean", "atmos"), ("atmos", "ocean")],
+ )
+
+ Bare ``(src, dst)`` name tuples become default Connectors; pass Connector
+ instances for regridding/time-policy/window options. Passing a sequence
+ (RunSequence object or DSL text) overrides the derived schedule — the
+ escape hatch for sequential coupling or hand-tuned action order.
+
+ Parameters
+ ----------
+ components : dict[str, Component]
+ All participants, mediators included; keys must match the names used
+ in the run sequence.
+ sequence : RunSequence | str, optional
+ Run sequence object or DSL text. None (default) derives the
+ canonical sequence from components + connectors via
+ :func:`~.sequence.derive_sequence`.
+ clock : Clock
+ Coupling clock; dt must divide every component timestep.
+ connectors : list[Connector | tuple[str, str]], optional
+ The coupling graph's edges: Connector instances or bare (src, dst)
+ name tuples (auto-built into default Connectors). With an explicit
+ sequence, any ConnectAction without a matching (src, dst) pair still
+ gets a default ``Connector(src, dst)`` built at initialize.
+ collect : bool
+ Keep per-ring export fields in memory for ``to_xarray()``
+ (default True; disable for long runs writing to real IO). Recorded
+ fields are detached clones — records are off the exchange path, so
+ collection never pins autograd graphs during gradient rollouts.
+ allow_unfed_imports : bool
+ By default ``initialize`` raises UnmatchedImportError when a
+ component advertises an import that no connector in the run sequence
+ delivers (the component would silently run on stale initial-condition
+ forcing). Pass True to downgrade this to a warning (default False).
+ io : dict[str, IOBackend], optional
+ Per-component IO backends (e.g. ``{"atmos": ZarrBackend()}``). Each
+ backend receives one array per export field of its component, with a
+ leading ``time`` dimension covering the component's ring times
+ including t0; the driver streams a write after the initial-condition
+ seed and after every component run. Independent of ``collect``.
+
+ Attributes
+ ----------
+ sequence_derived : bool
+ True when the run sequence was derived from the coupling graph
+ rather than passed in.
+ """
+
+ def __init__(
+ self,
+ components: dict[str, Component],
+ sequence: RunSequence | str | None = None,
+ clock: Clock | None = None,
+ connectors: "list[Connector | tuple[str, str]] | None" = None,
+ collect: bool = True,
+ io: "dict[str, IOBackend] | None" = None,
+ allow_unfed_imports: bool = False,
+ ):
+ if clock is None:
+ raise CouplingError(
+ "Driver needs a clock — Driver(components, sequence, clock) "
+ "or Driver(components, clock=Clock(start, stop, dt), "
+ "connectors=[...])"
+ )
+ self.components = dict(components)
+ self.clock = clock
+ self.collect = collect
+ self.allow_unfed_imports = allow_unfed_imports
+ self._connectors: dict[tuple[str, str], Connector] = {}
+ for item in connectors or []:
+ if not isinstance(item, Connector):
+ src, dst = item
+ unknown = [n for n in (src, dst) if n not in self.components]
+ if unknown:
+ raise CouplingError(
+ f"Connection ({src!r}, {dst!r}): {unknown} are not "
+ f"component names; known components: "
+ f"{sorted(self.components)}"
+ )
+ item = Connector(self.components[src], self.components[dst])
+ self._connectors[(item.src.name, item.dst.name)] = item
+ self.sequence_derived = sequence is None
+ if sequence is None:
+ self.sequence = derive_sequence(self.components, self._connectors.values())
+ else:
+ self.sequence = (
+ parse_run_sequence(sequence) if isinstance(sequence, str) else sequence
+ )
+ # per-component record of (time, {std_name: Field}) for to_xarray
+ self._records: dict[str, list[tuple[np.datetime64, dict[str, Field]]]] = {
+ name: [] for name in self.components
+ }
+ self._io: dict[str, "IOBackend"] = dict(io or {})
+ unknown = [n for n in self._io if n not in self.components]
+ if unknown:
+ raise CouplingError(
+ f"io= keys {unknown} are not component names; known "
+ f"components: {sorted(self.components)}"
+ )
+ self._io_ready: set[str] = set()
+ self._initialized = False
+
+ # -- setup -----------------------------------------------------------------
+ def initialize(
+ self, ics: dict[str, tuple[torch.Tensor, CoordSystem]] | None = None
+ ) -> None:
+ ics = ics or {}
+ self.sequence.validate(self.components, self.clock.dt)
+ # build connectors for every ConnectAction not covered by a prebuilt one
+ for action in self.sequence.connections():
+ key = (action.src, action.dst)
+ if key not in self._connectors:
+ self._connectors[key] = Connector(
+ self.components[action.src], self.components[action.dst]
+ )
+ for conn in self._connectors.values():
+ conn.match()
+ self._check_unfed_imports()
+ # realize + initialize components
+ for name, comp in self.components.items():
+ comp.realize(self.clock)
+ if name in ics:
+ comp.initialize(*ics[name])
+ elif not getattr(comp, "requires_ic", not isinstance(comp, Mediator)):
+ comp.initialize()
+ else:
+ raise CouplingError(
+ f"Component {name!r} needs an initial condition — pass "
+ f"ics={{{name!r}: (x, coords)}}"
+ )
+ self._record(name, self.clock.start, comp.export_state)
+ self._io_write(name, self.clock.start)
+ self._warn_unconsumed_exports()
+ self._warn_memory()
+ self._initialized = True
+
+ def _check_unfed_imports(self) -> None:
+ """Every advertised import must be delivered by some connector —
+ otherwise the component silently runs the whole simulation on its
+ stale initial-condition forcing."""
+ fed: dict[str, set[str]] = {name: set() for name in self.components}
+ for (_, dst), conn in self._connectors.items():
+ fed[dst] |= set(conn.match())
+ available = {n: list(c.export_names) for n, c in self.components.items()}
+ for name, comp in self.components.items():
+ for field in comp.import_names:
+ if field in fed[name]:
+ continue
+ if self.allow_unfed_imports:
+ logger.warning(
+ "Component {!r} imports {!r} but no connector in the "
+ "run sequence delivers it — it will run on stale "
+ "initial-condition forcing",
+ name,
+ field,
+ )
+ else:
+ raise UnmatchedImportError(name, field, available)
+
+ def _warn_unconsumed_exports(self) -> None:
+ consumed: set[tuple[str, str]] = set()
+ for (src, _), conn in self._connectors.items():
+ consumed |= {(src, f) for f in conn.match()}
+ for name, comp in self.components.items():
+ idle = [f for f in comp.export_names if (name, f) not in consumed]
+ if idle:
+ logger.warning(
+ "Component {!r} exports {} but no connector consumes them",
+ name,
+ idle,
+ )
+
+ def _warn_memory(self) -> None:
+ if not self.collect:
+ return
+ total = 0
+ for comp in self.components.values():
+ per_ring = sum(
+ f.data.numel() * f.data.element_size()
+ for f in comp.export_state.values()
+ )
+ rings = self.clock.n_steps * (
+ self.clock.dt.astype(np.int64) / comp.timestep.astype(np.int64)
+ )
+ total += per_ring * max(rings, 0)
+ if total > MEMORY_WARN_BYTES:
+ logger.warning(
+ "In-memory collection will hold ~{:.1f} GB of export fields; "
+ "pass collect=False and a real IO backend (e.g. ZarrBackend) "
+ "for runs of this size",
+ total / 1e9,
+ )
+
+ # -- record keeping ----------------------------------------------------------
+ def _record(self, name: str, time: np.datetime64, state: State) -> None:
+ """Append a snapshot of a component's export fields to _records.
+
+ Records are off the exchange path (they only feed ``to_xarray``), so
+ the field data is stored as detached clones — otherwise collection
+ would pin every step's autograd graph during gradient rollouts. The
+ Fields in component states and exchanges stay attached.
+ """
+ if not self.collect:
+ return
+ snapshot = {
+ std: replace(f, data=f.data.detach().clone()) for std, f in state.items()
+ }
+ self._records[name].append((time, snapshot))
+
+ def _strip_singletons(
+ self, name: str, field: Field
+ ) -> tuple[torch.Tensor, OrderedDict]:
+ """Drop size-1 'time'/'batch' dims (published by components whose
+ model coords carry them) so the ring-time axis can be prepended
+ without colliding with the field's own stale coordinate."""
+ data = field.data
+ coords: OrderedDict = OrderedDict(field.coords)
+ for key in ("time", "batch"):
+ if key not in coords:
+ continue
+ axis = list(coords).index(key)
+ if data.shape[axis] != 1:
+ raise CouplingError(
+ f"Component {name!r} export {field.standard_name!r} "
+ f"carries a {key!r} dimension of size {data.shape[axis]}; "
+ "the driver records one snapshot per ring and can only "
+ f"absorb a size-1 {key!r} dim — publish a single "
+ f"{key}-slice per run"
+ )
+ data = data.squeeze(axis)
+ del coords[key]
+ return data, coords
+
+ # -- IO streaming ------------------------------------------------------------
+ def _io_ring_times(self, comp: Component) -> np.ndarray:
+ """A component's ring times including t0 (its 'time' IO coordinate)."""
+ span_ns = (
+ (self.clock.stop - self.clock.start)
+ .astype("timedelta64[ns]")
+ .astype(np.int64)
+ )
+ n_rings = int(span_ns // comp.timestep.astype(np.int64))
+ return (self.clock.start + np.arange(n_rings + 1) * comp.timestep).astype(
+ "datetime64[ns]"
+ )
+
+ def _io_setup(self, name: str) -> bool:
+ """Allocate backend arrays for a component's export fields; returns
+ False when the component has not published any fields yet (mediators
+ at t0), in which case setup is retried on the next write."""
+ comp = self.components[name]
+ if not comp.export_state:
+ return False
+ backend = self._io[name]
+ times = self._io_ring_times(comp)
+ # group export fields sharing a coord structure into one add_array call
+ groups: dict[tuple, tuple[OrderedDict, list[str]]] = {}
+ for std_name, field in comp.export_state.items():
+ _, field_coords = self._strip_singletons(name, field)
+ key = tuple((k, np.asarray(v).tobytes()) for k, v in field_coords.items())
+ if key not in groups:
+ total_coords: OrderedDict = OrderedDict(time=times)
+ total_coords.update((k, np.asarray(v)) for k, v in field_coords.items())
+ groups[key] = (total_coords, [])
+ groups[key][1].append(std_name)
+ for total_coords, std_names in groups.values():
+ # NaN-initialize: zarr's default fill reads back as 0.0, which
+ # would let never-written rows (e.g. cadences coarser than the
+ # ring times, or a crashed run) masquerade as physical values
+ shape = tuple(len(v) for v in total_coords.values())
+ nan_init = [
+ torch.full(
+ shape,
+ float("nan"),
+ dtype=comp.export_state[n].data.dtype,
+ )
+ for n in std_names
+ ]
+ backend.add_array(total_coords, std_names, data=nan_init)
+ self._io_ready.add(name)
+ return True
+
+ def _io_write(self, name: str, time: np.datetime64) -> None:
+ """Stream a component's current export fields to its IO backend."""
+ if name not in self._io:
+ return
+ if name not in self._io_ready and not self._io_setup(name):
+ return
+ backend = self._io[name]
+ comp = self.components[name]
+ t = np.asarray([time], dtype="datetime64[ns]")
+ for std_name, field in comp.export_state.items():
+ data, field_coords = self._strip_singletons(name, field)
+ coords: OrderedDict = OrderedDict(time=t)
+ coords.update(field_coords)
+ # IO is off the exchange path: detaching here keeps backends
+ # (which convert to numpy) working during gradient rollouts
+ backend.write(data.detach().unsqueeze(0), coords, std_name)
+
+ # -- execution ----------------------------------------------------------------
+ def _slot_aligned(self, time: np.datetime64, interval: np.timedelta64) -> bool:
+ elapsed = (time - self.clock.start).astype("timedelta64[ns]")
+ elapsed_ns = elapsed.astype(np.int64)
+ interval_ns = as_timedelta(interval).astype(np.int64)
+ return elapsed_ns > 0 and elapsed_ns % interval_ns == 0
+
+ def _execute_time(self, time: np.datetime64) -> None:
+ for slot in self.sequence.slots:
+ if not self._slot_aligned(time, slot.interval):
+ continue
+ for action in slot.actions:
+ if isinstance(action, RunAction):
+ comp = self.components[action.component]
+ comp.run(time)
+ self._record(action.component, time, comp.export_state)
+ self._io_write(action.component, time)
+ elif isinstance(action, ConnectAction):
+ self._connectors[(action.src, action.dst)].execute(time)
+ elif isinstance(action, MediateAction):
+ med = self.components[action.mediator]
+ med.run(time)
+ self._record(action.mediator, time, med.export_state)
+ self._io_write(action.mediator, time)
+
+ def _check_not_exhausted(self) -> None:
+ if self.clock.done():
+ raise CouplingError(
+ f"Driver clock exhausted: already at stop time "
+ f"{self.clock.stop} — the run has completed. Call "
+ "driver.reset() and driver.initialize(ics) to run again"
+ )
+
+ def _steps_impl(self) -> Iterator[tuple[np.datetime64, dict[str, State]]]:
+ if not self._initialized:
+ raise CouplingError("Driver.initialize(ics) must be called before running")
+ for time in self.clock:
+ self._execute_time(time)
+ yield time, {n: c.export_state for n, c in self.components.items()}
+
+ def steps(self) -> Iterator[tuple[np.datetime64, dict[str, State]]]:
+ """Yield (time, {component: export State}) after every driver step —
+ the notebook-inspection path (computation runs under inference mode)."""
+ self._check_not_exhausted()
+ return self._steps_gen()
+
+ def _steps_gen(self) -> Iterator[tuple[np.datetime64, dict[str, State]]]:
+ it = self._steps_impl()
+ while True:
+ with torch.inference_mode():
+ try:
+ time, states = next(it)
+ except StopIteration:
+ return
+ yield time, states
+
+ def run(self) -> "dict[str, object]":
+ """Run to the clock's stop time; returns ``to_xarray()`` when
+ collection is on, else an empty dict."""
+ self._check_not_exhausted()
+ with torch.inference_mode():
+ for _ in self._steps_impl():
+ pass
+ return self.to_xarray() if self.collect else {}
+
+ def rollout(self, n_steps: int) -> dict[str, State]:
+ """Advance n driver steps keeping the autograd graph (when grad is
+ enabled) — the coupled fine-tuning entry point. Returns each
+ component's export State after the last step."""
+ self._check_not_exhausted()
+ it = self._steps_impl()
+ states: dict[str, State] = {}
+ for i in range(n_steps):
+ try:
+ _, states = next(it)
+ except StopIteration:
+ raise CouplingError(
+ f"rollout({n_steps}) ran past the clock stop time "
+ f"{self.clock.stop}: only {i} steps remained. Use "
+ f"n_steps <= clock.n_steps ({self.clock.n_steps}) or "
+ "call driver.reset() and re-initialize to run again"
+ ) from None
+ return states
+
+ def reset(self) -> None:
+ """Rewind for a fresh run: reset the clock, clear collected records,
+ connector transfer history, and IO-ready state. The driver must be
+ re-initialized (``initialize(ics)``) before running again."""
+ self.clock.reset()
+ self._records = {name: [] for name in self.components}
+ for conn in self._connectors.values():
+ conn.reset()
+ self._io_ready.clear()
+ self._initialized = False
+
+ # -- inspection ------------------------------------------------------------------
+ def describe(self) -> str:
+ """Terraform-plan-style preview of the coupled system (works before
+ initialize); see :func:`earth2studio.nvcoupler.api.describe`."""
+ from .api import describe
+
+ return describe(self)
+
+ def _repr_html_(self) -> str:
+ from .api import describe_html
+
+ return describe_html(self)
+
+ def probe(self, connector: str) -> dict[str, Field]:
+ """Last fields exchanged on a connector, addressed as "src->dst"."""
+ for conn in self._connectors.values():
+ if conn.name == connector.replace(" ", ""):
+ return dict(conn.last_transfer)
+ known = [c.name for c in self._connectors.values()]
+ raise KeyError(f"No connector {connector!r}; have {known}")
+
+ def to_xarray(self) -> "dict[str, object]":
+ """Collected export fields as one xarray.Dataset per component."""
+ import xarray as xr
+
+ out: dict[str, object] = {}
+ for name, records in self._records.items():
+ if not records:
+ continue
+ # size-1 time/batch dims on the fields themselves are stripped so
+ # the leading axis is always the RING times, never a field's own
+ # (stale) time coordinate
+ by_var: dict[str, tuple[list, list, OrderedDict]] = {}
+ for time, fields in records:
+ for std, f in fields.items():
+ data, field_coords = self._strip_singletons(name, f)
+ times, datas, _ = by_var.setdefault(std, ([], [], field_coords))
+ times.append(time)
+ datas.append(data.detach().cpu().numpy())
+ data_vars = {}
+ coords: dict[str, object] = {}
+ for std, (times, datas, field_coords) in by_var.items():
+ dims = ("time", *field_coords.keys())
+ data_vars[std] = (dims, np.stack(datas))
+ coords["time"] = np.asarray(times, dtype="datetime64[ns]")
+ coords.update(field_coords)
+ out[name] = xr.Dataset(data_vars, coords=coords)
+ return out
+
+ def __repr__(self) -> str:
+ return (
+ f"Driver(components={sorted(self.components)}, "
+ f"clock={self.clock!r}, connectors={[c.name for c in self._connectors.values()]})"
+ )
diff --git a/earth2studio/nvcoupler/errors.py b/earth2studio/nvcoupler/errors.py
new file mode 100644
index 000000000..4ae34118c
--- /dev/null
+++ b/earth2studio/nvcoupler/errors.py
@@ -0,0 +1,119 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Error hierarchy for the nvcoupler coupling framework.
+
+Every validation error names the components and fields involved and, where
+possible, the concrete fix — so a misconfigured coupled system fails at
+initialize time with an actionable message rather than mid-rollout.
+"""
+
+import difflib
+from collections.abc import Iterable
+
+
+def suggest(name: str, candidates: Iterable[str], n: int = 3) -> str:
+ """Format a 'did you mean' suffix from close matches, or empty string."""
+ matches = difflib.get_close_matches(name, list(candidates), n=n, cutoff=0.5)
+ if not matches:
+ return ""
+ return f" Did you mean: {', '.join(repr(m) for m in matches)}?"
+
+
+class CouplingError(Exception):
+ """Base class for all nvcoupler configuration and runtime errors."""
+
+
+class UnknownFieldError(CouplingError):
+ """A name could not be resolved in the field dictionary."""
+
+ def __init__(self, name: str, candidates: Iterable[str]):
+ super().__init__(
+ f"Field name {name!r} is not a registered standard name or alias."
+ + suggest(name, candidates)
+ + " Register it with FieldDictionary.register(FieldEntry(...)) or "
+ "add an alias with FieldDictionary.add_alias(...)."
+ )
+
+
+class UnmatchedImportError(CouplingError):
+ """A component advertises an import that no other component exports."""
+
+ def __init__(
+ self, component: str, field: str, available_exports: dict[str, list[str]]
+ ):
+ exports_flat = [f for fields in available_exports.values() for f in fields]
+ listing = (
+ "; ".join(
+ f"{comp} exports {', '.join(fields)}"
+ for comp, fields in available_exports.items()
+ if fields
+ )
+ or "no component exports anything"
+ )
+ super().__init__(
+ f"Component {component!r} imports {field!r} but no component exports it."
+ + suggest(field, exports_flat)
+ + f" Available exports: {listing}."
+ + " Add an alias, a Mediator producing the derived field, or a "
+ "DataComponent supplying it from a data source."
+ )
+
+
+class UnitsMismatchError(CouplingError):
+ """Matched fields disagree on units."""
+
+ def __init__(self, field: str, src: str, src_units: str, dst: str, dst_units: str):
+ super().__init__(
+ f"Field {field!r}: {src!r} exports units {src_units!r} but {dst!r} "
+ f"expects {dst_units!r}. Unit conversion is not performed in v1 — "
+ "convert in a Mediator or align the FieldDictionary entries."
+ )
+
+
+class IncompatibleFieldError(CouplingError):
+ """A connector could not reconcile matched fields (grid/units/vertical)."""
+
+
+class VerticalMismatchError(CouplingError):
+ """Source and destination vertical coordinates cannot be reconciled."""
+
+
+class CadenceError(CouplingError):
+ """Component or slot cadence does not align with the driver clock."""
+
+ def __init__(self, what: str, interval: str, dt: str):
+ super().__init__(
+ f"{what} interval {interval} is not a positive multiple of the "
+ f"driver clock dt {dt}. Choose a driver dt that divides every "
+ "component timestep (typically their GCD)."
+ )
+
+
+class AmbiguousCouplingError(CouplingError):
+ """couple() found more than one exporter for an imported field."""
+
+ def __init__(self, field: str, importer: str, exporters: list[str]):
+ super().__init__(
+ f"Import {field!r} of component {importer!r} is exported by multiple "
+ f"components: {', '.join(exporters)}. Auto-wiring cannot choose — "
+ "build the Driver explicitly with Connector(src, dst, fields=[...]) "
+ "or a run-sequence DSL."
+ )
+
+
+class SequenceError(CouplingError):
+ """A run sequence references unknown names or is otherwise invalid."""
diff --git a/earth2studio/nvcoupler/field.py b/earth2studio/nvcoupler/field.py
new file mode 100644
index 000000000..abb796134
--- /dev/null
+++ b/earth2studio/nvcoupler/field.py
@@ -0,0 +1,272 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Field and State: the exchange currency of the coupler (ESMF analogs).
+
+A :class:`Field` is one physical quantity — a torch tensor plus its
+CoordSystem, canonical identity (standard name + units), validity time, and
+optional mask / vertical-coordinate metadata. A :class:`State` is a named
+bag of Fields keyed by standard name; every component owns an import State
+and an export State, and Connectors move Fields between them.
+
+Field data stays a torch tensor end-to-end (never round-tripped through
+numpy) so autograd graphs survive the exchange — a hard requirement for
+coupled fine-tuning.
+"""
+
+from collections import OrderedDict
+from collections.abc import Iterable, Iterator, MutableMapping
+from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+
+from earth2studio.utils.coords import cat_coords, split_coords
+from earth2studio.utils.type import CoordSystem
+
+from .dictionary import FieldDictionary
+from .errors import CouplingError
+
+if TYPE_CHECKING:
+ from .vertical import VerticalCoordinate
+
+# Dims regarded as spatial when choosing where to (re)insert a variable axis.
+# "point" is a scattered sample-location dim (see .points.PointSet) rather
+# than a mesh dim, but it plays the same role here: it marks where a Field's
+# spatial content lives.
+_SPATIAL_DIMS = (
+ "level",
+ "face",
+ "lat",
+ "lon",
+ "hpx",
+ "height",
+ "width",
+ "y",
+ "x",
+ "point",
+)
+
+
+@dataclass
+class Field:
+ """One exchanged quantity.
+
+ Parameters
+ ----------
+ data : torch.Tensor
+ Field values; dimension order given by `coords` insertion order.
+ Must NOT contain a "variable" dimension — a Field is one variable.
+ coords : CoordSystem
+ earth2studio coordinate dictionary describing `data`.
+ standard_name : str
+ Canonical name from the FieldDictionary.
+ units : str
+ Units of `data` (checked, not converted, in v1).
+ valid_time : np.datetime64, optional
+ Time the data is valid for.
+ source : str, optional
+ Name of the producing component (provenance).
+ mask : torch.Tensor, optional
+ Boolean validity mask broadcastable to `data` (True = valid),
+ e.g. ocean points for SST.
+ vertical : VerticalCoordinate, optional
+ Vertical coordinate description when `coords` contains a "level"
+ dimension (see :mod:`earth2studio.nvcoupler.vertical`).
+ """
+
+ data: torch.Tensor
+ coords: CoordSystem
+ standard_name: str
+ units: str
+ valid_time: np.datetime64 | None = None
+ source: str | None = None
+ mask: torch.Tensor | None = None
+ vertical: "VerticalCoordinate | None" = None
+
+ def __post_init__(self) -> None:
+ if "variable" in self.coords:
+ raise CouplingError(
+ f"Field {self.standard_name!r} coords must not contain a "
+ "'variable' dimension; use State.from_tensor to split a "
+ "multi-variable tensor into Fields"
+ )
+ ndim_coords = len(self.coords)
+ if self.data.ndim != ndim_coords:
+ raise CouplingError(
+ f"Field {self.standard_name!r}: data has {self.data.ndim} dims "
+ f"but coords describe {ndim_coords} "
+ f"({list(self.coords.keys())})"
+ )
+
+ def to(self, device: Any) -> "Field":
+ return replace(
+ self,
+ data=self.data.to(device),
+ mask=self.mask.to(device) if self.mask is not None else None,
+ )
+
+ def clone(self) -> "Field":
+ return replace(
+ self,
+ data=self.data.clone(),
+ coords=OrderedDict({k: v.copy() for k, v in self.coords.items()}),
+ mask=self.mask.clone() if self.mask is not None else None,
+ )
+
+ def grid_signature(self) -> tuple:
+ """Hashable signature of the spatial grid, for regridder caching."""
+ parts: list[tuple] = []
+ for key, value in self.coords.items():
+ if key in _SPATIAL_DIMS:
+ parts.append((key, value.shape, value.tobytes()))
+ return tuple(parts)
+
+ def __repr__(self) -> str:
+ dims = ", ".join(
+ f"{k}: {len(v) if v.ndim else 0}" for k, v in self.coords.items()
+ )
+ t = f", valid_time={self.valid_time}" if self.valid_time is not None else ""
+ return f"Field({self.standard_name!r} [{self.units}], {dims}{t})"
+
+
+class State(MutableMapping):
+ """A named collection of Fields keyed by standard name (ESMF_State analog)."""
+
+ def __init__(self, name: str, fields: Iterable[Field] = ()):
+ self.name = name
+ self._fields: dict[str, Field] = {}
+ for f in fields:
+ self.add(f)
+
+ # -- MutableMapping interface -------------------------------------------
+ def __getitem__(self, key: str) -> Field:
+ try:
+ return self._fields[key]
+ except KeyError:
+ raise KeyError(
+ f"State {self.name!r} has no field {key!r}; "
+ f"present: {sorted(self._fields)}"
+ ) from None
+
+ def __setitem__(self, key: str, value: Field) -> None:
+ if key != value.standard_name:
+ raise CouplingError(
+ f"State key {key!r} must equal the field's standard_name "
+ f"{value.standard_name!r}"
+ )
+ self._fields[key] = value
+
+ def __delitem__(self, key: str) -> None:
+ del self._fields[key]
+
+ def __iter__(self) -> Iterator[str]:
+ return iter(self._fields)
+
+ def __len__(self) -> int:
+ return len(self._fields)
+
+ # -- convenience ---------------------------------------------------------
+ def add(self, field: Field, replace: bool = True) -> None:
+ if not replace and field.standard_name in self._fields:
+ raise CouplingError(
+ f"Field {field.standard_name!r} already in state {self.name!r}"
+ )
+ self._fields[field.standard_name] = field
+
+ def subset(self, names: Iterable[str]) -> "State":
+ return State(self.name, (self[n] for n in names))
+
+ def to(self, device: Any) -> "State":
+ return State(self.name, (f.to(device) for f in self._fields.values()))
+
+ def as_tensor(
+ self, names: list[str] | None = None
+ ) -> tuple[torch.Tensor, CoordSystem]:
+ """Stack fields along a new "variable" dimension.
+
+ All selected fields must share identical coords (same grid); use a
+ Connector to bring fields onto one grid first. The variable axis is
+ inserted immediately before the first spatial dimension, matching the
+ earth2studio convention (batch, time, lead_time, variable, spatial...).
+ """
+ names = list(names) if names is not None else sorted(self._fields)
+ if not names:
+ raise CouplingError(f"State {self.name!r}: no fields to stack")
+ fields = [self[n] for n in names]
+ ref = fields[0].coords
+ dims = list(ref.keys())
+ insert_at = next(
+ (i for i, d in enumerate(dims) if d in _SPATIAL_DIMS), len(dims)
+ )
+ tensors, coord_list = [], []
+ for f in fields:
+ c = OrderedDict()
+ for i, (k, v) in enumerate(f.coords.items()):
+ if i == insert_at:
+ c["variable"] = np.array([f.standard_name])
+ c[k] = v
+ if "variable" not in c:
+ c["variable"] = np.array([f.standard_name])
+ tensors.append(f.data.unsqueeze(insert_at))
+ coord_list.append(c)
+ # cat_coords validates all non-variable dims match across fields
+ return cat_coords(tuple(tensors), tuple(coord_list), dim="variable")
+
+ @classmethod
+ def from_tensor(
+ cls,
+ name: str,
+ x: torch.Tensor,
+ coords: CoordSystem,
+ dictionary: FieldDictionary,
+ valid_time: np.datetime64 | None = None,
+ source: str | None = None,
+ strict: bool = True,
+ ) -> "State":
+ """Split a multi-variable tensor into a State of Fields.
+
+ Raw variable names in ``coords["variable"]`` are resolved to standard
+ names (and canonical units) through the dictionary. Unknown names
+ raise unless ``strict=False``, in which case they are skipped.
+ """
+ if "variable" not in coords:
+ raise CouplingError(
+ f"from_tensor for state {name!r}: coords have no 'variable' dim"
+ )
+ tensors, reduced_coords, values = split_coords(x, coords, dim="variable")
+ state = cls(name)
+ for tensor, raw_name in zip(tensors, values):
+ if raw_name not in dictionary:
+ if strict:
+ dictionary.resolve(str(raw_name)) # raises UnknownFieldError
+ continue
+ entry = dictionary.resolve(str(raw_name))
+ state.add(
+ Field(
+ data=tensor,
+ coords=OrderedDict(reduced_coords),
+ standard_name=entry.standard_name,
+ units=entry.canonical_units,
+ valid_time=valid_time,
+ source=source,
+ )
+ )
+ return state
+
+ def __repr__(self) -> str:
+ return f"State({self.name!r}, fields={sorted(self._fields)})"
diff --git a/earth2studio/nvcoupler/mediator.py b/earth2studio/nvcoupler/mediator.py
new file mode 100644
index 000000000..9b4881829
--- /dev/null
+++ b/earth2studio/nvcoupler/mediator.py
@@ -0,0 +1,236 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Mediators: components that compute derived exchange fields.
+
+The NUOPC_Mediator analog. A Mediator sits between components of different
+cadence, accumulating fast-component fields as they arrive (every connector
+transfer into it) and, when its compute action runs, exporting a windowed reduction
+— the trailing 48 h mean an ocean model was trained on, a precipitation sum
+a flood model needs, a temperature max for impact indices.
+
+Reductions are running torch ops (add / maximum / minimum), so memory is one
+accumulator per field regardless of window length, and gradients flow
+through mean and sum (max/min propagate to the extremal sample).
+"""
+
+from collections import OrderedDict
+from typing import Any
+
+import numpy as np
+import torch
+
+from .component import Component
+from .dictionary import CellMethod
+from .errors import CouplingError
+from .field import Field, State
+
+
+class _RunningReduction:
+ """Running windowed reduction shared by Mediator and windowed Connector.
+
+ One accumulator tensor per key regardless of window length; mean divides
+ by the sample count at emit time. Wire in the same valid_time twice
+ (e.g. two connectors or a re-executed slot) and the second arrival is
+ ignored rather than double-counted. Pure torch ops, so gradients flow
+ through mean and sum (max/min propagate to the extremal sample).
+ """
+
+ def __init__(self) -> None:
+ self._acc: dict[str, torch.Tensor] = {}
+ self._count: dict[str, int] = {}
+ self._coords: dict[str, OrderedDict] = {}
+ self._last_time: dict[str, np.datetime64] = {}
+
+ def add(self, key: str, field: Field, method: str) -> None:
+ """Fold one field into the running reduction under `key`."""
+ if field.valid_time is not None and self._last_time.get(key) == (
+ field.valid_time
+ ):
+ return # duplicate arrival for the same time
+ self._last_time[key] = field.valid_time
+ if key not in self._acc:
+ self._acc[key] = field.data
+ self._count[key] = 1
+ else:
+ acc = self._acc[key]
+ if method in ("mean", "sum"):
+ self._acc[key] = acc + field.data
+ elif method == "max":
+ self._acc[key] = torch.maximum(acc, field.data)
+ else: # min
+ self._acc[key] = torch.minimum(acc, field.data)
+ self._count[key] += 1
+ self._coords[key] = OrderedDict(field.coords)
+
+ def __contains__(self, key: str) -> bool:
+ return key in self._acc
+
+ def counts(self) -> dict[str, int]:
+ """Samples folded in per key since the last reset."""
+ return dict(self._count)
+
+ def emit(self, key: str, method: str) -> tuple[torch.Tensor, OrderedDict]:
+ """Reduced (data, coords) for `key`; mean divides by the count."""
+ data = self._acc[key]
+ if method == "mean":
+ data = data / self._count[key]
+ return data, self._coords[key]
+
+ def reset(self) -> None:
+ self._acc.clear()
+ self._count.clear()
+ self._coords.clear()
+ self._last_time.clear()
+
+
+class _AccumulatingState(State):
+ """Import state that forwards every added field to the owning mediator."""
+
+ def __init__(self, name: str, mediator: "Mediator"):
+ super().__init__(name)
+ self._mediator = mediator
+
+ def add(self, field: Field, replace: bool = True) -> None:
+ super().add(field, replace=replace)
+ self._mediator.accumulate(field)
+
+
+class Mediator(Component):
+ """Base class: accumulate imports, reduce on ring.
+
+ Subclasses implement :meth:`accumulate` (called on every field arriving
+ in the import state) and :meth:`compute` (called when the mediator's
+ compute action runs in its slot; must populate ``export_state``).
+ """
+
+ requires_ic = False # mediators need no initial condition
+
+ def __init__(self, name: str, timestep: Any, imports=(), exports=(), **kwargs: Any):
+ super().__init__(name, timestep, imports, exports, **kwargs)
+ self.import_state = _AccumulatingState(f"{name}.imports", self)
+
+ def initialize(self, x: torch.Tensor | None = None, coords=None) -> None:
+ """Mediators need no initial condition."""
+
+ def accumulate(self, field: Field) -> None:
+ raise NotImplementedError
+
+ def compute(self, time: np.datetime64) -> None:
+ raise NotImplementedError
+
+ def run(self, time: np.datetime64) -> None:
+ self.compute(time)
+ self.run_count += 1
+
+
+class AccumulationMediator(Mediator):
+ """Windowed reduction of fast-component fields onto a slow cadence.
+
+ Parameters
+ ----------
+ name : str
+ fields : list[str]
+ *Derived* standard names to produce (e.g.
+ ``geopotential_at_1000hpa_48h_mean``). Each must be a dictionary
+ entry carrying a :class:`CellMethod`; the cell method supplies the
+ base field to import, the reduction, and the window (= the
+ mediator's timestep unless ``window`` overrides it).
+ window : optional
+ Override the reduction window; defaults to the (common) cell-method
+ window of `fields`.
+
+ This is the generalization of PhysicsNeMo's TrailingAverageCoupler and
+ DLESyM's ``_make_ocean_coupling`` chunk-mean, plus the sum/max/min
+ reductions impact chains need.
+ """
+
+ def __init__(self, name: str, fields: list[str], window: Any = None, **kwargs: Any):
+ from .dictionary import DEFAULT_DICTIONARY
+
+ dictionary = kwargs.get("dictionary") or DEFAULT_DICTIONARY
+ methods: dict[str, CellMethod] = {}
+ for derived in fields:
+ entry = dictionary.resolve(derived)
+ if entry.cell_method is None:
+ raise CouplingError(
+ f"AccumulationMediator {name!r}: {derived!r} has no "
+ "cell_method in the field dictionary — register a "
+ "FieldEntry(cell_method=CellMethod(base, method, window)) "
+ "describing how it derives from a base field"
+ )
+ methods[entry.standard_name] = entry.cell_method
+ windows = {cm.window for cm in methods.values()}
+ if window is None:
+ if len(windows) != 1:
+ raise CouplingError(
+ f"AccumulationMediator {name!r}: fields have differing "
+ f"windows {sorted(str(w) for w in windows)}; split them "
+ "across mediators or pass window= explicitly"
+ )
+ window = next(iter(windows))
+ # dedupe: several derived fields may reduce the same base import
+ imports = list(dict.fromkeys(cm.base for cm in methods.values()))
+ super().__init__(name, window, imports=imports, exports=list(methods), **kwargs)
+ self.methods = methods # derived std name -> CellMethod
+ # base std name -> ALL derived fields reducing it (e.g. the 24h max
+ # and 24h mean of t2m accumulate from the same delivered field)
+ self._base_to_derived: dict[str, list[str]] = {}
+ for derived, cm in methods.items():
+ self._base_to_derived.setdefault(cm.base, []).append(derived)
+ self._reduction = _RunningReduction()
+
+ def accumulate(self, field: Field) -> None:
+ for derived in self._base_to_derived.get(field.standard_name, ()):
+ self._reduction.add(derived, field, self.methods[derived].method)
+
+ def compute(self, time: np.datetime64) -> None:
+ for derived, cm in self.methods.items():
+ if derived not in self._reduction:
+ raise CouplingError(
+ f"Mediator {self.name!r}: no samples of {cm.base!r} "
+ f"accumulated before compute at {time} — is a connector "
+ "feeding this mediator in a faster slot?"
+ )
+ data, coords = self._reduction.emit(derived, cm.method)
+ entry = self.dictionary.resolve(derived)
+ self.export_state.add(
+ Field(
+ data=data,
+ coords=coords,
+ standard_name=derived,
+ units=entry.canonical_units,
+ valid_time=time,
+ source=self.name,
+ )
+ )
+ self.samples_last_window = self._reduction.counts()
+ self._reduction.reset()
+
+
+class TrailingAverageMediator(AccumulationMediator):
+ """AccumulationMediator restricted to mean reductions — the exact
+ semantics of DLESyM's ocean coupling and PhysicsNeMo's
+ TrailingAverageCoupler."""
+
+ def __init__(self, name: str, fields: list[str], window: Any = None, **kwargs: Any):
+ super().__init__(name, fields, window, **kwargs)
+ bad = [f for f, cm in self.methods.items() if cm.method != "mean"]
+ if bad:
+ raise CouplingError(
+ f"TrailingAverageMediator {name!r}: fields {bad} are not mean "
+ "reductions — use AccumulationMediator"
+ )
diff --git a/earth2studio/nvcoupler/points.py b/earth2studio/nvcoupler/points.py
new file mode 100644
index 000000000..c32c1abd8
--- /dev/null
+++ b/earth2studio/nvcoupler/points.py
@@ -0,0 +1,99 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""PointSet: an irregular collection of (lat, lon) sample locations.
+
+The "point" analog of :mod:`.vertical` — a component whose grid is a set of
+scattered coordinates (stations, sites, asset locations, arbitrary query
+points) rather than a lat/lon mesh. A component targeting points advertises
+a ``"point"`` dim in its :meth:`Component.grid_coords` (an index or name
+array, following the FieldDictionary convention that a coords value is the
+dim's own coordinate labels) and carries the actual locations on
+``Component.points``; the Connector's ``sample=`` path reads the latter to
+build an auto sampler (see ``connector.py``).
+"""
+
+from collections import OrderedDict
+from dataclasses import dataclass
+
+import numpy as np
+
+from earth2studio.utils.type import CoordSystem
+
+from .errors import CouplingError
+
+
+@dataclass(frozen=True)
+class PointSet:
+ """A fixed set of N sample locations.
+
+ Parameters
+ ----------
+ lat, lon : np.ndarray [N]
+ Latitude/longitude of each point, degrees.
+ names : tuple[str, ...], optional
+ Point identifiers (station IDs, site names, ...). Defaults to an
+ integer index ``0..N-1`` when omitted, which is what the "point" dim
+ coordinate carries in that case.
+ """
+
+ lat: np.ndarray
+ lon: np.ndarray
+ names: tuple[str, ...] | None = None
+
+ def __post_init__(self) -> None:
+ lat = np.asarray(self.lat, dtype=np.float64)
+ lon = np.asarray(self.lon, dtype=np.float64)
+ if lat.ndim != 1 or lon.ndim != 1:
+ raise CouplingError(
+ f"PointSet: lat/lon must be 1-D, got shapes {self.lat.shape} "
+ f"and {self.lon.shape}"
+ )
+ if lat.shape != lon.shape:
+ raise CouplingError(
+ f"PointSet: lat and lon must have the same length, got "
+ f"{lat.shape[0]} and {lon.shape[0]}"
+ )
+ if lat.shape[0] == 0:
+ raise CouplingError("PointSet: at least one point is required")
+ if self.names is not None and len(self.names) != lat.shape[0]:
+ raise CouplingError(
+ f"PointSet: names has {len(self.names)} entries but lat/lon "
+ f"have {lat.shape[0]}"
+ )
+ object.__setattr__(self, "lat", lat)
+ object.__setattr__(self, "lon", lon)
+
+ def __len__(self) -> int:
+ return self.lat.shape[0]
+
+ def labels(self) -> np.ndarray:
+ """The 'point' dim's own coordinate array: names if given, else
+ integer index — matches every other dim in a CoordSystem carrying
+ its own labels."""
+ if self.names is not None:
+ return np.array(self.names)
+ return np.arange(len(self))
+
+ def grid_coords(self) -> CoordSystem:
+ """A one-entry CoordSystem, the point-grid analog of a lat/lon
+ mesh's coords, for :meth:`Component.grid_coords`."""
+ return OrderedDict({"point": self.labels()})
+
+ def signature(self) -> tuple:
+ """Hashable signature for sampler caching (mirrors
+ Field.grid_signature)."""
+ return (self.lat.shape, self.lat.tobytes(), self.lon.tobytes())
diff --git a/earth2studio/nvcoupler/pull.py b/earth2studio/nvcoupler/pull.py
new file mode 100644
index 000000000..77438dbf6
--- /dev/null
+++ b/earth2studio/nvcoupler/pull.py
@@ -0,0 +1,199 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Pull-pattern coupling: adapters for models that fetch their own forcing.
+
+Some earth2studio models (StormCast is the canonical case) do not accept
+coupled fields as call arguments — they *pull* them, calling
+``fetch_data(self.conditioning_data_source, time, variables, ...)`` inside
+their own ``__call__``. The only injection point such a model exposes is the
+settable data-source attribute.
+
+:class:`PullAdapter` exploits exactly that point, without wrapping or
+modifying the model: before each step it installs a :class:`StateDataSource`
+— a tiny in-memory object satisfying the DataSource protocol that answers
+fetches from the component's import State. The model runs its unmodified
+production code path (fetch → interpolate → concatenate) and believes it is
+reading GFS; it is reading the coupler. This is the same masquerade the
+existing serve workflows play with ``InferenceOutputSource``, minus the
+store-and-replay staging: the "source" is this step's live exchange.
+
+Honest limitation: the pull path goes through the model's own
+``fetch_data``/xarray machinery, so field data crosses a numpy boundary —
+pull-coupled components are inference-only (no gradients through the
+exchange). Push-pattern adapters keep autograd intact.
+"""
+
+from collections import OrderedDict
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+import xarray as xr
+
+from earth2studio.utils.type import CoordSystem
+
+from .errors import CouplingError
+from .field import State
+
+if TYPE_CHECKING:
+ from .component import Exchange
+ from .dictionary import FieldDictionary
+
+# Fetch times farther than this from a served field's valid_time trigger a
+# warning: the pull is likely misaligned with the exchange cadence.
+DEFAULT_TIME_TOLERANCE = np.timedelta64(0, "h")
+
+
+class StateDataSource:
+ """A DataSource serving the current contents of an import State.
+
+ Answers ``__call__(time, variable)`` with an ``xr.DataArray`` of dims
+ ``(time, variable, lat, lon)`` built from the State's Fields. Variables
+ may be requested by the model's raw names (resolved through the owning
+ component's dictionary/aliases) or by standard name.
+
+ The source is a snapshot view: whatever the connector last delivered is
+ what every requested time receives. Cadence alignment is the run
+ sequence's job — a sequential ``global -> stormcast`` connect before the
+ pull guarantees the served fields are valid at the pulled time.
+ """
+
+ def __init__(
+ self,
+ state: State,
+ raw_to_std: Mapping[str, str] | None = None,
+ strict_time: bool = False,
+ dictionary: "FieldDictionary | None" = None,
+ ):
+ from .dictionary import DEFAULT_DICTIONARY
+
+ self.state = state
+ self.raw_to_std = dict(raw_to_std or {})
+ self.strict_time = strict_time
+ self.dictionary = dictionary or DEFAULT_DICTIONARY
+
+ def _resolve(self, name: str) -> str:
+ if name in self.state:
+ return name
+ if name in self.raw_to_std and self.raw_to_std[name] in self.state:
+ return self.raw_to_std[name]
+ # pulled conditioning names are usually model-raw (u10m, t2m): the
+ # Exchange map only covers state variables, so fall back to aliases
+ if name in self.dictionary:
+ std = self.dictionary.standard_name(name)
+ if std in self.state:
+ return std
+ raise CouplingError(
+ f"Pull-coupled model requested variable {name!r}, but the import "
+ f"state holds {sorted(self.state)} (raw-name map: "
+ f"{self.raw_to_std}). Add the field to the component's imports "
+ "and wire a connector delivering it before the model runs."
+ )
+
+ def __call__(self, time: Any, variable: Any) -> xr.DataArray:
+ times = np.atleast_1d(np.asarray(time, dtype="datetime64[ns]"))
+ variables = np.atleast_1d(np.asarray(variable))
+ fields = [self.state[self._resolve(str(v))] for v in variables]
+
+ grid = fields[0].coords
+ if list(grid.keys()) != ["lat", "lon"]:
+ raise CouplingError(
+ "StateDataSource serves (lat, lon) fields; got dims "
+ f"{list(grid.keys())} for {fields[0].standard_name!r} — "
+ "exchange-shaped Fields are expected (leading singleton dims "
+ "are squeezed by the publishing component)"
+ )
+ for f in fields:
+ if self.strict_time and f.valid_time is not None:
+ if any(t != f.valid_time for t in times):
+ raise CouplingError(
+ f"Pull for {f.standard_name!r} at {times} but the "
+ f"served field is valid at {f.valid_time} — check the "
+ "run-sequence ordering (the connector must run before "
+ "the pulling component in the same slot)"
+ )
+
+ # IO boundary of the pull path: the model's own fetch machinery is
+ # xarray-based, so this conversion is unavoidable (inference-only).
+ data = np.stack([f.data.detach().cpu().numpy() for f in fields], axis=0)[
+ np.newaxis
+ ].repeat(len(times), axis=0)
+ return xr.DataArray(
+ data,
+ dims=["time", "variable", "lat", "lon"],
+ coords={
+ "time": times,
+ "variable": variables,
+ "lat": np.asarray(grid["lat"]),
+ "lon": np.asarray(grid["lon"]),
+ },
+ )
+
+
+class PullAdapter:
+ """ImportAdapter for pull-pattern models (StormCast-style).
+
+ Before each model call, installs a :class:`StateDataSource` over the
+ current import State on the model's data-source attribute
+ (``conditioning_data_source`` by default), then calls
+ ``model(x, coords)`` unchanged. The model's internal fetch receives this
+ step's coupled forcing.
+
+ Parameters
+ ----------
+ attribute : str
+ Name of the model's settable data-source attribute.
+ strict_time : bool
+ Raise if the model pulls times that do not match the served fields'
+ valid_time (default False: serve the snapshot and let the run
+ sequence own alignment).
+ """
+
+ def __init__(
+ self,
+ attribute: str = "conditioning_data_source",
+ strict_time: bool = False,
+ dictionary: "FieldDictionary | None" = None,
+ ):
+ self.attribute = attribute
+ self.strict_time = strict_time
+ self.dictionary = dictionary
+
+ def __call__(
+ self, model: Any, exchange: "Exchange"
+ ) -> tuple[torch.Tensor, "CoordSystem"]:
+ if not hasattr(model, self.attribute):
+ raise CouplingError(
+ f"PullAdapter: model {type(model).__name__!r} has no "
+ f"attribute {self.attribute!r} — this adapter is for models "
+ "that fetch forcing from a settable data source (e.g. "
+ "StormCast's conditioning_data_source). For models taking "
+ "conditioning as an argument use ConditioningKwargAdapter."
+ )
+ raw_to_std = {raw: std for std, raw in exchange.std_to_raw.items()}
+ setattr(
+ model,
+ self.attribute,
+ StateDataSource(
+ exchange.imports,
+ raw_to_std,
+ strict_time=self.strict_time,
+ dictionary=self.dictionary,
+ ),
+ )
+ return model(exchange.x, OrderedDict(exchange.coords))
diff --git a/earth2studio/nvcoupler/sequence.py b/earth2studio/nvcoupler/sequence.py
new file mode 100644
index 000000000..2a41ea284
--- /dev/null
+++ b/earth2studio/nvcoupler/sequence.py
@@ -0,0 +1,358 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Run sequence: the ordered schedule of component runs and exchanges.
+
+The NUOPC runSeq analog. A RunSequence is a list of Slots, each with an
+interval and an ordered action list; the Driver executes a slot's actions,
+in order, at every clock time aligned with the slot interval. Coupling
+semantics are pure ordering: a ConnectAction placed *before* the destination
+component's RunAction in the same slot means the destination sees the
+source's previous state (lagged / NUOPC-explicit coupling); placed *after*,
+it sees the state just produced (sequential coupling).
+
+Sequences are derived, not required: :func:`derive_sequence` lays out the
+canonical schedule from the coupling graph (components + connections) alone,
+and the Driver calls it when no sequence is given. Hand-written sequences —
+via the string DSL mirroring NUOPC's runSeq — are the override for schedules
+the graph cannot express::
+
+ @6h
+ atmos -> med # ConnectAction (exports of atmos -> imports of med)
+ ocean -> atmos # lagged: before atmos runs
+ atmos # RunAction
+ @48h
+ med.compute # MediateAction
+ med -> ocean
+ ocean
+ @
+"""
+
+import re
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Literal
+
+import numpy as np
+
+from .clock import DeltaLike, as_timedelta, fmt_timedelta, is_multiple
+from .errors import CadenceError, SequenceError, suggest
+
+if TYPE_CHECKING:
+ from .component import Component
+ from .connector import Connector
+
+
+@dataclass(frozen=True)
+class RunAction:
+ component: str
+
+ def __str__(self) -> str:
+ return self.component
+
+
+@dataclass(frozen=True)
+class ConnectAction:
+ src: str
+ dst: str
+
+ def __str__(self) -> str:
+ return f"{self.src} -> {self.dst}"
+
+
+@dataclass(frozen=True)
+class MediateAction:
+ mediator: str
+ phase: str = "compute"
+
+ def __str__(self) -> str:
+ return f"{self.mediator}.{self.phase}"
+
+
+Action = RunAction | ConnectAction | MediateAction
+
+
+def _fmt_interval(d: np.timedelta64) -> str:
+ """Format a slot interval for the DSL: whole hours as NUOPC-style '@48h',
+ otherwise fall back to fmt_timedelta so sub-hour slots ('@90m', '@30m')
+ round-trip exactly instead of truncating to whole hours."""
+ ns = int(d.astype("timedelta64[ns]").astype(np.int64))
+ hour = 3_600_000_000_000
+ if ns % hour == 0:
+ return f"{ns // hour}h"
+ return fmt_timedelta(d)
+
+
+@dataclass
+class Slot:
+ interval: np.timedelta64
+ actions: list[Action] = field(default_factory=list)
+
+ def __post_init__(self) -> None:
+ self.interval = as_timedelta(self.interval)
+
+
+@dataclass
+class RunSequence:
+ slots: list[Slot]
+
+ def components_run(self) -> set[str]:
+ return {
+ a.component
+ for s in self.slots
+ for a in s.actions
+ if isinstance(a, RunAction)
+ } | {
+ a.mediator
+ for s in self.slots
+ for a in s.actions
+ if isinstance(a, MediateAction)
+ }
+
+ def connections(self) -> list[ConnectAction]:
+ return [
+ a for s in self.slots for a in s.actions if isinstance(a, ConnectAction)
+ ]
+
+ def validate(self, components: dict, dt: DeltaLike) -> None:
+ """Check name resolution, cadence alignment, and completeness."""
+ dt = as_timedelta(dt)
+ names = set(components)
+
+ def check_name(name: str, what: str) -> None:
+ if name not in names:
+ raise SequenceError(
+ f"Run sequence references unknown {what} {name!r}."
+ + suggest(name, names)
+ + f" Known components: {sorted(names)}"
+ )
+
+ for slot in self.slots:
+ if not is_multiple(slot.interval, dt):
+ raise CadenceError("Run-sequence slot", str(slot.interval), str(dt))
+ for action in slot.actions:
+ if isinstance(action, RunAction):
+ check_name(action.component, "component")
+ comp = components[action.component]
+ if comp.timestep != slot.interval:
+ raise CadenceError(
+ f"Component {action.component!r} (timestep "
+ f"{comp.timestep}) scheduled in a slot of",
+ str(slot.interval),
+ str(comp.timestep),
+ )
+ elif isinstance(action, ConnectAction):
+ check_name(action.src, "connector source")
+ check_name(action.dst, "connector destination")
+ else:
+ check_name(action.mediator, "mediator")
+ med = components[action.mediator]
+ if med.timestep != slot.interval:
+ raise CadenceError(
+ f"Mediator {action.mediator!r} (timestep "
+ f"{med.timestep}) scheduled in a slot of",
+ str(slot.interval),
+ str(med.timestep),
+ )
+ # every component must run somewhere
+ idle = names - self.components_run()
+ if idle:
+ raise SequenceError(
+ f"Components never run by the sequence: {sorted(idle)} — add a "
+ "RunAction (bare component name) to a slot matching their timestep"
+ )
+
+ def __str__(self) -> str:
+ lines = []
+ for slot in self.slots:
+ lines.append(f"@{_fmt_interval(slot.interval)}")
+ lines.extend(f" {a}" for a in slot.actions)
+ lines.append("@")
+ return "\n".join(lines)
+
+
+_SLOT_RE = re.compile(r"^@(\S+)?$")
+_CONNECT_RE = re.compile(r"^(\w[\w.-]*)\s*->\s*(\w[\w.-]*)$")
+_MEDIATE_RE = re.compile(r"^(\w[\w-]*)\.(\w+)$")
+_RUN_RE = re.compile(r"^(\w[\w-]*)$")
+
+
+def parse_run_sequence(text: str) -> RunSequence:
+ """Parse the NUOPC-flavored DSL into a RunSequence."""
+ slots: list[Slot] = []
+ current: Slot | None = None
+ for lineno, raw in enumerate(text.splitlines(), start=1):
+ line = raw.split("#", 1)[0].strip()
+ if not line:
+ continue
+ if m := _SLOT_RE.match(line):
+ if m.group(1): # "@6h" opens a slot; bare "@" closes the sequence
+ try:
+ current = Slot(as_timedelta(m.group(1)))
+ except ValueError as e:
+ raise SequenceError(f"Line {lineno}: {e}") from None
+ slots.append(current)
+ else:
+ current = None
+ continue
+ if current is None:
+ raise SequenceError(
+ f"Line {lineno}: action {line!r} outside any @interval slot"
+ )
+ if m := _CONNECT_RE.match(line):
+ current.actions.append(ConnectAction(m.group(1), m.group(2)))
+ elif m := _MEDIATE_RE.match(line):
+ current.actions.append(MediateAction(m.group(1), m.group(2)))
+ elif m := _RUN_RE.match(line):
+ current.actions.append(RunAction(m.group(1)))
+ else:
+ raise SequenceError(
+ f"Line {lineno}: cannot parse {line!r} — expected 'name', "
+ "'src -> dst', or 'mediator.compute'"
+ )
+ if not slots:
+ raise SequenceError("Run sequence is empty")
+ return RunSequence(slots)
+
+
+# ---------------------------------------------------------------------------
+# derive_sequence(): the schedule implied by the coupling graph
+# ---------------------------------------------------------------------------
+def _toposort(
+ nodes: list[str],
+ edges: list[tuple[str, str]],
+ interval: np.timedelta64,
+) -> list[str]:
+ """Order `nodes` so every sequential edge's source runs before its
+ destination, keeping declaration order among unconstrained nodes."""
+ node_set = set(nodes)
+ deps: dict[str, set[str]] = {n: set() for n in nodes}
+ for src, dst in edges:
+ if src in node_set and dst in node_set:
+ deps[dst].add(src)
+ out: list[str] = []
+ placed: set[str] = set()
+ pending = list(nodes)
+ while pending:
+ ready = [n for n in pending if deps[n] <= placed]
+ if not ready:
+ cycle_edges = [(s, d) for s, d in edges if s in pending and d in pending]
+ raise SequenceError(
+ f"Sequential coupling cycle among components {sorted(pending)} "
+ f"at cadence {fmt_timedelta(interval)}: connections "
+ f"{cycle_edges} require each destination to see state its "
+ "source produces in the same step, so no run order exists. "
+ f"Mark one edge lagged (e.g. lagged={{{cycle_edges[0]!r}}}) "
+ "or pass an explicit run sequence"
+ )
+ out.extend(ready)
+ placed.update(ready)
+ pending = [n for n in pending if n not in placed]
+ return out
+
+
+def derive_sequence(
+ components: "dict[str, Component]",
+ connectors: "Iterable[Connector | tuple[str, str]] | None" = None,
+ lagged: "set[tuple[str, str]] | Literal['all']" = "all",
+) -> RunSequence:
+ """Derive the canonical run sequence from the coupling graph.
+
+ One slot per distinct component cadence, fast to slow. Within a slot:
+
+ 1. lagged connects delivered at this cadence (the faster endpoint of
+ each connection) — destinations see the sources' previous exports;
+ 2. each mediator's compute followed by its outgoing connects (mediator
+ exports only exist after compute, so these are always sequential);
+ 3. component runs, ordered by the sequential (non-lagged) connections
+ among them, each run followed by its outgoing sequential connects.
+
+ Parameters
+ ----------
+ components : dict[str, Component]
+ All participants, mediators included.
+ connectors : iterable of Connector or (src, dst) name tuples, optional
+ The coupling graph's edges.
+ lagged : set[tuple[str, str]] | "all"
+ Connections whose destination consumes the source's *previous* state
+ (NUOPC-explicit coupling). Default "all" — the canonical shape.
+ Connections not in the set are sequential; a cycle of sequential
+ connections among same-cadence components raises SequenceError.
+ """
+ from .mediator import Mediator
+
+ names = set(components)
+ pairs: list[tuple[str, str]] = []
+ for item in connectors or []:
+ pair = (
+ (item.src.name, item.dst.name)
+ if hasattr(item, "src")
+ else (str(item[0]), str(item[1]))
+ )
+ if pair not in pairs:
+ pairs.append(pair)
+ for src, dst in pairs:
+ for name, what in ((src, "connection source"), (dst, "connection destination")):
+ if name not in names:
+ raise SequenceError(
+ f"Coupling graph references unknown {what} {name!r}."
+ + suggest(name, names)
+ + f" Known components: {sorted(names)}"
+ )
+
+ def is_lagged(pair: tuple[str, str]) -> bool:
+ return lagged == "all" or pair in lagged
+
+ def ns(td: np.timedelta64) -> int:
+ return int(td.astype("timedelta64[ns]").astype(np.int64))
+
+ order = {name: i for i, name in enumerate(components)}
+ cadence = {name: comp.timestep for name, comp in components.items()}
+ is_mediator = {
+ name: isinstance(comp, Mediator) for name, comp in components.items()
+ }
+ seq_edges = [p for p in pairs if not is_lagged(p) and not is_mediator[p[0]]]
+
+ slots: list[Slot] = []
+ for interval in sorted({cadence[n] for n in components}, key=ns):
+ here = [n for n in components if cadence[n] == interval]
+ actions: list[Action] = []
+ # 1. lagged connects delivered at this cadence (the faster endpoint),
+ # excluding mediator-sourced ones (delivered after compute below)
+ block = [
+ p
+ for p in pairs
+ if is_lagged(p)
+ and not is_mediator[p[0]]
+ and min(ns(cadence[p[0]]), ns(cadence[p[1]])) == ns(interval)
+ ]
+ block.sort(key=lambda p: (order[p[0]], order[p[1]]))
+ actions.extend(ConnectAction(s, d) for s, d in block)
+ # 2. mediators of this cadence: compute, then deliver
+ for med in (n for n in here if is_mediator[n]):
+ actions.append(MediateAction(med))
+ actions.extend(ConnectAction(s, d) for s, d in pairs if s == med)
+ # 3. runs in dependency order over the sequential connections, each
+ # followed by its outgoing sequential connects
+ runs = [n for n in here if not is_mediator[n]]
+ for name in _toposort(runs, seq_edges, interval):
+ actions.append(RunAction(name))
+ actions.extend(ConnectAction(s, d) for s, d in seq_edges if s == name)
+ if actions:
+ slots.append(Slot(interval, actions))
+ if not slots:
+ raise SequenceError("Cannot derive a run sequence from an empty component dict")
+ return RunSequence(slots)
diff --git a/earth2studio/nvcoupler/testing.py b/earth2studio/nvcoupler/testing.py
new file mode 100644
index 000000000..a47510fa3
--- /dev/null
+++ b/earth2studio/nvcoupler/testing.py
@@ -0,0 +1,120 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Synthetic toy components for tests and demos.
+
+A deterministic two-component system with the DLESyM cadence structure:
+
+- fake atmosphere, 6 h step, 32x64 lat/lon grid, state [z1000, sst]:
+ z1000 <- z1000 + 1.0 + gain * 0.1 * sst (sst = imported SST)
+- fake ocean, 48 h step, 16x32 lat/lon grid, state [sst, z48m]:
+ sst <- sst + gain * 0.01 * z48m (z48m = imported 48 h
+ mean of atmos z1000)
+
+With spatially-constant initial conditions every intermediate value is
+hand-computable, which the end-to-end driver tests rely on. Pass
+``gain=torch.tensor(1.0, requires_grad=True)`` to check that gradients flow
+across the exchange.
+"""
+
+from collections import OrderedDict
+
+import numpy as np
+import torch
+
+from earth2studio.utils.type import CoordSystem
+
+from .component import CallableComponent
+
+ATMOS_GRID = (32, 64)
+OCEAN_GRID = (16, 32)
+
+
+def grid_coords(nlat: int, nlon: int) -> CoordSystem:
+ return OrderedDict(
+ {
+ "lat": np.linspace(90.0, -90.0, nlat),
+ "lon": np.linspace(0.0, 360.0, nlon, endpoint=False),
+ }
+ )
+
+
+def fake_atmos(
+ gain: torch.Tensor | float = 1.0, timestep: str = "6h"
+) -> CallableComponent:
+ """Fast toy component: imports SST, exports z1000."""
+
+ def step(x: torch.Tensor, coords: CoordSystem):
+ z1000, sst = x[0], x[1]
+ z_next = z1000 + 1.0 + gain * 0.1 * sst
+ return torch.stack([z_next, sst]), coords
+
+ return CallableComponent(
+ "atmos",
+ step,
+ timestep=timestep,
+ imports=["sea_surface_temperature"],
+ exports=["geopotential_at_1000hpa"],
+ )
+
+
+def fake_ocean(
+ gain: torch.Tensor | float = 1.0,
+ timestep: str = "48h",
+ with_mask: bool = False,
+) -> CallableComponent:
+ """Slow toy component: imports the 48 h mean of z1000, exports SST.
+
+ With ``with_mask=True`` the exported SST carries a land mask covering the
+ northern half of the grid (True = valid ocean point).
+ """
+
+ def step(x: torch.Tensor, coords: CoordSystem):
+ sst, z48m = x[0], x[1]
+ sst_next = sst + gain * 0.01 * z48m
+ return torch.stack([sst_next, z48m]), coords
+
+ export_masks = None
+ if with_mask:
+ mask = torch.ones(*OCEAN_GRID, dtype=torch.bool)
+ mask[: OCEAN_GRID[0] // 2, :] = False # northern half is land
+ export_masks = {"sea_surface_temperature": mask}
+
+ return CallableComponent(
+ "ocean",
+ step,
+ timestep=timestep,
+ imports=["geopotential_at_1000hpa_48h_mean"],
+ exports=["sea_surface_temperature"],
+ variable_aliases={"z48m": "geopotential_at_1000hpa_48h_mean"},
+ export_masks=export_masks,
+ )
+
+
+def atmos_ic(z0: float = 0.0, sst0: float = 2.0) -> tuple[torch.Tensor, CoordSystem]:
+ coords = OrderedDict(
+ {"variable": np.array(["z1000", "sst"]), **grid_coords(*ATMOS_GRID)}
+ )
+ x = torch.stack([torch.full(ATMOS_GRID, z0), torch.full(ATMOS_GRID, sst0)])
+ return x, coords
+
+
+def ocean_ic(sst0: float = 2.0, z48m0: float = 0.0) -> tuple[torch.Tensor, CoordSystem]:
+ coords = OrderedDict(
+ {"variable": np.array(["sst", "z48m"]), **grid_coords(*OCEAN_GRID)}
+ )
+ x = torch.stack([torch.full(OCEAN_GRID, sst0), torch.full(OCEAN_GRID, z48m0)])
+ return x, coords
diff --git a/earth2studio/nvcoupler/vertical.py b/earth2studio/nvcoupler/vertical.py
new file mode 100644
index 000000000..8e128b049
--- /dev/null
+++ b/earth2studio/nvcoupler/vertical.py
@@ -0,0 +1,189 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Vertical coordinate descriptors and differentiable level interpolation.
+
+Most earth2studio models encode pressure levels in variable names (z500,
+t850) and never see this module. It exists for components with an explicit
+"level" dimension — chiefly chemistry emulators on hybrid sigma-pressure
+model levels (p_k = a_k + b_k * p_s) coupled to met components on pressure
+levels. Interpolation is linear in log-pressure via torch.searchsorted and
+gathers, so gradients flow through it (values, not indices).
+"""
+
+from collections import OrderedDict
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+from earth2studio.utils.type import CoordSystem
+
+from .errors import VerticalMismatchError
+
+
+@dataclass(frozen=True)
+class PressureLevels:
+ """Constant pressure levels in hPa, ordered top to bottom (increasing)."""
+
+ levels: tuple[float, ...]
+
+ def __post_init__(self) -> None:
+ if list(self.levels) != sorted(self.levels):
+ raise ValueError("Pressure levels must be increasing (top to bottom)")
+
+ def pressure_pa(self) -> np.ndarray:
+ return np.asarray(self.levels, dtype=np.float64) * 100.0
+
+
+@dataclass(frozen=True)
+class HybridLevels:
+ """Hybrid sigma-pressure levels: p_k = a_k + b_k * p_s.
+
+ `a` in Pa, `b` dimensionless, ordered top to bottom. `ps_field` names the
+ surface-pressure field (Pa) required to realize the levels; a Connector
+ performing a hybrid->pressure transform takes it from the source
+ component's exports automatically.
+ """
+
+ a: tuple[float, ...]
+ b: tuple[float, ...]
+ ps_field: str = "surface_pressure"
+
+ def __post_init__(self) -> None:
+ if len(self.a) != len(self.b):
+ raise ValueError("Hybrid coefficients a and b must have equal length")
+ # Levels must be strictly increasing in pressure (top to bottom) for
+ # every plausible surface pressure, else interpolation would pair
+ # level slices with wrong pressures. p_k(ps) = a_k + b_k * ps is
+ # linear in ps, so strict monotonicity at both ends of the plausible
+ # Earth surface-pressure range [50000, 110000] Pa (high terrain to
+ # strong anticyclone) is sufficient for every ps inside that range;
+ # ps values outside it are re-checked at interpolation time.
+ for ps in (50000.0, 110000.0):
+ p = (
+ np.asarray(self.a, dtype=np.float64)
+ + np.asarray(self.b, dtype=np.float64) * ps
+ )
+ if np.any(np.diff(p) <= 0):
+ raise ValueError(
+ f"Hybrid coefficients a={list(self.a)}, b={list(self.b)} "
+ f"produce non-increasing pressures {p.tolist()} Pa at "
+ f"surface pressure {ps:.0f} Pa. Order a and b top to "
+ "bottom so p_k = a_k + b_k * ps strictly increases for "
+ "all surface pressures in [50000, 110000] Pa."
+ )
+
+ def __len__(self) -> int:
+ return len(self.a)
+
+
+VerticalCoordinate = PressureLevels | HybridLevels
+
+
+def _log_source_pressure(
+ vertical: VerticalCoordinate,
+ ps: torch.Tensor | None,
+ like: torch.Tensor,
+) -> torch.Tensor:
+ """Log source pressure with shape (..., L) broadcastable to `like`
+ (which has the level axis moved to last)."""
+ if isinstance(vertical, PressureLevels):
+ p = torch.as_tensor(
+ vertical.pressure_pa(), dtype=like.dtype, device=like.device
+ )
+ return torch.log(p).expand(like.shape)
+ if ps is None:
+ raise VerticalMismatchError(
+ f"Hybrid->pressure interpolation requires the surface pressure "
+ f"field {vertical.ps_field!r}, which was not available"
+ )
+ a = torch.as_tensor(vertical.a, dtype=like.dtype, device=like.device)
+ b = torch.as_tensor(vertical.b, dtype=like.dtype, device=like.device)
+ p = a + b * ps.to(dtype=like.dtype, device=like.device).unsqueeze(-1)
+ if torch.any(p <= 0):
+ raise VerticalMismatchError("Non-positive pressure from hybrid coefficients")
+ if torch.any(p[..., 1:] <= p[..., :-1]):
+ raise VerticalMismatchError(
+ f"Hybrid levels a + b * ps are not strictly increasing along the "
+ f"level axis for the given surface pressure field "
+ f"{vertical.ps_field!r} — interpolation would pair level slices "
+ "with wrong pressures. Check that the hybrid coefficients a and b "
+ "are ordered top to bottom and that the surface pressure values "
+ "are physical (Pa)."
+ )
+ return torch.log(p).expand(like.shape)
+
+
+def interp_to_pressure(
+ x: torch.Tensor,
+ coords: CoordSystem,
+ src: VerticalCoordinate,
+ dst: PressureLevels,
+ ps: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, CoordSystem]:
+ """Interpolate a field with a "level" dim onto constant pressure levels.
+
+ Linear in log-pressure; clamped to the source column ends (no
+ extrapolation beyond top/bottom values). `ps` (Pa) must broadcast to the
+ field with the level dim removed and is required for hybrid sources. For
+ :class:`PressureLevels` sources the data's ``level`` coordinate (hPa)
+ must match `src.levels` exactly, so the level axis is guaranteed to be
+ paired with the declared pressures.
+ """
+ if "level" not in coords:
+ raise VerticalMismatchError(
+ f"interp_to_pressure: coords have no 'level' dim ({list(coords)})"
+ )
+ if isinstance(src, PressureLevels):
+ lev = np.asarray(coords["level"], dtype=np.float64)
+ src_lev = np.asarray(src.levels, dtype=np.float64)
+ if lev.shape != src_lev.shape or not np.allclose(lev, src_lev):
+ raise VerticalMismatchError(
+ f"interp_to_pressure: data 'level' coordinate {lev.tolist()} "
+ f"does not match the declared PressureLevels source "
+ f"{list(src.levels)} (hPa). Reorder the data so levels "
+ "increase top to bottom, or fix the source component's "
+ "export_vertical declaration to match its 'level' coordinate."
+ )
+ if isinstance(src, PressureLevels) and tuple(src.levels) == tuple(dst.levels):
+ return x, coords
+ lev_axis = list(coords).index("level")
+ n_src = len(coords["level"])
+ xp = x.movedim(lev_axis, -1) # (..., L)
+ if n_src != xp.shape[-1]:
+ raise VerticalMismatchError(
+ f"'level' coord length {n_src} != tensor level size {xp.shape[-1]}"
+ )
+ logp_src = _log_source_pressure(src, ps, xp)
+ logp_dst = torch.log(
+ torch.as_tensor(dst.pressure_pa(), dtype=xp.dtype, device=xp.device)
+ ).expand(*xp.shape[:-1], len(dst.levels))
+
+ idx_hi = torch.searchsorted(logp_src.contiguous(), logp_dst.contiguous())
+ idx_hi = idx_hi.clamp(1, xp.shape[-1] - 1)
+ idx_lo = idx_hi - 1
+ x_lo = torch.gather(xp, -1, idx_lo)
+ x_hi = torch.gather(xp, -1, idx_hi)
+ p_lo = torch.gather(logp_src, -1, idx_lo)
+ p_hi = torch.gather(logp_src, -1, idx_hi)
+ w = ((logp_dst - p_lo) / (p_hi - p_lo)).clamp(0.0, 1.0)
+ out = x_lo * (1.0 - w) + x_hi * w
+
+ out = out.movedim(-1, lev_axis)
+ new_coords = OrderedDict(coords)
+ new_coords["level"] = np.asarray(dst.levels, dtype=np.float64)
+ return out, new_coords
diff --git a/examples/09_nvcoupler/01_coupled_toy_workflow.py b/examples/09_nvcoupler/01_coupled_toy_workflow.py
new file mode 100644
index 000000000..a40ed212c
--- /dev/null
+++ b/examples/09_nvcoupler/01_coupled_toy_workflow.py
@@ -0,0 +1,195 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+# %%
+"""
+Running a Coupled Atmosphere-Ocean System
+=========================================
+
+The core nvcoupler workflow on synthetic components.
+
+This example builds the smallest complete coupled system: a fast "atmosphere"
+(6 h step, 32x64 grid) and a slow "ocean" (48 h step, 16x32 grid) exchanging
+fields through two connectors — one of them a windowed (trailing 48 h mean)
+reduction — exactly the cadence structure of DLESyM. The system is declared
+as components plus connections; the run sequence is derived from the coupling
+graph. Every number below is hand-computable.
+
+In this example you will learn:
+
+- How to declare components with imports/exports by standard name
+- How to declare the coupling graph and let the Driver derive the schedule
+- How a windowed connector (window=, reduce=) bridges a cadence gap
+- How to inspect exchanges (probe) and collect results as xarray
+"""
+
+# /// script
+# dependencies = [
+# "earth2studio @ git+https://github.com/NVIDIA/earth2studio.git",
+# "matplotlib",
+# ]
+# ///
+
+# %%
+# Set Up
+# ------
+# The toy components live in ``earth2studio.nvcoupler.testing``. The
+# atmosphere steps ``z1000 += 1 + 0.1 * sst`` and imports SST; the ocean
+# steps ``sst += 0.01 * z48m`` and imports the trailing 48 h mean of z1000.
+# Grids differ, so the connectors regrid automatically.
+
+import os
+
+os.makedirs("outputs", exist_ok=True)
+
+import earth2studio.nvcoupler as nvc
+from earth2studio.nvcoupler.testing import atmos_ic, fake_atmos, fake_ocean, ocean_ic
+
+atmos = fake_atmos() # 6h step, exports geopotential_at_1000hpa, imports SST
+ocean = fake_ocean() # 48h step, exports sea_surface_temperature
+
+print(atmos)
+print(ocean)
+
+# %%
+# Declare the Coupling Graph
+# --------------------------
+# Two edges. The SST hand-off is a plain connector — a bare ``(src, dst)``
+# tuple builds the default. The z1000 hand-off is a *windowed* connector:
+# ``window="48h", reduce="mean"`` folds the atmosphere's export into a
+# running mean every step and delivers it as the derived field
+# ``geopotential_at_1000hpa_48h_mean`` (declared by a CellMethod entry in the
+# ocean's dictionary) on each 48 h boundary. No mediator needed for a
+# single-source reduction.
+
+connectors = [
+ ("ocean", "atmos"), # lagged SST forcing
+ nvc.Connector(atmos, ocean, window="48h", reduce="mean"),
+]
+
+# %%
+# Build the Driver — No Run Sequence Required
+# -------------------------------------------
+# With no ``sequence=`` the Driver derives the canonical (lagged) schedule
+# from the coupling graph: one slot per cadence, connects before runs.
+# ``describe()`` shows the whole plan — components, connectors, and the
+# derived sequence — before anything runs. (An explicit run-sequence DSL
+# remains the escape hatch when the *ordering* is the experiment; see
+# example 02.)
+
+driver = nvc.Driver(
+ {"atmos": atmos, "ocean": ocean},
+ clock=nvc.Clock("2024-01-01", "2024-01-05", "6h"),
+ connectors=connectors,
+)
+print(driver.describe())
+
+# %%
+# Execute the Coupled Loop
+# ------------------------
+# The Driver validates everything at initialize (names, cadences, field
+# matching, units) and then runs 96 hours: 16 atmosphere steps, 2 ocean
+# steps, 2 window deliveries. We iterate with ``steps()`` to also capture
+# each 48 h mean as the windowed connector delivers it.
+
+driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+
+z48_means, seen = [], set()
+for time, states in driver.steps():
+ f = driver.probe("atmos->ocean").get("geopotential_at_1000hpa_48h_mean")
+ if f is not None and f.valid_time not in seen:
+ seen.add(f.valid_time)
+ z48_means.append(float(f.data.mean()))
+datasets = driver.to_xarray() # dict[str, xr.Dataset] (in-memory collection)
+
+print(f"atmos ran {atmos.run_count}x, ocean {ocean.run_count}x")
+
+# %%
+# Inspect the Results
+# -------------------
+# Expected values: z grows 1.2/step under SST=2, so z(48h)=9.6; the first
+# 48h mean is 4.2, giving sst=2.042; z then grows 1.2042/step to
+# z(96h)=19.2336, and sst(96h)=2.180147.
+
+import numpy as np
+
+z = datasets["atmos"]["geopotential_at_1000hpa"]
+sst = datasets["ocean"]["sea_surface_temperature"]
+
+print("\ntime series (area means):")
+print(f"{'time':>20} {'z1000':>10} {'sst':>10}")
+for t in z.time.values:
+ z_t = float(z.sel(time=t).mean())
+ row = f"{str(t)[:16]:>20} {z_t:>10.4f}"
+ if t in sst.time.values:
+ row += f" {float(sst.sel(time=t).mean()):>10.6f}"
+ print(row)
+print(f"\n48h means delivered by the windowed connector: {z48_means}")
+
+if not np.isclose(float(z.values[-1].mean()), 19.2336, atol=1e-4):
+ raise ValueError("z1000(96h) does not match the hand-computed value")
+if not np.isclose(float(sst.values[-1].mean()), 2.180147, atol=1e-6):
+ raise ValueError("sst(96h) does not match the hand-computed value")
+if not np.allclose(z48_means, [4.2, 13.8147], atol=1e-4):
+ raise ValueError("48h means do not match the hand-computed values")
+
+# %%
+# Plot the Coupled Time Series
+# ----------------------------
+# The same area means, visualized: z1000 grows a little faster after each
+# 48 h coupling event because the ocean warms in response to the mean z1000
+# it received — the coupling feedback is visible as the kinks at 48 h/96 h.
+
+import matplotlib.pyplot as plt
+
+fig, (ax0, ax1) = plt.subplots(2, 1, figsize=(8, 6), sharex=True)
+
+ax0.plot(z.time, z.mean(("lat", "lon")), "o-", color="tab:blue")
+ax0.set_ylabel("z1000 area mean")
+ax0.set_title("Coupled toy system: 6 h atmosphere, 48 h ocean")
+
+ax1.plot(sst.time, sst.mean(("lat", "lon")), "s-", color="tab:red")
+ax1.set_ylabel("sst area mean")
+ax1.set_xlabel("time")
+
+fig.autofmt_xdate()
+plt.tight_layout()
+plt.savefig("outputs/01_coupled_toy_timeseries.jpg")
+
+# %%
+# Probe an Exchange
+# -----------------
+# Every connector remembers the last fields it moved — useful when a coupled
+# run misbehaves and you need to see what actually crossed the interface.
+# The windowed connector's probe carries the *derived* standard name.
+
+f = driver.probe("ocean->atmos")["sea_surface_temperature"]
+print(f"last ocean->atmos transfer: {f}")
+print(f"regridded to the atmos grid: {tuple(f.data.shape)}")
+z48 = driver.probe("atmos->ocean")["geopotential_at_1000hpa_48h_mean"]
+print(f"last atmos->ocean transfer: {z48}")
+
+# %%
+# One-Call Auto-Wiring
+# --------------------
+# ``couple()`` goes one step further: it discovers both edges (including the
+# windowed one, from the ocean's derived import) by matching standard names,
+# so the whole system above is:
+
+driver2 = nvc.couple(fake_atmos(), fake_ocean(), start="2024-01-01", stop="2024-01-05")
+driver2.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+z96 = float(driver2.run()["atmos"]["geopotential_at_1000hpa"].values[-1].mean())
+print(f"couple() reproduces z1000(96h) = {z96:.4f}")
diff --git a/examples/09_nvcoupler/02_lagged_vs_sequential.py b/examples/09_nvcoupler/02_lagged_vs_sequential.py
new file mode 100644
index 000000000..bb719ddb3
--- /dev/null
+++ b/examples/09_nvcoupler/02_lagged_vs_sequential.py
@@ -0,0 +1,117 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+# %%
+"""
+Coupling Order as a One-Line Experiment
+=======================================
+
+Lagged vs sequential coupling by reordering run-sequence actions.
+
+In physical coupled modeling, whether the atmosphere sees the ocean state
+from *this* coupling step (sequential/implicit-ish) or the *previous* one
+(lagged/explicit) is an architectural decision baked deep into the coupler.
+In nvcoupler it is one line of the run sequence: a ConnectAction placed
+before the destination's RunAction in a slot is lagged; after, sequential.
+The sequence a Driver derives from the coupling graph is always the lagged
+canonical form — passing an explicit run-sequence DSL is the escape hatch
+that makes the ordering itself the experiment, as here.
+
+In this example you will learn:
+
+- How action order inside a slot defines coupling semantics
+- How to run the same system twice under both orderings
+- Why the two differ by exactly one coupling window (hand-computable)
+"""
+
+# /// script
+# dependencies = [
+# "earth2studio @ git+https://github.com/NVIDIA/earth2studio.git",
+# ]
+# ///
+
+# %%
+# Two Sequences, One Line Apart
+# -----------------------------
+# The SST hand-off to the atmosphere is confined to the 48 h slot so the
+# orderings are cleanly comparable. Only the position of ``ocean -> atmos``
+# changes.
+
+import numpy as np
+
+import earth2studio.nvcoupler as nvc
+from earth2studio.nvcoupler.testing import atmos_ic, fake_atmos, fake_ocean, ocean_ic
+
+LAGGED = """
+@6h
+ atmos -> med
+ atmos
+@48h
+ med.compute
+ ocean -> atmos # BEFORE the ocean runs: atmos gets the OLD sst
+ med -> ocean
+ ocean
+@
+"""
+
+SEQUENTIAL = """
+@6h
+ atmos -> med
+ atmos
+@48h
+ med.compute
+ med -> ocean
+ ocean
+ ocean -> atmos # AFTER the ocean runs: atmos gets the FRESH sst
+@
+"""
+
+
+def run(sequence: str) -> float:
+ driver = nvc.Driver(
+ {
+ "atmos": fake_atmos(),
+ "ocean": fake_ocean(),
+ "med": nvc.TrailingAverageMediator(
+ "med", ["geopotential_at_1000hpa_48h_mean"]
+ ),
+ },
+ sequence=sequence,
+ clock=nvc.Clock("2024-01-01", "2024-01-05", "6h"),
+ )
+ driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+ ds = driver.run()
+ return float(ds["atmos"]["geopotential_at_1000hpa"].values[-1].mean())
+
+
+# %%
+# Compare
+# -------
+# Hand computation: at 48 h the ocean updates SST from 2.0 to 2.042. Under
+# the lagged ordering the atmosphere is forced by SST=2.0 for the next 8
+# steps; under the sequential ordering by 2.042. The difference in final
+# z1000 is exactly 8 steps x 0.1 x 0.042 = 0.0336.
+
+z_lagged = run(LAGGED)
+z_sequential = run(SEQUENTIAL)
+
+print(f"z1000(96h) lagged: {z_lagged:.4f}")
+print(f"z1000(96h) sequential: {z_sequential:.4f}")
+print(f"difference: {z_sequential - z_lagged:.4f}")
+print(f"expected 8*0.1*0.042 = {8 * 0.1 * 0.042:.4f}")
+if not np.isclose(z_sequential - z_lagged, 8 * 0.1 * 0.042, atol=1e-5):
+ raise ValueError("coupling-order difference does not match the analytic value")
+print("\ncoupling-order experiment reproduced the analytic difference ✓")
diff --git a/examples/09_nvcoupler/03_impact_chain.py b/examples/09_nvcoupler/03_impact_chain.py
new file mode 100644
index 000000000..e12a6c70c
--- /dev/null
+++ b/examples/09_nvcoupler/03_impact_chain.py
@@ -0,0 +1,176 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+# %%
+"""
+Weather-to-Impact Chains with Windowed Reductions
+=================================================
+
+Feeding an impact model with accumulated weather quantities.
+
+Impact models (flood, crop, energy, fire) rarely want instantaneous fields —
+they want precipitation *sums*, temperature *maxima*, degree days. Derived
+fields are declared through CF-style cell-method entries in the field
+dictionary, and nvcoupler offers two mechanisms to produce them: a **windowed
+connector** (``window=``, ``reduce=``) when one source feeds one destination,
+and an **AccumulationMediator** for terminal or multi-source reductions. Any
+Python function can be the impact model via CallableComponent (it does not
+need to be ML).
+
+In this example you will learn:
+
+- How to register derived fields with a CellMethod (no suffix parsing)
+- How a windowed connector reduces a fast field for a slow consumer
+- When an AccumulationMediator is still the right tool
+- How the run sequence is derived from the coupling graph
+"""
+
+# /// script
+# dependencies = [
+# "earth2studio @ git+https://github.com/NVIDIA/earth2studio.git",
+# ]
+# ///
+
+# %%
+# A Toy Weather Component
+# -----------------------
+# Exports 6 h precipitation (1.0 kg m-2 every step) and 2 m temperature
+# (steps upward 1 K per step from 280 K) on a 6 h cadence.
+
+from collections import OrderedDict
+
+import numpy as np
+import torch
+
+import earth2studio.nvcoupler as nvc
+from earth2studio.nvcoupler.testing import grid_coords
+
+GRID = (16, 32)
+
+
+def weather_step(x, coords):
+ tp06, t2m = x[0], x[1]
+ return torch.stack([tp06, t2m + 1.0]), coords
+
+
+weather = nvc.CallableComponent(
+ "weather",
+ weather_step,
+ timestep="6h",
+ exports=["total_precipitation_6h", "air_temperature_2m"],
+)
+
+# %%
+# Derived Fields via CellMethod
+# -----------------------------
+# ``total_precipitation_48h_sum`` and ``air_temperature_2m_24h_max`` ship in
+# the default dictionary. Windowed connectors and mediators alike read the
+# base field, reduction, and window off the entry — nothing is inferred from
+# name strings.
+
+from earth2studio.nvcoupler.dictionary import DEFAULT_DICTIONARY
+
+entry = DEFAULT_DICTIONARY.resolve("total_precipitation_48h_sum")
+print(f"{entry.standard_name}: {entry.cell_method}")
+
+# %%
+# A (Non-ML) Impact Model
+# -----------------------
+# A trivial flood index: 0.1 x the 48 h precipitation sum. The index field is
+# registered in a per-component dictionary copy; the imported sum arrives as
+# a state variable (the default VariableOverwriteAdapter pattern).
+
+d = nvc.FieldDictionary(DEFAULT_DICTIONARY)
+d.register(nvc.FieldEntry("flood_risk_index", "", "toy flood index"))
+
+
+def flood_step(x, coords):
+ _index, p48 = x[0], x[1]
+ return torch.stack([0.1 * p48, p48]), coords
+
+
+flood = nvc.CallableComponent(
+ "flood",
+ flood_step,
+ timestep="48h",
+ imports=["total_precipitation_48h_sum"],
+ exports=["flood_risk_index"],
+ variable_aliases={
+ "findex": "flood_risk_index",
+ "p48": "total_precipitation_48h_sum",
+ },
+ dictionary=d,
+)
+
+# %%
+# Wire the Chain
+# --------------
+# The precip sum has one source and one destination, so it is a **windowed
+# connector** — no mediator, no extra component. The 24 h t2m max is a
+# *terminal* product (nothing imports it), so it needs an
+# **AccumulationMediator**: mediators are components with export states that
+# land in the collected output. The run sequence is derived from the graph —
+# weather feeds both reductions every 6 h, the max reduces every 24 h, the
+# flood model runs every 48 h on the freshly delivered sum.
+
+t2m_max = nvc.AccumulationMediator("tmax", ["air_temperature_2m_24h_max"])
+
+driver = nvc.Driver(
+ {"weather": weather, "tmax": t2m_max, "flood": flood},
+ clock=nvc.Clock("2024-01-01", "2024-01-05", "6h"),
+ connectors=[
+ nvc.Connector(weather, flood, window="48h", reduce="sum"),
+ ("weather", "tmax"),
+ ],
+)
+print(driver.describe())
+
+# %%
+# Run It
+# ------
+
+ic_weather = (
+ torch.stack([torch.full(GRID, 1.0), torch.full(GRID, 280.0)]),
+ OrderedDict({"variable": np.array(["tp06", "t2m"]), **grid_coords(*GRID)}),
+)
+ic_flood = (
+ torch.zeros(2, *GRID),
+ OrderedDict({"variable": np.array(["findex", "p48"]), **grid_coords(*GRID)}),
+)
+
+# Note: initialize logs warnings that tmax's and flood's exports have no
+# consumer — correct here, they are the chain's terminal outputs.
+driver.initialize({"weather": ic_weather, "flood": ic_flood})
+ds = driver.run()
+
+# %%
+# Check the Numbers
+# -----------------
+# 8 samples of 1.0 kg m-2 per 48 h window -> sum 8.0 -> flood index 0.8.
+# t2m rises 1 K/step; the max over each 24 h window is the last sample.
+
+p48 = driver.probe("weather->flood")["total_precipitation_48h_sum"]
+tmax_series = ds["tmax"]["air_temperature_2m_24h_max"].mean(("lat", "lon")).values
+flood_series = ds["flood"]["flood_risk_index"].mean(("lat", "lon")).values
+
+print(f"last 48h precip sum: {float(p48.data.mean())}")
+print(f"24h t2m maxima: {tmax_series}")
+print(f"flood risk index: {flood_series}")
+if not np.isclose(float(p48.data.mean()), 8.0) or not np.allclose(
+ flood_series[1:], 0.8
+):
+ raise ValueError("impact chain did not reproduce the analytic values")
+print("\nimpact chain produced the analytic values ✓")
diff --git a/examples/09_nvcoupler/04_vertical_chemistry.py b/examples/09_nvcoupler/04_vertical_chemistry.py
new file mode 100644
index 000000000..6337bb11d
--- /dev/null
+++ b/examples/09_nvcoupler/04_vertical_chemistry.py
@@ -0,0 +1,158 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+# %%
+"""
+Vertical Coupling for Chemistry Models
+======================================
+
+Hybrid sigma-pressure to pressure-level interpolation in a connector.
+
+Chemistry emulators typically live on hybrid model levels
+(p_k = a_k + b_k * p_s) while meteorology components export pressure-level
+fields — or vice versa. When source and destination declare different
+vertical coordinates for a field with a "level" dimension, the connector
+interpolates linearly in log-pressure, pulling the surface-pressure field
+from the source's exports automatically.
+
+Most earth2studio models encode levels in variable names (z500, t850) and
+never touch this machinery.
+
+In this example you will learn:
+
+- How components declare export/import vertical coordinates
+- How the connector auto-resolves the surface-pressure dependency
+- What the failure mode looks like when ps is missing
+"""
+
+# /// script
+# dependencies = [
+# "earth2studio @ git+https://github.com/NVIDIA/earth2studio.git",
+# ]
+# ///
+
+# %%
+# Source on Hybrid Levels
+# -----------------------
+# A "met" component exports ozone on 3 hybrid levels. With surface pressure
+# 1000 hPa the levels realize at 300 / 700 / 1000 hPa. The ozone profile is
+# f = log(p), so interpolation results are exact and checkable.
+
+from collections import OrderedDict
+
+import numpy as np
+import torch
+
+import earth2studio.nvcoupler as nvc
+from earth2studio.nvcoupler.dictionary import DEFAULT_DICTIONARY
+from earth2studio.nvcoupler.testing import grid_coords
+
+NLAT, NLON = 8, 16
+PS = 100000.0 # Pa
+
+d = nvc.FieldDictionary(DEFAULT_DICTIONARY)
+d.register(nvc.FieldEntry("ozone_mixing_ratio", "kg kg-1", aliases=frozenset({"o3"})))
+
+hybrid = nvc.HybridLevels(a=(30000.0, 20000.0, 0.0), b=(0.0, 0.5, 1.0))
+target = nvc.PressureLevels((500.0, 850.0))
+
+
+def identity(x, coords):
+ return x, coords
+
+
+met = nvc.CallableComponent(
+ "met",
+ identity,
+ "6h",
+ exports=["ozone_mixing_ratio"],
+ export_vertical={"ozone_mixing_ratio": hybrid},
+ dictionary=d,
+)
+chem = nvc.CallableComponent(
+ "chem",
+ identity,
+ "6h",
+ imports=["ozone_mixing_ratio"],
+ import_vertical={"ozone_mixing_ratio": target},
+ dictionary=d,
+)
+
+# %%
+# Initialize and Exchange
+# -----------------------
+# The met component publishes ozone = log(p) columns plus the surface
+# pressure field the hybrid transform needs.
+
+grid = grid_coords(NLAT, NLON)
+p_src = np.array(hybrid.a) + np.array(hybrid.b) * PS # [30000, 70000, 100000] Pa
+o3 = (
+ torch.tensor(np.log(p_src), dtype=torch.float64)
+ .view(1, 3, 1, 1)
+ .expand(1, 3, NLAT, NLON)
+ .clone()
+)
+clock = nvc.Clock("2024-01-01", "2024-01-02", "6h")
+met.realize(clock)
+chem.realize(clock)
+met.initialize(
+ o3, OrderedDict({"variable": np.array(["o3"]), "level": np.arange(3.0), **grid})
+)
+met.export_state.add(
+ nvc.Field(
+ torch.full((NLAT, NLON), PS, dtype=torch.float64),
+ OrderedDict(grid),
+ "surface_pressure",
+ "Pa",
+ valid_time=clock.start,
+ source="met",
+ )
+)
+chem.initialize(
+ torch.zeros(1, 2, NLAT, NLON, dtype=torch.float64),
+ OrderedDict(
+ {"variable": np.array(["o3"]), "level": np.array([500.0, 850.0]), **grid}
+ ),
+)
+
+conn = nvc.Connector(met, chem, fields=["ozone_mixing_ratio"])
+conn.execute(clock.start)
+
+# %%
+# Verify Exactness
+# ----------------
+# Linear-in-log-p interpolation of f = log(p) must return log(p_target).
+
+got = chem.import_state["ozone_mixing_ratio"]
+expected = np.log(np.array([50000.0, 85000.0]))
+print(f"received on levels {list(got.coords['level'])} hPa")
+print(f"column values: {got.data[:, 0, 0].numpy()}")
+print(f"expected: {expected}")
+if not np.allclose(got.data[:, 0, 0].numpy(), expected):
+ raise ValueError("hybrid -> pressure interpolation did not match log(p)")
+print("hybrid -> pressure interpolation exact ✓")
+
+# %%
+# The Failure Mode
+# ----------------
+# Remove surface pressure from the source and the connector refuses with a
+# fix, at the exchange — not as NaNs three days into a rollout.
+
+del met.export_state["surface_pressure"]
+try:
+ nvc.Connector(met, chem, fields=["ozone_mixing_ratio"]).execute(clock.start)
+except nvc.VerticalMismatchError as e:
+ print(f"\nVerticalMismatchError: {e}")
diff --git a/examples/09_nvcoupler/05_coupled_finetuning.py b/examples/09_nvcoupler/05_coupled_finetuning.py
new file mode 100644
index 000000000..cb80d1501
--- /dev/null
+++ b/examples/09_nvcoupler/05_coupled_finetuning.py
@@ -0,0 +1,128 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+# %%
+"""
+Gradients Across the Exchange: Coupled Fine-Tuning
+==================================================
+
+Backpropagating a coupled-rollout loss into both components.
+
+Separately-trained emulators drift when coupled: each was trained against
+truth forcing, not against the other model's imperfect output. The remedy is
+coupled fine-tuning — optimizing both models jointly on coupled rollouts.
+nvcoupler's exchange path (regrid gathers, windowed reductions, functional
+import injection) is autograd-clean end to end, and ``driver.rollout()``
+keeps the graph, so a loss on one component's final state reaches the
+parameters of every component upstream through the exchanges.
+
+This example fine-tunes the two scalar "physics" parameters of the toy
+system so the 96 h coupled forecast hits a target. The point is not the toy
+optimization — it is that gradients cross the coupler.
+
+In this example you will learn:
+
+- How rollout() differs from run()/steps() (graph kept vs inference mode)
+- How to verify gradients reach parameters through the exchange
+- The shape of a minimal coupled training step
+"""
+
+# /// script
+# dependencies = [
+# "earth2studio @ git+https://github.com/NVIDIA/earth2studio.git",
+# ]
+# ///
+
+# %%
+# Trainable Components
+# --------------------
+# Each toy takes a gain parameter inside its step function. gain_ocean can
+# influence the atmosphere ONLY through the exchange chain:
+# ocean sst -> connector regrid -> import injection -> atmos step.
+
+import torch
+
+import earth2studio.nvcoupler as nvc
+from earth2studio.nvcoupler.testing import atmos_ic, fake_atmos, fake_ocean, ocean_ic
+
+gain_atmos = torch.tensor(1.0, requires_grad=True)
+gain_ocean = torch.tensor(1.0, requires_grad=True)
+
+
+def coupled_forecast() -> torch.Tensor:
+ """One 96 h coupled rollout; returns the final mean z1000 (graph kept).
+
+ couple() auto-wires the two components: lagged SST forcing one way, a
+ windowed (trailing 48 h mean) connector the other — the running
+ reduction is differentiable like every other exchange stage.
+ """
+ driver = nvc.couple(
+ fake_atmos(gain=gain_atmos),
+ fake_ocean(gain=gain_ocean),
+ start="2024-01-01",
+ stop="2024-01-05",
+ collect=False,
+ )
+ driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+ states = driver.rollout(16) # full 96h, autograd graph retained
+ return states["atmos"]["geopotential_at_1000hpa"].data.mean()
+
+
+# %%
+# Gradients Cross the Coupler
+# ---------------------------
+# With gains at 1.0 the forecast lands at 19.2336 (see example 01). A loss on
+# the atmosphere's final state produces nonzero gradients for BOTH gains —
+# the ocean's only route to the loss is through the two connectors (one of
+# them a windowed reduction).
+
+with torch.enable_grad():
+ z96 = coupled_forecast()
+ z96.backward()
+
+print(f"z1000(96h) = {z96.item():.4f}")
+print(f"d(loss)/d(gain_atmos) = {gain_atmos.grad.item():.4f}")
+print(f"d(loss)/d(gain_ocean) = {gain_ocean.grad.item():.6f}")
+if gain_atmos.grad == 0 or gain_ocean.grad == 0:
+ raise ValueError("a gradient failed to cross the exchange path")
+print("gradients reached both components through the exchange ✓\n")
+
+# %%
+# A Minimal Coupled Training Loop
+# -------------------------------
+# Fine-tune both gains so the coupled forecast hits z1000(96h) = 25. Each
+# iteration rebuilds the system from the same initial conditions (a fresh
+# clock) and takes one optimizer step — the skeleton of coupled fine-tuning.
+
+TARGET = 25.0
+optimizer = torch.optim.Adam([gain_atmos, gain_ocean], lr=0.1)
+
+for it in range(60):
+ optimizer.zero_grad()
+ with torch.enable_grad():
+ z96 = coupled_forecast()
+ loss = (z96 - TARGET) ** 2
+ loss.backward()
+ optimizer.step()
+ if it % 10 == 0 or it == 59:
+ print(
+ f"iter {it:2d}: z96 = {z96.item():7.4f} loss = {loss.item():9.5f} "
+ f"gains = ({gain_atmos.item():.4f}, {gain_ocean.item():.4f})"
+ )
+
+if abs(z96.item() - TARGET) >= 0.5:
+ raise ValueError("coupled fine-tuning failed to reach the target forecast")
+print(f"\ncoupled system fine-tuned to the target ({z96.item():.3f} ≈ {TARGET}) ✓")
diff --git a/examples/09_nvcoupler/06_pull_conditioning.py b/examples/09_nvcoupler/06_pull_conditioning.py
new file mode 100644
index 000000000..8f3eabadf
--- /dev/null
+++ b/examples/09_nvcoupler/06_pull_conditioning.py
@@ -0,0 +1,283 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+# %%
+"""
+Pull-Pattern Coupling: Feeding a Model That Fetches Its Own Forcing
+===================================================================
+
+Coupling a StormCast-style regional model without temp-file staging.
+
+Most models take coupled fields as call arguments — the coupler pushes.
+Some do not: StormCast holds a settable ``conditioning_data_source`` and,
+inside its own ``__call__``, *pulls* its conditioning by calling the real
+``earth2studio.data.utils.fetch_data`` on that attribute. The production
+workflow (``serve/server/example_workflows/stormcast_conus_workflow.py``)
+copes today by running the full global conditioning forecast first, staging
+it to temporary NetCDF files, and handing StormCast an
+``InferenceOutputSource`` over those files — a data source masquerading as
+GFS. It works, but the entire conditioning forecast must be materialized to
+disk before the regional model takes a single step.
+
+nvcoupler's :class:`~earth2studio.nvcoupler.pull.PullAdapter` plays the
+same masquerade minus the staging: before each step it installs a
+:class:`~earth2studio.nvcoupler.pull.StateDataSource` — an in-memory
+DataSource answering fetches from the component's live import State — on
+the model's data-source attribute. The model runs its unmodified production
+fetch path and receives THIS step's coupled forcing, delivered by a
+connector moments earlier in the same run-sequence slot.
+
+In this example you will learn:
+
+- What the pull pattern is and which models need it (StormCast)
+- How PullAdapter + StateDataSource replace the temp-file staging
+- How an explicit sequential run sequence guarantees fresh conditioning
+- How to verify, step by step, that the model pulled the current exchange
+
+One honest caveat up front: the pull path crosses the model's own
+xarray/numpy fetch machinery, so pull-coupled components are
+**inference-only** — no gradients flow through this exchange (contrast
+example 05, where the push-pattern exchange is autograd-clean).
+"""
+
+# /// script
+# dependencies = [
+# "earth2studio @ git+https://github.com/NVIDIA/earth2studio.git",
+# ]
+# ///
+
+# %%
+# A Pull-Pattern Mock Regional Model
+# ----------------------------------
+# We stand in for StormCast with a mock that is *protocol-faithful* to its
+# coupling surface: a settable ``conditioning_data_source`` attribute, a
+# declared list of raw conditioning variables, and — crucially — a fetch
+# through the REAL ``earth2studio.data.utils.fetch_data``, the same code
+# path the production model runs. If the shim satisfies fetch_data here, it
+# satisfies StormCast's mechanics. Its "physics" is a hand-computable
+# update of a single radar-reflectivity state:
+#
+# refc <- refc + 1 + mean(u10m) + 0.1 * mean(t2m)
+#
+# and it logs every pull so we can audit exactly what conditioning it saw.
+
+from collections import OrderedDict
+
+import numpy as np
+import torch
+
+import earth2studio.nvcoupler as nvc
+from earth2studio.data.utils import fetch_data
+from earth2studio.nvcoupler.pull import PullAdapter
+from earth2studio.nvcoupler.testing import grid_coords
+
+GRID = (8, 16)
+
+
+class MockPullModel:
+ """StormCast stand-in: pulls u10m/t2m via fetch_data inside __call__.
+
+ Mirrors the production coupling surface exactly — the coupler never
+ calls anything but ``model(x, coords)``; the conditioning arrives
+ through the data source the model itself fetches from.
+ """
+
+ conditioning_variables = np.array(["u10m", "t2m"])
+
+ def __init__(self):
+ self.conditioning_data_source = None # PullAdapter sets this
+ self.pull_log: list[np.ndarray] = []
+
+ def input_coords(self):
+ return OrderedDict(
+ {
+ "time": np.empty(0),
+ "lead_time": np.array([np.timedelta64(0, "h")]),
+ "variable": np.array(["refc"]),
+ **grid_coords(*GRID),
+ }
+ )
+
+ def output_coords(self, input_coords):
+ out = OrderedDict({k: v.copy() for k, v in input_coords.items()})
+ out["lead_time"] = input_coords["lead_time"] + np.timedelta64(1, "h")
+ return out
+
+ def __call__(self, x, coords):
+ if self.conditioning_data_source is None:
+ raise RuntimeError("conditioning_data_source not set")
+ # The REAL fetch path StormCast uses — not a shortcut around it.
+ cond, _ = fetch_data(
+ self.conditioning_data_source,
+ time=np.atleast_1d(coords["time"]),
+ variable=self.conditioning_variables,
+ )
+ self.pull_log.append(cond.numpy().copy())
+ u_mean = cond[0, 0, 0].mean()
+ t_mean = cond[0, 0, 1].mean()
+ return x + 1.0 + u_mean + 0.1 * t_mean, self.output_coords(coords)
+
+ def to(self, device):
+ return self
+
+
+# %%
+# The Global Conditioning Component
+# ---------------------------------
+# A toy "global" model in the Jussi-workflow shape: it exports the two
+# fields StormCast conditions on. Its 10 m wind grows by 1 m/s per step
+# (so freshness is detectable — each hour's conditioning differs from the
+# last) while temperature holds constant at 280 K.
+
+
+def global_step(x, coords):
+ # x stacks [u10m, t2m]; u10m grows 1 m/s per step, t2m constant
+ return torch.stack([x[0] + 1.0, x[1]]), coords
+
+
+glob = nvc.CallableComponent(
+ "global",
+ global_step,
+ timestep="1h",
+ exports=["eastward_wind_10m", "air_temperature_2m"],
+)
+
+# %%
+# Wire the Regional Component with PullAdapter
+# --------------------------------------------
+# The regional component imports what the global one exports. Its
+# ``import_adapter=PullAdapter()`` is the whole pull-pattern story: before
+# each step, the adapter installs a StateDataSource over the import State
+# on the model's ``conditioning_data_source``, then calls the model
+# unchanged. The model pulls "u10m"/"t2m" by its raw names; the shim
+# resolves them to the standard names in the State. ``refc`` is not in the
+# default field dictionary, so we register it.
+
+dictionary = nvc.FieldDictionary(nvc.DEFAULT_DICTIONARY)
+dictionary.register(
+ nvc.FieldEntry("radar_reflectivity", "dBZ", aliases=frozenset({"refc"}))
+)
+
+model = MockPullModel()
+stormcast = nvc.PrognosticComponent(
+ "stormcast",
+ model,
+ imports=["eastward_wind_10m", "air_temperature_2m"],
+ exports=["radar_reflectivity"],
+ import_adapter=PullAdapter(),
+ variable_aliases={"refc": "radar_reflectivity"},
+ dictionary=dictionary,
+)
+
+# %%
+# Explicit Sequential Run Sequence
+# --------------------------------
+# Freshness is an *ordering* property, so we write the sequence explicitly
+# rather than letting the Driver derive the lagged default: in every 1 h
+# slot the global model steps, the connector delivers its new exports, and
+# only then does the regional model run — so each pull sees conditioning
+# valid at the pulled time. This is the sequential-coupling half of
+# example 02, applied to the pull pattern.
+
+T0 = np.datetime64("2024-01-01")
+
+driver = nvc.Driver(
+ {"global": glob, "stormcast": stormcast},
+ sequence="""
+ @1h
+ global
+ global -> stormcast
+ stormcast
+ @
+ """,
+ clock=nvc.Clock(T0, "2024-01-01T04:00", "1h"),
+ connectors=[nvc.Connector(glob, stormcast)],
+)
+print(driver.describe())
+
+# %%
+# Initialize and Run 4 Hours
+# --------------------------
+# The global state starts at u10m = 2, t2m = 280; refc starts at 0. Hand
+# computation: at hour k the global model has already stepped, so the pull
+# sees u = 2 + k, and the refc increment is 1 + (2 + k) + 0.1 * 280 =
+# 31 + k — i.e. 32, 33, 34, 35 over four hours, cumulative 134. We iterate
+# with ``steps()`` to record refc after every slot.
+
+ic_glob = (
+ torch.stack([torch.full(GRID, 2.0), torch.full(GRID, 280.0)]),
+ OrderedDict({"variable": np.array(["u10m", "t2m"]), **grid_coords(*GRID)}),
+)
+ic_sc = model.input_coords()
+ic_sc["time"] = np.array([T0])
+driver.initialize(
+ {"global": ic_glob, "stormcast": (torch.zeros(1, 1, 1, *GRID), ic_sc)}
+)
+
+refc_series = []
+for time, states in driver.steps():
+ refc_series.append(float(stormcast.export_state["radar_reflectivity"].data.mean()))
+
+# %%
+# Prove Every Step Pulled Fresh Conditioning
+# ------------------------------------------
+# The model's pull log is the audit trail: each entry is what its own
+# fetch_data returned that step. The pulled u10m must be 3, 4, 5, 6 — the
+# global state *after* that hour's step, never a stale or staged value —
+# and each refc increment must match the arithmetic above.
+
+print(
+ f"{'hour':>6} {'pulled u10m':>12} {'pulled t2m':>11} "
+ f"{'refc incr':>10} {'refc':>8}"
+)
+prev = 0.0
+for k, (pull, refc) in enumerate(zip(model.pull_log, refc_series), start=1):
+ u, t = float(pull[0, 0, 0].mean()), float(pull[0, 0, 1].mean())
+ incr = refc - prev
+ prev = refc
+ print(f"{k:>5}h {u:>12.1f} {t:>11.1f} {incr:>10.1f} {refc:>8.1f}")
+
+pulled_u = [float(p[0, 0, 0].mean()) for p in model.pull_log]
+if pulled_u != [3.0, 4.0, 5.0, 6.0]:
+ raise ValueError("a pull saw stale conditioning — sequencing is broken")
+if not np.isclose(refc_series[-1], 32.0 + 33.0 + 34.0 + 35.0):
+ raise ValueError("refc(4h) does not match the hand-computed value")
+print("\nevery regional step pulled that step's fresh conditioning ✓")
+
+# %%
+# Probe the Exchange
+# ------------------
+# The connector's probe shows the last fields it moved — the same fields
+# the StateDataSource then served to the model's fetch. Note the u10m value
+# matches the final pull (6.0): what crossed the interface is what the
+# model consumed, with no files in between.
+
+u10 = driver.probe("global->stormcast")["eastward_wind_10m"]
+t2 = driver.probe("global->stormcast")["air_temperature_2m"]
+print(f"last global->stormcast transfer: {u10}")
+print(f" {t2}")
+
+# %%
+# Honest Closing Notes
+# --------------------
+# Two limits worth stating plainly. First, the pull path goes through the
+# model's own fetch_data/xarray machinery, so field data crosses a numpy
+# boundary: pull-coupled components are inference-only, and no gradients
+# flow through this exchange (push-pattern adapters keep autograd intact —
+# see example 05). Second, this example validates the *mechanics* — the
+# mock is protocol-faithful, fetching through the real fetch_data — but a
+# run against real StormCast weights, replacing the temp-file staging in
+# ``stormcast_conus_workflow.py`` end to end, remains future validation.
diff --git a/examples/09_nvcoupler/07_point_sampling.py b/examples/09_nvcoupler/07_point_sampling.py
new file mode 100644
index 000000000..d082de5b0
--- /dev/null
+++ b/examples/09_nvcoupler/07_point_sampling.py
@@ -0,0 +1,142 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+# %%
+"""
+Grid-to-Point Sampling for Station-Level Applications
+======================================================
+
+Delivering a gridded field to a scattered set of locations.
+
+Many applications need a coarse gridded forecast at specific, non-gridded
+points rather than on a mesh: verification against station observations,
+site-level agriculture fields, or energy-asset locations. A destination
+component whose grid is a :class:`~earth2studio.nvcoupler.PointSet` (arbitrary
+(lat, lon) locations, not a mesh) tells the Connector to sample instead of
+regrid — ``sample="nearest"`` or ``sample="bilinear"``.
+
+In this example you will learn:
+
+- How to declare a point-target component with ``points=PointSet(...)``
+- How ``sample="nearest"`` and ``sample="bilinear"`` differ
+- What the failure mode looks like when neither ``sample=`` nor a custom
+ ``regridder=`` is given for a point target
+"""
+
+# /// script
+# dependencies = [
+# "earth2studio @ git+https://github.com/NVIDIA/earth2studio.git",
+# ]
+# ///
+
+# %%
+# Source: a Gridded "Atmosphere"
+# -------------------------------
+# A 32x64 lat/lon component exporting a field that is exactly linear in
+# (lat, lon): ``temperature = lat + 0.1 * lon``. Linearity makes both nearest
+# and bilinear sampling hand-checkable — bilinear must reproduce the formula
+# exactly anywhere inside the grid, and nearest must reproduce it exactly at
+# any point that coincides with a grid cell.
+
+from collections import OrderedDict
+
+import numpy as np
+import torch
+
+import earth2studio.nvcoupler as nvc
+from earth2studio.nvcoupler.testing import grid_coords
+
+NLAT, NLON = 32, 64
+
+
+def identity(x, coords):
+ return x, coords
+
+
+atmos = nvc.CallableComponent("atmos", identity, "6h", exports=["air_temperature_2m"])
+grid = grid_coords(NLAT, NLON)
+lat, lon = np.asarray(grid["lat"]), np.asarray(grid["lon"])
+temperature = torch.as_tensor(lat).view(-1, 1) + 0.1 * torch.as_tensor(lon).view(1, -1)
+
+clock = nvc.Clock("2024-01-01", "2024-01-02", "6h")
+atmos.realize(clock)
+atmos.initialize(
+ temperature.unsqueeze(0).double(),
+ OrderedDict({"variable": np.array(["air_temperature_2m"]), **grid}),
+)
+
+# %%
+# Destination: Named Stations
+# ----------------------------
+# Three stations: two sit exactly on grid points (so nearest is exact there
+# too), one sits at the midpoint between four cells (nearest and bilinear
+# will disagree there).
+
+stations = nvc.PointSet(
+ lat=np.array([lat[4], lat[10], (lat[6] + lat[7]) / 2]),
+ lon=np.array([lon[3], lon[20], (lon[12] + lon[13]) / 2]),
+ names=("boulder", "denver", "midpoint"),
+)
+site = nvc.CallableComponent(
+ "stations", identity, "6h", imports=["air_temperature_2m"], points=stations
+)
+site.realize(clock)
+# initialize() always publishes through State.from_tensor, which requires a
+# 'variable' dim even for a component with no exports; a placeholder value
+# is skipped (nothing in export_names asks for it).
+site.initialize(
+ torch.zeros(1, len(stations)),
+ OrderedDict({"variable": np.array(["_ic"]), "point": stations.labels()}),
+)
+
+# %%
+# Sample: Nearest vs. Bilinear
+# ------------------------------
+nvc.Connector(atmos, site, sample="nearest").execute(clock.start)
+nearest = site.import_state["air_temperature_2m"].data.clone()
+
+nvc.Connector(atmos, site, sample="bilinear").execute(clock.start)
+bilinear = site.import_state["air_temperature_2m"].data.clone()
+
+expected_exact = lat[[4, 10]] + 0.1 * lon[[3, 20]] # the two on-grid stations
+expected_midpoint = float(
+ (lat[6] + lat[7]) / 2 + 0.1 * (lon[12] + lon[13]) / 2
+) # linear field: exact anywhere under bilinear
+
+print(f"stations: {stations.labels()}")
+print(f"nearest sample: {nearest.numpy()}")
+print(f"bilinear sample: {bilinear.numpy()}")
+print(f"expected (on-grid): {expected_exact}")
+print(f"expected (midpoint): {expected_midpoint:.4f}")
+
+if not np.allclose(nearest.numpy()[:2], expected_exact):
+ raise ValueError("nearest sampling at on-grid stations did not match")
+if not np.allclose(bilinear.numpy(), [*expected_exact, expected_midpoint]):
+ raise ValueError("bilinear sampling did not exactly reproduce the linear field")
+if np.isclose(nearest.numpy()[2], bilinear.numpy()[2]):
+ raise ValueError("nearest and bilinear were expected to disagree at the midpoint")
+print("nearest and bilinear both correct; they disagree at the midpoint as expected ✓")
+
+# %%
+# The Failure Mode
+# -----------------
+# A point-target destination with neither ``sample=`` nor a custom
+# ``regridder=`` refuses at the exchange, with the fix in the message.
+
+try:
+ nvc.Connector(atmos, site).execute(clock.start)
+except nvc.CouplingError as e:
+ print(f"\nCouplingError: {e}")
diff --git a/examples/09_nvcoupler/README.rst b/examples/09_nvcoupler/README.rst
new file mode 100644
index 000000000..19200b59d
--- /dev/null
+++ b/examples/09_nvcoupler/README.rst
@@ -0,0 +1,16 @@
+.. _nvcoupler_examples:
+
+Coupling (nvcoupler)
+--------------------
+
+Examples for nvcoupler, the NUOPC/ESMF-inspired coupling framework for AI
+Earth-system inference. All examples run on synthetic toy components — no
+model weights, GPU, or network access required — and each prints
+hand-verifiable numbers. They cover the coupled atmos/ocean loop declared as
+a coupling graph (with a windowed-reduction connector), coupling order
+experiments via explicit run sequences, impact chains mixing windowed
+connectors and accumulation mediators, vertical (hybrid to pressure)
+coupling for chemistry, gradient flow across the exchange for coupled
+fine-tuning, pull-pattern (StormCast-style) conditioning via
+``PullAdapter``, and grid-to-point sampling (``sample=``) for
+station-level applications.
diff --git a/test/nvcoupler/test_api.py b/test/nvcoupler/test_api.py
new file mode 100644
index 000000000..81ea2a4f1
--- /dev/null
+++ b/test/nvcoupler/test_api.py
@@ -0,0 +1,257 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Tests for the UX layer: couple() auto-wiring, coupled(), describe().
+
+The auto-wired toy system must reproduce the hand-computed expectations of
+test_driver.py: z(96h) = 19.2336 and sst(96h) = 2.180147 — with the derived
+import carried by a windowed connector instead of a mediator.
+"""
+
+import numpy as np
+import pytest
+import torch
+import xarray as xr
+
+from earth2studio.nvcoupler.api import couple, coupled, describe, describe_html
+from earth2studio.nvcoupler.component import CallableComponent
+from earth2studio.nvcoupler.errors import (
+ AmbiguousCouplingError,
+ UnmatchedImportError,
+)
+from earth2studio.nvcoupler.mediator import AccumulationMediator
+from earth2studio.nvcoupler.testing import (
+ ATMOS_GRID,
+ OCEAN_GRID,
+ atmos_ic,
+ fake_atmos,
+ fake_ocean,
+ ocean_ic,
+)
+
+T0 = "2024-01-01"
+T96 = "2024-01-05"
+
+
+def test_couple_synthesizes_windowed_connector():
+ # No mediator: couple() carries the derived import on a windowed
+ # connector built from the CellMethod of geopotential_at_1000hpa_48h_mean.
+ driver = couple(fake_atmos(), fake_ocean(), start=T0, stop=T96)
+
+ assert not any(
+ isinstance(c, AccumulationMediator) for c in driver.components.values()
+ )
+ conn = driver._connectors[("atmos", "ocean")]
+ assert conn.window == np.timedelta64(48, "h").astype("timedelta64[ns]")
+ assert conn.reduce == "mean"
+
+ # dt defaults to the GCD of 6h and 48h
+ assert driver.clock.dt == np.timedelta64(6, "h").astype("timedelta64[ns]")
+
+ # sequence is derived from the graph in the canonical lagged shape
+ assert driver.sequence_derived
+ expected = "\n".join(
+ [
+ "@6h",
+ " atmos -> ocean",
+ " ocean -> atmos",
+ " atmos",
+ "@48h",
+ " ocean",
+ "@",
+ ]
+ )
+ assert str(driver.sequence) == expected
+
+ driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+ driver.run()
+ z = driver.components["atmos"].export_state["geopotential_at_1000hpa"]
+ assert torch.allclose(z.data, torch.full(ATMOS_GRID, 19.2336), atol=1e-4)
+ sst = driver.components["ocean"].export_state["sea_surface_temperature"]
+ assert torch.allclose(sst.data, torch.full(OCEAN_GRID, 2.180147), atol=1e-4)
+
+
+def test_couple_synthesizes_mediator_when_pair_also_transfers_plainly():
+ """When the (src, dst) pair already carries a plain transfer, the derived
+ import cannot share the connector — a mediator is genuinely needed."""
+
+ def step(x, coords):
+ return x, coords
+
+ greedy = CallableComponent(
+ "greedy",
+ step,
+ timestep="48h",
+ imports=["geopotential_at_1000hpa", "geopotential_at_1000hpa_48h_mean"],
+ exports=["sea_surface_temperature"],
+ variable_aliases={
+ "z1000": "geopotential_at_1000hpa",
+ "z48m": "geopotential_at_1000hpa_48h_mean",
+ },
+ )
+ driver = couple(fake_atmos(), greedy, start=T0, stop=T96)
+ mediators = [
+ c for c in driver.components.values() if isinstance(c, AccumulationMediator)
+ ]
+ assert len(mediators) == 1
+ med = mediators[0]
+ assert med.export_names == ["geopotential_at_1000hpa_48h_mean"]
+ # plain z1000 rides the direct connector; the mean goes via the mediator
+ assert ("atmos", "greedy") in driver._connectors
+ assert ("atmos", med.name) in driver._connectors
+ assert (med.name, "greedy") in driver._connectors
+ assert driver._connectors[("atmos", "greedy")].window is None
+
+
+def test_ambiguous_exports_raise():
+ def step(x, coords):
+ return x, coords
+
+ ocean2 = CallableComponent(
+ "ocean2",
+ step,
+ timestep="48h",
+ imports=["geopotential_at_1000hpa_48h_mean"],
+ exports=["sea_surface_temperature"],
+ variable_aliases={"z48m": "geopotential_at_1000hpa_48h_mean"},
+ )
+ with pytest.raises(AmbiguousCouplingError, match="sea_surface_temperature"):
+ couple(fake_atmos(), fake_ocean(), ocean2, start=T0, stop=T96)
+
+
+def test_unmatched_import_raises_with_available_exports():
+ def step(x, coords):
+ return x, coords
+
+ # imports air_temperature_2m: no exporter, and no cell_method to
+ # synthesize a mediator from
+ lonely = CallableComponent(
+ "lonely",
+ step,
+ timestep="6h",
+ imports=["air_temperature_2m"],
+ exports=["mean_sea_level_pressure"],
+ )
+ other = CallableComponent(
+ "other", step, timestep="6h", exports=["geopotential_at_1000hpa"]
+ )
+ with pytest.raises(UnmatchedImportError) as err:
+ couple(other, lonely, start=T0, stop=T96)
+ msg = str(err.value)
+ assert "air_temperature_2m" in msg
+ assert "geopotential_at_1000hpa" in msg # available exports listed
+
+
+def test_unmatched_derived_import_without_base_exporter_raises():
+ def step(x, coords):
+ return x, coords
+
+ # ocean imports the 48h mean but nobody exports the base field
+ with pytest.raises(UnmatchedImportError, match="geopotential_at_1000hpa_48h_mean"):
+ couple(
+ fake_ocean(),
+ CallableComponent(
+ "other", step, timestep="6h", exports=["mean_sea_level_pressure"]
+ ),
+ start=T0,
+ stop=T96,
+ )
+
+
+def test_describe_pre_initialize():
+ driver = couple(fake_atmos(), fake_ocean(), start=T0, stop=T96)
+ text = describe(driver) # must work BEFORE initialize
+ assert "atmos" in text
+ assert "ocean" in text
+ assert "6h" in text
+ assert "2D" in text # 48h cadence formatted by fmt_timedelta
+ assert "->" in text
+ assert "sea_surface_temperature" in text
+ assert "constant" in text # time policy column
+ assert "lagged" in text
+ # run sequence section present
+ assert "@6h" in text and "@48h" in text
+
+ html = describe_html(driver)
+ assert "atmos" in html and "ocean" in html
+ assert "nvc-box" in html and "→" in html
+
+ # still works after initialize
+ driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+ assert "atmos" in describe(driver)
+
+
+def test_describe_labels_mediator_delivery_sequential():
+ """Mode is per exchange: med -> ocean follows med.compute in the same
+ slot, so ocean consumes state produced in this very iteration."""
+ from earth2studio.nvcoupler.clock import Clock
+ from earth2studio.nvcoupler.driver import Driver
+ from earth2studio.nvcoupler.mediator import TrailingAverageMediator
+
+ driver = Driver(
+ {
+ "atmos": fake_atmos(),
+ "ocean": fake_ocean(),
+ "med": TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"]),
+ },
+ clock=Clock(T0, T96, "6h"),
+ connectors=[("atmos", "med"), ("ocean", "atmos"), ("med", "ocean")],
+ )
+ text = describe(driver)
+ rows = [
+ line.strip()
+ for line in text.splitlines()
+ if " -> " in line and ("lagged" in line or "sequential" in line)
+ ]
+
+ def mode_of(name: str) -> str:
+ row = next(r for r in rows if r.startswith(name))
+ return "sequential" if "sequential" in row else "lagged"
+
+ assert mode_of("atmos -> med") == "lagged"
+ assert mode_of("ocean -> atmos") == "lagged"
+ assert mode_of("med -> ocean") == "sequential"
+
+
+def test_coupled_end_to_end():
+ ds = coupled(
+ T0,
+ T96,
+ [fake_atmos(), fake_ocean()],
+ ics={"atmos": atmos_ic(), "ocean": ocean_ic()},
+ verbose=False,
+ )
+ assert isinstance(ds["atmos"], xr.Dataset)
+ z = ds["atmos"]["geopotential_at_1000hpa"]
+ assert z.dims == ("time", "lat", "lon")
+ assert z.shape == (17, 32, 64)
+ assert np.allclose(z.values[-1], 19.2336, atol=1e-4)
+ sst = ds["ocean"]["sea_surface_temperature"]
+ assert sst.shape == (3, 16, 32)
+ assert np.allclose(sst.values[-1], 2.180147, atol=1e-4)
+
+
+def test_coupled_accepts_nsteps_and_dict():
+ ds = coupled(
+ T0,
+ 16, # 16 x 6h = 96h
+ {"atmos": fake_atmos(), "ocean": fake_ocean()},
+ ics={"atmos": atmos_ic(), "ocean": ocean_ic()},
+ verbose=False,
+ )
+ assert np.allclose(
+ ds["atmos"]["geopotential_at_1000hpa"].values[-1], 19.2336, atol=1e-4
+ )
diff --git a/test/nvcoupler/test_clock.py b/test/nvcoupler/test_clock.py
new file mode 100644
index 000000000..484ee9f7a
--- /dev/null
+++ b/test/nvcoupler/test_clock.py
@@ -0,0 +1,66 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import numpy as np
+import pytest
+
+from earth2studio.nvcoupler.clock import Clock, as_timedelta, is_multiple
+from earth2studio.nvcoupler.errors import CadenceError
+
+
+def test_as_timedelta_parsing():
+ assert as_timedelta("6h") == np.timedelta64(6, "h")
+ assert as_timedelta("2D") == np.timedelta64(2, "D")
+ assert as_timedelta("2d") == np.timedelta64(2, "D")
+ assert as_timedelta(np.timedelta64(30, "m")) == np.timedelta64(30, "m")
+ with pytest.raises(ValueError):
+ as_timedelta("h6")
+
+
+def test_bare_int_timedelta_rejected():
+ # finding 2e: a unit-less number is ambiguous (hours? steps?) — reject it
+ # with an actionable message instead of a TypeError deep in numpy
+ with pytest.raises(ValueError, match="'6h'"):
+ as_timedelta(6)
+ with pytest.raises(ValueError, match="np.timedelta64"):
+ Clock("2024-01-01", "2024-01-02", 6)
+
+
+def test_clock_iteration():
+ clock = Clock("2024-01-01", "2024-01-05", "6h")
+ assert clock.n_steps == 16
+ times = list(clock)
+ assert len(times) == 16
+ assert times[0] == np.datetime64("2024-01-01T06:00")
+ assert times[-1] == np.datetime64("2024-01-05T00:00")
+ assert clock.done()
+ with pytest.raises(StopIteration):
+ clock.advance()
+ clock.reset()
+ assert clock.current == np.datetime64("2024-01-01")
+ assert len(clock.times()) == 17 # includes start
+
+
+def test_clock_validation():
+ with pytest.raises(CadenceError):
+ Clock("2024-01-01", "2024-01-02T01:00", "6h") # span not multiple of dt
+ with pytest.raises(ValueError):
+ Clock("2024-01-02", "2024-01-01", "6h") # stop before start
+
+
+def test_is_multiple():
+ assert is_multiple(np.timedelta64(48, "h"), np.timedelta64(6, "h"))
+ assert not is_multiple(np.timedelta64(7, "h"), np.timedelta64(6, "h"))
diff --git a/test/nvcoupler/test_component.py b/test/nvcoupler/test_component.py
new file mode 100644
index 000000000..d40cbe735
--- /dev/null
+++ b/test/nvcoupler/test_component.py
@@ -0,0 +1,320 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+from collections import OrderedDict
+
+import numpy as np
+import pytest
+import torch
+
+from earth2studio.nvcoupler.clock import Clock
+from earth2studio.nvcoupler.component import (
+ CallableComponent,
+ ConditioningKwargAdapter,
+ Exchange,
+ ExtraTensorAdapter,
+ PrognosticComponent,
+)
+from earth2studio.nvcoupler.errors import CadenceError, CouplingError
+from earth2studio.nvcoupler.field import Field, State
+from earth2studio.nvcoupler.testing import (
+ atmos_ic,
+ fake_atmos,
+ fake_ocean,
+ grid_coords,
+ ocean_ic,
+)
+
+
+def _sst_import(value=3.0, nlat=32, nlon=64):
+ return State(
+ "imports",
+ [
+ Field(
+ torch.full((nlat, nlon), value),
+ grid_coords(nlat, nlon),
+ "sea_surface_temperature",
+ "K",
+ )
+ ],
+ )
+
+
+class TestExchangeInject:
+ def test_inject_overwrites_slice(self):
+ x, coords = atmos_ic(z0=0.0, sst0=2.0)
+ ex = Exchange(x, coords, _sst_import(5.0), {"sea_surface_temperature": "sst"})
+ out = ex.inject()
+ assert torch.all(out[1] == 5.0) # sst slice overwritten
+ assert torch.all(out[0] == 0.0) # z1000 untouched
+ assert torch.all(x[1] == 2.0) # original not mutated
+
+ def test_missing_variable_raises(self):
+ x, coords = atmos_ic()
+ ex = Exchange(x, coords, _sst_import(), {"sea_surface_temperature": "nope"})
+ with pytest.raises(CouplingError, match="not a state variable"):
+ ex.inject()
+
+ def test_gradient_flows_through_injection(self):
+ x, coords = atmos_ic()
+ sst = torch.full((32, 64), 5.0, requires_grad=True)
+ imports = State(
+ "imports",
+ [Field(sst, grid_coords(32, 64), "sea_surface_temperature", "K")],
+ )
+ ex = Exchange(x, coords, imports, {"sea_surface_temperature": "sst"})
+ ex.inject().sum().backward()
+ assert sst.grad is not None and torch.all(sst.grad == 1.0)
+
+
+class TestOtherAdapters:
+ def test_conditioning_kwarg(self):
+ captured = {}
+
+ class Model:
+ def call_with_conditioning(
+ self, x, coords, conditioning, conditioning_coords
+ ):
+ captured["conditioning"] = conditioning
+ captured["coords"] = conditioning_coords
+ return x, coords
+
+ x, coords = atmos_ic()
+ adapter = ConditioningKwargAdapter()
+ adapter(Model(), Exchange(x, coords, _sst_import(7.0)))
+ assert torch.all(captured["conditioning"] == 7.0)
+ assert list(captured["coords"]) == ["variable", "lat", "lon"]
+
+ def test_extra_tensor(self):
+ captured = {}
+
+ def model(x, coords, coupling):
+ captured["coupling"] = coupling
+ return x, coords
+
+ x, coords = atmos_ic()
+ ExtraTensorAdapter()(model, Exchange(x, coords, _sst_import(9.0)))
+ assert torch.all(captured["coupling"] == 9.0)
+
+ def test_multiple_imports_require_explicit_field_order(self):
+ # channel order is model-sensitive; alphabetical stacking would run
+ # fine and predict garbage, so >1 import without field_order= must fail
+ from earth2studio.nvcoupler.field import Field
+ from earth2studio.nvcoupler.testing import grid_coords
+
+ imports = State(
+ "imports",
+ [
+ Field(
+ torch.ones(8, 16),
+ grid_coords(8, 16),
+ "sea_surface_temperature",
+ "K",
+ ),
+ Field(torch.ones(8, 16), grid_coords(8, 16), "air_temperature_2m", "K"),
+ ],
+ )
+ x, coords = atmos_ic()
+ exchange = Exchange(x, coords, imports)
+ model = lambda x, coords, coupling: (x, coords) # noqa: E731
+ with pytest.raises(CouplingError, match="field_order"):
+ ExtraTensorAdapter()(model, exchange)
+ # explicit order works, and is honored (not alphabetical)
+ captured = {}
+
+ def capture(x, coords, coupling):
+ captured["coupling"] = coupling
+ return x, coords
+
+ ExtraTensorAdapter(
+ field_order=["sea_surface_temperature", "air_temperature_2m"]
+ )(capture, exchange)
+ assert captured["coupling"].shape[0] == 2
+ # unknown name in field_order is rejected
+ with pytest.raises(CouplingError, match="not in the"):
+ ExtraTensorAdapter(field_order=["nope"])(model, exchange)
+
+
+class TestCustomAdapter:
+ def test_custom_adapter_receives_exchange(self):
+ """The user-facing pattern: any callable (model, exchange) ->
+ (x, coords) plugs in as import_adapter=."""
+ seen = {}
+
+ def scaled_overwrite(model, exchange: Exchange):
+ seen["time"] = exchange.time
+ seen["std_to_raw"] = dict(exchange.std_to_raw)
+ return model(0.5 * exchange.inject(), exchange.coords)
+
+ def step(x, coords):
+ z1000, sst = x[0], x[1]
+ return torch.stack([z1000 + 1.0 + 0.1 * sst, sst]), coords
+
+ atmos = CallableComponent(
+ "atmos",
+ step,
+ timestep="6h",
+ imports=["sea_surface_temperature"],
+ exports=["geopotential_at_1000hpa"],
+ import_adapter=scaled_overwrite,
+ )
+ clock = Clock("2024-01-01", "2024-01-02", "6h")
+ atmos.realize(clock)
+ atmos.initialize(*atmos_ic(z0=0.0, sst0=2.0))
+ atmos.import_state.add(_sst_import(10.0)["sea_surface_temperature"])
+ t1 = clock.advance()
+ atmos.run(t1)
+ # inject() overwrote sst=10, adapter halved the state: z = 0.5*0 + 1
+ # + 0.1 * (0.5*10) = 1.5
+ z = atmos.export_state["geopotential_at_1000hpa"]
+ assert torch.allclose(z.data, torch.full((32, 64), 1.5))
+ assert seen["time"] == t1
+ assert seen["std_to_raw"]["sea_surface_temperature"] == "sst"
+
+
+class TestCallableComponent:
+ def test_toy_atmos_step_arithmetic(self):
+ atmos = fake_atmos()
+ clock = Clock("2024-01-01", "2024-01-02", "6h")
+ atmos.realize(clock)
+ atmos.initialize(*atmos_ic(z0=0.0, sst0=2.0))
+ # export seeded at t0 for lagged coupling
+ assert atmos.export_state["geopotential_at_1000hpa"].valid_time == clock.start
+
+ t1 = clock.advance()
+ atmos.run(t1)
+ # z = 0 + 1 + 0.1*2 = 1.2 (no import set; state sst used)
+ z = atmos.export_state["geopotential_at_1000hpa"]
+ assert torch.allclose(z.data, torch.full((32, 64), 1.2))
+ assert z.valid_time == t1 and z.source == "atmos"
+
+ # inject an import and step again: z = 1.2 + 1 + 0.1*10 = 3.2
+ atmos.import_state.add(_sst_import(10.0)["sea_surface_temperature"])
+ atmos.run(clock.advance())
+ assert torch.allclose(
+ atmos.export_state["geopotential_at_1000hpa"].data,
+ torch.full((32, 64), 3.2),
+ )
+
+ def test_ocean_mask_export(self):
+ ocean = fake_ocean(with_mask=True)
+ ocean.realize(Clock("2024-01-01", "2024-01-03", "48h"))
+ ocean.initialize(*ocean_ic())
+ sst = ocean.export_state["sea_surface_temperature"]
+ assert sst.mask is not None
+ assert not sst.mask[0, 0] and sst.mask[-1, -1]
+
+ def test_cadence_validation(self):
+ atmos = fake_atmos(timestep="6h")
+ with pytest.raises(CadenceError):
+ atmos.realize(Clock("2024-01-01", "2024-01-02", "4h"))
+
+ def test_run_before_initialize_raises(self):
+ atmos = fake_atmos()
+ atmos.realize(Clock("2024-01-01", "2024-01-02", "6h"))
+ with pytest.raises(CouplingError, match="not initialized"):
+ atmos.run(np.datetime64("2024-01-01T06:00"))
+
+ def test_advertise(self):
+ ocean = fake_ocean()
+ imports, exports = ocean.advertise()
+ assert imports == ["geopotential_at_1000hpa_48h_mean"]
+ assert exports == ["sea_surface_temperature"]
+
+
+class MockPrognostic:
+ """Minimal PrognosticModel: two variables, 6h step, +1 per step."""
+
+ def __init__(self):
+ self._in = OrderedDict(
+ {
+ "lead_time": np.array([np.timedelta64(0, "h")]),
+ "variable": np.array(["z1000", "sst"]),
+ **grid_coords(8, 16),
+ }
+ )
+
+ def input_coords(self):
+ return OrderedDict({k: v.copy() for k, v in self._in.items()})
+
+ def output_coords(self, input_coords):
+ out = OrderedDict({k: v.copy() for k, v in input_coords.items()})
+ out["lead_time"] = input_coords["lead_time"] + np.timedelta64(6, "h")
+ return out
+
+ def __call__(self, x, coords):
+ return x + 1.0, self.output_coords(coords)
+
+ def to(self, device):
+ return self
+
+
+class TestPrognosticComponent:
+ def test_timestep_and_exports_inferred(self):
+ comp = PrognosticComponent("mock", MockPrognostic())
+ assert comp.timestep == np.timedelta64(6, "h")
+ assert set(comp.export_names) == {
+ "geopotential_at_1000hpa",
+ "sea_surface_temperature",
+ }
+
+ def test_rollout_and_publish(self):
+ comp = PrognosticComponent(
+ "mock", MockPrognostic(), imports=["sea_surface_temperature"]
+ )
+ clock = Clock("2024-01-01", "2024-01-02", "6h")
+ comp.realize(clock)
+ ic = MockPrognostic().input_coords()
+ comp.initialize(torch.zeros(1, 2, 8, 16), ic)
+ for t in clock:
+ comp.run(t)
+ assert comp.run_count == 4
+ z = comp.export_state["geopotential_at_1000hpa"]
+ # exports are exchange-shaped: singleton lead_time squeezed away
+ assert list(z.coords) == ["lat", "lon"]
+ assert z.data.shape == (8, 16)
+ assert torch.all(z.data == 4.0)
+ assert z.valid_time == np.datetime64("2024-01-02")
+ # the internal model state keeps the full model dims
+ x, coords = comp.state
+ assert list(coords) == ["lead_time", "variable", "lat", "lon"]
+ assert x.shape == (1, 2, 8, 16)
+
+ def test_multi_window_needs_next_input(self):
+ class TwoWindow(MockPrognostic):
+ def output_coords(self, input_coords):
+ out = super().output_coords(input_coords)
+ out["lead_time"] = np.array(
+ [np.timedelta64(6, "h"), np.timedelta64(12, "h")]
+ )
+ return out
+
+ def __call__(self, x, coords):
+ return torch.cat([x, x]), self.output_coords(coords)
+
+ comp = PrognosticComponent("two", TwoWindow(), timestep="6h")
+ comp.realize(Clock("2024-01-01", "2024-01-02", "6h"))
+ comp.initialize(torch.zeros(1, 2, 8, 16), MockPrognostic().input_coords())
+ with pytest.raises(CouplingError, match="next_input"):
+ comp.run(np.datetime64("2024-01-01T06:00"))
+
+ def test_publish_missing_export_raises(self):
+ comp = PrognosticComponent(
+ "mock", MockPrognostic(), exports=["air_temperature_2m"]
+ )
+ comp.realize(Clock("2024-01-01", "2024-01-02", "6h"))
+ with pytest.raises(CouplingError, match="advertises export"):
+ comp.initialize(torch.zeros(1, 2, 8, 16), MockPrognostic().input_coords())
diff --git a/test/nvcoupler/test_config.py b/test/nvcoupler/test_config.py
new file mode 100644
index 000000000..af6ce8d83
--- /dev/null
+++ b/test/nvcoupler/test_config.py
@@ -0,0 +1,302 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""YAML config round-trip tests on the synthetic coupled system."""
+
+import numpy as np
+import pytest
+import yaml
+
+from earth2studio.nvcoupler.clock import Clock, as_timedelta
+from earth2studio.nvcoupler.config import from_yaml, to_yaml
+from earth2studio.nvcoupler.dictionary import (
+ DEFAULT_DICTIONARY,
+ CellMethod,
+ FieldDictionary,
+ FieldEntry,
+)
+from earth2studio.nvcoupler.driver import Driver
+from earth2studio.nvcoupler.errors import CouplingError
+from earth2studio.nvcoupler.mediator import (
+ AccumulationMediator,
+ TrailingAverageMediator,
+)
+from earth2studio.nvcoupler.testing import atmos_ic, fake_atmos, fake_ocean, ocean_ic
+
+T0 = "2024-01-01"
+T96 = "2024-01-05"
+
+DSL = """
+@6h
+ atmos -> med
+ ocean -> atmos
+ atmos
+@48h
+ med.compute
+ med -> ocean
+ ocean
+@
+"""
+
+
+def make_tagged_driver():
+ """The test_driver.py toy system, with yaml_spec-tagged toy components."""
+ atmos = fake_atmos(gain=1.0)
+ atmos.yaml_spec = {
+ "class": "earth2studio.nvcoupler.testing.fake_atmos",
+ "kwargs": {"gain": 1.0, "timestep": "6h"},
+ }
+ ocean = fake_ocean(gain=1.0)
+ ocean.yaml_spec = {
+ "class": "earth2studio.nvcoupler.testing.fake_ocean",
+ "kwargs": {"gain": 1.0, "timestep": "48h"},
+ }
+ med = TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"])
+ return Driver(
+ {"atmos": atmos, "ocean": ocean, "med": med}, DSL, Clock(T0, T96, "6h")
+ )
+
+
+def run_all(driver):
+ driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+ return driver.run()
+
+
+def test_round_trip_identical_outputs(tmp_path):
+ driver = make_tagged_driver()
+ path = tmp_path / "system.yaml"
+ text = to_yaml(driver, path=path)
+ assert path.read_text() == text
+
+ rebuilt = from_yaml(path)
+ assert sorted(rebuilt.components) == sorted(driver.components)
+ assert str(rebuilt.sequence) == str(driver.sequence)
+ assert rebuilt.clock.start == driver.clock.start
+ assert rebuilt.clock.stop == driver.clock.stop
+ assert rebuilt.clock.dt == driver.clock.dt
+
+ ds_a = run_all(driver)
+ ds_b = run_all(rebuilt)
+ for comp in ("atmos", "ocean", "med"):
+ for var in ds_a[comp].data_vars:
+ assert np.array_equal(ds_a[comp][var].values, ds_b[comp][var].values)
+ # sanity anchor from the hand-computed driver tests
+ assert np.allclose(
+ ds_b["atmos"]["geopotential_at_1000hpa"].values[-1], 19.2336, atol=1e-4
+ )
+
+
+def test_from_yaml_accepts_text_and_path(tmp_path):
+ driver = make_tagged_driver()
+ text = to_yaml(driver)
+ rebuilt = from_yaml(text) # text, not a path
+ assert sorted(rebuilt.components) == ["atmos", "med", "ocean"]
+
+
+def test_closure_component_without_yaml_spec_raises():
+ driver = make_tagged_driver()
+ del driver.components["atmos"].yaml_spec
+ with pytest.raises(CouplingError, match="yaml_spec"):
+ to_yaml(driver)
+
+
+def test_mediator_window_serialized_as_string():
+ driver = make_tagged_driver()
+ doc = yaml.safe_load(to_yaml(driver))
+ window = doc["components"]["med"]["kwargs"]["window"]
+ assert isinstance(window, str)
+ assert as_timedelta(window) == np.timedelta64(48, "h")
+
+
+def test_custom_dictionary_entry_round_trip():
+ """A non-default derived entry (12 h mean, np.timedelta64 window) survives
+ the YAML round-trip as a '12h'-style string."""
+ dictionary = FieldDictionary(DEFAULT_DICTIONARY)
+ dictionary.register(
+ FieldEntry(
+ "geopotential_at_1000hpa_12h_mean",
+ "m2 s-2",
+ "trailing 12 h mean of z1000",
+ frozenset({"z12m"}),
+ CellMethod("geopotential_at_1000hpa", "mean", np.timedelta64(12, "h")),
+ )
+ )
+ atmos = fake_atmos()
+ atmos.yaml_spec = {
+ "class": "earth2studio.nvcoupler.testing.fake_atmos",
+ "kwargs": {"gain": 1.0, "timestep": "6h"},
+ }
+ med = AccumulationMediator(
+ "med", ["geopotential_at_1000hpa_12h_mean"], dictionary=dictionary
+ )
+ dsl = "@6h\n atmos -> med\n atmos\n@12h\n med.compute\n@"
+ driver = Driver({"atmos": atmos, "med": med}, dsl, Clock(T0, "2024-01-02", "6h"))
+ # the toy atmos intentionally runs on constant IC forcing (its SST import
+ # is fed by nothing in this two-component system)
+ driver.allow_unfed_imports = True
+ text = to_yaml(driver)
+ doc = yaml.safe_load(text)
+ entries = {e["standard_name"]: e for e in doc["dictionary"]}
+ entry = entries["geopotential_at_1000hpa_12h_mean"]
+ assert isinstance(entry["cell_method"]["window"], str)
+ assert as_timedelta(entry["cell_method"]["window"]) == np.timedelta64(12, "h")
+ # default entries stay out of the file
+ assert "geopotential_at_1000hpa" not in entries
+
+ rebuilt = from_yaml(text)
+ rebuilt.allow_unfed_imports = True
+ med2 = rebuilt.components["med"]
+ assert med2.timestep == np.timedelta64(12, "h").astype("timedelta64[ns]")
+ driver.initialize({"atmos": atmos_ic()})
+ rebuilt.initialize({"atmos": atmos_ic()})
+ ds_a, ds_b = driver.run(), rebuilt.run()
+ assert np.array_equal(
+ ds_a["med"]["geopotential_at_1000hpa_12h_mean"].values,
+ ds_b["med"]["geopotential_at_1000hpa_12h_mean"].values,
+ )
+
+
+def test_add_alias_round_trip():
+ """Aliases added via FieldDictionary.add_alias() after registration must
+ survive the YAML round-trip even though the FieldEntry itself still
+ equals the default (regression: only non-default entries were dumped)."""
+ dictionary = FieldDictionary(DEFAULT_DICTIONARY)
+ dictionary.add_alias("geopotential_at_1000hpa", "phi1000")
+ atmos = fake_atmos()
+ atmos.yaml_spec = {
+ "class": "earth2studio.nvcoupler.testing.fake_atmos",
+ "kwargs": {"gain": 1.0, "timestep": "6h"},
+ }
+ med = TrailingAverageMediator(
+ "med", ["geopotential_at_1000hpa_48h_mean"], dictionary=dictionary
+ )
+ dsl = "@6h\n atmos -> med\n atmos\n@48h\n med.compute\n@"
+ driver = Driver({"atmos": atmos, "med": med}, dsl, Clock(T0, "2024-01-03", "6h"))
+
+ text = to_yaml(driver)
+ doc = yaml.safe_load(text)
+ assert doc["aliases"] == {"phi1000": "geopotential_at_1000hpa"}
+ # the default entry itself stays out of the dictionary section
+ assert not any(
+ e["standard_name"] == "geopotential_at_1000hpa"
+ for e in doc.get("dictionary", [])
+ )
+
+ rebuilt = from_yaml(text)
+ med2 = rebuilt.components["med"]
+ assert med2.dictionary.standard_name("phi1000") == "geopotential_at_1000hpa"
+ # and the round-trip is stable: dumping again re-emits the alias
+ # (re-tag atmos: closures never carry yaml_spec across from_yaml)
+ rebuilt.components["atmos"].yaml_spec = atmos.yaml_spec
+ doc2 = yaml.safe_load(to_yaml(rebuilt))
+ assert doc2["aliases"] == {"phi1000": "geopotential_at_1000hpa"}
+
+
+def test_connectors_round_trip():
+ from earth2studio.nvcoupler.connector import Connector
+
+ driver = make_tagged_driver()
+ atmos, med = driver.components["atmos"], driver.components["med"]
+ driver._connectors[("atmos", "med")] = Connector(
+ atmos, med, fields=["geopotential_at_1000hpa"], time_policy="constant"
+ )
+ doc = yaml.safe_load(to_yaml(driver))
+ assert doc["connectors"] == [
+ {
+ "src": "atmos",
+ "dst": "med",
+ "time_policy": "constant",
+ "fill": "none",
+ "fields": ["geopotential_at_1000hpa"],
+ }
+ ]
+ rebuilt = from_yaml(to_yaml(driver))
+ conn = rebuilt._connectors[("atmos", "med")]
+ assert conn._fields == ["geopotential_at_1000hpa"]
+
+
+def make_tagged_declarative_driver():
+ """The toy system declared as a graph (derived sequence, windowed
+ connector for the 48 h mean) with yaml_spec-tagged components."""
+ from earth2studio.nvcoupler.api import couple
+
+ atmos = fake_atmos(gain=1.0)
+ atmos.yaml_spec = {
+ "class": "earth2studio.nvcoupler.testing.fake_atmos",
+ "kwargs": {"gain": 1.0, "timestep": "6h"},
+ }
+ ocean = fake_ocean(gain=1.0)
+ ocean.yaml_spec = {
+ "class": "earth2studio.nvcoupler.testing.fake_ocean",
+ "kwargs": {"gain": 1.0, "timestep": "48h"},
+ }
+ return couple(atmos, ocean, start=T0, stop=T96)
+
+
+def test_derived_sequence_round_trip():
+ driver = make_tagged_declarative_driver()
+ assert driver.sequence_derived
+ text = to_yaml(driver)
+ doc = yaml.safe_load(text)
+ # derived sequences serialize with the flag plus the (informational) text
+ assert doc["sequence"]["derived"] is True
+ assert doc["sequence"]["text"] == str(driver.sequence)
+ # the windowed connector carries its reduction options
+ conn_doc = next(
+ c for c in doc["connectors"] if (c["src"], c["dst"]) == ("atmos", "ocean")
+ )
+ assert as_timedelta(conn_doc["window"]) == np.timedelta64(48, "h")
+ assert conn_doc["reduce"] == "mean"
+
+ rebuilt = from_yaml(text)
+ assert rebuilt.sequence_derived # re-derived, not replayed
+ assert str(rebuilt.sequence) == str(driver.sequence)
+ conn = rebuilt._connectors[("atmos", "ocean")]
+ assert conn.window == np.timedelta64(48, "h").astype("timedelta64[ns]")
+ assert conn.reduce == "mean"
+
+ ds_a = run_all(driver)
+ ds_b = run_all(rebuilt)
+ for comp in ("atmos", "ocean"):
+ for var in ds_a[comp].data_vars:
+ assert np.array_equal(ds_a[comp][var].values, ds_b[comp][var].values)
+ assert np.allclose(
+ ds_b["atmos"]["geopotential_at_1000hpa"].values[-1], 19.2336, atol=1e-4
+ )
+ assert np.allclose(
+ ds_b["ocean"]["sea_surface_temperature"].values[-1], 2.180147, atol=1e-4
+ )
+
+
+def test_explicit_sequence_still_serializes_as_text():
+ doc = yaml.safe_load(to_yaml(make_tagged_driver()))
+ assert isinstance(doc["sequence"], str)
+ assert "@6h" in doc["sequence"]
+
+
+def test_bad_sequence_mapping_raises():
+ text = to_yaml(make_tagged_declarative_driver())
+ bad = text.replace("derived: true", "derived: false")
+ with pytest.raises(CouplingError, match="derived"):
+ from_yaml(bad)
+
+
+def test_helpful_error_on_bad_import_path():
+ text = to_yaml(make_tagged_driver()).replace(
+ "earth2studio.nvcoupler.testing.fake_atmos", "no.such.module.fake"
+ )
+ with pytest.raises(CouplingError, match="no.such.module"):
+ from_yaml(text)
diff --git a/test/nvcoupler/test_connector.py b/test/nvcoupler/test_connector.py
new file mode 100644
index 000000000..ac4edc6be
--- /dev/null
+++ b/test/nvcoupler/test_connector.py
@@ -0,0 +1,732 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+from collections import OrderedDict
+
+import numpy as np
+import pytest
+import torch
+
+from earth2studio.nvcoupler.clock import Clock
+from earth2studio.nvcoupler.component import CallableComponent
+from earth2studio.nvcoupler.connector import Connector
+from earth2studio.nvcoupler.dictionary import (
+ DEFAULT_DICTIONARY,
+ FieldDictionary,
+ FieldEntry,
+)
+from earth2studio.nvcoupler.errors import (
+ CouplingError,
+ IncompatibleFieldError,
+ VerticalMismatchError,
+)
+from earth2studio.nvcoupler.points import PointSet
+from earth2studio.nvcoupler.testing import (
+ atmos_ic,
+ fake_atmos,
+ fake_ocean,
+ ocean_ic,
+)
+from earth2studio.nvcoupler.vertical import HybridLevels, PressureLevels
+
+T0 = np.datetime64("2024-01-01")
+
+
+@pytest.fixture
+def caplog(caplog):
+ """Extend caplog to also capture loguru log records (the connector logs
+ via loguru, which does not propagate to stdlib logging by default)."""
+ from loguru import logger as loguru_logger
+
+ handler_id = loguru_logger.add(caplog.handler, format="{message}", level=0)
+ yield caplog
+ loguru_logger.remove(handler_id)
+
+
+def _realized_pair(with_mask=False):
+ atmos, ocean = fake_atmos(), fake_ocean(with_mask=with_mask)
+ clock = Clock(T0, "2024-01-05", "6h")
+ atmos.realize(clock)
+ ocean.realize(clock)
+ atmos.initialize(*atmos_ic())
+ ocean.initialize(*ocean_ic(sst0=3.0))
+ return atmos, ocean, clock
+
+
+def test_match_and_units():
+ atmos, ocean, _ = _realized_pair()
+ conn = Connector(ocean, atmos)
+ assert conn.match() == ["sea_surface_temperature"]
+ # no overlap the other way (atmos exports z1000, ocean imports 48h mean)
+ with pytest.raises(IncompatibleFieldError, match="no fields match"):
+ Connector(atmos, ocean).match()
+ # explicit fields not present on both sides
+ with pytest.raises(IncompatibleFieldError, match="not in both"):
+ Connector(ocean, atmos, fields=["air_temperature_2m"]).match()
+
+
+def test_execute_regrids_to_destination_grid():
+ atmos, ocean, _ = _realized_pair()
+ Connector(ocean, atmos).execute(T0)
+ sst = atmos.import_state["sea_surface_temperature"]
+ assert sst.data.shape == (32, 64) # ocean 16x32 -> atmos grid
+ assert torch.allclose(
+ sst.data, torch.full((32, 64), 3.0)
+ ) # constant stays constant
+ assert sst.source == "ocean"
+ assert list(sst.coords) == ["lat", "lon"]
+ assert np.array_equal(sst.coords["lat"], atmos.grid_coords()["lat"])
+
+
+def test_identity_when_grids_match():
+ a1, a2 = fake_atmos(), fake_atmos()
+ a2.name = "atmos2"
+ a2.import_names = ["geopotential_at_1000hpa"]
+ a2.export_names = ["sea_surface_temperature"]
+ clock = Clock(T0, "2024-01-02", "6h")
+ for c in (a1, a2):
+ c.realize(clock)
+ c.initialize(*atmos_ic())
+ conn = Connector(a1, a2) # a1 exports z1000, a2 imports it; same grid
+ conn.execute(T0)
+ z = a2.import_state["geopotential_at_1000hpa"]
+ assert torch.equal(z.data, a1.export_state["geopotential_at_1000hpa"].data)
+ assert conn._regridders == {} # identity fast path built no regridder
+
+
+def test_regrid_matches_direct_kernel_call():
+ from earth2studio.utils.interp import latlon_interpolation_regular
+
+ atmos, ocean, _ = _realized_pair()
+ # put a non-constant field in the ocean export
+ lat = torch.as_tensor(ocean.grid_coords()["lat"])
+ lon = torch.as_tensor(ocean.grid_coords()["lon"])
+ data = lat.view(-1, 1) + 0.1 * lon.view(1, -1)
+ f = ocean.export_state["sea_surface_temperature"]
+ f.data = data.to(torch.float32)
+ Connector(ocean, atmos).execute(T0)
+ got = atmos.import_state["sea_surface_temperature"].data
+
+ lat1, lon1 = np.meshgrid(
+ atmos.grid_coords()["lat"], atmos.grid_coords()["lon"], indexing="ij"
+ )
+ expected = latlon_interpolation_regular(
+ torch.flip(data.to(torch.float32), dims=(-2,)), # ascending lat
+ torch.as_tensor(np.asarray(ocean.grid_coords()["lat"])[::-1].copy()).float(),
+ torch.as_tensor(ocean.grid_coords()["lon"]).float(),
+ torch.as_tensor(lat1).float(),
+ torch.as_tensor(lon1).float(),
+ )
+ assert torch.allclose(got, expected)
+
+
+def _points_destination(points, imports=("sea_surface_temperature",)):
+ """Toy destination component whose grid is a scattered PointSet."""
+
+ def step(x, coords):
+ return x, coords
+
+ return CallableComponent(
+ "stations",
+ step,
+ timestep="48h",
+ imports=list(imports),
+ exports=[],
+ points=points,
+ )
+
+
+def _init_no_export(component, x, coords):
+ """initialize() a component with no exports.
+
+ Component.publish (called unconditionally by initialize()) always
+ round-trips through State.from_tensor, which requires a 'variable' dim
+ even when export_names is empty — so every IC here needs a placeholder
+ variable axis; publish()'s strict=False path skips the unresolved name
+ since nothing in export_names asks for it.
+ """
+ coords = OrderedDict({"variable": np.array(["_ic"]), **coords})
+ component.initialize(x.unsqueeze(0), coords)
+
+
+def test_sample_nearest_recovers_exact_grid_points():
+ atmos, ocean, clock = _realized_pair()
+ lat = np.asarray(ocean.grid_coords()["lat"])
+ lon = np.asarray(ocean.grid_coords()["lon"])
+ data = torch.as_tensor(lat).view(-1, 1) + 0.1 * torch.as_tensor(lon).view(1, -1)
+ ocean.export_state["sea_surface_temperature"].data = data.to(torch.float32)
+
+ points = PointSet(lat=np.array([lat[3], lat[5]]), lon=np.array([lon[2], lon[10]]))
+ stations = _points_destination(points)
+ stations.realize(clock)
+ _init_no_export(
+ stations, torch.zeros(len(points)), OrderedDict({"point": points.labels()})
+ )
+
+ Connector(ocean, stations, sample="nearest").execute(T0)
+ got = stations.import_state["sea_surface_temperature"]
+ assert list(got.coords) == ["point"]
+ assert got.data.shape == (2,)
+ assert torch.allclose(got.data, torch.tensor([data[3, 2], data[5, 10]]).float())
+ assert np.array_equal(got.coords["point"], points.labels())
+
+
+def test_sample_bilinear_matches_direct_kernel_call():
+ atmos, ocean, clock = _realized_pair()
+ lat = np.asarray(ocean.grid_coords()["lat"])
+ lon = np.asarray(ocean.grid_coords()["lon"])
+ data = torch.as_tensor(lat).view(-1, 1) + 0.1 * torch.as_tensor(lon).view(1, -1)
+ ocean.export_state["sea_surface_temperature"].data = data.to(torch.float32)
+
+ # off-grid points, midway between cells
+ points = PointSet(
+ lat=np.array([(lat[3] + lat[4]) / 2, (lat[6] + lat[7]) / 2]),
+ lon=np.array([(lon[2] + lon[3]) / 2, (lon[10] + lon[11]) / 2]),
+ )
+ stations = _points_destination(points)
+ stations.realize(clock)
+ _init_no_export(
+ stations, torch.zeros(len(points)), OrderedDict({"point": points.labels()})
+ )
+
+ Connector(ocean, stations, sample="bilinear").execute(T0)
+ got = stations.import_state["sea_surface_temperature"].data
+
+ from earth2studio.utils.interp import latlon_interpolation_regular
+
+ flip = lat[0] > lat[-1]
+ lat0 = torch.as_tensor(lat[::-1].copy() if flip else lat).float()
+ lon0 = torch.as_tensor(lon).float()
+ src = torch.flip(data.float(), dims=(-2,)) if flip else data.float()
+ lat1 = torch.as_tensor(points.lat).unsqueeze(-1).float()
+ lon1 = torch.as_tensor(points.lon).unsqueeze(-1).float()
+ expected = latlon_interpolation_regular(src, lat0, lon0, lat1, lon1).squeeze(-1)
+ assert torch.allclose(got, expected)
+
+
+def test_sample_and_regridder_mutually_exclusive():
+ with pytest.raises(CouplingError, match="mutually exclusive"):
+ Connector(*_realized_pair()[:2], sample="nearest", regridder=lambda x: x)
+
+
+def test_sample_missing_choice_raises_actionable_error():
+ atmos, ocean, clock = _realized_pair()
+ points = PointSet(lat=np.array([0.0]), lon=np.array([0.0]))
+ stations = _points_destination(points)
+ stations.realize(clock)
+ _init_no_export(stations, torch.zeros(1), OrderedDict({"point": points.labels()}))
+ with pytest.raises(CouplingError, match="neither sample= nor regridder="):
+ Connector(ocean, stations).execute(T0)
+
+
+def test_sample_without_points_metadata_raises():
+ atmos, ocean, clock = _realized_pair()
+
+ def step(x, coords):
+ return x, coords
+
+ # a "point" dim without a registered PointSet (points=None) — reachable
+ # if a component hand-builds coords with a "point" key directly
+ stations = CallableComponent(
+ "stations",
+ step,
+ timestep="48h",
+ imports=["sea_surface_temperature"],
+ exports=[],
+ )
+ stations.realize(clock)
+ _init_no_export(stations, torch.zeros(1), OrderedDict({"point": np.array([0])}))
+ with pytest.raises(CouplingError, match="no points= location metadata"):
+ Connector(ocean, stations, sample="nearest").execute(T0)
+
+
+def test_sample_requires_latlon_source():
+ """A source without lat/lon (e.g. a mediator pass-through) cannot be
+ auto-sampled onto points."""
+ clock = Clock(T0, "2024-01-02", "6h")
+
+ def step(x, coords):
+ return x, coords
+
+ src = CallableComponent(
+ "src", step, timestep="6h", exports=["sea_surface_temperature"]
+ )
+ src.realize(clock)
+ src.initialize(
+ torch.zeros(1, 4),
+ OrderedDict(
+ {"variable": np.array(["sst"]), "y": np.arange(4)}
+ ), # non-lat/lon spatial dim
+ )
+ points = PointSet(lat=np.array([0.0]), lon=np.array([0.0]))
+ stations = _points_destination(points)
+ stations.realize(clock)
+ _init_no_export(stations, torch.zeros(1), OrderedDict({"point": points.labels()}))
+ with pytest.raises(IncompatibleFieldError, match="needs lat/lon"):
+ Connector(src, stations, sample="nearest").execute(T0)
+
+
+def test_user_regridder_can_target_points():
+ """A custom regridder= still works for a point destination — the auto
+ sample= path is a convenience, not the only way in."""
+ atmos, ocean, clock = _realized_pair()
+ points = PointSet(lat=np.array([0.0, 0.0]), lon=np.array([0.0, 0.0]))
+ stations = _points_destination(points)
+ stations.realize(clock)
+ _init_no_export(
+ stations, torch.zeros(len(points)), OrderedDict({"point": points.labels()})
+ )
+ picked = lambda data: data[..., :2, 0] # trivial deterministic "sampler"
+ Connector(ocean, stations, regridder=picked).execute(T0)
+ got = stations.import_state["sea_surface_temperature"]
+ assert list(got.coords) == ["point"]
+ assert got.data.shape == (2,)
+
+
+def test_mask_fill_nearest_and_zero():
+ atmos, ocean, _ = _realized_pair(with_mask=True)
+ # poison the invalid (northern land) half; valid half stays 3.0
+ f = ocean.export_state["sea_surface_temperature"]
+ f.data = f.data.clone()
+ f.data[:8, :] = 999.0
+ Connector(ocean, atmos, fill="nearest").execute(T0)
+ sst = atmos.import_state["sea_surface_temperature"]
+ assert torch.all(sst.data < 4.0) # no 999 leaked through regrid
+ assert torch.allclose(sst.data, torch.full((32, 64), 3.0))
+ assert sst.mask is None # consumed
+
+ zero_conn = Connector(ocean, atmos, fill="zero")
+ zero_conn.execute(T0)
+ sst0 = atmos.import_state["sea_surface_temperature"]
+ assert sst0.data.min() == 0.0 # land became zero (then regridded)
+
+
+def test_time_policy_linear_extrapolates():
+ atmos, ocean, clock = _realized_pair()
+ conn = Connector(ocean, atmos, time_policy="linear")
+ f0 = ocean.export_state["sea_surface_temperature"]
+ conn.execute(T0) # first transfer: only one export, constant fallback
+ assert torch.allclose(
+ atmos.import_state["sea_surface_temperature"].data,
+ torch.full((32, 64), 3.0),
+ )
+ # ocean produces a new export at +48h with value 5.0
+ f1 = f0.clone()
+ f1.data.fill_(5.0)
+ f1.valid_time = T0 + np.timedelta64(48, "h")
+ ocean.export_state.add(f1)
+ # driver asks at +72h: linear extrapolation = 5 + (5-3) * 24/48 = 6
+ conn.execute(T0 + np.timedelta64(72, "h"))
+ got = atmos.import_state["sea_surface_temperature"]
+ assert torch.allclose(got.data, torch.full((32, 64), 6.0))
+ # constant policy would have given 5.0
+ const = Connector(ocean, atmos, time_policy="constant")
+ const.execute(T0 + np.timedelta64(72, "h"))
+ assert torch.allclose(
+ atmos.import_state["sea_surface_temperature"].data,
+ torch.full((32, 64), 5.0),
+ )
+
+
+def test_time_policy_linear_holds_slope_across_repeated_executes():
+ """Regression: repeated executes between source updates must keep
+ extrapolating along the (prev, latest) slope instead of silently
+ degrading to constant after the first extrapolated step."""
+ atmos, ocean, _ = _realized_pair()
+ conn = Connector(ocean, atmos, time_policy="linear")
+ f0 = ocean.export_state["sea_surface_temperature"]
+ f0.data = f0.data.clone().fill_(1.0)
+ conn.execute(T0) # seed history with 1.0 @ t0
+ f1 = f0.clone()
+ f1.data.fill_(2.0)
+ f1.valid_time = T0 + np.timedelta64(48, "h")
+ ocean.export_state.add(f1)
+ # slope is (2-1)/48h; every 6h step past 48h adds 0.125
+ for hours, expected in [(54, 2.125), (60, 2.25), (66, 2.375), (72, 2.5)]:
+ conn.execute(T0 + np.timedelta64(hours, "h"))
+ got = atmos.import_state["sea_surface_temperature"]
+ assert torch.allclose(
+ got.data, torch.full((32, 64), expected)
+ ), f"at +{hours}h expected {expected}, got {got.data.flatten()[0]}"
+
+
+def test_time_policy_linear_falls_back_for_lead_time_fields(caplog):
+ from earth2studio.nvcoupler.field import Field
+
+ atmos, ocean, _ = _realized_pair()
+ conn = Connector(ocean, atmos, time_policy="linear")
+
+ def lead_field(value, hours):
+ return Field(
+ torch.full((2, 16, 32), value),
+ OrderedDict(
+ {
+ "lead_time": np.array(
+ [np.timedelta64(48, "h"), np.timedelta64(96, "h")]
+ ),
+ **{k: v for k, v in ocean.grid_coords().items()},
+ }
+ ),
+ "sea_surface_temperature",
+ "K",
+ valid_time=T0 + np.timedelta64(hours, "h"),
+ )
+
+ ocean.export_state.add(lead_field(3.0, 0))
+ conn.execute(T0)
+ ocean.export_state.add(lead_field(5.0, 48))
+ conn.execute(T0 + np.timedelta64(72, "h"))
+ # linear would extrapolate to 6.0; lead_time fields must hold constant
+ got = atmos.import_state["sea_surface_temperature"]
+ assert torch.allclose(got.data, torch.full((2, 32, 64), 5.0))
+ assert any("undefined" in r.message for r in caplog.records)
+
+
+def test_probe_last_transfer():
+ atmos, ocean, _ = _realized_pair()
+ conn = Connector(ocean, atmos)
+ conn.execute(T0)
+ assert "sea_surface_temperature" in conn.last_transfer
+ assert conn.name == "ocean->atmos"
+
+
+def test_unproduced_export_raises():
+ atmos, ocean, _ = _realized_pair()
+ del ocean.export_state["sea_surface_temperature"]
+ with pytest.raises(CouplingError, match="has not produced"):
+ Connector(ocean, atmos).execute(T0)
+
+
+# ---------------------------------------------------------------------------
+# Vertical stage through the connector (chemistry-style coupling)
+# ---------------------------------------------------------------------------
+def _chem_dictionary():
+ d = FieldDictionary(DEFAULT_DICTIONARY)
+ d.register(FieldEntry("ozone_mixing_ratio", "kg kg-1", aliases=frozenset({"o3"})))
+ return d
+
+
+def _vertical_pair(export_ps=True):
+ """met (hybrid levels) -> chem (pressure levels) toy pair."""
+ d = _chem_dictionary()
+ hybrid = HybridLevels((30000.0, 20000.0, 0.0), (0.0, 0.5, 1.0))
+ pressure = PressureLevels((500.0, 850.0))
+ nlat, nlon = 4, 8
+
+ def met_step(x, coords):
+ return x, coords
+
+ met = CallableComponent(
+ "met",
+ met_step,
+ "6h",
+ exports=["ozone_mixing_ratio"],
+ dictionary=d,
+ export_vertical={"ozone_mixing_ratio": hybrid},
+ )
+ chem = CallableComponent(
+ "chem",
+ met_step,
+ "6h",
+ imports=["ozone_mixing_ratio"],
+ dictionary=d,
+ import_vertical={"ozone_mixing_ratio": pressure},
+ )
+ clock = Clock(T0, "2024-01-02", "6h")
+ met.realize(clock)
+ chem.realize(clock)
+
+ grid = OrderedDict(
+ {
+ "lat": np.linspace(90, -90, nlat),
+ "lon": np.linspace(0, 360, nlon, endpoint=False),
+ }
+ )
+ ps_value = 100000.0
+ p_src = np.array([30000.0, 70000.0, 100000.0])
+ o3 = (
+ torch.tensor(np.log(p_src), dtype=torch.float64)
+ .view(1, 3, 1, 1)
+ .expand(1, 3, nlat, nlon)
+ .clone()
+ )
+ # sp has no level dim so it cannot share o3's tensor; it is hand-added to
+ # the export state below (the connector reads export_state directly)
+ met_coords = OrderedDict(
+ {"variable": np.array(["o3"]), "level": np.arange(3.0), **grid}
+ )
+ met.initialize(o3, met_coords)
+ if export_ps:
+ from earth2studio.nvcoupler.field import Field
+
+ met.export_state.add(
+ Field(
+ torch.full((nlat, nlon), ps_value, dtype=torch.float64),
+ OrderedDict(grid),
+ "surface_pressure",
+ "Pa",
+ valid_time=T0,
+ source="met",
+ )
+ )
+ chem_coords = OrderedDict(
+ {"variable": np.array(["o3"]), "level": np.array([500.0, 850.0]), **grid}
+ )
+ chem.initialize(torch.zeros(1, 2, nlat, nlon, dtype=torch.float64), chem_coords)
+ return met, chem
+
+
+def test_vertical_stage_hybrid_to_pressure():
+ met, chem = _vertical_pair(export_ps=True)
+ Connector(met, chem, fields=["ozone_mixing_ratio"]).execute(T0)
+ o3 = chem.import_state["ozone_mixing_ratio"]
+ assert list(o3.coords["level"]) == [500.0, 850.0]
+ expected = np.log(np.array([50000.0, 85000.0]))
+ assert o3.data.shape == (2, 4, 8) # (level, lat, lon) after variable split
+ assert np.allclose(o3.data[:, 0, 0].numpy(), expected)
+ assert o3.vertical == PressureLevels((500.0, 850.0))
+
+
+def test_vertical_missing_ps_raises():
+ met, chem = _vertical_pair(export_ps=False)
+ with pytest.raises(VerticalMismatchError, match="surface_pressure"):
+ Connector(met, chem, fields=["ozone_mixing_ratio"]).execute(T0)
+
+
+# ---------------------------------------------------------------------------
+# HEALPix-style face grids: identity when identical, error when different
+# ---------------------------------------------------------------------------
+def _hpx_pair(dst_nside=4, src_nside=4):
+ """src/dst CallableComponents on (face, height, width) HEALPix-style
+ grids; dst_nside controls whether the grids match (4) or differ."""
+
+ def step(x, coords):
+ return x, coords
+
+ src = CallableComponent("hpx_src", step, "6h", exports=["sea_surface_temperature"])
+ dst = CallableComponent(
+ "hpx_dst",
+ step,
+ "6h",
+ imports=["sea_surface_temperature"],
+ exports=["geopotential_at_1000hpa"],
+ )
+ clock = Clock(T0, "2024-01-02", "6h")
+
+ def hpx_coords(nside):
+ return OrderedDict(
+ {
+ "variable": np.array(["sst"]),
+ "face": np.arange(12),
+ "height": np.arange(nside),
+ "width": np.arange(nside),
+ }
+ )
+
+ src.realize(clock)
+ src.initialize(
+ torch.full((1, 12, src_nside, src_nside), 3.0), hpx_coords(src_nside)
+ )
+ dst.realize(clock)
+ dst_coords = hpx_coords(dst_nside)
+ dst_coords["variable"] = np.array(["z1000"])
+ dst.initialize(torch.zeros(1, 12, dst_nside, dst_nside), dst_coords)
+ return src, dst
+
+
+def test_identical_face_grids_pass_through_identity():
+ src, dst = _hpx_pair(dst_nside=4)
+ conn = Connector(src, dst)
+ conn.execute(T0)
+ got = dst.import_state["sea_surface_temperature"]
+ assert torch.equal(got.data, src.export_state["sea_surface_temperature"].data)
+ assert list(got.coords) == ["face", "height", "width"]
+ assert torch.allclose(got.data, torch.full((12, 4, 4), 3.0))
+ assert conn._regridders == {} # identity path built no regridder
+
+
+def test_differing_face_grids_without_regridder_raise():
+ src, dst = _hpx_pair(dst_nside=8)
+ with pytest.raises(IncompatibleFieldError, match="regridder"):
+ Connector(src, dst).execute(T0)
+
+
+# ---------------------------------------------------------------------------
+# Windowed reductions (window= / reduce=): the connector-level mediator path
+# ---------------------------------------------------------------------------
+H6 = np.timedelta64(6, "h")
+
+WINDOWED_DSL = """
+@6h
+ atmos -> ocean
+ ocean -> atmos
+ atmos
+@48h
+ ocean
+@
+"""
+
+
+def test_windowed_connector_replaces_mediator_end_to_end():
+ """Acceptance: the canonical toy system WITHOUT a mediator — the windowed
+ connector in the fast slot reproduces test_driver.py's hand-computed
+ numbers exactly (first-window mean 4.2, z96 = 19.2336, sst = 2.180147)."""
+ from earth2studio.nvcoupler.driver import Driver
+
+ atmos, ocean = fake_atmos(), fake_ocean()
+ conn = Connector(atmos, ocean, window="48h", reduce="mean")
+ driver = Driver(
+ {"atmos": atmos, "ocean": ocean},
+ WINDOWED_DSL,
+ Clock(T0, "2024-01-05", "6h"),
+ connectors=[conn],
+ )
+ driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+ for time, _ in driver.steps():
+ if time == T0 + np.timedelta64(48, "h"):
+ # first window delivered: mean of z(0..42h) on the OCEAN grid
+ z48 = ocean.import_state["geopotential_at_1000hpa_48h_mean"]
+ assert z48.data.shape == (16, 32)
+ assert torch.allclose(z48.data, torch.full((16, 32), 4.2), atol=1e-5)
+ assert atmos.run_count == 16 and ocean.run_count == 2
+ z = atmos.export_state["geopotential_at_1000hpa"]
+ assert torch.allclose(z.data, torch.full((32, 64), 19.2336), atol=1e-4)
+ sst = ocean.export_state["sea_surface_temperature"]
+ assert torch.allclose(sst.data, torch.full((16, 32), 2.180147), atol=1e-6)
+ # probes carry the DERIVED name
+ assert "geopotential_at_1000hpa_48h_mean" in conn.last_transfer
+
+
+def _drive_windowed(conn, atmos, hours=48, executes_per_step=1):
+ """Manually drive the fast side: connector (lagged) then atmos, each 6h."""
+ for h in range(6, hours + 1, 6):
+ for _ in range(executes_per_step):
+ conn.execute(T0 + np.timedelta64(h, "h"))
+ atmos.run(T0 + np.timedelta64(h, "h"))
+
+
+def test_windowed_delivery_only_on_boundaries():
+ atmos, ocean, _ = _realized_pair()
+ conn = Connector(atmos, ocean, window="48h", reduce="mean")
+ assert conn.match() == [
+ "geopotential_at_1000hpa",
+ "geopotential_at_1000hpa_48h_mean",
+ ]
+ _drive_windowed(conn, atmos, hours=42)
+ # mid-window: nothing delivered, destination import untouched
+ assert "geopotential_at_1000hpa_48h_mean" not in ocean.import_state
+ assert conn.last_transfer == {}
+ conn.execute(T0 + np.timedelta64(48, "h")) # boundary (origin = t0 lineage)
+ z48 = ocean.import_state["geopotential_at_1000hpa_48h_mean"]
+ # atmos held at sst=2 (never fed): z(t) = 1.2*t/6h, mean(z(0..42h)) = 4.2
+ assert torch.allclose(z48.data, torch.full((16, 32), 4.2), atol=1e-5)
+ assert z48.valid_time == T0 + np.timedelta64(48, "h")
+ assert z48.source == "atmos"
+
+
+def test_windowed_duplicate_valid_time_not_double_counted():
+ atmos, ocean, _ = _realized_pair()
+ conn = Connector(atmos, ocean, window="48h", reduce="mean")
+ # execute the connector twice per step: same source valid_time, one sample
+ _drive_windowed(conn, atmos, hours=42, executes_per_step=2)
+ conn.execute(T0 + np.timedelta64(48, "h"))
+ z48 = ocean.import_state["geopotential_at_1000hpa_48h_mean"]
+ assert torch.allclose(z48.data, torch.full((16, 32), 4.2), atol=1e-5)
+
+
+def test_windowed_max_reduction():
+ from earth2studio.nvcoupler.field import Field
+ from earth2studio.nvcoupler.testing import grid_coords
+
+ def step(x, coords):
+ return x, coords
+
+ met = CallableComponent("met", step, "6h", exports=["air_temperature_2m"])
+ impact = CallableComponent(
+ "impact", step, "24h", imports=["air_temperature_2m_24h_max"]
+ )
+ clock = Clock(T0, "2024-01-02", "6h")
+ met.realize(clock)
+ impact.realize(clock)
+ coords = OrderedDict({"variable": np.array(["t2m"]), **grid_coords(4, 8)})
+ met.initialize(torch.full((1, 4, 8), 280.0), coords)
+ impact.initialize(torch.zeros(1, 4, 8), coords)
+ conn = Connector(met, impact, window="24h", reduce="max")
+ for i, v in enumerate([280.0, 295.0, 290.0, 285.0]):
+ met.export_state.add(
+ Field(
+ torch.full((4, 8), v),
+ grid_coords(4, 8),
+ "air_temperature_2m",
+ "K",
+ valid_time=T0 + i * H6,
+ source="met",
+ )
+ )
+ conn.execute(T0 + (i + 1) * H6)
+ tmax = impact.import_state["air_temperature_2m_24h_max"]
+ assert torch.allclose(tmax.data, torch.full((4, 8), 295.0))
+
+
+def test_window_and_reduce_must_come_together():
+ atmos, ocean, _ = _realized_pair()
+ with pytest.raises(CouplingError, match="set together"):
+ Connector(atmos, ocean, window="48h")
+ with pytest.raises(CouplingError, match="set together"):
+ Connector(atmos, ocean, reduce="mean")
+ with pytest.raises(CouplingError, match="unsupported reduce"):
+ Connector(atmos, ocean, window="48h", reduce="median")
+
+
+def test_windowed_without_derived_import_raises():
+ atmos, ocean, _ = _realized_pair()
+ # atmos imports plain sst — no CellMethod entry derives it, no invention
+ with pytest.raises(CouplingError, match="register a FieldEntry"):
+ Connector(ocean, atmos, window="48h", reduce="mean").match()
+ # window mismatch (24h vs the 48h the entry declares) must not match
+ with pytest.raises(CouplingError, match="register a FieldEntry"):
+ Connector(atmos, ocean, window="24h", reduce="mean").match()
+
+
+def test_connector_reset_clears_windowed_state():
+ atmos, ocean, _ = _realized_pair()
+ conn = Connector(atmos, ocean, window="48h", reduce="mean")
+ _drive_windowed(conn, atmos, hours=48)
+ assert conn.last_transfer
+ conn.reset()
+ assert conn.last_transfer == {} and conn._origin is None
+ assert "geopotential_at_1000hpa" not in conn._reduction
+
+
+def test_user_regridder_on_differing_face_grids():
+ """Regression: a user-supplied regridder= must work for non-latlon
+ (face, height, width) grids and rebuild coords from the destination."""
+ src, dst = _hpx_pair(dst_nside=4, src_nside=8)
+
+ def pool(x: torch.Tensor) -> torch.Tensor:
+ # mean-pool the trailing (height, width) dims 8x8 -> 4x4
+ return x.reshape(*x.shape[:-2], 4, 2, 4, 2).mean(dim=(-3, -1))
+
+ conn = Connector(src, dst, regridder=pool)
+ conn.execute(T0)
+ got = dst.import_state["sea_surface_temperature"]
+ assert got.data.shape == (12, 4, 4)
+ assert torch.allclose(got.data, torch.full((12, 4, 4), 3.0))
+ # coords carry the destination's face/height/width grid
+ assert list(got.coords) == ["face", "height", "width"]
+ dst_grid = dst.grid_coords()
+ for k in ("face", "height", "width"):
+ assert np.array_equal(got.coords[k], np.asarray(dst_grid[k]))
diff --git a/test/nvcoupler/test_data_diagnostic.py b/test/nvcoupler/test_data_diagnostic.py
new file mode 100644
index 000000000..2ff9204b6
--- /dev/null
+++ b/test/nvcoupler/test_data_diagnostic.py
@@ -0,0 +1,314 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""DataComponent (prescribed forcing) and DiagnosticComponent tests.
+
+Everything is mocked — no network, no model registries:
+
+- MockDataSource returns the exact xr.DataArray shape fetch_data expects
+ from a non-forecast DataSource: dims (time, variable, lat, lon).
+- MockDiagnostic implements the models/dx/base.py interface and computes
+ z500 = 0.5 * z1000, a hand-checkable single-step transform.
+"""
+
+from collections import OrderedDict
+
+import numpy as np
+import pytest
+import torch
+import xarray as xr
+
+from earth2studio.nvcoupler.clock import Clock
+from earth2studio.nvcoupler.component import DataComponent, DiagnosticComponent
+from earth2studio.nvcoupler.driver import Driver
+from earth2studio.nvcoupler.errors import CouplingError
+from earth2studio.nvcoupler.field import Field
+from earth2studio.nvcoupler.testing import (
+ ATMOS_GRID,
+ atmos_ic,
+ fake_atmos,
+ grid_coords,
+)
+
+T0 = "2024-01-01"
+T24 = "2024-01-02"
+T48 = "2024-01-03"
+
+
+# ---------------------------------------------------------------------------
+# Mocks
+# ---------------------------------------------------------------------------
+class MockDataSource:
+ """Deterministic in-memory DataSource: constant value per variable on a
+ small lat/lon grid, dims (time, variable, lat, lon)."""
+
+ def __init__(self, values: dict[str, float], nlat: int = 16, nlon: int = 32):
+ self.values = values
+ self.grid = grid_coords(nlat, nlon)
+ self.calls: list[tuple[np.ndarray, np.ndarray]] = []
+
+ def __call__(self, time, variable) -> xr.DataArray:
+ time = np.atleast_1d(np.asarray(time, dtype="datetime64[ns]"))
+ variable = np.atleast_1d(np.asarray(variable))
+ self.calls.append((time.copy(), variable.copy()))
+ lat, lon = self.grid["lat"], self.grid["lon"]
+ data = np.empty((len(time), len(variable), len(lat), len(lon)))
+ for j, v in enumerate(variable):
+ data[:, j] = self.values[str(v)]
+ return xr.DataArray(
+ data,
+ dims=["time", "variable", "lat", "lon"],
+ coords={"time": time, "variable": variable, "lat": lat, "lon": lon},
+ )
+
+
+class MockDiagnostic:
+ """z500 = 0.5 * z1000 on the atmos grid (models/dx/base.py interface)."""
+
+ def __init__(self, nlat: int = 32, nlon: int = 64):
+ self.grid = grid_coords(nlat, nlon)
+ self.call_count = 0
+
+ def input_coords(self):
+ return OrderedDict(
+ {
+ "batch": np.empty(0),
+ "variable": np.array(["z1000"]),
+ "lat": self.grid["lat"],
+ "lon": self.grid["lon"],
+ }
+ )
+
+ def output_coords(self, input_coords):
+ out = OrderedDict(input_coords)
+ out["variable"] = np.array(["z500"])
+ return out
+
+ def __call__(self, x, coords):
+ # the component must present variables in the model's raw vocabulary
+ assert list(coords["variable"]) == ["z1000"]
+ assert x.ndim == len(coords)
+ self.call_count += 1
+ return 0.5 * x, self.output_coords(coords)
+
+ def to(self, device):
+ return self
+
+
+# ---------------------------------------------------------------------------
+# DataComponent
+# ---------------------------------------------------------------------------
+def make_data_ocean(values=None, **kwargs):
+ source = MockDataSource(values or {"sst": 3.0})
+ comp = DataComponent(
+ "ocean",
+ source=source,
+ exports=["sea_surface_temperature"],
+ timestep="24h",
+ **kwargs,
+ )
+ return comp, source
+
+
+def test_data_component_fetch_and_publish():
+ comp, source = make_data_ocean()
+ comp.realize(Clock(T0, T48, "6h"))
+ comp.initialize() # no IC: fetches at clock.start
+ sst = comp.export_state["sea_surface_temperature"]
+ assert list(sst.coords) == ["lat", "lon"] # time/lead squeezed away
+ assert sst.data.shape == (16, 32)
+ assert torch.allclose(sst.data, torch.full((16, 32), 3.0))
+ assert sst.valid_time == np.datetime64(T0)
+ assert sst.units == "K"
+ # raw name resolved through dictionary aliases, fetched at clock.start
+ assert list(source.calls[0][1]) == ["sst"]
+ assert source.calls[0][0][0] == np.datetime64(T0)
+ # grid is known after the first fetch
+ grid = comp.grid_coords()
+ assert grid is not None and len(grid["lat"]) == 16 and len(grid["lon"]) == 32
+
+ t1 = np.datetime64(T24)
+ comp.run(t1)
+ assert comp.run_count == 1
+ assert comp.export_state["sea_surface_temperature"].valid_time == t1
+ assert source.calls[-1][0][0] == t1
+
+
+def test_data_component_variable_map():
+ comp, source = make_data_ocean(
+ values={"analysed_sst": 5.0},
+ variable_map={"sea_surface_temperature": "analysed_sst"},
+ )
+ comp.realize(Clock(T0, T24, "6h"))
+ comp.initialize()
+ assert list(source.calls[0][1]) == ["analysed_sst"]
+ sst = comp.export_state["sea_surface_temperature"]
+ assert torch.allclose(sst.data, torch.full((16, 32), 5.0))
+
+
+def test_data_component_initialize_before_realize_raises():
+ comp, _ = make_data_ocean()
+ with pytest.raises(CouplingError, match="realize"):
+ comp.initialize()
+
+
+def test_data_component_replaces_modeled_ocean_in_driver():
+ """Prescribed forcing: swap fake_ocean for a DataComponent; the atmos,
+ connector, and sequence stay as-is. z += 1 + 0.1 * sst each step."""
+ ocean, source = make_data_ocean(values={"sst": 3.0})
+ components = {"atmos": fake_atmos(), "ocean": ocean}
+ dsl = """
+@6h
+ ocean -> atmos
+ atmos
+@24h
+ ocean
+@
+"""
+ driver = Driver(components, dsl, Clock(T0, T48, "6h"))
+ # DataComponent needs no IC tensor; (None, None) means "fetch at t0"
+ driver.initialize({"atmos": atmos_ic(), "ocean": (None, None)})
+ ds = driver.run()
+
+ atmos = driver.components["atmos"]
+ assert atmos.run_count == 8 # 48h at 6h
+ assert ocean.run_count == 2 # 24h and 48h
+ # constant sst=3 from t0 on: z_n = n * (1 + 0.1 * 3) = 1.3 n
+ z = ds["atmos"]["geopotential_at_1000hpa"]
+ assert np.allclose(z.values, 1.3 * np.arange(9)[:, None, None], atol=1e-5)
+ # the connector regridded the 16x32 source field onto the atmos grid
+ sst_in = driver.probe("ocean->atmos")["sea_surface_temperature"]
+ assert sst_in.data.shape == ATMOS_GRID
+ assert torch.allclose(sst_in.data, torch.full(ATMOS_GRID, 3.0), atol=1e-6)
+
+
+def test_data_component_no_arg_initialize_standalone():
+ """DataComponent needs no IC: requires_ic is False and initialize()
+ with no arguments fetches at clock.start."""
+ comp, _ = make_data_ocean()
+ assert comp.requires_ic is False
+ comp.realize(Clock(T0, T24, "6h"))
+ comp.initialize() # no arguments
+ assert "sea_surface_temperature" in comp.export_state
+
+
+# ---------------------------------------------------------------------------
+# DiagnosticComponent
+# ---------------------------------------------------------------------------
+def test_diagnostic_defaults_from_model_coords():
+ comp = DiagnosticComponent("diag", MockDiagnostic(), timestep="6h")
+ assert comp.import_names == ["geopotential_at_1000hpa"]
+ assert comp.export_names == ["geopotential_at_500hpa"]
+ # grid_coords available straight from the model's input_coords
+ grid = comp.grid_coords()
+ assert grid is not None and len(grid["lat"]) == 32 and len(grid["lon"]) == 64
+
+
+def test_diagnostic_run_standalone():
+ model = MockDiagnostic()
+ comp = DiagnosticComponent("diag", model, timestep="6h")
+ comp.realize(Clock(T0, T24, "6h"))
+ comp.initialize() # tolerant of no state tensor
+ assert len(comp.export_state) == 0
+
+ t1 = np.datetime64("2024-01-01T06")
+ comp.import_state.add(
+ Field(
+ data=torch.full(ATMOS_GRID, 4.0),
+ coords=grid_coords(*ATMOS_GRID),
+ standard_name="geopotential_at_1000hpa",
+ units="m2 s-2",
+ valid_time=t1,
+ )
+ )
+ comp.run(t1)
+ out = comp.export_state["geopotential_at_500hpa"]
+ assert list(out.coords) == ["lat", "lon"] # batch dim squeezed back off
+ assert torch.allclose(out.data, torch.full(ATMOS_GRID, 2.0))
+ assert out.valid_time == t1
+ assert model.call_count == 1
+
+
+def test_diagnostic_no_arg_initialize_standalone():
+ """DiagnosticComponent needs no IC: requires_ic is False and a no-arg
+ initialize() derives its grid from the model's input_coords()."""
+ comp = DiagnosticComponent("diag", MockDiagnostic(), timestep="6h")
+ assert comp.requires_ic is False
+ comp.initialize() # no arguments, not even realized
+ grid = comp.grid_coords()
+ assert grid is not None and len(grid["lat"]) == 32 and len(grid["lon"]) == 64
+
+
+def test_stateful_components_require_ic():
+ assert fake_atmos().requires_ic is True
+
+
+def test_diagnostic_missing_import_raises():
+ comp = DiagnosticComponent("diag", MockDiagnostic(), timestep="6h")
+ comp.realize(Clock(T0, T24, "6h"))
+ comp.initialize()
+ with pytest.raises(CouplingError, match="missing imports"):
+ comp.run(np.datetime64("2024-01-01T06"))
+
+
+def test_diagnostic_chain_in_driver():
+ """fake_atmos exports z1000 -> diagnostic derives z500 = 0.5 * z1000."""
+ model = MockDiagnostic()
+ components = {
+ "atmos": fake_atmos(),
+ "diag": DiagnosticComponent("diag", model, timestep="6h"),
+ }
+ dsl = """
+@6h
+ atmos
+ atmos -> diag
+ diag
+@
+"""
+ # atmos's sst import is intentionally unfed (held at its IC value)
+ driver = Driver(components, dsl, Clock(T0, T24, "6h"), allow_unfed_imports=True)
+ driver.initialize({"atmos": atmos_ic(), "diag": (None, None)})
+ ds = driver.run()
+
+ diag = driver.components["diag"]
+ assert diag.run_count == 4 # every 6h over 24h
+ assert model.call_count == 4
+ # atmos: sst held at IC (2.0), so z_n = 1.2 n; diagnostic halves it
+ z500 = ds["diag"]["geopotential_at_500hpa"]
+ assert z500.dims == ("time", "lat", "lon")
+ assert np.allclose(
+ z500.values, 0.5 * 1.2 * np.arange(1, 5)[:, None, None], atol=1e-5
+ )
+ # exports valid at the driver time, same cadence as the source
+ out = diag.export_state["geopotential_at_500hpa"]
+ assert out.valid_time == np.datetime64(T24)
+
+
+def test_diagnostic_gradient_flows_through_chain():
+ gain = torch.tensor(1.0, requires_grad=True)
+ components = {
+ "atmos": fake_atmos(gain=gain),
+ "diag": DiagnosticComponent("diag", MockDiagnostic(), timestep="6h"),
+ }
+ dsl = "@6h\n atmos\n atmos -> diag\n diag\n@"
+ # atmos's sst import is intentionally unfed (held at its IC value)
+ driver = Driver(components, dsl, Clock(T0, T24, "6h"), allow_unfed_imports=True)
+ driver.initialize({"atmos": atmos_ic(), "diag": (None, None)})
+ with torch.enable_grad():
+ states = driver.rollout(4)
+ loss = states["diag"]["geopotential_at_500hpa"].data.sum()
+ loss.backward()
+ assert gain.grad is not None and gain.grad != 0
diff --git a/test/nvcoupler/test_dictionary.py b/test/nvcoupler/test_dictionary.py
new file mode 100644
index 000000000..eb376644a
--- /dev/null
+++ b/test/nvcoupler/test_dictionary.py
@@ -0,0 +1,129 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import numpy as np
+import pytest
+
+from earth2studio.nvcoupler.dictionary import (
+ DEFAULT_DICTIONARY,
+ CellMethod,
+ FieldDictionary,
+ FieldEntry,
+ normalize_units,
+)
+from earth2studio.nvcoupler.errors import UnitsMismatchError, UnknownFieldError
+
+
+def test_alias_resolution():
+ entry = DEFAULT_DICTIONARY.resolve("z1000")
+ assert entry.standard_name == "geopotential_at_1000hpa"
+ assert entry.canonical_units == "m2 s-2"
+ # standard name resolves to itself
+ assert DEFAULT_DICTIONARY.standard_name("geopotential_at_1000hpa") == (
+ "geopotential_at_1000hpa"
+ )
+
+
+def test_unknown_name_suggestions():
+ with pytest.raises(UnknownFieldError) as err:
+ DEFAULT_DICTIONARY.resolve("z1000h")
+ assert "z1000" in str(err.value)
+ assert "register" in str(err.value).lower()
+ # plain CouplingError, not a KeyError (no repr-quoting of the message)
+ assert not isinstance(err.value, KeyError)
+ assert not str(err.value).startswith('"')
+
+
+def test_register_and_alias():
+ d = FieldDictionary(DEFAULT_DICTIONARY)
+ d.register(FieldEntry("sea_ice_fraction", "", aliases=frozenset({"sic"})))
+ assert d.standard_name("sic") == "sea_ice_fraction"
+ # remapping an alias to a different standard name is an error
+ with pytest.raises(ValueError):
+ d.add_alias("sea_surface_temperature", "sic")
+ # alias colliding with a standard name is an error
+ with pytest.raises(ValueError):
+ d.add_alias("sea_surface_temperature", "sea_ice_fraction")
+ # copy did not pollute the default dictionary
+ assert "sic" not in DEFAULT_DICTIONARY
+
+
+def test_units_check():
+ DEFAULT_DICTIONARY.check_units(
+ "geopotential_at_1000hpa", "m**2 s**-2", src="a", dst="b"
+ ) # synonym normalizes, no raise
+ with pytest.raises(UnitsMismatchError) as err:
+ DEFAULT_DICTIONARY.check_units(
+ "geopotential_at_1000hpa", "K", src="ocean", dst="atmos"
+ )
+ assert "ocean" in str(err.value) and "K" in str(err.value)
+
+
+def test_normalize_units():
+ assert normalize_units("m/s") == "m s-1"
+ assert normalize_units("Kelvin") == "K"
+ # unknown units pass through in collapsed comparison form
+ assert normalize_units("W m-2") == normalize_units("w m**-2")
+
+
+def test_normalize_units_agrees_with_earthmover_lexicon():
+ """nvcoupler's normalizer must not disagree with the earthmover lexicon
+ normalizer (earth2studio/lexicon/earthmover.py): every pair of spellings
+ the lexicon treats as equivalent compares equal here too."""
+ equal_pairs = [
+ ("m s**-1", "m s-1"),
+ ("m s^-1", "m s-1"),
+ ("M/S", "m s**-1"),
+ ("m2 s-2", "m**2 s**-2"),
+ ("m^2/s^2", "m2s-2"),
+ ("(0-1)", "dimensionless"),
+ ("0-1", "fraction"),
+ ("1", "Dimensionless"),
+ ("%", "percent"),
+ ("degC", "celsius"),
+ ("degree_celsius", "Degrees Celsius"),
+ ("m of water equivalent", "mwe"),
+ ("KELVIN", "k"),
+ ]
+ for a, b in equal_pairs:
+ assert normalize_units(a) == normalize_units(b), (a, b)
+
+ # cross-check directly against the lexicon normalizer: anything it
+ # equates, we equate (nvcoupler may equate more, e.g. 'm/s' vs 'm s-1')
+ from earth2studio.lexicon.earthmover import (
+ normalize_units as lexicon_normalize_units,
+ )
+
+ for a, b in equal_pairs:
+ if lexicon_normalize_units(a) == lexicon_normalize_units(b):
+ assert normalize_units(a) == normalize_units(b), (a, b)
+
+ # distinct units stay distinct
+ assert normalize_units("K") != normalize_units("degC")
+ assert normalize_units("Pa") != normalize_units("hPa")
+
+
+def test_cell_method_derived_field():
+ entry = DEFAULT_DICTIONARY.resolve("geopotential_at_1000hpa_48h_mean")
+ cm = entry.cell_method
+ assert cm is not None
+ assert cm.base == "geopotential_at_1000hpa"
+ assert cm.method == "mean"
+ assert cm.window == np.timedelta64(48, "h")
+ # non-derived fields have none
+ assert DEFAULT_DICTIONARY.derived_from("sea_surface_temperature") is None
+ with pytest.raises(ValueError):
+ CellMethod("x", "median", np.timedelta64(1, "h"))
diff --git a/test/nvcoupler/test_dlesym_split.py b/test/nvcoupler/test_dlesym_split.py
new file mode 100644
index 000000000..33584d533
--- /dev/null
+++ b/test/nvcoupler/test_dlesym_split.py
@@ -0,0 +1,424 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Structural tests for the DLESyM split adapter (no real weights).
+
+MockDLESyM replicates the attribute surface split_dlesym relies on with
+tiny deterministic sub-models on an nside=8 HEALPix grid; sub-model inputs
+and outputs are captured so tests can assert exactly what crossed the
+coupling seam.
+"""
+
+from collections import OrderedDict
+
+import numpy as np
+import pytest
+import torch
+
+from earth2studio.nvcoupler.clock import Clock
+from earth2studio.nvcoupler.component import CallableComponent
+from earth2studio.nvcoupler.connector import Connector
+from earth2studio.nvcoupler.dlesym_split import (
+ DLESyMAtmosComponent,
+ DLESyMOceanComponent,
+ build_dlesym_driver,
+ split_dlesym,
+)
+from earth2studio.nvcoupler.errors import CouplingError
+
+NSIDE = 8
+START = np.datetime64("2024-01-01", "ns")
+STEP = np.timedelta64(96, "h")
+
+
+# ---------------------------------------------------------------------------
+# Mock DLESyM
+# ---------------------------------------------------------------------------
+class MockAtmosModel:
+ """HEALPixRecUNet-shaped callable: inputs [state (B,F,Tin,C,H,W),
+ insolation (B,F,Tsol,1,H,W), constants (F,Cc,H,W),
+ coupling (Tc,B,Cv,F,H,W)] -> (B,F,Tout,C,H,W)."""
+
+ input_time_dim = 4
+ output_time_dim = 16
+
+ def __init__(self, log: list):
+ self.calls: list[dict] = []
+ self._log = log
+
+ def __call__(self, inputs: list[torch.Tensor]) -> torch.Tensor:
+ state, sol, const, coupling = inputs
+ B, F, tin, C, H, W = state.shape
+ assert (F, H, W) == (12, NSIDE, NSIDE)
+ assert tin == self.input_time_dim
+ assert sol.shape == (B, F, tin + self.output_time_dim, 1, H, W)
+ assert const.shape[0] == F
+ assert coupling.shape[0] == 1 + self.output_time_dim // self.input_time_dim
+ leads = torch.arange(1, self.output_time_dim + 1, dtype=state.dtype)
+ leads = leads.view(1, 1, -1, 1, 1, 1)
+ cp = coupling.mean(dim=0).mean(dim=1) # (B, F, H, W)
+ out = (
+ state[:, :, -1].unsqueeze(2)
+ + 0.01 * leads
+ + 0.1 * cp[:, :, None, None, :, :]
+ )
+ self.calls.append({"inputs": inputs, "output": out})
+ self._log.append("atmos")
+ return out
+
+
+class MockOceanModel:
+ """HEALPixUNet-shaped callable; coupling arrives as
+ (lead=1, B, n_window * n_coupling_vars, F, H, W)."""
+
+ input_time_dim = 2
+ output_time_dim = 2
+
+ def __init__(self, log: list):
+ self.calls: list[dict] = []
+ self._log = log
+
+ def __call__(self, inputs: list[torch.Tensor]) -> torch.Tensor:
+ state, sol, const, coupling = inputs
+ B, F, tin, C, H, W = state.shape
+ assert (F, H, W) == (12, NSIDE, NSIDE)
+ assert coupling.shape[:2] == (1, B) and coupling.shape[3] == F
+ leads = torch.arange(1, self.output_time_dim + 1, dtype=state.dtype)
+ leads = leads.view(1, 1, -1, 1, 1, 1)
+ cp = coupling.mean(dim=(0, 2)) # (B, F, H, W)
+ out = (
+ state[:, :, -1].unsqueeze(2)
+ + 0.05 * leads
+ + 0.2 * cp[:, :, None, None, :, :]
+ )
+ self.calls.append({"inputs": inputs, "output": out})
+ self._log.append("ocean")
+ return out
+
+
+class MockDLESyM:
+ """Attribute-compatible stand-in for earth2studio.models.px.DLESyM."""
+
+ def __init__(self, nside: int = NSIDE):
+ self.nside = nside
+ self.atmos_variables = ["z1000", "ws10m", "t2m"]
+ self.ocean_variables = ["sst"]
+ self.atmos_coupling_variables = ["sst"]
+ self.ocean_coupling_variables = ["z1000", "ws10m"]
+
+ self.atmos_input_times = np.array([-18, -12, -6, 0], dtype="timedelta64[h]")
+ self.ocean_input_times = np.array([-48, 0], dtype="timedelta64[h]")
+ self.atmos_output_times = np.arange(6, 97, 6).astype("timedelta64[h]")
+ self.ocean_output_times = np.array([48, 96], dtype="timedelta64[h]")
+ self.full_input_times = np.arange(
+ self.ocean_input_times[0],
+ self.atmos_input_times[-1] + 1,
+ self.atmos_input_times[1] - self.atmos_input_times[0],
+ )
+ self.atmos_sol_times = np.concatenate(
+ [self.atmos_input_times, self.atmos_output_times]
+ )
+ self.ocean_sol_times = np.concatenate(
+ [self.ocean_input_times, self.ocean_output_times]
+ )
+
+ self.call_log: list[str] = []
+ self.atmos_model = MockAtmosModel(self.call_log)
+ self.ocean_model = MockOceanModel(self.call_log)
+
+ n_atmos_steps = 1 + max(
+ self.atmos_model.output_time_dim // self.atmos_model.input_time_dim, 1
+ )
+ full = list(self.full_input_times)
+ out = list(self.atmos_output_times)
+ self.atmos_input_lt_idx = [full.index(t) for t in self.atmos_input_times]
+ self.ocean_input_lt_idx = [full.index(t) for t in self.ocean_input_times]
+ self.atmos_coupled_input_lt_idx = [
+ full.index(self.atmos_input_times[-1])
+ ] * n_atmos_steps
+ self.ocean_output_lt_idx = [out.index(t) for t in self.ocean_output_times]
+
+ variables = self.atmos_variables + self.ocean_variables
+ self.atmos_var_idx = [variables.index(v) for v in self.atmos_variables]
+ self.ocean_var_idx = [variables.index(v) for v in self.ocean_variables]
+ self.atmos_coupling_var_idx = [
+ variables.index(v) for v in self.atmos_coupling_variables
+ ]
+ self.ocean_coupling_var_idx = [
+ variables.index(v) for v in self.ocean_coupling_variables
+ ]
+
+ nvar = len(variables)
+ self.center = torch.zeros(1, 1, 1, nvar, 1, 1, 1)
+ self.scale = torch.ones(1, 1, 1, nvar, 1, 1, 1)
+ self.atmos_constants = torch.zeros(12, 2, nside, nside)
+ self.ocean_constants = torch.zeros(12, 1, nside, nside)
+
+ def _make_insolation_tensor(
+ self, anchor_times: np.ndarray, timedeltas: np.ndarray
+ ) -> torch.Tensor:
+ return torch.zeros(
+ len(anchor_times), 12, len(timedeltas), 1, self.nside, self.nside
+ )
+
+ # Copied verbatim from earth2studio.models.px.DLESyM — dlesym_split now
+ # calls these parent methods instead of re-implementing their math, so
+ # the mock must expose them with identical semantics.
+ def _make_atmos_coupling(self, x: torch.Tensor, coords) -> torch.Tensor:
+ atmos_coupling = x[:, self.atmos_coupled_input_lt_idx][
+ ..., self.atmos_coupling_var_idx, :, :, :
+ ].permute(1, 0, 2, 3, 4, 5)
+ return atmos_coupling
+
+ def _make_ocean_coupling(self, x: torch.Tensor, coords) -> torch.Tensor:
+ ocean_coupling = x[:, :, :, self.ocean_coupling_var_idx, :, :]
+ slices = ocean_coupling.chunk(len(self.ocean_output_times), dim=2)
+ ocean_coupling = torch.concat(
+ [s.mean(dim=2, keepdim=True) for s in slices], dim=3
+ )
+ ocean_coupling = ocean_coupling.permute(2, 0, 3, 1, 4, 5)
+ return ocean_coupling
+
+
+def make_ic(mock: MockDLESyM, seed: int = 0) -> tuple[torch.Tensor, OrderedDict]:
+ g = torch.Generator().manual_seed(seed)
+ nvar = len(mock.atmos_variables + mock.ocean_variables)
+ nlead = len(mock.full_input_times)
+ x = torch.rand(1, 1, nlead, nvar, 12, mock.nside, mock.nside, generator=g)
+ coords = OrderedDict(
+ {
+ "batch": np.arange(1),
+ "time": np.array([START]),
+ "lead_time": mock.full_input_times.copy(),
+ "variable": np.array(mock.atmos_variables + mock.ocean_variables),
+ "face": np.arange(12),
+ "height": np.arange(mock.nside),
+ "width": np.arange(mock.nside),
+ }
+ )
+ return x, coords
+
+
+def expected_ocean_coupling(atmos_out: torch.Tensor, mock: MockDLESyM) -> torch.Tensor:
+ """Replicate DLESyM._make_ocean_coupling on a (B,F,L,C,H,W) atmos output."""
+ oc = atmos_out[:, :, :, mock.ocean_coupling_var_idx, :, :]
+ slices = oc.chunk(len(mock.ocean_output_times), dim=2)
+ oc = torch.concat([s.mean(dim=2, keepdim=True) for s in slices], dim=3)
+ return oc.permute(2, 0, 3, 1, 4, 5)
+
+
+# ---------------------------------------------------------------------------
+# Connector face-identity probe (a concurrent change makes identical HEALPix
+# grids pass through as identity; xfail integration tests until it lands)
+# ---------------------------------------------------------------------------
+def _face_identity_supported() -> bool:
+ coords = OrderedDict(
+ {
+ "variable": np.array(["sst"]),
+ "face": np.arange(12),
+ "height": np.arange(2),
+ "width": np.arange(2),
+ }
+ )
+ src = CallableComponent("src", lambda x, c: (x, c), "6h", exports=["sst"])
+ dst = CallableComponent("dst", lambda x, c: (x, c), "6h", imports=["sst"])
+ x = torch.zeros(1, 12, 2, 2)
+ src.initialize(x, coords)
+ dst.initialize(x, coords)
+ try:
+ Connector(src, dst).execute(np.datetime64("2024-01-01"))
+ except CouplingError:
+ return False
+ return True
+
+
+face_xfail = pytest.mark.xfail(
+ condition=not _face_identity_supported(),
+ reason="pending connector face fix",
+ strict=False,
+)
+
+
+# ---------------------------------------------------------------------------
+# (1) structural: split produces the right components
+# ---------------------------------------------------------------------------
+def test_split_structure():
+ mock = MockDLESyM()
+ atmos, ocean = split_dlesym(mock)
+
+ assert isinstance(atmos, DLESyMAtmosComponent)
+ assert isinstance(ocean, DLESyMOceanComponent)
+ assert atmos.timestep == np.timedelta64(96, "h").astype("timedelta64[ns]")
+ assert ocean.timestep == atmos.timestep
+
+ assert atmos.import_names == ["sea_surface_temperature"]
+ derived = [
+ "geopotential_at_1000hpa_48h_mean",
+ "wind_speed_10m_48h_mean",
+ ]
+ for name in ["geopotential_at_1000hpa", "wind_speed_10m", "air_temperature_2m"]:
+ assert name in atmos.export_names
+ for name in derived:
+ assert name in atmos.export_names
+
+ assert ocean.import_names == derived
+ assert ocean.export_names == ["sea_surface_temperature"]
+
+
+def test_split_registers_unknown_variables():
+ mock = MockDLESyM()
+ mock.atmos_variables = ["z1000", "ws10m", "mystery_var"]
+ atmos, _ = split_dlesym(mock)
+ assert "mystery_var" in atmos.dictionary
+ assert "mystery_var" in atmos.export_names
+
+
+def test_split_rejects_uneven_windows():
+ mock = MockDLESyM()
+ mock.ocean_output_times = np.array([32, 64, 96], dtype="timedelta64[h]")
+ with pytest.raises(CouplingError, match="chunk evenly"):
+ split_dlesym(mock)
+
+
+# ---------------------------------------------------------------------------
+# (2a) one step, standalone (no connectors): coupling tensor math
+# ---------------------------------------------------------------------------
+def _standalone_pair(mock, n_steps=1):
+ atmos, ocean = split_dlesym(mock)
+ clock = Clock(START, START + n_steps * STEP, dt=STEP)
+ atmos.realize(clock)
+ ocean.realize(clock)
+ x, coords = make_ic(mock)
+ atmos.initialize(x, coords)
+ ocean.initialize(x, coords)
+ return atmos, ocean, x
+
+
+def _manual_step(atmos, ocean, time):
+ # ocean -> atmos (lagged SST), atmos, atmos -> ocean, ocean
+ atmos.import_state.add(ocean.export_state["sea_surface_temperature"])
+ atmos.run(time)
+ for name in ocean.import_names:
+ ocean.import_state.add(atmos.export_state[name])
+ ocean.run(time)
+
+
+def test_standalone_step_delivers_chunk_mean_coupling():
+ mock = MockDLESyM()
+ atmos, ocean, x = _standalone_pair(mock)
+ _manual_step(atmos, ocean, START + STEP)
+
+ assert mock.call_log == ["atmos", "ocean"]
+
+ # atmos coupling = IC SST at lead 0, persisted over all internal sub-steps
+ atmos_coupling = mock.atmos_model.calls[-1]["inputs"][3]
+ sst0 = x[:, :, -1, 3].reshape(-1, 12, NSIDE, NSIDE) # normalized == physical
+ assert atmos_coupling.shape[0] == len(mock.atmos_coupled_input_lt_idx)
+ for k in range(atmos_coupling.shape[0]):
+ torch.testing.assert_close(atmos_coupling[k, :, 0], sst0)
+
+ # ocean coupling tensor == _make_ocean_coupling chunk means of atmos output
+ atmos_out = mock.atmos_model.calls[-1]["output"]
+ expected = expected_ocean_coupling(atmos_out, mock)
+ received = mock.ocean_model.calls[-1]["inputs"][3]
+ torch.testing.assert_close(received, expected)
+
+ # SST export is the ocean output at the 96 h lead
+ ocean_out = mock.ocean_model.calls[-1]["output"]
+ sst_export = ocean.export_state["sea_surface_temperature"]
+ torch.testing.assert_close(
+ sst_export.data.reshape(-1, 12, NSIDE, NSIDE), ocean_out[:, :, -1, 0]
+ )
+ assert sst_export.valid_time == START + STEP
+
+
+def test_standalone_two_steps_lagged_sst_and_window():
+ mock = MockDLESyM()
+ atmos, ocean, _ = _standalone_pair(mock, n_steps=2)
+ _manual_step(atmos, ocean, START + STEP)
+ _manual_step(atmos, ocean, START + 2 * STEP)
+
+ # lagged feedback: atmos step-2 coupling == ocean step-1 SST at 96 h
+ sst_96h = mock.ocean_model.calls[0]["output"][:, :, -1, 0]
+ coupling2 = mock.atmos_model.calls[1]["inputs"][3]
+ for k in range(coupling2.shape[0]):
+ torch.testing.assert_close(coupling2[k, :, 0], sst_96h)
+
+ # sliding window: atmos step-2 state == step-1 outputs at 78/84/90/96 h
+ out1 = mock.atmos_model.calls[0]["output"] # (B, F, 16, C, H, W)
+ state2 = mock.atmos_model.calls[1]["inputs"][0]
+ torch.testing.assert_close(state2, out1[:, :, 12:16])
+
+ # ocean window: step-2 state == step-1 outputs at 48/96 h
+ oout1 = mock.ocean_model.calls[0]["output"]
+ ostate2 = mock.ocean_model.calls[1]["inputs"][0]
+ torch.testing.assert_close(ostate2, oout1)
+
+
+def test_ocean_run_without_imports_raises():
+ mock = MockDLESyM()
+ _, ocean, _ = _standalone_pair(mock)
+ with pytest.raises(CouplingError, match="atmos -> ocean connector"):
+ ocean.run(START + STEP)
+
+
+def test_gradients_flow_across_split():
+ mock = MockDLESyM()
+ atmos, ocean, _ = _standalone_pair(mock)
+ x, coords = make_ic(mock)
+ x = x.clone().requires_grad_(True)
+ atmos.initialize(x, coords)
+ ocean.initialize(x, coords)
+ _manual_step(atmos, ocean, START + STEP)
+ loss = ocean.export_state["sea_surface_temperature"].data.sum()
+ loss.backward()
+ assert x.grad is not None
+ assert torch.isfinite(x.grad).all()
+
+
+# ---------------------------------------------------------------------------
+# (2b)/(3) driver integration through connectors (identity on HEALPix grid)
+# ---------------------------------------------------------------------------
+@face_xfail
+def test_driver_one_step_atmos_then_ocean():
+ mock = MockDLESyM()
+ driver = build_dlesym_driver(mock, START, START + STEP)
+ x, coords = make_ic(mock)
+ driver.initialize({"atmos": (x, coords), "ocean": (x, coords)})
+ for _time, _states in driver.steps():
+ pass
+
+ assert mock.call_log == ["atmos", "ocean"]
+ atmos_out = mock.atmos_model.calls[-1]["output"]
+ expected = expected_ocean_coupling(atmos_out, mock)
+ received = mock.ocean_model.calls[-1]["inputs"][3]
+ torch.testing.assert_close(received, expected)
+
+
+@face_xfail
+def test_driver_two_steps_sst_feedback():
+ mock = MockDLESyM()
+ driver = build_dlesym_driver(mock, START, START + 2 * STEP)
+ x, coords = make_ic(mock)
+ driver.initialize({"atmos": (x, coords), "ocean": (x, coords)})
+ for _time, _states in driver.steps():
+ pass
+
+ assert mock.call_log == ["atmos", "ocean", "atmos", "ocean"]
+ sst_96h = mock.ocean_model.calls[0]["output"][:, :, -1, 0]
+ coupling2 = mock.atmos_model.calls[1]["inputs"][3]
+ for k in range(coupling2.shape[0]):
+ torch.testing.assert_close(coupling2[k, :, 0], sst_96h)
diff --git a/test/nvcoupler/test_dlesym_weights_equivalence.py b/test/nvcoupler/test_dlesym_weights_equivalence.py
new file mode 100644
index 000000000..a3f22cce7
--- /dev/null
+++ b/test/nvcoupler/test_dlesym_weights_equivalence.py
@@ -0,0 +1,173 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""The real-weights equivalence gate for the DLESyM split adapter.
+
+Every other dlesym_split test runs against a MockDLESyM authored from a
+reading of dlesym.py — which makes them structurally circular: a misreading
+of the model (normalization order, insolation times, window chunking) would
+pass all of them and fail on real weights. THIS test is the actual proof:
+driving the split components through nvcoupler must reproduce the native
+``DLESyM.__call__`` output on the same input, with real checkpoints.
+
+It is skipped unless the environment can run it (physicsnemo installed and
+the hf://nvidia/dlesym-v1-era5 package fetchable — several GB; set
+NVCOUPLER_DLESYM_WEIGHTS=1 to opt in). Until it has passed somewhere,
+treat "nvcoupler can host DLESyM" as unverified.
+"""
+
+import os
+from collections import OrderedDict
+
+import numpy as np
+import pytest
+import torch
+
+requires_weights = pytest.mark.skipif(
+ os.environ.get("NVCOUPLER_DLESYM_WEIGHTS") != "1",
+ reason=(
+ "real-weights equivalence gate: set NVCOUPLER_DLESYM_WEIGHTS=1 with "
+ "physicsnemo installed and network/cache access to the "
+ "hf://nvidia/dlesym-v1-era5 package (several GB)"
+ ),
+)
+
+
+def native_final_lead_slab(
+ y_native: torch.Tensor, y_coords: OrderedDict, var_index: int
+) -> torch.Tensor:
+ """Select variable ``var_index`` at the LAST lead time from a native
+ DLESyM output tensor laid out per ``y_coords``.
+
+ The first ``select`` removes the variable axis, which renumbers every
+ axis AFTER it by -1 while axes before it keep their indices. In the
+ native layout (batch, time, lead_time, variable, ...) lead_time sits
+ BEFORE variable, so its index must NOT be shifted — ``lead_axis - 1``
+ would silently select the time axis instead.
+ """
+ var_axis = list(y_coords).index("variable")
+ lead_axis = list(y_coords).index("lead_time")
+ slab = y_native.select(var_axis, var_index)
+ reduced_lead_axis = lead_axis if lead_axis < var_axis else lead_axis - 1
+ return slab.select(reduced_lead_axis, -1)
+
+
+def test_native_final_lead_slab_axis_arithmetic():
+ """Cheap structural check of the comparison helper: the selected slab
+ must be exactly (var, last-lead) for the native DLESyM axis order, with
+ the lead axis NOT shifted after the variable select removes a later axis.
+ Runs without weights so the gate's axis math is verified by execution.
+ """
+ coords = OrderedDict(
+ {
+ "batch": np.arange(1),
+ "time": np.arange(2),
+ "lead_time": np.arange(3),
+ "variable": np.arange(2),
+ "face": np.arange(2),
+ "height": np.arange(2),
+ "width": np.arange(2),
+ }
+ )
+ shape = tuple(len(v) for v in coords.values())
+ x = torch.zeros(shape)
+ for lead in range(shape[2]):
+ for var in range(shape[3]):
+ # encode (var, lead) so any wrong-axis select is detectable
+ x[:, :, lead, var] = 100 * var + lead
+
+ for var in range(shape[3]):
+ got = native_final_lead_slab(x, coords, var)
+ assert got.shape == (1, 2, 2, 2, 2) # variable and lead_time removed
+ assert torch.all(got == 100 * var + (shape[2] - 1))
+
+ # a layout with variable BEFORE lead_time exercises the shifted branch
+ coords_vl = OrderedDict(
+ {
+ "batch": np.arange(1),
+ "variable": np.arange(2),
+ "lead_time": np.arange(3),
+ "face": np.arange(2),
+ }
+ )
+ y = torch.zeros(1, 2, 3, 2)
+ for var in range(2):
+ for lead in range(3):
+ y[:, var, lead] = 100 * var + lead
+ for var in range(2):
+ got = native_final_lead_slab(y, coords_vl, var)
+ assert got.shape == (1, 2)
+ assert torch.all(got == 100 * var + 2)
+
+
+@requires_weights
+def test_split_adapter_matches_native_dlesym():
+ pytest.importorskip("physicsnemo", reason="DLESyM checkpoints need physicsnemo")
+ from earth2studio.models.px import DLESyM
+ from earth2studio.nvcoupler.dlesym_split import build_dlesym_driver
+
+ model = DLESyM.load_model(DLESyM.load_default_package())
+ model.eval()
+
+ # Equivalence must hold for ANY input, so a random (but finite,
+ # reasonably-scaled) state on the native input coords suffices — no data
+ # source needed. Batch/time dims of size 1 replace the empty wildcards.
+ ic = model.input_coords()
+ coords = OrderedDict(
+ {
+ "batch": np.array([0]),
+ "time": np.array([np.datetime64("2024-01-01")]),
+ **{k: v for k, v in ic.items() if k not in ("batch", "time")},
+ }
+ )
+ shape = tuple(len(v) for v in coords.values())
+ torch.manual_seed(0)
+ x = torch.randn(shape)
+
+ # native: one coupled 96h step
+ with torch.inference_mode():
+ y_native, y_coords = model(x.clone(), coords)
+
+ # nvcoupler: the split components driven for one coupling step
+ driver = build_dlesym_driver(model, start="2024-01-01", stop="2024-01-05")
+ driver.initialize({"atmos": (x.clone(), coords), "ocean": (x.clone(), coords)})
+ states = driver.rollout(driver.clock.n_steps)
+
+ atmos = states["atmos"]
+ ocean = states["ocean"]
+
+ # compare every variable at the final lead time (96 h): the atmos
+ # component exports its prognostics at the 96 h window end, which is
+ # native lead_time[-1]; 96 h is also an ocean output time, so the
+ # native SST slice at lead_time[-1] is valid (other atmos leads hold
+ # uninitialized memory for ocean variables — never index those).
+ for i, raw in enumerate(y_coords["variable"]):
+ raw = str(raw)
+ if raw in model.atmos_variables:
+ std = driver.components["atmos"].dictionary.standard_name(raw)
+ got = atmos[std].data
+ else:
+ std = driver.components["ocean"].dictionary.standard_name(raw)
+ got = ocean[std].data
+ # align on the last valid lead time for a like-for-like check
+ native_last = native_final_lead_slab(y_native, y_coords, i)
+ torch.testing.assert_close(
+ got.reshape(native_last.shape).float(),
+ native_last.float(),
+ rtol=1e-4,
+ atol=1e-4,
+ msg=f"nvcoupler output diverges from native DLESyM for {raw!r}",
+ )
diff --git a/test/nvcoupler/test_driver.py b/test/nvcoupler/test_driver.py
new file mode 100644
index 000000000..5f159dd5e
--- /dev/null
+++ b/test/nvcoupler/test_driver.py
@@ -0,0 +1,441 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""End-to-end coupled-system tests on the synthetic atmos/ocean toys.
+
+Hand-computed expectations (gain = 1, atmos z0 = 0, sst0 = 2, dt = 6 h):
+
+- atmos step: z += 1 + 0.1 * sst. With sst = 2 held for the first 48 h,
+ z(t) = 1.2 * t/6h, so z(42h) = 8.4 and z(48h) = 9.6.
+- mediator window 1 accumulates z at t = 0..42h (lagged transfer before each
+ atmos run): mean = 1.2 * (0+..+7)/8 = 4.2.
+- ocean at 48 h: sst = 2 + 0.01 * 4.2 = 2.042.
+- atmos then steps at 1.2042 for 48..96 h: z(96h) = 9.6 + 8 * 1.2042 = 19.2336.
+- mediator window 2: mean of z(48..90h) = (9.6 + 18.0294)/2 = 13.8147.
+- ocean at 96 h: sst = 2.042 + 0.138147 = 2.180147.
+"""
+
+from collections import OrderedDict
+from contextlib import contextmanager
+
+import numpy as np
+import pytest
+import torch
+
+from earth2studio.nvcoupler.clock import Clock
+from earth2studio.nvcoupler.component import CallableComponent
+from earth2studio.nvcoupler.driver import Driver
+from earth2studio.nvcoupler.errors import CouplingError, UnmatchedImportError
+from earth2studio.nvcoupler.mediator import TrailingAverageMediator
+from earth2studio.nvcoupler.testing import (
+ ATMOS_GRID,
+ atmos_ic,
+ fake_atmos,
+ fake_ocean,
+ ocean_ic,
+)
+
+T0 = "2024-01-01"
+T96 = "2024-01-05"
+
+LAGGED_DSL = """
+@6h
+ atmos -> med
+ ocean -> atmos
+ atmos
+@48h
+ med.compute
+ med -> ocean
+ ocean
+@
+"""
+
+
+@contextmanager
+def capture_loguru(level="WARNING"):
+ """Collect loguru messages emitted inside the block (repo loguru pattern)."""
+ from loguru import logger
+
+ messages: list[str] = []
+ handler_id = logger.add(messages.append, level=level, format="{message}")
+ try:
+ yield messages
+ finally:
+ logger.remove(handler_id)
+
+
+def make_driver(dsl=LAGGED_DSL, gain_atmos=1.0, gain_ocean=1.0, stop=T96):
+ components = {
+ "atmos": fake_atmos(gain=gain_atmos),
+ "ocean": fake_ocean(gain=gain_ocean),
+ "med": TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"]),
+ }
+ driver = Driver(components, dsl, Clock(T0, stop, "6h"))
+ driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+ return driver
+
+
+def make_declarative_driver(stop=T96):
+ """The same trio declared as a coupling graph — no sequence given."""
+ components = {
+ "atmos": fake_atmos(),
+ "ocean": fake_ocean(),
+ "med": TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"]),
+ }
+ driver = Driver(
+ components,
+ clock=Clock(T0, stop, "6h"),
+ connectors=[("atmos", "med"), ("ocean", "atmos"), ("med", "ocean")],
+ )
+ driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+ return driver
+
+
+def test_declarative_driver_derives_the_canonical_sequence():
+ driver = make_declarative_driver()
+ assert driver.sequence_derived
+ from earth2studio.nvcoupler.sequence import parse_run_sequence
+
+ assert str(driver.sequence) == str(parse_run_sequence(LAGGED_DSL))
+
+
+def test_declarative_driver_matches_hand_computed_values():
+ driver = make_declarative_driver()
+ driver.run()
+ z = driver.components["atmos"].export_state["geopotential_at_1000hpa"]
+ assert torch.allclose(z.data, torch.full(ATMOS_GRID, 19.2336), atol=1e-4)
+ sst = driver.components["ocean"].export_state["sea_surface_temperature"]
+ assert torch.allclose(sst.data, torch.full((16, 32), 2.180147), atol=1e-6)
+ zmean = driver.components["med"].export_state["geopotential_at_1000hpa_48h_mean"]
+ assert torch.allclose(zmean.data, torch.full(ATMOS_GRID, 13.8147), atol=1e-4)
+
+
+def test_explicit_sequence_marks_driver_not_derived():
+ driver = make_driver()
+ assert not driver.sequence_derived
+
+
+def test_declarative_unknown_connection_name_raises():
+ with pytest.raises(CouplingError, match="atmso"):
+ Driver(
+ {"atmos": fake_atmos()},
+ clock=Clock(T0, "2024-01-02", "6h"),
+ connectors=[("atmso", "atmos")],
+ )
+
+
+def test_driver_without_clock_raises():
+ with pytest.raises(CouplingError, match="clock"):
+ Driver({"atmos": fake_atmos()})
+
+
+def test_cadence_and_hand_computed_values():
+ driver = make_driver()
+ driver.run()
+ atmos = driver.components["atmos"]
+ ocean = driver.components["ocean"]
+ med = driver.components["med"]
+ assert atmos.run_count == 16
+ assert ocean.run_count == 2
+ assert med.run_count == 2
+
+ z = atmos.export_state["geopotential_at_1000hpa"]
+ assert torch.allclose(z.data, torch.full(ATMOS_GRID, 19.2336), atol=1e-4)
+ assert z.valid_time == np.datetime64("2024-01-05")
+
+ sst = ocean.export_state["sea_surface_temperature"]
+ assert torch.allclose(sst.data, torch.full((16, 32), 2.180147), atol=1e-6)
+
+ zmean = med.export_state["geopotential_at_1000hpa_48h_mean"]
+ assert torch.allclose(zmean.data, torch.full(ATMOS_GRID, 13.8147), atol=1e-4)
+ assert med.samples_last_window["geopotential_at_1000hpa_48h_mean"] == 8
+
+
+def test_lagged_vs_sequential_within_slot():
+ # Both variants isolate the sst hand-off to the 48h slot; they differ only
+ # in whether atmos receives the sst produced in that same slot (sequential)
+ # or the one from before the ocean ran (lagged).
+ lagged = """
+@6h
+ atmos -> med
+ atmos
+@48h
+ med.compute
+ ocean -> atmos
+ med -> ocean
+ ocean
+@
+"""
+ sequential = """
+@6h
+ atmos -> med
+ atmos
+@48h
+ med.compute
+ med -> ocean
+ ocean
+ ocean -> atmos
+@
+"""
+ z_lagged = make_driver(lagged).run()["atmos"]["geopotential_at_1000hpa"].values[-1]
+ z_sequential = (
+ make_driver(sequential).run()["atmos"]["geopotential_at_1000hpa"].values[-1]
+ )
+ # lagged: atmos runs 48-96h forced by the IC sst (2.0) -> z96 = 19.2
+ # sequential: forced by the fresh sst (2.042) -> z96 = 19.2336
+ assert np.allclose(z_lagged, 19.2, atol=1e-4)
+ assert np.allclose(z_sequential, 19.2336, atol=1e-4)
+ # difference = 8 steps * 0.1 * (2.042 - 2.0) = one coupling window's worth
+ assert np.allclose(z_sequential - z_lagged, 8 * 0.1 * 0.042, atol=1e-5)
+
+
+def test_gradient_flows_across_the_exchange():
+ gain_atmos = torch.tensor(1.0, requires_grad=True)
+ gain_ocean = torch.tensor(1.0, requires_grad=True)
+ driver = make_driver(gain_atmos=gain_atmos, gain_ocean=gain_ocean)
+ with torch.enable_grad():
+ states = driver.rollout(16) # full 96h
+ loss = states["atmos"]["geopotential_at_1000hpa"].data.sum()
+ loss.backward()
+ # atmos gain obviously in the graph; ocean gain reaches the loss only
+ # THROUGH the exchange: ocean sst -> connector regrid -> atmos injection
+ assert gain_atmos.grad is not None and gain_atmos.grad != 0
+ assert gain_ocean.grad is not None and gain_ocean.grad != 0
+
+
+def test_steps_iteration_and_probe():
+ driver = make_driver(stop="2024-01-03") # 48h
+ seen = []
+ for time, states in driver.steps():
+ seen.append(time)
+ assert "atmos" in states and "ocean" in states
+ assert len(seen) == 8
+ transfer = driver.probe("ocean->atmos")
+ assert "sea_surface_temperature" in transfer
+ assert transfer["sea_surface_temperature"].data.shape == ATMOS_GRID
+
+
+def test_to_xarray_output():
+ ds = make_driver().run()
+ z = ds["atmos"]["geopotential_at_1000hpa"]
+ assert z.dims == ("time", "lat", "lon")
+ assert z.shape == (17, 32, 64) # IC + 16 rings
+ assert np.allclose(z.values[0], 0.0)
+ assert np.allclose(z.values[-1], 19.2336, atol=1e-4)
+ sst = ds["ocean"]["sea_surface_temperature"]
+ assert sst.shape == (3, 16, 32) # IC + 2 rings
+ assert np.allclose(sst.values[-1], 2.180147, atol=1e-6)
+
+
+def test_missing_ic_raises():
+ driver = Driver(
+ {
+ "atmos": fake_atmos(),
+ "ocean": fake_ocean(),
+ "med": TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"]),
+ },
+ LAGGED_DSL,
+ Clock(T0, T96, "6h"),
+ )
+ with pytest.raises(CouplingError, match="initial condition"):
+ driver.initialize({"atmos": atmos_ic()})
+
+
+def test_run_before_initialize_raises():
+ driver = Driver(
+ {"atmos": fake_atmos()}, "@6h\n atmos\n@", Clock(T0, "2024-01-02", "6h")
+ )
+ with pytest.raises(CouplingError, match="initialize"):
+ driver.run()
+
+
+def test_unconsumed_export_warning():
+ driver = Driver(
+ {"atmos": fake_atmos()},
+ "@6h\n atmos\n@",
+ Clock(T0, "2024-01-02", "6h"),
+ allow_unfed_imports=True, # atmos's sst import is deliberately unfed
+ )
+ with capture_loguru() as messages:
+ driver.initialize({"atmos": atmos_ic()})
+ assert any("no connector consumes" in m for m in messages)
+
+
+# -- finding 2b: unfed imports must fail loudly at initialize -------------------
+FORGOT_OCEAN_TO_ATMOS_DSL = """
+@6h
+ atmos -> med
+ atmos
+@48h
+ med.compute
+ med -> ocean
+ ocean
+@
+"""
+
+
+def make_forgetful_driver(**kwargs):
+ from earth2studio.nvcoupler.mediator import TrailingAverageMediator
+
+ components = {
+ "atmos": fake_atmos(),
+ "ocean": fake_ocean(),
+ "med": TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"]),
+ }
+ return Driver(components, FORGOT_OCEAN_TO_ATMOS_DSL, Clock(T0, T96, "6h"), **kwargs)
+
+
+def test_forgotten_connector_raises_at_initialize():
+ driver = make_forgetful_driver()
+ with pytest.raises(
+ UnmatchedImportError, match=r"'atmos' imports 'sea_surface_temperature'"
+ ):
+ driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+
+
+def test_allow_unfed_imports_warns_and_runs():
+ driver = make_forgetful_driver(allow_unfed_imports=True)
+ with capture_loguru() as messages:
+ driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+ assert any(
+ "atmos" in m and "sea_surface_temperature" in m and "no connector" in m
+ for m in messages
+ )
+ driver.run() # runs on stale IC forcing, by explicit opt-in
+ assert driver.components["atmos"].run_count == 16
+
+
+# -- finding 2c: collected records must not pin autograd graphs ------------------
+def test_records_are_detached_during_gradient_rollout():
+ gain_atmos = torch.tensor(1.0, requires_grad=True)
+ gain_ocean = torch.tensor(1.0, requires_grad=True)
+ driver = make_driver(gain_atmos=gain_atmos, gain_ocean=gain_ocean)
+ with torch.enable_grad():
+ states = driver.rollout(16)
+ # exchange path stays attached...
+ loss = states["atmos"]["geopotential_at_1000hpa"].data.sum()
+ loss.backward()
+ assert gain_atmos.grad is not None and gain_ocean.grad is not None
+ # ...but records (off the exchange path) are detached clones
+ for records in driver._records.values():
+ for _, fields in records:
+ for field in fields.values():
+ assert not field.data.requires_grad
+ assert field.data.grad_fn is None
+
+
+# -- finding 2d: exhausted clock fails loudly; reset() enables a rerun -----------
+def make_atmos_only_driver():
+ driver = Driver(
+ {"atmos": fake_atmos()},
+ "@6h\n atmos\n@",
+ Clock(T0, "2024-01-02", "6h"),
+ allow_unfed_imports=True,
+ )
+ driver.initialize({"atmos": atmos_ic()})
+ return driver
+
+
+def test_second_run_raises_clock_exhausted():
+ driver = make_atmos_only_driver()
+ driver.run()
+ with pytest.raises(CouplingError, match="exhausted.*reset"):
+ driver.run()
+
+
+def test_steps_after_exhaustion_raises():
+ driver = make_atmos_only_driver()
+ for _ in driver.steps():
+ pass
+ with pytest.raises(CouplingError, match="exhausted.*reset"):
+ driver.steps()
+
+
+def test_rollout_exhaustion_raises_coupling_error_not_stopiteration():
+ driver = make_atmos_only_driver()
+ driver.run()
+ with pytest.raises(CouplingError, match="exhausted.*reset"):
+ driver.rollout(1)
+ # mid-iteration exhaustion: more steps requested than remain
+ driver.reset()
+ driver.initialize({"atmos": atmos_ic()})
+ with pytest.raises(CouplingError, match=r"rollout\(5\)"):
+ driver.rollout(5) # 24h clock has only 4 steps
+
+
+def test_reset_requires_reinitialize_then_reruns_identically():
+ driver = make_atmos_only_driver()
+ first = driver.run()
+ driver.reset()
+ assert all(not records for records in driver._records.values())
+ with pytest.raises(CouplingError, match="initialize"):
+ driver.run() # reset invalidates initialization
+ driver.initialize({"atmos": atmos_ic()})
+ second = driver.run()
+ z1 = first["atmos"]["geopotential_at_1000hpa"].values
+ z2 = second["atmos"]["geopotential_at_1000hpa"].values
+ assert np.array_equal(z1, z2)
+ assert z1.shape[0] == 5 # IC + 4 rings, not stale/duplicated
+
+
+# -- finding 2a: fields carrying their own size-1 time/batch coords ---------------
+def make_time_coord_component(n_time: int = 1):
+ """A component whose published fields keep size-1 (or larger) time and
+ batch dims — the DLESyM-split style of model coords."""
+
+ def step(x, coords):
+ return x + 1.0, coords
+
+ comp = CallableComponent(
+ "comp",
+ step,
+ timestep="6h",
+ exports=["geopotential_at_1000hpa"],
+ )
+ coords = OrderedDict(
+ {
+ "batch": np.array([0]),
+ "time": np.array([np.datetime64(T0, "ns")] * n_time),
+ "variable": np.array(["z1000"]),
+ "lat": np.linspace(90.0, -90.0, 8),
+ "lon": np.linspace(0.0, 360.0, 16, endpoint=False),
+ }
+ )
+ x = torch.zeros(1, n_time, 1, 8, 16)
+ return comp, (x, coords)
+
+
+def test_to_xarray_with_field_time_coord_uses_ring_times():
+ comp, ic = make_time_coord_component()
+ driver = Driver({"comp": comp}, "@6h\n comp\n@", Clock(T0, "2024-01-01T12", "6h"))
+ driver.initialize({"comp": ic})
+ ds = driver.run()
+ z = ds["comp"]["geopotential_at_1000hpa"]
+ # squeezed batch/time dims, ring-time axis prepended
+ assert z.dims == ("time", "lat", "lon")
+ assert z.shape == (3, 8, 16)
+ # RING times, not the stale IC time carried on every field
+ expected = np.datetime64(T0, "ns") + np.arange(3) * np.timedelta64(6, "h")
+ assert np.array_equal(z["time"].values, expected)
+ assert np.allclose(z.values, np.arange(3)[:, None, None])
+
+
+def test_non_size1_time_dim_raises():
+ comp, ic = make_time_coord_component(n_time=2)
+ driver = Driver({"comp": comp}, "@6h\n comp\n@", Clock(T0, "2024-01-01T12", "6h"))
+ driver.initialize({"comp": ic})
+ with pytest.raises(CouplingError, match="'time' dimension of size 2"):
+ driver.run()
diff --git a/test/nvcoupler/test_field.py b/test/nvcoupler/test_field.py
new file mode 100644
index 000000000..b35454a64
--- /dev/null
+++ b/test/nvcoupler/test_field.py
@@ -0,0 +1,137 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+from collections import OrderedDict
+
+import numpy as np
+import pytest
+import torch
+
+from earth2studio.nvcoupler import (
+ DEFAULT_DICTIONARY,
+ CouplingError,
+ Field,
+ State,
+ UnknownFieldError,
+)
+
+
+def _grid_coords(nlat=8, nlon=16):
+ return OrderedDict(
+ {
+ "lat": np.linspace(90, -90, nlat),
+ "lon": np.linspace(0, 360, nlon, endpoint=False),
+ }
+ )
+
+
+def _field(name="sea_surface_temperature", units="K", value=1.0, nlat=8, nlon=16):
+ return Field(
+ data=torch.full((nlat, nlon), value),
+ coords=_grid_coords(nlat, nlon),
+ standard_name=name,
+ units=units,
+ )
+
+
+def test_field_validation():
+ # dims mismatch
+ with pytest.raises(CouplingError):
+ Field(torch.zeros(4), _grid_coords(), "sea_surface_temperature", "K")
+ # variable dim forbidden
+ coords = OrderedDict({"variable": np.array(["sst"]), **_grid_coords()})
+ with pytest.raises(CouplingError):
+ Field(torch.zeros(1, 8, 16), coords, "sea_surface_temperature", "K")
+
+
+def test_field_clone_and_grid_signature():
+ f = _field()
+ g = f.clone()
+ g.data += 1.0
+ assert torch.all(f.data == 1.0) # clone is independent
+ assert f.grid_signature() == g.grid_signature()
+ assert f.grid_signature() != _field(nlat=4, nlon=8).grid_signature()
+
+
+def test_state_mapping_and_subset():
+ s = State("imports", [_field(), _field("air_temperature_2m", "K", 2.0)])
+ assert len(s) == 2
+ assert s["sea_surface_temperature"].data.mean() == 1.0
+ sub = s.subset(["air_temperature_2m"])
+ assert list(sub) == ["air_temperature_2m"]
+ with pytest.raises(KeyError) as err:
+ s["geopotential_at_500hpa"]
+ assert "imports" in str(err.value)
+ # key must equal standard_name
+ with pytest.raises(CouplingError):
+ s["wrong_key"] = _field()
+
+
+def test_as_tensor_from_tensor_roundtrip():
+ s = State(
+ "exports",
+ [
+ _field("sea_surface_temperature", "K", 1.0),
+ _field("air_temperature_2m", "K", 2.0),
+ ],
+ )
+ x, coords = s.as_tensor(["sea_surface_temperature", "air_temperature_2m"])
+ # variable inserted before spatial dims
+ assert list(coords) == ["variable", "lat", "lon"]
+ assert x.shape == (2, 8, 16)
+ assert torch.all(x[0] == 1.0) and torch.all(x[1] == 2.0)
+
+ s2 = State.from_tensor("roundtrip", x, coords, DEFAULT_DICTIONARY)
+ assert sorted(s2) == ["air_temperature_2m", "sea_surface_temperature"]
+ assert torch.equal(
+ s2["sea_surface_temperature"].data, s["sea_surface_temperature"].data
+ )
+ assert list(s2["air_temperature_2m"].coords) == ["lat", "lon"]
+
+
+def test_as_tensor_grid_mismatch_raises():
+ s = State("bad", [_field(), _field("air_temperature_2m", nlat=4, nlon=8)])
+ with pytest.raises(ValueError):
+ s.as_tensor(["sea_surface_temperature", "air_temperature_2m"])
+
+
+def test_from_tensor_alias_resolution_and_strict():
+ coords = OrderedDict(
+ {"variable": np.array(["sst", "mystery_var"]), **_grid_coords()}
+ )
+ x = torch.zeros(2, 8, 16)
+ with pytest.raises(UnknownFieldError):
+ State.from_tensor("s", x, coords, DEFAULT_DICTIONARY)
+ s = State.from_tensor("s", x, coords, DEFAULT_DICTIONARY, strict=False)
+ assert list(s) == ["sea_surface_temperature"] # alias resolved, unknown skipped
+
+
+def test_from_tensor_with_leading_dims():
+ coords = OrderedDict(
+ {
+ "time": np.array([np.datetime64("2024-01-01")]),
+ "variable": np.array(["z1000", "t2m"]),
+ **_grid_coords(),
+ }
+ )
+ x = torch.arange(2 * 8 * 16, dtype=torch.float32).reshape(1, 2, 8, 16)
+ s = State.from_tensor("s", x, coords, DEFAULT_DICTIONARY)
+ f = s["geopotential_at_1000hpa"]
+ assert list(f.coords) == ["time", "lat", "lon"]
+ assert f.data.shape == (1, 8, 16)
+ # stacking back inserts variable before lat (after time)
+ y, ycoords = s.as_tensor()
+ assert list(ycoords) == ["time", "variable", "lat", "lon"]
diff --git a/test/nvcoupler/test_io.py b/test/nvcoupler/test_io.py
new file mode 100644
index 000000000..5e30a2a84
--- /dev/null
+++ b/test/nvcoupler/test_io.py
@@ -0,0 +1,159 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Streaming IO tests: the coupled toy system writing into ZarrBackends."""
+
+import numpy as np
+import pytest
+
+from earth2studio.io import ZarrBackend
+from earth2studio.nvcoupler.clock import Clock
+from earth2studio.nvcoupler.driver import Driver
+from earth2studio.nvcoupler.errors import CouplingError
+from earth2studio.nvcoupler.mediator import TrailingAverageMediator
+from earth2studio.nvcoupler.testing import atmos_ic, fake_atmos, fake_ocean, ocean_ic
+
+T0 = "2024-01-01"
+T96 = "2024-01-05"
+
+LAGGED_DSL = """
+@6h
+ atmos -> med
+ ocean -> atmos
+ atmos
+@48h
+ med.compute
+ med -> ocean
+ ocean
+@
+"""
+
+
+def make_driver(io=None, collect=True):
+ components = {
+ "atmos": fake_atmos(gain=1.0),
+ "ocean": fake_ocean(gain=1.0),
+ "med": TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"]),
+ }
+ driver = Driver(
+ components, LAGGED_DSL, Clock(T0, T96, "6h"), collect=collect, io=io
+ )
+ driver.initialize({"atmos": atmos_ic(), "ocean": ocean_ic()})
+ return driver
+
+
+def test_zarr_streaming_matches_to_xarray():
+ io = {"atmos": ZarrBackend(), "ocean": ZarrBackend()}
+ driver = make_driver(io=io)
+ ds = driver.run()
+
+ z = io["atmos"]["geopotential_at_1000hpa"][:]
+ assert z.shape == (17, 32, 64) # IC + 16 rings
+ assert np.array_equal(z, ds["atmos"]["geopotential_at_1000hpa"].values)
+ assert np.allclose(z[0], 0.0)
+ assert np.allclose(z[-1], 19.2336, atol=1e-4)
+
+ sst = io["ocean"]["sea_surface_temperature"][:]
+ assert sst.shape == (3, 16, 32) # IC + 2 rings
+ assert np.array_equal(sst, ds["ocean"]["sea_surface_temperature"].values)
+ assert np.allclose(sst[-1], 2.180147, atol=1e-6)
+
+ # time coords are each component's ring times INCLUDING t0
+ atmos_times = io["atmos"]["time"][:].astype("datetime64[ns]")
+ expected = np.datetime64(T0, "ns") + np.arange(17) * np.timedelta64(6, "h")
+ assert np.array_equal(atmos_times, expected)
+ ocean_times = io["ocean"]["time"][:].astype("datetime64[ns]")
+ expected = np.datetime64(T0, "ns") + np.arange(3) * np.timedelta64(48, "h")
+ assert np.array_equal(ocean_times, expected)
+
+ # spatial coords round-trip
+ assert np.allclose(io["atmos"]["lat"][:], np.linspace(90.0, -90.0, 32))
+ assert np.allclose(io["ocean"]["lat"][:], np.linspace(90.0, -90.0, 16))
+
+
+def test_mediator_io_deferred_setup():
+ """A mediator exports nothing at t0, so its arrays are allocated at first
+ compute; the t0 row stays unwritten (zarr's default fill value)."""
+ io = {"med": ZarrBackend()}
+ driver = make_driver(io=io)
+ driver.run()
+ zm = io["med"]["geopotential_at_1000hpa_48h_mean"][:]
+ assert zm.shape == (3, 32, 64)
+ # unwritten t0 row must be NaN (poisonous), not zarr's 0.0 default —
+ # never-written data must not masquerade as physical values
+ assert np.all(np.isnan(zm[0]))
+ assert np.allclose(zm[1], 4.2, atol=1e-6)
+ assert np.allclose(zm[2], 13.8147, atol=1e-4)
+
+
+def test_io_independent_of_collect():
+ io = {"atmos": ZarrBackend()}
+ driver = make_driver(io=io, collect=False)
+ result = driver.run()
+ assert result == {} # nothing collected in memory
+ z = io["atmos"]["geopotential_at_1000hpa"][:]
+ assert z.shape == (17, 32, 64)
+ assert np.allclose(z[-1], 19.2336, atol=1e-4)
+
+
+def test_zarr_rows_land_under_ring_times_for_time_coord_fields():
+ """A component whose fields keep size-1 batch/time coords (DLESyM-split
+ style) must stream rows under the RING times — the field's own stale IC
+ time coord is squeezed out, never overwriting the ring time."""
+ from collections import OrderedDict
+
+ import torch
+
+ from earth2studio.nvcoupler.component import CallableComponent
+
+ def step(x, coords):
+ return x + 1.0, coords
+
+ comp = CallableComponent(
+ "comp", step, timestep="6h", exports=["geopotential_at_1000hpa"]
+ )
+ coords = OrderedDict(
+ {
+ "batch": np.array([0]),
+ "time": np.array([np.datetime64(T0, "ns")]),
+ "variable": np.array(["z1000"]),
+ "lat": np.linspace(90.0, -90.0, 8),
+ "lon": np.linspace(0.0, 360.0, 16, endpoint=False),
+ }
+ )
+ io = {"comp": ZarrBackend()}
+ driver = Driver(
+ {"comp": comp}, "@6h\n comp\n@", Clock(T0, "2024-01-01T12", "6h"), io=io
+ )
+ driver.initialize({"comp": (torch.zeros(1, 1, 1, 8, 16), coords)})
+ driver.run()
+
+ z = io["comp"]["geopotential_at_1000hpa"][:]
+ assert z.shape == (3, 8, 16) # (ring time, lat, lon) — batch/time squeezed
+ assert np.allclose(z, np.arange(3)[:, None, None])
+ times = io["comp"]["time"][:].astype("datetime64[ns]")
+ expected = np.datetime64(T0, "ns") + np.arange(3) * np.timedelta64(6, "h")
+ assert np.array_equal(times, expected)
+
+
+def test_unknown_io_key_raises():
+ with pytest.raises(CouplingError, match="not component names"):
+ Driver(
+ {"atmos": fake_atmos()},
+ "@6h\n atmos\n@",
+ Clock(T0, "2024-01-02", "6h"),
+ io={"atmso": ZarrBackend()},
+ )
diff --git a/test/nvcoupler/test_mediator.py b/test/nvcoupler/test_mediator.py
new file mode 100644
index 000000000..fb185a354
--- /dev/null
+++ b/test/nvcoupler/test_mediator.py
@@ -0,0 +1,215 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import numpy as np
+import pytest
+import torch
+
+from earth2studio.nvcoupler.dictionary import (
+ DEFAULT_DICTIONARY,
+ CellMethod,
+ FieldDictionary,
+ FieldEntry,
+)
+from earth2studio.nvcoupler.errors import CouplingError
+from earth2studio.nvcoupler.field import Field
+from earth2studio.nvcoupler.mediator import (
+ AccumulationMediator,
+ TrailingAverageMediator,
+)
+from earth2studio.nvcoupler.testing import grid_coords
+
+T0 = np.datetime64("2024-01-01")
+H = np.timedelta64(1, "h")
+
+
+def _z1000(value, hours):
+ return Field(
+ torch.full((4, 8), float(value)),
+ grid_coords(4, 8),
+ "geopotential_at_1000hpa",
+ "m2 s-2",
+ valid_time=T0 + hours * H,
+ )
+
+
+def test_trailing_average_over_window():
+ med = TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"])
+ assert med.timestep == np.timedelta64(48, "h")
+ assert med.import_names == ["geopotential_at_1000hpa"]
+ # 8 x 6h samples with values 1..8 -> mean 4.5
+ for i in range(1, 9):
+ med.import_state.add(_z1000(i, 6 * i))
+ med.run(T0 + np.timedelta64(48, "h"))
+ out = med.export_state["geopotential_at_1000hpa_48h_mean"]
+ assert torch.allclose(out.data, torch.full((4, 8), 4.5))
+ assert out.valid_time == T0 + np.timedelta64(48, "h")
+ assert med.samples_last_window["geopotential_at_1000hpa_48h_mean"] == 8
+ # accumulator reset: next window with values 10, 20 -> mean 15
+ med.import_state.add(_z1000(10, 54))
+ med.import_state.add(_z1000(20, 60))
+ med.run(T0 + np.timedelta64(96, "h"))
+ assert torch.allclose(
+ med.export_state["geopotential_at_1000hpa_48h_mean"].data,
+ torch.full((4, 8), 15.0),
+ )
+
+
+def test_duplicate_valid_time_ignored():
+ med = TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"])
+ f = _z1000(4, 6)
+ med.import_state.add(f)
+ med.import_state.add(f) # same valid_time -> not double counted
+ med.run(T0 + np.timedelta64(48, "h"))
+ assert torch.allclose(
+ med.export_state["geopotential_at_1000hpa_48h_mean"].data,
+ torch.full((4, 8), 4.0),
+ )
+
+
+def test_sum_and_max_reductions():
+ d = FieldDictionary(DEFAULT_DICTIONARY)
+ med_sum = AccumulationMediator(
+ "psum", ["total_precipitation_48h_sum"], dictionary=d
+ )
+ for i, v in enumerate([1.0, 2.0, 3.0]):
+ med_sum.import_state.add(
+ Field(
+ torch.full((4, 8), v),
+ grid_coords(4, 8),
+ "total_precipitation_6h",
+ "kg m-2",
+ valid_time=T0 + i * H,
+ )
+ )
+ med_sum.run(T0 + np.timedelta64(48, "h"))
+ assert torch.allclose(
+ med_sum.export_state["total_precipitation_48h_sum"].data,
+ torch.full((4, 8), 6.0),
+ )
+
+ med_max = AccumulationMediator("tmax", ["air_temperature_2m_24h_max"], dictionary=d)
+ for i, v in enumerate([280.0, 295.0, 290.0]):
+ med_max.import_state.add(
+ Field(
+ torch.full((4, 8), v),
+ grid_coords(4, 8),
+ "air_temperature_2m",
+ "K",
+ valid_time=T0 + i * H,
+ )
+ )
+ med_max.run(T0 + np.timedelta64(24, "h"))
+ assert torch.allclose(
+ med_max.export_state["air_temperature_2m_24h_max"].data,
+ torch.full((4, 8), 295.0),
+ )
+
+
+def test_two_reductions_of_same_base_field():
+ """Two derived fields of one base (24h max AND 24h mean of t2m) must
+ each accumulate every delivery — regression for the base->derived map
+ silently keeping only the last derived field."""
+ d = FieldDictionary(DEFAULT_DICTIONARY)
+ d.register(
+ FieldEntry(
+ "air_temperature_2m_24h_mean",
+ "K",
+ "24 h mean 2 m temperature",
+ frozenset(),
+ CellMethod("air_temperature_2m", "mean", np.timedelta64(24, "h")),
+ )
+ )
+ med = AccumulationMediator(
+ "med",
+ ["air_temperature_2m_24h_max", "air_temperature_2m_24h_mean"],
+ dictionary=d,
+ )
+ # the shared base import is deduped, not advertised twice
+ assert med.import_names == ["air_temperature_2m"]
+ values = [280.0, 295.0, 290.0]
+ for i, v in enumerate(values):
+ med.import_state.add(
+ Field(
+ torch.full((4, 8), v),
+ grid_coords(4, 8),
+ "air_temperature_2m",
+ "K",
+ valid_time=T0 + i * H,
+ )
+ )
+ med.run(T0 + np.timedelta64(24, "h"))
+ assert torch.allclose(
+ med.export_state["air_temperature_2m_24h_max"].data,
+ torch.full((4, 8), 295.0),
+ )
+ assert torch.allclose(
+ med.export_state["air_temperature_2m_24h_mean"].data,
+ torch.full((4, 8), sum(values) / len(values)),
+ )
+ assert med.samples_last_window["air_temperature_2m_24h_max"] == 3
+ assert med.samples_last_window["air_temperature_2m_24h_mean"] == 3
+
+
+def test_mediator_requires_no_ic():
+ med = TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"])
+ assert med.requires_ic is False
+ med.initialize() # no-arg initialize is safe
+
+
+def test_compute_without_samples_raises():
+ med = TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"])
+ with pytest.raises(CouplingError, match="no samples"):
+ med.run(T0 + np.timedelta64(48, "h"))
+
+
+def test_field_without_cell_method_rejected():
+ with pytest.raises(CouplingError, match="cell_method"):
+ AccumulationMediator("med", ["sea_surface_temperature"])
+
+
+def test_trailing_average_rejects_non_mean():
+ with pytest.raises(CouplingError, match="not mean"):
+ TrailingAverageMediator("med", ["total_precipitation_48h_sum"])
+
+
+def test_mixed_windows_need_explicit_window():
+ d = FieldDictionary(DEFAULT_DICTIONARY)
+ with pytest.raises(CouplingError, match="differing"):
+ AccumulationMediator(
+ "med",
+ ["geopotential_at_1000hpa_48h_mean", "air_temperature_2m_24h_max"],
+ dictionary=d,
+ )
+
+
+def test_gradient_flows_through_mean():
+ med = TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"])
+ xs = [torch.full((4, 8), float(i), requires_grad=True) for i in (1, 3)]
+ for i, x in enumerate(xs):
+ med.import_state.add(
+ Field(
+ x,
+ grid_coords(4, 8),
+ "geopotential_at_1000hpa",
+ "m2 s-2",
+ valid_time=T0 + i * H,
+ )
+ )
+ med.run(T0 + np.timedelta64(48, "h"))
+ med.export_state["geopotential_at_1000hpa_48h_mean"].data.sum().backward()
+ for x in xs:
+ assert x.grad is not None and torch.allclose(x.grad, torch.full((4, 8), 0.5))
diff --git a/test/nvcoupler/test_points.py b/test/nvcoupler/test_points.py
new file mode 100644
index 000000000..5cab94382
--- /dev/null
+++ b/test/nvcoupler/test_points.py
@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import numpy as np
+import pytest
+
+from earth2studio.nvcoupler.errors import CouplingError
+from earth2studio.nvcoupler.points import PointSet
+
+
+def test_default_labels_are_integer_index():
+ p = PointSet(lat=np.array([1.0, 2.0, 3.0]), lon=np.array([4.0, 5.0, 6.0]))
+ assert len(p) == 3
+ assert np.array_equal(p.labels(), np.arange(3))
+ assert list(p.grid_coords()) == ["point"]
+ assert np.array_equal(p.grid_coords()["point"], np.arange(3))
+
+
+def test_names_become_labels():
+ p = PointSet(lat=np.array([1.0, 2.0]), lon=np.array([3.0, 4.0]), names=("a", "b"))
+ assert np.array_equal(p.labels(), np.array(["a", "b"]))
+
+
+def test_mismatched_lat_lon_length_raises():
+ with pytest.raises(CouplingError, match="same length"):
+ PointSet(lat=np.array([1.0, 2.0]), lon=np.array([3.0]))
+
+
+def test_non_1d_raises():
+ with pytest.raises(CouplingError, match="1-D"):
+ PointSet(lat=np.array([[1.0, 2.0]]), lon=np.array([[3.0, 4.0]]))
+
+
+def test_empty_raises():
+ with pytest.raises(CouplingError, match="at least one point"):
+ PointSet(lat=np.array([]), lon=np.array([]))
+
+
+def test_names_length_mismatch_raises():
+ with pytest.raises(CouplingError, match="names has"):
+ PointSet(lat=np.array([1.0, 2.0]), lon=np.array([3.0, 4.0]), names=("a",))
+
+
+def test_signature_stable_and_distinguishes_locations():
+ p1 = PointSet(lat=np.array([1.0]), lon=np.array([2.0]))
+ p2 = PointSet(lat=np.array([1.0]), lon=np.array([2.0]))
+ p3 = PointSet(lat=np.array([1.0]), lon=np.array([3.0]))
+ assert p1.signature() == p2.signature()
+ assert p1.signature() != p3.signature()
diff --git a/test/nvcoupler/test_prognostic_real.py b/test/nvcoupler/test_prognostic_real.py
new file mode 100644
index 000000000..71c03f435
--- /dev/null
+++ b/test/nvcoupler/test_prognostic_real.py
@@ -0,0 +1,138 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Smoke tests: PrognosticComponent against a real earth2studio prognostic.
+
+Uses the Persistence model (models/px/persistence.py) — a real
+PrognosticModel implementing the full protocol (input_coords /
+output_coords / batch_func-decorated __call__) that returns its input
+unchanged and needs no weights. This exercises the seam between
+nvcoupler's component phases and the earth2studio model conventions:
+'batch' dims of np.empty(0) in input_coords, lead_time windows, and
+batch_func's compress/decompress behavior.
+"""
+
+from collections import OrderedDict
+
+import numpy as np
+import pytest
+import torch
+
+from earth2studio.models.px import Persistence
+from earth2studio.nvcoupler.clock import Clock
+from earth2studio.nvcoupler.component import PrognosticComponent
+from earth2studio.nvcoupler.testing import grid_coords
+
+T0 = np.datetime64("2024-01-01")
+VARIABLES = ["t2m", "z1000"]
+NLAT, NLON = 8, 16
+
+
+def _make_component() -> PrognosticComponent:
+ model = Persistence(VARIABLES, grid_coords(NLAT, NLON))
+ return PrognosticComponent("persist", model)
+
+
+def _make_ic(model: Persistence) -> tuple[torch.Tensor, OrderedDict]:
+ """Build an IC tensor matching model.input_coords() with the empty
+ 'batch' dim scrubbed (batch_func re-inserts it on call, the pattern
+ run.py uses when preparing model inputs)."""
+ coords = model.input_coords()
+ del coords["batch"] # np.empty(0) placeholder, not a real axis
+ shape = tuple(len(v) for v in coords.values())
+ torch.manual_seed(0)
+ x = torch.randn(shape)
+ return x, coords
+
+
+def test_timestep_inferred_from_output_coords():
+ comp = _make_component()
+ assert comp.timestep == np.timedelta64(6, "h")
+ # exports resolved from raw variable names to dictionary standard names
+ assert comp.export_names == [
+ "air_temperature_2m",
+ "geopotential_at_1000hpa",
+ ]
+
+
+def test_realize_initialize_run_loop_persists_values():
+ comp = _make_component()
+ clock = Clock(T0, "2024-01-02", "6h")
+ comp.realize(clock)
+
+ x0, coords0 = _make_ic(comp.model)
+ comp.initialize(x0, coords0)
+
+ # initialize seeds the exports at t0 for lagged coupling
+ for std in comp.export_names:
+ assert std in comp.export_state
+ assert comp.export_state[std].valid_time == clock.start
+ assert list(comp.export_state[std].coords) == ["lat", "lon"]
+
+ var_axis = list(coords0).index("variable")
+ # exchange-shaped IC slices: the singleton lead_time dim dropped, since
+ # published Fields must carry plain (lat, lon) coords
+ ic_slices = {
+ "air_temperature_2m": x0.select(var_axis, 0)[0],
+ "geopotential_at_1000hpa": x0.select(var_axis, 1)[0],
+ }
+
+ for i in range(1, 4): # three 6 h steps
+ time = T0 + i * np.timedelta64(6, "h")
+ comp.run(time)
+ for std, ic in ic_slices.items():
+ field = comp.export_state[std]
+ assert field.standard_name == std
+ assert field.valid_time == time
+ assert field.source == "persist"
+ # exports are exchange-shaped: no batch/time/lead_time singletons
+ assert list(field.coords) == ["lat", "lon"]
+ # persistence: values are the IC, unchanged, on the same grid
+ assert torch.equal(field.data, ic)
+ assert np.array_equal(field.coords["lat"], coords0["lat"])
+ assert np.array_equal(field.coords["lon"], coords0["lon"])
+ assert comp.run_count == 3
+
+
+def test_internal_state_window_stays_model_shaped():
+ comp = _make_component()
+ clock = Clock(T0, "2024-01-02", "6h")
+ comp.realize(clock)
+ x0, coords0 = _make_ic(comp.model)
+ comp.initialize(x0, coords0)
+ comp.run(T0 + np.timedelta64(6, "h"))
+ x, coords = comp.state
+ # next_input rewound lead_time to the model's input window
+ assert np.array_equal(coords["lead_time"], comp.model.input_coords()["lead_time"])
+ assert list(coords) == ["lead_time", "variable", "lat", "lon"]
+ assert x.shape == (1, len(VARIABLES), NLAT, NLON)
+
+
+def test_multi_history_needs_next_input_hook():
+ """A history-2 Persistence outputs 1 lead time but takes 2 as input:
+ the default next_input cannot manage the window and must say so."""
+ from earth2studio.nvcoupler.errors import CouplingError
+
+ model = Persistence(VARIABLES, grid_coords(NLAT, NLON), history=2)
+ comp = PrognosticComponent("persist2", model)
+ clock = Clock(T0, "2024-01-02", "6h")
+ comp.realize(clock)
+ coords = model.input_coords()
+ del coords["batch"]
+ x = torch.randn(tuple(len(v) for v in coords.values()))
+ comp.initialize(x, coords)
+ with pytest.raises(CouplingError, match="next_input"):
+ comp.run(T0 + np.timedelta64(6, "h"))
diff --git a/test/nvcoupler/test_pull.py b/test/nvcoupler/test_pull.py
new file mode 100644
index 000000000..2df40cfb8
--- /dev/null
+++ b/test/nvcoupler/test_pull.py
@@ -0,0 +1,255 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Pull-pattern coupling tests (PullAdapter / StateDataSource).
+
+MockPullModel is protocol-faithful to StormCast's coupling surface: it holds
+a settable ``conditioning_data_source`` and, inside its own ``__call__``,
+fetches its conditioning variables through the REAL
+``earth2studio.data.utils.fetch_data`` — the same code path the production
+model runs. If the shim satisfies fetch_data, it satisfies StormCast's
+mechanics (the model-weights physics is exactly what these tests cannot
+cover, and says so).
+"""
+
+from collections import OrderedDict
+
+import numpy as np
+import pytest
+import torch
+
+from earth2studio.data.utils import fetch_data
+from earth2studio.nvcoupler.clock import Clock
+from earth2studio.nvcoupler.component import CallableComponent, PrognosticComponent
+from earth2studio.nvcoupler.connector import Connector
+from earth2studio.nvcoupler.driver import Driver
+from earth2studio.nvcoupler.errors import CouplingError
+from earth2studio.nvcoupler.field import Field, State
+from earth2studio.nvcoupler.pull import PullAdapter, StateDataSource
+from earth2studio.nvcoupler.testing import grid_coords
+
+T0 = np.datetime64("2024-01-01")
+GRID = (8, 16)
+
+
+class MockPullModel:
+ """StormCast-shaped mock: pulls u10m/t2m via fetch_data inside __call__,
+ then steps its single state variable by the conditioning mean:
+ refc <- refc + 1 + mean(u10m) + 0.1 * mean(t2m)
+ Deterministic and hand-computable."""
+
+ conditioning_variables = np.array(["u10m", "t2m"])
+
+ def __init__(self):
+ self.conditioning_data_source = None
+ self.pull_log: list[np.ndarray] = []
+
+ def input_coords(self):
+ return OrderedDict(
+ {
+ "time": np.empty(0),
+ "lead_time": np.array([np.timedelta64(0, "h")]),
+ "variable": np.array(["refc"]),
+ **grid_coords(*GRID),
+ }
+ )
+
+ def output_coords(self, input_coords):
+ out = OrderedDict({k: v.copy() for k, v in input_coords.items()})
+ out["lead_time"] = input_coords["lead_time"] + np.timedelta64(1, "h")
+ return out
+
+ def __call__(self, x, coords):
+ if self.conditioning_data_source is None:
+ raise RuntimeError("conditioning_data_source not set")
+ # the REAL fetch path StormCast uses
+ cond, _ = fetch_data(
+ self.conditioning_data_source,
+ time=np.atleast_1d(coords["time"]),
+ variable=self.conditioning_variables,
+ )
+ self.pull_log.append(cond.numpy().copy())
+ u_mean = cond[0, 0, 0].mean()
+ t_mean = cond[0, 0, 1].mean()
+ return x + 1.0 + u_mean + 0.1 * t_mean, self.output_coords(coords)
+
+ def to(self, device):
+ return self
+
+
+def _field(name, value, valid=T0):
+ return Field(
+ torch.full(GRID, float(value)),
+ grid_coords(*GRID),
+ name,
+ "m s-1" if "wind" in name else "K",
+ valid_time=valid,
+ )
+
+
+def _imports(u=2.0, t=280.0):
+ return State(
+ "imports",
+ [_field("eastward_wind_10m", u), _field("air_temperature_2m", t)],
+ )
+
+
+class TestStateDataSource:
+ def test_serves_import_state_via_real_fetch_data(self):
+ src = StateDataSource(
+ _imports(u=3.0, t=290.0),
+ raw_to_std={"u10m": "eastward_wind_10m", "t2m": "air_temperature_2m"},
+ )
+ x, coords = fetch_data(
+ src, time=np.array([T0]), variable=np.array(["u10m", "t2m"])
+ )
+ assert x.shape == (1, 1, 2, *GRID)
+ assert torch.all(x[0, 0, 0] == 3.0) and torch.all(x[0, 0, 1] == 290.0)
+ assert list(coords["variable"]) == ["u10m", "t2m"]
+
+ def test_unknown_variable_actionable_error(self):
+ src = StateDataSource(_imports(), raw_to_std={})
+ with pytest.raises(CouplingError, match="wire a connector"):
+ src(np.array([T0]), np.array(["msl"]))
+
+ def test_strict_time_mismatch_raises(self):
+ src = StateDataSource(
+ _imports(),
+ raw_to_std={"u10m": "eastward_wind_10m", "t2m": "air_temperature_2m"},
+ strict_time=True,
+ )
+ with pytest.raises(CouplingError, match="run-sequence ordering"):
+ src(np.array([T0 + np.timedelta64(6, "h")]), np.array(["u10m"]))
+
+
+class TestPullAdapter:
+ def test_redirects_model_pull_to_imports(self):
+ model = MockPullModel()
+ comp = PrognosticComponent(
+ "stormcast",
+ model,
+ imports=["eastward_wind_10m", "air_temperature_2m"],
+ exports=["radar_reflectivity"],
+ import_adapter=PullAdapter(),
+ variable_aliases={"refc": "radar_reflectivity"},
+ dictionary=_dict_with_refc(),
+ )
+ clock = Clock(T0, "2024-01-01T04:00", "1h")
+ comp.realize(clock)
+ ic = model.input_coords()
+ ic["time"] = np.array([T0])
+ comp.initialize(torch.zeros(1, 1, 1, *GRID), ic)
+ comp.import_state.add(_field("eastward_wind_10m", 2.0))
+ comp.import_state.add(_field("air_temperature_2m", 280.0))
+ comp.run(clock.advance())
+ # refc = 0 + 1 + 2 + 28 = 31, pulled through the model's own fetch
+ got = comp.export_state["radar_reflectivity"]
+ assert torch.allclose(got.data, torch.full(GRID, 31.0))
+ assert len(model.pull_log) == 1
+
+ def test_missing_attribute_actionable_error(self):
+ class NoPull:
+ def input_coords(self):
+ return MockPullModel().input_coords()
+
+ def output_coords(self, c):
+ return MockPullModel().output_coords(c)
+
+ comp = PrognosticComponent(
+ "x",
+ NoPull(),
+ imports=["eastward_wind_10m"],
+ exports=["radar_reflectivity"],
+ import_adapter=PullAdapter(),
+ variable_aliases={"refc": "radar_reflectivity"},
+ dictionary=_dict_with_refc(),
+ )
+ comp.realize(Clock(T0, "2024-01-01T02:00", "1h"))
+ ic = comp.model.input_coords()
+ ic["time"] = np.array([T0])
+ comp.initialize(torch.zeros(1, 1, 1, *GRID), ic)
+ comp.import_state.add(_field("eastward_wind_10m", 1.0))
+ with pytest.raises(CouplingError, match="ConditioningKwargAdapter"):
+ comp.run(np.datetime64("2024-01-01T01:00"))
+
+
+def _dict_with_refc():
+ from earth2studio.nvcoupler.dictionary import (
+ DEFAULT_DICTIONARY,
+ FieldDictionary,
+ FieldEntry,
+ )
+
+ d = FieldDictionary(DEFAULT_DICTIONARY)
+ d.register(FieldEntry("radar_reflectivity", "dBZ", aliases=frozenset({"refc"})))
+ return d
+
+
+class TestCoupledPullWorkflow:
+ """The Jussi-workflow shape end-to-end: a 'global' component's exports
+ flow per-step into a pull-pattern regional model — no staging, no
+ InferenceOutputSource, sequential exchange in one slot."""
+
+ def test_per_step_conditioning_updates(self):
+ def global_step(x, coords):
+ # u10m grows 1 m/s per step; t2m constant
+ return torch.stack([x[0] + 1.0, x[1]]), coords
+
+ glob = CallableComponent(
+ "global",
+ global_step,
+ timestep="1h",
+ exports=["eastward_wind_10m", "air_temperature_2m"],
+ )
+ model = MockPullModel()
+ stormcast = PrognosticComponent(
+ "stormcast",
+ model,
+ imports=["eastward_wind_10m", "air_temperature_2m"],
+ exports=["radar_reflectivity"],
+ import_adapter=PullAdapter(),
+ variable_aliases={"refc": "radar_reflectivity"},
+ dictionary=_dict_with_refc(),
+ )
+ driver = Driver(
+ {"global": glob, "stormcast": stormcast},
+ sequence="""
+ @1h
+ global
+ global -> stormcast
+ stormcast
+ @
+ """,
+ clock=Clock(T0, "2024-01-01T03:00", "1h"),
+ connectors=[Connector(glob, stormcast)],
+ )
+ ic_glob = (
+ torch.stack([torch.full(GRID, 2.0), torch.full(GRID, 280.0)]),
+ OrderedDict({"variable": np.array(["u10m", "t2m"]), **grid_coords(*GRID)}),
+ )
+ ic_sc = MockPullModel().input_coords()
+ ic_sc["time"] = np.array([T0])
+ driver.initialize(
+ {"global": ic_glob, "stormcast": (torch.zeros(1, 1, 1, *GRID), ic_sc)}
+ )
+ driver.run()
+ # sequential coupling: stormcast at hour k pulls u = 2 + k (fresh)
+ # refc increments: k=1: 1+3+28=32; k=2: 1+4+28=33; k=3: 1+5+28=34
+ got = stormcast.export_state["radar_reflectivity"]
+ assert torch.allclose(got.data, torch.full(GRID, 32.0 + 33.0 + 34.0))
+ # three pulls, each seeing the CURRENT global state (3, 4, 5)
+ pulled_u = [p[0, 0, 0].mean() for p in model.pull_log]
+ assert pulled_u == [3.0, 4.0, 5.0]
diff --git a/test/nvcoupler/test_seams.py b/test/nvcoupler/test_seams.py
new file mode 100644
index 000000000..9a3064997
--- /dev/null
+++ b/test/nvcoupler/test_seams.py
@@ -0,0 +1,255 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Cross-component seam tests: exports must be exchange-shaped.
+
+PrognosticComponent wraps models whose state carries singleton batch /
+time / lead_time dims. Its published Fields must nevertheless be plain
+spatial (lat, lon) tensors, or they break every consumer downstream:
+VariableOverwriteAdapter's slice broadcast, State.as_tensor stacking with
+fields from other components, and DiagnosticComponent's input conformance.
+These tests exercise those three seams end to end.
+"""
+
+from collections import OrderedDict
+
+import numpy as np
+import torch
+import xarray as xr
+
+from earth2studio.nvcoupler.clock import Clock
+from earth2studio.nvcoupler.component import (
+ CallableComponent,
+ DataComponent,
+ DiagnosticComponent,
+ PrognosticComponent,
+)
+from earth2studio.nvcoupler.connector import Connector
+from earth2studio.nvcoupler.driver import Driver
+from earth2studio.nvcoupler.testing import grid_coords
+
+T0 = np.datetime64("2024-01-01")
+NLAT, NLON = 8, 16
+GRID = (NLAT, NLON)
+
+
+class BatchTimePrognostic:
+ """MockPrognostic with explicit singleton batch and time dims.
+
+ Mirrors real earth2studio prognostics whose state tensors are
+ (batch, time, lead_time, variable, lat, lon); +1.0 per 6 h step.
+ """
+
+ def __init__(self):
+ self._in = OrderedDict(
+ {
+ "batch": np.array([0]),
+ "time": np.array([T0]),
+ "lead_time": np.array([np.timedelta64(0, "h")]),
+ "variable": np.array(["z1000", "sst"]),
+ **grid_coords(NLAT, NLON),
+ }
+ )
+
+ def input_coords(self):
+ return OrderedDict({k: v.copy() for k, v in self._in.items()})
+
+ def output_coords(self, input_coords):
+ out = OrderedDict({k: v.copy() for k, v in input_coords.items()})
+ out["lead_time"] = input_coords["lead_time"] + np.timedelta64(6, "h")
+ return out
+
+ def __call__(self, x, coords):
+ return x + 1.0, self.output_coords(coords)
+
+ def to(self, device):
+ return self
+
+
+def _prognostic_ic(z0=0.0, sst0=3.0):
+ """IC tensor matching BatchTimePrognostic.input_coords()."""
+ coords = BatchTimePrognostic().input_coords()
+ x = torch.empty(1, 1, 1, 2, NLAT, NLON)
+ x[..., 0, :, :] = z0
+ x[..., 1, :, :] = sst0
+ return x, coords
+
+
+def _sink_atmos():
+ """CallableComponent on the same grid: z += 1 + 0.1 * sst."""
+
+ def step(x, coords):
+ z1000, sst = x[0], x[1]
+ return torch.stack([z1000 + 1.0 + 0.1 * sst, sst]), coords
+
+ return CallableComponent(
+ "sink",
+ step,
+ timestep="6h",
+ imports=["sea_surface_temperature"],
+ exports=["geopotential_at_1000hpa"],
+ )
+
+
+def _sink_ic(z0=0.0, sst0=0.0):
+ coords = OrderedDict(
+ {"variable": np.array(["z1000", "sst"]), **grid_coords(NLAT, NLON)}
+ )
+ x = torch.stack([torch.full(GRID, z0), torch.full(GRID, sst0)])
+ return x, coords
+
+
+class HalfDiagnostic:
+ """z500 = 0.5 * z1000 on the seam grid (models/dx/base.py interface)."""
+
+ def __init__(self):
+ self.grid = grid_coords(NLAT, NLON)
+
+ def input_coords(self):
+ return OrderedDict(
+ {
+ "batch": np.empty(0),
+ "variable": np.array(["z1000"]),
+ "lat": self.grid["lat"],
+ "lon": self.grid["lon"],
+ }
+ )
+
+ def output_coords(self, input_coords):
+ out = OrderedDict(input_coords)
+ out["variable"] = np.array(["z500"])
+ return out
+
+ def __call__(self, x, coords):
+ return 0.5 * x, self.output_coords(coords)
+
+ def to(self, device):
+ return self
+
+
+class ConstantDataSource:
+ """In-memory DataSource: constant value per variable on the seam grid."""
+
+ def __init__(self, values):
+ self.values = values
+ self.grid = grid_coords(NLAT, NLON)
+
+ def __call__(self, time, variable) -> xr.DataArray:
+ time = np.atleast_1d(np.asarray(time, dtype="datetime64[ns]"))
+ variable = np.atleast_1d(np.asarray(variable))
+ lat, lon = self.grid["lat"], self.grid["lon"]
+ data = np.empty((len(time), len(variable), len(lat), len(lon)))
+ for j, v in enumerate(variable):
+ data[:, j] = self.values[str(v)]
+ return xr.DataArray(
+ data,
+ dims=["time", "variable", "lat", "lon"],
+ coords={"time": time, "variable": variable, "lat": lat, "lon": lon},
+ )
+
+
+def test_prognostic_export_feeds_callable_component():
+ """(a) PrognosticComponent -> Connector -> CallableComponent import."""
+ clock = Clock(T0, "2024-01-02", "6h")
+ prog = PrognosticComponent("prog", BatchTimePrognostic())
+ sink = _sink_atmos()
+ prog.realize(clock)
+ sink.realize(clock)
+ prog.initialize(*_prognostic_ic(z0=0.0, sst0=3.0))
+ sink.initialize(*_sink_ic())
+
+ # exports must be exchange-shaped: plain (lat, lon), no batch/time/lead
+ sst0 = prog.export_state["sea_surface_temperature"]
+ assert list(sst0.coords) == ["lat", "lon"]
+ assert sst0.data.shape == GRID
+
+ conn = Connector(prog, sink)
+ t1 = clock.advance()
+ prog.run(t1) # sst: 3 -> 4
+ conn.execute(t1)
+ sink.run(t1) # z = 0 + 1 + 0.1 * 4 = 1.4
+
+ z = sink.export_state["geopotential_at_1000hpa"]
+ assert torch.allclose(z.data, torch.full(GRID, 1.4))
+
+
+def test_data_and_prognostic_exports_stack_in_one_import_state():
+ """(b) DataComponent + PrognosticComponent exports into one component's
+ imports, stacked via State.as_tensor without cat_coords errors."""
+ clock = Clock(T0, "2024-01-02", "6h")
+ prog = PrognosticComponent("prog", BatchTimePrognostic())
+ data = DataComponent(
+ "data",
+ source=ConstantDataSource({"t2m": 280.0}),
+ exports=["air_temperature_2m"],
+ timestep="6h",
+ )
+
+ def step(x, coords):
+ return x, coords
+
+ dst = CallableComponent(
+ "dst",
+ step,
+ timestep="6h",
+ imports=["geopotential_at_1000hpa", "air_temperature_2m"],
+ exports=[],
+ )
+ for comp in (prog, data, dst):
+ comp.realize(clock)
+ prog.initialize(*_prognostic_ic(z0=0.0, sst0=3.0))
+ data.initialize()
+ dst.initialize(
+ torch.zeros(1, NLAT, NLON),
+ OrderedDict({"variable": np.array(["z1000"]), **grid_coords(NLAT, NLON)}),
+ )
+
+ t1 = clock.advance()
+ prog.run(t1) # z1000: 0 -> 1
+ data.run(t1)
+ Connector(prog, dst, fields=["geopotential_at_1000hpa"]).execute(t1)
+ Connector(data, dst, fields=["air_temperature_2m"]).execute(t1)
+
+ names = ["geopotential_at_1000hpa", "air_temperature_2m"]
+ stacked, coords = dst.import_state.as_tensor(names)
+ assert list(coords["variable"]) == names
+ assert stacked.shape == (2, NLAT, NLON)
+ assert torch.allclose(stacked[0], torch.full(GRID, 1.0))
+ assert torch.allclose(stacked[1], torch.full(GRID, 280.0))
+
+
+def test_diagnostic_consumes_prognostic_export_in_driver():
+ """(c) DiagnosticComponent consuming a PrognosticComponent export."""
+ components = {
+ "prog": PrognosticComponent("prog", BatchTimePrognostic()),
+ "diag": DiagnosticComponent("diag", HalfDiagnostic(), timestep="6h"),
+ }
+ dsl = """
+@6h
+ prog
+ prog -> diag
+ diag
+@
+"""
+ driver = Driver(components, dsl, Clock(T0, "2024-01-02", "6h"))
+ driver.initialize({"prog": _prognostic_ic(z0=0.0, sst0=0.0), "diag": (None, None)})
+ ds = driver.run()
+
+ assert driver.components["diag"].run_count == 4
+ z500 = ds["diag"]["geopotential_at_500hpa"]
+ assert z500.dims == ("time", "lat", "lon")
+ # z1000_n = n (starts 0, +1 each step); diagnostic halves it
+ assert np.allclose(z500.values, 0.5 * np.arange(1, 5)[:, None, None], atol=1e-6)
diff --git a/test/nvcoupler/test_sequence.py b/test/nvcoupler/test_sequence.py
new file mode 100644
index 000000000..7ce43bf13
--- /dev/null
+++ b/test/nvcoupler/test_sequence.py
@@ -0,0 +1,247 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import numpy as np
+import pytest
+
+from earth2studio.nvcoupler.errors import CadenceError, SequenceError
+from earth2studio.nvcoupler.mediator import TrailingAverageMediator
+from earth2studio.nvcoupler.sequence import (
+ ConnectAction,
+ MediateAction,
+ RunAction,
+ RunSequence,
+ Slot,
+ derive_sequence,
+ parse_run_sequence,
+)
+from earth2studio.nvcoupler.testing import fake_atmos, fake_ocean
+
+DSL = """
+@6h
+ atmos -> med # accumulate into mediator
+ ocean -> atmos # lagged coupling
+ atmos
+@48h
+ med.compute
+ med -> ocean
+ ocean
+@
+"""
+
+
+def _components():
+ return {
+ "atmos": fake_atmos(),
+ "ocean": fake_ocean(),
+ "med": TrailingAverageMediator("med", ["geopotential_at_1000hpa_48h_mean"]),
+ }
+
+
+def test_parse_dsl():
+ seq = parse_run_sequence(DSL)
+ assert len(seq.slots) == 2
+ assert seq.slots[0].interval == np.timedelta64(6, "h")
+ assert seq.slots[0].actions == [
+ ConnectAction("atmos", "med"),
+ ConnectAction("ocean", "atmos"),
+ RunAction("atmos"),
+ ]
+ assert seq.slots[1].actions == [
+ MediateAction("med", "compute"),
+ ConnectAction("med", "ocean"),
+ RunAction("ocean"),
+ ]
+
+
+def test_parse_dsl_object_equivalence_and_str_roundtrip():
+ seq = parse_run_sequence(DSL)
+ manual = RunSequence(
+ [
+ Slot(
+ "6h",
+ [
+ ConnectAction("atmos", "med"),
+ ConnectAction("ocean", "atmos"),
+ RunAction("atmos"),
+ ],
+ ),
+ Slot(
+ "48h",
+ [
+ MediateAction("med"),
+ ConnectAction("med", "ocean"),
+ RunAction("ocean"),
+ ],
+ ),
+ ]
+ )
+ assert seq == manual
+ assert parse_run_sequence(str(seq)) == seq # str() emits valid DSL
+
+
+def test_str_preserves_subhour_intervals():
+ """'90m' and '30m' slots must not truncate to whole hours in str().
+
+ Regression: str() used astype('timedelta64[h]'), silently turning a 90 m
+ slot into '@1h' and a 30 m slot into '@0h' through YAML round-trips.
+ """
+ seq = parse_run_sequence("@90m\n atmos\n@30m\n ocean\n@")
+ text = str(seq)
+ assert "@90m" in text
+ assert "@30m" in text
+ rebuilt = parse_run_sequence(text)
+ assert rebuilt == seq
+ assert rebuilt.slots[0].interval == np.timedelta64(90, "m")
+ assert rebuilt.slots[1].interval == np.timedelta64(30, "m")
+
+
+def test_validate_mediator_cadence_mismatch():
+ # med (48h window) scheduled via med.compute in the 6h slot must be
+ # rejected, exactly like a RunAction cadence mismatch
+ seq = parse_run_sequence(
+ "@6h\n atmos -> med\n atmos\n med.compute\n@48h\n ocean\n@"
+ )
+ with pytest.raises(CadenceError, match="med"):
+ seq.validate(_components(), "6h")
+
+
+def test_parse_errors():
+ with pytest.raises(SequenceError, match="outside any"):
+ parse_run_sequence("atmos\n@6h\n")
+ with pytest.raises(SequenceError, match="cannot parse"):
+ parse_run_sequence("@6h\n atmos -> \n")
+ with pytest.raises(SequenceError, match="empty"):
+ parse_run_sequence("# just a comment\n")
+
+
+def test_validate_ok():
+ seq = parse_run_sequence(DSL)
+ seq.validate(_components(), "6h")
+
+
+def test_validate_unknown_name_with_suggestion():
+ seq = parse_run_sequence("@6h\n atmoss\n@")
+ with pytest.raises(SequenceError, match="atmos"):
+ seq.validate(_components(), "6h")
+
+
+def test_validate_cadence_mismatch():
+ # ocean (48h) scheduled in the 6h slot
+ seq = parse_run_sequence("@6h\n atmos\n ocean\n@48h\n med.compute\n@")
+ with pytest.raises(CadenceError):
+ seq.validate(_components(), "6h")
+ # slot interval not a multiple of driver dt
+ seq = parse_run_sequence("@7h\n atmos\n@")
+ with pytest.raises(CadenceError):
+ seq.validate(_components(), "6h")
+
+
+def test_validate_idle_component():
+ seq = parse_run_sequence("@6h\n atmos\n@")
+ with pytest.raises(SequenceError, match="never run"):
+ seq.validate(_components(), "6h")
+
+
+def test_helpers():
+ seq = parse_run_sequence(DSL)
+ assert seq.components_run() == {"atmos", "ocean", "med"}
+ assert ConnectAction("ocean", "atmos") in seq.connections()
+
+
+# -- derive_sequence: the schedule implied by the coupling graph -----------------
+TRIO_EDGES = [("atmos", "med"), ("ocean", "atmos"), ("med", "ocean")]
+
+
+def test_derive_matches_hand_written_dsl_exactly():
+ """Acceptance: the derived sequence for the classic toy trio equals the
+ hand-written canonical DSL, string for string."""
+ derived = derive_sequence(_components(), TRIO_EDGES)
+ assert str(derived) == str(parse_run_sequence(DSL))
+ derived.validate(_components(), "6h")
+
+
+def test_derive_edge_declaration_order_is_irrelevant():
+ a = derive_sequence(_components(), TRIO_EDGES)
+ b = derive_sequence(_components(), list(reversed(TRIO_EDGES)))
+ assert str(a) == str(b)
+
+
+def test_derive_accepts_connector_objects():
+ from earth2studio.nvcoupler.connector import Connector
+
+ comps = _components()
+ conns = [
+ Connector(comps["atmos"], comps["med"]),
+ Connector(comps["ocean"], comps["atmos"]),
+ Connector(comps["med"], comps["ocean"]),
+ ]
+ assert str(derive_sequence(comps, conns)) == str(parse_run_sequence(DSL))
+
+
+def test_derive_sequential_edges_reorder_runs():
+ """Non-lagged (sequential) connects land after their source's run; the
+ trio with a sequential ocean->atmos edge reproduces test_driver.py's
+ hand-written 'sequential' DSL."""
+ seq = derive_sequence(_components(), TRIO_EDGES, lagged={("atmos", "med")})
+ expected = parse_run_sequence(
+ "@6h\n atmos -> med\n atmos\n"
+ "@48h\n med.compute\n med -> ocean\n ocean\n ocean -> atmos\n@"
+ )
+ assert str(seq) == str(expected)
+
+
+def test_derive_same_cadence_sequential_dependency_order():
+ a, b = fake_atmos(), fake_atmos()
+ b.name = "atmos2"
+ comps = {"atmos": a, "atmos2": b}
+ seq = derive_sequence(comps, [("atmos2", "atmos")], lagged=set())
+ assert [str(x) for s in seq.slots for x in s.actions] == [
+ "atmos2",
+ "atmos2 -> atmos",
+ "atmos",
+ ]
+
+
+def test_derive_sequential_cycle_raises():
+ a, b = fake_atmos(), fake_atmos()
+ b.name = "atmos2"
+ comps = {"atmos": a, "atmos2": b}
+ with pytest.raises(SequenceError, match="lagged"):
+ derive_sequence(comps, [("atmos", "atmos2"), ("atmos2", "atmos")], lagged=set())
+ # marking one edge lagged breaks the cycle
+ seq = derive_sequence(
+ comps,
+ [("atmos", "atmos2"), ("atmos2", "atmos")],
+ lagged={("atmos2", "atmos")},
+ )
+ assert [str(x) for s in seq.slots for x in s.actions] == [
+ "atmos2 -> atmos",
+ "atmos",
+ "atmos -> atmos2",
+ "atmos2",
+ ]
+
+
+def test_derive_unknown_name_raises_with_suggestion():
+ with pytest.raises(SequenceError, match="atmos"):
+ derive_sequence(_components(), [("atmoss", "ocean")])
+
+
+def test_derive_no_connections_runs_everything():
+ comps = {"atmos": fake_atmos(), "ocean": fake_ocean()}
+ seq = derive_sequence(comps, [])
+ assert str(seq) == "@6h\n atmos\n@48h\n ocean\n@"
diff --git a/test/nvcoupler/test_vertical.py b/test/nvcoupler/test_vertical.py
new file mode 100644
index 000000000..b4f22589b
--- /dev/null
+++ b/test/nvcoupler/test_vertical.py
@@ -0,0 +1,157 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+from collections import OrderedDict
+
+import numpy as np
+import pytest
+import torch
+
+from earth2studio.nvcoupler.errors import VerticalMismatchError
+from earth2studio.nvcoupler.vertical import (
+ HybridLevels,
+ PressureLevels,
+ interp_to_pressure,
+)
+
+
+def _column(values, nlat=4, nlon=8, levels=None):
+ """(level, lat, lon) tensor with a horizontally-uniform column.
+
+ `levels` sets the 'level' coordinate (hPa for pressure sources, model
+ level index for hybrid sources); defaults to a 0..N-1 index.
+ """
+ lev = torch.tensor(values, dtype=torch.float64)
+ x = lev.view(-1, 1, 1).expand(len(values), nlat, nlon).clone()
+ if levels is None:
+ levels = np.arange(len(values), dtype=np.float64)
+ coords = OrderedDict(
+ {
+ "level": np.asarray(levels, dtype=np.float64),
+ "lat": np.linspace(90, -90, nlat),
+ "lon": np.linspace(0, 360, nlon, endpoint=False),
+ }
+ )
+ return x, coords
+
+
+def test_pressure_to_pressure_linear_in_logp():
+ # Field varies linearly in log-p: f = log(p). Interpolation onto any
+ # intermediate level must be exact.
+ src_levels = (100.0, 300.0, 500.0, 850.0, 1000.0)
+ f = np.log(np.array(src_levels) * 100.0)
+ x, coords = _column(f, levels=src_levels)
+ dst = PressureLevels((200.0, 400.0, 700.0))
+ out, out_coords = interp_to_pressure(x, coords, PressureLevels(src_levels), dst)
+ expected = np.log(np.array(dst.levels) * 100.0)
+ assert np.allclose(out[:, 0, 0].numpy(), expected)
+ assert list(out_coords["level"]) == [200.0, 400.0, 700.0]
+
+
+def test_hybrid_to_pressure_with_surface_pressure():
+ # Hybrid column: p_k = a_k + b_k * ps. With ps = 1000 hPa the levels land
+ # at 300/700/1000 hPa; field = log(p) again for exactness.
+ a = (30000.0, 20000.0, 0.0)
+ b = (0.0, 0.5, 1.0)
+ ps_value = 100000.0
+ p_src = np.array(a) + np.array(b) * ps_value # [30000, 70000, 100000] Pa
+ x, coords = _column(np.log(p_src))
+ ps = torch.full((4, 8), ps_value, dtype=torch.float64)
+ dst = PressureLevels((500.0, 850.0))
+ out, _ = interp_to_pressure(x, coords, HybridLevels(a, b), dst, ps=ps)
+ expected = np.log(np.array([50000.0, 85000.0]))
+ assert np.allclose(out[:, 0, 0].numpy(), expected)
+
+
+def test_hybrid_requires_ps():
+ x, coords = _column([1.0, 2.0, 3.0])
+ with pytest.raises(VerticalMismatchError, match="surface pressure"):
+ interp_to_pressure(
+ x,
+ coords,
+ HybridLevels((1.0, 2.0, 3.0), (0.0, 0.0, 0.0)),
+ PressureLevels((500.0,)),
+ )
+
+
+def test_clamped_at_column_ends():
+ x, coords = _column([10.0, 20.0], levels=(300.0, 700.0))
+ src = PressureLevels((300.0, 700.0))
+ out, _ = interp_to_pressure(x, coords, src, PressureLevels((100.0, 1000.0)))
+ assert torch.all(out[0] == 10.0) # above top -> top value
+ assert torch.all(out[1] == 20.0) # below bottom -> bottom value
+
+
+def test_identity_shortcut_and_missing_level_dim():
+ x, coords = _column([1.0, 2.0], levels=(500.0, 850.0))
+ src = PressureLevels((500.0, 850.0))
+ out, out_coords = interp_to_pressure(x, coords, src, PressureLevels((500.0, 850.0)))
+ assert out is x # no-op
+ with pytest.raises(VerticalMismatchError, match="no 'level' dim"):
+ bad = OrderedDict((k, v) for k, v in coords.items() if k != "level")
+ interp_to_pressure(x[0], bad, src, PressureLevels((500.0,)))
+
+
+def test_pressure_source_rejects_mismatched_level_coord():
+ """Descending level data against an ascending PressureLevels source must
+ raise instead of silently pairing slices with wrong pressures."""
+ src = PressureLevels((500.0, 850.0, 1000.0))
+ dst = PressureLevels((700.0,))
+ # data ordered bottom-to-top, source declared top-to-bottom
+ x, coords = _column([3.0, 2.0, 1.0], levels=(1000.0, 850.0, 500.0))
+ with pytest.raises(VerticalMismatchError, match="1000.*850.*500"):
+ interp_to_pressure(x, coords, src, dst)
+ # arbitrary index coords (not the declared pressures) are rejected too
+ x, coords = _column([1.0, 2.0, 3.0]) # level = [0, 1, 2]
+ with pytest.raises(VerticalMismatchError, match="PressureLevels"):
+ interp_to_pressure(x, coords, src, dst)
+ # ... even on the identity (src == dst levels) shortcut path
+ x, coords = _column([1.0, 2.0, 3.0])
+ with pytest.raises(VerticalMismatchError, match="level"):
+ interp_to_pressure(x, coords, src, PressureLevels(src.levels))
+
+
+def test_hybrid_crossing_coefficients_rejected_at_construction():
+ # crossing at low surface pressure: p = [50000, 30000] at ps = 50000 Pa
+ with pytest.raises(ValueError, match="non-increasing"):
+ HybridLevels((0.0, 30000.0), (1.0, 0.0))
+ # increasing at ps = 50000 Pa but crossing at ps = 110000 Pa
+ with pytest.raises(ValueError, match="non-increasing"):
+ HybridLevels((0.0, 60000.0), (1.0, 0.0))
+ # non-crossing across the whole plausible ps range is accepted
+ HybridLevels((30000.0, 20000.0, 0.0), (0.0, 0.5, 1.0))
+
+
+def test_hybrid_crossing_at_runtime_rejected():
+ """Coefficients valid for plausible ps can still cross for extreme ps
+ values; the interpolation-time monotonicity check must catch that."""
+ hybrid = HybridLevels((0.0, 115000.0), (1.0, 0.0)) # ok in [50k, 110k] Pa
+ x, coords = _column([1.0, 2.0])
+ ps = torch.full((4, 8), 140000.0, dtype=torch.float64) # crosses level 2
+ with pytest.raises(VerticalMismatchError, match="strictly increasing"):
+ interp_to_pressure(x, coords, hybrid, PressureLevels((1000.0,)), ps=ps)
+
+
+def test_gradient_flows_through_interpolation():
+ x, coords = _column([1.0, 2.0, 3.0], levels=(300.0, 500.0, 1000.0))
+ x.requires_grad_(True)
+ out, _ = interp_to_pressure(
+ x, coords, PressureLevels((300.0, 500.0, 1000.0)), PressureLevels((400.0,))
+ )
+ out.sum().backward()
+ assert x.grad is not None
+ assert torch.all(x.grad[2] == 0) # bottom level unused for 400 hPa target
+ assert torch.all(x.grad[:2] > 0)