feat: add CosmoDownscaling diagnostic model (ERA5 -> COSMO-REA6/REA2)#939
feat: add CosmoDownscaling diagnostic model (ERA5 -> COSMO-REA6/REA2)#939gertln wants to merge 6 commits into
Conversation
Add CosmoDownscaling, a diagnostic model that downscales ERA5 to high-resolution COSMO-REA regional reanalysis -- COSMO-REA6 (~6 km) and COSMO-REA2 (~2.2 km) -- each with a deterministic regression (mean) and a generative EDM/Karras diffusion mode, selected via `mode` on `load_model`. Supports movable sub-domains via `set_domain` and optional hub-height wind components for wind-energy use. - earth2studio/models/dx/cosmo_downscaling.py: the wrapper (physicsnemo DiT with axial RoPE + NATTEN attention, EDM diffusion sampler, channel transforms + physical-space constraints incl. a solar-zenith day/night gate, sub-domain cropping, derived hub-height wind components) - test/models/dx/test_cosmo_downscaling.py: construction + behavior tests with mocked nets, plus a GPU `--package` test against a real package - examples/03_downscaling/04_cosmo_rea_downscaling.py: end-to-end example (SFNO -> downscale full domain, sub-domain, rollout, diffusion ensemble vs regression mean, hub-height wind, COSMO-REA2) - register the model; add the `cosmo` extra; install/CHANGELOG/API-doc entries Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: gertl <gertl@nvidia.com>
Output domains: native and extended grid footprints for both resolutions. REA6 spans broad Europe at ~6 km; REA2 is central-European natively at ~2.2 km, with an extended grid reaching a comparable broad domain.
|
Greptile SummaryAdds
|
| Filename | Overview |
|---|---|
| earth2studio/models/dx/cosmo_downscaling.py | New 1995-line CosmoDownscaling diagnostic model with regression and diffusion modes, rotated-pole output grid, set_domain sub-region slicing, hub-height wind derivation, and full package-loading logic; well-engineered but shares mutable DiT state across sub-domain instances and hardcodes 0.25° ERA5 spacing in set_domain |
| test/models/dx/test_cosmo_downscaling.py | Comprehensive 1324-line CPU test suite covering coordinate contracts, constraint parsing, domain slicing, halo/patch-snap edge cases, hub-wind derivation, diffusion schedule, and a constructor-signature guard that forces new args to be threaded through set_domain |
| examples/03_downscaling/04_cosmo_rea_downscaling.py | End-to-end example covering SFNO to CosmoDownscaling rollout, sub-domain, diffusion ensemble vs regression mean, hub-height wind, and COSMO-REA2; correctly listed in expected_failing_examples since weights are not yet hosted |
| pyproject.toml | Adds cosmo extra (einops, natten, nvidia-physicsnemo>=2.0, nvtx) and includes it in the all extra; all permissive licenses |
| earth2studio/models/dx/init.py | Adds CosmoDownscaling import and all entry |
| docs/conf.py | Adds the COSMO-REA example to expected_failing_examples with a clear comment; appropriate short-term measure until weights are hosted |
| test/conftest.py | Adds the new test file to _TEST_DEPENDENCIES under the cosmo extra, consistent with other model test registrations |
Reviews (1): Last reviewed commit: "feat: add CosmoDownscaling diagnostic mo..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
…currency - set_domain: derive the ERA5 input-axis spacing from the package's own validated grid instead of a hardcoded 0.25 deg (behavior-preserving for the shipped 0.25 deg package; supports a non-0.25 deg package). - interp_levels_to_height: rename local `interp` -> `interp_val` so it no longer shadows the module-level `interp` import. - set_domain / _rebind_latent: document that the network is shared by reference and rebinds its latent-grid state in place, so sub-domains are safe sequentially but not concurrently; and that a second set_domain on a returned sub-domain cannot reach back to the full/extended footprint. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
| cosmo = [ | ||
| "einops>=0.8.1", | ||
| "natten ; python_version < '3.14'", | ||
| "nvidia-physicsnemo>=2.0", |
There was a problem hiding this comment.
I believe we need a stricter constraint on this if we need RoPE stuff from what I recently merged
There was a problem hiding this comment.
Pinned nvidia-physicsnemo>=2.2.0. natten2d_rope is only on main (2.2.0a0) and the latest release (2.1.1) doesn't have it. So it is intentionally uninstallable until 2.2.0 ships. (Also removed cosmo from the all extra because >=2.2.0 conflicts with da-healda <=2.1.1 )
There was a problem hiding this comment.
So it is intentionally uninstallable until 2.2.0 ships
We can just have the cosmo dependency pin physicsnemo to a specific git commit from main and then it will be usable in the meantime. Just add a TODO comment to update to 2.2.0 when ready, and drop cosmo from the all group if there are any conflicts that arise
| "CLCT": "tcc", # total cloud cover; % -> 0-1 fraction (COSMO_OUTPUT_UNIT_SCALE) | ||
| "PS": "sp", | ||
| "PMSL": "msl", | ||
| } |
There was a problem hiding this comment.
Can we open and merge a COSMO-REA data source before this PR so that the model wrapper can simply reuse the COSMO lexicon? Adding a lexicon to the model wrapper file is not standard
There was a problem hiding this comment.
And, I don't see a particular for the model wrapper to be doing any lexicon interaction in the first place. Wrappers operate based on input_coords and output_coords, which should specify variables that exist in the E2S lexicon. If there are things in the model package using these other names, ideally we rename them to avoid having to deal with lexicon conversions in model code. But, if absolutely necessary, then we should pull from a COSMO lexicon as I suggested above
There was a problem hiding this comment.
Added a CosmoLexicon in earth2studio/lexicon/ and the wrapper uses it
|
|
||
| @check_optional_dependencies() | ||
| class CosmoDownscaling(torch.nn.Module, AutoModelMixin): | ||
| """COSMO-REA downscaling model: ERA5 -> high-resolution COSMO-REA. |
There was a problem hiding this comment.
This should eventually get filled out with details of the model, use-cases, intended usage, etc
There was a problem hiding this comment.
added details about the two modes, use cases and intended usage
There was a problem hiding this comment.
Nice, looks good. Can we add an additional Note with references/links to the dataset?
…, standard sampler - Add CosmoLexicon (earth2studio/lexicon/cosmo.py); the wrapper maps COSMO output names to Earth2Studio names via it, replacing the in-module mapping dicts. - Load checkpoints with the standard DiT.from_checkpoint / EDMPreconditioner.from_checkpoint; drop the manual .mdlus zip extraction, the hand-built DiT, and the forced NATTEN backend (native modules let NATTEN auto-select the kernel per GPU). - Replace the inline EDM/Karras sampler loop with physicsnemo's EDMNoiseScheduler + sample(). - Literal[...] for fixed-option string args; move physicsnemo/json/xarray imports to the top (physicsnemo in the protected optional-dependency block); resolve a package-root config.json in load_model for HuggingFace download tracking. - Pin nvidia-physicsnemo>=2.2.0 (RoPE); drop cosmo from the `all` extra (it conflicts with da-healda's <=2.1.1 cap until 2.2.0 ships). - Docstrings/comments: "crop-size agnostic at the fixed resolution" wording, remove internal/dangling references, expand the class docstring, TODO(cosmo) markers, badge year 2026. - Tests: add CosmoLexicon tests + an euler-vs-heun solver test; route the output-name test through CosmoLexicon; drop the obsolete DiT-builder test. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: gertl <gertl@nvidia.com>
…xtra Address the latest PR review round (Peter Harrington): - pyproject: pin nvidia-physicsnemo to a physicsnemo `main` commit (natten2d_rope is not in a tagged release yet) so `earth2studio[cosmo]` is installable now, with a TODO to revert to `>=2.2.0` at release. Note in-file that PyPI rejects direct-URL deps, so the pin must not reach a tagged release. Refresh the `all`-exclusion note (cosmo stays out of `all` while the git-pin is in place). - Add COSMO-REA6 / REA2 (DWD) and ERA5 (ECMWF) dataset links to the CosmoDownscaling class docstring and the CosmoLexicon docstring, matching the model card. - Document the example's runtime (~3 min) and peak GPU memory (~18 GB) in its header note. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: gertl <gertl@nvidia.com>
Rename the downscaling model to the corrdiff-cosmo-era5 brand: class CosmoDownscaling -> CorrDiffCosmoEra5, module/test files -> corrdiff_cosmo_era5, and HF/cache/URI/doc brand strings updated. Kept CosmoLexicon and the cosmo pip extra. CPU suite: 60 passed / 4 skipped. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: gertl <gertl@nvidia.com>
…o 2.2.0) Point corrdiff-cosmo-era5 at its hosted Hugging Face package, now that it ships alongside physicsnemo 2.2.0 (which contains the RoPE natten2d_rope backend): - Set DEFAULT_PACKAGE_URI = "hf://nvidia/corrdiff-cosmo-era5" and drop the NotImplementedError placeholder in load_default_package. - pyproject: the cosmo extra requires nvidia-physicsnemo>=2.2.0; cosmo stays out of `all` only because da-healda still caps physicsnemo at <=2.1.1. - Example loads the hosted default (with $COSMO_REA_PACKAGE as a local override); remove the docs-gallery whitelist so it runs like other GPU examples. - Tests: run the real-weight --package test unconditionally (drop the now-obsolete natten2d_rope availability guard); update the load_default_package test to assert the hosted URI. CPU suite: 60 passed / 4 skipped. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: gertl <gertl@nvidia.com>
2 participants
Description
Add
CosmoDownscaling, a diagnostic model that downscales global ERA5 to high-resolution COSMO-REA regional reanalysis over Europe — COSMO-REA6 (~6 km) and COSMO-REA2 (~2.2 km). Its input is an ERA5 state, so it can downscale an ERA5 analysis directly or run behind a global forecast model (e.g. SFNO → CosmoDownscaling). A runnable example is included.A single class loads one of four checkpoints via two
load_modelselectors:mode="mean"— a deterministic regression that predicts the conditional mean (the expected field,E[y | x]): a single, smooth field in one forward pass. Fast and cheap — a good first high-resolution look.mode="diffusion"— a generative diffusion model that samples the conditional distributionp(y | x): an ensemble of realizations that also captures the spread the mean cannot represent (~N forward passes; seeded for reproducibility).resolution="rea6"|"rea2"— selects the per-resolution weights and grid from one package.Both modes use a PhysicsNeMo DiT (diffusion transformer) at a fixed output resolution; the network is crop-size agnostic, so
set_domain(...)returns a new instance for lat/lon sub-region. From a small box up to the full grid. Optional hub-height wind adds derivedu{H}m/v{H}mcomponents (wind speed via the stockDerivedWS)Blocked by:
Closes: #940
Diagnostic details
modeinput_coords())u{H}m/v{H}msampledim always present: 1 formean,number_of_samplesfordiffusion(per-sample seedseed + i)Dependencies added
New
cosmomodel extra (pip install earth2studio[cosmo]/uv add earth2studio --extra cosmo), also added to theallextra.nvidia-physicsnemo>=2.0natten; python_version < '3.14'einops>=0.8.1nvtx>=0.2.11All permissive (Apache-2.0 / MIT); no non-permissive licenses introduced.
Validation
test/models/dx/test_cosmo_downscaling.py): 51 passed, 5 skipped. The 5 skips are intentionally resource-gated tests — the 4 real-weight GPU--packagecases and 1 real-DiT-builder smoke test — which require a GPU, a built package, and/or physicsnemo's not-yet-released RoPE op, so they cannot run in the default CI lane (the standard@pytest.mark.packagepattern used by every model).{rea6, rea2} × {mean, diffusion}combos load 0 missing / 0 unexpected keys and run a fulldx(x, coords)forward (output shape, finiteness, and every metadata min/max bound asserted in both modes), plus the DiT-builder smoke test.black+ruffclean; SPDX headers present.examples/03_downscaling/04_cosmo_rea_downscaling.pyverified end-to-end on real weights (full domain, sub-domain, rollout, diffusion ensemble vs regression mean, hub-height wind, COSMO-REA2).The CPU tests are construction-driven and GPU-free: they pin the wrapper's contract — coords/handshake + lexicon output names, channel-transform round-trips, physical-space min/max bounds and a mode-identical solar-zenith day/night gate,
set_domainsub-domain slicing (incl. NATTEN-minimum and patch-snap edges), the diffusion sampler's schedule termination + ensemble seeding, and the hub-wind derivation +DerivedWSnaming handshake.Notes for reviewers
Two known, bounded items, each guarded so CI stays green today:
natten2d_rope, merged on physicsnemomainbut not in a tagged release. Thecosmoextra therefore pinsnvidia-physicsnemo>=2.0, and the real-DiT-builder +--packagetests are skip-guarded when the op is unavailable (they skip, not fail). A follow-up bumps the pin and drops the guards once a release includes it.load_default_packageraisesNotImplementedError; the example is inexpected_failing_examplesindocs/conf.pyso the gallery build stays green. FlippingDEFAULT_PACKAGE_URIon and removing the whitelist entry is the only change needed once hosting lands. (The model was validated against real weights locally.)Checklist