StormScopeMRMS

NANWC202680 GBNVIDIAPyTorch

Import path: earth2studio.models.px.StormScopeMRMS

View source on GitHub View install commands

Documentation

Bases: StormScopeBase

StormScope model forecasting MRMS data on the HRRR grid.

This model supports multiple variants at different temporal resolutions, selected by passing model_name to load_model (default: "3km_10min"). Variant names are semantic (<resolution>_<cadence>):

• 3km_10min: 3km resolution, 10 minute timestep, MRMS+GLM nowcasting (default)
• 6km_1hr: 6km resolution, 60 minute timestep (legacy nearcasting)

Use list_available_models to inspect the variants in a given package. Legacy training-style names are still accepted as aliases.

Variants whose input cadence is finer than their output cadence use a sliding window of input timesteps and predict one output timestep; others use a single input timestep and predict one output timestep. All StormScopeMRMS models by default expect GOES-East data as conditioning; typically in a forecasting run this can be provided by passing the predictions from a StormScopeGOES model to this model's call_with_conditioning method. Otherwise, the user must provide a conditioning data source for the model to use during inference.

Parameters:

• model_spec (list[dict[str, Any]]) –

  Sequence of stage specifications; see StormScopeBase.

• means (Tensor) –

  Per-variable mean for normalization, shape [1, C, 1, 1].

• stds (Tensor) –

  Per-variable std for normalization, shape [1, C, 1, 1].

• latitudes (Tensor) –

  Latitudes of the grid, expected shape [H, W].

• longitudes (Tensor) –

  Longitudes of the grid, expected shape [H, W].

• variables (ndarray, default: array(['refc']) ) –

  MRMS input variables. Default is ["refc"].

• conditioning_variables (ndarray, default: array(['abi01c', 'abi02c', 'abi03c', 'abi07c', 'abi08c', 'abi09c', 'abi10c', 'abi13c']) ) –

  Auxiliary conditioning variables (typically GOES channels). Default is ["abi01c", "abi02c", "abi03c", "abi07c", "abi08c", "abi09c", "abi10c", "abi13c"].

• conditioning_means (Tensor | None, default: None ) –

  Means to normalize any external conditioning data. Default is None.

• conditioning_stds (Tensor | None, default: None ) –

  Stds to normalize any external conditioning data. Default is None.

• conditioning_data_source (Any | None, default: None ) –

  Data source for external conditioning. Default is None.

• sampler_args (dict[str, float | int] | None, default: {'num_steps': 100, 'S_churn': 10} ) –

  Default sampler arguments passed to the diffusion sampler. Default is {"num_steps": 100, "S_churn": 10}.

• y_coords (ndarray | None, default: None ) –

  Y coordinates of the grid, expected shape [H, W]. Default is None, in which case the model uses the enumerated indices inferred from the latitude and longitude grid shapes.

• x_coords (ndarray | None, default: None ) –

  X coordinates of the grid, expected shape [H, W]. Default is None, in which case the model uses the enumerated indices inferred from the latitude and longitude grid shapes.

• input_times (ndarray, default: array([timedelta64(0, 'h')]) ) –

  Input timesteps, of type timedelta64. Default is [0 h] (i.e., the current time).

• output_times (ndarray, default: array([timedelta64(1, 'h')]) ) –

  Output timesteps, of type timedelta64. Default is [1 h] (i.e., 1 hour from the current time).

• input_interp_max_dist_km (float, default: 12.0 ) –

  Maximum distance in kilometers for nearest neighbor interpolation of input data. Points beyond this distance are masked as invalid. Default is 12.0.

• conditioning_interp_max_dist_km (float, default: 12.0 ) –

  Maximum distance in kilometers for nearest neighbor interpolation of conditioning data. Points beyond this distance are masked as invalid. Default is 26.0.

• mrms_coverage_mask (Tensor | None, default: None ) –

  Boolean NEXRAD-coverage mask of shape [H, W] on the model grid, True where MRMS data is considered valid (inside NEXRAD circular coverage). When provided, it is used as the initial valid_mask and is ANDed with any interpolator-derived mask built by build_input_interpolator. Loaded automatically from the package for non-deprecated variants. Default is None.

• glm_data_source (DataSource | None, default: None ) –

  Gridded GLM source (e.g. earth2studio.data.GOESGLMGrid) for variants with a glm_density state channel (3km_10min only). When set, __call__ (and create_iterator) fetch, regrid, and inject GLM into the state automatically on every step. Not used by the coupled path (call_with_conditioning), where the caller is responsible for populating GLM channels. Default is None.

