An ML plugin for Pywr-DRB, focused on water temperature and salt front prediction.

Coupled Pywr-DRB = Pywr-DRB v2.0+ + PywrDRB-ML plugin.

See the demo in Tutorial 05 Run Coupled PywrDRB.ipynb.

Note: Please install Pywr-DRB v2.0+ first, which contains all dependencies required to run Coupled Pywr-DRB. Then clone this repo to run run_coupled_pywrdrb.py. In Tutorial 05 Run Coupled PywrDRB.ipynb or run_coupled_pywrdrb.py, you need to change root_dir to the directory where you cloned PywrDRB-ML.

The figure below shows a snapshot of the trained LSTMs' performance.

Model evaluation of daily maximum stream temperature at Lordville ($T_max$) and 7-day average salt front location in the estuary ($Saltfront$) dynamics in the DRB from 1979 to 2023. a) and c) shows Root Mean Square Error (RMSE) across different ranges of T_max and Saltfront, respectively. b) and d) present time series comparison of observed (obs; dots) and simulated (sim; solid lines) T_max and Saltfront in 2007.

Designed to run through 2023, as the reconstruction data end in that year.

• src/data_retrievel/pull_gridmet.py: Downloads and processes GridMET meteorological data for use as model input features.
  • data\raw\gridmet_lordville.csv
• src/pull_nwis.py: Downloads and processes streamflow and water quality data from the USGS National Water Information System (NWIS) for use as model input or validation data.
  • data\raw\nwis_Cannonsville_degC_mgd.csv
  • data\raw\nwis_Lordville_degC_mgd.csv
  • data\raw\salt_front_data.csv

• data\raw\sf_data.xlsx (.csv)

  • https://drbc.net/Sky/hydro/saltfront.html#header2

• data\raw\drb_reservoir_storage_mg_2000_2024.csv

  • From Marilyn (should be identical to what we have in PywrDRB)

• src/data_retrievel/pull_pywrdrb_simulated_data.py
  • Simulates pub_nhmv10_BC_withObsScaled, the reconstruction data (median).
  • From 1945 to 2023

• src\form_TempLSTM_db\use_lstm_to_fill_gaps.py
  • data\raw\lstm_simed_T_degC.csv
  • Later used to fill gaps via linear interpolation over 3 consecutive days
• src/form_TempLSTM_db/decompose_Ti_Qi.py
  • data\decomposed_Ti_Qi\df_QobsTavg.csv
  • data\decomposed_Ti_Qi\df_QbcTavg.csv
  • data\decomposed_Ti_Qi\df_QobsTmax.csv
  • data\decomposed_Ti_Qi\df_QbcTmax.csv
• src\form_TempLSTM_db\create_TempLSTM_database.py
  • data\database\TempLSTM_database.csv

• src\form_SalinityLSTM_db\create_SalinityLSTM_database.py
  • data\database\SalinityLSTM_database.csv

• src/prep_data_utils.py: Utility functions for data preparation and preprocessing.
• src/prep_data.py: Main script for preparing datasets, likely orchestrates data loading and processing.
• src/torch_bmi.py: Implements BMI (Basic Model Interface) integration for PyTorch models.
• src/torch_models.py: Defines PyTorch model architectures, including LSTM models for temperature and salinity prediction.
• src/sampling_utils.py: Provides sampling functions for uncertainty quantification in LSTM model outputs, such as sampling from Gaussian Mixture Models (GMM).
• src/training_utils.py: Contains training utilities for LSTM models, including custom loss functions (e.g., MaskedGMMLoss, MaskedCMALLoss, MaskedUMALLoss), evaluation metrics (such as rmse_masked), and model fitting routines (e.g., fit_torch_model).

• lstm_train_TempLSTM.py
  • Creates the model, trains it, and runs basic analysis
• lstm_train_SalinityLSTM.py
  • Creates the model, trains it, and runs basic analysis

• lstm_evaluation.py
  • Evaluates model performance and feature importance

Lin, C.-Y., Reed, P., Amestoy, T., Zwart, J., Gorski, G., & Sunkara, S. V. (2026). Integrating deep learning and thermal control into water management modeling of the Delaware River Basin to better capture stream temperature and salt front dynamics. Environmental Modelling & Software (in review).

Code contributions by Jacob Zwart and Galen Gorski, who are government employees, to this repository are in the public domain. Refer to the mixed LICENSE file for further details.