earth2studio-create-prognostic

Create Earth2Studio prognostic (time-stepping forecast) model wrappers. Do NOT use for diagnostic models, data sources, or installation.

By nvidia · 582 installs

npx skills add nvidia/skills --skill earth2studio-create-prognostic

Source repository · Upstream listing

Quick Start Checklist Do these steps IN ORDER. Do not skip any step. [ ] Read this SKILL.md completely first [ ] Get reference script (Step 0) [ ] Create earth2studio/models/px/<name .py with triple inheritance [ ] Create test/models/px/test <name .py with mock tests [ ] Run: uv run pytest test/models/px/test <name .py v [ ] Add/update model extra, install docs, API docs, and changelog (Steps 1 2, 9) [ ] Run: make format && make lint ⚠️ CRITICAL: Always use uv run for Python commands: ✅ uv run pytest ... / uv run python ... ❌ pytest ... / python ... (missing dependencies) Stuck or wrong output: Do not keep retrying the same fix. Follow [Self Improvement]( self improvement) to patch this skill before continuing. Purpose Implement a prognostic model wrapper connecting third party ML weather models to Earth2Studio. Prognostic models time integrate forward—given initial state, they predict future states by stepping through time (e.g., 6 hour increments). Workspace Context Location Harbor eval Write to /workspace/output/earth2studio/models/px/... Harbor + copy repo Full checkout at /workspace/repo Local clone Directory with pyproject.toml Never read evals/targets/ — grader references only. Reference Files Load on demand during the matching step: File Content Load at references/skeleton template.py Full model skeleton with FILL comments Steps 3–6 references/method templates.py Canonical method implementations Steps 4–6 references/testing guide.py Test skeleton and mock patterns Step 7 references/validation guide.md Comparison scripts, PR, code review Steps 10–11 Workflow Steps Step 0 — Get Reference Script If $ARGUMENTS provided, use it. Otherwise ask: Please provide a reference inference script URL/path. Step 1 — Analyze & Propose Dependencies Analyze: packages, architecture, I/O shapes, time step, resolution, checkpoint. Propose pyproject.toml group (alphabetical, add to all ). Every prognostic model must have an optional dependency extra, even when no packages are required: [CONFIRM] Present dependencies and ask user to approve. Step 2 — Add Dependencies Edit pyproject.toml : add the model extra alphabetically, even if it is empty, and update the all aggregate. Step 3 — Create Model File File: earth2studio/models/px/<lowercase .py Required inheritance (all three): Required imports: SPDX header (required at top of every .py file): Canonical method order: 1. init 2. input coords 3. output coords (@batch coords) 4. load default package 5. load model 6. to (optional) 7. Private methods 8. call (@batch func) 9. default generator 10. create iterator Step 4 — Implement Coordinates input coords rules: batch : np.empty(0) time : np.empty(0) (dynamic) lead time : starts at np.timedelta64(0, "h") lat : 90 to 90 (north to south); this is the public Earth2Studio convention even if the source model uses the opposite order lon : 0 to 360 If a checkpoint/model core expects south to north latitude, flip tensors internally before/after the core model; do not expose flipped latitude in input coords or output coords Map variables to E2STUDIO VOCAB (282 entries in earth2studio/lexicon/base.py ) output coords: Use handshake dim / handshake coords for input validation, then increment lead time . Prefer a shared coordinate check helper and call it from output coords , call , and iterator setup before model execution. Step 5 — Implement Forward Pass call : @batch func decorated, shape (batch, time, lead time, var, lat, lon). Reshape to model format → call model → reshape back. create iterator : MUST yield initial condition first (step 0). Use front hook / rear hook for perturbation injection. Step 6 — Implement Model Loading load default package : Lock HuggingFace URLs: hf://org/repo@commit load model : Use package.resolve() , map location="cpu" , eval() mode, decorate with @check optional dependencies() . Step 7 — Write Tests File: test/models/px/test <name .py Required tests: Function Purpose test <model call Single forward pass (parametrize device/time) test <model iter Iterator produces sequence test <model exceptions Invalid coords raise errors test <model package Real weights ( @pytest.mark.package ) Create PhooModelName dummy matching interface for mock tests. Run tests: Do not omit the package test. If arbitrary random inputs are not physically valid for the real checkpoint, use a stable model appropriate synthetic input while still loading real weights and running a forward pass. Step 8 — Register Model (if requested) Add to earth2studio/models/px/ init .py (alphabetical) Verify deps in pyproject.toml Step 9 — Documentation Add to docs/modules/models px.rst (alphabetical). This is required for every new prognostic model so the API docs include the generated page. Add to docs/userguide/about/install.md (alphabetical tab) for the model extra, even when the extra is empty. Include model specific notes plus both pip install earth2studio[model name] and uv add earth2studio extra model name instructions. Update CHANGELOG.md under Added . This is required for every new prognostic model. Format and lint: Step 10 Validation (if requested) Follow references/validation guide.md . Create uncommitted vanilla, E2S, comparison, and sanity check scripts; do not commit generated outputs or images. Use PR safe placeholders for plots so the user can upload images manually. [CONFIRM] User must visually inspect plots before proceeding. Step 11 PR (if requested) Follow references/validation guide.md and use: references/pr body template.md references/pr comment template.md Before creating the PR, verify pyproject.toml has the model extra, the all extra includes it, install docs include both pip and uv commands, and docs/modules/models px.rst plus CHANGELOG.md are updated. Do not include machine names, absolute paths, device inventory, or uploaded image links in PR text. Use plot placeholders instead. Examples Simple Identity Model External Model (Pangu) Key Patterns Coordinate Template Iterator Template Troubleshooting Error Solution OptionalDependencyFailure uv add optional <group <pkg Coordinate handshake fails Check handshake dim indices match dim position Iterator wrong shapes Debug reshape logic with random input ModuleNotFoundError: pytest Use uv run pytest not pytest Reminders DO: Use uv run python for ALL Python commands Use loguru.logger , never print() Inherit torch.nn.Module + AutoModelMixin + PrognosticMixin Yield initial condition first in create iterator Use front hook() / rear hook() in default generator Include SPDX header in every .py file DON'T: Create general base classes for reuse Commit API keys or comparison scripts Read from evals/targets/