Earth2studio Create Prognostic

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 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

Never read `evals/targets/` — grader references only.

Reference Files

Load on demand during the matching step:

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:

model-name = ["package1>=version", "package2"]
# or, when no additional packages are required:
model-name = []

[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):

class ModelName(torch.nn.Module, AutoModelMixin, PrognosticMixin):

Required imports:

import numpy as np
import torch
from earth2studio.models.auto import AutoModelMixin, Package
from earth2studio.models.batch import batch_coords, batch_func
from earth2studio.models.px.base import PrognosticMixin
from earth2studio.models.utils import create_coords_from_lat_lon, handshake_dim
from earth2studio.lexicon import E2STUDIO_VOCAB
from earth2studio.utils import check_optional_dependencies
from loguru import logger

SPDX header (required at top of every .py file):

# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES.
# SPDX-License-Identifier: Apache-2.0

Canonical method order:

1. __init__ 2. input_coords 3. output_coords (@batch_coords)
2. load_default_package 5. load_model 6. to (optional)
3. Private methods 8. __call__ (@batch_func) 9. _default_generator
4. 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:

uv run pytest test/models/px/test_<name>.py -m "not package" -v
uv run pytest test/models/px/test_<name>.py::test_<model>_package --package -v

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:

make format && make lint && make license

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

User: Create IdentityModel - returns input unchanged, 6h step, 181x360, vars: t2m, u10m, v10m, msl

Agent: [reads SKILL.md, creates identity.py with triple inheritance,
        creates test_identity.py, runs pytest, runs make format && lint]

External Model (Pangu)

User: Add Pangu-Weather wrapper
      GitHub: https://github.com/198808xc/Pangu-Weather

Agent: [reads SKILL.md, fetches inference.py, creates pangu.py,