Earth2Studio：英伟达开源的 AI 天气气候预测与分析工具包

一、Earth2Studio是什么？

Earth2Studio是一款由英伟达基于Python开发的开源AI推理流水线工具包，核心定位是为天气和气候科学领域提供“开箱即用”的AI开发与应用解决方案。其核心目标是打破AI地球系统模型的使用壁垒，让不同背景的用户（从科研人员到工程开发者）都能快速上手，构建属于自己的天气气候AI工作流。

从本质来看，Earth2Studio并非单一的AI模型，而是一个“生态级工具平台”——它搭建了统一的API接口，能够兼容不同的AI框架（如PyTorch、JAX、ONNX）、模型架构（图神经网络、Transformer、神经算子等）、数据存储系统（本地文件、云存储）和科学计算工具（SciML、Xarray），将气象AI开发中的“数据源接入、模型调用、数据处理、结果存储、性能评估”等环节模块化，用户可像“搭积木”一样组合各类组件，无需关注底层兼容性问题。

其核心使命可概括为三点：一是降低AI天气气候模型的使用门槛，通过预训练模型库和极简API，让非AI专家也能快速实现高精度预测；二是赋能科研创新，支持自定义组件扩展，为新模型、新算法的验证提供高效流水线；三是推动行业标准化，通过统一的数据格式、坐标系统和接口设计，解决气象数据分散、模型兼容性差的行业痛点。

二、功能特色

Earth2Studio的核心竞争力源于其“模块化、全栈式、高兼容”的功能设计，具体特色可分为以下七大核心亮点：

1. 模块化可组合，支持复杂流水线构建

工具包的所有核心组件（数据源、模型、IO存储、扰动方法、统计指标）均采用模块化设计，支持灵活组合与链式调用。用户可根据需求搭建从“数据获取-模型推理-结果处理-评估分析”的全流程流水线，无论是简单的单模型预报，还是复杂的多模型串联、集合预报，都能通过组件组合快速实现。例如：可将“GFS数据源+GraphCast预测模型+CBottle超分诊断模型+Zarr存储后端”组合，实现“全球预报→区域高分辨率优化→云存储”的端到端流程。

2. 全球最大的天气气候AI预训练模型库

Earth2Studio内置了当前最全面的天气气候AI预训练模型库，覆盖“预测型”和“诊断型”两大类别，适配不同分辨率、覆盖范围和应用场景：

• 预测模型支持从6小时到3个月的时间尺度预测，分辨率从全球0.25°（约28km）到区域3km，架构涵盖图神经网络（GraphCast系列）、Transformer（Pangu、Aurora、FuXi等）、神经算子（SFNO）、扩散模型（StormCast）等主流类型；

• 诊断模型专注于特定气象衍生量计算，无需时间积分，可快速输出降水、太阳辐射、台风轨迹、最大阵风等关键指标，填补了传统预测模型的功能缺口。

3. 标准化多源数据访问，无需关注底层格式

工具包提供了统一的数据源API，整合了运营气象模型（GFS、HRRR、IFS）、再分析数据（ARCO ERA5、CDS）、AI生成数据（CBottle3D）等各类气象数据，支持GRIB2、Zarr、NetCDF、HEALPix等多种格式。用户无需手动处理数据下载、格式转换、坐标对齐等复杂操作，只需调用对应数据源接口，即可获得标准化的Xarray数据数组，直接用于模型输入或验证。例如：调用GFS()接口即可自动获取全球0.25°分辨率的运营气象数据，无需关心数据来自AWS还是NOMADS平台。

4. 高效IO存储与多场景适配

内置多种IO存储后端，覆盖从本地存储到云存储的全场景需求：

• 支持传统科学格式（NetCDF4）、云优化格式（Zarr）和内存存储（Xarray、Key-Value）；

• 提供异步写入、并行IO等高级功能（AsyncZarrBackend），适配大规模数据的高效存储与访问；

