README.md

---
license: apache-2.0
language:
- en
tags:
- OneScience
- Earth Science
- Global Atmospheric Emulation
- Neural Operator
- Climate Simulation
- Physical Constraints
frameworks: PyTorch
---

<p align="center"><strong><span style="font-size: 30px;">ACE2</span></strong></p>

# 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](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)

## 2. Download and Installation

```bash
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**

```bash
# 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**

```bash
# 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.

```bash
python scripts/fake_data.py
```

### Training

For single-GPU training, use:

```bash
python scripts/train.py
```

For multi-GPU training, use:

```bash
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:

```text
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

```bash
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:

```text
result/output/predictions.npz
```

### Evaluation and Visualization

```bash
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:

```text
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.