Hydropower Valley JAX

A clean, high-performance JAX implementation of the "Hydropower Valley" control benchmark.

This repository provides a pure Python/JAX port of the original HD-MPC benchmark (Savorgnan & Diehl, 2011), featuring cascaded lakes, river channels modeled by Saint-Venant equations, and hydropower turbines.

MPC Baselines

Method	File	RMSE	Constraints	Solve Time
Multiple Shooting SQP	mpc_sqp_multiple_shooting.py	3.45 MW	Hard	~9.5s/step
Single Shooting SQP	mpc_sqp_single_shooting.py	8.10 MW	Soft	~1.4s/step
MPPI (sampling)	mpc_mppi.py	~12 MW	Penalty	~0.5s/step

All methods use the full 249-state physics model (no reduced-order approximations).

Installation

# Clone and setup environment
python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Quick Start

# Best result (Multiple Shooting SQP, 3.45 MW RMSE)
python -m examples.mpc_sqp_multiple_shooting

# Faster alternative (Single Shooting SQP, 8.10 MW RMSE)
python -m examples.mpc_sqp_single_shooting

# Sampling-based (MPPI)
python -m examples.mpc_mppi

# Run verification tests
python -m pytest tests/ -v

Repository Structure

hpv_jax/
├── physics/           # EXACT reimplementation of PDF equations (readable)
├── model/             # Physical constants and state/control definitions
├── dynamics.py        # Clean interface to the dynamics engine
├── dynamics_generated.py # Fast, numerically verified engine (from auto-gen C)
├── simulation.py      # Stiff ODE integration using Diffrax
└── data.py            # Loading of steady-state benchmark data

examples/
├── mpc_sqp_multiple_shooting.py  # BEST: 3.45 MW RMSE, hard constraints
├── mpc_sqp_single_shooting.py    # Fast: 8.10 MW RMSE, soft constraints
└── mpc_mppi.py                   # Sampling: ~12 MW RMSE

tests/
├── test_physics.py    # Unit tests for physical components
└── test_vs_generated.py # Cross-validation between clean and fast engines

Documentation

• hpv_jax/physics/: Each file contains citations to the specific equations in HPV_model.pdf.

References

• Savorgnan, C., & Diehl, M. (2011). Control benchmark of a Hydro Power Plant.
• HD-MPC Project: Original EU FP7 project on Hierarchical and Distributed Model Predictive Control.
• MATLAB benchmark included in HD-MPC_hydropower_valley_public_benchmark/ folder.