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)