Differentiable and scalable LTI routing operator integrated into PyTorch
This repository contains the code for the diffroute model presented in the paper:
Differentiable river routing for end-to-end learning of hydrological processes at diverse scales, ESS Open Archive . May 27, 2025.
diffroute provides a differentiable and GPU-accelerated formulation of classical Linear Time Invariant (LTI) river routing models.
Every component is written in PyTorch and triton, which keeps routing differentiable and easy to integrate in deep learning workflows.
For advanced use-cases, we recommend using diffroute through the companion module diffhydro,
which layers additional data structures and utilities to integrate DiffRoute into more complex differentiable hydrological pipelines, with minimal dependency.
However, diffroute can also be used as a standalone component as showcased in the documentation and examples.
Please note that diffroute is in an early development stage and thus subject to changes.
Nevertheless, the explanations included in the documentation are limited to API and structural components that are unlikely to change in any short term.
- GPU Acceleration: Formulating LTI River Routing operations as 1D Convolution layers allow for efficient GPU execution.
- Differentiable: DiffRoute is integrated to Pytorch, which allows for efficient gradient computations through Automatic Differentiation.
- Scalable: Computations scale to very large river graphs (up to millions of river channels) using staged computations.
- Generality: DiffRoute allows to formulate any LTI River Routing scheme. Currenlty implemented routing schemes include the Muskingum, Linear Diffusive Wave, Linear Storage, Nash Cascade and Pure Lag schemes. Custom schemes can easily be added as shown in the examples.
- Batching: Accepts batched runoff tensors with shape
[batch, catchments, time]and routes them efficiently. This simplifies batched training of parameters and inference of ensemble predictions. - Composability: Integration to Pytorch Automatic Differentiation framework also aims to combine
diffroutewith other learnable components. The companion modulediffhydroaims to simplify the assembly of more complex hydrological pipelines. - Pure Python: DiffRoute is written in pure python. As such it is easily hackable for users interested in experimenting with variations.
You can install diffroute from pip repository:
pip install diffrouteOr install the latest version from GitHub:
pip install git+https://github.com/TristHas/DiffRoute.gitOr clone and install locally:
git clone https://github.com/TristHas/DiffRoute.git
cd DiffRoute; pip install .The package depends on:
torch>=2.0networkxpandastqdm
These are automatically installed when using pip.
import numpy as np
import pandas as pd
import networkx as nx
import torch
from diffroute import RivTree, LTIRouter
b = 2 # Batch (or ensemble size)
n = 20 # Number of channels
T = 100 # Number of time steps
device = "cuda:0" # GPU device to use
G = nx.gn_graph(n) # Toy example tree with a unique outlet.
# Define per-node routing parameters. For a Linear Storage scheme, only one parameter "tau"
params = pd.DataFrame({"tau": np.random.rand(n)}, index=G.nodes)
# River Tree data structure
riv_tree = RivTree(G, params=params, irf_fn="linear_storage").to(device)
# Generate random input runoff
runoff = torch.rand(b, n, T, device=device)
# Instantiate the routing model with desired parameters
router = LTIRouter(max_delay=48, dt=1)
# Compute output discharges from input graph and runoffs
discharge = router(runoff, riv_tree)Early documentation is accessible here.
The motivation for, mathematical derivations behind, and illustrative use-cases of DiffRoute are presented in the paper "Differentiable river routing for end-to-end learning of hydrological processes".
"Differentiable river routing for end-to-end learning of hydrological processes at diverse scales"
ESS Open Archive . May 27, 2025. DOI: 10.22541/essoar.174835108.87664030/v1
If you use diffroute in your academic work, please cite the above reference.
The original code used for the experiment in the paper at the time of submission can be found in a standalone repository
However, We recommend using diffroute through the companion package diffhydro, which layers additional data structures and utilities to ease the use of DiffRoute, with minimal dependency.
The original experiments of the paper have been included as example notebooks within the diffhydro package.
This project is licensed under the terms of the MIT license.
See the LICENSE file for details.