DeepPrim: A Physics-Driven 3D Deep Weather Forecaster via Primitive Equation Learning

This is the official implementation of DeepPrim in ICLR 2026.

⛰️ DeepPrim Framework | ✅ Data Preparation | 🚀 Setup | 📊 Train and Evaluation

⛰️ DeepPrim Framework

DeepPrim is a physics-informed 3D deep weather forecaster designed to learn primitive equations of the Earth’s atmospheric dynamics for short-term global and regional weather forecasting based on the ERA5 dataset. This repository includes preprocessing and training pipelines and supports global and regional forecasting for various areas.

Framework of DeepPrim and key modules, including (a) overall pipeline and architectures of (b) force network based on 3D-BiViT, (c) initialization network, and (d) source-sink network.

✅ Data Preparation

Download ERA5 Data

Download the ERA5 reanalysis dataset from the WeatherBench. Organize the data directory as follows:

5.625deg
   ├── 10m_u_component_of_wind
   ├── 10m_v_component_of_wind
   ├── 2m_temperature
   ├── constants.nc
   ├── geopotential
   ├── relative_humidity
   ├── specific_humidity
   ├── temperature
   ├── toa_incident_solar_radiation
   ├── total_precipitation
   ├── u_component_of_wind
   └── v_component_of_wind

Preprocessing

Convert the raw NetCDF files into smaller, more manageable NumPy files and compute essential statistical measures. Execute the following script:

python src/data_preprocessing/nc2np_equally_era5.py \
    --root_dir /mnt/data/5.625deg \
    --save_dir /mnt/data/5.625deg_npz \
    --start_train_year 1979 --start_val_year 2016 \
    --start_test_year 2017 --end_year 2019 --num_shards 8

The preprocessed data directory will have the following structure:

5.625deg_npz
   ├── train
   ├── val
   ├── test
   ├── normalize_mean.npz
   ├── normalize_std.npz
   ├── lat.npy
   └── lon.npy

🚀 Setup

Environment Configuration

1. Create conda environment(Optional)

conda env create -f environment.yml
conda activate deepprim

Installing deepprim

Install the depprim package in editable mode:

pip install -e .

📊 Train and Evaluation

Global Forecasting

To train a global forecasting model with a forecasting lead time of 6-hour , use the following command:

bash ./scripts/global/train_6h.sh

Scripts for 12-hour, 18-hour, and 24-hour forecast models are also available.

Regional Forecasting

For regional forecasting in Australia with a 6-hour prediction horizon, use the following command:

bash ./scripts/regional/Australia/train_6h.sh

Scripts for 12-hour, 18-hour, and 24-hour forecast models, and for other regions (e.g., North America, South America) are also available.

Citation

If you find this repo useful, please cite:

@inproceedings{
chen2026deepprim,
title={DeepPrim: a Physics-Driven 3D Short-term Weather Forecaster via Primitive Equation Learning},
author={Jiawei Chen and Weiqi Chen and Rong Hu and Peiyuan Liu and Haifan Zhang and Liang Sun},
booktitle={The Fourteenth International Conference on Learning Representations},
year={2026},
url={https://openreview.net/forum?id=EyyWd0hH0q}
}