earth2studio-create-prognostic

Compare original and translation side by side

🇺🇸

Original

English

🇨🇳

Translation

Chinese

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.

请按顺序完成以下步骤，请勿跳过任何步骤。

先完整阅读本SKILL.md文档
获取参考脚本（步骤0）
创建
```
earth2studio/models/px/<name>.py
```
文件，使用三重继承
创建
```
test/models/px/test_<name>.py
```
文件，编写模拟测试

运行命令：

uv run pytest test/models/px/test_<name>.py -v

添加/更新模型扩展依赖、安装文档、API文档及变更日志（步骤1-2、9）
运行命令：
```
make format && make lint
```

⚠️ 重要提示： 所有Python命令请务必使用
uv run
执行：
✅
uv run pytest ...
/
uv run python ...
❌
pytest ...
/
python ...
（会缺失依赖）
遇到问题或输出异常： 请勿重复尝试相同的修复方案。请遵循自我改进部分的内容完善本技能后再继续。

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

实现一个预测型模型包装器，将第三方机器学习气象模型与Earth2Studio对接。预测型模型会进行时间积分推演——给定初始状态，它们通过时间步长（例如6小时增量）预测未来状态。

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.

上下文	位置
Harbor 评估	写入至 `/workspace/output/earth2studio/models/px/...`
Harbor + `--copy-repo`	完整代码检出至 `/workspace/repo`
本地克隆	包含 `pyproject.toml` 的目录

请勿读取
evals/targets/
目录 —— 该目录仅用于 grader 参考。

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

匹配步骤中按需加载：

文件	内容	加载时机
`references/skeleton-template.py`	带有FILL注释的完整模型骨架	步骤3–6
`references/method-templates.py`	标准方法实现示例	步骤4–6
`references/testing-guide.py`	测试骨架与模拟模式	步骤7
`references/validation-guide.md`	对比脚本、PR、代码评审指南	步骤10–11

Workflow Steps

工作流步骤

Step 0 — Get Reference Script

步骤0 — 获取参考脚本

$ARGUMENTS

provided, use it. Otherwise ask:

Please provide a reference inference script URL/path.

若提供了

$ARGUMENTS

，则直接使用。否则请询问：

请提供参考推理脚本的URL/路径。

Step 1 — Analyze & Propose Dependencies

步骤1 — 分析并提出依赖建议

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

分析内容：包、架构、输入输出形状、时间步长、分辨率、检查点。

在

pyproject.toml

中提出分组建议（按字母排序，添加至

all

分组）。每个预测型模型都必须有一个可选依赖扩展，即使不需要额外包：

toml

model-name = ["package1>=version", "package2"]

or, when no additional packages are required:

若无需额外包：

model-name = []


**[CONFIRM]** Present dependencies and ask user to approve.

model-name = []


**[确认]** 列出依赖并请求用户批准。

Step 2 — Add Dependencies

步骤2 — 添加依赖

Edit

pyproject.toml

: add the model extra alphabetically, even if it is empty, and update the

all

aggregate.

编辑

pyproject.toml

：按字母顺序添加模型扩展（即使为空），并更新

all

聚合分组。

Step 3 — Create Model File

步骤3 — 创建模型文件

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

undefined

文件路径：

earth2studio/models/px/<lowercase>.py

必须使用三重继承：

python

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

必须导入的模块：

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头部（每个.py文件顶部必须添加）：

python

