Skip to content

DiffRoute Overview

diffroute provides a differentiable and GPU-accelerated formulation of classical Linear Time Invariant (LTI) river routing models. Every component is written in PyTorch, 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 this documentation.

Please note that diffroute is in an early development stage and thus subject to changes. Nevertheless, the explanations included in this documentation are limited to API and structural components that are unlikely to change in the short to middle term.

Key Features

  • 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 this example.
  • 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 diffroute with other learnable components. The companion module diffhydro aims 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.

Installation

You can install diffroute from pip repository:

pip install diffroute

Or install the latest version from GitHub: 

pip install git+https://github.com/TristHas/DiffRoute.git

Or clone and install locally:

git clone https://github.com/TristHas/DiffRoute.git
cd DiffRoute; pip install .

Example Usage

Below is a toy example for routing 2 realisations of 100 time steps random input runoff time series through a synthetic river graph of 20 channels.

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)

Documentation

We recommend first having a look at the Quickstart page for additional explanations on the above code snippet.

The Concepts section introduces the main data structures and workflows provided by DiffRoute.

Details on repository structure and API are laid out in Code Structure and the API Reference.

The Examples section provides two tutorial notebooks for (i) accelerating existing RAPID simulations on GPU and (ii) integrating custom LTI schemes.

For advanced usage examples and integration of DiffRoute into larger differentiable hydrological pipelines, please refer to the DiffHydro documentation.

Contribution

This project is under early development efforts. Please be aware that some changes may be brought to the API in future versions. We appreciate all contributions, from reporting bugs to implementing new features. Please use the GitHub Issues for any contribution, comment or request.

License and citation

This project is licensed under the terms of the MIT license.

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.

Next

Continue with the Quickstart to see a minimal runnable example and recommended configuration patterns.