返回 文章 apply CMS 文章

NVIDIA Earth-2 教程:用 AI 数据同化将最新观测转化为及时天气决策

用 AI 数据同化把实时观测融入天气模型,让预报更及时、更本地化。

AI数据同化天气预报NVIDIA Earth-2Earth2Studio
成长分 / 100 81 综合收获、行动、留存与影响

NVIDIA Earth-2 教程:用 AI 数据同化将最新观测转化为及时天气决策
为什么值得读了解如何用观测数据约束扩散模型,在不重新训练的情况下改进区域天气预报。

掌握 CorrDiff-SDA、StormCast-SDA 和 HealDA 三种技术的实际代码实现。

关键洞察
  1. 基于分数的数据同化(SDA)通过扩散模型的多步去噪过程,在每一步将中间预测与观测比较,引导模型做出与观测一致的预测,无需重新训练。
  2. SDA 解锁两大能力:更快速地更新预报(持续融入观测而非按固定时间表),以及融入专有的、区域的或特定领域的观测数据。
  3. 在 CorrDiff-COSMO 降尺度示例中,SDA 将留出站点的风速 RMSE 降低了 54%;在 StormCast-CONUS 预报示例中,SDA 在六个时间步长上将留出站点的风速 RMSE 平均降低了 7.2%。
转成行动

深入阅读

正文与原文对照

原文保真覆盖:全文原文字符:24737

受天气影响的行业越来越多地能够获取观测数据,从而更早、更本地化地了解不断变化的条件。能源公司收集风能和太阳能资产的测量数据,应急管理团队依赖雷达和本地传感器,卫星提供商持续观测地球。这些数据帮助资本市场、保险、农业和物流等行业的组织理解和管理物理风险。

借助 NVIDIA Earth-2 中的 AI 数据同化工具,您可以更高效地处理这些观测数据。通过纳入专有或第三方数据,您可以使用这些工具更频繁地发布预报,使估计值与实时条件保持一致,并针对特定区域和应用定制您的预报流程。

本教程涵盖两种技术:

  • 用点观测约束扩散模型,通常用于区域模型。
  • 将不同数据集同化为一致状态,通常用于全球模型。

前提条件

本教程需要:

  • 已安装 Earth2Studio的开发环境 - 一块 NVIDIA RTX PRO 或数据中心 GPU
  • 基本的 Python 知识
  • 大约 30 分钟

用观测改进区域预报

您可能运行一个区域天气预报流程来管理能源生产和需求,从风电场和太阳能电站、输电走廊或人口密集地区获取观测数据。借助 AI 数据同化,您可以使用这些观测数据在本地精度最重要的地方约束您的预报,帮助您改进运营决策。同样的技术也可以支持其他行业,使用来自生产现场、活动场所、物流网络或其他资产的观测数据,这些地方的本地条件直接驱动决策。

您可以使用基于分数的数据同化(SDA)将观测数据纳入基于扩散的 AI 降尺度和预报模型,如 CorrDiff 和 StormCast。SDA 引导模型做出与您的观测一致的预测,而无需重新训练模型。

下面的图 1 展示了该过程在底层的工作原理。扩散模型通过一系列去噪步骤生成高分辨率预测。在每一步,SDA 将中间预测与您的观测进行比较,并将模型推向正确的方向。SDA 的输出是概率性的,观测位置附近的不确定性较小,而距离较远处分布更广,那里的预测越来越多地由其他模型输入和底层 AI 模拟决定。

SDA 过程示意图。两个输入从左侧进入模型:顶部是一叠天气场,底部是由分散的彩色点表示的点观测。顶行展示了六个步骤的扩散过程,从嘈杂的、类似静电的场逐步过渡到清晰的高分辨率天气场。底行显示了相应的 SDA 步骤,为六个条形图,条形高度逐渐降低,表示点观测与扩散状态之间差异的减小。该过程产生一叠高分辨率天气场,如右侧所示。

图 1. SDA 利用扩散模型的多步去噪过程。每一步都将中间结果与观测进行比较,以指导下一步去噪。最终输出是融合观测信息的高分辨率天气预报

为了引导模型,你需要定义一个观测算子,它将模型输出映射到你在每个测量位置期望观测到的量。对于温度或风速等物理量的原位测量,这尤其简单。在这种情况下,最简单的算子形式是将附近的网格值插值到每个观测位置。也可以为代理测量或观测影响创建算子。例如,风力涡轮机的功率输出可以作为风速的代理测量。

