Streamflow Prediction with Sequence Models

This project implements HWRS640 Assignment 4.

Structure

• main.py - root CLI entry point
• cli/ - command-line parsing and orchestration
• dataset/ - MiniCAMELS access, preprocessing, sequence dataset, dataloaders
• model/ - sequence models such as LSTM and Transformer
• training/ - trainer, losses, checkpointing, early stopping
• evaluation/ - test-time evaluation workflows
• util/ - config, metrics, data utilities, logging utilities
• visualization.py - plotting functions
• configs/ - reusable experiment configuration files
• outputs/ - checkpoints, metrics, and figures

The assignment lists data.py, train.py, utils.py, and model.py as simple module names. This repository keeps the folder structure you requested. Thin compatibility modules are included for data.py, train.py, and utils.py. The model code lives in the model/ package, because a repository cannot have both a root model.py file and a model/ directory with the same name.

Setup

python3 -m venv .venv
source .venv/bin/activate
python3 -m pip install -r requirements.txt

Google Colab

!git clone https://github.com/mfarmani95/Streamflow_Prediction.git
%cd Streamflow_Prediction
!pip install -r requirements.txt

import torch
print(torch.cuda.is_available())
print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else "No GPU")

CLI Command Cookbook

All commands should be run from the repository root.

Environment Setup

Using uv:

uv venv --python 3.11 .venv
source .venv/bin/activate
uv pip install -r requirements.txt

Using standard pip:

python3 -m venv .venv
source .venv/bin/activate
python3 -m pip install -r requirements.txt

If MiniCAMELS is missing, install it directly from GitHub:

python3 -m pip install git+https://github.com/BennettHydroLab/minicamels.git

Inspect and Analyze the Dataset

Quick dataset summary:

python3 main.py summarize-data

Dataset summary plus the assignment exploratory plots:

python3 main.py summarize-data --make-plots --output-dir outputs

Split-aware data analysis from the YAML config:

python3 main.py analyze-data --config configs/default.yaml --output-dir outputs/data_analysis

If the MiniCAMELS data are stored in a custom location, add --data-dir:

python3 main.py analyze-data \
  --config configs/default.yaml \
  --data-dir /path/to/minicamels \
  --output-dir outputs/data_analysis

Train One Model

Train using the default YAML:

python3 main.py train --config configs/default.yaml

Train an LSTM with explicit arguments:

python3 main.py train \
  --config configs/default.yaml \
  --model lstm \
  --seq-len 120 \
  --window-stride 120 \
  --hidden-size 128 \
  --num-layers 1 \
  --dropout 0.1 \
  --batch-size 32 \
  --lr 0.001 \
  --loss nse \
  --epochs 50 \
  --split-strategy stratified \
  --split-stratify-attribute aridity \
  --train-basin-count 30 \
  --val-basin-count 10 \
  --test-basin-count 10 \
  --output-dir outputs/single_lstm \
  --checkpoint outputs/single_lstm/best_model.pt

Train a Transformer with explicit arguments:

python3 main.py train \
  --config configs/transformer_sweep.yaml \
  --model transformer \
  --seq-len 120 \
  --window-stride 120 \
  --hidden-size 64 \
  --num-layers 2 \
  --nhead 4 \
  --dim-feedforward 128 \
  --dropout 0.2 \
  --batch-size 128 \
  --lr 0.001 \
  --loss mse \
  --epochs 100 \
  --split-strategy stratified \
  --split-stratify-attribute aridity \
  --train-basin-count 30 \
  --val-basin-count 10 \
  --test-basin-count 10 \
  --output-dir outputs/single_transformer \
  --checkpoint outputs/single_transformer/best_model.pt

Evaluate and Plot One Trained Model

Evaluate a saved checkpoint on the held-out test basins:

python3 main.py evaluate \
  --checkpoint outputs/single_lstm/best_model.pt \
  --output-dir outputs/single_lstm

Create the basic evaluation plots from the evaluation outputs:

