用途:为 Earth2Studio 创建诊断模型封装(diagnostic model wrapper),把第三方或派生的机器学习变换接入框架。诊断模型只做单步数据变换:输入场进、输出场出,不涉及预报时间积分;时间步进预报模型应改用 earth2studio-create-prognostic。
主要能力与分类:按实现方式分为三类——简单派生诊断(仅继承 torch.nn.Module,通常无需依赖 extra,如恒等、风速);打包 AutoModel 诊断(torch.nn.Module + AutoModelMixin,必须声明命名可选依赖 extra,即使为空,如 PrecipitationAFNO);生成式诊断(同样需 extra,如 CorrDiff,输出含 sample 维度)。
使用方式:先完整阅读 SKILL.md,取得参考推理脚本、仓库、论文或模型文档,记录输入/输出变量、张量形状、归一化、网格、checkpoint 来源与许可。分类后先提出依赖 extra 并获用户批准,再编辑 pyproject.toml(按字母序加入 [project.optional-dependencies] 并更新 all 聚合)。随后创建 earth2studio/models/dx/<name>.py,按规范顺序实现 init、input_coords、output_coords(加 @batch_coords())、load_default_package/load_model(AutoModel 与生成式)、call(加 @torch.inference_mode() 与 @batch_func())。坐标中 batch 用 np.empty(0) 且置于首位,lat 通常 90 到 -90,lon 0 到 360,不得暴露 lead_time;生成式需在 batch 后加 sample 维度。测试文件为 test/models/dx/test_<name>.py,需含 call、exceptions 及 package 测试。
必要命令:所有 Python 命令必须用 uv run,例如 uv run pytest test/models/dx/test_<name>.py -m "not package" -v;格式化与检查用 make format && make lint && make license。
适用场景:新增派生诊断量、封装外部 checkpoint 的降水等诊断模型、封装扩散/超分等生成式诊断。
限制:仅处理单步变换;真实包测试可能需要访问 NGC、HuggingFace、S3 等网络;生成式校验可能需要 GPU 与固定随机种子;不得读取 evals/targets/;不得提交密钥、校验脚本、图表或生成产物。