README.md

---
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
---
<p align="center">
  <strong>
    <span style="font-size: 30px;">Prithvi WxC</span>
  </strong>
</p>


# 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.