--- license: mit language: - en - zh tags: - OneScience - Earth science - Weather forecasting - Medium- to long-range weather forecasting - Foundation models - Vision Transformer - ERA5 frameworks: PyTorch datasets: - OneScience/ERA5 ---

Prithvi WxC

# Model Introduction Prithvi WxC (Weather and Climate) was proposed by NASA-IMPACT, IBM, and other teams. It is a weather and climate foundation model based on a Vision Transformer (alternating local/global attention with Hiera and MaxViT), supporting forecasting (6-hour-step rollout) and climate simulation (internal error growth). Paper:Prithvi WxC: Foundation Model for Weather and Climate https://arxiv.org/abs/2409.13598 # Model Description Prithvi WxC is a deterministic global weather foundation model: it takes atmospheric states at two consecutive 6-hour time steps, optionally with static fields, and outputs the target state. Longer lead times are obtained through autoregressive rollout.This repository is organized from the official `NASA-IMPACT/Prithvi-WxC` implementation and integrated with the OneScience data loading and training workflow. # Use Cases | Scenario | Description | | :---: | :--- | | Global weather and climate foundation model research | Train or fine-tune a Vision Transformer forecasting model on ERA5 data. | | Long-horizon autoregressive rollout | Generate medium- to long-range forecasts autoregressively at 6-hour intervals. | | Local quick validation | Use synthetic data to check data loading, training, inference, and result scripts. | | ModelScope/OneCode execution | Download the model package, install dependencies, and run the scripts directly. | | Multi-card training | Launch multi-process training with `torchrun`. | # Usage ## 1. OneCode Usage Use the OneCode online environment for intelligent one-click AI4S programming: [Try intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home) ## 2. Manual Installation and Usage **Hardware Requirements** - GPU or DCU is recommended. - CPU can be used for imports and small-configuration connectivity validation, but full training and inference are slower. - DCU users must install DTK beforehand. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended. - The paper-level configuration (`embed_dim=2560`, 25 encoder blocks, 5 decoder blocks, and approximately 2.3 billion parameters) requires substantial GPU memory. ### Download the Model Package ```bash hf download OneScience-Group/PrithviWxC --local-dir ./PrithviWxC cd PrithviWxC ``` ### Install the Runtime Environment **DCU Environment** ```bash # Activate DTK and CONDA first conda create -n onescience311 python=3.11 -y conda activate onescience311 # uv installation is supported 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 # uv installation is supported pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai ``` ### Training Data The OneScience community provides ERA5 data for training (the current repository contains complete data slices subject to data-file size limits). Download it with the command below and confirm that the data path in `conf/config.yaml` is correct: ```bash hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data ``` For a quick workflow validation, run the synthetic data script first: ```bash python scripts/fake_data.py ``` > Note: `scripts/fake_data.py` generates the `[T, C, H, W]` HDF5 data required by the two input time steps and generates `data/static/static.npy` (currently `[4, 32, 64]`) for training and inference. ### Training Single card： ```bash python scripts/train.py ``` Multiple cards： ```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 outputs: ```text data/checkpoints/model_bak.pth data/checkpoints/trloss.npy data/checkpoints/valoss.npy ``` ### Training Weights The `weight/` folder is reserved for model weights. Official weights with approximately 2.3 billion parameters (such as PrithviWxC_160_13b_2t_0p5d_v1.pt) are published on Hugging Face, but their structure differs from this repository's small configuration. Align the channel count and grid size before loading; weights are not provided by default, and users may train the model using the paper configuration. ### Inference Inference reads `data/checkpoints/model_bak.pth`： ```bash python scripts/inference.py ``` Prediction results are written to: ```text result/output/ ``` ### Evaluation and Visualization ```bash python scripts/result.py ``` Outputs include： - `result/rmse.npy` - `result/acc.npy` - `result/loss.png` - Forecast comparison plots for the specified date and variables # Official Source and Reproduction Notes - The model implementation comes from the official `NASA-IMPACT/Prithvi-WxC` (MIT License). The official implementation is embedded unchanged in `model/prithvi_wxc_official.py`; `model/prithvi_wxc.py` is only a YAML-driven thin wrapper (with identity normalization parameters for small-configuration connectivity validation). - Commit fetched for the current case directory: `79dabfcd17abe77e2d5c696707c0164a04f2ec01` (2026-02-05). - `conf/config.yaml` uses a small configuration (`embed_dim=32`, `n_blocks_encoder=1`, `n_blocks_decoder=1`) for connectivity validation by default; paper-level reproduction requires a 0.5°×0.625° grid, 160 channels, `embed_dim=2560`, and 13+12 encoder blocks/3+2 decoder blocks as specified in the paper. - The following details are not disclosed in the paper and are assumptions in this reproduction:data normalization statistics (identity normalization is currently used; real statistics will be supplied with the data), masked-training details and pretraining schedule, and some hyperparameters (such as the relative positional encoding implementation). # 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 organization and adaptation of Prithvi WxC. The model source is based on the official `NASA-IMPACT/Prithvi-WxC` implementation by Schmude et al. (2024) and follows the MIT License. - Please cite:Schmude et al. Prithvi WxC: Foundation Model for Weather and Climate. arXiv:2409.13598, 2024.