| --- |
| license: other |
| language: |
| - en |
| - zh |
| tags: |
| - OneScience |
| - Earth Science |
| - Weather Forecasting |
| - Spherical Neural Operators |
| - ERA5 |
| frameworks: PyTorch |
| datasets: |
| - OneScience/ERA5 |
| --- |
| |
| <p align="center"> |
| <strong><span style="font-size: 30px;">Spherical Fourier Neural Operator</span></strong> |
| </p> |
|
|
| # 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](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home) |
|
|
| ## 2. Manual Installation & Usage |
|
|
| **Hardware Requirements** |
|
|
| - CPU can run the current small configuration. |
| - GPU is recommended for full ERA5 training. |
|
|
| ### Download the Model Package |
|
|
| ```bash |
| hf download --model OneScience-Group/SFNO --local-dir ./SFNO |
| cd SFNO |
| ``` |
|
|
| ### Set Up the Runtime Environment |
|
|
| **DCU Environment** |
|
|
| ```bash |
| 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** |
|
|
| ```bash |
| 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 |
|
|
| ```bash |
| python scripts/train.py |
| ``` |
|
|
| Training now performs multi-epoch pair Dataset training, time-sequential validation split, learning rate scheduling, and early stopping: |
|
|
| ```bash |
| python scripts/train.py --epochs 10 |
| python scripts/train.py --resume weight/training/latest.pth --epochs 20 |
| ``` |
|
|
| ### Inference |
|
|
| ```bash |
| python scripts/inference.py |
| ``` |
|
|
| Output files: |
|
|
| ```text |
| 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 |
|
|
| ```bash |
| 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. |
|
|