Conversation
A Python-native coupling layer letting independent AI Earth-system components exchange fields through a Driver, Connectors, and Mediators, instead of coupling baked into datapipes or single-model forward passes. Core: Field/State exchange (torch-native, autograd-clean), CF-style field dictionary with CellMethod derived fields, Clock/Alarm cadences, Components (Prognostic/Data/Diagnostic/Callable) with pluggable ImportAdapters matching real model call shapes, Connector pipeline (time policies, hybrid-to-pressure vertical interpolation, mask fill, cached regridding), accumulation mediators, NUOPC-runSeq-style DSL. UX: couple() auto-wiring, describe() dry-run preview, notebook steps()/probe(), xarray or streaming Zarr output (NaN-poisoned unwritten rows), YAML round-trip, rollout() keeping gradients for coupled fine-tuning. Real-model integration: DLESyM split adapter exposing atmos/ocean sub-models as separate components, with an opt-in real-weights equivalence gate (NVCOUPLER_DLESYM_WEIGHTS=1, not yet run). 146 tests on deterministic synthetic components (hand-computed expectations, no weights or network), including cross-component seam tests; survived a design review and an execution-verified adversarial code review with all 12 confirmed findings fixed with regression tests. Five runnable examples under examples/09_nvcoupler. No existing earth2studio files are modified. 🤖 Generated with [Claude Code](https://claude.com/claude-code)
Six pages under earth2studio/nvcoupler/docs/, linked from the README: concepts (exchange contracts, adapters, connector pipeline), user guide (11 task recipes), DSL & YAML reference, API reference (signatures read from source), errors & troubleshooting (every exception with verbatim greppable messages), and design & roadmap (decisions with alternatives, verification story, honest limitations). Every code snippet was executed against the synthetic toy components before inclusion; all error messages in the catalog are quoted verbatim from triggered exceptions. 🤖 Generated with [Claude Code](https://claude.com/claude-code)
Elegance refactors from the post-review cleanup pass: - Single cadence mechanism: delete the vestigial Alarm class — the driver's slot alignment was already the only mechanism consulted; components now just declare a timestep. - Windowed connectors: Connector(window=, reduce=) folds a trailing reduction (mean/sum/max/min) into the exchange itself, so the common fast-to-slow coupling no longer needs a third Mediator component and extra run-sequence lines; Mediator remains for multi-source computations, sharing the same accumulator core. - UnknownFieldError drops the KeyError dual inheritance and its __str__ workaround. 151 tests pass (windowed-connector tests added; Alarm-mechanism tests removed with the class). 🤖 Generated with [Claude Code](https://claude.com/claude-code)
couple() now derives the run sequence from component cadences and the connector graph (lagged by default), synthesizing windowed connectors for derived imports — a mediator is only created when a pair also carries a plain transfer of the same fields. The canonical toy system reproduces its hand-computed values (z96=19.2336, sst=2.180147) with a 5-line derived sequence and no mediator component. Docs alignment pass across all six pages and the README: Alarm removal, windowed connectors documented as the preferred fast-to-slow path with Mediator repositioned as the multi-source general form, YAML window/reduce no-round-trip limitation stated, UnknownFieldError bases corrected, test counts refreshed. Every new snippet executed before inclusion; stale mediator docstrings fixed. 170 tests pass. 🤖 Generated with [Claude Code](https://claude.com/claude-code)
Couples models that fetch their own forcing (StormCast-style) without modifying them: PullAdapter installs a StateDataSource shim on the model's conditioning_data_source attribute before each step, so the model's unmodified production fetch path (fetch_data -> interpolate -> concatenate) receives this step's live coupled forcing — the same masquerade the serve workflows play with InferenceOutputSource, minus the store-and-replay staging. Verified against a protocol-faithful mock pulling through the real earth2studio.data.utils.fetch_data, including an end-to-end global -> regional sequential run proving each step pulls that step's fresh conditioning. Honest limitation documented: the pull path crosses fetch_data's numpy boundary, so pull-coupled components are inference-only. New example 06_pull_conditioning.py (self-checking, hand-computable), docs coverage in concepts/user guide/API reference/error catalog, and package exports. 176 tests pass. 🤖 Generated with [Claude Code](https://claude.com/claude-code)
Closes the elegance workflow's final integrator notes: delete the temporary Component.should_run cadence shim (its only caller in test_prognostic_real.py dropped), and reword the last alarm-era docstring in mediator.py. Driver.reset() already delegates to Connector.reset(). 🤖 Generated with [Claude Code](https://claude.com/claude-code)
negin513
marked this pull request as draft
August 27, 2026 17:38
for more information, see https://pre-commit.ci
Contributor
Greptile SummaryThe PR introduces the nvcoupler prototype for composing independently scheduled Earth-system components through field connectors, mediators, configurable run sequences, and YAML configuration.
|
| Filename | Overview |
|---|---|
| earth2studio/nvcoupler/driver.py | Implements coupled lifecycle, scheduling, collection, and IO, but does not reject spans that are misaligned with coarse component cadences. |
| earth2studio/nvcoupler/sequence.py | Implements DSL parsing, derived scheduling, and cadence validation; its validation does not account for the clock’s total span. |
| earth2studio/nvcoupler/connector.py | Implements field matching and the time, vertical, fill, regrid, and windowed-reduction transfer pipeline with substantial focused coverage. |
| earth2studio/nvcoupler/component.py | Adds component lifecycle abstractions and import adapters for state overwrite, conditioning, extra-tensor, data, diagnostic, and prognostic workflows. |
| earth2studio/nvcoupler/config.py | Adds YAML serialization and reconstruction using dotted factory paths; no current untrusted configuration path was established. |
| earth2studio/nvcoupler/vertical.py | Adds hybrid-to-pressure interpolation with metadata validation and gradient-preserving tensor operations. |
| earth2studio/nvcoupler/api.py | Adds public auto-wiring, one-call coupled execution, and text or HTML plan inspection. |
| test/nvcoupler/test_driver.py | Covers aligned mixed-cadence execution, exhaustion, reset, rollout, collection, and IO, but not a clock span misaligned with a component cadence. |
Reviews (1): Last reviewed commit: "Remove should_run shim and last alarm-er..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| self, ics: dict[str, tuple[torch.Tensor, CoordSystem]] | None = None | ||
| ) -> None: | ||
| ics = ics or {} | ||
| self.sequence.validate(self.components, self.clock.dt) |
Contributor
There was a problem hiding this comment.
When the clock span is divisible by dt but not by a component timestep, initialization succeeds while _slot_aligned never runs that component at the requested stop time. For example, a 120-hour run with a 6-hour clock and a 48-hour component leaves that component and its persisted or collected output at 96 hours, so the run completes with stale state for its final 24 hours.
Adds the grid-to-point sampling primitive that was previously only a "# roadmap" comment in the impact-modeling design doc's Tessera downscaling example (needed by that application, and by any site-level use case: station verification, agriculture field points, energy asset locations). - points.py: new PointSet(lat, lon, names=None) — a scattered N-location target grid, the "point" analog of a lat/lon mesh. - component.py: Component gains a points= constructor kwarg; grid_coords() reports the point target when set. - connector.py: Connector gains sample="nearest"|"bilinear", mutually exclusive with regridder=. Auto-builds a cached sampler when the destination is a point target: bilinear reuses the existing regular-grid kernel via a degenerate [N,1] mesh, nearest does a great-circle KDTree lookup. Raises actionable errors when the destination is a point target but neither sample= nor a custom regridder= is given, or when it has no points= metadata, or when the source lacks lat/lon. - config.py: sample= now round-trips through to_yaml/from_yaml (previously missing entirely for connector kwargs added after fill=/time_policy=). - field.py: "point" added to _SPATIAL_DIMS. - 14 new tests (7 PointSet unit tests, 7 Connector integration tests); zero regressions across the pre-existing nvcoupler suite. - New worked example: examples/09_nvcoupler/07_point_sampling.py. - Docs: api_reference.md, concepts.md, dsl_and_yaml_reference.md, errors_and_troubleshooting.md.
for more information, see https://pre-commit.ci
negin513
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Earth2Studio Pull Request
Description
Checklist
Dependencies