Stormcast-CONUS

[jleinonen]

Earth2Studio Pull Request

Description

Adds StormCast-CONUS model: A CONUS-wide StormCast-like model with a DiT backbone.

Main features/additions:

1. models.nn.stormcastconus.StormCastCONUSBase model wrapper:
   • Uses _SplitModelWrapper to handle per-sigma and per-variable expert split
   • Support for SDA
   • Support for splitting model to subregion of full CONUS
   • Optimizations: toggleable model compilation and AMP inference
2. models.px.stormcastconus.StormCastCONUS and models.da.sda_stormcastconus.StormCastCONUSSDA specialize the above as forecasting and DA models
3. utils.obs.ObsGridMapping for converting between point obs and gridded variables, with CuPy support
4. utils.cos_zenith for a torch port of cos zenith calculations (Stormscope uses version from PhysicsNeMo, which is implemented in NumPy and runs on CPU)
5. Example, following the structure of the StormCast example
6. SDA example, following the structure of the StormCast SDA example
7. Earth2Studio API example utilizing StormCast-CONUS
8. Unit tests
9. Docs updated

TBA/TBD:

1. Do we want to further generalize SplitModelWrapper?
2. Discuss what to do with utils.cos_zenith which admittedly adds code duplication
3. Add example using SDA
4. Add workflow for Earth2Studio API for using the model with various conditioning models (e.g. GFS, FCN3, AIFS)
5. Add unit tests

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• Assess and address Greptile feedback (AI code review bot for guidance; use discretion, addressing all feedback is not required).

Dependencies

[pzharrington]

Given the functionality, this feels like it would belong more alongside the other interpolation utils we have in earth2studio.utils.interp. In fact this whole file in general is maybe questionable, I know I was the one who approved its location here in earth2sutdio.models.da but maybe we should migrate the da utils to the general utils folder? Maybe I'm being too nitpicky... @NickGeneva any thoughts?

[jleinonen]

This now lives in earth2studio.utils.obs, is it ok there? I don't think we got @NickGeneva's opinion.

[greptile-apps]

Greptile Summary

This PR adds the StormCast-CONUS model: a 3 km HRRR-resolution DiT-based generative forecast model for the Continental United States, with support for score-based data assimilation (SDA) via DPS guidance and optional subdomain cropping.

• StormCastCONUSBase (models/nn/stormcastconus.py) orchestrates a sigma-split ensemble of four sub-models (_SplitModelWrapper), handles normalisation/denormalisation, cos-zenith conditioning, EDM diffusion sampling, and the obs-guided SDA path; StormCastCONUS and StormCastCONUSSDA specialise it for forecasting and DA respectively.
• ObsGridMapping (utils/obs.py) is a new utility that bins sparse point observations onto regular or irregular 2-D grids with CuPy GPU support and bilinear grid_to_obs sampling; filter_time_range is refactored out of models/da/utils.py and re-exported for backward compatibility.
• New examples (nowcasting and SDA), server workflow, and unit tests are included; the Sphinx docs reference for utils.ObsGridMapping points to the wrong module path and would break the docs build.

Confidence Score: 3/5

The model and utility code is functionally sound, but the docs change introduces a broken autosummary reference that would fail the documentation build.

The core StormCastCONUSBase, ObsGridMapping, and their tests look correct and well-structured. The one concrete problem is in docs/modules/models_da.rst: the utils.ObsGridMapping autosummary entry resolves to earth2studio.models.da.utils.ObsGridMapping, which does not exist. ObsGridMapping lives in earth2studio.utils.obs and is never re-exported from the DA utils module. Sphinx will fail to generate that page. The mutable default arguments for variables and conditioning_variables in the constructor are a low-risk code smell rather than a present defect.

docs/modules/models_da.rst needs the utils.ObsGridMapping reference fixed before the docs build will pass.

Important Files Changed

Filename	Overview
earth2studio/models/nn/stormcastconus.py	Core model implementation for StormCastCONUS with EDM diffusion, sigma-split sub-model routing, SDA support, and subdomain cropping. Mutable default arguments for variables and conditioning_variables parameters are a minor concern; otherwise well-structured.
earth2studio/utils/obs.py	New ObsGridMapping utility for binning point observations onto 2-D model grids with CuPy GPU support. Logic is correct; the filter_time_range function signature type hint is narrowed to pd.DataFrame but the body and docstring still handle cuDF DataFrames.
docs/modules/models_da.rst	Adds utils.ObsGridMapping to the DA autosummary, but ObsGridMapping is not present in earth2studio.models.da.utils — this would break the Sphinx docs build.
earth2studio/models/px/stormcastconus.py	Thin forecasting wrapper around StormCastCONUSBase; correctly delegates to super().__call__(x, coords, obs=None) and provides a standalone create_generator with PrognosticMixin hooks.
earth2studio/models/da/init.py	Exports StormCastCONUSBase as StormCastCONUSSDA; intentional — base class supports obs-guided diffusion. Does not inherit AssimilationModel, but the SDA example uses the generator interface directly rather than a generic DA framework.
earth2studio/models/da/utils.py	Removes the duplicate filter_time_range implementation and re-exports it from earth2studio.utils.obs. Clean refactor, though ObsGridMapping is missing as a re-export despite being referenced from the docs.
test/models/px/test_stormcastconus.py	Comprehensive unit tests covering call, iterator, AMP, value clamping, multi-time, and ensemble dimensions. Uses a well-structured stub diffusion model and correct subdomain alignment.
test/utils/test_obs.py	Good coverage of ObsGridMapping: construction, obs_coords, obs_to_grid (empty input, averaging, variable/time filtering), grid_to_obs (nearest and bilinear), irregular grids, and CPU/GPU parity checks.

_{Reviews (1): Last reviewed commit: "ObsGridMapping unit tests + support miss..." | Re-trigger Greptile}

[jleinonen]

/blossom-ci

[jleinonen]

/blossom-ci

[jleinonen]

/blossom-ci

[pzharrington]

TODO drop this when HF package is ready

[pzharrington]

TODO update when HF is available

[pzharrington]

Just curious did someone ask for this explicitly?

[jleinonen]

I have a user in mind for this, yes. Also generally I find these workflows are better starting points for building end user workflows (even if they're not using the API) than the examples in the examples directory.