--- 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](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home) ## 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 ```bash hf download --model OneScience-Group/AtmoRep --local-dir ./AtmoRep cd AtmoRep ``` ### Set Up the Runtime Environment **DCU Environment** ```bash 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 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: ```bash pip install zarr wandb cfgrib xarray dask netCDF4 torchinfo ``` ### Weights & Data The current directory preserves: ```text 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 bash scripts/download_official_resources.sh . ``` ### Tiny Training ```bash 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: ```bash 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 ```bash python scripts/inference.py ``` Inference results are saved as: ```text result/prediction.pt result/target.pt ``` These include the ensemble, ensemble mean, ensemble std, and mask. ### Result Inspection ```bash 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.