HealDA

GlobalDA202640 GBNVIDIAPyTorch

Import path: earth2studio.models.da.HealDA

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin

HealDA data assimilation model for global weather analysis from sparse observations on a HEALPix grid.

HealDA is a stateless assimilation model that produces a single global weather analysis from conventional and satellite observations. It operates on a HEALPix level-6 padded XY grid and outputs ERA5-compatible atmospheric variables.

The model accepts pre-processed observation DataFrames from earth2studio.data.UFSObsConv and earth2studio.data.UFSObsSat and produces a global analysis field.

Parameters:

• model (Module) –

  The underlying HealDA neural network

• condition (Tensor) –

  Static conditioning fields (e.g. orography, land fraction) on the HEALPix grid of size [1, n_static, 1, npix]

• era5_mean (Tensor) –

  ERA5 per-channel mean for output denormalization [1, out_variables, 1, 1]

• era5_std (Tensor) –

  ERA5 per-channel std for output denormalization [1, out_variables, 1, 1]

• sensor_stats (dict[str, dict[str, ndarray]]) –

  Per-sensor normalization statistics loaded from the package

• lat_lon (bool, default: False ) –

  If True the model output is regridded from the native HEALPix grid to a regular equiangular lat-lon grid using earth2grid. If False the raw HEALPix output is returned with an npix dimension, by default False

• output_resolution (tuple[int, int], default: (181, 360) ) –

  (nlat, nlon) size of the output lat-lon grid. Only used when lat_lon=True, by default (181, 360) (1° resolution)

• time_tolerance (TimeTolerance, default: (timedelta64(-21, 'h'), timedelta64(3, 'h')) ) –

  Time tolerance for filtering observations, by default (-21 hours, 3 hours)

__call__

__call__(
    conv_obs: DataFrame | None = None,
    sat_obs: DataFrame | None = None,
) -> DataArray

Run HealDA inference from conventional and/or satellite observations.

At least one of the two observation DataFrames must be provided. Each DataFrame must carry a request_time entry in its .attrs.

Parameters:

• conv_obs (DataFrame | None, default: None ) –

  Conventional observation DataFrame from earth2studio.data.UFSObsConv, by default None

• sat_obs (DataFrame | None, default: None ) –

  Satellite observation DataFrame from earth2studio.data.UFSObsSat, by default None

Returns:

• DataArray –

  Global analysis on the HEALPix grid with dimensions [time, variable, npix]. Data is on the same device as the model (cupy array for GPU, numpy for CPU).

Raises:

• ValueError –

  If both conv_obs and sat_obs are None

create_generator

create_generator() -> Generator[
    DataArray,
    tuple[DataFrame | None, DataFrame | None],
    None,
]

Creates a generator which accepts collection of input observations and yields the output global assimilated data.

Yields:

• DataArray –

  Global analysis on the HEALPix grid

Receives:

• tuple[DataFrame | None, DataFrame | None] –

  A (conv_obs, sat_obs) tuple sent via generator.send(). Either element may be None but not both.

load_default_package classmethod

load_default_package() -> Package

Load the default HealDA model package from HuggingFace.

Returns:

• Package –

  Model package pointing to the HuggingFace repository

load_model classmethod

load_model(
    package: Package,
    lat_lon: bool = False,
    output_resolution: tuple[int, int] = (181, 360),
    time_tolerance: TimeTolerance = (
        timedelta64(-21, "h"),
        timedelta64(3, "h"),
    ),
) -> AssimilationModel

Load HealDA model from package.

Parameters:

• package (Package) –

  Package containing model checkpoint and statistics

• lat_lon (bool, default: False ) –

  If True the output is regridded to a regular lat-lon grid, by default False

• output_resolution (tuple[int, int], default: (181, 360) ) –

  (nlat, nlon) size of the output lat-lon grid. Only used when lat_lon=True, by default (181, 360)

• time_tolerance (TimeTolerance, default: (timedelta64(-21, 'h'), timedelta64(3, 'h')) ) –

  Time tolerance for filtering observations, by default (-21 hours, 3 hours)

Returns:

• AssimilationModel –

  Loaded HealDA assimilation model

Examples using earth2studio.models.da.HealDA

HealDA Global Data AssimilationProducing a global weather analysis from satellite and in-situ observations.