Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere

Boris Bonev; Thorsten Kurth; Christian Hundt; Jaideep Pathak; Maximilian Baust; Karthik Kashinath; Anima Anandkumar

Source-linked AI summary

Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere

Boris Bonev, Thorsten Kurth, Christian Hundt, Jaideep Pathak, Maximilian Baust, Karthik Kashinath, Anima Anandkumar

arXiv:2306.03838v1Jun 6, 2023cs.LGmath.NAphysics.ao-phphysics.comp-ph

TL;DR

Flat-geometry Fourier transforms can create artifacts and dissipation when FNOs learn operators in spherical coordinates. The paper introduces equivariant Spherical FNOs using spherical geometry and demonstrates stable, physically plausible year-long atmospheric rollouts.

Problem
FNOs use discrete Fourier transforms that assume flat geometry, causing artifacts and dissipation when learning operators in spherical coordinates.
Method
The paper extends FNOs to spherical geometry with generalized Fourier transforms, spherical convolutions, and rotationally equivariant operations.
Results
SFNOs generate stable, physically plausible atmospheric forecasts for a year-long rollout, with predictions computed in about 13 minutes on a single GPU.
Takeaways & Limitations
Respecting spherical geometry supports stable long-range weather forecasting without predictions degenerating to implausible physical distributions.
Takeaways & Limitations
Point-wise nonlinearities provide exact equivariance only in the continuous limit, because discrete nonlinearities can introduce arbitrary frequency components.

Abstract

from arXiv · show

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.

1. Introduction

SFNOs extend Fourier Neural Operators to spherical geometry, targeting efficient, resolution-independent modeling of atmospheric dynamics with rotational equivariance and stable long-term forecasts.

Motivation: Climate and weather simulations are computationally expensive because fine grids are needed to resolve multiple spatial and temporal scales and nonlinear interactions.Physics-based methods discretize governing PDEs on grids, while purely data-driven methods have struggled to faithfully represent long-range spatio-temporal dynamics.
Approach: SFNOs generalize FNOs to the sphere using spherical harmonics and a spherical convolution formulation that respects spherical geometry and symmetries.The resulting operator is rotationally equivariant, so rotating the input commutes with applying the operator.
Experiments: The method is evaluated on ERA5 atmospheric data and rotating spherical shallow-water equations for forecasting physical dynamics.The experiments assess both predictive accuracy and long-term autoregressive stability.
Contributions: The SFNO architecture is designed for nonlinear chaotic dynamical systems on the sphere and retains FNO grid-invariance.Point-wise spatial operations allow retraining at different configurations and resolutions, including changing resolution at inference time.
Results: SFNOs demonstrate exceptional stability in year-long autoregressive rollouts while maintaining plausible predicted physical quantities.The paper identifies stable, physically plausible long-term forecasting as a central result.

2. Related Work and Broader Context

The paper situates SFNOs within recent machine-learning approaches to weather and climate modeling and contributes an equivariant extension of FNOs to the sphere.

Related Work: Machine-learning weather and climate modeling has rapidly advanced through deep learning and includes CNNs, Fourier-based models, and vision-transformer architectures.Prior work also explored cubed-sphere representations to address limitations of latitude-longitude grids near the poles.
Broader Context: Equivariance and symmetry provide broader machine-learning inductive biases that constrain representations or parameters according to transformations of the input.The paper places its approach within geometric deep learning and symmetry-based modeling.
Contribution: The paper contributes an equivariant extension of Fourier Neural Operators to the sphere based on a convolution theorem on the sphere.This extends global operator learning while incorporating spherical symmetries.

3. Background

The background formulates atmospheric evolution as operator learning on the sphere and develops the geometric and spectral foundations needed for equivariant spherical operators.

Dynamical Systems: Atmospheric dynamics are represented as a discrete transition map from state u_n at one time to state u_{n+1}, learned from data rather than inferred from PDE discretization.The observed state may omit hidden variables and is sampled at discrete times and positions.
Fourier Neural Operators: FNOs learn resolution-independent global convolution operators using Fourier representations and efficient FFT implementations.This enables modeling long-range dependencies with quasi-log-linear asymptotic complexity.
Geometric Generalization: The spherical FNO layer adapts Fourier operations to respect the symmetry of the underlying manifold.On general manifolds, the generalized Fourier transform uses eigenfunctions of the Laplace-Beltrami operator.
Geometric Generalization: Flat-space translations do not define a global convolution on curved manifolds, motivating geometry-aware operations for physical systems on general manifolds.The paper emphasizes that intrinsic geometric constraints must be incorporated into both transforms and feature operations.
Sphere and Symmetry: For the sphere S2, the relevant symmetry group is SO(3), acting on square-integrable functions in L2(S2).Rotations are represented through passive coordinate transformations on spherical functions.
Sphere and Symmetry: Equivariance requires applying a transformation before or after the operator to produce equivalent results, expressed as F ◦ R = R ◦ F.This formalizes compatibility between the operator and rotational group actions.
Fourier Transform on the Sphere: On the sphere, the spherical harmonic transform decomposes functions in L2(S2) into coefficients in a basis adapted to rotational symmetries.Spherical harmonics form an orthogonal basis and organize rotated functions by degree.
Fourier Transform on the Sphere: Spherical harmonics uniquely exploit sphere symmetries among bases of L2(S2), with rotations mixing functions of the same degree.The spherical Fourier transformation can be viewed as decomposition into minimal rotation-invariant subspaces.

