Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere (2306.03838v1)

Abstract: Fourier Neural Operators (FNOs) have proven to be an efficient and effective method for resolution-independent operator learning in a broad variety of application areas across scientific machine learning. A key reason for their success is their ability to accurately model long-range dependencies in spatio-temporal data by learning global convolutions in a computationally efficient manner. To this end, FNOs rely on the discrete Fourier transform (DFT), however, DFTs cause visual and spectral artifacts as well as pronounced dissipation when learning operators in spherical coordinates since they incorrectly assume a flat geometry. To overcome this limitation, we generalize FNOs on the sphere, introducing Spherical FNOs (SFNOs) for learning operators on spherical geometries. We apply SFNOs to forecasting atmospheric dynamics, and demonstrate stable auto-regressive rollouts for a year of simulated time (1,460 steps), while retaining physically plausible dynamics. The SFNO has important implications for machine learning-based simulation of climate dynamics that could eventually help accelerate our response to climate change.

• The paper introduces SFNOs that replace DFT with spherical harmonic transforms to accurately simulate global climate dynamics.
• It demonstrates enhanced long-term stability through year-long autoregressive forecasts and outperforms traditional FFT-based methods.
• Results on Spherical SWEs and ERA5 data highlight SFNOs' ability to maintain physical symmetries while enabling rapid, GPU-accelerated inference.

Analyzing Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere

This review explores the Spherical Fourier Neural Operators (SFNOs) introduced by Bonev et al., as a geometric extension to Fourier Neural Operators (FNOs) for spherical topologies. This paper presents an innovative approach to machine learning-based simulation of climate dynamics, primarily targeting atmospheric weather prediction on a global scale.

Rationale and Approach

Fourier Neural Operators have been effective in representing long-range dependencies within spatio-temporal datasets through global convolution operations. These operators, however, traditionally utilize Discrete Fourier Transform (DFT), which presupposes a Euclidean geometry, inappropriately applying it to spherical environments. Such use leads to artifacts and instabilities, especially around poles in spherical datasets.

The authors propose a solution by introducing SFNOs, leveraging a Spherical Harmonic Transform (SHT) instead of DFT, thus aligning the neural operators with the natural symmetries and geometries of the sphere, specifically the manifold $S^2$. These spherical harmonics provide a suitable orthogonal basis to decompose functions while maintaining rotational equivariance, a desirable property reflecting how physical laws remain unchanged under rotations.

Architecture and Features

SFNO's architecture retains several advantages of standard FNOs, such as computational efficiency and grid-independence. Unlike traditional methods, SFNOs respect spherical symmetries, resulting in enhanced long-term stability over extended forecast periods. This characteristic is demonstrated through successful auto-regressive rollouts spanning up to a year without introducing non-physical artifacts.

The network design comprises:

• Spherical Convolutions: Utilizing a convolution theorem derived on the sphere allows for an efficient and equivariant treatment of spherical data, maintaining rotational symmetry.
• Point-wise Nonlinearities: These are applied to manage localized transformations in the spatial domain.
• Grid-Invariance: By not binding transformations to specific grid points, SFNOs offer flexibility concerning input data grids and resolutions.

Numerical Results and Practical Implications

The SFNO's application on both the Spherical Shallow Water Equations (SWEs) and ERA5 weather data showcases its efficacy. The paper reports that the SFNO outperforms traditional FFT-based approaches, especially regarding long-term stability and preservation of physical dynamics over extensive time steps. For instance, eigenmetric evaluations confirm SFNO’s improved prediction fidelity for atmospheric phenomena, aligning closely with IFS benchmarks over medium-range (5-10 day) forecasts.

The computational efficiency is noteworthy; a year-long atmospheric simulation using SFNOs is realized within minutes on a single GPU, demonstrating a significant speed advantage over classical methods. This rapid inference capability emphasizes their potential role in large-scale climate studies, reducing the barrier for conducting long-time horizon climate simulations.

Future Prospects

The integration of SFNOs may herald significant advancements in data-driven atmospheric modeling, enabling more accurate climate predictions and contributing to enhanced modeling of Earth's climate systems. The promising results suggest potential applications ranging from seasonal forecasting to climate change modeling, where stable, long-term predictive capabilities are crucial.

From a theoretical perspective, exploring generalized Fourier transforms across different manifolds could unveil new machine learning frontiers, extending beyond mere atmospheric applications to diverse scientific fields requiring spherical or non-Euclidean modeling.