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](#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: ```toml 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):** ```python class ModelName(torch.nn.Module, AutoModelMixin, PrognosticMixin): ``` **Required imports:** ```python 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):** ```python # 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) 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:** ```bash 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.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:** ```bash 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, `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 ```text 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) ```text User: Add Pangu-Weather wrapper GitHub: https://github.com/198808xc/Pangu-Weather Agent: [reads SKILL.md, fetches inference.py, creates pangu.py, creates test_pangu.py, runs pytest] ``` --- ## Key Patterns ### Coordinate Template ```python @property def input_coords(self) -> CoordSystem: return CoordSystem({ "batch": np.empty(0), "time": np.empty(0), "lead_time": np.array([np.timedelta64(0, "h")]), "variable": np.array(["t2m", "u10m", ...]), # Public Earth2Studio convention is north-to-south latitude. "lat": np.linspace(90, -90, 181), "lon": np.linspace(0, 359, 360), }) @batch_coords() def output_coords(self, input_coords: CoordSystem) -> CoordSystem: output = input_coords.copy() output["lead_time"] = input_coords["lead_time"] + np.timedelta64(6, "h") return output ``` ### Iterator Template ```python def create_iterator(self, x, coords): yield x, coords # Initial condition (step 0) while True: x, coords = self.front_hook(x, coords) x, coords = self(x, coords) x, coords = self.rear_hook(x, coords) yield x, coords ``` --- ## 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/`