ddrs

Differentiable distributed routing. A BURN-based Rust port of the Muskingum-Cunge routing solver from DDR (Python/PyTorch), gradient-exact against the reference at single precision.

Getting started

Install

cargo install --path .

This puts the ddrs binary in ~/.cargo/bin/. If that directory isn't on your PATH:

echo 'export PATH="$HOME/.cargo/bin:$PATH"' >> ~/.bashrc
source ~/.bashrc

First-time setup

From your project root:

ddrs plan      # probes GPU + smoke test (first run), opens $EDITOR on
               # ddrs.yaml if missing, locks data sources, validates,
               # builds adjacency/baseline caches, prints the plan
ddrs run       # executes the workflow, writes manifest + outputs

The first ddrs plan runs a 5-reach RAPID sandbox parity check on CUDA when available and falls back to CPU otherwise — so the install path works on laptops and CI. The verdict is cached; later plans are fast. When no ddrs.yaml exists, plan asks whether to start from your last successful run's config or the clean bundled template (config/merit_training.yaml). To start fresh at any time: rm ddrs.yaml && ddrs plan.

The adjacency stores are managed, so data_sources only needs the raw inputs:

data_sources:
  geospatial_fabric: /path/to/riv_pfaf_7_..._bugfix1.shp  # MERIT flowlines (.shp/.dbf/.gpkg)
  attributes:        /path/to/merit_global_attributes_v2.nc
  streamflow:        /path/to/merit_dhbv2_UH_retrospective.ic
  observations:      /path/to/usgs_daily_observations
  gages:             /path/to/gages_3000.csv

On the first ddrs plan, the CONUS and per-gauge adjacency zarr stores are built from the fabric's attribute table into .ddrs/adjacency/<key>/ (~10 s for the CONUS dbf, content-addressed and reused afterwards). The fabric may also be a GeoPackage — e.g. a merged global MERIT flowlines .gpkg — in which case attributes are read via SQL and geometry is never touched; if the gpkg holds more than one feature layer, pick one with geospatial_fabric_layer: <name>. If you already have pre-built zarr stores, drop geospatial_fabric and set both conus_adjacency and gages_adjacency to their paths instead.

What lives where

Path	Written by	Purpose
ddrs.yaml	ddrs plan (via $EDITOR)	Workflow + experiment config
.ddrs/system.json	ddrs plan	GPU/driver/smoke-test record
.ddrs/sources.lock	ddrs plan	Fingerprints of data_sources paths
.ddrs/adjacency/<key>/	ddrs plan (managed build)	Cached CONUS + gauges adjacency zarr stores
.ddrs/runs/<id>/manifest.json	ddrs run	Per-run manifest (config + sources + git SHA + outputs)
.ddrs/runs/<id>/run.log	ddrs run	Timestamped tee of everything the run printed (stdout + stderr, incl. CUDA messages)
.ddrs/runs/<id>/eval/predictions.zarr	ddrs run --workflow eval / train-and-test Phase 2	Predictions for plotting
.ddrs/runs/<id>/checkpoints/epoch_*_mb_*/	ddrs run --workflow train / train-and-test Phase 1	KAN checkpoints
.ddrs/runs/<id>/kan_parameters.nc	train-and-test Phase 2 (leakance runs only)	Per-reach eval-window zeta diagnostic (see Leakance below)

Run ids are <UTC timestamp>-[<group>-]<workflow> — e.g. 2026-06-12T14-02-10Z-global-train-and-test. The <group> segment appears when the config's data_sources matches a saved group (see ddrs sources below), so run dirs say which dataset they were trained on.

Override workflow on the command line

The workflow: key in ddrs.yaml is what plan/run use by default. To override for a single invocation:

ddrs plan --workflow eval
ddrs run --workflow train

mode: and workflow: must agree (mode: training ↔ workflow ∈ {train, train-and-test}; mode: testing ↔ workflow: eval). ddrs plan will reject contradictions at load time.

The top-level device: key in ddrs.yaml selects the CUDA device ordinal (default 0, mirrors DDR's device: key) — on multi-GPU hosts set e.g. device: 1 to keep training off the display/shared GPU.

Data-source groups

Named "save files" for the data_sources: block, stored under config/sources/<name>.yaml (tracked in git). Groups that ship in-repo:

Group	Streamflow source	Notes
conus	dHBV2 UH retrospective (daily)	Default CONUS source
conus-hourly	same + AORC precip	Enables daily→hourly disaggregation head
global	Global zarr-v2 Q' stores	2.94M reaches
daily-lstm	CudaLSTM unit-catchment forwards (daily)	NH LSTM output
hourly-lstm	MTS-LSTM unit-catchment forwards (hourly-native)	NH LSTM output; no disagg

Switching datasets never requires hand-editing ddrs.yaml:

ddrs sources list                # '*' marks the group matching ddrs.yaml
ddrs sources save <name>         # snapshot current data_sources (--force to overwrite)
ddrs sources use  <name>         # splice group into ddrs.yaml + refresh sources.lock

Starting a global train from a CONUS workspace is just ddrs sources use global && ddrs plan && ddrs run.

Importing a Q' store

Any icechunk store that follows the DDR Q' contract (Qr(divide_id, time) f32 m³/s, CF days since/hours since time axis) can be registered as a source group in one command:

ddrs import <store-path> --dry-run          # validate + coverage report only
ddrs import <store-path> --name <group>     # validate + register config/sources/<group>.yaml

The reader auto-detects resolution from the CF time units — hourly stores are sliced natively (no repeat-24 interpolation, no disaggregation head — combining a disaggregation config with an hourly-native source is rejected at dataset open). The daily-lstm and hourly-lstm groups were registered this way; they can serve as reference examples. After import:

ddrs sources use <group> && ddrs plan && ddrs run --workflow train

See docs/nh-qprime-store-contract.md for the full producer/consumer contract.

Leakance (experimental, NOT promotable)

An optional groundwater–surface-water loss term zeta = leakance_factor · area_z · K_D · (depth − d_gw) subtracted from the routing right-hand side each timestep (positive zeta = losing stream), with the three parameters learned per reach by the KAN head. Off by default; enabling it takes params.use_leakance: true, the three extra learnable_parameters, and their parameter_ranges — see the "Leakance" section in CLAUDE.md for the exact keys, and config/experiments/leakance_hourly_on.yaml for a complete working config.

Research status (2026-07-06): NOT PROMOTABLE. A full identifiability campaign established that zeta cannot be constrained from gauged discharge alone — the observation operator (gauge = Σ upstream zeta) is not invertible for the per-reach distribution. The term is code-complete and gradient-exact; it is left in place as the anchor for the paper's selective-equifinality axis. Do not attempt to promote it without reading docs/2026-07-06-leakance-nogo-scientific-summary.md.

Advanced

• ddrs show <run_id> — inspect a past run's manifest
• ddrs status — list runs
• ddrs gc — clean up old run directories
• ddrs <cmd> --help for full flag list