基于機器學習的天氣預測已經成為傳統數值天氣預測（NWP）模型的有力補充。例如，NVIDIA FourCastNet已經證明，生成天氣預報的計算時間可以從幾個小時縮短到幾秒鐘，這對當前基于NWP的工作流程帶來了重大改進。
傳統方法通常根據第一性原理制定，并需要時間步限制以保證基礎數值方法的準確性。而基于 ML 的方法沒有這樣的限制，其統一的內存訪問模式非常適合 GPU。
然而，這些方法純粹是數據驅動的，你可能會問：
- 我們怎么能信任這些模型呢？
- 他們的概括能力如何？
- 如果它們不是根據第一原則制定的，我們如何進一步提高它們的技能、可信度和可解釋性？
在這篇文章中，我們討論了球面傅立葉神經算子（SFNO）、球面上的物理系統、對稱性的重要性，以及如何使用SFNO的實現和球面諧波變換（SHT）。如果你想了解更多關于這個數學主題的信息，請參閱 ICML 論文，Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere。
對稱性的重要性
創建有原則和值得信賴的模型的一種潛在方法是以類似于物理定律的方式來制定模型。
物理定律通常是根據對稱性考慮制定的：
- 我們不希望物理學依賴于參考系。
- 我們進一步預計，如果參照系發生變化，基本物理定律將保持不變。
在球體上的物理系統的上下文中，參考系的變化是通過旋轉來實現的。因此，我們努力建立一個在輪換下保持不變的公式。
當前基于 ML 的天氣預測模型將大氣狀態視為一系列離散的矢量，這些矢量表示隨著時間的推移在不同空間位置感興趣的物理量。這些向量中的任何一個都由學習函數更新，該函數將當前狀態映射到序列中的下一個狀態。
簡單地說，我們要求神經網絡在顯示今天的天氣時，連續預測下一個時間步長的天氣。這相當于使用傳統方法對物理系統進行集成，但需要注意的是，要以純粹的數據驅動方式學習動力學，而不是從物理定律中推導動力學。與傳統方法相比，這種方法能夠實現更大的時間步長。
因此，手頭的任務可以理解為學習有限維向量空間之間的圖像到圖像的映射。
雖然 U-Nets 等多種神經網絡拓撲結構適用于這項任務，但這種方法忽略了問題的功能性質。輸入和輸出都是函數，它們的演化由偏微分方程控制。
傳統的 ML 方法（如 U-Nets ）忽略了這一點，因為它們以固定的分辨率學習地圖。神經算子推廣神經網絡來解決這個問題。他們學習的不是有限維空間之間的映射，而是可以直接將一個函數映射到另一個函數的算子。
像Fourier neural operators（ FNO ）這樣的技術，為學習函數空間之間的映射和逼近偏微分方程的解算子提供了一個強大的框架，該算子將一種狀態映射到下一種狀態。
然而，經典的 FNO 是在笛卡爾空間中定義的，其相關的對稱性與球體的對稱性不同。在實踐中，由于模型的自回歸性質，忽略幾何結構并假裝地球是一個周期性矩形會導致偽影，這些偽影會在長時間的推出中積累。這種偽影通常發生在極點周圍，并導致模型崩潰（圖 2 ）。
你現在可能會想，球體上的 FNO 會是什么樣子？
圖 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，請參考以下代碼示例：
import torchimport torch_harmonics as thdevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')# parametersnlat = 512nlon = 2*nlatbatch_size = 32signal = torch.randn(batch_size, nlat, nlon)# create SHT instancesht = th.RealSHT(nlat, nlon).to(device).float()# execute transformcoeffs = sht(signal) |
我們建議您開始使用torch-harmonics，并參考getting-started notebook，它將指導您如何計算火星高程圖的球面諧波系數（圖 3 ）。該示例展示了如何使用 SHT 和 ISHT 的可微性來計算系數。
對基于 ML 的天氣預報的啟示
我們使用了由歐洲中期天氣預報中心（ECMWF）提供的ERA5數據集。這個數據集代表了我們對過去44年地球大氣層狀況的最佳理解。圖2顯示，SFNO在極點上沒有顯示出偽影的跡象，并且在長達一年的時間里，在數千個自回歸步驟上，表現出了顯著的穩定性（見圖1）。
這些結果為部署基于 ML 的天氣預測方法鋪平了道路。它們讓我們得以一窺基于 ML 的方法如何成為彌合天氣預報和氣候預報之間差距的關鍵，成為亞季節性到季節性預報的圣杯。
在一臺 NVIDIA RTX A6000 上，在 13 分鐘內計算出一年內單次推出 SFNO ，包括 1460 個自回歸步驟。這比傳統的數值天氣預報方法快一千多倍。
這種速度快得多的預測工具為在使用傳統 NWP 進行單一預測的同時計算數千種可能的情景打開了大門，從而能夠對罕見但影響大的極端天氣事件的風險進行更高的置信度預測。
更多關于 SFNO 和 NVIDIA Earth-2 計劃的信息
要了解 SFNO 是如何產生數千名樂團成員并預測 2018 年阿爾及利亞熱浪的，請觀看以下視頻：
有關 SFNO 的更多信息，請參閱以下資源：
- 我們的 ICML 海報
- 查閱Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere的白皮書。
- NVIDIA/torch-harmonics 的 GitHub 項目
- SFNO的實現
- 開始使用 torch-harmonics
- 在球形淺水方程上訓練 SFNO
- /neuraloperator 的 GitHub 倉庫
- 在neuraloperator中對淺水方程進行SFNO訓練
- 在 PhysicsNeMo 中的 SFNO 實現
有關 NVIDIA Earth-2 計劃的更多信息，請參閱以下資源：
?