4. Spherical Fourier Neural Operator

SFNO extends Fourier Neural Operators to spherical geometry by combining spherical convolution, spherical harmonic transforms, and equivariant network components. Its architecture learns global dependencies while preserving spherical symmetries, grid-invariance, and efficient multiresolution processing.

Spherical convolution: SFNO replaces a naive spherical Fourier extension with a convolution-based formulation designed to obtain equivariance on the sphere.The spherical convolution is compatible with rotations in SO(3), unlike directly substituting spherical harmonic transforms without accounting for the convolution theorem.
Spherical Fourier Neural Layer: The spherical Fourier neural layer forms the core of SFNO by applying learned filter weights in spherical harmonic space.The generalized Fourier transform and its inverse enable efficient learning of global spatial correlations.
SFNO block: SFNO blocks combine point-wise MLPs for channel mixing with spherical harmonic filtering for long-range dependencies.The block also supports vector-valued inputs through learned complex-valued matrices and can perform equivariant up- and down-scaling by frequency truncation and resolution changes.
Network topology: The full architecture uses point-wise encoder and decoder MLPs, multiple SFNO blocks, optional position embeddings, and a skip connection for maps close to identity.The first and last blocks perform up- or down-scaling, while instance normalization is selected because it applies the same operation globally.
Grid-invariance: SFNO is grid-invariant because its learned operations, except the positional embedding, do not depend on the input grid.It can therefore operate on arbitrary grids and resolutions when the spherical harmonic transform and its inverse are computable there.

5. Numerical Experiments

The experiments evaluate SFNOs on rotating-sphere shallow-water dynamics and ERA5 weather forecasting. SFNOs achieve similar medium-range accuracy to FFT-based models while avoiding pole artifacts and maintaining stable long-term rollouts.

5.1. Spherical Shallow Water Equations: The rotating-sphere shallow-water experiments train two SHT-based and two FFT-based FNO variants using data generated by a classical spectral solver.Models use four (S)FNO blocks, a down-scaling factor of 3, and embedding dimensions of 256.
5.1. Spherical Shallow Water Equations: At 5- and 10-hour rollouts, FFT-based models develop increasingly severe artifacts at the poles, whereas SHT-based SFNOs avoid them.Figure 4 compares geopotential-height predictions with a classical spectral-solver ground truth from a south-pole-centered view.
5.1. Spherical Shallow Water Equations: Training each shallow-water model takes 240 seconds on one NVIDIA A6000, mainly because the numerical solver requires 150 steps to generate a one-hour-ahead sample.
5.2. Weather Prediction / Atmospheric Dynamics: ERA5 experiments use 40 years of data, with 1979–2015 for training, 2016–2017 for validation and selection, and 2018 held out for testing.Models forecast atmospheric variables sampled every 6 hours at the native 0.25-degree latitude–longitude resolution.
5.2. Weather Prediction / Atmospheric Dynamics: FFT- and SHT-based models attain similar medium-range forecast accuracy, but respecting spherical geometry improves behavior during longer rollouts.The reported evaluation uses relative L2 and L1 losses and ACC, with lower losses and higher ACC indicating better performance.

6. Conclusion

SFNO extends Fourier Neural Operators to spherical geometries using spherical harmonics, preserving equivariance and grid-invariance while enabling stable, efficient long-range atmospheric forecasts.

Conclusion: SFNO extends FNOs to spherical geometry through the spherical harmonic transform, supporting equivariance under spherical symmetries.The construction generalizes Fourier transforms from flat domains to functions on the sphere.
Conclusion: Respecting spherical geometry realizes topological boundary conditions correctly and supports stable rollouts without implausible wind-speed distributions.The paper links spherical geometry to physically plausible long-term behavior.
Conclusion: 1,450 autoregressive steps produce a year-long SFNO rollout with physically and statistically consistent behavior, whereas the FFT-based architecture shows artifacts and excessive diffusion.SFNO captures weather dynamics within the two-week predictability horizon and remains plausible for up to a year.
Conclusion: Grid-invariance permits application on arbitrary grids when the spherical harmonic transform and its inverse are available.The architecture can be fine-tuned across new grids and resolutions.
Conclusion: SFNOs show promise for sub-seasonal-to-seasonal forecasting and potentially for machine-learning-based climate prediction.The paper presents these as future applications supported by accuracy, stability, and computational efficiency.

B. Implementation Details

The implementation provides differentiable, parallelizable spherical harmonic transforms for scalable SFNO computation while preserving compatibility between distributed and serial results.