Note

To have a unified coordinate system over CONUS for convenience, the model uses the HRRR grid. As a result, there are portions of the domain which go beyond the extent of the MRMS data, so these portions are masked as invalid (set to NaN).

Note

GLM state channel. The 3km_10min variant includes a glm_density channel (gridded GLM lightning counts, normalized with log1p) as part of its state — both an input observation and a predicted output (the 6km_1hr variant has no GLM channel). Because the GLM source lives on a different native grid from MRMS, glm_density is handled separately from the radar channels and is the GLM analogue of the GOES conditioning:

• Auto path (__call__ / create_iterator): pass glm_data_source (e.g. earth2studio.data.GOESGLMGrid) to load_model and GLM is fetched, bilinearly regridded, and injected into the state automatically on every step — exactly as conditioning_data_source is fetched via fetch_conditioning. The GLM bilinear interpolator is built lazily on the first call. The input state x only needs its radar channels populated (the GLM channels are overwritten); a zero placeholder is fine. In this case, the model will be using ground-truth GLM observations during the rollout, so is not doing pure forecasting (and can only be run for dates in the past where the full timeseries of GLM observations is available).

• Coupled path (call_with_conditioning): just as this method takes conditioning from the caller rather than the data source, it leaves the entire state — GLM included — to the caller and never touches glm_data_source. Populate the GLM channels of x yourself (e.g. via fetch_glm for the initial state); during the rollout GLM then flows autoregressively from the model's own predictions, like the radar channels. This is the more typical pure-forecast use case.

__call__

__call__(
    x: Tensor, coords: CoordSystem
) -> tuple[Tensor, CoordSystem]

Runs the prognostic model one step. Assumes the last two dimensions of the input tensor are the spatial dimensions.

Parameters:

• x (Tensor) –

  Input tensor.

• coords (CoordSystem) –

  Input coordinate system.

Returns:

• tuple[Tensor, CoordSystem] –

  Output tensor and coordinate system.

create_iterator

create_iterator(
    x: Tensor, coords: CoordSystem
) -> Iterator[tuple[Tensor, CoordSystem]]

Creates an iterator to perform time-integration of the prognostic model.

Parameters:

• x (Tensor) –

  Input tensor.

• coords (CoordSystem) –

  Input coordinate system.

Yields:

• Iterator[tuple[Tensor, CoordSystem]] –

  Iterator that generates time-steps of the prognostic model containing the output data tensor and coordinate system dictionary.

load_default_package classmethod

load_default_package() -> Package

Load the default StormScope package from Hugging Face.

load_model classmethod

load_model(
    package: Package,
    model_name: Literal[
        "3km_10min", "6km_1hr"
    ] = "3km_10min",
    conditioning_data_source: (
        DataSource | ForecastSource | None
    ) = None,
    glm_data_source: DataSource | None = None,
    amp: bool = True,
    compile: bool = False,
) -> PrognosticModel

Load model from package.

Parameters:

• package (Package) –

  Package to load model from

• model_name (Literal['3km_10min', '6km_1hr'], default: '3km_10min' ) –

  Variant to load. Available variants (see list_available_models):

  • "3km_10min": 3km resolution, 10 minute timestep, MRMS+GLM nowcasting
  • "6km_1hr": 6km resolution, 60 minute timestep, MRMS+GLM nearcasting

  Legacy training-style names are accepted as aliases. Default is "3km_10min".

• conditioning_data_source (DataSource | ForecastSource | None, default: None ) –

  Data source to use for conditioning (GOES), by default None.

• glm_data_source (DataSource | None, default: None ) –

  Gridded GLM source (e.g. earth2studio.data.GOESGLMGrid) used for variants with a glm_density state channel (3km_10min only — the 6km_1hr variant has no GLM channel). The GLM analogue of conditioning_data_source: when set, __call__ (and create_iterator) fetch, regrid, and inject GLM into the state automatically. The coupled path (call_with_conditioning) does not use it — there the caller populates the GLM channels of x (e.g. via fetch_glm). By default None.

• amp (bool, default: True ) –

  Enable automatic mixed precision (autocast) for the sampler's network forward passes. Default is True.

• compile (bool, default: False ) –

  Compile each staged expert with torch.compile ("reduce-overhead"). Default is False.

Returns:

• PrognosticModel –

  Instantiated StormScopeMRMS model

Examples using earth2studio.models.px.StormScopeMRMS

StormScope Satellite and Radar NowcastingStormScope inference workflow with GOES satellite imagery and MRMS radar data.