OneScience-Group/Spherical-Fourier-Neural-Operators

Spherical Fourier Neural Operator

Model Overview

SFNO learns dynamical system evolution on the sphere using spherical harmonic transforms, and can be applied to global weather forecasting and spherical shallow-water equation prediction.

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

https://proceedings.mlr.press/v202/bonev23a.html

Model Description

This model package invokes NVIDIA's official torch-harmonics linear SFNO implementation and supports SHT on fake spherical fields, one parameter update step, checkpoint recovery, and short-term autoregressive rollout. It is an operator-level smoke package, not a reproduction of the paper's SWE/ERA5 experiments.

Use Cases

Scenario	Description
Spherical Operator Research	Verify SHT, spectral filtering, and inverse SHT.
Local Rapid Verification	Run through training and inference with fake spherical data.
ERA5 Weather Forecasting	Subsequently interface with 26- or 73-channel ERA5 data.

Usage

1. OneCode

Click to experience intelligent one-click AI4S programming

2. Manual Installation & Usage

Hardware Requirements

• CPU can run the current small configuration.
• GPU is recommended for full ERA5 training.

Download the Model Package

hf download --model OneScience-Group/SFNO --local-dir ./SFNO
cd SFNO

Set Up the Runtime Environment

DCU Environment

conda create -n onescience311 python=3.11 -y
conda activate onescience311
pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
pip install torch-harmonics==0.8.0

GPU Environment

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
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
pip install torch-harmonics==0.8.0

This directory also retains torch-harmonics==0.8.0 under .deps/.

Data

The current scripts generate low-order smooth fake spherical fields in memory and split consecutive time frames into T-1 input/target pairs; no external data download is required.

Training

python scripts/train.py

Training now performs multi-epoch pair Dataset training, time-sequential validation split, learning rate scheduling, and early stopping:

python scripts/train.py --epochs 10
python scripts/train.py --resume weight/training/latest.pth --epochs 20

Inference

python scripts/inference.py

Output files:

weight/model.pth
weight/training/latest.pth
weight/training/best.pth
weight/training/history.json
result/prediction.pt
result/target.pt
result/inference.json

Result Inspection

python scripts/result.py

The current test only verifies that the model executes. Randomly-initialized rollouts do not represent the paper's long-term stability results.

The result script generates result/metrics.json and result/comparison.png. The current RMSE does not incorporate spherical integration weights, and the ACC uses the sample's own spatial mean rather than a long-term training-set climatology; therefore the metrics are not comparable with those reported in the paper.

Paper vs. Current Implementation I/O

Item	Paper SWE / ERA5	Current Smoke Configuration
Input / Output	SWE 3 fields 256×512 / ERA5 26 or 73 channels	[B,2,17,32] smooth synthetic fields
Time Step	SWE 1 hour / ERA5 6 hours	Consecutive indices with no physical units
Architecture	SWE 4×256; weather model 8×384	2 blocks, embed dim 8
Training	Single-step training followed by two-step autoregressive fine-tuning	Multi-epoch single-step pair training and validation; rollout used for inference analysis
Analysis	Spherically weighted relative error and climatological ACC	Unweighted smoke RMSE/ACC

The complete execution flow is train.py -> inference.py -> result.py. Training generates z-scored [T,C,Nlat,Nlon] in memory and forms pairs from consecutive frames. The checkpoint config stores only the model configuration, while training parameters are stored separately under train_config, allowing inference to reconstruct the model architecture directly from the checkpoint. The model package is distributed without local training weights or result/ artifacts. A production SWE/ERA5 mode further requires data loading, variable tables, formal train/validation splits, area-weighted loss, and the paper's two-stage training loop.

Real Data

Real-data training requires ERA5 26/73-channel data, 6-hour temporal pairing, training-set statistics, and spherical grid resampling configuration.

OneScience Official 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 & License

• Official Implementation: https://github.com/NVIDIA/torch-harmonics
• This directory is an independent runnable adaptation of SFNO; see THIRD_PARTY.md for third-party terms.