• 所有输出数据均保留完整元数据，符合CF规范，可直接用于后续分析或可视化工具（如Panoply、Matplotlib）。

5. 扰动方法与集合预报支持，捕捉预测不确定性

天气气候预测的核心挑战之一是不确定性量化，Earth2Studio内置了多种标准化扰动方法（高斯噪声、相关球面高斯、 Bred Vector等），支持构建集合预报流水线。通过对初始数据添加合理扰动，结合多模型或多参数推理，可输出概率性预报结果，并用内置指标（如CRPS、Rank Histogram）评估集合的可靠性与技巧，显著提升预测结果的实用性。

6. 内置丰富统计指标，实时评估预报性能

工具包集成了气象预报领域常用的统计评估指标，涵盖误差度量（RMSE）、相关性分析（ACC）、集合验证（CRPS、Rank Histogram）、离散度分析（标准差、Spread-Skill Ratio）等，支持在流水线中实时计算，无需额外编写评估代码。用户可直接通过API调用指标函数，快速对比不同模型、不同参数的预报效果，加速模型优化与验证过程。

7. GPU加速与多框架兼容，兼顾效率与灵活性

Earth2Studio原生支持GPU加速，可充分利用NVIDIA GPU的计算能力，大幅提升模型推理速度。同时，工具包兼容PyTorch、JAX、ONNX等主流AI框架，支持PhysicsNeMo训练的Earth-2模型无缝部署，用户无需重构现有模型代码，即可接入工具包的流水线生态，兼顾计算效率与开发灵活性。

三、技术细节

1. 核心模型架构与参数

Earth2Studio的模型库涵盖多种先进AI架构，适配不同气象预测场景，下表汇总了核心模型的关键技术参数：

注：“-”表示诊断模型无需时间步长（无时间积分）

从技术原理来看，预测模型的核心是“时间序列建模”，通过 autoregressive（自回归）方式，基于当前气象场预测未来多个时间步的气象场，例如GraphCast利用图神经网络捕捉全球气象要素的空间关联，Pangu则通过Transformer架构建模长时程依赖；诊断模型则聚焦“空间映射”，通过深度学习模型将输入的气象场直接映射为目标衍生量，无需考虑时间维度的演变，例如PrecipitationAFNO通过神经算子高效处理高维气象数据，快速输出降水结果。

2. 数据处理核心技术

• 统一坐标系统：所有数据源输出的数据均采用标准化的经纬度坐标、气压层高度和时间格式，确保不同来源、不同格式的数据可直接拼接或对比，避免坐标不匹配导致的模型输入错误；

• 数据格式兼容层：工具包内置格式转换中间件，自动处理GRIB2的二进制编码、Zarr的分块存储、NetCDF的多维数组结构等底层差异，向上提供统一的Xarray接口，降低用户数据处理成本；

• 云数据访问优化：针对ARCO ERA5（GCS）、GFS（AWS）等云数据源，采用异步下载、数据缓存、按需加载等技术，减少网络传输延迟，提升数据获取效率，尤其适用于大规模批量推理场景。

3. 流水线执行机制

Earth2Studio的流水线执行遵循“组件解耦-接口统一-链式调用”的设计原则，核心流程如下：

1. 数据输入：通过DataSource组件获取标准化数据（Xarray格式），支持单时间点或时间序列输入；

2. 扰动处理（可选）：通过Perturbation组件对输入数据添加噪声，生成多个扰动样本；

3. 模型推理：PrognosticModel（预测模型）进行时间积分预测，或DiagnosticModel（诊断模型）进行空间映射计算，输出预测结果；

4. 结果评估（可选）：通过Statistics组件计算预报指标，评估模型性能；

5. 数据输出：通过IOBackend组件将结果存储为指定格式（NetCDF、Zarr等），支持本地或远程存储。

