OneScience-Group/RISE-UNet

RISE-UNet

Model Introduction

RISE-UNet combines deep learning and dynamical forecasts for subseasonal root-zone soil-moisture prediction. It recursively predicts five weekly anomalies and evaluates drought probabilities through ensembles.

Paper: Skillful subseasonal soil moisture drought forecasts with deep learning-dynamic models
https://doi.org/10.1038/s41467-025-62761-3

Model Description

The model was proposed by researchers at Auburn University. It was trained with GLEAM root-zone soil moisture, ERA5 reanalysis, and GEFSv12 and ECMWF S2S reforecasts. By combining residual, inception, squeeze-and-excitation, and UNet++ operations with recursive predictions, it supports weekly root-zone soil-moisture and flash-drought forecasting over the contiguous United States, China, and Australia.

Use Cases

Use Case	Description
Subseasonal soil moisture	Predict root-zone soil-moisture anomalies for weeks 1–5.
Drought forecasting	Identify events below the twentieth percentile.
Ensemble forecasting	Use 11 dynamical members and stochastic inference dropout.
Hybrid modeling	Fuse reanalysis and dynamical reforecasts.
ModelScope/OneCode execution	Validate structured data, training, inference, probabilistic precipitation metrics, and visualization in ModelScope or OneCode.
Multi-GPU training	Start multi-process training through torchrun.

Usage Instructions

Download

hf download OneScience-Group/RISE-UNet --local-dir ./RISE-UNet
cd RISE-UNet

Environment Dependencies

Hardware Requirements

• A GPU or DCU is recommended.
• A CPU can run the default small-sample connectivity configuration.
• DCU users should install DTK 25.04.2 or a compatible OneScience-recommended version first.

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

Synthetic Data

The paper uses a 48x96 0.5-degree regional grid, 11 ensemble members, weekly historical and forecast variables, and GLEAM 0–100 cm root-zone soil-moisture anomalies as targets. Synthetic data preserve the grid, member count, recursive five-week protocol, and RISE operators while reducing initialization count, width, and epochs. Results verify the workflow only and do not represent paper performance.

python scripts/fake_data.py

Training

python scripts/train.py
torchrun --nproc_per_node=2 --nnodes=1 --master_addr="localhost" --master_port=29500 scripts/train.py

The default synthetic run completes five-week recursive optimization, deep-supervision losses, and the ensemble-spread constraint, and both single-process and two-process DDP training have been verified. It produces one recoverable checkpoint and records the CRPSexp training result. Training results are saved to:

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

Weights

The paper's code is available at https://osf.io/6y4kh/, but an independently licensed official pretrained checkpoint was not confirmed.

Inference

python scripts/inference.py

Inference restores the checkpoint, retains stochastic dropout, and recursively generates weeks 1–5 for 11 members. The output shape is [11,5,48,96] and has passed finite-value checks. Inference results are saved to:

result/output/predictions.npz

Evaluation

python scripts/result.py

Evaluation computes weekly ACC, CRPS, and drought GSS and creates a week-3 spatial error figure. All metrics are finite, and the PNG has passed format and non-empty-pixel checks; synthetic results do not represent paper performance. Evaluation results are saved to:

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

Official OneScience Information

Platform	OneScience	OneSkills
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 RISE-UNet specifications, with code licensed under the Apache License 2.0.

The original paper is licensed under CC BY-NC-ND 4.0; the paper and GLEAM, ERA5, GEFSv12, and ECMWF S2S data remain subject to their respective licenses and terms.