API Reference

The sections below are generated from the package docstrings using mkdocstrings. Every time the site is built, the reference stays up to date with the latest type hints and Google-style docstrings.

LTIRouter

Bases: Module

Linear time-invariant river routing module.

Combines an impulse response function aggregator with a block-sparse convolution to transform runoff into downstream discharge.

__init__

__init__(max_delay=100, dt=1, sampling_mode='avg', block_size=16, block_f=128, cascade=1, **kwargs)

Initialize the router with aggregation and convolution settings.

Parameters:

Name	Type	Description	Default
max_delay	int	Maximum impulse response length in time-steps.	100
dt	float	Temporal resolution of the runoff inputs.	1
sampling_mode	str	Strategy for sampling cascade parameters.	'avg'
block_size	int	Spatial block size used for block-sparse kernels.	16
block_f	int	Hidden dimensionality for kernel factorization.	128
cascade	int	Number of cascaded IRFs combined by the aggregator.	1
**kwargs		Unused keyword arguments kept for legacy compatibility.	{}

forward

forward(runoff, g, params=None)

Compute routed discharge for a set of runoff inputs.

Parameters:

Name	Type	Description	Default
runoff	Tensor	Tensor shaped [B, C, T] with batch, channel (node), and time dimensions.	required
g	RivTree	River network containing kernel parameters.	required
params	Tensor | None	Optional per-cluster parameters; defaults to attributes stored on g.	None

Returns:

Type	Description
Tensor	torch.Tensor: Routed discharge with shape [B, C, T].

Raises:

Type	Description
ValueError	If runoff does not have three dimensions.

LTIStagedRouting

Bases: Module

Staged router that orchestrates block-sparse routing over clusters.

Wraps LTIRouter while managing inter-cluster transfers entirely in PyTorch tensors for batched execution.

__init__

__init__(max_delay=32, dt=1.0, sampling_mode='avg', block_size=16, block_f=128, cascade=1)

Configure the staged router and construct its base LTI model.

Parameters:

Name	Type	Description	Default
max_delay	int	Maximum impulse response length in time-steps.	32
dt	float	Temporal resolution of the runoff inputs.	1.0
sampling_mode	str	Strategy for sampling cascade parameters.	'avg'
block_size	int | None	Optional override for kernel block size.	16
block_f	int	Hidden dimensionality for kernel factorization.	128
cascade	int	Number of cascaded IRFs combined by the aggregator.	1

route_one_cluster

route_one_cluster(runoff, gs, cluster_idx, params=None, transfer_bucket=None)

Route a single cluster and update the transfer bucket.

Parameters:

Name	Type	Description	Default
runoff	Tensor	Cluster runoff shaped [B, n_c, T].	required
gs		Clustered river structure providing routing metadata.	required
cluster_idx	int	Identifier of the cluster to route.	required
params	Any | None	Optional parameter bundle for the cluster.	None
transfer_bucket	Tensor | None	Global transfer storage.	None

Returns:

Type	Description
Tensor	Tuple[torch.Tensor, torch.Tensor | None]: Routed discharge for
Tensor | None	the cluster and the updated transfer bucket.

route_all_clusters

route_all_clusters(x, gs, params=None, display_progress=False)

Route all clusters sequentially and assemble the full discharge.

Parameters:

Name	Type	Description	Default
x	Tensor	Full runoff tensor shaped [B, N_nodes, T].	required
gs		Clustered river structure with index metadata.	required
params	List[Any] | None	Optional per-cluster parameter list.	None
display_progress	bool	Whether to wrap the loop with a tqdm bar.	False

Returns:

Type	Description
Tensor	torch.Tensor: Routed discharge tensor shaped [B, N_nodes, T].

route_all_clusters_yield

route_all_clusters_yield(xs, gs, params=None, display_progress=False)

Yield per-cluster discharges lazily for streamed routing.

Parameters:

Name	Type	Description	Default
xs	List[Tensor]	Sequence of cluster runoff tensors.	required
gs		Clustered river structure with transfer metadata.	required
params	List[Any] | None	Optional per-cluster parameter list.	None
display_progress	bool	Whether to wrap iteration in tqdm.	False

Yields:

Type	Description
	torch.Tensor: Discharge tensor for each cluster in order.

init_upstream_discharges

init_upstream_discharges(xs, gs, cluster_idx, params=None, display_progress=False)

Warm up the staged router until the target cluster.

Parameters:

Name	Type	Description	Default
xs	List[Tensor]	Sequence of cluster runoff tensors.	required
gs		Clustered river structure with transfer metadata.	required
cluster_idx	int	Cluster index to stop before routing.	required
params	List[Any] | None	Optional per-cluster parameter list.	None
display_progress	bool	Whether to wrap iteration in tqdm.	False

Returns:

Type	Description
Tensor	torch.Tensor: Transfer bucket capturing upstream discharges.

Bases: Module

Collection of river subgraphs with optional inter-cluster transfers.

__init__

__init__(clusters_g, node_transfer, irf_fn=None, include_index_diag=True, param_df=None, nodes_idx=None)

Assemble clustered river networks and transfer bookkeeping.

Parameters:

Name	Type	Description	Default
clusters_g	Sequence[DiGraph]	Clustered river graphs.	required
node_transfer	Dict[int, List[Tuple[int, int, int]]] | None	Mapping describing inter-cluster transfers.	required
irf_fn	str | None	Name of the IRF parameterization to use.	None
include_index_diag	bool	Whether kernels include self-loops.	True
param_df	DataFrame | None	Optional parameter table.	None
nodes_idx	Sequence[Series] | None	Custom node orderings.	None

Graph Utilities