license: other
language:
- en
- zh
tags:
- OneScience
- Earth Science
- Weather Forecasting
- ERA5
- Representation Learning
frameworks: PyTorch
datasets:
- OneScience/ERA5
AtmoRep
Model Overview
AtmoRep is a stochastic atmospheric dynamics model based on large-scale representation learning that captures the distribution of atmospheric states through masked-token training and ensemble outputs.
Paper: AtmoRep: A stochastic model of atmosphere dynamics using large scale representation learning
https://arxiv.org/abs/2308.13280
Model Description
This directory retains the official vorticity single-field model weights, configuration, and normalization, and provides a tiny AtmoRep-style model for local training and inference verification.
Use Cases
| Scenario | Description |
|---|---|
| Official Resource Validation | Load official .mod weights and verify configuration. |
| Local Rapid Verification | Run masked-token training and inference with the tiny model. |
| ERA5 Atmospheric Representation Learning | Subsequently interface with official GRIB or Zarr data. |
Usage
1. OneCode
Click to experience intelligent one-click AI4S programming
2. Manual Installation & Usage
Hardware Requirements
- The tiny model runs on CPU.
- GPU is recommended for the official model and real-data inference.
Download the Model Package
hf download --model OneScience-Group/AtmoRep --local-dir ./AtmoRep
cd AtmoRep
Set Up the Runtime Environment
DCU Environment
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
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
If dependencies are missing from the official paths, install them with:
pip install zarr wandb cfgrib xarray dask netCDF4 torchinfo
Weights & Data
The current directory preserves:
resources/id4nvwbetz/AtmoRep_id4nvwbetz.mod
resources/id4nvwbetz/model_id4nvwbetz.json
resources/data/normalization/vorticity/
Before re-downloading official resources please note: the current vendor snapshot is not a Git checkout, and the existing download script cannot re-clone into a non-empty vendor/atmorep-official. The resources bundled with this package do not require re-downloading; the script is only suitable for an empty target directory.
bash scripts/download_official_resources.sh .
Tiny Training
python scripts/train.py
This command performs multi-epoch masked-token training using independent train/validation fake Datasets, including DataLoader, AdamW, validation, learning rate scheduling, early stopping, best/latest checkpointing, and training history. Default parameters reside in conf/config.yaml.
Resuming training:
python scripts/train.py --resume weight/training/latest.pth --epochs 20
Training artifacts are weight/training/latest.pth, best.pth, and history.json; the inference-compatible weights weight/tiny_atmorep.pth are also updated. Each Dataset sample consists of fields [T,V,H,W] and a non-empty mask [N]; train/validation use different seeds. This is a complete training pipeline for the tiny model, not a reproduction of the paper's 3.5-billion-parameter official model training.
Tiny Inference
python scripts/inference.py
Inference results are saved as:
result/prediction.pt
result/target.pt
These include the ensemble, ensemble mean, ensemble std, and mask.
Result Inspection
python scripts/result.py
This command generates result/metrics.json and result/comparison.png. The reported metrics are ensemble/mean/spread RMSE in normalized token space, not the paper's physical-unit RMSE, ACC, CRPS, or spread-skill.
Paper vs. Current Implementation I/O
| Item | Paper / Official Model | Tiny Smoke Model |
|---|---|---|
| Input | ERA5 local 4D neighborhoods, 5 mode levels, multiple physical fields | [B,4,1,8,8] single-field random tensor |
| Token | Variable-correlated 4D tokens with absolute space-time and level conditioning | 1×4×4 patches, 16 tokens, relative coordinates with single-level conditioning |
| Output | Multi-head ensembles supporting reconstruction, nowcasting, and interpolation | 4-member masked-token ensemble |
| Training | Large-scale masked-token distribution learning | Multi-epoch Dataset training with independent validation and checkpoint resumption |
| Weights | resources/id4nvwbetz official vorticity weights |
weight/tiny_atmorep.pth; the two are mutually incompatible |
The complete tiny execution flow is train.py -> inference.py -> result.py. The result analysis reads weight/training/history.json when present; inference output includes ensemble, ensemble_mean, ensemble_std, mask, and target. Random data is generated by the Dataset by index; there is currently no independent fake dataset that can be used with the official Zarr sampler. The model package is distributed without local training weights or result/ artifacts — these are created at the paths described above after running the commands.
Official Real-Data Inference
A directly runnable official real-data inference command is not yet available. It further requires ERA5 vorticity GRIB/Zarr data, the ecCodes environment, and parameterization of the file paths in the official evaluate.py.
Real Data
The official vorticity model requires ERA5 vorticity, model levels 96/105/114/123/137, hourly temporal axes, and 0.25° global GRIB/Zarr data.
OneScience Official 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 & License
- Official code is MIT License.
- Official model weights are declared CC BY 4.0.
- ERA5 is subject to Copernicus/ECMWF data terms.