SDA 解锁了两大能力:

  • 更快速地更新预报。数值分析需要大量处理时间,并按固定时间表发布。SDA 使你能够持续融入观测。下面的图 2 展示了一个具体示例:一个以一小时为间隔进行预报的流水线,依赖于每六小时分发一次的全球分析。
  • 融入专有的、区域的或特定领域的观测。数值分析利用广泛的观测数据。SDA 允许你融入来自自己来源的数据,使预报聚焦于特定资产位置或下游应用。

SDA 的有效性取决于几个因素:观测的数量、空间分布和准确性;你正在预测的场的特征长度尺度;以及观测算子的质量和适定性。

示例 SDA 流水线在五个逐小时时间步上的流程图,从最新分析时间到当前时间之后一步。顶行显示低分辨率全球预报(FCN3 + 插值),它为中间行的高分辨率预报提供条件。底行显示观测,过去四个时间步可用,但未来最后一个时间步不可用。SDA 应用于流水线的两个部分:在第一个时间步对分析进行降尺度(CorrDiff + SDA),以及在第二至第四个时间步为高分辨率预报提供信息(StormCast + SDA)。在最后的未来时间步及之后,高分辨率预报在没有观测的情况下进行(StormCast)。

图 2. 使用 SDA 进行降尺度(CorrDiff + SDA)和高分辨率预报(StormCast + SDA)的示例流水线。CorrDiff 为 StormCast 创建高分辨率初始条件。SDA 通过将观测融入过去的降尺度和预报步骤,弥合分析参考时间与当前时间之间的差距

如何在 Earth2Studio 中运行 CorrDiff-SDA

CorrDiff 是一种基于 AI 的降尺度技术。Earth2Studio 提供了一个在欧洲预训练的 CorrDiff 模型,可将 0.25° 天气场转换为 2.2 公里预测。使用 AI 数据同化,你可以在局部精度重要的地方用观测改进这些预测。利用精炼后的输出,你可以初始化区域预报或创建用于校准下游模型的再分析数据集。

首先加载预训练模型。我们将区域限定为荷兰的一部分和德国西北部,并选择同化10米风速。

from datetime import datetime
from earth2studio.data import GHCNHourly
from earth2studio.models.da import CorrDiffCosmoEra5SDA
domain = dict(lat_min=50.2, lat_max=53.8, lon_min=4.6, lon_max=10.4)
sda = CorrDiffCosmoEra5SDA.load_model(
CorrDiffCosmoEra5SDA.load_default_package(),
assimilate_variables=("u10m", "v10m"),
resolution="rea2",
domain=domain,
number_of_samples=1,
sampler_steps=12,
amp=True,
).to("cuda")

获取用于低分辨率条件化的 ERA5 数据以及该区域内的 GHCN 风观测数据。

# 获取 ERA5 输入数据并将其重网格化到高分辨率区域网格上
# 完整实现请参考下方链接中的示例
init_time = datetime(2024, 1, 26)
x = fetch_and_regrid_era5(init_time, domain)
# 获取模型区域内的 GHCN 逐小时 10 米风观测数据
lat, lon = sda.model.lat_output_numpy, sda.model.lon_output_numpy
bbox = lat.min(), lon.min(), lat.max(), lon.max()
ghcn = GHCNHourly(stations=GHCNHourly.get_stations_bbox(bbox))
obs = ghcn(init_time, ["u10m", "v10m"]).dropna(subset=["observation"])

最后,使用输入数据运行模型。我们执行两次运行,以衡量额外观测值对结果的影响。

prior = sda(x) # 自由降尺度,无观测数据
analysis = sda(x, obs) # 引导扩散过程趋向观测数据

完整实现请参见 Earth2Studio 中的示例,上述代码即改编自该示例。

三张并排的荷兰与德国西北部地图,展示了有无 SDA 的风速降尺度示例。左图显示先验,即在不使用额外观测的情况下用 CorrDiff-COSMO 对 ERA5 进行降尺度得到的结果。中图显示使用 SDA 并利用 GHCN-Hourly 站点观测生成的观测引导分析。左图和中图中的风速以从蓝色到黄色的色标显示,范围从 0 到约 10 米/秒。中图中,57 个同化站点用黑点标记,24 个留出站点用白点标记。SDA 的效果在西北部北海区域最为明显,SDA 在那里预测了更高的风速。右图显示 SDA 与无 SDA 实验之间的差异,使用从蓝色到红色的色标,西北部的红色区域表示使用 SDA 时风速更高。在留出站点位置,绿色向上三角形表示误差减小,橙色向下三角形表示误差增大。三角形大小与误差变化幅度成正比。

