README.md

---
license: mit
language:
- en
- zh
tags:
- OneScience
- Earth science
- Weather forecasting
- Short-term precipitation nowcasting
- Deep generative models
- Generative adversarial networks
frameworks: PyTorch
datasets:
  - OneScience/ERA5
---
<p align="center">
  <strong>
    <span style="font-size: 30px;">DGMR</span>
  </strong>
</p>


# Model Introduction

DGMR (Deep Generative Model of Radar) was proposed by DeepMind in 2021. It is a short-term precipitation nowcasting model based on a conditional generative adversarial network (cGAN). The generator consists of a latent conditioning stack, a context conditioning stack, and an autoregressive sampler based on multiple ConvGRU layers. The discriminator evaluates both spatial and temporal dimensions, and the training objective combines hinge GAN loss with a grid-cell regularizer.

Paper:Skillful Precipitation Nowcasting using Deep Generative Models of Radar

https://arxiv.org/abs/2104.00954

# Model Description

DGMR is a probabilistic short-term precipitation nowcasting model: it takes four consecutive radar echo frames and generates 18 future frames at once (5-minute steps, 90 minutes in total, in the paper). Its output is a sample rather than a deterministic estimate. This repository is organized from the official Open Climate Fix PyTorch implementation (`openclimatefix/dgmr`, MIT License) and integrated with the OneScience data loading and training workflow.

# Use Cases

| Scenario | Description |
| :---: | :--- |
| Short-term precipitation nowcasting research | Train a cGAN on radar echo sequences to generate future radar echo frames. |
| Probabilistic forecast output | Obtain multiple generated samples of future fields through latent-variable sampling. |
| 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.


### Download the Model Package

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

### 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 a single-channel radar sequence from `num_context`, `forecast_steps`, batch size, and radar grid size. The current small configuration uses 4 input frames, 6 output frames, and a 128×128 grid.

### 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. Pretrained weights are not provided by default; users may train the model using the paper configuration. Official DeepMind pretrained weights require access approval and differ from the current configuration in channel count and grid size, so they must be aligned before loading.

### Inference

Inference reads `data/checkpoints/model_bak.pth`：

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

Prediction results are written frame by frame 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 Open Climate Fix PyTorch package `openclimatefix/dgmr` (MIT License). The related network modules (common/layers/generators/discriminators/losses) are embedded unchanged in `model/dgmr_official/` (with the HuggingFace hub mixin and pytorch_lightning training loop removed); `model/dgmr.py` is only a YAML-driven thin wrapper.
- `conf/config.yaml` uses a small configuration (`forecast_steps=6`, `output_shape=128`, `latent_channels=384`, `context_channels=192`) for connectivity validation by default; paper-level reproduction requires 4→18 frames, a 256×256 grid, `latent_channels=768`, and `context_channels=384`.
- The paper-level configuration (`num_context=4`, `forecast_steps=18`) requires at least `T >= num_context + forecast_steps + 1 = 23` frames per year; the current synthetic data uses `T=10` (`time_step=6h`), so the default configuration uses 4→6 frames.
- The discriminator contains BatchNorm1d, so the training batch must be `>= 2` (with synthetic data, the training set uses two samples from two years, exactly one batch).
- 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), and the number of randomly sampled time frames in the discriminator.

# 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 DGMR. The model source is based on the Open Climate Fix `openclimatefix/dgmr` implementation and follows the MIT License.
- Please cite:Ravuri et al. Skilful Precipitation Nowcasting using Deep Generative Models of Radar. Nature 597, 672-677, 2021.