earth2studio-create-diagnostic

快速入门清单

请按顺序完成以下步骤。不要跳过步骤。编辑前，请完整阅读此 SKILL.md，并加载与模型类型相关的参考文件。实现后，在声称测试通过前，请先运行针对性的 pytest 命令。如果 无法运行测试，请报告确切的命令和失败信息，不要声称测试成功。

• 首先完整阅读此 SKILL.md
• 获取参考脚本、代码库、论文或模型文档（步骤 0）
• 将诊断模型归类为简单型、AutoModel 型或生成型（步骤 1）
• 在编辑依赖文件前，先提出依赖额外项（步骤 1）
• 创建 earth2studio/models/dx/<name>.py，其中仅包含诊断模型 API
• 创建 test/models/dx/test_<name>.py，并编写模拟测试
• 运行：uv run pytest test/models/dx/test_<name>.py -m "not package" -v
• 必要时添加或更新模型额外项、安装文档、API 文档和变更日志
• 运行：make format && make lint && make license

关键命令规则：Python 命令始终使用 uv run：

• 使用 uv run pytest ... 和 uv run python ...
• 不要在代码库工作流中直接使用 pytest 或 python

如果生成的模型有误，不要反复尝试相同的修复。请遵循 自我改进，修订此技能或其参考文件，然后 继续执行修正后的工作流。

目的

实现诊断模型封装器，将第三方或派生的 ML 转换操作连接到 Earth2Studio。诊断模型在单个时间点转换数据：输入字段进入，输出字段返回，不涉及预报积分。

前提条件

• 通过 uv 安装 Earth2Studio，并包含开发依赖（uv sync --all-extras）
• Python 3.10+ 环境
• 参考推理脚本、代码库、论文或模型文档
• 已打包模型的检查点来源和许可证信息

限制

• 仅处理单步转换
• 不支持随时间步进的预报模型；请使用 earth2studio-create-prognostic
• 真实软件包测试可能需要访问 NGC、HuggingFace、S3 或其他注册表的网络
• 有意义的生成式验证可能需要 GPU 和固定随机种子

诊断模型类型

类型	继承	依赖额外项	示例
简单派生诊断模型	仅 torch.nn.Module	通常无	Identity、风速
已打包的 AutoModel 诊断模型	torch.nn.Module, AutoModelMixin	必需，即使为空也需声明	PrecipitationAFNO
生成式诊断模型	torch.nn.Module, AutoModelMixin	必需，即使为空也需声明	CorrDiff

工作区

环境	位置
Harbor 评测	写入 /workspace/output/earth2studio/models/dx/...
Harbor + --copy-repo	完整检出位于 /workspace/repo
本地克隆	包含 pyproject.toml 的目录

切勿读取 evals/targets/；其中的文件仅供评分器参考。

参考文件

在相应工作流中按需加载以下文件：

文件	内容	加载时机
references/skeleton-template.py	简单型、AutoModel 型和生成型封装器的完整诊断模型骨架	步骤 3-6
references/method-templates.py	坐标、加载、前向传播和设备方法的精简代码片段	步骤 4-6
references/testing-guide.py	模拟、软件包、异常、样本和随机种子测试模式	步骤 7
references/validation-guide.md	参考比较、图表、PR 规范和审查后续工作	步骤 10-11
references/pr-body-template.md	PR 正文模板	步骤 11
references/pr-comment-template.md	验证评论模板	步骤 11

说明

步骤 0 - 获取参考资料

如果 $ARGUMENTS 提供了 URL 或本地路径，请使用该内容。否则询问：

请提供参考推理脚本、代码库、论文或模型文档。

记录参考模型的输入变量、输出变量、张量形状、归一化方式、网格、检查点来源、依赖要求和许可证。

步骤 1 - 分析类型并提出依赖项

编辑文件前，先对请求的诊断类型进行分类：

如果模型……	则使用……
计算派生量且不使用检查点	简单诊断
从 Package 或外部检查点加载权重	AutoModel 诊断
生成多个样本、扩散输出、VAE 样本或随机超分辨率结果	生成式诊断

依赖项政策：

• 简单派生诊断通常不需要 pyproject.toml 额外依赖组。
• AutoModel 和生成式诊断必须有命名的可选依赖组，即使依赖列表为空。
• 在 [project.optional-dependencies] 下按字母顺序添加额外依赖组，并将其纳入 all 聚合组。
• 在 OptionalDependencyFailure("model-extra") 和 @check_optional_dependencies() 中使用模型额外依赖组的名称。

