earth2studio.models.da.HealDA#

class earth2studio.models.da.HealDA(model, condition, era5_mean, era5_std, sensor_stats, lat_lon=False, output_resolution=(181, 360), time_tolerance=(numpy.timedelta64(-21, 'h'), numpy.timedelta64(3, 'h')))[source]#

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 (torch.nn.Module) – The underlying HealDA neural network

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

  • era5_mean (torch.Tensor) – ERA5 per-channel mean for output denormalization [1, out_variables, 1, 1]

  • era5_std (torch.Tensor) – ERA5 per-channel std for output denormalization [1, out_variables, 1, 1]

  • sensor_stats (dict[str, dict[str, np.ndarray]]) – Per-sensor normalization statistics loaded from the package

  • lat_lon (bool, optional) – 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], optional) – (nlat, nlon) size of the output lat-lon grid. Only used when lat_lon=True, by default (181, 360) (1° resolution)

  • time_tolerance (TimeTolerance, optional) – Time tolerance for filtering observations, by default (-21 hours, 3 hours)

__call__(conv_obs=None, sat_obs=None)[source]#

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:
Returns:

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).

Return type:

xr.DataArray

Raises:

ValueError – If both conv_obs and sat_obs are None

create_generator()[source]#

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

Yields:

xr.DataArray – Global analysis on the HEALPix grid

Parameters:

  • None (tuple[pd.DataFrame |) – A (conv_obs, sat_obs) tuple sent via generator.send(). Either element may be None but not both.

  • None] (pd.DataFrame |) – A (conv_obs, sat_obs) tuple sent via generator.send(). Either element may be None but not both.

Return type:

Generator[DataArray, tuple[DataFrame | None, DataFrame | None], None]

classmethod load_default_package()[source]#

Load the default HealDA model package from HuggingFace.

Returns:

Model package pointing to the HuggingFace repository

Return type:

Package

classmethod load_model(package, lat_lon=False, output_resolution=(181, 360), time_tolerance=(numpy.timedelta64(-21, 'h'), numpy.timedelta64(3, 'h')))[source]#

Load HealDA model from package.

Parameters:

  • package (Package) – Package containing model checkpoint and statistics

  • lat_lon (bool, optional) – If True the output is regridded to a regular lat-lon grid, by default False

  • output_resolution (tuple[int, int], optional) – (nlat, nlon) size of the output lat-lon grid. Only used when lat_lon=True, by default (181, 360)

  • time_tolerance (TimeTolerance, optional) – Time tolerance for filtering observations, by default (-21 hours, 3 hours)

Returns:

Loaded HealDA assimilation model

Return type:

AssimilationModel

Examples using earth2studio.models.da.HealDA#

HealDA Global Data Assimilation

HealDA Global Data Assimilation