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 (gitignored)
.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/baselines/<key>/	ddrs plan	Cached summed-Q' baseline
.ddrs/runs/<id>/manifest.json	ddrs run	Per-run manifest (config + sources + git SHA + outputs)
.ddrs/runs/<id>/config.yaml	ddrs run	Snapshot of the config that produced this run
.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	train-and-test Phase 2	Predictions for plotting
.ddrs/runs/<id>/baseline/	train-and-test	Copy of the cached summed-Q' baseline (*.f32 + manifest.json)
.ddrs/runs/<id>/checkpoints/epoch_*_mb_*/	ddrs run --workflow train / train-and-test Phase 1	KAN checkpoints (head.mpk, optim.mpk, state.json)
.ddrs/runs/<id>/plot/kan_parameters.nc	ddrs run --plot	Per-COMID learned KAN parameters (full CONUS)
.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 train-and-test
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.

ddrs run --workflow eval is not implemented. It returns "standalone --workflow eval needs a --from-run <run-id> flag", and --from-run does not exist yet. To produce eval output use --workflow train-and-test; to evaluate an existing checkpoint without retraining, use the legacy eval binary (see --help). ddrs init is also a dead stub — it exits 2 and tells you to run ddrs plan.

Add --backend cpu to any ddrs run invocation to route on the CPU. The default is cuda, and a training workflow under --backend cuda exits 5 if the system probe finds no GPU.

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; MHPI cluster paths + managed adjacency
conus-hourly	dHBV2 UH retrospective + AORC precip	Feeds the daily→hourly disaggregation head. Not simply conus plus one key — it uses workstation paths and explicit adjacency zarr
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 · max(0, 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. The max(0, ·) clamp is params.leakance_losing_only, which defaults to true, so gaining reaches produce zeta ≡ 0; params.leakance_impervious_threshold (default 0.7) additionally masks reaches whose corridor_impervious exceeds it. 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.

Daily→hourly disaggregation head

The Muskingum-Cunge solver steps hourly, but the Q' forcing is daily. Adding a kan_head.disaggregation: block enables a mass-preserving KAN sub-head that learns the within-day shape from hourly AORC precip, conserving the daily mean exactly. This is what makes routing's within-day effect visible to the gradient — without it, flat repeat-24 upsampling plus daily-mean aggregation puts that effect in the gradient null space and training loss goes flat.

kan_head:
  disaggregation:
    hidden_size: 16          # architecture fields must match any pretrained
    num_hidden_layers: 2     #   checkpoint or load_record fails loudly
    grid: 20
    k: 3
    boundary_blend: 0.0      # day-boundary shape continuity; used when chunk_days <= 1
    chunk_days: 1            # > 1 relaxes mass balance to the chunk aggregate
    pretrained_checkpoint: output/disagg_pretrain/capacity_chunk1.mpk
    freeze: true             # requires pretrained_checkpoint

The block's presence makes data_sources.aorc_precip mandatory — the head always consumes precip, and a missing source is a hard error at dataset open rather than a silent degradation to flat repeat-24. (There is no use_precip key; it was removed when the head became unconditionally precip-driven.)

Choosing the training objective

experiment.loss.kind selects the objective — l1 (default), nnse-kge, kge, or — once PR #31 lands — nse-batch (dHBV's batch-NSE, which pairs with experiment.optimizer: adadelta). See CLAUDE.md §"Training objective" for why the menu exists and what each term buys.

Advanced

• ddrs show <run_id> — inspect a past run's manifest
• ddrs status — list runs
• ddrs gc — clean up old run directories
• ddrs run flags: --backend {cuda|cpu}, --plot, --strict, --max-mini-batches N, --batch-order-from PATH, --json
• Global flags: --config PATH, --workspace PATH. Pass --workspace whenever you pass --config — the workspace otherwise defaults to .ddrs beside the config file, so --config config/experiments/x.yaml silently creates config/experiments/.ddrs/.
• ddrs <cmd> --help for full flag list