README.md

---
license: cc-by-nc-nd-4.0
language:
- en
tags:
- OneScience
- Earth Science
- Rainfall Correction
- UnetDif
frameworks: PyTorch
---

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

# Model Introduction

UnetDif uses the paper's bias-targeted multitask U-Net to correct short-range heavy rainfall forecasts.

Paper: Bias-targeted deep learning enhances short-range heavy rainfall forecasts  
https://doi.org/10.1038/s41612-026-01366-z

# Model Description

The method reproduced by UnetDif was proposed by research teams from Zhejiang Meteorological Observatory, the National Meteorological Center, and Zhejiang University. The paper constructs Yangtze River Delta samples from European Centre for Medium-Range Weather Forecasts (ECMWF) fields, China Meteorological Administration Multi-source Precipitation Analysis (CMPA), and terrain data. The model learns `ECMWF - CMPA` rainfall bias and corrects three-hour precipitation through a shared U-Net backbone with four heads for dry areas, false alarms, positive bias, and negative bias.

# Use Cases

| Use Case | Description |
| :---: | :--- |
| Three-hour rainfall-bias correction | Generate corrected ECMWF rainfall from 39-channel inputs on the fixed Yangtze River Delta grid. |
| ModelScope/OneCode execution | Validate synthetic-data generation, six-loss training, inference, evaluation, and visualization in ModelScope or OneCode environments. |
| Hugging Face/OneCode execution | Run the complete engineering workflow in Hugging Face or OneCode environments. |
| Multi-GPU training | Launch distributed data-parallel training with `torchrun`. |

# Usage Instructions

## 1.OneCode

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

[Try intelligent, one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)

## 2. Download and Installation

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

### Environment Dependencies

**Hardware Requirements**

- A GPU or DCU is recommended.
- A CPU can be used for connectivity validation with the default small-sample configuration.
- DCU users must install DTK first. DTK 25.04.2 or later, or the OneScience-recommended version matching the cluster, is recommended.

**DCU Environment**

```bash
# Activate DTK and Conda first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
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
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
```

### Training Data

Training data contain ECMWF weather forecasts, CMPA precipitation observations, and static terrain. Each sample uses 39 input channels, with corresponding CMPA precipitation and ECMWF precipitation bias as targets. The data use a `56×56` Yangtze River Delta grid at `0.125°` resolution. Each rain day contains eight consecutive three-hour periods, with precipitation in `mm/3h`. Synthetic data validate training, inference, and evaluation only and do not represent the distributions of ECMWF, CMPA, real terrain, or formal paper performance.

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

### Training

For single-device training, use:

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

For multi-GPU training, use:

```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 learns dry-area, false-alarm, and signed precipitation-bias tasks and saves the model checkpoint and training metrics. Results are saved to:

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

### Trained Weights

This repository does not include weights under `weight/`. The paper does not provide a confirmed official checkpoint, and the engineering checkpoint trained on synthetic data is for workflow validation only and is not claimed to be compatible with external weights.

### Inference

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

Inference loads the checkpoint and applies all four output heads under `bias = ECMWF - CMPA` to produce nonnegative corrected rainfall. Results are saved to:

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

### Evaluation and Visualization

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

Evaluation produces categorical and spatial-skill metrics for individual three-hour periods and 24-hour accumulated precipitation and saves structured results. It also generates comparisons of ECMWF, CMPA, corrected precipitation, and errors. Synthetic-data results validate the engineering workflow only and do not represent formal paper performance.

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

# 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 engineering reproduction of the public UnetDif specifications.

Use of this repository's code, official model weights, and data remains subject to the licenses and terms of their respective projects.