AI 平台/部署

使用 NVIDIA Earth-2 预测两周以上的天气

能够预测极端天气事件至关重要,因为此类条件变得更加常见且更具破坏性。次季节性气候预测 (预测未来两周或两周以上的天气) 是对天气波动敏感的各行业主动决策和风险管理的基础。

它可以帮助农民更好地选择种植哪些作物,并管理干旱易发地区的水资源。电力公司可以平衡能源供需,而渔业可以保护自己免受海洋热浪的影响。此外,政府可以为自然灾害和公共卫生威胁做好准备,例如在亚季节性前景较差的地区预先配置移动消防和缓解热风险的基础设施。

在过去两年中,使用 AI 模型预测天气和气候在研究中获得了巨大的动力,现在在操作环境中越来越受欢迎。NVIDIA Earth-2 平台通过提供高性能和可扩展的工具堆栈,一直在为科学和企业社区提供支持。无论是希望评估和验证模型技能的气象专家,还是尝试为各种用例和数据集开发、定制和扩展模型的 AI/ ML 专家,它都能让所有人受益。

在本文中,我们将概述 Earth-2 平台为气象领域专家提供的实用程序,这些实用程序用于开发和验证用于概率亚季节预测的大型集成,且计算成本远低于传统的非机器学习技术。

借助 AI 进行次季节预测

AI 天气模型的主要优势之一是,与传统方法相比,它能够以低几个数量级的计算成本运行更大的操作集成。今年早些时候,加州大学伯克利分校的研究人员展示了一种使用育种向量/ 多检查点 (BVMC) 方法生成经过良好校准的数千人集成 (“Huge Ensemble”,简称 HENS) 的有效方法。JBA 和 AXA 等企业正在将这种 HENS 方法与 FourCastNet V2 (SFNO) 模型结合使用,以在保险应用程序中进行后投射。

最新版本的 Earth2Studio 引入了在深度学习地球系统模型 (DLESym) 背景下演示的新的次季节到季节 (S2S) 预测功能。这是一个简洁的深度学习模型,将多层大气 AI 模型与预测海面温度变化的单独海洋 AI 模型相结合。

模型架构是一个带有填充操作的 U-Net,经过修改后,可支持使用分辨率约为 1 度的 HEALPix 网格。由于此架构基于不使用位置嵌入的本地模板,因此具有泛化潜力。该模型已证明能够逼真地在多个月的时间尺度上接近预期的气候错误率,华盛顿大学的研究人员也已证明它能够在气候尺度模拟中实现出色的自回归稳定性

在以下代码段中,您可以看到使用该模型生成亚季节预报的难度。此处可在 Earth2Studio 中获取完整实现。

# Prepare model, data source, and I/O backend
package = DLESyMLatLon.load_default_package()
model = DLESyMLatLon.load_model(package).to(device)
data = ARCO()
io = KVBackend()

# 60-day forecast, initialized in June 2021
ic_date = np.datetime64("2021-06-15")
n_steps = 16

# Prepare coordinates for forecast outputs
input_coords = model.input_coords()
output_coords = model.output_coords(input_coords)
inp_lead_time = input_coords["lead_time"]
out_lead_times = [
    output_coords["lead_time"] + output_coords["lead_time"][-1] * i
    for i in range(n_steps)
]
output_coords["lead_time"] = np.concatenate([inp_lead_time, *out_lead_times])

# Run forecast
io = run.deterministic(
    [ic_date], n_steps, model, data, io, output_coords=output_coords
)

使用集成进行概率预测

但是,S2S 预测本质上是概率性的,而非确定性的。它们不会提前预测特定日期的确切天气,而是提供季节条件偏离标准的可能性。这些预测通常用万亿次概率表示:对于温度或降水等变量,即将到来的季节处于历史气候分布的上三分之一 (高于正常水平) 、中三分之一 (接近正常水平) 或下三分之一 (低于正常水平) 的可能性。

在这一新模型推出之前,企业已经通过 FourCastNet V2 (SFNO) 模型扩展了 HENS 方法,以进行 S2S 预测。加利福尼亚大学尔湾分校的研究人员已经表明,它与 ECMWF 预测系统在预测麦登 – Julian 振动 (MJO) 方面的能力一样出色;MJO 是大气中 S2S 可预测性的主要来源。

现在,Earth2Studio 为有兴趣尝试 HENS-SFNO、DLESyM 或其他模型以进行 S2S 预测的用户提供了新的 S2S recipe。为反映对更大集成和更长预测时间尺度的需求,此 recipe 支持多 GPU 分布式推理以及并行 I/ O,以便在生成预测数据时高效保存。如果存储空间受限,它还允许仅保存预测输出的子集。为了简化此 recipe 的使用,就像 Earth2Studio 中的 HENS recipe 一样,运行集成的复杂方面已经得到了处理。控制行为相当于指定配置:

# DLESyM ensemble with 16 total checkpoint combinations
nperturbed: 4
ncheckpoints: 16
batch_size: 4
defaults:
    - forecast_model: dlesym
    - perturbation: gaussian

