Update StormScope model package

[pzharrington]

Earth2Studio Pull Request

Description

Updates the StormScope model in Earth2Studio to target the refreshednvidia/stormscope-goes-mrms package (commit 62f0fd2), replacing the legacy DiT-wrapping approach with direct physicsnemo.Module loading.
New models enable GOES, MRMS, and GLM Lightning predictions at CONUS-wide 3km / 10 minute spatiotemporal resolution. Example output:

Summary

• Drop legacy shim layer — deleted earth2studio/models/nn/stormscope_util.py (EDMPrecond, DropInDiT, model_wrap). Checkpoints now load via physicsnemo.Module.from_checkpoint; the returned module is stored directly in model_spec. Sampler calls updated to the new preconditioner signature (forward(x, t, condition=None), no round_sigma, no class_labels).

• Registry-driven model loading — load_model resolves variant names (and aliases) from registry.json rather than hard-coded indices. Normalization statistics and variables/conditioning_vars lists are now selected by name from the registry's top-level normalization block.

• Mixed normalization (GLM channels) — per-channel glm_mask / conditioning_glm_mask buffers gate between log1p/expm1 (GLM channels) and standard mean/std normalization, applied in _forward and normalize_conditioning.

• Static invariant channels — topo and nexrad_proximity arrays are loaded from the package and registered as buffers; appended in build_condition (nexrad then topo, broadcast over batch), enabling MRMS+GLM variants.

• GOESGLMGrid data source — new class in earth2studio/data/goes_glm.py wrapping GOESGLM event parsing into a gridded 0.1° product via 5-minute binning and histogram2d, matching preprocessing that was applied to model training data. Emits xr.DataArray[time, variable, lat, lon] with glm_density and glm_energy_density variables; exported from earth2studio.data. Also fixes a session-reuse bug in GOESGLM._fetch where the aiobotocore client could be torn down between repeated calls.

• AMP + torch.compile support — StormScopeBase gains amp and compile constructor arguments. When amp=True, DiT forward passes run under torch.autocast; sampler latent math stays in fp64. compile_experts() wraps each staged expert with torch.compile(mode="reduce-overhead"). Roughly 5x speedup for inference

• Dependency bump — nvidia-physicsnemo>=2.1.0 (was >=2.0).

• Tests — new parametrized tests cover AMP/compile flag propagation, mrms_coverage_mask AND-combination with the interpolator mask, and invalid fill-constant behavior for both the default and legacy paths.

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

[greptile-apps]

Greptile Summary

This PR refreshes the StormScope model package to target the new nvidia/stormscope-goes-mrms checkpoint (commit 62f0fd2), replacing the legacy EDMPrecond/DropInDiT shim with direct physicsnemo.Module.from_checkpoint loading and introducing registry-driven variant selection, mixed GLM log1p normalization, static invariant buffers (topo, nexrad_proximity), AMP/torch.compile support, and a new GOESGLMGrid data source.

• Model refactor — deletes stormscope_util.py; load_model now resolves variants by name from registry.json, builds normalization stats per channel (with log1p gating for GLM channels), and loads invariant buffers from the package. _INPUT_INVALID_FILL_CONSTANT for deprecated variants is set as a post-construction instance attribute in load_model, which does not survive state_dict save/load (see inline comment).
• GOESGLMGrid — wraps GOESGLM into a gridded 0.1-degree product via 5-min binning and histogram2d, matching training preprocessing; also fixes the session-reuse bug in GOESGLM._fetch.
• Example + tests — updated example uses the 3km_10min CONUS nowcasting variant with GLM channel assembly; new parametrized tests cover AMP/compile flags, mrms_coverage_mask AND-combination, and fill-constant behavior.

Confidence Score: 4/5

Safe to merge once the deprecated-variant fill-constant lifetime is addressed; all other changes are well-structured and covered by tests.

The deprecated 6km_1hr MRMS variant relies on _INPUT_INVALID_FILL_CONSTANT being set as a post-construction instance attribute in load_model. This attribute is invisible to state_dict, so any roundtrip that rebuilds the object and reloads weights will silently revert to 0.0 fill, corrupting inference for legacy checkpoints. All other changes are correct and well-tested.

earth2studio/models/px/stormscope.py - specifically StormScopeMRMS.load_model lines 2266-2268 and the absence of an input_invalid_fill_constant constructor argument.

Important Files Changed

Filename	Overview
earth2studio/models/px/stormscope.py	Major refactor: drops legacy EDMPrecond/DropInDiT shim in favour of physicsnemo.Module.from_checkpoint; adds registry-driven model loading, mixed GLM normalization, static invariant buffers, AMP/torch.compile support. One issue: deprecated variant's fill constant is a post-construction instance attribute not preserved through state_dict save/load.
earth2studio/data/goes_glm.py	Adds GOESGLMGrid: wraps GOESGLM into a gridded 0.1-degree product via 5-min binning and histogram2d; also fixes session-reuse bug in GOESGLM._fetch. Looks correct.
earth2studio/models/nn/stormscope_util.py	File deleted - removes the legacy EDMPrecond, DropInDiT, and model_wrap shim layer. No longer needed.
test/models/px/test_stormscope.py	Adds parametrized tests for AMP/compile flags, mrms_coverage_mask AND-combination, and fill-constant behavior for default and legacy paths.
examples/04_nowcasting/03_stormscope_goes_example.py	Updated to 3km_10min CONUS nowcasting variant; adds GLM channel fetch/stack, removes GFS_FX conditioning, fixes output filename and plot coordinates.
pyproject.toml	Bumps nvidia-physicsnemo minimum from 2.0 to 2.1.0 for the stormscope extra.
earth2studio/data/init.py	Exports GOESGLMGrid alongside existing GOESGLM.

_{Reviews (3): Last reviewed commit: "nits" | Re-trigger Greptile}

[pzharrington]

@greptileai

[pzharrington]

@greptileai

[pzharrington]

/blossom-ci

[NickGeneva]

one last comment, approving!

[pzharrington]

/blossom-ci

[pzharrington]

/blossom-ci