earth2studio-create-prognostic

快速入门检查清单

请按顺序完成这些步骤。不要跳过任何步骤。

• 先完整阅读此 SKILL.md
• 获取参考脚本（步骤 0）
• 使用三重继承创建 earth2studio/models/px/<name>.py
• 使用模拟测试创建 test/models/px/test_<name>.py
• 运行：uv run pytest test/models/px/test_<name>.py -v
• 添加/更新模型 extra、安装文档、API 文档和变更日志（步骤 1–2、9）
• 运行：make format && make lint

⚠️ 重要： Python 命令始终使用 uv run：

• ✅ uv run pytest ... / uv run python ...
• ❌ pytest ... / python ...（缺少依赖）

遇到困难或输出错误： 不要反复尝试相同的修复方法。继续之前，请按照 自我改进的说明修补此 skill。

目的

实现一个预报模型封装器，将第三方机器学习天气模型接入 Earth2Studio。预报模型通过时间积分向前推演——给定初始状态后，它们会通过逐步推进时间来预测未来状态（例如，以 6 小时为增量）。

工作区

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

切勿读取 evals/targets/——该目录仅供评分器使用。

参考文件

在对应步骤中按需加载：

文件	内容	加载步骤
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

---

工作流步骤

步骤 0 — 获取参考脚本

如果提供了 $ARGUMENTS，则使用它。否则请询问：

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

步骤 1 — 分析并提出依赖项

分析：软件包、架构、I/O 形状、时间步长、分辨率、检查点。

提出 pyproject.toml 分组（按字母顺序排列，并添加到 all）。每个 预报模型都必须有一个可选依赖 extra，即使不需要任何软件包：

model-name = ["package1>=version", "package2"]
# or, when no additional packages are required:
model-name = []

[确认] 列出依赖项并征求用户批准。

步骤 2 — 添加依赖项

编辑 pyproject.toml：按字母顺序添加模型 extra，即使该 extra 为空，并更新 all 汇总项。

步骤 3 — 创建模型文件

文件： earth2studio/models/px/<lowercase>.py

必需的继承（全部三项）：

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

必需的导入：

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 文件顶部均为必需）：

# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES.
# SPDX-License-Identifier: Apache-2.0

规范方法顺序：

1. __init__ 2. input_coords 3. output_coords (@batch_coords)
2. load_default_package 5. load_model 6. to（可选）
3. 私有方法 8. __call__ (@batch_func) 9. _default_generator
4. create_iterator

步骤 4 — 实现坐标

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__ 和迭代器初始化中调用它。

步骤 5 — 实现前向传播

__call__： 使用 @batch_func 装饰，形状为 (batch, time, lead_time, var, lat, lon)。 重塑为模型格式 → 调用模型 → 重塑回原格式。

create_iterator： 必须先产出初始条件（step 0）。 使用 front_hook/rear_hook 注入扰动。

步骤 6 — 实现模型加载

load_default_package： 锁定 HuggingFace URL：hf://org/repo@commit

load_model： 使用 package.resolve()、map_location="cpu" 和 eval() 模式， 并使用 @check_optional_dependencies() 装饰。

步骤 7 — 编写测试

文件： test/models/px/test_<name>.py

必需的测试：

函数	用途
test_<model>_call	单次前向传播（按 device/time 参数化）
test_<model>_iter	迭代器产出序列
test_<model>_exceptions	无效坐标会引发错误
test_<model>_package	真实权重（@pytest.mark.package）

创建与接口匹配的 PhooModelName 虚拟模型，用于模拟测试。

运行测试：

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

不要省略 package 测试。如果任意随机输入不符合真实检查点的物理要求，请使用稳定且适合该模型的合成输入，同时仍加载真实权重并运行一次前向传播。

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

• 将模型添加到 earth2studio/models/px/__init__.py（按字母顺序）
• 在 pyproject.toml 中确认依赖项

步骤 9 — 文档

• 按字母顺序添加到 docs/modules/models_px.rst。每个新的预报模型都必须执行此操作，以便 API 文档包含生成的页面。
• 按字母顺序添加到 docs/userguide/about/install.md 中对应的选项卡，填写模型 extra；即使 extra 为空，也必须添加。包含模型专属说明，以及 pip install earth2studio[model-name] 和 uv add earth2studio --extra model-name 两种安装指令。
• 在 CHANGELOG.md 的 ### Added 下更新。这是每个新预报模型的必需操作。

格式与 lint：

make format && make lint && make license

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

遵循 references/validation-guide.md。创建未提交的 vanilla、E2S、comparison 和 sanity-check 脚本；不要提交生成的输出或图像。 使用适合 PR 的占位符表示图表，方便用户手动上传图片。

[CONFIRM] 用户必须先目视检查图表，才能继续。

步骤 11 - PR（如有要求）

遵循 references/validation-guide.md，并使用：

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

创建 PR 前，请确认 pyproject.toml 中有该模型 extra，all extra 包含它，安装文档同时包含 pip 和 uv 指令，并且已更新 docs/modules/models_px.rst 和 CHANGELOG.md。

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

---

示例

简单恒等模型

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]

外部模型（Pangu）

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]

---

关键模式

坐标模板

@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

迭代器模板

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

---

故障排除

错误	解决方案
OptionalDependencyFailure	uv add --optional <group> <pkg>
坐标握手失败	检查 handshake_dim 索引是否与维度位置匹配
迭代器形状错误	使用随机输入调试重塑逻辑
ModuleNotFoundError: pytest	使用 uv run pytest，而非 pytest

---

提醒

应做：

• 所有 Python 命令都使用 uv run python
• 使用 loguru.logger，绝不使用 print()
• 继承 torch.nn.Module + AutoModelMixin + PrognosticMixin
• 在 create_iterator 中先产出初始条件
• 在 _default_generator 中使用 front_hook()/rear_hook()
• 在每个 .py 文件中包含 SPDX 标头

不要做：

• 创建通用基类以供复用
• 提交 API 密钥或 comparison 脚本
• 从 evals/targets/ 读取内容