undefined

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`


**标准方法顺序：**
1. `__init__` 2. `input_coords` 3. `output_coords` (@batch_coords)
4. `load_default_package` 5. `load_model` 6. `to`（可选）
7. 私有方法 8. `__call__` (@batch_func) 9. `_default_generator`
10. `create_iterator`

Step 4 — Implement Coordinates

步骤4 — 实现坐标系统

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.

input_coords规则：

```
batch
```
:
```
np.empty(0)
```
```
time
```
:
```
np.empty(0)
```
（动态）
```
lead_time
```
: 起始值为
```
np.timedelta64(0, "h")
```
```
lat
```
: 90到-90（北到南）；这是Earth2Studio的公开约定，即使源模型使用相反顺序
```
lon
```
: 0到360
如果检查点/模型核心要求纬度从南到北，则在核心模型执行前后内部翻转张量；请勿在
```
input_coords
```
或
```
output_coords
```
中暴露翻转后的纬度

将变量映射至

E2STUDIO_VOCAB

（

earth2studio/lexicon/base.py

中有282个条目）

output_coords： 使用

handshake_dim

handshake_coords

进行输入验证，然后递增

lead_time

。建议使用共享的坐标检查辅助函数，并在

output_coords

、

__call__

和迭代器初始化（模型执行前）中调用它。

Step 5 — Implement Forward Pass

步骤5 — 实现前向传播

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

__call__
：带有@batch_func装饰器，形状为(batch, time, lead_time, var, lat, lon)。将形状转换为模型格式 → 调用模型 → 转换回原形状。

create_iterator
：必须先返回初始条件（步骤0）。使用

front_hook

rear_hook

注入扰动。

Step 6 — Implement Model Loading

步骤6 — 实现模型加载

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

load_default_package
：锁定HuggingFace URL：

hf://org/repo@commit

load_model
：使用

package.resolve()

、

map_location="cpu"

、

eval()

模式，添加

@check_optional_dependencies()

装饰器。

Step 7 — Write Tests

步骤7 — 编写测试

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.

文件路径：

test/models/px/test_<name>.py

必须包含的测试：

函数	目的
`test_<model>_call`	单次前向传播（参数化设备/时间）
`test_<model>_iter`	迭代器生成序列
`test_<model>_exceptions`	非法坐标触发错误
`test_<model>_package`	真实权重测试（ `@pytest.mark.package` ）

创建与接口匹配的

PhooModelName

虚拟模型用于模拟测试。

运行测试：

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

请勿省略包测试。如果真实检查点不接受任意随机输入，请使用稳定的、适合模型的合成输入，同时仍需加载真实权重并执行前向传播。

Step 8 — Register Model (if requested)

步骤8 — 注册模型（若有要求）

Add to
```
earth2studio/models/px/__init__.py
```
(alphabetical)
Verify deps in pyproject.toml

将模型添加至
```
earth2studio/models/px/__init__.py
```
（按字母排序）
验证
```
pyproject.toml
```
中的依赖

Step 9 — Documentation

步骤9 — 文档编写

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:

bash

make format && make lint && make license

将模型添加至
```
docs/modules/models_px.rst
```
（按字母排序）。每个新预测型模型都必须添加，这样API文档才会包含生成的页面。
将模型添加至
```
docs/userguide/about/install.md
```
（按字母排序的标签页）中的模型扩展部分，即使扩展为空。包含模型特定说明，以及
```
pip install earth2studio[model-name]
```
和
```
uv add earth2studio --extra model-name
```
两种安装指令。
在
```
CHANGELOG.md
```
的
```
### Added
```
部分更新内容。每个新预测型模型都必须更新。

格式化与代码检查：

bash

make format && make lint && make license

Step 10 - Validation (if requested)

步骤10 - 验证（若有要求）

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.

遵循

references/validation-guide.md

。创建未提交的原生版本、E2S版本、对比版本和合理性检查脚本；请勿提交生成的输出或图像。使用PR安全的占位符表示图表，以便用户手动上传图像。

[确认] 用户必须先可视化检查图表，然后才能继续。

Step 11 - PR (if requested)

步骤11 - 创建PR（若有要求）

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.

遵循

references/validation-guide.md

并使用：

```
references/pr-body-template.md
```
```
references/pr-comment-template.md
```

创建PR前，验证

pyproject.toml

包含模型扩展、

all

扩展包含该模型、安装文档包含pip和uv两种命令，且

docs/modules/models_px.rst

和

CHANGELOG.md

已更新。

PR文本中请勿包含机器名称、绝对路径、设备清单或上传的图像链接。请使用图表占位符替代。

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]

text

用户：创建IdentityModel - 原样返回输入，6小时步长，181x360分辨率，变量：t2m, u10m, v10m, msl

Agent：[阅读SKILL.md，创建带有三重继承的identity.py，创建test_identity.py，运行pytest，运行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]

text

用户：添加Pangu-Weather包装器
      GitHub：https://github.com/198808xc/Pangu-Weather

Agent：[阅读SKILL.md，获取inference.py，创建pangu.py，创建test_pangu.py，运行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

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", ...]),
        # Earth2Studio公开约定为纬度从北到南
        "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

python

def create_iterator(self, x, coords):
    yield x, coords  # 初始条件（步骤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`

错误	解决方案
`OptionalDependencyFailure`	执行 `uv add --optional <group> <pkg>`
坐标握手失败	检查 `handshake_dim` 索引是否与维度位置匹配
迭代器形状错误	使用随机输入调试形状转换逻辑
`ModuleNotFoundError: pytest`	使用 `uv run pytest` 而非 `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()

_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/
```

请务必：

所有Python命令使用
```
uv run python
```
执行
使用
```
loguru.logger
```
，禁止使用
```
print()
```

继承

torch.nn.Module + AutoModelMixin + PrognosticMixin

在
```
create_iterator
```
中先返回初始条件

在

_default_generator

中使用

front_hook()

rear_hook()

每个.py文件都添加SPDX头部

请勿：

创建通用基类用于复用
提交API密钥或对比脚本
读取
```
evals/targets/
```
目录