earth2studio.models.px.StormCast#
- class earth2studio.models.px.StormCast(regression_model, diffusion_model, lat, lon, means, stds, invariants, variables=array(['u10m', 'v10m', 't2m', 'mslp', '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_means=None, conditioning_stds=None, conditioning_variables=array(['u10m', 'v10m', 't2m', 'tcwv', 'mslp', 'sp', '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={}, interp_method='linear')[source]#
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.
- Parameters:
regression_model (torch.nn.Module) – Deterministic model used to make an initial prediction
diffusion_model (torch.nn.Module) – Generative model correcting the deterministic prediciton
lat (np.array) – Latitude array (2D) of the domain
lon (np.array) – Longitude array (2D) of the domain
means (torch.Tensor) – Mean value of each input high-resolution variable
stds (torch.Tensor) – Standard deviation of each input high-resolution variable
invariants (torch.Tensor) – Static invariant quantities
variables (np.array, optional) – High-resolution variables, by default np.array(VARIABLES)
conditioning_means (torch.Tensor | None, optional) – Means to normalize conditioning data, by default None
conditioning_stds (torch.Tensor | None, optional) – Standard deviations to normalize conditioning data, by default None
conditioning_variables (np.array, optional) – Global variables for conditioning, by default np.array(CONDITIONING_VARIABLES)
conditioning_data_source (DataSource | None, optional) – Data Source to use for global conditoining. Required for running in iterator mode, by default None
sampler_args (dict[str, float | int], optional) – Arguments to pass to the diffusion sampler, by default {}
interp_method (str, optional) – Interpolation method to use when regridding coarse conditoining data, by default “linear”
- __call__(x, coords)[source]#
Forward pass of diagnostic
- Parameters:
x (Tensor)
coords (OrderedDict[str, ndarray])
- Return type:
tuple[Tensor, OrderedDict[str, ndarray]]
- create_iterator(x, coords)[source]#
Creates a iterator which can be used to perform time-integration of the prognostic model. Will return the initial condition first (0th step). :param x: Input tensor :type x: torch.Tensor :param coords: Input coordinate system :type coords: CoordSystem
- Yields:
Iterator[tuple[torch.Tensor, CoordSystem]] – Iterator that generates time-steps of the prognostic model container the output data tensor and coordinate system dictionary.
- Parameters:
x (Tensor)
coords (OrderedDict[str, ndarray])
- Return type:
Iterator[tuple[Tensor, OrderedDict[str, ndarray]]]