提出依赖项额外组方案，并在编辑 pyproject.toml 前征求用户批准：

model-name = ["package1>=version", "package2"]
# or, when the packaged diagnostic needs no extra runtime packages:
model-name = []

步骤 2 - 添加依赖项

获得批准后，编辑 pyproject.toml：

• 按字母顺序添加额外依赖组。
• 更新 all 聚合组。
• 优先采用参考软件包文档中支持的最低版本。
• 除非参考模型要求且用户批准，否则不要添加未固定版本的宽泛 Git 依赖项。

步骤 3 - 创建模型文件

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

使用现有模型文件中所示的仓库标准 SPDX/许可证标头。

简单诊断常见的导入项包括：

from collections import OrderedDict
import numpy as np
import torch
from earth2studio.models.batch import batch_coords, batch_func
from earth2studio.utils import handshake_coords, handshake_dim
from earth2studio.utils.type import CoordSystem

打包诊断和生成式诊断通常还会包括：

from earth2studio.models.auto import AutoModelMixin, Package
from earth2studio.models.dx.base import DiagnosticModel
from earth2studio.utils.imports import OptionalDependencyFailure, check_optional_dependencies
from loguru import logger

标准方法顺序：

1. __init__
2. input_coords
3. 使用 @batch_coords() 装饰的 output_coords
4. 如有用，则添加 __str__
5. AutoModel/生成式诊断的 load_default_package
6. AutoModel/生成式诊断的 load_model
7. 仅当非 PyTorch 状态需要迁移设备时，才添加 to
8. 私有/辅助方法
9. 使用 @torch.inference_mode() 和 @batch_func() 装饰的 __call__

除非包装器自然地包含多个密切相关的变体，且一个小型基类能减少重复，否则避免使用共享基类或宽泛抽象。

步骤 4 - 实现坐标

诊断输入坐标通常采用此公开 Earth2Studio 顺序：

1. batch：np.empty(0)，并作为 OrderedDict 中的第一项
2. variable：使用 Earth2Studio 词汇名称的输入变量名称
3. lat：公开纬度约定为从北到南，通常为 90 到 -90
4. lon：公开经度约定为 0 到 360，通常不包含端点

任何诊断包装器都不应公开 lead_time。如果诊断需要有效时间元数据，请将其记录为 coords["time"] 中的逐样本元数据；除非现有 dx 模式有此要求，否则不要将其设为张量维度。

output_coords 必须使用 handshake_dim 和 handshake_coords 验证输入。然后更新输出变量；必要时更新输出纬度/经度分辨率。生成式诊断必须在 batch 后添加 sample 维度。

步骤 5 - 实现前向传递

使用单步 __call__；绝不要创建迭代器。执行模型前验证坐标，然后返回 (output_tensor, output_coords)。

@torch.inference_mode()
@batch_func()
def __call__(self, x: torch.Tensor, coords: CoordSystem) -> tuple[torch.Tensor, CoordSystem]:
    output_coords = self.output_coords(coords)
    x = (x - self.center) / self.scale
    out = self.core_model(x)
    return out, output_coords

对于生成式诊断，遍历批次维度，并为每个输入项生成 number_of_samples 个样本。如果参考实现支持设置采样种子，请使用显式种子以确保可复现性。

步骤 6 - 实现模型加载

对于打包诊断：

• load_default_package 应将 HuggingFace URL 锁定到某个提交（hf://org/repo@commit），或将 NGC/S3 版本锁定到不可变的发布版本。
• load_model 应调用 package.resolve(...)，先在 CPU 上加载检查点，将模块设为 eval()，并在适当情况下禁用梯度。
• 只有在必须加载完整的 PyTorch 序列化对象时，才使用 weights_only=False。
• 使用 @check_optional_dependencies() 装饰可选模型类和 load_model。
• 使用 loguru.logger 输出有用的加载信息；不要在 earth2studio/ 中使用 print()。

步骤 7 - 编写测试

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

必需的测试：

函数	用途
test_<model>_call	使用模拟模型或简单模型进行前向传播
test_<model>_exceptions	无效的坐标顺序、值或变量会引发错误
test_<model>_package	使用 @pytest.mark.package 和真实权重测试 AutoModel/生成式诊断