图 3. 有无 SDA 的 CorrDiff-COSMO 降尺度。在此示例中,SDA 将留出站点的风速 RMSE 降低了 54%。左:无 SDA 降尺度得到的风速。中:使用 GHCN-Hourly 站点观测进行 SDA 降尺度得到的风速。右:SDA 与无 SDA 实验之间的风速差异

如何在 Earth2Studio 中运行 StormCast-SDA

StormCast 是一种类似于 CorrDiff 的技术,但专为高分辨率区域预报而设计。Earth2Studio 包含一个在美国本土(CONUS)上预训练的 StormCast 模型,该模型使用 HRRR 初始化,并以 3 公里分辨率进行预测。新的 HRRR 分析的计算和分发需要一些时间,但您可以使用 SDA 将当前可用的分析与最新观测相结合,以更新您的预报。

首先加载预训练模型。我们将范围限制在美国中部。

import numpy as np
from earth2studio.data import GHCNHourly
from earth2studio.models.px import StormCastCONUS
# Limit the domain to the central U.S.
hrrr_lat_lim, hrrr_lon_lim = (305, 785), (595, 1203)
model = StormCastCONUS.load_model(
StormCastCONUS.load_default_package(),
hrrr_lat_lim=hrrr_lat_lim, # comment out for full CONUS domain
hrrr_lon_lim=hrrr_lon_lim, # comment out for full CONUS domain
num_diffusion_steps=18,
num_sda_diffusion_steps=96, # more steps for SDA for better stability
sda_std_obs=0.15,
sda_gamma=1e-3,
).to("cuda")

接下来,获取用于模型初始化的 HRRR 分析场,并在模型域上定义观测数据源。

# 获取 HRRR 初始条件
# 完整实现请参考下方链接中的示例
init_time = datetime(2026, 4, 17, 18)
x, coords = fetch_hrrr(init_time)
# 为模型区域定义 GHCN 逐小时数据源
lat, lon = model.lat, model.lon
bbox = lat.min(), lon.min(), lat.max(), lon.max()
ghcn = GHCNHourly(
stations=GHCNHourly.get_stations_bbox(bbox),
time_tolerance=timedelta(minutes=15),
)

我们现在可以在初始展开步骤中使用观测运行模型,然后再过渡到无需额外观测的预报。与上文图2中的示意类似,该方法利用观测来弥合最新分析与当前条件之间的差距,此后预报便独立进行。

对于以HRRR初始化的流程,在新分析到达之前,通常只有一个受SDA影响的步骤是相关的。当使用全球分析进行初始化时,多个展开步骤可以从SDA中受益。

# 初始化生成器并获取第一个输出(分析透传)
gen = model.create_generator(x.clone(), coords.copy())
x, coords = next(gen)
# 使用 SDA 运行 rollout 的第一部分
for step in range(nsteps_sda):
valid_time = np.array(
[coords["time"][0] + coords["lead_time"][0] + np.timedelta64(1, "h")]
)
obs = ghcn(valid_time, ["u10m", "v10m", "t2m"])
x, coords = gen.send(obs) # 使用观测推进一步
# 在不使用 SDA 的情况下运行剩余的 rollout
for step in range(nsteps_non_sda):
x, coords = next(gen) # 不使用观测推进一步

你可以在 Earth2Studio 示例库中找到该示例的完整实现

六张地图以两行三列排列在美国中部上空,展示了使用和不使用 SDA 进行风速预测的示例。顶部和底部行分别显示 3 小时和 6 小时预测步骤。左列显示先验结果,通过运行 StormCast-CONUS 而未添加额外观测获得。中间列显示使用 SDA 并利用 GHCN-Hourly 站点观测生成的观测引导输出。左列和中间列的风速以从 0 到约 12 米/秒的蓝到黄色标尺显示。在中间列中,70% 的站点被同化并用黑点标记,其余 30% 被留出并用白点标记。右列显示 SDA 与无 SDA 实验之间的差异,使用蓝到红色标尺。最大差异出现在俄克拉荷马州、堪萨斯州和密苏里州的中部低地沿线。在留出站点位置,绿色向上三角形表示误差减小,橙色向下三角形表示误差增大。三角形大小与误差变化幅度成正比。

图 4. 使用和不使用 SDA 的 StormCast-CONUS 预测。在此示例中,SDA 在六个时间步长上将留出站点的风速 RMSE 平均降低了 7.2%。顶部和底部行分别显示 3 小时和 6 小时预测。左:无 SDA 的风速预测。中:使用 GHCN-Hourly 站点观测的 SDA 风速预测。右:SDA 与无 SDA 预测之间的风速差异

