OneScience-Group/MetNet-2

MetNet-2

Model Introduction

MetNet-2 generates high-resolution probabilistic precipitation forecasts up to 12 hours ahead from radar, satellite, and atmospheric-state inputs for short-range forecasting and extreme-precipitation risk analysis.

Paper: Deep learning for twelve hour precipitation forecasts
https://doi.org/10.1038/s41467-022-32483-x

Model Description

The method was proposed by the Google Research team. The paper constructs its 2017-2020 training and test data from MRMS radar, GOES satellite imagery, and HRRR assimilated atmospheric states. The model predicts grid-cell precipitation probability distributions at two-minute intervals up to 12 hours ahead.

Use Cases

Use Case	Description
Categorical precipitation forecasting	Validate the 512-category conditional distribution and 12-hour lead-time protocol.
Core-method validation	Validate ConvLSTM, lead-time FiLM, and multiscale dilated residual stacks.
Local engineering validation	Exercise the complete logical 641×512×512 contract through a deterministic procedural field without materializing the full input.
ModelScope/OneCode execution	Validate structured data, training, inference, probabilistic precipitation metrics, and visualization in ModelScope or OneCode environments.
Multi-GPU training	Validate distributed training and the checkpoint workflow through torchrun.

Usage Instructions

1.OneCode

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

Try intelligent, one-click AI4S programming

2. Download and Installation

hf download OneScience-Group/MetNet-2 --local-dir ./MetNet-2
cd MetNet-2

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 in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.

DCU Environment

# 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

# 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

The paper training inputs combine MRMS, HRRR, GOES, static geography, and temporal information under a 641-channel, 512×512 spatial-domain protocol. This repository creates eight deterministic window records and constructs only the selected 32×32 windows plus their halos at runtime while retaining the complete logical shape and channel grouping. The synthetic data validate engineering connectivity only and do not represent real meteorological distributions, paper-scale training, or paper performance.

python scripts/fake_data.py

Training

For single-GPU training, use:

python scripts/train.py

For multi-GPU training, use:

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

The default reduces sample count, network width, residual blocks, and training steps without reducing the 641 channels, 512 precipitation categories, or 12-hour lead protocol. Formal experiments require real MRMS, GOES, and HRRR data and full computing resources, with artifacts saved to:

result/checkpoints/metnet_2.pt
result/training/metrics.json

Trained Weights

The paper does not provide directly loadable official model weights, and this repository bundles no weights under weight/. The locally trained checkpoint is saved to result/checkpoints/metnet_2.pt and must not be represented as an official pretrained weight.

Inference

python scripts/inference.py

Inference loads the training checkpoint and generates 512-category precipitation probabilities and their CDF for selected spatial windows while preserving coverage and completeness metadata. Complete numerical results are saved to:

result/output/predictions.npz

Evaluation and Visualization

python scripts/result.py

Evaluation computes discrete CRPS, Brier Score, and CSI at multiple precipitation thresholds and generates target, expected-rate, and error comparisons. Synthetic-data results validate engineering only and do not represent paper performance; outputs are saved to:

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 MetNet-2 specifications.

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