earth2studio-deterministic-forecast

Earth2Studio 确定性预报技能

指导用户使用 earth2studio.run.deterministic 构建确定性（单成员）天气预报推理脚本。

前置条件

• 已安装 Earth2Studio，并配备支持 CUDA 的 GPU
• Python 3.10+，可访问网络以获取模型权重和数据

实时文档参考

在推荐组件前，获取相关文档以核实当前 API：

组件	URL
预报模型	https://nvidia.github.io/earth2studio/modules/models_px.html
数据源（分析场）	https://nvidia.github.io/earth2studio/modules/datasources_analysis.html
数据源（预报）	https://nvidia.github.io/earth2studio/modules/datasources_forecast.html
IO 后端	https://nvidia.github.io/earth2studio/modules/io.html
run.deterministic	https://github.com/NVIDIA/earth2studio/blob/main/earth2studio/run.py

工作流

1. 收集需求（跳过已提供的信息）

• 预报时长（小时/天/周）
• 感兴趣的变量（t2m、风、位势等）
• 区域（全球或特定区域，如 CONUS）
• 可用 GPU/显存

2. 选择模型

获取预报模型页面。根据预报时长、区域和显存进行筛选。记录模型的：

• 输入变量（input_coords["variable"]）
• 时间步长（output_coords["lead_time"]）

3. 选择数据源

数据源必须提供模型所需的所有输入变量。通过 earth2studio/lexicon/<source>.py 中的词汇表进行核实。常见搭配：全球模型 → GFS/ARCO/IFS；区域模型 → HRRR。

4. 选择 IO 后端

默认：ZarrBackend。传统工具使用 NetCDF4Backend，内存或小规模运行使用 XarrayBackend。

5. 计算 nsteps

nsteps = forecast_hours / model_step_hours

示例：预报 5 天，时间步长为 6h → nsteps = 120 / 6 = 20

6. 决定：筛选 output_coords

• 筛选变量（output_coords）：用户要求特定变量（例如“t2m 和风”）时使用，可减小输出大小
• 保存所有变量（省略 output_coords）：用户表示“所有变量”或未指定变量时使用，可保留模型的完整输出

7. 生成脚本

from collections import OrderedDict
import numpy as np
import torch
from earth2studio.models.px import <ModelClass>
from earth2studio.data import <DataSourceClass>
from earth2studio.io import <IOBackendClass>
from earth2studio.run import deterministic

model = <ModelClass>.load_model(<ModelClass>.load_default_package())
data = <DataSourceClass>()
io = <IOBackendClass>("<output_path>")

# Include output_coords ONLY if user requested specific variables
output_coords = OrderedDict({"variable": np.array(["t2m", "u10m"])})

io = deterministic(
    time=["YYYY-MM-DDTHH:MM:SS"],
    nsteps=<N>,
    prognostic=model,
    data=data,
    io=io,
    output_coords=output_coords,  # omit if saving all variables
    device=torch.device("cuda"),
)

8. 手动循环替代方案

当用户明确要求手动实现（不使用 earth2studio.run.deterministic）时，按顺序执行以下清单：

1. fetch_data - 获取初始条件：x, coords = fetch_data(data, time, model.input_coords, device)
2. 设置 total_coords - 构建 time 和 lead_time 维度的坐标数组
3. io.add_array - 在循环前使用 total_coords 初始化 IO 后端
4. create_iterator - 创建预报迭代器：model_iter = model.create_iterator(x, coords)
5. 循环 nsteps 次 - for step, (x, coords) in enumerate(model_iter): if step >= nsteps: break
6. map_coords - 根据需要筛选输出变量：x_out, coords_out = map_coords(x, coords, output_coords)
7. split_coords - 为 IO 写入做准备：x_out, coords_out = split_coords(x_out, coords_out)
8. io.write - 将每个时间步写入后端

9. 说明后续步骤

• 如何更改预报时间或运行多个初始化
• 如何读取输出（xr.open_zarr(...)）
• 指向用于后处理的诊断工作流

职责范围

负责： 模型选择、数据源兼容性、IO 后端选择、nsteps 计算，以及生成 earth2studio.run.deterministic 脚本。

不负责： 集合预报工作流、诊断、仅获取数据、安装和模型训练。

故障排除

常见错误及解决方案请参阅 references/troubleshooting.md。

提醒

• 在推荐模型或数据源之前，务必获取最新文档——API 会随版本更新而变化
• 验证词汇表兼容性——模型输入变量必须存在于数据源的 VOCAB 中
• 使用 load_default_package()——这是加载模型权重的标准模式
• 时间格式为 ISO 8601——time 参数使用 "YYYY-MM-DDTHH:MM:SS" 格式
• 风速需要两个分量——如果用户询问“风速”，请同时包含 u10m 和 v10m
• nsteps 使用整数除法——nsteps = total_hours // model_step_hours
• ZarrBackend 是默认选项——只有用户有特定需求时才建议其他选项
• 需要 GPU——所有预报模型都需要 CUDA；不支持 CPU 推理