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 DAModel.run_analysis_unified the transform alone on a unified table; both return 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: AnalysisTimes,
    *,
    satellite_tables: Mapping[str, CycleTableSource]
    | None = None,
    gpsro_tables: CycleTableSource | None = None,
    satwnd_tables: CycleTableSource | None = None,
    prepbufr_tables: CycleTableSource | None = None,
) -> Tensor

Analyses from per-family observation tables, run through the recipe's loaders.

Each analysis is the last frame of a time_length window fed with the observations of every frame. The loaders quality control, thin and normalize the tables as in training. Analyses run one at a time.

Parameters:

analysis_times (AnalysisTimes) –

Analysis times, UTC, anything pd.DatetimeIndex accepts.
satellite_tables (Mapping[str, CycleTableSource] | None, default: None ) –

Per sensor name, {cycle: table} in schema.table_schema(sensor). A sensor not passed contributes no observations.
gpsro_tables (CycleTableSource | None, default: None ) –

{cycle: table} in schema.table_schema("gpsro").
satwnd_tables (CycleTableSource | None, default: None ) –

{cycle: table} in schema.table_schema("satwnd").
prepbufr_tables (CycleTableSource | None, default: None ) –

{cycle: table} in schema.table_schema("prepbufr"). A family not passed contributes no observations.

Returns:

Tensor –

(len(analysis_times), len(self.channels), 721, 1440) float32 analyses in
Tensor –

physical units on the 0.25 degree lat/lon grid, on self.device. Bounded
Tensor –

channels (humidity, water columns, snow depth, soil moisture, sea ice and
Tensor –

cloud fractions) are clipped to state_transforms.PHYSICAL_BOUNDS.

run_analysis_unified ¶

run_analysis_unified(
    analysis_times: AnalysisTimes, observations: Table
) -> Tensor

Analyses from a unified observation table, skipping the loaders.

Each frame reads the rows whose DA_window is one of its windows and passes them to featurization unchanged. Analyses run one at a time.

Parameters:

analysis_times (AnalysisTimes) –

Analysis times, UTC, anything pd.DatetimeIndex accepts.
observations (Table) –

One table in schema.table_schema("unified"), rows already thinned, quality controlled and normalized; rows outside the analyses' windows are not read.

Returns:

Tensor –

(len(analysis_times), len(self.channels), 721, 1440) float32 analyses in
Tensor –

physical units on the 0.25 degree lat/lon grid, on self.device. Bounded
Tensor –

channels (humidity, water columns, snow depth, soil moisture, sea ice and
Tensor –

cloud fractions) are clipped to state_transforms.PHYSICAL_BOUNDS.

Raises:

ValueError –

observations does not match the unified schema, or a read row holds a sensor_id the transform does not group or an id outside the network's embedding tables.

to_physical ¶

to_physical(prediction: Tensor) -> Tensor

Model space (b, C, lat, lon) -> physical units, bounded channels clipped to their physical range, fills re-imposed.

load_da_model ¶

load_da_model(
    checkpoint: str | PathLike,
    device: device | str,
    *,
    extra_filters: ObsFilters | None = None,
) -> DAModel

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

The recipe is the loop.json stored in the checkpoint; it must be an NNJA lat/lon recipe. The first call downloads the ERA5 static fields into the healda cache.

Parameters:

checkpoint (str | PathLike) –

A .checkpoint file written by healda-train.
device (device | str) –

Device the network runs on.
extra_filters (ObsFilters | None, default: None ) –

Observations removed beyond the recipe's own filters. Defaults to inference_filters(): the dated channel denials, and the provider quality flags of tables that carry them.

Returns:

DAModel –

The model, ready for run_analysis and run_analysis_unified.

Raises:

ValueError –

The checkpoint has no loop.json or is not an NNJA lat/lon recipe.

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), NNJAObsSatwnd (u, v) and NNJAObsSat return raw decoded observations in one of two layouts (layout):

long: one row per observation, its variable in variable and its value in observation;
wide: one row per report level, profile level, wind or footprint, each variable's value in observation_<variable> and, for conventional data, its quality mark in quality_<variable>. A satellite frame holds one sensor, each footprint's channels as fixed-size lists quality and observation_<variable> over the sensor's channel axis, which the frame carries in attrs["channels"].

analysis_tables takes wide frames and analysis_tables_long long ones. Both restate them in the column contract of healda's NNJA PrepBUFR, GPS-RO, SATWND and wide satellite 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, and both give the same tables. They return DAModel.run_analysis's table arguments: each a {cycle: table} CycleTables, satellite_tables one per sensor.

The module holds the wide layout and the archive builders both layouts share; the long layout's regrouping is its last section.

Satellite values stay in kelvin as earth2studio returns them. The IR sounders' columns are brightness_temperature__ch_NNNNN with field metadata nnja.value_unit = K, which NNJAWideLoader reads without a Planck inversion.

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, every stream at once from wide frames::

