用球面傅立叶神经算子模拟地球大气层

这篇文章最初发表在 NVIDIA 技术博客上。

基于机器学习的天气预测已经成为传统数值天气预测（NWP）模型的有力补充。例如，NVIDIA FourCastNet已经证明，生成天气预报的计算时间可以从几个小时缩短到几秒钟，这对当前基于NWP的工作流程带来了重大改进。

传统方法通常根据第一性原理制定，并需要时间步限制以保证基础数值方法的准确性。而基于 ML 的方法没有这样的限制，其统一的内存访问模式非常适合 GPU。

然而，这些方法纯粹是数据驱动的，你可能会问：

我们怎么能信任这些模型呢？
他们的概括能力如何？
如果它们不是根据第一原则制定的，我们如何进一步提高它们的技能、可信度和可解释性？

在这篇文章中，我们讨论了球面傅立叶神经算子（SFNO）、球面上的物理系统、对称性的重要性，以及如何使用SFNO的实现和球面谐波变换（SHT）。如果你想了解更多关于这个数学主题的信息，请参阅 ICML 论文，Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere。

对称性的重要性

图 1 。 SFNO 推出 5 个月。将 SFNO 和地面实况数据的地面风速预测相互比较。

创建有原则和值得信赖的模型的一种潜在方法是以类似于物理定律的方式来制定模型。

物理定律通常是根据对称性考虑制定的：

我们不希望物理学依赖于参考系。
我们进一步预计，如果参照系发生变化，基本物理定律将保持不变。

在球体上的物理系统的上下文中，参考系的变化是通过旋转来实现的。因此，我们努力建立一个在轮换下保持不变的公式。

当前基于 ML 的天气预测模型将大气状态视为一系列离散的矢量，这些矢量表示随着时间的推移在不同空间位置感兴趣的物理量。这些向量中的任何一个都由学习函数更新，该函数将当前状态映射到序列中的下一个状态。

简单地说，我们要求神经网络在显示今天的天气时，连续预测下一个时间步长的天气。这相当于使用传统方法对物理系统进行集成，但需要注意的是，要以纯粹的数据驱动方式学习动力学，而不是从物理定律中推导动力学。与传统方法相比，这种方法能够实现更大的时间步长。

因此，手头的任务可以理解为学习有限维向量空间之间的图像到图像的映射。

虽然 U-Nets 等多种神经网络拓扑结构适用于这项任务，但这种方法忽略了问题的功能性质。输入和输出都是函数，它们的演化由偏微分方程控制。

传统的 ML 方法（如 U-Nets ）忽略了这一点，因为它们以固定的分辨率学习地图。神经算子推广神经网络来解决这个问题。他们学习的不是有限维空间之间的映射，而是可以直接将一个函数映射到另一个函数的算子。

像Fourier neural operators（ FNO ）这样的技术，为学习函数空间之间的映射和逼近偏微分方程的解算子提供了一个强大的框架，该算子将一种状态映射到下一种状态。

然而，经典的 FNO 是在笛卡尔空间中定义的，其相关的对称性与球体的对称性不同。在实践中，由于模型的自回归性质，忽略几何结构并假装地球是一个周期性矩形会导致伪影，这些伪影会在长时间的推出中积累。这种伪影通常发生在极点周围，并导致模型崩溃（图 2 ）。

你现在可能会想，球体上的 FNO 会是什么样子？

图 2 :使用 AFNO 与 SFNO 的温度预测

图 2 显示了与球面傅立叶神经算子（ SFNO ）相比，使用自适应傅立叶神经算子的温度预测。尊重球形几何体和相关的对称性可以避免伪影，并实现稳定的卷展栏。

球面傅立叶神经算子

为了尊重地球的球形几何形状，我们提出了球面傅立叶神经算子（ SFNO ），这是一种直接在球面坐标中公式化的傅立叶神经算子。为了实现这一点，我们使用了在球体上公式化的卷积定理。

FNO 的核心构建块是全局卷积。通过 FFT 的计算，我们可以利用卷积定理这个强大的数学工具，将卷积与傅立叶变换连接起来。

类似地，球面上的卷积定理将球面卷积与球面上傅立叶变换的推广联系起来：球面谐波变换（ SHT ）。

实现可微球面谐波变换

为了实现 SFNO，我们需要一个可微分的 SHT。为此，我们开发了torch-harmonics，这是一个用于可微 SHT 的 PyTorch 库。该库能够在单个和多个 GPU 以及 CPU 上进行 SHT 计算，以实现模型的可扩展并行性。torch-harmonics可以通过运行以下命令进行简单安装：

复制代码

pip install torch-harmonics

torch-harmonics 可以与 PyTorch 无缝集成。可微分 SHT 可以轻松地作为一个模块集成到任何现有的 ML 架构中。如果要计算随机函数的 SHT，请参考以下代码示例：

ini 复制代码

import torch
import torch_harmonics as th

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# parameters
nlat = 512
nlon = 2*nlat
batch_size = 32
signal = torch.randn(batch_size, nlat, nlon)

# create SHT instance
sht = th.RealSHT(nlat, nlon).to(device).float()

# execute transform
coeffs = sht(signal)

我们建议您开始使用torch-harmonics，并参考getting-started notebook，它将指导您如何计算火星高程图的球面谐波系数（图 3 ）。该示例展示了如何使用 SHT 和 ISHT 的可微性来计算系数。

图 3 。火星高程图的球面谐波系数，用火炬谐波计算（左）。用反球面谐波变换计算的重构信号（右）。

对基于 ML 的天气预报的启示

我们使用了由欧洲中期天气预报中心（ECMWF）提供的ERA5数据集。这个数据集代表了我们对过去44年地球大气层状况的最佳理解。图2显示，SFNO在极点上没有显示出伪影的迹象，并且在长达一年的时间里，在数千个自回归步骤上，表现出了显著的稳定性（见图1）。

这些结果为部署基于 ML 的天气预测方法铺平了道路。它们让我们得以一窥基于 ML 的方法如何成为弥合天气预报和气候预报之间差距的关键，成为亚季节性到季节性预报的圣杯。

在一台 NVIDIA RTX A6000 上，在 13 分钟内计算出一年内单次推出 SFNO ，包括 1460 个自回归步骤。这比传统的数值天气预报方法快一千多倍。

这种速度快得多的预测工具为在使用传统 NWP 进行单一预测的同时计算数千种可能的情景打开了大门，从而能够对罕见但影响大的极端天气事件的风险进行更高的置信度预测。

用球面傅立叶神经算子模拟地球大气层

对称性的重要性

球面傅立叶神经算子

实现可微球面谐波变换

对基于 ML 的天气预报的启示

更多关于 SFNO 和 NVIDIA Earth-2 计划的信息