File size: 5,134 Bytes
45b074b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | ---
license: apache-2.0
language:
- en
tags:
- OneScience
- Earth Science
- Soil Moisture
- Drought Forecasting
- Subseasonal Forecasting
- Ensemble Forecasting
frameworks: PyTorch
---
<p align="center"><strong><span style="font-size: 30px;">RISE-UNet</span></strong></p>
# 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
```bash
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**
```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
```
## 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.
```bash
python scripts/fake_data.py
```
## Training
```bash
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:
```text
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
```bash
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:
```text
result/output/predictions.npz
```
## Evaluation
```bash
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:
```text
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.
|