zhangrenchao commited on
Commit
393fb7e
·
verified ·
1 Parent(s): 380b161

Add English model card

Browse files
Files changed (1) hide show
  1. README.md +146 -0
README.md ADDED
@@ -0,0 +1,146 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ language:
4
+ - en
5
+ tags:
6
+ - OneScience
7
+ - Earth Science
8
+ - Global Atmospheric Emulation
9
+ - Neural Operator
10
+ - Climate Simulation
11
+ - Physical Constraints
12
+ frameworks: PyTorch
13
+ ---
14
+
15
+ <p align="center"><strong><span style="font-size: 30px;">ACE2</span></strong></p>
16
+
17
+ # Model Introduction
18
+
19
+ 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.
20
+
21
+ Paper: ACE2: accurately learning subseasonal to decadal atmospheric variability and forced responses
22
+ https://doi.org/10.1038/s41612-025-01090-0
23
+
24
+ # Model Description
25
+
26
+ 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.
27
+
28
+ # Use Cases
29
+
30
+ | Use Case | Description |
31
+ | :---: | :--- |
32
+ | Global atmospheric emulation | Run six-hour autoregressive forecasts on a one-degree grid. |
33
+ | Physically constrained simulation | Validate hard dry-air-mass and global-moisture constraints. |
34
+ | Local engineering validation | Validate full-grid, 50-channel, eight-level structured synthetic data. |
35
+ | ModelScope/OneCode execution | Validate data, training, inference, atmospheric metrics, and evaluation. |
36
+ | Multi-GPU training | Validate distributed training and checkpoint workflows through `torchrun`. |
37
+
38
+ # Usage Instructions
39
+
40
+ ## 1.OneCode
41
+
42
+ [Try intelligent, one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
43
+
44
+ ## 2. Download and Installation
45
+
46
+ ```bash
47
+ hf download OneScience-Group/ACE2 --local-dir ./ACE2
48
+ cd ACE2
49
+ ```
50
+
51
+ ### Environment Dependencies
52
+
53
+ **Hardware Requirements**
54
+
55
+ - A GPU or DCU is recommended.
56
+ - A CPU can be used for connectivity validation with the default small-sample configuration.
57
+ - DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.
58
+
59
+ **DCU Environment**
60
+
61
+ ```bash
62
+ # Activate DTK and Conda first
63
+ conda create -n onescience311 python=3.11 -y
64
+ conda activate onescience311
65
+ pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
66
+ ```
67
+
68
+ **GPU Environment**
69
+
70
+ ```bash
71
+ # Activate Conda first
72
+ conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
73
+ conda activate onescience311
74
+ pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
75
+ ```
76
+
77
+ ### Training Data
78
+
79
+ 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.
80
+
81
+ ```bash
82
+ python scripts/fake_data.py
83
+ ```
84
+
85
+ ### Training
86
+
87
+ For single-GPU training, use:
88
+
89
+ ```bash
90
+ python scripts/train.py
91
+ ```
92
+
93
+ For multi-GPU training, use:
94
+
95
+ ```bash
96
+ 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
97
+ ```
98
+
99
+ 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:
100
+
101
+ ```text
102
+ result/checkpoints/ace2.pt
103
+ result/training/metrics.json
104
+ ```
105
+
106
+ ### Trained Weights
107
+
108
+ 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.
109
+
110
+ ### Inference
111
+
112
+ ```bash
113
+ python scripts/inference.py
114
+ ```
115
+
116
+ 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:
117
+
118
+ ```text
119
+ result/output/predictions.npz
120
+ ```
121
+
122
+ ### Evaluation and Visualization
123
+
124
+ ```bash
125
+ python scripts/result.py
126
+ ```
127
+
128
+ 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:
129
+
130
+ ```text
131
+ result/evaluation/metrics.json
132
+ result/evaluation/comparison.png
133
+ ```
134
+
135
+ # Official OneScience Information
136
+
137
+ | Platform | OneScience Main Repository | Skills Repository |
138
+ | --- | --- | --- |
139
+ | Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
140
+ | GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
141
+
142
+ # Citation and License
143
+
144
+ This repository is an independent engineering reproduction of the public ACE2 specifications, with code licensed under the Apache License 2.0.
145
+
146
+ 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.