这一机制确保了各组件的独立性，用户可替换任意环节的组件（如用HRRR数据源替换GFS，用ZarrBackend替换NetCDF4Backend），而无需修改其他代码。

4. 扰动方法与不确定性量化

工具包提供的扰动方法可分为两类：

• 噪声型扰动：如Gaussian（高斯噪声）、Spherical Gaussian（球面高斯噪声），通过添加随机噪声模拟初始场的观测误差，适用于简单集合预报；

• 动力型扰动：如Bred Vector（繁殖向量）、Hemispheric Centred Bred Vector（半球中心繁殖向量），基于气象模型的动力特性生成扰动，更贴合真实气象系统的不确定性来源，适用于高精度集合预报。

5. 统计指标计算逻辑

内置指标均遵循气象预报评估的行业标准：

• RMSE（均方根误差）：衡量预报值与观测值的平均偏差，数值越小表示预报越准确；

• ACC（异常相关系数）：评估预报场与观测场的空间分布相似性，取值范围[-1,1]，越接近1表示模式一致性越好；

• CRPS（连续等级概率评分）：评估概率预报的准确性，数值越小表示概率预报越可靠；

• Spread-Skill Ratio（离散度-技巧评分比）：衡量集合预报的离散度与预报技巧的匹配程度，接近1表示集合校准良好。

四、应用场景

Earth2Studio的模块化设计和丰富功能，使其适用于气象、气候、环境等多个领域的实际应用，具体场景包括：

1. 短期天气预报（0-14天）

适用于气象部门、应急管理机构的日常预报工作。用户可选用GraphCast Operational、Pangu 6hr等全球高精度预测模型，结合GFS或HRRR数据源，快速输出未来14天的气温、气压、风速、降水等关键气象要素。对于区域高影响天气（如美国中部的强对流），可搭配StormCast模型（3km分辨率），实现小时级的精准预报，为防灾减灾提供决策支持。

2. 亚季节到季节（S2S）预报（2周-3个月）

针对农业、能源、水资源等行业的中长期规划需求，工具包新增的S2S预报流水线可实现2周至3个月的气候趋势预测。用户可选用AIFS、FuXi等适用于长周期预测的模型，结合ARCO ERA5再分析数据初始化，输出季节尺度的温度异常、降水偏多/偏少等趋势，帮助农业部门制定种植计划、能源企业优化电力调度。

3. 区域高精度气象预测

面向城市气象、航空航天、交通运输等场景的精细化需求，可组合“全球模型+区域诊断模型”实现高精度分析。例如：用GraphCast输出全球0.25°预报结果，再通过CBottleSR诊断模型（5km分辨率）进行超分优化，获得城市级的气温、风速分布；或用CorrDiffTaiwan模型（2km分辨率），专门针对台湾地区输出精细化天气数据，支持当地交通、电力等行业的运营调度。

4. 气候数据补全与超分

气候科研中常面临数据缺失、分辨率不足等问题，Earth2Studio的CBottle系列诊断模型可提供解决方案：

• CBottleInfill（100km分辨率）可补全缺失的全球气候数据样本，生成完整的气候场；

• CBottleSR（5km分辨率）可将低分辨率气候数据提升至区域高精度，助力气候极端事件（如历史高温、暴雨）的复盘分析。

5. 热带气旋（台风）监测与追踪

针对海洋气象领域的台风预警需求，TCTrackerVitart诊断模型可直接从预测模型输出的气象场中提取台风轨迹、强度、气压等关键参数，无需人工识别。用户可将其与GraphCast、Pangu等全球模型组合，实现台风路径的提前预报与实时追踪，为沿海地区的防台防汛工作提供技术支持。

6. 科研与教学应用

对于高校、科研机构而言，Earth2Studio是理想的气象AI研究与教学平台：

• 科研人员可基于工具包的模块化架构，快速验证新模型（如自定义Transformer气象模型），对比现有预训练模型的性能，加速科研成果转化；