tables = analysis_tables(
    prepbufr=prepbufr, gpsro=gpsro, satwnd=satwnd, sat={"atms": atms},
    ir_channels=...,
)
model.run_analysis(times, **tables)

analysis_tables ¶

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

DAModel.run_analysis table arguments from wide Earth2Studio NNJA frames.

Each stream is its own frame: prepbufr and gpsro from NNJAObsConv, satwnd from NNJAObsSatwnd, and sat a mapping from sensor to that sensor's NNJAObsSat frame. A stream or sensor not passed gives no tables, which the loaders read as no observations. Long frames go to analysis_tables_long.

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 CPUs this process may run on); the result equals adapting each whole stream. The function returns once the streams are cut. Each {cycle: table} is a CycleTables whose lookups wait for that cycle's adaptation.

analysis_tables_long ¶

analysis_tables_long(
    *,
    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 long frames: conv 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.

Observation table schemas¶

healda.observations.schema ¶

Arrow schemas of the observation tables healda reads and emits.

The documented API is table_schema and validate. They are the contracts of two ways in:

the NNJA archive tables, per family (table_schema("prepbufr"), ...) and per satellite sensor (table_schema(sensor, ir_channels)), which the NNJA loaders and adapters.e2s_nnja.analysis_tables read;
the unified table, table_schema("unified"), one row per observation, which DAModel.run_analysis_unified takes.

footprint_fields and channel_field are the satellite tables' fields, which the NNJA ETL writes the archive with.

Each row of the unified table is one scalar measurement: a radiosonde level with u, v, t and q is four rows, and a satellite footprint is one row per channel.

table_schema ¶

table_schema(
    kind: str,
    ir_channels: str
    | Mapping[str, Sequence[int]]
    | None = None,
) -> Schema

The schema of kind: "prepbufr", "gpsro", "satwnd", a satellite sensor, or "unified".

For a satellite sensor, ir_channels is a preset name ("ir32", "ir48") or a mapping of sounder to channel numbers; the channels it names are required, and without it every channel of the sensor's vocabulary is optional and at least one is required. IASI and CrIS columns are brightness temperature in K.

"unified" is the table DAModel.run_analysis_unified takes and the NNJA loaders emit: one row per observation, thinned, quality controlled and normalized.

validate ¶

validate(
    table: Table,
    kind: str,
    ir_channels: str
    | Mapping[str, Sequence[int]]
    | None = None,
) -> None

Raise ValueError naming every column of table that kind's reader cannot read.

kind is as in table_schema. For a satellite sensor, the value columns are checked as brightness temperature when table has any, as the loader reads them.

Training presets¶

One YAML file per preset in src/healda/configs/loops/, each extending another, read by healda-train --config; --verbose prints the resolved loop in the same form.

healda.training.presets ¶

Training presets as YAML files, one per preset, each a few changes over another.

A preset file has up to three top-level keys. extends names the preset this one starts from; loop holds the fields that differ from it, nested the way the training loop dataclass is (a mapping under obs_config changes only the keys it lists, a list or scalar replaces the parent's value). Without extends, loop is the whole loop and anything missing takes the dataclass default. aliases lists other names that resolve to the same preset. The file's stem is the preset name.

The shipped presets live in healda.configs.loops and may only extend each other. A file passed to healda-train --config may extend any shipped preset, so a new experiment is a file with extends and the handful of fields it changes. healda-train --verbose prints the resolved loop in this form, every field written out.

The loop type is passed in rather than imported, so the training CLI can import this module to fill its registry without a cycle.

load_preset ¶

load_preset(
    path: str | PathLike,
    loop_type: Type[T],
    presets: Mapping[str, T] | None = None,
) -> tuple[T, list[str]]

Return the loop in path and its alias names.

presets is where extends is looked up, normally the training CLI's registry.

load_preset_dir ¶

load_preset_dir(
    directory: str | PathLike, loop_type: Type[T]
) -> dict[str, T]

Every *.yaml in directory under its stem and each of its aliases.

extends resolves to another file in the directory, by stem or alias, in any order.

to_yaml ¶

to_yaml(
    loop, aliases: Iterable[str] = (), header: str = ""
) -> str

loop as a preset file with every field written out.

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.

API reference¶

Inference¶

healda.inference ¶

DAModel dataclass ¶

ir_channels property ¶

to ¶

frame_times ¶

observation_window ¶

run_analysis ¶

run_analysis_unified ¶

to_physical ¶

load_da_model ¶

Earth2Studio observation adapter¶

healda.observations.adapters.e2s_nnja ¶

analysis_tables ¶

analysis_tables_long ¶

Observation table schemas¶

healda.observations.schema ¶

table_schema ¶

validate ¶

Training presets¶

healda.training.presets ¶

load_preset ¶

load_preset_dir ¶

to_yaml ¶

Command-line entry points¶

healda.cli.train ¶

healda.cli.inference_local ¶

DAModel `dataclass` ¶

ir_channels `property` ¶