credit.datasets.gen_2.mrms

MRMSDataset: PyTorch Dataset for MRMS data with nested input/target structure.

Sample structure returned by __getitem__:

{
    "input":    {<user_provided_name>: {"<user_provided_name>/prognostic/2d/MultiSensor_QPE_01H_Pass2_00.00": tensor,
                                        "<user_provided_name>/prognostic/2d/MultiSensor_QPE_06H_Pass2_00.00": tensor}},
    "target":   {<user_provided_name>: {"<user_provided_name>/prognostic/2d/MultiSensor_QPE_01H_Pass2_00.00": tensor,
                                        "<user_provided_name>/prognostic/2d/MultiSensor_QPE_06H_Pass2_00.00": tensor}},  # only populated when return_target=True
    "metadata": {<user_provided_name>: {"input_datetime": int, "target_datetime": int}},
}

All MRMS variables are 2D. Tensor shape (no batch dimension):

(1, 1, lat, lon)   — singleton level dim, consistent with ERA5 2D convention

After DataLoader collation the batch dimension is prepended:

(batch, 1, 1, lat, lon)

Modes:

local — load from NetCDF (.nc) or Zarr (.zarr) files on disk using

the same filename_time_format strftime convention as ERA5.

remote — stream directly from AWS S3 (noaa-mrms-pds, anonymous access)

via s3fs + pygrib.

File naming (local mode):

Controlled by the optional filename_time_format config key. Defaults to "%Y%m%d-%H%M%S" (one file per timestamp).

Examples:

filename_time_format: "%Y%m%d-%H%M%S"   # MRMS_20240601-060000.nc
filename_time_format: "%Y%m%d"           # MRMS_20240601.nc  (daily)
filename_time_format: "%Y%m"             # MRMS_202406.nc    (monthly)

If only a single file matches the glob pattern, filename_time_format is ignored and that file is used for all timestamps.

Classes

MRMSDataset

PyTorch Dataset for MRMS data with nested input/target structure.

Module Contents

class credit.datasets.gen_2.mrms.MRMSDataset(data_config: dict[str, Any], return_target: bool = False)

Bases: credit.datasets.gen_2.base_dataset.BaseDataset

PyTorch Dataset for MRMS data with nested input/target structure.

Field types follow CREDIT Gen2 conventions: prognostic variables appear in both input (at step 0) and target; dynamic_forcing appears in input at every step; diagnostic appears in target only. At step i > 0 the model’s own prognostic predictions are fed back — no disk read occurs for prognostic fields at those steps.

Supports loading directly from AWS S3 (remote mode) or from local NetCDF / Zarr files (local mode). Spatial subsetting via extent is applied at load time on the native MRMS grid.

See module docstring for full description of output format and file naming.

Example YAML configuration (local mode):

data:
  source:
    Example_MRMS:  # User-provided name (arbitrary key)
      dataset_type: "mrms"
      mode: "local"
      variables:
        prognostic:                         # input at step 0 + target
          vars_2D:
            - "MultiSensor_QPE_01H_Pass2_00.00"
          path: "/data/MRMS_*.nc"
          filename_time_format: "%Y%m%d-%H%M%S"
        dynamic_forcing:                    # input every step
          vars_2D:
            - "MultiSensor_QPE_06H_Pass2_00.00"
          path: "/data/MRMS_*.nc"
          filename_time_format: "%Y%m%d-%H%M%S"
      extent: [-130, -60, 20, 55]   # [min_lon, max_lon, min_lat, max_lat]

  start_datetime: "2024-06-01"
  end_datetime:   "2024-07-01"
  timestep:       "6h"
  forecast_len:   0

Example YAML configuration (remote mode):

data:
  source:
    Example_MRMS:  # User-provided name (arbitrary key)
      dataset_type: "mrms"
      mode: "remote"
      region: "CONUS"
      variables:
        prognostic:
          vars_2D:
            - "MultiSensor_QPE_01H_Pass2_00.00"
      extent: [-130, -60, 20, 55]

Assumptions:

1. Local files have time, lat, lon dimensions/coordinates.

2. Longitude coordinates are in the 0–360 convention (both local and remote).

3. extent is specified as [min_lon, max_lon, min_lat, max_lat] in either -180-180 or 0-360 format; it is normalised to 0-360 internally.

dataset_type: str = 'mrms'

region: str

extent: list[float] | None

static_metadata: dict