`StormCastCONUS`¶

NANWC202624 GBNVIDIAPyTorch

Import path: earth2studio.models.px.StormCastCONUS

View source on GitHub View install commands

Documentation¶

Bases: Module, AutoModelMixin, PrognosticMixin

StormCast-CONUS generative convection-allowing model for the full CONUS domain.

High-resolution (3km) HRRR state over the Continental 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:

Parameters:

diffusion_model (Module) –

Configured diffusion model (e.g. a _SplitModelWrapper instance created by load_model). Must be a _SplitModelWrapper instance to set hrrr_lat_lim or hrrr_lon_lim to non-default values.
means (Tensor) –

Per-channel mean for normalising the high-resolution state.
stds (Tensor) –

Per-channel standard deviation for normalising the high-resolution state.
invariants (Tensor) –

Static invariant fields (e.g. land-sea mask, orography).
conditioning_means (Tensor) –

Per-channel mean for normalising the low-resolution conditioning.
conditioning_stds (Tensor) –

Per-channel standard deviation for normalising the low-resolution conditioning.
hrrr_lat_lim (tuple[int, int], default: (17, 1041) ) –

HRRR grid latitude limits, defaults to be the StormCastCONUS region in Continental United States, by default (17, 1041)
hrrr_lon_lim (tuple[int, int], default: (3, 1795) ) –

HRRR grid longitude limits, defaults to be the StormCastCONUS region in Continental United States, by default (3, 1795)
variables (ndarray, default: array(VARIABLES) ) –

High-resolution variable names, by default np.array(VARIABLES).
conditioning_variables (ndarray, default: array(CONDITIONING_VARIABLES) ) –

Low-resolution conditioning variable names, by default np.array(CONDITIONING_VARIABLES).
conditioning_data_source (DataSource or ForecastSource or None, default: None ) –

Data source for global conditioning. Required for inference, by default None.
conditioning_init_time (datetime, list[datetime], TimeArray, or None, default: None ) –

Fixed initialisation time for conditioning_data_source. When None, uses the same conditioning initialization time as the forecast. Useful when the conditioning is from a forecast source available on a coarser cycle (e.g. GFS_FX every 6 h while HRRR initialises hourly), by default None.
sampler_args (dict, default: None ) –

Overrides for the EDM sampler/scheduler. Recognised keys: sigma_min, sigma_max, rho (scheduler), and S_churn, S_min, S_max, S_noise (solver). Unspecified keys use sensible defaults.
num_diffusion_steps (int, default: 18 ) –

Number of diffusion sampling steps for the EDM (no-obs) path, by default 18.
num_sda_diffusion_steps (int, default: 36 ) –

Number of diffusion sampling steps for the SDA (obs-guided) path, by default 36.
batch_size (int, default: 1 ) –

Maximum batch size processed in one forward pass, by default 1.
time_tolerance (np.timedelta64, tuple[np.timedelta64, np.timedelta64], or None, default: None ) –

Time window for filtering observations around each target time. A single np.timedelta64 creates a symmetric window; a 2-tuple (lower, upper) is passed directly to ObsGridMapping.obs_to_grid. None disables time filtering, by default None.
sda_std_obs (float or dict[str, float], default: 0.1 ) –

Observation noise standard deviation (in physical units) used by the SDA denoiser when observations are provided. A scalar applies uniformly; a dict maps variable names to per-variable values (default 0.1 for unlisted variables). Dict values are converted to normalised units internally, by default 0.1.
sda_dps_norm (float, default: 2 ) –

Gradient normalisation factor for DPS guidance, by default 2.
sda_gamma (float, default: 0.001 ) –

DPS guidance step size / scale, by default 0.001.
use_amp (bool, default: True ) –

Whether to run the diffusion forward pass under torch.autocast with bfloat16, by default True.
clamp_values (bool, default: True ) –

Whether to apply reflectivity clipping in _forward. When the model is loaded via load_model, this flag is also forwarded to _SplitModelWrapper to enable per-variable physical-minimum clamping, by default True.

call ¶

__call__(
    x: Tensor,
    coords: CoordSystem,
    obs: DataFrame | tuple[Tensor, Tensor] | None = None,
) -> tuple[Tensor, CoordSystem]

Runs prognostic model 1 step

Parameters:

x (Tensor) –

Input tensor
coords (CoordSystem) –

Input coordinate system
obs (pd.DataFrame, tuple[torch.Tensor, torch.Tensor], or None, default: None ) –

Observations for SDA guidance. Either a dataframe with columns variable, lat, lon, observation, time, or a pre-gridded (y_obs, mask) tuple. None runs unconditional diffusion.

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]]

Iterator wrapper around create_generator without observation input.

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(),
    **model_kwargs: Any
) -> StormCastCONUS

Load a StormCastCONUS model from a 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
**model_kwargs (Any, default: {} ) –

Additional keyword arguments forwarded to the model constructor (e.g. hrrr_lat_lim, num_diffusion_steps, use_amp, clamp_values). clamp_values is additionally applied to the internal _SplitModelWrapper.