OneScience-Group/SphericalFourierNeuralOperators

Spherical Fourier Neural Operators

Model Introduction

SFNO（Spherical Fourier Neural Operator）was proposed jointly by NVIDIA, Caltech, and other institutions，replaces the planar FFT in standard FNOs with spherical harmonic transforms (SHT), making spectral convolution respect spherical geometry and mitigating polar artifacts, spectral artifacts, and instability during long autoregressive rollouts.

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

https://arxiv.org/abs/2306.03838

Model Description

SFNO is a deterministic, single-state global weather dynamics model: it takes the atmospheric state at one 6-hour time step and outputs the same set of variables at the next 6-hour time step. Longer lead times are obtained through autoregressive rollout.This repository is organized from the official reference implementation in NVIDIA/torch-harmonics and integrated with the OneScience data loading and training workflow.

Use Cases

Scenario	Description
Global weather dynamics research	Train a spherical-equivariant neural operator forecasting model on ERA5 data.
Long-horizon autoregressive stability research	Evaluate polar artifacts and dissipation during multi-step rollout.
Local quick validation	Use synthetic data to check data loading, training, inference, and result scripts.
ModelScope/OneCode execution	Download the model package, install dependencies, and run the scripts directly.
Multi-card training	Launch multi-process training with torchrun.

Usage

1. OneCode Usage

Use the OneCode online environment for intelligent one-click AI4S programming:

Try intelligent one-click AI4S programming

2. Manual Installation and Usage

Hardware Requirements

• GPU or DCU is recommended.
• CPU can be used for imports and small-configuration connectivity validation, but full training and inference are slower.
• DCU users must install DTK beforehand. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.
• The model depends on torch-harmonics(spherical harmonic transforms)，Confirm that its version matches PyTorch/CUDA before installation.

Download the Model Package

hf download OneScience-Group/SphericalFourierNeuralOperators --local-dir ./SphericalFourierNeuralOperators
cd SphericalFourierNeuralOperators

Install the Runtime Environment

DCU Environment

# Activate DTK and CONDA first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
# uv installation is supported
pip install onescience[earth-dcu] torch-harmonics -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

GPU Environment

# Activate CONDA first
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
conda activate onescience311
# uv installation is supported
pip install onescience[earth-gpu] torch-harmonics -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

Training Data

The OneScience community provides ERA5 data for training (the current repository contains complete data slices subject to data-file size limits). Download it with the command below and confirm that the data path in conf/config.yaml is correct:

hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data

For a quick workflow validation, run the synthetic data script first:

python scripts/fake_data.py

Note: scripts/fake_data.py generates [T, C, H, W] HDF5 data from the model configuration. The current small configuration uses 6 channels and a 32×64 grid, and automatically calculates a time length sufficient for the batch.

Training

Single card：

python scripts/train.py

Multiple cards：

torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py

Training outputs:

data/checkpoints/model_bak.pth
data/checkpoints/trloss.npy
data/checkpoints/valoss.npy

Training Weights

The weight/ folder is reserved for model weights. The paper does not clearly release weights for the 26/73-channel weather model; this repository does not currently provide official weights, and users may train the model using the paper configuration.

Inference

Inference reads data/checkpoints/model_bak.pth：

python scripts/inference.py

Prediction results are written to:

result/output/

Evaluation and Visualization

python scripts/result.py

Outputs include：

• result/rmse.npy
• result/acc.npy
• result/loss.png
• Forecast comparison plots for the specified date and variables

Official Source and Reproduction Notes

• The model implementation comes from the SFNO reference implementation in the official NVIDIA/torch-harmonics (BSD-3-Clause).
• Commit fetched for the current case directory: 49bac755cd8306fbd27a3604acafa65adf7ca202 (2026-08-14).
• conf/config.yaml uses a small configuration (img_size=[32, 64], embed_dim=16, num_layers=2) for connectivity validation by default; paper-level reproduction requires a 0.25° (721×1440) grid, 26/73 channels, and a larger network as specified in the paper.
• The following details are not disclosed in the paper and are assumptions in this reproduction:the internal spectral downsampling factor of the weather model, the positional embedding form, per-variable normalization statistics, and training batch size.

Official OneScience Information

Platform	OneScience Main Repository	Skills Repository
Gitee	https://gitee.com/onescience-ai/onescience	https://gitee.com/onescience-ai/oneskills
GitHub	https://github.com/onescience-ai/OneScience	https://github.com/onescience-ai/oneskills

Citation and License

• This repository is an independent organization and adaptation of SFNO. The model source is based on the official torch-harmonics implementation by Bonev et al. (2023) and follows BSD-3-Clause.
• Please cite:Bonev, Kurth, Hundt, Pathak, Baust, Kashinath, Anandkumar. Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere. ICML 2023.