Examples

The scripts below live in the repository and are included here verbatim.

Minimal training run

examples/training/minimal.sh runs the single-GPU smoke recipe. It is minimal in duration, not in data requirements: it expects the prepared state Zarr and the NNJA observation archive used by the selected recipe.

examples/training/minimal.sh

#!/usr/bin/env bash
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

set -euo pipefail

# The smoke preset is short but still requires compatible prepared assets.
: "${ERA5_HPX64_104CH_ZARR:?set the packed state Zarr path}"
: "${NNJA_ROOT:?set the NNJA observation archive root}"

OUTPUT_DIR="${OUTPUT_DIR:-./training-runs}"

torchrun --standalone --nproc-per-node=1 \
  -m healda.cli.train \
  --name debug-single-gpu \
  --output_dir "${OUTPUT_DIR}"

# Run this command again with the same name and output directory to resume the
# newest valid training-state checkpoint automatically.

Fine-tune from a checkpoint

examples/fine_tuning/from_checkpoint.sh initializes model weights from a checkpoint with --finetune_from. Optimizer and iterator state are not restored, so give a fine-tuning preset its own learning-rate schedule.

examples/fine_tuning/from_checkpoint.sh

#!/usr/bin/env bash
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

set -euo pipefail

# Experimental only: this workflow has not been scientifically validated.
: "${ERA5_HPX64_104CH_ZARR:?set the packed state Zarr path}"
: "${NNJA_ROOT:?set the NNJA observation archive root}"
: "${BASE_CHECKPOINT:?set a compatible training-state checkpoint}"

OUTPUT_DIR="${OUTPUT_DIR:-./fine-tuning-runs}"
RUN_NAME="${RUN_NAME:-debug-single-gpu}"

torchrun --standalone --nproc-per-node=1 \
  -m healda.cli.train \
  --name "${RUN_NAME}" \
  --output_dir "${OUTPUT_DIR}" \
  --finetune_from "${BASE_CHECKPOINT}"

# If OUTPUT_DIR/RUN_NAME already contains a valid checkpoint, normal resume
# takes precedence and BASE_CHECKPOINT is not loaded.

NNJA observations through Earth2Studio

docs/examples/e2s_nnja_observations.py loads prepared NNJA observations with the Earth2Studio adapter.

docs/examples/e2s_nnja_observations.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""Feed earth2studio NNJA observations through the HealDA NNJA loaders.

earth2studio returns raw decoded rows; ``e2s_nnja`` restates them in the NNJA archive
layout; the loaders then apply the same QC, thinning and normalization as training.

usage: python docs/examples/e2s_nnja_observations.py 2022-01-01T00
"""

import asyncio
import sys
from datetime import timedelta

import pandas as pd
from earth2studio.data import NNJAObsConv, NNJAObsSatwnd

from healda.observations.adapters import e2s_nnja
from healda.observations.loaders.combined import NNJAConventionalLoader


def main(target: pd.Timestamp) -> None:
    # The loader's (-3, 3) h context is one NCEP cycle file.
    tolerance = (timedelta(hours=-3), timedelta(hours=3))
    # Refractivity rows carry what the adapter derives height and pressure from.
    gps = NNJAObsConv(time_tolerance=tolerance)(
        target.to_pydatetime(), ["gps", "gps_refractivity"]
    )
    winds = NNJAObsSatwnd(time_tolerance=tolerance)(target.to_pydatetime(), ["u", "v"])

    loader = NNJAConventionalLoader(
        obs_context_hours=(-3, 3),
        gpsro_saids="full",
        include_satwnd=True,
        gpsro_table_source=e2s_nnja.gpsro_tables(gps),
        satwnd_table_source=e2s_nnja.satwnd_tables(winds),
    )
    obs = asyncio.run(loader.sel_time(pd.DatetimeIndex([target])))
    for name, tables in obs.items():
        print(name, [table.num_rows for table in tables])


if __name__ == "__main__":
    main(pd.Timestamp(sys.argv[1]))