Hugging Face's logo

OneScience-Group
/
Ai2_Climate_Emulator

English
ace
OneScience
Earth Science
Climate Simulation
Global Atmospheric State Simulation
SFNO
FV3GFS
Model card Files
xet
 Community
Ai2_Climate_Emulator
File size: 4,573 Bytes
5c365c5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
---
frameworks: PyTorch
language:
- en
license: apache-2.0
tags:
- OneScience
- Earth Science
- Climate Simulation
- Global Atmospheric State Simulation
- SFNO
- FV3GFS
tasks: []
datasets:
- FV3GFS
---
<p align="center">
  <strong>
    <span style="font-size: 30px;">Ai2_Climate_Emulator</span>
  </strong>
</p>

# Model Introduction

The AI2 Climate Emulator (ACE) is a global atmospheric state emulator proposed by the Allen Institute for AI (AI2).

Paper: ACE: A fast, scalable foundation model for the atmosphere

https://arxiv.org/abs/2310.02074

# Model Description

This project implements the spherical Fourier neural operator (SFNO) forward graph with PyTorch and `torch_harmonics`. It takes the atmospheric state and external forcings at the current six-hour time step as input, predicts the state at the next time step, and can generate multi-step climate or weather fields autoregressively.

# Use Cases

| Scenario | Description |
| :---: | :--- |
| Global atmospheric state simulation | Train a one-step ACE model with FV3GFS data following the 40/44-channel protocol. |
| Local quick validation | Generate synthetic NPZ files with `scripts/fake_data.py` to check the training, inference, and result-visualization pipeline. |
| ModelScope / OneCode execution | Download the standalone model package, install dependencies, and run the scripts directly. |
| Multi-GPU training | Launch PyTorch DDP 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](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)

## 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

```bash
hf download OneScience-Group/Ai2_Climate_Emulator --local-dir ./Ai2_Climate_Emulator
cd Ai2_Climate_Emulator
```

### Install the Runtime Environment

**DCU Environment**

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

### Training Data Introduction

The ACE paper uses an ensemble of 11 FV3GFS initial conditions: 10 members for training and one member for validation. The simulations are written at six-hour intervals and regridded to a Gaussian latitude-longitude grid. The original FV3GFS files and NOAA `fregrid` are not included in this model package; users must prepare and convert them to the NPZ format required by the project:

```text
inputs:  [N, 40, H, W]
targets: [N, 44, H, W]
```

When real data is unavailable, generate synthetic data for pipeline validation:

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

### Training

Single GPU:

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

Multi-GPU:

```bash
torchrun --nproc_per_node=8 scripts/train.py
```

Training checkpoints are written to `data/checkpoint/model_bak.pt` by default.

### Training Weights

This repository provides weights trained on FV3GFS data in the `weight/` folder. The weight files will be uploaded soon and are expected to be available in the near future.

### Inference

```bash
python scripts/inference.py
```

Inference results are saved to `output/infer/rollout.npz` by default.

### Evaluation and Visualization

```bash
python scripts/result.py
```

Area-weighted RMSE, global mean bias, and PNG figures are written to `output/pic/` by default.

# 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 reproduction of the ACE model.