Add Solar Model to StormScope

[albertocarpentieri]

Earth2Studio Pull Request

Description

Adds StormScopeNSRDB, a StormScope-family diagnostic that estimates surface
Global Horizontal Irradiance (GHI) from StormScope GOES satellite imagery on the
3-km HRRR grid. It is a same-time estimator (given GOES at time T it estimates
GHI at T) and is designed to chain after StormScopeGOES to produce solar
irradiance nowcasts over the Continental US.
The model supports three inference modes:

• Pure diffusion — sample GHI entirely from noise (default when no regression
  model is present).
• Regression only — deterministic prediction from the regression model.
• Regression + SDEdit — warm-start diffusion from the regression prediction by
  injecting noise at sdedit_sigma and denoising from there.
  It also supports clearness-index normalization (state_norm_mode="clearness_index"),
  where GHI is normalized by insolation (max(cos_zenith, 0) * solar_constant) rather
  than z-score, using the physicsnemo 1995 orbital insolation model to match training.

What's included

• earth2studio/models/px/stormscope.py — new StormScopeNSRDB class
  (sharing StormScopeBase), including estimate_from_goes(...) as the intended
  entry point when chaining StormScopeGOES -> StormScopeNSRDB, plus regression +
  SDEdit support and clearness-index normalization.
• earth2studio/models/nn/stormscope_util.py — MaskedModel denoiser wrapper
  that zeros outputs at invalid (off-grid / ocean) pixels after every sampling step
  while preserving the round_sigma interface used by the EDM sampler.
• earth2studio/models/px/__init__.py — export StormScopeNSRDB.
• examples/04_nowcasting/04_stormscope_solar_example.py — minimal example
  chaining StormScopeGOES -> StormScopeNSRDB to nowcast GHI from GOES-16 imagery
  and plot the result.
• test/models/px/test_stormscope.py — unit tests for StormScopeNSRDB
  (estimate_from_goes output shape/coords/masking and output_coords handshakes),
  skipped when physicsnemo.utils.insolation is unavailable.
• docs/modules/models_px.rst + CHANGELOG.md — documentation and changelog
  entries.
• earth2studio/statistics/fss.py — minor unrelated fix: cast
  np.prod(...) to int before Tensor.view(...) (happy to drop if preferred).

Notes / follow-ups

• StormScopeNSRDB currently has no load_default_package(); weights are loaded
  via an explicit Package(...). A default public Hugging Face package can be wired
  up once the official weights repo is available.

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

No new dependencies. StormScopeNSRDB relies on packages already required by the
stormscope extra (e.g. physicsnemo for the insolation model, earth2grid,
scipy).

[greptile-apps]

Greptile Summary

This PR adds StormScopeNSRDB, a new StormScopeBase subclass that estimates surface GHI from StormScope GOES satellite imagery using a regression-then-SDEdit warm-start diffusion pipeline on the 3-km HRRR grid. It also adds a MaskedModel denoiser wrapper to zero invalid pixels after each sampling step, extends the existing GOES example to chain StormScopeGOES → StormScopeNSRDB, and includes a minor fss.py fix for NumPy integer type handling.

• StormScopeNSRDB normalizes GHI via clearness index (insolation denominator), runs a single fixed regression + warm-start diffusion step, and uses physicsnemo's 1995 orbital insolation model to match training; the estimate_from_goes entry point builds a zero state and delegates to call_with_conditioning.
• Test suite (test_stormscope.py): the new create_spoof_nsrdb_model helper passes diffusion_model= and sampler_args= to StormScopeNSRDB.__init__, which accepts neither, and omits the required model_spec and regression_model positional arguments — tests will raise TypeError in environments where physicsnemo.utils.insolation is installed.
• fss.py fix (cast np.prod to int) is correct and uncontroversial.

Confidence Score: 4/5

Safe to merge once the test helper constructor call is corrected; the model implementation itself is sound.

The StormScopeNSRDB model logic, clearness-index normalization, MaskedModel wrapper, and SDEdit warm-start are all consistent with the described design. The only concrete defect is in the test file: create_spoof_nsrdb_model uses parameter names that don't exist in StormScopeNSRDB.__init__ and omits required arguments — every test guarded by @_needs_insolation will fail with a TypeError in any environment with physicsnemo.utils.insolation installed.

