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 one wide DataFrame per stream, passed positionally in a fixed order, model(prepbufr, gpsro, satwnd, *sensors): prepbufr (u, v, q, t, pres), gpsro (gps, gps_refractivity), satwnd (u, v), then one frame per satellite sensor in the model's sensor order (a subset of atms, amsua, amsub, mhs, airs, iasi, cris). input_coords() returns one schema per argument in the same order; a sensor schema's variable entry is its NNJAObsSat variable. Any source producing these schemas works; the NNJA sources return them with layout="wide", one stream per call. To reproduce training from NNJA, use NNJAObsConv(original_event=True, exclude_message_types=("SATWND",), layout="wide") for prepbufr and gpsro, NNJAObsSatwnd(layout="wide") and NNJAObsSat(sensor_indices=model.sensor_indices, layout="wide") for each sensor; earth2studio.data.fetch_dataframes fetches them concurrently.

Any stream may be None or empty, and trailing streams may be left out; 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 satellite streams' provider quality columns.

Recommended environment, set before Python starts:

  • MALLOC_ARENA_MAX=16 keeps host memory level over repeated analyses in one process.
  • PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True avoids GPU out-of-memory from fragmentation over repeated analyses.
  • OMP_NUM_THREADS=1.

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__(*obs: DataFrame | None) -> DataArray

Run HealDA v2 from NNJA observations.

Each frame 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:

  • *obs (DataFrame | None, default: () ) –

    Wide observation frames in input_coords() order; a stream may be None or empty and trailing streams may be left out, but at least one frame must be given

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 no frame is given, more frames are given than streams, no frame carries request_time, or a frame is in the long layout

create_generator 

create_generator() -> (
    Generator[DataArray, tuple[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, ...] –

    Observation frames in input_coords() order, as for __call__, sent via generator.send()

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.