OneScience-Group/Spherical_DYffusion

Spherical DYffusion

Model Introduction

Spherical DYffusion was proposed by Salva Ruhling Cachay and collaborators for probabilistic simulation of a global climate model.

Paper: Probabilistic Emulation of a Global Climate Model with Spherical DYffusion

https://arxiv.org/abs/2406.14798

Model Description

The original method models spherical dynamics with an SFNO and uses the DYffusion interpolator and forecaster in a two-stage training procedure for probabilistic ensemble simulation. This repository contains a compact local implementation that preserves the project's tensor and data contracts for smoke testing; it is not a full paper-scale SFNO/DYffusion implementation.

Use Cases

Scenario	Description
Local pipeline validation	Use synthetic 37-channel global-grid data to check training, inference, and visualization.
FV3GFS protocol checks	Validate NetCDF variables, spatial dimensions, and consecutive time frames.
ModelScope / OneCode execution	Download the standalone model package and run the compact local pipeline.
Multi-GPU training	Launch PyTorch DistributedDataParallel with torchrun.

Usage Guide

1. OneCode Usage

Experience intelligent one-click AI4S programming through the OneCode online environment:

Click to Experience Intelligent One-Click AI4S Programming

2. Manual Installation and Usage

Hardware Requirements

• A GPU or DCU is recommended.
• CPU can be used for import and small-scale connectivity verification; full training and inference will be slow.
• DCU users must install DTK in advance. DTK 25.04.2 or above, or the OneScience recommended version matching your cluster, is recommended.

Download the Model Package

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

Install the Runtime Environment

DCU Environment

# Please 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] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

GPU Environment

# Please 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] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

Generate Synthetic Data

Generate a deterministic NetCDF FV3GFS-contract fixture at data/data/synthetic_fv3gfs.nc:

python scripts/fake_data.py

The fixture contains 37 protocol variables, including surface pressure and temperature, eight vertical levels of temperature, total water, and wind components, plus DSWRFtoa, HGTsfc, and ocean_fraction. It is intended only for protocol checks. The local training pipeline creates its own learnable data/data/virtual_fv3gfs.npz fixture.

Training

Single GPU:

python scripts/train.py

Multi-GPU:

torchrun --nproc_per_node=8 scripts/train.py

Training starts from random initialization and saves data/checkpoint/model_bak.pt and data/checkpoint/last.pt.

The complete local smoke workflow can also be run with:

python scripts/local_pipeline.py all

Training Weights

This repository provides a weight/ directory for FV3GFS-compatible checkpoints. The weight files will be uploaded soon and are expected to be available in the near future.

Inference

Inference reads data/checkpoint/model_bak.pt and writes output/inference/prediction.npz:

python scripts/inference.py

Evaluation and Visualization

python scripts/result.py

The script computes per-variable and overall diagnostics and writes output/visualization/diagnostic_dashboard.png and output/visualization/variable_metrics.png, with machine-readable summaries under output/metrics/.

Official OneScience Resources

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 a compact local reproduction of the original Spherical DYffusion paper and does not claim to reproduce the paper-scale training setup or metrics.