earth2studio.models.px.StormScopeMRMS#
- class earth2studio.models.px.StormScopeMRMS(model_spec, means, stds, latitudes, longitudes, variables=array(['refc'], dtype='<U4'), conditioning_variables=array(['abi01c', 'abi02c', 'abi03c', 'abi07c', 'abi08c', 'abi09c', 'abi10c', 'abi13c'], dtype='<U6'), conditioning_means=None, conditioning_stds=None, conditioning_data_source=None, sampler_args={'S_churn': 10, 'num_steps': 100}, y_coords=None, x_coords=None, input_times=array([0], dtype='timedelta64[h]'), output_times=array([1], dtype='timedelta64[h]'), input_interp_max_dist_km=12.0, conditioning_interp_max_dist_km=12.0)[source]#
StormScope model forecasting MRMS data on the HRRR grid.
- This model supports multiple variants at different temporal resolutions:
6km resolution, 60 minute timestep
6km resolution, 10 minute timestep
Selection between these can be made by passing the
model_name argumentto this class’sload_modelmethod.The 6km/10min model uses a sliding window of 6 input timesteps and predicts one output timestep; other models use a single input timestep and predict one output timestep. All StormScopeMRMS models by default expect GOES-East data as conditioning; typically in a forecasting run this can be provided by passing the predictions from a StormScopeGOES model to this model’s
call_with_conditioningmethod. Otherwise, the user must provide a conditioning data source for the model to use during inference.- Parameters:
model_spec (list[dict[str, Any]]) – Sequence of stage specifications; see StormScopeBase.
means (torch.Tensor) – Per-variable mean for normalization, shape [1, C, 1, 1].
stds (torch.Tensor) – Per-variable std for normalization, shape [1, C, 1, 1].
latitudes (torch.Tensor) – Latitudes of the grid, expected shape [H, W].
longitudes (torch.Tensor) – Longitudes of the grid, expected shape [H, W].
variables (np.ndarray, optional) – MRMS input variables. Default is [“refc”].
conditioning_variables (np.ndarray, optional) – Auxiliary conditioning variables (typically GOES channels). Default is [“abi01c”, “abi02c”, “abi03c”, “abi07c”, “abi08c”, “abi09c”, “abi10c”, “abi13c”].
conditioning_means (torch.Tensor | None, optional) – Means to normalize any external conditioning data. Default is None.
conditioning_stds (torch.Tensor | None, optional) – Stds to normalize any external conditioning data. Default is None.
conditioning_data_source (Any | None, optional) – Data source for external conditioning. Default is None.
sampler_args (dict[str, float | int] | None, optional) – Default sampler arguments passed to the diffusion sampler. Default is {“num_steps”: 100, “S_churn”: 10}.
y_coords (np.ndarray | None, optional) – 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 (np.ndarray | None, optional) – 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_times (np.ndarray, optional) – Input timesteps, of type timedelta64. Default is [0 h] (i.e., the current time).
output_times (np.ndarray, optional) – Output timesteps, of type timedelta64. Default is [1 h] (i.e., 1 hour from the current time).
input_interp_max_dist_km (float, optional) – 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, optional) – 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 MRMS data, so these portions are masked as invalid (set to NaN).
- __call__(x, coords)#
Runs the prognostic model one step. Assumes the last two dimensions of the input tensor are the spatial dimensions.
- Parameters:
x (torch.Tensor) – Input tensor.
coords (CoordSystem) – Input coordinate system.
- Returns:
Output tensor and coordinate system.
- Return type:
tuple[torch.Tensor, CoordSystem]
- create_iterator(x, coords)#
Creates an iterator to perform time-integration of the prognostic model.
- Parameters:
x (torch.Tensor) – Input tensor.
coords (CoordSystem) – Input coordinate system.
- Yields:
Iterator[tuple[torch.Tensor, CoordSystem]] – Iterator that generates time-steps of the prognostic model containing the output data tensor and coordinate system dictionary.
- Return type:
Iterator[tuple[Tensor, OrderedDict[str, ndarray]]]
- classmethod load_default_package()#
Load a default local package for StormScope models.
- Return type:
- classmethod load_model(package, model_name='6km_60min_natten_cos_zenith_input_mrms_eoe', conditioning_data_source=None)[source]#
Load model from package.
- Parameters:
package (Package) – Package to load model from
model_name (str, optional) –
- Model name to load; allows for selection between different variants of the model:
”6km_60min_natten_cos_zenith_input_mrms_eoe”: 6km resolution, 60 minute timestep
”6km_10min_natten_pure_obs_mrms_obs_6steps”: 6km resolution, 10 minute timestep
conditioning_data_source (DataSource | ForecastSource | None, optional) – Data source to use for conditioning, by default None.
- Returns:
Instantiated StormScopeMRMS model
- Return type:
PrognosticModel