如何在自己的模型中使用 SDA

你可以通过扩展自定义模型的 Earth2Studio 模型包装器来同化观测。为此,请使用 PhysicsNeMo 中的扩散实用程序。我们从 x0_predictor 开始

,这是一个预训练的去噪扩散模型,它接收噪声样本及其噪声水平作为输入,并预测无噪声样本。在没有 SDA 的情况下,模型的扩散采样将这样实现:

from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler
from physicsnemo.diffusion.samplers import sample
# Construct diffusion scheduler
sigma_min, sigma_max = 0.01, 100
scheduler = EDMNoiseScheduler(sigma_min=sigma_min, sigma_max=sigma_max)
# Get denoiser from scheduler
denoiser = scheduler.get_denoiser(x0_predictor=x0_predictor)
# Generate sample
latents = sigma_max * torch.randn(shape)
sample(denoiser, latents, noise_scheduler=scheduler, num_steps=num_steps)

要使用 SDA,我们将 x0_predictor

转换为带有 SDA 引导的分数预测模型。我们使用 DataConsistencyDPSGuidance

,它将每个被掩码的像素与对应的观测值关联起来。你可以用它来同化来自气象站、专有传感器或类似基于点的来源的观测数据。

from physicsnemo.diffusion.guidance import (
DataConsistencyDPSGuidance,
DPSScorePredictor,
)
# Setup SDA guidance
guidance = DataConsistencyDPSGuidance(
mask=mask, # binary mask that identifies pixels with observations
y=y_obs, # gridded observations
std_y=sda_std_obs, # the remaining parameters are SDA settings
norm=sda_dps_norm,
gamma=sda_gamma,
sigma_fn=scheduler.sigma,
alpha_fn=scheduler.alpha,
)
# Convert x0_predictor to score predictor
score_predictor = DPSScorePredictor(
x0_predictor=x0_predictor,
x0_to_score_fn=scheduler.x0_to_score,
guidances=guidance,
)
denoiser = scheduler.get_denoiser(score_predictor=score_predictor)
# Generate sample (identical to non-SDA example)
latents = sigma_max * torch.randn(shape)
sample(denoiser, latents, noise_scheduler=scheduler, num_steps=num_steps)

对于更高级的 SDA 流水线,可使用 ModelConsistencyDPSGuidance 从多个网格点推导模拟观测值。此方法要求你提供一个 PyTorch 模型,将每个样本映射到对应的模拟观测值。借助自定义 PyTorch 模型,你还可以同化观测到的影响。例如,你可以使用风电模型来同化风机输出测量值。

如需完整实现,请查看 Earth2Studio 中的 StormCast CONUS 封装器,上述示例即基于此实现。

从观测计算全球天气

大多数全球天气预报流水线都使用通过数值数据同化得到的当前天气估计值进行初始化。数值数据同化的计算需求很高,这降低了预报的时效性和刷新率,也使得集成自定义观测变得更加困难。

借助一种名为 HealDA 的基于 AI 的技术,你可以在几秒钟内估算出全球大气状态。这使你能够发布更接近当前状况的预报,或计算自定义的再分析。

HealDA 将时间窗口内的遥感和原位观测映射到全球网格化大气状态。它由两个主要组件构成:观测编码器和视觉 Transformer(ViT)主干。编码器将异构观测作为点云摄入,将每个标量值与地理位置和时间等元数据一起嵌入为一个 token。这些 token 随后被聚合到目标网格上,并由 ViT 主干进行处理。

HealDA 示意图。左侧以一组地球堆叠图展示输入观测,其中包含卫星扫幅和点观测。观测被传递到观测编码器,该编码器由独立的传感器嵌入器及随后的传感器融合步骤组成。得到的表示被传递到 ViT 主干。右侧以一组地球堆叠图展示最终分析输出,其中包含稠密天气场。

图 5. HealDA 结合来自不同平台的原位和遥感观测,提供一致的全球天气估计。观测编码器将每个数据流视为空间和时间中的点云,并将测量值转换为 token,再由 ViT 主干进行处理

你可以使用预训练的全球数据同化模型作为起点。如果你有自定义的常规观测,通常无需修改模型即可将其纳入。

对于专有卫星数据,你可以调整编码器以支持你的数据源。这种灵活性让你能够针对自己的区域或应用定制数据同化系统。你可以使用相同的技术来训练区域系统而非全球系统。要开始使用,请参阅开源 Python 库 PhysicsNeMo 中的 HealDA 训练流水线。