python3 main.py plot \
  --checkpoint outputs/single_lstm/best_model.pt \
  --output-dir outputs/single_lstm

Create paper-style analysis plots and tables for one completed run directory:

python3 main.py analyze-run --run-dir outputs/single_lstm

This writes files such as metrics_summary.csv, parity_test_all.png, timeseries_best_basin.png, timeseries_worst_basin.png, basin_kge_cdf.png, flow_regime_metrics.csv, and discussion_notes.txt under outputs/single_lstm/paper_analysis/.

Run the LSTM Sweep

Check what will run without training:

python3 main.py sweep --config configs/default.yaml --dry-run

Run or resume the LSTM sweep:

python3 main.py sweep \
  --config configs/default.yaml \
  --output-root outputs/sweeps_mse \
  --skip-existing

Override a sweep from the command line:

python3 main.py sweep \
  --config configs/default.yaml \
  --seq-lens 30 60 90 120 \
  --hidden-sizes 32 64 128 256 \
  --batch-sizes 32 64 128 \
  --loss nse \
  --lr 0.001 \
  --output-root outputs/sweeps_lstm_custom \
  --skip-existing

Analyze the Sweep Results

Rank sweep runs and select the best model by validation NSE:

python3 main.py analyze-sweep \
  --sweep-root outputs/sweeps_mse \
  --selection-metric best_val_nse

Create grouped training-history plots for the sweep:

python3 main.py sweep-plots --sweep-root outputs/sweeps_mse

Plot only selected hyperparameter effects:

python3 main.py sweep-plots \
  --sweep-root outputs/sweeps_mse \
  --effects seq_len hidden_size batch_size lr

Analyze the best LSTM run:

python3 main.py analyze-run \
  --run-dir outputs/sweeps_mse/seq120_hidden128_batch032_lr0p001

Run and Analyze the Transformer Sweep

Dry run:

python3 main.py sweep --config configs/transformer_sweep.yaml --dry-run

Run or resume:

python3 main.py sweep \
  --config configs/transformer_sweep.yaml \
  --output-root output/transformer_MSE_T1 \
  --skip-existing

Analyze the Transformer sweep:

python3 main.py analyze-sweep \
  --sweep-root output/transformer_MSE_T1 \
  --selection-metric best_val_nse

Create Transformer sweep training-history plots:

python3 main.py sweep-plots --sweep-root output/transformer_MSE_T1

Analyze the best Transformer run:

python3 main.py analyze-run \
  --run-dir output/transformer_MSE_T1/seq120_hidden064_batch128_lr0p001_drop0p2

Compare LSTM and Transformer

Create side-by-side model comparison plots and tables:

python3 main.py compare-runs \
  --run-dirs \
    outputs/sweeps_mse/seq120_hidden128_batch032_lr0p001 \
    output/transformer_MSE_T1/seq120_hidden064_batch128_lr0p001_drop0p2 \
  --labels LSTM Transformer \
  --output-dir reports/problem4/model_comparison

This writes overall metric comparison plots, parity overlays, per-basin NSE comparison, KGE CDF comparison, seasonal metrics, flow-regime comparison, and best/worst basin hydrographs.

Create the Problem 4 Report

Generate the Markdown report, PDF report, and all supporting figures:

python3 main.py problem4-report \
  --lstm-run-dir outputs/sweeps_mse/seq120_hidden128_batch032_lr0p001 \
  --transformer-run-dir output/transformer_MSE_T1/seq120_hidden064_batch128_lr0p001_drop0p2 \
  --output-dir reports/problem4

Outputs:

reports/problem4/problem4_report.md
reports/problem4/problem4_report.pdf
reports/problem4/lstm_analysis/
reports/problem4/transformer_analysis/
reports/problem4/model_comparison/

Useful Help Commands

python3 main.py --help
python3 main.py train --help
python3 main.py evaluate --help
python3 main.py sweep --help
python3 main.py analyze-sweep --help
python3 main.py analyze-run --help
python3 main.py compare-runs --help
python3 main.py problem4-report --help

