OneScience-Group/SmaAtUNet

SmaAt-UNet

Model Introduction

SmaAt-UNet is a lightweight convolutional neural network for precipitation nowcasting. It takes consecutive radar precipitation maps as input and predicts the precipitation distribution over multiple future time steps. The model incorporates a Convolutional Block Attention Module and depthwise-separable convolutions into U-Net to maintain competitive forecasting performance with fewer parameters.

Paper: SmaAt-UNet: Precipitation Nowcasting using a Small Attention-UNet Architecture
https://arxiv.org/abs/2007.04417

Model Description

SmaAt-UNet was proposed by researchers at Maastricht University. The model uses approximately 420,000 Dutch KNMI radar precipitation maps from 2016 to 2019 at 5-minute intervals and is also evaluated on French binary cloud-cover data. The model is suitable for short-term precipitation nowcasting, cloud-cover prediction, and lightweight meteorological image regression tasks.

Use Cases

Scenario	Description
Precipitation nowcasting	Predict precipitation for the next 30 minutes from radar maps of the previous 60 minutes.
Multi-step image regression	Output 6 consecutive precipitation prediction maps at once.
Attention-based feature extraction	Use CBAM to enhance important channels and spatial regions.
Local workflow validation	Use synthetic radar sequences to check training, inference, evaluation, visualization, and checkpoint workflows.
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

2. Download and Installation

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

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

This repository uses normalized synthetic radar sequences to validate the engineering workflow. Each sample retains the actual data dimensions of the official precipitation task: the input is 12×288×288, the target is 6×288×288, and adjacent images are 5 minutes apart.

python scripts/fake_data.py

Training

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 configuration only reduces the number of samples, base feature width, and training epochs; it does not change the number of input frames, number of output frames, or spatial grid dimensions.

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

Training Weights

This repository does not include synthetic or trained weights, and the authors have not currently released a directly downloadable pretrained checkpoint.

Inference

python scripts/inference.py

Inference loads the training checkpoint, generates precipitation predictions and attention maps for the next 6 time steps from 12 historical precipitation maps, and saves them to:

result/output/predictions.npz

Evaluation and Visualization

python scripts/result.py

The evaluation computes MSE, MAE, Precision, Recall, F1, CSI, FAR, and HSS, and generates comparison plots of the targets, predictions, and absolute errors for the next 30 minutes. Results on synthetic data are only for validating the engineering workflow and do not represent metrics on the paper's actual radar test set.

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 publicly available specifications in the SmaAt-UNet paper.

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