README.md

---
license: apache-2.0
language:
- en
tags:
- OneScience
- Earth Science
- Weather Downscaling
- Diffusion Model
frameworks: PyTorch
datasets:
- ERA5
- CWA-WRF
---

<p align="center">
  <strong>
    <span style="font-size: 30px;">CorrDiff</span>
  </strong>
</p>

# Model Introduction

CorrDiff is a two-stage generative model for kilometer-scale atmospheric downscaling. It first predicts a high-resolution mean with conditional regression and then uses a residual diffusion model to generate local stochastic details and ensemble forecasts.

Paper: Residual Corrective Diffusion Modeling for Km-scale Atmospheric Downscaling  
https://arxiv.org/abs/2309.15214

# Model Description

CorrDiff was proposed by an NVIDIA research team. The model is trained using ERA5 reanalysis data and high-resolution regional CWA-WRF simulation data, and performs weather-field downscaling through regression and residual diffusion. It is suitable for high-resolution weather-field generation, kilometer-scale downscaling, and ensemble uncertainty analysis.

# Use Cases

| Scenario | Description |
| :---: | :--- |
| Weather downscaling | Downscale `36x36` coarse-resolution condition fields to `448x448`. |
| Ensemble forecasting | Generate multiple possible high-resolution results through diffusion sampling. |
| Radar reflectivity generation | Generate high-resolution maximum radar reflectivity that is absent from the ERA5 condition fields, supporting analysis of fine-scale precipitation-system structures. |
| Extreme-weather risk analysis | Use ensemble members and spatial uncertainty to characterize the possible evolution of local heavy precipitation and other high-impact weather. |
| Local engineering validation | Use a small amount of synthetic data to check training, inference, and evaluation workflows. |
| Multi-GPU training | Launch distributed training with `torchrun`. |

# Usage Guide

## 1.OneCode

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

[Click to Experience Intelligent One-Click AI4S Programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)

## 2.Download and Installation

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

### Environment Dependencies

**Hardware Requirements**

- A GPU or DCU is recommended.
- CPU can be used for connectivity validation with a small configuration; full training and inference are slower.
- DCU users must install DTK in advance. DTK 25.04.2 or above, or the OneScience-recommended version matching the current cluster, is recommended.

**DCU Environment**

```bash
# 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**

```bash
# 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

By default, two synthetic samples following the same protocol are used to validate the engineering workflow. The data are saved as `data/corrdiff.npz`, and the synthetic data do not represent real weather distributions.

The synthetic data retain the input and output specifications of the Taiwan downscaling experiment in the paper: 12-channel `36x36` condition fields and 4-channel `448x448` target fields.

Real data must be preprocessed and converted to the following NPZ training protocol. This protocol is consistent with the model input specification but is not equivalent to the original datasets' download format.

```text
input: float32 [N,12,36,36]
target: float32 [N,4,448,448]
```

`fake_data.py` automatically writes the `protocol` and `data_source` protocol metadata. These fields must be retained when using real data.

```bash
python scripts/fake_data.py
```

### Training

```bash
python scripts/train.py
```

Multi-GPU training can use:

```bash
torchrun --nproc_per_node=8 scripts/train.py
```

Training sequentially optimizes the conditional-regression and residual-diffusion stages and saves a checkpoint and overall training metrics. The default configuration is intended for quick workflow validation. Formal experiments should use the data scale, model configuration, and training schedule corresponding to the paper.

```text
result/checkpoints/corrdiff.pt
result/training/metrics.json
```

### Training Weights

This repository will provide CorrDiff training weights in the `weight/` folder. The weight files will be uploaded soon and are expected to be available in the near future.

### Inference

```bash
python scripts/inference.py
```

Inference loads the training checkpoint, generates high-resolution ensemble forecasts according to the configuration, and saves the results to:

```text
result/output/predictions.npz
```

### Evaluation and Visualization

```bash
python scripts/result.py
```

Evaluation reports overall MAE, RMSE, ensemble CRPS, and ensemble spread, and generates an ensemble diagnostic figure. Synthetic-data results are only for validating the engineering workflow and do not represent the full performance reported in the paper.

```text
result/evaluation/metrics.json
result/evaluation/ensemble_diagnostics.png
```

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

The use of the code and data in this repository remains subject to the licenses and terms of use of their respective projects.