• 教师可利用工具包的极简API和丰富示例，开展气象AI教学，让学生直观理解AI模型在气象预测中的应用，降低教学门槛。

五、使用方法

1. 环境准备与安装

Earth2Studio基于Python开发，支持主流Python版本（建议3.8+），且需依赖GPU环境（NVIDIA GPU，支持CUDA 11.0+）以实现加速推理。

安装命令：

pip install earth2studio[dlwp]

注：[dlwp]表示安装DLWP模型相关依赖，若需使用其他模型（如GraphCast、Pangu），可替换为对应依赖标签（如[graphcast]、[pangu]），或直接安装完整依赖：pip install earth2studio[all]。

2. 快速入门：确定性天气预测

以下代码示例展示如何用3行核心代码实现2024年1月1日的10步（每步6小时，共60小时）确定性天气预报：

# 1. 导入核心模块
from earth2studio.models.px import DLWP # 导入DLWP预测模型
from earth2studio.data import GFS # 导入GFS数据源
from earth2studio.io import NetCDF4Backend # 导入NetCDF4存储后端
from earth2studio.run import deterministic as run # 导入确定性运行函数

# 2. 初始化组件
model = DLWP.load_model(DLWP.load_default_package()) # 加载预训练DLWP模型
ds = GFS() # 初始化GFS数据源（自动获取AWS/NOMADS上的GFS数据）
io = NetCDF4Backend("output.nc") # 初始化NetCDF4存储后端，输出文件为output.nc

# 3. 运行预报流水线
run(
  start_times=["2024-01-01"], # 起始时间
  n_steps=10, # 预报步数（每步6小时，共60小时）
  model=model, # 预测模型
  data_source=ds, # 数据源
  io=io # 存储后端
)

运行结果：工具包会自动下载2024年1月1日的GFS初始数据，通过DLWP模型预测未来10步的气象场，将结果保存为output.nc文件，可直接用Xarray打开分析：

import xarray as xr
result = xr.open_dataset("output.nc")
print(result) # 查看预报的气象要素（温度、风速、气压等）和维度信息

3. 组件替换示例：切换模型与数据源

Earth2Studio支持灵活替换组件，无需修改整体代码结构：

（1）替换预测模型（如用GraphCast替代DLWP）

# 安装GraphCast依赖
# pip install earth2studio[graphcast]

# 导入GraphCast模型
from earth2studio.models.px import GraphCast

# 初始化GraphCast模型（默认加载1.0°分辨率版本）
model = GraphCast.load_model(GraphCast.load_default_package())

# 其余代码（数据源、存储后端、运行函数）保持不变
ds = GFS()
io = NetCDF4Backend("graphcast_output.nc")
run(start_times=["2024-01-01"], n_steps=10, model=model, data_source=ds, io=io)

（2）替换数据源（如用ARCO ERA5再分析数据替代GFS）

from earth2studio.data import ARCO_ERA5

# 初始化ARCO ERA5数据源（GCS存储的再分析数据）
ds = ARCO_ERA5()

# 其余代码保持不变
model = DLWP.load_model(DLWP.load_default_package())
io = NetCDF4Backend("era5_output.nc")
run(start_times=["2024-01-01"], n_steps=10, model=model, data_source=ds, io=io)

4. 进阶使用：构建集合预报流水线

以下示例展示如何结合扰动方法构建集合预报，捕捉预测不确定性：

# 导入扰动模块
from earth2studio.perturbation import Gaussian

# 初始化组件
model = GraphCast.load_model(GraphCast.load_default_package())
ds = GFS()
io = ZarrBackend("ensemble_output.zarr") # 用Zarr存储集合结果（支持并行写入）
perturb = Gaussian(scale=0.1) # 初始化高斯扰动（噪声强度0.1）