生成式诊断还需要样本数量和确定性种子的测试。 使用 references/testing-guide.py。创建一个符合真实核心模型接口并能生成确定性输出的 Phoo<ModelName> 虚拟模型。

运行针对性测试：

uv run pytest test/models/dx/test_<name>.py -m "not package" -v
uv run pytest test/models/dx/test_<name>.py::test_<model>_package --package -v

对于已打包的模型，不要省略包测试。如果任意随机输入不符合真实检查点的物理要求，请构造稳定且适合该模型的输入，同时仍加载真实权重并运行前向传播。

步骤 8 - 注册模型

对于公开模型，按字母顺序更新 earth2studio/models/dx/__init__.py。 只有当用户明确要求创建不应导出的内部或实验性文件时，才跳过注册。

步骤 9 - 文档

对于公开模型：

• 按字母顺序添加到 docs/modules/models_dx.rst，以便 API 文档包含生成的页面。
• 如果存在模型 extra，则添加到 docs/userguide/about/install.md。加入模型说明，以及 pip install earth2studio[model-name] 和 uv add earth2studio --extra model-name 两种安装命令。
• 在 CHANGELOG.md 的 ### Added 下更新。

格式化并检查代码风格：

make format && make lint && make license

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

遵循 references/validation-guide.md。创建未提交的原生、Earth2Studio、对比和健全性检查脚本。不要提交生成的输出、检查点、图像或本地验证脚本。

对于生成式诊断，请固定种子并比较匹配的样本；如果无法实现完全相等，则报告基于统计或容差的吻合情况。 继续之前，请用户目视检查图表。

步骤 11 - PR（如有要求）

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

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

创建 PR 之前，请确认依赖 extra、all、安装文档、API 文档、更新日志、测试和验证产物彼此一致。PR 文本中不要包含机器名称、主机名、绝对路径、缓存路径、设备清单或已上传图像的链接。手动上传图像时，请使用图表占位符。

示例

简单诊断

User: Create a diagnostic that computes wind speed from u10m and v10m.
Agent: Reads SKILL.md, classifies as simple, creates windspeed.py with only
       torch.nn.Module, writes call and exception tests, runs focused pytest.

AutoModel 诊断

User: Add a precipitation estimator from this reference script.
Agent: Reads SKILL.md and references, proposes dependency extra, creates a
       torch.nn.Module + AutoModelMixin wrapper, writes mock/package tests,
       updates docs/changelog/dependencies, and runs validation commands.

生成式诊断

User: Wrap this diffusion super-resolution model.
Agent: Classifies as generative, adds sample output coordinates, supports seed
       handling, writes sample and deterministic-seed tests, and prepares seeded
       validation comparisons.

故障排除

错误	解决方案
OptionalDependencyFailure	使用 uv sync --extra <model-extra> 安装，或修正 extra 名称
坐标握手失败	检查 OrderedDict 顺序和 handshake_dim 索引
输出形状错误	确认 output_coords 的长度与返回张量形状匹配
ModuleNotFoundError: pytest	使用 uv run pytest，不要直接运行 pytest
包测试在随机输入上失败	使用稳定且符合物理要求的输入，同时仍加载真实权重

提醒事项

应当：

• 所有 Python 命令均使用 uv run python 和 uv run pytest。
• 在 output_coords 上使用 @batch_coords()。
• 在 __call__ 上使用 @torch.inference_mode() 和 @batch_func()。
• 在 input_coords 中，将 batch 作为第一个坐标，并使用 np.empty(0)。
• 使用 handshake_dim() 和 handshake_coords() 验证坐标。
• 在生成式 output_coords 中添加 sample。
• 在每个 Python 文件中包含仓库标准的 SPDX/许可证头。
• 在 earth2studio/ 中使用 loguru.logger，绝不使用 print()。

不得：

• 继承 PrognosticMixin。
• 包含 lead_time 坐标。
• 创建 create_iterator。
• 除非明确需要支持多个变体，否则不要为单个包装器创建通用基类。
• 提交 API 密钥、凭据、验证脚本、绘图或生成的输出。
• 从 evals/targets/ 读取内容。

自我改进

如果此技能产生了错误输出，请在继续之前更新它：

1. 找出生成的代码或工作流中的问题。
2. 编辑 SKILL.md 或 references/ 中的相关文件，修正指南。
3. 对修改后的技能文件运行有针对性的验证。
4. 在需要提交的分支上工作时，单独提交技能修复。
5. 使用修正后的工作流继续模型实现。