StreamflowAI Modelers

StreamflowAI Modelers is a reproducible, station-based workflow for diagnosing and correcting local bias in NOAA's National Water Model (NWM). It pairs hourly NWM streamflow with USGS observations and ERA5 meteorology, diagnoses seasonal errors, and evaluates sequence-learning and ensemble methods across contrasting hydrologic regimes.

Paper

Advancing ML and AI Frameworks for Enhanced Hydrologic Prediction

Basit Akinade, Ahmed Omar, Sanjeev Panta, Saddy Pineda-Castellanos, Mohammad Ali Javidian, and Sushant Mehan.

The project asks two questions:

1. Can machine learning reduce local NWM streamflow bias?
2. Where does correction add the most value as baseline NWM skill changes?

Study design

The same workflow was applied to three unregulated USGS gauges from 2010–2020 at hourly resolution:

Each station directory contains station-specific notebooks, clean data, modeling modules, and documentation.

Figure 1. Hourly observed USGS and simulated NWM hydrographs with ERA5 precipitation at the largest observed event for each site.

Data and predictors

The clean datasets combine:

• USGS observed discharge, used as the training target and never as a predictor
• NWM simulated discharge from the co-located river reach
• ERA5 total precipitation and volumetric soil moisture
• Cyclic encodings of hour, day of year, and month

All records are aligned to a continuous hourly index and converted to cubic meters per second. Missing USGS observations account for 3.4–6.9% of the records. A masked-gap benchmark found linear interpolation to be the most accurate method for 79–99% of test gaps. Imputed hours remain flagged and are excluded from model loss and evaluation metrics.

Modeling framework

The report evaluates four sequence models:

• Recurrent neural network (RNN)
• Gated recurrent unit (GRU)
• Long short-term memory network (LSTM)
• Transformer encoder

Each model is trained under two prediction setups:

• Residual: learn USGS − NWM, then add the predicted correction back to NWM.
• Direct: predict USGS discharge directly without relying on NWM as the baseline.

Models use lookback windows of past hourly inputs. Optuna selects lookback length, hidden size, layers, dropout, learning rate, and attention heads. Training uses Adam, gradient clipping, early stopping, and validation KGE for checkpoint selection.

Three ensemble methods combine the base models:

• Equal-weight averaging
• Error-weighted averaging
• Non-negative constrained stacking with weights that sum to one

Stacking weights are fit only on out-of-fold predictions to prevent leakage. Flow-conditioned split-conformal intervals provide 90% predictive uncertainty bands.

Evaluation

Performance is measured with:

• Kling–Gupta Efficiency (KGE)
• Nash–Sutcliffe Efficiency (NSE)
• Log-transformed NSE for low-flow performance
• Percent bias (PBIAS)

The evaluation uses expanding-window, walk-forward cross-validation inside a final sealed test block. Scalers are refit within each fold, and a lookback-length gap is left at fold boundaries so future information cannot leak into earlier predictions.

Main results

Machine-learning correction improved KGE at all three gauges, with larger nominal gains as baseline NWM skill decreased:

Figure 2. Seasonal KGE for raw NWM, the best single model, and the best ensemble.

Figure 3. KGE and NSE leaderboard on the sealed test block. Stars identify the recommended models.

Watauga and New River passed the study's robustness criteria (logNSE ≥ 0.5 and |PBIAS| ≤ 15%). Little Spearfish did not: its large KGE gain mainly reflects an extremely poor NWM baseline, while its corrected low-flow and volume performance remained unreliable.

The results support a practical framing rule:

Correct NWM residuals where NWM already carries useful signal; predict discharge directly where it does not.

Figure 4. Residual versus direct modeling KGE by architecture and study site.

The selected memory windows also matched basin response time: about 6 hours for flashy Watauga and about 72 hours for the slower New River. Ensembles did not always produce the highest single KGE, but they generally offered more balanced performance across the full metric set.

Figure 5. Recommended-model hydrographs at the largest observed events with flow-conditioned 90% uncertainty bands.

Figure 6. Model rankings across the complete evaluation scorecard; green indicates stronger performance and red weaker performance.

Limits

The surface-oriented predictors do not capture deeper groundwater storage. At Little Spearfish, seasonal timing correlation was near zero and the model could not recover a missing subsurface signal. The current study covers only three gauges; future work should add intermediate-skill sites, groundwater-storage proxies, antecedent multi-month indices, and multi-seed experiments.

Requirements

• Python 3.10 or newer; Python 3.12 is recommended
• venv
• Git

PyTorch is installed from PyPI by default. For a CUDA-specific build, follow the instructions at pytorch.org before installing the remaining requirements.

Installation

bash git clone https://github.com/NWC-CUAHSI-Summer-Institute/StreamflowAI-Modelers.git cd StreamflowAI-Modelers python3 -m venv .venv source .venv/bin/activate python -m pip install --upgrade pip python -m pip install -r requirements.txt

On Windows PowerShell, create and activate the environment with:

powershell py -m venv .venv .venv\Scripts\Activate.ps1

Optionally register it as a Jupyter kernel:

bash python -m ipykernel install --user --name streamflow-ai --display-name "Streamflow AI"

Running the workflow

Start JupyterLab from the repository root:

bash jupyter lab

Open a station directory and run its notebooks in stage order:

1. Stage1a_Imputation_Method_Selection_<site>.ipynb
2. Stage1b_Build_Clean_Dataset_From_Raw_<site>.ipynb
3. Stage2_Diagnose_Bias_and_Select_Inputs_<site>.ipynb
4. Stage3a_Seasonality_Regime_<site>.ipynb
5. Stage3b_Seasonality_WaterYear_<site>.ipynb
6. Stage4_Modelling/Stage4_Master_Modeling_<site>.ipynb

Stage 1a benchmarks imputation methods against a seed clean dataset. Restore an existing clean dataset before running it, or begin with Stage 1b when rebuilding from raw inputs.

Stages 2–4 expect:

text <station>/team_clean_dataset_<site>/data/clean_<site>.parquet

Raw source data is not versioned because of its size. If only the clean CSV is present, either run Stage 1b with the original inputs or convert the CSV:

bash python -c "import pandas as pd; p='STATION/team_clean_dataset_SITE/data/clean_SITE'; pd.read_csv(p+'.csv').to_parquet(p+'.parquet', index=False)"

Replace STATION and SITE with the station directory and USGS site number.

To execute Stage 4 non-interactively, run this from the station directory:

bash jupyter nbconvert --to notebook --execute --inplace \ --ExecutePreprocessor.timeout=-1 \ Stage4_Modelling/Stage4_Master_Modeling_SITE.ipynb

Stage 4 supports SARIMAX, RNN, GRU, LSTM, Transformer, residual and direct prediction, hyperparameter tuning, and ensemble evaluation. Generated artifacts are written to:

text <station>/Stage4_Modelling/artifacts/<site>/

Generated artifacts, notebook checkpoints, local environments, raw downloads, and workflow output directories are excluded from Git.

Project team

• Basit Akinade — The University of Alabama
• Ahmed Omar — Texas A&M University–Corpus Christi
• Sanjeev Panta — University of Louisiana at Lafayette
• Saddy Pineda-Castellanos — Utah State University
• Mohammad Ali Javidian — Appalachian State University, theme leader
• Sushant Mehan — South Dakota State University, theme leader

This work was supported by CUAHSI and the Cooperative Institute for Research to Operations in Hydrology (CIROH) under NOAA award NA22NWS4320003.