test/models/px/test_stormscope.py — the create_spoof_nsrdb_model helper needs its constructor call updated to match StormScopeNSRDB.__init__.

Important Files Changed

Filename	Overview
earth2studio/models/px/stormscope.py	Adds StormScopeNSRDB and MaskedModel; core logic appears sound (SDEdit warm-start, clearness-index de-normalization, valid-pixel masking) with no blocking logic errors found.
test/models/px/test_stormscope.py	create_spoof_nsrdb_model passes unknown kwargs (diffusion_model, sampler_args) and omits required params (model_spec, regression_model); both new NSRDB tests will raise TypeError when insolation is available.
earth2studio/statistics/fss.py	Minor fix: casts np.prod(...) result to int before passing to Tensor.view() to avoid dtype errors in newer NumPy.
examples/04_nowcasting/03_stormscope_goes_example.py	Extends the GOES example to load StormScopeNSRDB, call estimate_from_goes in the rollout loop, and plot the resulting GHI field; straightforward additions.
earth2studio/models/px/init.py	Adds StormScopeNSRDB to the public exports.
docs/modules/models_px.rst	Adds StormScopeNSRDB to the model documentation index.
CHANGELOG.md	Adds changelog entry for StormScopeNSRDB.

Comments Outside Diff (1)

1. test/models/px/test_stormscope.py, line 661-677 (link)

   Test constructor arguments don't match StormScopeNSRDB.__init__

   create_spoof_nsrdb_model passes diffusion_model=diffusion and sampler_args={"num_steps": 2}, but StormScopeNSRDB.__init__ accepts neither parameter. It also omits the two required positional parameters model_spec (a list of stage dicts) and regression_model (the regression network). Any test environment that has physicsnemo.utils.insolation installed will fail immediately with TypeError: __init__() got an unexpected keyword argument 'diffusion_model' — the @_needs_insolation guard only skips in environments without that package.

   To fix, replace diffusion_model=diffusion with a proper model_spec=[{"model": diffusion, "sigma_min": ..., "sigma_max": ...}], add a regression_model=<spoof regression model>, and drop sampler_args (NSRDB hardcodes {"num_steps": self._NUM_STEPS} internally).

_{Reviews (2): Last reviewed commit: "Merge branch 'NVIDIA:main' into acarpent..." | Re-trigger Greptile}

[pzharrington]

I think a separate file / wrapper for the NSRDB model might be better than adding it into this one, but maybe update the branch first and see how it looks with the newer stormscope base class/models. If in a separate file, can still import and subclass the base stormscope wrapper class if it makes sense and if a lot of the code would get reused.

[albertocarpentieri]

I updated the PR, I believe that since this model works only on stormscope goes it is appropriate here, but also open to move it to its own py script

[CharlelieLrt]

I don't think we want an example for this new wrapper. @NickGeneva to conform, but I think we've had this discussion before for another wrapper and the conclusion was "no example". However the Pr discussion section should have a code snippet + images/figures to help validate the model wrapper.

[albertocarpentieri]

I added the nsrdb model in the present example and removed the new one

[CharlelieLrt]

I am just unsure if this is okay to do that. Please check with @NickGeneva

[CharlelieLrt]

This is using none of the physicsnemo diffusion module, which I am not approving at all

[albertocarpentieri]

This is more a global dependency for the full stormscope model, maybe for a future update of the model? I do not think it is part of this PR

[CharlelieLrt]

Left a few comments. Main things:

• Relies entirely on legacy APIs from physicnemo, that's a clear no-go for me
• Over-engineered
• No need for an example

My suggestion it's to create a wrapper class from scratch (don't inherit the StormScope one), because I am in any case not sure how much code you are really saving by composing the classes here. Also, keep the parameters from the model package that are absolutely necessary. You should have an idea of what works best, so we should stick to that in this wrapper. This will help simplify a lot of things.

I'll review again after these changes are implemented.

[CharlelieLrt]

Are you initializing the latent x_N with a noised prediction from the regression model? Does it work better than starting from plain noise? Just trying to understand the rationale for having a separate regresion model

[albertocarpentieri]

Yes correct, i tested using the diffusion model only but it gave worse results. The rationale is that the regression model can make you skip the high sigma diffusion steps, running the diffusion model only for the final part. It's an effort to make the model faster and temporally coherent.