Differentiable hybrid hydrological model for CAMELS-US

A physics-based, differentiable hydrological modeling framework that combines neural networks with process-based hydrological models on the CAMELS-US dataset.

Overview

This repository implements a hybrid modeling approach where neural networks (MLP or LSTM) estimate static or dynamic parameters for a differentiable HBV hydrological model. The entire framework is end-to-end differentiable, enabling gradient-based optimization.

This repository supports four hybrid model configurations:

• MLP+1HBV: MLP estimates static parameters for a single HBV model
• MLP+16HBV: MLP estimates static parameters for 16 parallel HBV models
• LSTM+1HBV: LSTM estimates dynamic parameters for a single HBV model
• LSTM+16HBV: LSTM estimates dynamic parameters for 16 parallel HBV models

Publication: [DOI link to be added]

Repository Structure

├── models/multi_hbv.py          # Modular implementations of MLP, LSTM, and differentiable HBV
├── 01preprocessing/             # Data preprocessing scripts
├── 02hyperparameter_tuning/     # Hyperparameter optimization scripts for hybrid models
├── 01best_mlp_hbv.py           # Train/test MLP+HBV models
├── 02best_lstm_hbv.py          # Train/test LSTM+HBV models
├── 03get_params_lstm_hbv.py    # Extract dynamic HBV parameters
├── 04get_dynamic_simulation_lstm_hbv.py  # Run inference with sequence-varying parameters
├── 05get_integrated_gradients_lstm_hbv.py  # Compute feature importance
└── 06plot_figures.ipynb        # Generate all paper figures

Main Scripts

• 01best_mlp_hbv.py — Train, validate, and test MLP+HBV or MLP+16HBV models using optimized hyperparameters.

• 02best_lstm_hbv.py — Train, validate, and test LSTM+HBV or LSTM+16HBV models using optimized hyperparameters.

• 03get_params_lstm_hbv.py — Extract dynamic (sequence-varying) HBV parameters generated by the LSTM+1HBV model.

• 04get_dynamic_simulation_lstm_hbv.py — Run LSTM+1HBV inference with sequence-varying parameters and save streamflow simulations.

• 05get_integrated_gradients_lstm_hbv.py — Compute Integrated Gradients feature importance for LSTM+1HBV model inputs.

• 06plot_figures.ipynb — Generate all figures presented in the paper, including supplementary materials.

How to Run This Repository

Follow these steps to reproduce the results:

1. Setup Environment → Create a virtual environment and install required packages from requirements.txt

2. Prepare Data → Download CAMELS-US dataset and run preprocessing script in 01preprocessing/

3. Tune Models → Run hyperparameter optimization in 02hyperparameter_tuning/

4. Train & Evaluate → Use best hyperparameter and execute 01best_mlp_hbv.py and 02best_lstm_hbv.py to train and make inference on all study basins. Results will be saved as csv files.

5. Extract Analysis → Run scripts 03, 04, 05 to run LSTM+1HBV hybrid model and extract dynamic parameters, simulate streamflow with dynamic parameters, and compute feature importance with Integrated Gradients, respectively. Results will be saved as csv files.

6. Generate Figures → Run 06plot_figures.ipynb to reproduce all figures using saved csv results.