OneScience-Group/AtmoRep

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.