Skip to content

HealDAv2

GlobalDA202680 GBNVIDIAPyTorch

Import path: earth2studio.models.da.HealDAv2

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin

HealDA v2, a global machine-learning data-assimilation model that maps a 48-hour sequence of satellite and conventional observations to a 104-channel, 0.25 degree atmospheric analysis on a regular latitude-longitude grid.

HealDA v2 is a stateless, deterministic, observation-only assimilation model. One analysis at time t reads an eight-frame, six-hourly window ending at t and every observation within three hours of each frame, i.e. observations spanning [t - 45h, t + 3h], and returns the analysis of the last frame on the 721 x 1440 equiangular grid.

The model takes the NNJA observing system it was trained on as three DataFrames: conv_obs (PrepBUFR u, v, q, t, pres and GPS-RO gps, gps_refractivity), satwnd_obs (u, v) and sat_obs (radiances of atms, amsua, amsub, mhs, airs, iasi and cris). Any source producing these schemas works. To reproduce training from NNJA, use NNJAObsConv(original_event=True, exclude_message_types=("SATWND",)), NNJAObsSatwnd and NNJAObsSat(sensor_indices=model.sensor_indices).

Any stream may be omitted; the analysis then uses the others alone. Beyond the trained recipe, HealDA removes observations by its inference policy (healda.observations.system.inference_filters): channels inside dated denial windows, and footprints flagged by the provider quality columns of sat_obs.

Parameters:

  • model (DAModel) –

    The trained network with its observation pipeline, as returned by healda.inference.load_da_model

Note

For more information see the following references:

__call__ 

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

Run HealDA v2 from NNJA observations.

At least one DataFrame must be provided. Each must carry a request_time entry in its .attrs (earth2studio.data.fetch_dataframe sets it) and should cover [t - 45h, t + 3h] around each analysis time.

Parameters:

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

    PrepBUFR and GPS-RO observations from NNJAObsConv(original_event=True), by default None

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

    Satellite winds from NNJAObsSatwnd, by default None

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

    Satellite radiances from NNJAObsSat, by default None

Returns:

  • DataArray –

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

Raises:

  • ValueError –

    If every DataFrame is None or none carries request_time

create_generator 

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

Creates a generator which accepts observations and yields the analysis.

Yields:

  • DataArray –

    Global analysis on the 0.25 degree grid

Receives:

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

    A (conv_obs, satwnd_obs, sat_obs) tuple sent via generator.send(). Any element may be None but not all.

load_default_package classmethod 

load_default_package() -> Package

Load the default HealDA v2 model package from HuggingFace.

Returns:

  • Package –

    Model package pointing to the HuggingFace repository

load_model classmethod 

load_model(
    package: Package,
    checkpoint_name: str = "healda_v2.checkpoint",
) -> AssimilationModel

Load HealDA v2 from package.

The model is built on the CPU; move it with .to("cuda"), as the network runs on a CUDA device only, as a single unsharded process. The 0.25 degree recipe also fetches ERA5 static fields into the healda cache the first time it is built.

Parameters:

  • package (Package) –

    Package containing the .checkpoint archive written by healda-train

  • checkpoint_name (str, default: 'healda_v2.checkpoint' ) –

    Name of the checkpoint file inside the package, by default "healda_v2.checkpoint"

Returns:

  • AssimilationModel –

    Loaded HealDA v2 assimilation model

Examples using earth2studio.models.da.HealDAv2

HealDA v2 Global Data AssimilationProducing a 0.25 degree global analysis from the NNJA observing system.