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().