⚡

创建Earth2Studio预测模型Skill earth2studio-create-prognostic

该技能用于在 Earth2Studio 中创建预测（时间步进）模型包装器，连接第三方机器学习天气模型，涵盖坐标处理、模型加载、测试、文档、验证与 PR 流程。关键词：Earth2Studio、气候预测、天气模型、prognostic model、PyTorch、时间积分、AI 气象预测。

气候预测 0 次安装 1 次浏览 更新于 9/6/2026
名称 earth2studio-create-prognostic
版本 0.16.0
开源协议 Apache-2.0 metadata:
作者 NVIDIA Earth-2 团队 agent-skills@nvidia.com tags: [earth2studio, prognostic-model, python]
描述 > 创建 Earth2Studio 预测（时间步进预报）模型包装器。 请勿用于诊断模型、数据源或安装。 argument-hint: 参考推理脚本的 URL 或本地路径（可选）

快速入门清单

按顺序执行这些步骤。不要跳过任何步骤。

  • [ ] 首先完整阅读本 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 ...（会缺少依赖）

遇到卡住或错误输出： 不要反复尝试同样的修复。请按照自我改进部分先修补此技能，然后再继续。

目的

实现一个预测模型包装器，将第三方机器学习天气模型连接到 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 — 分析并提议依赖项

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

在 pyproject.toml 中提议分组（按字母顺序，并添加到 all）。每个预测模型必须有一个可选依赖 extra，即使不需要任何包：

model-name = ["package1>=version", "package2"]
# 或者，当不需要额外包时：
model-name = []

[确认] 展示依赖并请求用户批准。

步骤 2 — 添加依赖项

编辑 pyproject.toml：按字母顺序添加模型 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
  • 如果 checkpoint/模型核心期望南到北纬度，请在核心模型前后翻转张量；不要在 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： 必须先产生初始条件（步骤 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 单次前向传递（参数化设备/时间）
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

不要省略包测试。如果任意随机输入对于真实检查点不物理有效，请使用稳定的、模型合适的合成输入，同时仍加载真实权重并执行前向传递。

步骤 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、对比和 sanity-check 脚本；不要提交生成的输出或图像。使用 PR 安全的图占位符，以便用户手动上传图片。

[确认] 用户必须先目视检查绘图，然后再继续。

步骤 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 文本中包含机器名称、绝对路径、设备清单或上传的图像链接。请使用图占位符。

示例

简单恒等模型

用户：创建 IdentityModel - 返回输入不变，6 小时间隔，181x360，变量：t2m, u10m, v10m, msl

助手：[阅读 SKILL.md，创建带三重继承的 identity.py，
       创建 test_identity.py，运行 pytest，运行 make format && lint]

外部模型（Pangu）

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

助手：[阅读 SKILL.md，获取 inference.py，创建 pangu.py，
       创建 test_pangu.py，运行 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", ...]),
        # 公开的 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

迭代器模板

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

故障排查

错误 解决方案
OptionalDependencyFailure uv add --optional <group> <pkg>
坐标 handshake 失败 检查 handshake_dim 索引是否匹配维度位置
迭代器形状错误 使用随机输入调试 reshape 逻辑
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 密钥或对比脚本
  • 读取 evals/targets/