StormCast

NANWC202440 GBNVIDIAPyTorch

Import path: earth2studio.models.px.StormCast

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin, PrognosticMixin

StormCast generative convection-allowing model for regional forecasts consists of two core models: a regression and diffusion model. Model time step size is 1 hour, taking as input:

• High-resolution (3km) HRRR state over the central United States (99 vars)
• High-resolution land-sea mask and orography invariants
• Coarse resolution (25km) global state (26 vars)

The high-resolution grid is the HRRR Lambert conformal projection Coarse-resolution inputs are regridded to the HRRR grid internally.

Note

For more information see the following references:

• arxiv.org/abs/2408.10958
• huggingface.co/nvidia/stormcast-v1-era5-hrrr

Parameters:

• regression_model (Module) –

  Deterministic model used to make an initial prediction

• diffusion_model (Module) –

  Generative model correcting the deterministic prediciton

• means (Tensor) –

  Mean value of each input high-resolution variable

• stds (Tensor) –

  Standard deviation of each input high-resolution variable

• invariants (Tensor) –

  Static invariant quantities

• hrrr_lat_lim (tuple[int, int], default: (273, 785) ) –

  HRRR grid latitude limits, defaults to be the StormCastV1 region in central United States, by default (273, 785)

• hrrr_lon_lim (tuple[int, int], default: (579, 1219) ) –

  HRRR grid longitude limits, defaults to be the StormCastV1 region in central United States,, by default (579, 1219)

• variables (array, default: array(VARIABLES) ) –

  High-resolution variables, by default np.array(VARIABLES)

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

  Means to normalize conditioning data, by default None

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

  Standard deviations to normalize conditioning data, by default None

• conditioning_variables (array, default: array(CONDITIONING_VARIABLES) ) –

  Global variables for conditioning, by default np.array(CONDITIONING_VARIABLES)

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

  Data Source to use for global conditioning. Required for running in iterator mode, by default None

• sampler_steps (int, default: 18 ) –

  Number of diffusion sampler steps, by default 36

• sampler_args (dict[str, float | int], default: None ) –

  Arguments to pass to the diffusion sampler, by default None

__call__

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

Runs prognostic model 1 step

Parameters:

• x (Tensor) –

  Input tensor

• coords (CoordSystem) –

  Input coordinate system

Returns:

• tuple[Tensor, CoordSystem] –

  Output tensor and coordinate system

Raises:

• RuntimeError –

  If conditioning data source is not initialized

create_iterator

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

Creates a iterator which can be used to perform time-integration of the prognostic model. Will return the initial condition first (0th step).

Parameters:

• x (Tensor) –

  Input tensor

• coords (CoordSystem) –

  Input coordinate system

Yields:

• Iterator[tuple[Tensor, CoordSystem]] –

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

load_default_package classmethod

load_default_package() -> Package

Load prognostic package

load_model classmethod

load_model(
    package: Package,
    conditioning_data_source: (
        DataSource | ForecastSource
    ) = GFS_FX(),
    sampler_steps: int = 18,
) -> PrognosticModel

Load prognostic from package

Parameters:

• package (Package) –

  Package to load model from

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

  Data source to use for global conditioning, by default GFS_FX

• sampler_steps (int, default: 18 ) –

  Number of diffusion sampler steps, by default 18

Returns:

• PrognosticModel –

  Prognostic model