earth2studio-create-prognostic

Files (18)

• BENCHMARK.md
• SKILL.md
• config.yml · evals/
• Dockerfile · evals/environment/
• bootstrap.sh · evals/environment/setup/
• evals.json · evals/
• eval_1_target.py · evals/targets/
• eval_2_target.py · evals/targets/
• eval_1_test_target.py · evals/targets/test/
• eval_2_test_target.py · evals/targets/test/
• method-templates.py · references/
• pr-body-template.md · references/
• pr-comment-template.md · references/
• skeleton-template.py · references/
• testing-guide.py · references/
• validation-guide.md · references/
• skill-card.md
• skill.oms.sig

SKILL.md (first lines)

## 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`)

From NVIDIA/skills, Apache-2.0.

More from NVIDIA/skills