# 运行集合预报（生成5个扰动样本）
from earth2studio.run import ensemble as run_ensemble
run_ensemble(
  start_times=["2024-01-01"],
  n_steps=10,
  model=model,
  data_source=ds,
  io=io,
  perturbation=perturb,
  n_members=5 # 集合成员数（5个扰动样本）
)

# 计算集合评估指标（CRPS）
from earth2studio.statistics import crps
obs_ds = ARCO_ERA5().load_data(start_times=["2024-01-01"], n_steps=10) # 加载观测数据（ARCO ERA5）
ensemble_result = xr.open_zarr("ensemble_output.zarr") # 加载集合预报结果
crps_score = crps(ensemble_result, obs_ds) # 计算CRPS
print("集合预报CRPS评分：", crps_score)

六、常见问题解答（FAQ）

1. Earth2Studio需要什么硬件配置？

核心要求是NVIDIA GPU（支持CUDA 11.0+），推荐显存8GB以上（全球高分辨率模型如GraphCast Operational（0.25°）需16GB+显存）；CPU无特殊要求，建议4核以上以支持数据预处理；内存建议16GB以上，避免大规模数据加载时内存不足。

2. 工具包支持哪些AI框架？

兼容PyTorch、JAX、ONNX等主流框架：GraphCast基于JAX开发，Pangu基于ONNX，DLWP、StormCast等基于PyTorch，工具包通过统一API屏蔽了框架差异，用户无需手动适配。

3. 如何接入自定义模型？

工具包支持自定义模型扩展，需遵循以下步骤：① 继承earth2studio.models.PrognosticModel（预测模型）或DiagnosticModel（诊断模型）抽象类；② 实现forward（推理）、input_spec（输入规格）等核心方法；③ 注册模型到工具包的模型库。具体可参考官方贡献指南中的“自定义模型接入”章节。

4. 数据源需要手动下载吗？

不需要。工具包的数据源组件会自动从官方服务器（如AWS、GCS、NOMADS、ECMWF Opendata）下载数据，用户只需调用对应数据源接口（如GFS()、ARCO_ERA5()），工具包会处理下载、格式转换、坐标对齐等所有步骤。部分数据源（如CDS）可能需要用户注册API密钥并配置环境变量，具体参考数据源文档。

5. 输出数据可以用什么工具分析？

输出数据支持多种主流分析与可视化工具：① 科学计算：Xarray、Pandas（数据处理）；② 可视化：Matplotlib、Cartopy（绘图）、Panoply（NetCDF/Zarr可视化）；③ 批量处理：Dask（大规模并行分析）。所有输出数据均保留完整元数据，可直接导入上述工具使用。

6. Earth2Studio与其他气象AI工具（如WeatherBench2）有什么区别？

WeatherBench2是气象AI的基准测试数据集与评估框架，核心聚焦“模型评估”；而Earth2Studio是全栈式推理工具包，涵盖“数据获取-模型推理-存储-评估”全流程，且内置大量预训练模型，支持实际业务部署，两者可互补使用（如用WeatherBench2数据训练模型，在Earth2Studio中部署推理）。

七、相关链接

• 项目GitHub仓库：https://github.com/NVIDIA/earth2studio

八、总结

Earth2Studio是NVIDIA推出的一款面向天气与气候领域的开源Python工具包，以模块化设计和统一API为核心优势，整合了全球最大的气象AI预训练模型库、多源标准化数据源、灵活IO存储、扰动方法及评估指标，支持GPU加速与多框架兼容。它不仅降低了AI天气气候模型的使用门槛，让非专业用户也能快速实现高精度预报，还为科研人员提供了高效的模型验证与创新平台，适用于短期天气预报、中长期气候趋势预测、区域精细化分析、科研教学等多种场景。作为连接AI技术与气象行业的桥梁，Earth2Studio通过标准化、模块化的设计，解决了传统气象AI开发中数据分散、模型兼容差、开发效率低的痛点，为气象气候领域的AI应用普及与技术创新提供了强大支撑。