StormCastCONUS#

class earth2studio.models.px.StormCastCONUS( diffusion_model, means, stds, invariants, conditioning_means, conditioning_stds, hrrr_lat_lim=(17, 1041), hrrr_lon_lim=(3, 1795), variables=array(['u10m', 'v10m', 't2m', 'msl', 'u1hl', 'u2hl', 'u3hl', 'u4hl', 'u5hl', 'u6hl', 'u7hl', 'u8hl', 'u9hl', 'u10hl', 'u11hl', 'u13hl', 'u15hl', 'u20hl', 'u25hl', 'u30hl', 'v1hl', 'v2hl', 'v3hl', 'v4hl', 'v5hl', 'v6hl', 'v7hl', 'v8hl', 'v9hl', 'v10hl', 'v11hl', 'v13hl', 'v15hl', 'v20hl', 'v25hl', 'v30hl', 't1hl', 't2hl', 't3hl', 't4hl', 't5hl', 't6hl', 't7hl', 't8hl', 't9hl', 't10hl', 't11hl', 't13hl', 't15hl', 't20hl', 't25hl', 't30hl', 'q1hl', 'q2hl', 'q3hl', 'q4hl', 'q5hl', 'q6hl', 'q7hl', 'q8hl', 'q9hl', 'q10hl', 'q11hl', 'q13hl', 'q15hl', 'q20hl', 'q25hl', 'q30hl', 'Z1hl', 'Z2hl', 'Z3hl', 'Z4hl', 'Z5hl', 'Z6hl', 'Z7hl', 'Z8hl', 'Z9hl', 'Z10hl', 'Z11hl', 'Z13hl', 'Z15hl', 'Z20hl', 'Z25hl', 'Z30hl', 'p1hl', 'p2hl', 'p3hl', 'p4hl', 'p5hl', 'p6hl', 'p7hl', 'p8hl', 'p9hl', 'p10hl', 'p11hl', 'p13hl', 'p15hl', 'p20hl', 'refc'], dtype='<U5'), conditioning_variables=array(['u10m', 'v10m', 't2m', 'tcwv', 'sp', 'msl', 'u1000', 'u850', 'u500', 'u250', 'v1000', 'v850', 'v500', 'v250', 'z1000', 'z850', 'z500', 'z250', 't1000', 't850', 't500', 't250', 'q1000', 'q850', 'q500', 'q250'], dtype='<U5'), conditioning_data_source=None, sampler_args=None, num_diffusion_steps=18, num_sda_diffusion_steps=36, batch_size=1, time_tolerance=None, sda_std_obs=0.1, sda_dps_norm=2, sda_gamma=0.001, use_amp=True, clamp_values=True, )[source]#

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 (torch.nn.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 (torch.Tensor) – Per-channel mean for normalising the high-resolution state.
stds (torch.Tensor) – Per-channel standard deviation for normalising the high-resolution state.
invariants (torch.Tensor) – Static invariant fields (e.g. land-sea mask, orography).
conditioning_means (torch.Tensor) – Per-channel mean for normalising the low-resolution conditioning.
conditioning_stds (torch.Tensor) – Per-channel standard deviation for normalising the low-resolution conditioning.
hrrr_lat_lim (tuple[int, int], optional) – HRRR grid latitude limits, defaults to be the StormCastCONUS region in Continental United States, by default (17, 1041)
hrrr_lon_lim (tuple[int, int], optional) – HRRR grid longitude limits, defaults to be the StormCastCONUS region in Continental United States, by default (3, 1795)
variables (np.ndarray, optional) – High-resolution variable names, by default np.array(VARIABLES).
conditioning_variables (np.ndarray, optional) – Low-resolution conditioning variable names, by default np.array(CONDITIONING_VARIABLES).
conditioning_data_source (DataSource or ForecastSource or None, optional) – Data source for global conditioning. Required for inference, by default None.
sampler_args (dict, optional) – 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, optional) – Number of diffusion sampling steps for the EDM (no-obs) path, by default 18.
num_sda_diffusion_steps (int, optional) – Number of diffusion sampling steps for the SDA (obs-guided) path, by default 36.
batch_size (int, optional) – Maximum batch size processed in one forward pass, by default 1.
time_tolerance (np.timedelta64, tuple[np.timedelta64, np.timedelta64], or None, optional) – 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], optional) – 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, optional) – Gradient normalisation factor for DPS guidance, by default 2.
sda_gamma (float, optional) – DPS guidance step size / scale, by default 0.001.
use_amp (bool, optional) – Whether to run the diffusion forward pass under torch.autocast with bfloat16, by default True.
clamp_values (bool, optional) – 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__(x, coords, obs=None)[source]#

Runs prognostic model 1 step

Parameters:

x (torch.Tensor) – Input tensor
coords (CoordSystem) – Input coordinate system
obs (pd.DataFrame, tuple[torch.Tensor, torch.Tensor], or None, optional) – 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:

Output tensor and coordinate system

Return type:

tuple[torch.Tensor, CoordSystem]

Raises:

RuntimeError – If conditioning data source is not initialized

create_iterator(x, coords)[source]#

Iterator wrapper around create_generator() without observation input.

Parameters:

x (Tensor)
coords (OrderedDict[str, ndarray])

Return type:

Iterator[tuple[Tensor, OrderedDict[str, ndarray]]]

classmethod load_default_package()[source]#

Load prognostic package

Return type:: Package

classmethod load_model(

package,

conditioning_data_source=<earth2studio.data.gfs.GFS_FX object>,

**model_kwargs,

)[source]#

Load a StormCastCONUS model from a package.

Parameters:

package (Package) – Package to load model from
conditioning_data_source (DataSource | ForecastSource, optional) – Data source to use for global conditioning, by default GFS_FX
**model_kwargs – 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.

Returns: