OneScience-Group/CorrDiff

CorrDiff

Model Introduction

CorrDiff (Residual Corrective Diffusion Model) is a two-stage generative model for kilometer-scale regional weather downscaling. It downscales approximately 25 km ERA5 fields to approximately 2 km regional weather fields over Taiwan and synthesizes radar reflectivity that is not present in the input.

Paper: Residual Corrective Diffusion Modeling for Km-scale Atmospheric Downscaling

https://arxiv.org/abs/2309.15214

Model Description

CorrDiff was proposed by NVIDIA and its collaborators and was trained with coarse-resolution ERA5 reanalysis and high-resolution WRF regional model data from Taiwan's Central Weather Administration. The model is suitable for converting coarse-resolution global weather fields into high-resolution regional weather fields and producing probabilistic weather predictions.

Use Cases

Scenario	Description
Regional weather downscaling training	Train CorrDiff with time-aligned coarse-resolution ERA5 inputs and high-resolution CWA-WRF targets.
Local quick validation	Use synthetic paired data to check data loading, two-stage training, ensemble inference, and result visualization.
Hugging Face / OneCode execution	Download the standalone model package, install dependencies, and run the scripts directly.
Multi-GPU training	Launch multi-process training with torchrun after adapting the real-data pipeline for distributed training.

Usage Guide

1. OneCode Usage

Experience intelligent one-click AI4S programming through the OneCode online environment:

Click to Experience Intelligent One-Click AI4S Programming

2. Manual Installation and Usage

Hardware Requirements

• A GPU or DCU is recommended.
• CPU can be used for import and small-scale connectivity verification; full training and inference will be slow.
• DCU users must install DTK in advance. DTK 25.04.2 or above, or the OneScience recommended version matching your cluster, is recommended.

Download the Model Package

hf download OneScience-Group/CorrDiff --local-dir ./CorrDiff
cd CorrDiff

Install the Runtime Environment

DCU Environment

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

# Please 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 Introduction

This repository uses synthetic paired data by default to validate training, ensemble inference, and evaluation:

python scripts/fake_data.py --output data/era5_corrdiff.npz

The synthetic file contains input with shape [N, 12, 36, 36] and target with shape [N, 4, 448, 448]. It is intended for pipeline validation only and does not represent real weather predictions.

ERA5 inputs can be downloaded from the OneScience community:

hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data/era5

The paired CWA-WRF resource can be downloaded from NVIDIA NGC:

ngc registry resource download-version "nvidia/modulus/modulus_datasets_cwa:v1"

ERA5 alone is insufficient for supervised CorrDiff training. Real data must be time-aligned and preprocessed into the paired input and target interface expected by the scripts. Verify the data path in conf/config.yaml before training.

Training

Single GPU:

python scripts/train.py

Multi-GPU:

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:

data/checkpoints/regression_model.pth
data/checkpoints/diffusion_model.pth
data/checkpoints/training_history.npz

Training first fits the conditional-mean regression model and then freezes it while fitting the residual diffusion model. training_history.npz stores the regression and diffusion denoising losses.

Training Weights

The training command generates regression_model.pth and diffusion_model.pth from data/era5_corrdiff.npz. Weights produced with synthetic data validate the training and inference pipeline only and do not provide real weather forecasting skill. This repository does not present them as pretrained weights reproducing the paper.

Inference

python scripts/inference.py

Prediction output:

result/output/predictions.npz

The file stores the ensemble predictions, ensemble mean, ensemble standard deviation, input, and target generated by the current trained weights.

Evaluation and Visualization

python scripts/result.py

Evaluation outputs:

result/output/metrics.json
result/output/prediction_comparison.png

metrics.json stores MAE and CRPS for the four output variables. prediction_comparison.png compares the ensemble mean, ensemble standard deviation, and target. Metrics and figures generated from synthetic data are pipeline checks, not paper results.

Official OneScience Resources

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 a reproduction of the original CorrDiff paper.