When using --config, command-line flags override YAML values. For example:

python3 main.py train --config configs/default.yaml --epochs 2 --limit-basins 10

The default loss is masked MSE, available as --loss mse or --loss masked_mse. You can also use --loss mae, --loss masked_mae, or --loss kge / --loss nse. The KGE and NSE training losses are minimized as 1 - KGE and 1 - NSE, while validation and test reports still show the actual KGE and NSE metric values. For NSE/KGE validation, val_loss is computed from the full validation set metric, not as an average of per-batch NSE/KGE.

The sweep grid is controlled from the sweep section of the YAML config. The current LSTM sweep in configs/default.yaml is:

sweep:
  output_root: outputs/sweeps_mse/
  loss: nse
  learning_rate: 0.001
  grid:
    seq_len: [30, 60, 90, 120]
    hidden_size: [32, 64, 128, 256]
    batch_size: [32, 64, 128, 256]
    learning_rate: [0.001, 0.003, 0.0005]

Any training argument can be swept by adding it to sweep.grid, for example learning_rate: [0.0001, 0.0005, 0.001] or dropout: [0.0, 0.1, 0.2]. Each run is saved under the sweep output_root with its own checkpoint, history, metrics, and plots. Sweep runs use non-overlapping windows by setting the stride equal to each sequence length when window_stride is null.

After the sweep finishes, use python3 main.py sweep-plots --sweep-root <sweep-root> to create combined training-history figures under <sweep-root>/comparison_plots/. These compare one hyperparameter at a time for the variables that changed across the sweep runs.

Use python3 main.py analyze-sweep --sweep-root <sweep-root> to rank runs by validation performance, write top-run tables, create parameter-effect summary plots, and save best_model_summary.txt. By default the best run is selected by best_val_nse; test NSE/KGE are reported after selection, not used to choose the hyperparameters.

configs/transformer_sweep.yaml mirrors the LSTM sweep layout but sets model.name: transformer, so it can be launched with the same sweep command when you are ready to compare architectures.

Use python3 main.py analyze-run --run-dir <completed-run-dir> for paper-style test-set figures: overall metric tables, parity plots, best/worst basin hydrographs, flow-duration curves, per-basin NSE/KGE plots, KGE CDF over catchments, seasonal skill, flow-regime bias, and discussion notes. Once both LSTM and Transformer runs are evaluated, use python3 main.py compare-runs --run-dirs <lstm-run-dir> <transformer-run-dir> --labels LSTM Transformer to put their metrics, parity plots, basin NSE, catchment KGE CDFs, seasonal skill, and selected hydrographs in the same comparison folder.

For the assignment write-up, python3 main.py problem4-report --lstm-run-dir <lstm-run-dir> --transformer-run-dir <transformer-run-dir> creates reports/problem4/problem4_report.md, reports/problem4/problem4_report.pdf, and the supporting analysis figures.

The default YAML uses a fixed basin split of 30 training basins, 10 validation basins, and 10 test basins. The split is stratified by aridity with seed: 42, so validation and test basins better cover the same aridity range as training basins while still avoiding basin leakage. You can switch back to a plain random basin split with:

data:
  split_strategy: random

The split-aware exploratory plots are written to outputs/data_analysis/. They include streamflow and forcing distributions, static attribute plots such as aridity and q_mean, basin locations by split, missing-value diagnostics, correlation heatmaps, non-overlapping sequence target distributions, and example hydrographs.

Current Implementation Notes

The default data split is by basin: 30 training basins, 10 validation basins, and 10 test basins, stratified by aridity. Normalization statistics are fitted only from the training basins, then reused for validation and test. The default static attributes are physical/climate catchment descriptors and exclude discharge-derived signatures such as q_mean, runoff_ratio, hfd_mean, and baseflow_index to avoid leakage.

MiniCAMELS is installed directly from GitHub because it is not published on PyPI:

python3 -m pip install git+https://github.com/BennettHydroLab/minicamels.git