B. Implementation Details: The PyTorch spherical harmonic transform implementation supports differentiable computation and scalable model parallelism.The library is implemented specifically to enable differentiable SHT operations in SFNOs.
B. Implementation Details: The forward SHT combines a one-dimensional real-to-complex DFT over azimuth with matrix multiplication involving associated Legendre polynomials.The transform operates over discrete angular coordinates and spherical-harmonic modes.
B. Implementation Details: Negative-m spherical-harmonic modes need not be stored because they can be recovered from the symmetry relation for real-to-complex transforms.This reduces the modes required by the implementation.
B. Implementation Details: Legendre weights depend on the spatial grid and can be precomputed, while reducing the maximum wave number L provides effective downsampling.The implementation assumes static input and output grids for storing the weight matrices.
B. Implementation Details: The SHT can be distributed with pencil decomposition, keeping spatial decompositions similar between inputs and outputs.This simplifies end-to-end spatial parallelization of SFNO models.
B. Implementation Details: Distributed transforms using the described local-transform strategy are bit-wise identical to serial transforms, unlike fork-join approaches.The result follows from performing the transforms locally after communication and transposition.

B.2. Training

SFNO training first optimizes single-step accuracy, then fine-tunes autoregressive performance by backpropagating through increasingly long forecast sequences.

B.2. Training: The geometric loss approximates the Lp error on the sphere between the predicted next state and the target next state.The loss is evaluated over spherical grid points and weighted by geometric and quadrature factors.
B.2. Training: The loss is computed separately for each predicted channel, normalized by the target norm, averaged across channels, and trained with p = 2.Grid-point contributions use weights formed from the spherical Jacobian and quadrature weights.
B.2. Training: Training begins with single-step optimization to obtain the best possible one-step performance.The learning rate follows a cosine schedule starting at 10^-3.
B.2. Training: Autoregressive fine-tuning accumulates loss across multiple forecast steps and backpropagates through the entire unrolled sequence.The sequence length starts at nsteps = 2 and increases during training.
B.2. Training: Each autoregressive training length uses five epochs with a constant learning rate of 10^-5.The schedule increases nsteps progressively during fine-tuning.

B.3. Performance evaluation

Performance is evaluated with relative L2 and L1 losses and latitude-weighted anomaly correlation, averaged across forecasts from varied initial conditions.

B.3. Performance evaluation: Relative L2 and L1 losses are used alongside anomaly correlation coefficient to assess model performance.ACC is a common weather-prediction metric used for forecast evaluation.
B.3. Performance evaluation: Latitude-weighted ACC compares long-term-mean-subtracted predicted and true anomalies across grid locations.The weighting factor accounts for latitude at each coordinate.
B.3. Performance evaluation: Mean ACC is reported across forecasts from different initial conditions, with first- and third-quartile values showing variability.This reporting convention is used for the ACC plots unless otherwise stated.

C.1. Shallow Water Equations on the Rotating Sphere

The rotating spherical shallow water equations model a thin fluid layer on a rotating sphere as hyperbolic PDEs, with state variables and fluxes defined on S2. Training data are generated from randomized spherical initial conditions, normalized, and optimized with Adam.

C.1. Shallow Water Equations on the Rotating Sphere: The rotating shallow water equations describe a thin, incompressible fluid layer covering a rotating sphere.They are derived from three-dimensional Navier–Stokes equations by integrating over the fluid depth.
C.1. Shallow Water Equations on the Rotating Sphere: The governing system evolves geopotential depth and momentum through continuity and momentum equations with prescribed initial conditions.The displayed equations specify φ and u at t = 0 and include a source term S.
C.1. Shallow Water Equations on the Rotating Sphere: The state vector contains geopotential layer depth φ and tangential momentum φu, while the flux tensor includes the outer product φu ⊗u.The source term includes flux contributions such as the Coriolis force.
C.1. Shallow Water Equations on the Rotating Sphere: Training data are produced by sampling Gaussian random fields for initial height and velocity, then advancing each realization with a classical numerical solver.The initial layer depth has average φavg = 103 · g and standard deviation 120 · g; velocity has zero mean and standard deviation 0.2 ∗√φavg.
C.1. Shallow Water Equations on the Rotating Sphere: The data are z-score normalized, and model weights are optimized with Adam using a learning rate of 2 · 10−3.Training epochs contain 256 samples each.

C.2. Weather prediction/ERA5 Data

ERA5 is a temporally consistent, multi-decadal atmospheric reanalysis used as machine-learning training data, while experiments use 26- and 73-variable subsets. The paper also compares its use with prior data-driven weather models.

C.2. Weather prediction/ERA5 Data: ERA5 is a multi-decadal, high-frequency estimate of Earth’s atmospheric state produced through data assimilation.It combines modern numerical weather models with historical observational records.
C.2. Weather prediction/ERA5 Data: Reanalysis maintains temporal consistency because it uses dynamics from an unchanging modern numerical model across multiple decades.This distinguishes it from operational analysis datasets whose numerical models are periodically updated.
C.2. Weather prediction/ERA5 Data: Previous data-driven numerical weather models have also used ERA5 as training data.The cited studies include work published from 2019 through 2022.
C.2. Weather prediction/ERA5 Data: The experiments use 26-channel and 73-channel ERA5 subsets to train models of different sizes and compare them with IFS.The 73-variable dataset adds additional pressure-level variables to the 26-variable set.