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 argument to this class’s load_model method.

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_conditioning method. 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:

Package

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

Examples using earth2studio.models.px.StormScopeMRMS#

StormScope Satellite and Radar Nowcasting

StormScope Satellite and Radar Nowcasting