借助这一新方法,领域专家现在可以从 HENS FourCastNet V2 (SFNO) 和 DLESym 生成大型集成预报,以了解和验证这些模型的技能。例如,您可以探索预测不确定性如何由干扰驱动到初始条件或替代模型检查点权重。这样,您就可以生成一个巧妙、经过校准的亚季节预报集合。以此为基础,您可以探索其他策略,在 S2S 时间尺度上实现最佳 AI 预测校正。

作为一个示范性示例,该方法可用于为 2021 年太平洋西北热浪生成 S2S 预报,如图 1 所示。《地球物理研究快报》 ( Geophysical Research Letters) 上发表的一篇论文指出,这一前所未有的事件在极端高温的强度和持续时间方面非常显著,并且难以在 S2S 时间尺度上进行预测。虽然没有一个模型能够完美地捕捉热浪的位置和强度,但我们可以看到,所有模型早在三周前就开始预测北美的某种程度的暖异常,准确度在 HENS-SFNO、IFS ENS 和 DLESym 之间存在差异。

Four maps comparing ECMWF IFS and earth2studio forecast models evaluated on S2S timescales for the 2021 Pacific Northwest heatwave. The ground truth ERA5 shows extreme heat anomalies centered over western Canada, and each of the IFS, HENS-SFNO, and DLESyM models show warm anomalies in the general vicinity but none fully capture the location and intensity of the event.
图 1。2021 年太平洋西北热浪期间 (左上方为逆时针方向) IFS ENS ( 11 个后投成员,已下载 ECMWF API) 、SFNO-HENS 和 DLESym 之间的每周平均 S2S 预报的样本比较,以及预报第 3 周的相应 ERA5 数据。所有模型都能预测出北美的一定程度的暖异常,但由于交付周期过长,我们很难同时捕获到这种极端高温的确切位置和强度。

下一步是什么?

加速采用 AI 进行 S2S 预测需要领域专家对此类模型及其功能进行更可靠的评估。提供开源库可以降低进入 AI 领域所需技能的门槛。它还向 AI/ ML 研究社区提供有关模型未来开发的反馈。

欧洲中期天气预报中心 (ECMWF) 举办的 AI Weather Quest 竞赛旨在加速社区参与推进 S2S 预测。在 NVIDIA 工程师 (以及华盛顿大学的研究人员) 准备参与本次竞赛的同时,我们正在研究 Earth-2 工具与 ECMWF 为 Weather Quest 比赛提供的工具的可组合性,以使社区能够参与其中。这将允许使用 ECMWF 的 AI-WQ 软件包直接与 Earth2Studio 中生成的预测数据更快地评估模型,并能够在 PhysicsNeMo 中训练自定义模型。这些是 NVIDIA 研究团队使用的相同工具,我们希望分享这些工具将使其他研究人员能够快速迭代他们的想法。

# Score an earth2studio-generated forecast using ECMWF AI-WQ routines
for var in aiwq_variables:
    fcst_data, fcst_coords = load_forecast_for_aiwq(io_backend, ic, var)
    rpss_wk3, rpss_wk4 = compute_aiwq_rpss(fcst_data, fcst_coords, var)
    write_aiwq_scores(score_io_dict, rpss_wk3, rpss_wk4, var)

一般来说,大型 S2S 集成预报的高效推理和评分是科学过程的重要组成部分。正确评估模型需要对多个预测进行评分,以确定其技能。为了加速这一资源密集型流程,Earth2Studio 现在可以高效运行大型 S2S 集成并为其打分。例如,在 8 个 GPU 上,使用多个大气和海洋模型在全年运行 DLESym 集成预报并进行评分只需不到两个小时。

有关这些评分结果的示例,请参阅图 2,该示例还表明,DLESym 模型具有 S2S 技能,在第三周到第五周内即可与 ECMWF IFS 竞争,这是一个强大的物理基准。我们将在 Earth2Studio 的 S2S recipe 中发布这些通用评分功能以及 AI 天气任务特有的评分功能。这为从业者提供了各种方法来评估他们有兴趣尝试的模型的性能。

A chart comparing DLESyM and IFS ENS for fair CRPS of z500 evaluated in 2018. In weeks three through five, the DLESyM model nearly matches the IFS ENS fCRPS score, and in the second week, it lags behind.
图 2。Fair CRPS z500 在 2018 年针对 IFS ENS 和 DLESym 模型进行了每周平均 S2S 预测评估。总体而言,DLESym 与 IFS 竞争激烈,但由于模型分布缩小,其技能在前几周有所落后。

要点

S2S 预测对于众多气候敏感领域至关重要。本文讨论了 Earth2Studio 中的关键新功能,使企业能够评估和验证预训练的大气 – 海洋合 AI 预测模型 (例如 DLESym) ,以生成集成预报。

以下是入门资源:

通过这些 GTC 会议详细了解 Earth-2 平台。这些资源提供了有关企业如何使用 AI 生成大型集合预报的更多见解:

 

 

标签