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

Required imports:

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

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:

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.md (alphabetical). This is required for every new prognostic model so the API docs include the generated page.
• Add an entry to docs/userguide/about/install_options.yml for the model extra, even when the extra is empty. Include model-specific notes, source-specific preinstall commands, and api_refs for any model classes that should show a View Install Notes button in the generated API docs.
• 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, docs/userguide/about/install_options.yml has the install entry and api_refs, and docs/modules/models_px.md 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

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/