OneScience-Group/Streamflow-LSTM

<p align="center"><strong><span style="font-size: 30px;">Streamflow-LSTM</span></strong></p>

Model Introduction

Streamflow-LSTM is an engineering reproduction of the gauge-specific LSTM river-flow forecasting method proposed by Hunt et al. It uses the preceding seven days of six-hourly meteorological and hydrological sequences to forecast 40 six-hourly leads, or ten days, for ten western US gauges.

Paper: Using a long short-term memory (LSTM) neural network to boost river streamflow forecasts over the western United States https://doi.org/10.5194/hess-26-5449-2022

Model Description

The model was proposed by researchers from the University of Reading, the European Centre for Medium-Range Weather Forecasts, Loughborough University, and the UK Centre for Ecology and Hydrology. It was trained with streamflow observations from ten western US gauges together with catchment-averaged meteorological forecasts and hydrological histories. The model is suitable for improving gauge-specific river-flow forecasts over the following ten days and for comparison with persistence and GloFAS baselines. Its key feature is independent temporal learning for each gauge, combined with member selection and ensemble averaging to improve forecast robustness.

Use Cases

Usage Instructions

1.OneCode

Try intelligent, one-click AI4S programming

2. Download and Installation

hf download OneScience-Group/Streamflow-LSTM --local-dir ./Streamflow-LSTM
cd Streamflow-LSTM

Environment Dependencies

Hardware Requirements

• —A GPU or DCU is recommended for the paper-scale ensemble.
• —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 dataset combines sources including ERA5, IFS, GloFAS, and USGS for ten gauges, arranging the previous seven days as 28 six-hourly steps with 23 variables per step. The repository's deterministic synthetic data retain ten gauges, [B,28,23] inputs, and 40 six-hourly leads while reducing the numbers of training, validation, and forecast samples. These data validate engineering only and do not represent official distributions, paper-scale training, or paper performance; observed flow is frozen after issue time in forecast inputs to prevent future-observation leakage.

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

Training uses MSE, Adam at 0.001, and 0.1 dropout for every gauge member, with member tasks divided among processes. Pass --paper to restore 50 hidden units, 100 members per gauge, and selection of the best five members. All gauges and members are consolidated into one checkpoint, and the artifacts are:

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

Trained Weights

The paper does not provide directly loadable official weights, and weight/ contains only a status note. Local training stores every ensemble member for all ten gauges, normalization statistics, and validation NSE values in one result/checkpoints/streamflow_lstm.pt; it must not be represented as an official pretrained checkpoint.

Inference

python scripts/inference.py

Inference restores every gauge member from the single checkpoint and validates its format version and ten-gauge order. It selects the best two members by validation NSE in default mode or the best five in paper mode, evaluates each [C,40,28,23] input, and averages the ensemble. Outputs are clipped to non-negative streamflow and retain 40 six-hourly leads in m3 s-1. Inference results are saved to:

result/output/predictions.npz

Evaluation and Visualization

python scripts/result.py

Evaluation computes gauge-level KGE, NSE, and RMSE at two-, five-, and eight-day leads, including the correlation, variability, and bias components of KGE. Forecasts are compared with issue-time persistence and a synthetic GloFAS proxy. The figure presents mean RMSE over all leads and five-day KGE for each gauge; synthetic-data results validate engineering only. Evaluation artifacts are saved to:

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

Official OneScience Information

Citation and License

This repository is an independent engineering reproduction of the public Streamflow-LSTM specifications, with code licensed under the Apache License 2.0.

The original paper is licensed under CC BY 4.0; the paper, model weights, and ERA5, IFS, GloFAS, and USGS data remain subject to the licenses and terms of their respective projects.