如何在 Earth2Studio 中运行 HealDA

Earth2Studio 提供了一个用于研究目的的预训练全球数据同化模型。它将来自微波探测仪、无线电掩星、地面站、飞机、浮标及其他来源的数据集成到 1° HEALPix 网格(HPX64)上。

首先,加载模型。

from datetime import timedelta
import numpy as np
from earth2studio.data import UFSObsConv, UFSObsSat, fetch_dataframe
from earth2studio.models.da import HealDA
model = HealDA.load_model(
HealDA.load_default_package(),
lat_lon=True, # regrid from HEALPix to regular lat/lon
).to("cuda")

接下来,从 NOAA UFS 回放仓库中获取输入观测数据。我们使用常规观测和卫星观测。

# HealDA 是在 UFS 回放窗口上训练的:分析时间前 21 小时到后 3 小时
time_tolerance = (timedelta(hours=-21), timedelta(hours=3))
analysis_time = np.array([np.datetime64("2024-01-01T00:00")])
# input_coords() 返回两个观测 DataFrame 必须满足的模式
conv_schema, sat_schema = model.input_coords()
# fetch_dataframe 附加模型所需的 request_time 元数据
conv_df = fetch_dataframe(
UFSObsConv(time_tolerance=time_tolerance),
time=analysis_time,
variable=np.array(conv_schema["variable"]),
fields=np.array(list(conv_schema.keys())),
)
sat_df = fetch_dataframe(
UFSObsSat(time_tolerance=time_tolerance),
time=analysis_time,
variable=np.array(sat_schema["variable"]),
fields=np.array(list(sat_schema.keys())),
)

然后使用观测数据帧调用模型。

# 无状态模型 - 直接调用以进行一次性分析,或使用
# create_generator 进行循环同化
analysis = model(conv_obs=conv_df, sat_obs=sat_df)

你可以在 Earth2Studio 示例库中找到运行 HealDA 的扩展示例

使用 Earth2Studio 访问观测数据

Earth2Studio 让你能够访问广泛的数据源,用于开发、初始化和验证天气模型,包括来自不同平台和传感器类型的观测数据。

其中包括来自地球静止卫星(GOES、Himawari、Meteosat)和雷达网络(MRMS、OPERA)的网格化数据,你可以直接使用这些数据来训练和快速更新区域高分辨率预报模型,例如 StormScope。当你想要预报依赖于日照或降水的量(如太阳能发电、冷却过程和水库入流)时,这些数据源尤其有用。

为了开发和基准测试数据同化系统,Earth2Studio 还让你能够访问常规观测数据的存档,如 GHCN/ISD、NNJA 和 UFS,以及来自 GDAS 和 ASOS 的业务观测数据。这些数据源以数据帧的形式提供温度和风速等变量。来自极轨卫星系统(包括 MetOp 和 JPSS)的观测数据也可供使用。

Earth2Studio 为所有数据源提供统一的接口。你实例化一个数据源对象,然后用时间步长和变量名称的列表来调用它。预报数据源还接受一个预报时效列表。这种一致的接口使得在同一工作流中组合多个数据源或将你自己的观测数据连接到流水线变得非常容易。

era5 = NCAR_ERA5()
da_era5 = era5(datetime(2025, 7, 15), ["t2m", "z500"])
print(da_era5.shape) # (1, 2, 721, 1440)
ifs = IFS_FX()
da_ifs = ifs(datetime(2026, 7, 15), timedelta(hours=48), ["t2m"])
print(da_ifs.shape) # (1, 1, 1, 721, 1440)
goes = GOES(satellite="goes19", scan_mode="C")
da_goes = goes(datetime(2026, 7, 15), ["abi01c", "abi02c", "abi03c"])
print(da_goes.shape) # (1, 3, 1500, 2500)
ghcn = GHCNHourly(stations=["USW00013301"])
df_ghcn = ghcn(datetime(2026, 6, 15), ["t2m", "ws10m"])
print(df_ghcn.shape) # (10, 7)

有关支持的数据源的完整列表,请参阅用户指南中的 API 参考

Earth2Studio 入门

在 Earth2Studio 示例库中探索端到端 AI 数据同化示例。要将您自己的观测数据接入流水线,请遵循自定义数据源示例

AI 数据同化让您能够纳入对您所在地区或组织至关重要的观测数据,从而发布更准确、更及时的预报。

访问 Earth2Studio 用户指南,开始使用 AI 数据同化,并探索 AI 天气模型的更广泛能力。