Add HealDAv2 video data assimilation model wrapper#975
Draft
negin513 wants to merge 1 commit into
Draft
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
HealDAwrapper: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 configurationPCACodecartifactslead_timedimension covering the full 8-frame windowNew files:
earth2studio/models/da/healda_v2.py— the wrapperearth2studio/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 checksDraft status / remaining work
.mdlusconversion;load_default_packageis stubbed until the HF package is publishednvidia-physicsnemopin bump in theda-healdaextra (released<=2.1.1predatesVideoHealDA)docs/modules/models_da.rstentry + example scripts3://noaa-reanalyses-pds) Nspect + DGPTT registration in progress; ERA5/UFS Replay coverage under the HealDA v1 submission being confirmedPreprocessing/API review is welcome now — the above only affect release, not the code shape.
Checklist
🤖 Generated with Claude Code