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

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](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.