Skip to content

Add HealDAv2 video data assimilation model wrapper#975

Draft
negin513 wants to merge 1 commit into
NVIDIA:mainfrom
negin513:healda-v2
Draft

Add HealDAv2 video data assimilation model wrapper#975
negin513 wants to merge 1 commit into
NVIDIA:mainfrom
negin513:healda-v2

Conversation

@negin513

@negin513 negin513 commented Jul 17, 2026

Copy link
Copy Markdown
Member

Description

Adds HealDAv2, the earth2studio wrapper for the HealDA-v2 video data assimilation model (physicsnemo.experimental.models.healda.VideoHealDA, dit-5B). The model jointly produces an 8-frame / 48-hour window of global analyses at 6-hour spacing on the HEALPix level-6 grid; the final frame (lead_time == 0) is the present-time analysis.

Key differences from the v1 HealDA wrapper:

  • Pressure-level conventional obs normalization (conv-plevel): vertically structured conv channels are binned to the nearest of 13 ERA5 pressure levels and z-scored with per-level statistics, using the expanded 92-channel vocabulary — matching the current healda training configuration
  • Infrared hyperspectral sounders (IASI, CrIS-FSR, AIRS): footprints are pivoted and PCA-encoded to 32 latent observations each via bundled PCACodec artifacts
  • 50-dim unified observation metadata (features v2) as a pure-torch port (no Triton dependency)
  • Output carries a lead_time dimension covering the full 8-frame window

New files:

  • earth2studio/models/da/healda_v2.py — the wrapper
  • earth2studio/models/da/healda_v2_utils.py — self-contained ports of the sensor vocabulary, conv-plevel expansion, PCA codec, and metadata featurization (no dependency on the internal healda package)
  • test/models/da/test_da_healda_v2.py — 23 unit tests incl. hand-computed normalization/channel-mapping checks

Draft status / remaining work

  • Checkpoint: awaiting the pressure-level-normalization training checkpoint and its DCP → physicsnemo .mdlus conversion; load_default_package is stubbed until the HF package is published
  • nvidia-physicsnemo pin bump in the da-healda extra (released <=2.1.1 predates VideoHealDA)
  • docs/modules/models_da.rst entry + example script
  • Package-marked integration test
  • Dataset registration: NNJA (NOAA-NASA Joint Archive, CC BY 4.0, s3://noaa-reanalyses-pds) Nspect + DGPTT registration in progress; ERA5/UFS Replay coverage under the HealDA v1 submission being confirmed

Preprocessing/API review is welcome now — the above only affect release, not the code shape.

Checklist

  • Unit tests pass (mocked model, CPU)
  • v1 HealDA tests unaffected
  • black / ruff / mypy clean

🤖 Generated with Claude Code

Wraps physicsnemo VideoHealDA (8-frame 48h video-DA, dit-5B). Implements
the pressure-level conventional obs normalization from healda main
(conv-plevel expanded channel vocabulary), microwave sounder prep, IR PCA
footprint encoding (IASI/CrIS-FSR/AIRS), and the 50-dim unified metadata
featurization as a pure-torch port.

load_default_package is stubbed pending checkpoint conversion and HF
package publication.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant