OneScience-Group/ACE2

ACE2

Model Introduction

ACE2 autoregressively simulates weather, climate variability, and forced responses at six-hour intervals from global atmospheric states and external forcing such as sea-surface temperature. Hard physical corrections constrain dry-air mass and moisture budgets, supporting fast simulations from weather scales to long-term climate statistics.

Paper: ACE2: accurately learning subseasonal to decadal atmospheric variability and forced responses
https://doi.org/10.1038/s41612-025-01090-0

Model Description

The method was proposed by teams from the Allen Institute for AI, the Geophysical Fluid Dynamics Laboratory, and collaborating institutions. The paper trains separate models on ERA5 reanalysis and SHiELD historical atmospheric simulations. ACE2 uses a spherical Fourier neural operator to learn global atmospheric transitions with hard dry-air-mass and moisture-budget corrections. The model supports six-hour autoregressive atmospheric emulation and climate-statistics analysis on a one-degree global grid.

Use Cases

Use Case	Description
Global atmospheric emulation	Run six-hour autoregressive forecasts on a one-degree grid.
Physically constrained simulation	Validate hard dry-air-mass and global-moisture constraints.
Local engineering validation	Validate full-grid, 50-channel, eight-level structured synthetic data.
ModelScope/OneCode execution	Validate data, training, inference, atmospheric metrics, and evaluation.
Multi-GPU training	Validate distributed training and checkpoint workflows through torchrun.

Usage Instructions

1.OneCode

Try intelligent, one-click AI4S programming

2. Download and Installation

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

Environment Dependencies

Hardware Requirements

• A GPU or DCU is recommended.
• A CPU can be used for connectivity validation with the default small-sample configuration.
• DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.

DCU Environment

# Activate DTK and Conda first
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

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

Training Data

This repository uses a small synthetic dataset for engineering validation, including global atmospheric states, external forcing, 50 state channels, eight vertical levels, six-hour temporal relationships, and the real [N,T,50,180,360] dimensions. The synthetic data retain the paper's grid, channels, levels, and temporal specification while reducing samples, model scale, and epochs; the 50-channel ledger is an engineering interpretation of the main text. These data validate SFNO, hard corrections, training, inference, and evaluation only and do not represent official ERA5 or SHiELD distributions and scale.

python scripts/fake_data.py

Training

For single-GPU training, use:

python scripts/train.py

For multi-GPU training, use:

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 uses a two-step six-hour autoregressive MSE objective followed by hard dry-air-mass and moisture corrections. The default reduces SFNO width, spectral modes, and epochs without reducing data dimensions or the two-step target. Training artifacts are saved to:

result/checkpoints/ace2.pt
result/training/metrics.json

Trained Weights

The paper provides an official ACE2-ERA5 checkpoint at https://doi.org/10.57967/hf/5377. It is not bundled under weight/, and this compact model does not claim checkpoint compatibility.

Inference

python scripts/inference.py

Inference loads the local checkpoint and takes an initial global atmospheric state plus time-varying external forcing. It autoregressively generates atmospheric states at six-, twelve-, and eighteen-hour leads and applies physical correction after every step. Outputs retain 50 channels, the 180×360 grid, and lead-time ordering. Inference results are saved to:

result/output/predictions.npz

Evaluation and Visualization

python scripts/result.py

Evaluation computes latitude-weighted RMSE, global-mean R², persistence skill, and dry-mass and moisture-closure errors. It saves structured model, baseline, and conservation results and generates a comparison of forecast errors and conservation residuals. Synthetic-data results validate engineering only and do not represent formal paper performance. Evaluation results are saved to:

result/evaluation/metrics.json
result/evaluation/comparison.png

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 engineering reproduction of the public ACE2 specifications, with code licensed under the Apache License 2.0.

The original paper is licensed under CC BY 4.0; the paper, official model weights, ERA5 data, and SHiELD data remain subject to the licenses and terms of their respective projects.