API reference¶

The entry points HealDA supports for use from other code. Everything else under healda may change without notice.

Inference¶

healda.inference ¶

Analyses from a trained HealDA checkpoint and observation tables.

load_da_model rebuilds the network the way the run trained it, from the loop.json the checkpoint carries; DAModel.run_analysis runs the recipe's own observation loaders and transform on in-memory cycle tables and returns physical fields. Downstream wrappers (Earth2Studio) build on this and add nothing to the numerics.

DAModel `dataclass` ¶

ir_channels `property` ¶

ir_channels: dict[str, list[int]]

Channel numbers the recipe reads per IR sounder; it reads no others.

to ¶

to(device) -> DAModel

Move the network and the device observations are prepared on.

frame_times ¶

frame_times(analysis_time) -> DatetimeIndex

The time_length frames ending at analysis_time.

observation_window ¶

observation_window(
    analysis_time,
) -> tuple[Timestamp, Timestamp]

Span of observation times one analysis reads, [start, end).

run_analysis ¶

run_analysis(
    analysis_times,
    *,
    satellite_tables: Mapping[str, CycleTableSource]
    | None = None,
    gpsro_tables: CycleTableSource | None = None,
    satwnd_tables: CycleTableSource | None = None,
    prepbufr_tables: CycleTableSource | None = None,
) -> Tensor

Physical-space analyses (b, C, lat, lon), one per analysis time.

Each analysis is the last frame of a time_length window fed with the observations of every frame. satellite_tables is keyed by sensor; a stream or sensor not passed contributes no observations.

to_physical ¶

to_physical(prediction: Tensor) -> Tensor

Model space (b, C, lat, lon) -> physical units, fills re-imposed.

load_da_model ¶

load_da_model(
    checkpoint,
    device,
    *,
    extra_filters: ObsFilters | None = None,
    fill_uv_sonde_heights: bool = True,
    fp32_output_head: bool = True,
) -> DAModel

Rebuild the trained network and its observation pipeline from a checkpoint.

Only the NNJA lat/lon recipe is supported: a 0.25 degree decode fed by NNJA satellite and NNJA conventional observations, the streams run_analysis takes as tables. checkpoint is a .checkpoint zip written by healda-train; extra_filters defaults to inference_filters(), as the inference CLI does. Provider flags act only where a table carries them; e2s_nnja tables carry none. fill_uv_sonde_heights and fp32_output_head are as in set_inference_recipe. Building the recipe fetches ERA5 statics into the healda cache on first use.

Earth2Studio observation adapter¶

healda.observations.adapters.e2s_nnja ¶

Earth2Studio NNJA frames as the canonical NNJA archive tables.

earth2studio.data.NNJAObsConv (PrepBUFR u, v, q, t, pres; GPS-RO gps, gps_refractivity) and NNJAObsSatwnd (u, v) return raw decoded rows. These functions restate them in the column contract of healda's NNJA PrepBUFR, GPS-RO and SATWND archive tables, deriving what the ETL derives with the same code, so the loaders screen, thin and emit them exactly as they do archive rows. Each returns {cycle: table}, a CycleTableSource for the loaders' table_source argument.

sat_tables restates NNJAObsSat radiances as the wide satellite archive's rows; NNJAWideLoader takes them per sensor as its table_source.

Rows are keyed by cycle_time, the NCEP file each was decoded from, which the Earth2Studio NNJA sources fill: an observation at the shared edge of two files' windows is in both files, and stays in both, as in the archive.

Usage, for the cycle file around cycle (dumps cover [cycle - 3h, cycle + 3h])::

frame = NNJAObsConv(time_tolerance=window)(cycle, ["gps", "gps_refractivity"])
loader = NNJAGpsroLoader(table_source=gpsro_tables(frame))
await loader.sel_time(times)  # or via NNJAConventionalLoader(gpsro_table_source=...)

conv = NNJAObsConv(time_tolerance=window, original_event=True)
frame = conv(cycle, ["u", "v", "q", "t", "pres"])
NNJAConventionalLoader(prepbufr_table_source=prepbufr_tables(frame))

frame = NNJAObsSat(time_tolerance=window)(cycle, ["atms_antenna_temperature"])
loader = NNJAWideLoader(["atms"], table_source={"atms": sat_tables(frame, "atms")})

analysis_tables ¶

analysis_tables(
    conv: DataFrame | None = None,
    satwnd: DataFrame | None = None,
    sat: DataFrame | None = None,
    *,
    sensors: Sequence[str],
    ir_channels: str | Mapping[str, Sequence[int]] | None,
    workers: int | None = None,
) -> dict

DAModel.run_analysis table arguments from Earth2Studio NNJA frames.

conv is an NNJAObsConv(original_event=True) frame of any PrepBUFR and GPS-RO variables, satwnd an NNJAObsSatwnd frame and sat an NNJAObsSat frame of any of sensors. A frame not passed gives no argument.

Every footprint, report, occultation and wind comes from one file, so each stream is cut by cycle_time and the pieces adapted in parallel on workers threads (default the smaller of 32 and the CPU count); the result equals adapting each whole stream.

Command-line entry points¶

healda-train and healda-inference-local are the installed console scripts. Run either with --help for the authoritative option list.

healda.cli.train ¶

Train HealPix Data Assimilation model

healda.cli.inference_local ¶

Run a trained checkpoint over a date range of local data and write the analysis.

Reads the observations and targets the checkpoint's recipe names from local storage, loads a checkpoint written by healda-train, and writes the analysis zarr. The checkpoint's loop.json supplies the geometry. Launch under torchrun -m healda.cli.inference_local.