StormScopeGOES

NANWC202680 GBNVIDIAPyTorch

Import path: earth2studio.models.px.StormScopeGOES

View source on GitHub View install commands

Documentation

Bases: StormScopeBase

StormScope model forecasting GOES data on the HRRR grid.

This model supports multiple variants at different spatiotemporal resolutions, selected by passing model_name to load_model (default: "3km_10min"). The primary focus is CONUS nowcasting at 3km resolution; coarser 6km nearcasting variants are retained as legacy checkpoints. Variant names are semantic (<resolution>_<cadence>):

• 3km_10min : 3km resolution, 10 minute timestep (CONUS nowcasting)
• 6km_1hr : 6km resolution, 60 minute timestep (legacy nearcasting)

Use list_available_models to inspect the variants in a given package (including any added after this release). 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.

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(['abi01c', 'abi02c', 'abi03c', 'abi07c', 'abi08c', 'abi09c', 'abi10c', 'abi13c']) ) –

  GOES input variables. Default is ["abi01c", "abi02c", "abi03c", "abi07c", "abi08c", "abi09c", "abi10c", "abi13c"].

• conditioning_variables (ndarray, default: array(['z500']) ) –

  Auxiliary conditioning variables. Default is ["z500"].

• 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, Any] | None, default: {'num_steps': 100, 'S_churn': 10} ) –

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

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

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

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

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

• 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_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: 26.0 ) –

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

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 GOES-East data, so these portions are masked as invalid (set to NaN).

__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,
    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, by default "3km_10min" (the recommended CONUS nowcasting variant). Available variants (see list_available_models):

  • "3km_10min": 3km resolution, 10 minute timestep (CONUS nowcasting)
  • "6km_1hr": 6km resolution, 60 minute timestep (legacy nearcasting)

  Legacy training-style names are accepted as aliases.

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

  Data source to use for conditioning, 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 StormScopeGOES model

Examples using earth2studio.models.px.StormScopeGOES

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