From 289163bd0acfb6d1818e71bbd6fb691034e8c1a5 Mon Sep 17 00:00:00 2001 From: Tadd Bindas Date: Wed, 16 Sep 2026 22:23:20 -0400 Subject: [PATCH 1/3] feat(leakance): temperature-annealed gate replaces the continuous factor zeta multiplies leakance_factor by K_D, so only their product reaches the physics: an exact scaling degeneracy. Measured over 346,321 CONUS reaches on a trained arm, rho(leakance_factor, K_D) = +0.9986 and all four learned outputs of that head hold 97 percent of their variance in one direction. Physically a reach either sits above a losing aquifer or it does not, so the factor should select where leakance acts rather than scale how much. A gate and a conductance are not degenerate with each other. The gate is a sharpening sigmoid on the head's normalized output, g = sigmoid(logit(clamp(u, eps, 1-eps)) / tau), annealed by an epoch-keyed schedule that mirrors the learning rate. tau = 1 is special-cased to return u untouched, so every config without params.leakance_gate is bit-exact unchanged. Differentiable at every positive tau, so autograd gives the exact gradient and no straight-through estimator is introduced. Applied to the head output before setup_inputs, so no routing operator or backward changes. Replicated across all three hand-written head readers rather than unified, per the WET preference, with the eval reader and dump_parameters using the FINAL temperature so the scored model is the one training ended on. A previous bug shipped because an eval reader silently dropped a learned head output. Gates: lib 377 passed; leakance_gate 11; gamma_eval_parity 4; sp8_gradcheck 20; leakance_gradcheck 16; leakance_off_parity 3; zeta_accum 11; leakance_reference_match 5; ddr_sandbox_match 1. compare_ddr_sandbox ABSOLUTE MATCH at 1.525879e-4 m3/s, unchanged. verify_doc_paths 0 unresolved. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_0143bx7gVz8ug2oc6DTy5jtz --- .claude/skills/ddrs-dev/references/config.md | 59 +++ .claude/skills/ddrs-dev/references/testing.md | 7 +- src/bin/probe_zeta_gradient.rs | 6 + src/config.rs | 161 +++++++ src/dump_parameters.rs | 26 +- src/training/driver.rs | 13 +- src/training/eval.rs | 10 + src/training/forward.rs | 36 +- src/training/gate.rs | 64 +++ src/training/mod.rs | 2 + src/training/probe.rs | 8 + tests/gamma_eval_parity.rs | 27 +- tests/leakance_gate.rs | 398 ++++++++++++++++++ tests/leakance_off_parity.rs | 3 + 14 files changed, 807 insertions(+), 13 deletions(-) create mode 100644 src/training/gate.rs create mode 100644 tests/leakance_gate.rs diff --git a/.claude/skills/ddrs-dev/references/config.md b/.claude/skills/ddrs-dev/references/config.md index 70ee42ae..9f1415d0 100644 --- a/.claude/skills/ddrs-dev/references/config.md +++ b/.claude/skills/ddrs-dev/references/config.md @@ -229,6 +229,7 @@ and bounds every topology from above. | `use_leakance` | false | | | `leakance_losing_only` | **true** | Clamps `head = max(0, depth − d_gw)`, so gaining reaches produce `zeta ≡ 0` | | `leakance_impervious_threshold` | 0.7 | Masks reaches whose `corridor_impervious` is **`>`** this value (not `≥`) | +| `leakance_gate` | absent | Temperature-annealed 0/1 gate on the head's `leakance_factor` output; see §Leakance gate. Absent ⇒ byte-identical to the ungated readers. Requires `use_leakance: true` | | `tau` | 9 | **Not** a routing sub-step count. Since 2026-08-08: hours the routed output is ADVANCED before daily scoring (translation-only inverse routing, dMC-Juniata's sign). Slice `[tau : -(24-tau)]`, pooled day i ↔ obs day i, valid range [0, 24) (`src/training/loss.rs`). Nothing in `src/routing/` reads it. Default 9 = the measured CONUS optimum (findings §5g). **Legacy scale (pre-2026-08-08): old = new + 11**, slice `[13+tau : -11+tau]`, day i ↔ obs day i+1; DDR-Python still uses it, all older configs/checkpoints carry it (old shipped 3 ≡ new −8, wrong direction). Never copy a `tau:` value across the convention boundary. | | `log_space_parameters` | `["p_spatial"]` | | | `defaults` | `{p_spatial: 21.0}` | Value used when a parameter is not in `learnable_parameters` | @@ -265,6 +266,9 @@ when subdivision is enabled), plus one at dataset open. | | nested `validate_subdivision_reaches_the_builder`: `enabled: true` + explicit `conus_adjacency`/`gages_adjacency` pointing at a non-subdivided store | `"params.subdivision"` + `"conflicts with the explicit"` | | `validate_geodataset` | `geodataset:` contradicting the adjacency source (`ddm30` with `geospatial_fabric`, `merit` with `gridded_network`) | `"geodataset"` + the source key. Absent ⇒ inferred; explicit adjacency paths ⇒ any label allowed | | `validate_leakance` | `use_leakance` + `use_cuda_graphs` | both key names | +| `validate_leakance_gate` | `leakance_gate` block without `use_leakance: true` | `"leakance_gate"` + `"use_leakance"` | +| | empty `leakance_gate.temperature` | `"temperature"` | +| | a temperature that is not positive and finite | `"epoch N"` + `"positive"` | | `validate_ddr_match` | `use_cuda_graphs: true` without the deprecated `ddr_match: true` | `"use_cuda_graphs: true` requires the DEPRECATED `ddr_match: true"` | | `validate_enforce_positivity` | `enforce_positivity: true` + `ddr_match: true` | `"requires \`ddr_match: false\`"` | | `validate_disagg_pretrained` | `freeze: true` without `pretrained_checkpoint` | `"freeze: true requires pretrained_checkpoint"` | @@ -316,3 +320,58 @@ target/release/eval --config config/experiments/leakance_hourly_on.yaml \ ``` Takes about 10 minutes, no retrain. + +## Leakance gate (`params.leakance_gate`, added 2026-09-16) + +`zeta = leakance_factor · area_z · K_D · (depth − d_gw)` multiplies +`leakance_factor` and `K_D`, so only their product reaches the physics: an +exact scaling degeneracy (measured `rho(leakance_factor, K_D) = +0.9986` over +346,321 CONUS reaches on a trained arm, with 97 % of the four outputs' variance +in one direction). Physically a reach either sits above a losing aquifer or it +does not, so `leakance_factor` is meant to be a SELECTOR of where leakance +acts, not a continuous multiplier. The gate makes it one without a +straight-through estimator: the head's normalized output `u ∈ (0, 1)` becomes + +```text +g = sigmoid( logit(clamp(u, 1e-6, 1 − 1e-6)) / tau ), logit(x) = ln(x / (1 − x)) +``` + +`tau = 1` is the identity (`g = u`; the code returns `u` untouched, so it is +BIT-exact, not round-tripped through `ln`/`exp`). Smaller `tau` sharpens toward +a step at `u = 0.5`; `0.5` is a fixed point at every temperature. The clamp +margin `1e-6` sits ~17 f32 ulps below 1.0 so `1 − u` keeps ~3 % relative +accuracy at the edge, and it only alters values the head has already saturated +past `|pre-activation| > 13.8`. Autograd through `clamp → log → div → sigmoid` +gives the exact gradient; `TimestepLeakanceOp` and `src/routing/` are untouched +because the gate is applied to the head OUTPUT before `setup_inputs` +denormalizes it (`src/training/gate.rs::leakance_gate`). + +```yaml +params: + use_leakance: true + leakance_gate: + temperature: {1: 1.0, 6: 0.5, 11: 0.2, 16: 0.05} # keyed by 1-indexed epoch +``` + +| Key | Default | Notes | +|---|---|---| +| `temperature` | required | Epoch-keyed schedule resolved like `experiment.learning_rate` (`LeakanceGate::resolve`: largest key `<= epoch`, first value before it). Every value must be positive and finite; `1.0` = identity, `> 1` softens | + +Which temperature each reader applies: + +| Reader | Temperature | Log line | +|---|---|---| +| `src/training/forward.rs::forward` (training) | `resolve(epoch)`, threaded by the driver as `gate_tau` | `epoch E lr=… leakance_gate_tau=…` | +| `src/training/forward.rs::forward_eval_core` (eval) | `final_temperature()` | `leakance gate: final temperature tau=…` | +| `src/training/probe.rs::probe_forward` | `final_temperature()` (after lifting, so leaf grads carry the gate's Jacobian) | same, from `probe_zeta_gradient` | +| `src/dump_parameters.rs::dump` / `dump_init` | `final_temperature()` (the dumped factor is the gated value) | same | + +`forward` asserts `gate_tau.is_some() == params.leakance_gate.is_some()`, so a +caller cannot train an ungated model against a gated config or vice versa. In +`forward_eval_core` the gate is applied BEFORE `LeakanceOverride`, so an +override still replaces the value that reaches denormalization. Tests: +`tests/leakance_gate.rs` (tau = 1 bit-exact identity through all three readers, +central-difference gradcheck at five temperatures, monotonicity, limits, +saturation without NaN/Inf) and the gated rows of `tests/gamma_eval_parity.rs`. +This does not re-open the leakance verdict in `research-status.md`; it removes +one degeneracy so a future arm can be judged on the gate, not on the product. diff --git a/.claude/skills/ddrs-dev/references/testing.md b/.claude/skills/ddrs-dev/references/testing.md index 0f7f5928..f4718ed8 100644 --- a/.claude/skills/ddrs-dev/references/testing.md +++ b/.claude/skills/ddrs-dev/references/testing.md @@ -74,7 +74,10 @@ cargo test --lib && cargo test && \ mkdir -p output && cargo run --release --example compare_ddr_sandbox ``` If you touched `src/training/forward.rs` (disagg / leakance threading), also run -`cargo test --test leakance_off_parity`. +`cargo test --test leakance_off_parity`. If you touched any of the three head +readers (`forward`, `forward_eval_core`, `probe_forward`) or +`src/training/gate.rs`, also run +`cargo test --features fixtures --test gamma_eval_parity --test leakance_gate`. If you touched `src/routing/leakance.rs`, `src/geometry.rs`'s depth inversion, or the S6/S25 leakance call site in `src/routing/mmc_op.rs::forward_chain_inner`, run @@ -148,7 +151,7 @@ line citation outright rather than trusting it to stay pinned. | Gridded (DDM30) ingestion | `gridded_bundle` (sub-reach store → subdivided layout, gauge at the cell's last piece, cache hit, dataset opens), `cargo test --lib adjacency::gridded` (synthetic-store validation), `hourly_streamflow::time_major_store_reads_identically_to_divide_major`, `cargo test --lib zarr::tests::upstream_comids_names_each_subdivided_parent_once` | | Sparse / autograd | `sparse_gradcheck`, `sp8_gradcheck` | | KAN head | the 4 `kan_head_*` fixture tests (need `--features fixtures`) | -| Leakance | `leakance_reference_match` (cross-implementation, vs DDR `_compute_zeta` @ `c2bd0f9`), `leakance_gradcheck`, `leakance_off_parity`, `zeta_accum` (incl. multi-timestep volume accounting) | +| Leakance | `leakance_reference_match` (cross-implementation, vs DDR `_compute_zeta` @ `c2bd0f9`), `leakance_gradcheck`, `leakance_off_parity`, `zeta_accum` (incl. multi-timestep volume accounting), `leakance_gate` (tau = 1 bit-exact identity through all three readers, gate gradcheck, saturation) | | Subdivision | `subdivide`, `subdivision_integration`, `gauge_mass_conservation`; `compare_ddr_sandbox` must still report ABSOLUTE MATCH | | Adjacency | `adjacency_parity` (managed builder byte-identical to the petgraph engine on `order`/`indices_0`/`indices_1`), `adjacency_build`, `data_zarr_store::conus_adjacency_loads_real_merit_zarr` (invariant 3 on real CONUS data) | | CLI / data | `data_dataset`, `data_static`, `cli_manifest`, `cli_lockfile`, `cli_json_contract` | diff --git a/src/bin/probe_zeta_gradient.rs b/src/bin/probe_zeta_gradient.rs index 0a1014fe..f538cb23 100644 --- a/src/bin/probe_zeta_gradient.rs +++ b/src/bin/probe_zeta_gradient.rs @@ -472,6 +472,12 @@ fn run(cfg: Config, cli: Cli, device: I::Device) -> Result<(), Box, +} + +impl LeakanceGate { + /// Temperature in force at `epoch` (1-indexed). Mirrors + /// `src/training/optimizer.rs::resolve_lr`: largest key `<= epoch`, else + /// the first value. Panics on an empty schedule, which + /// `validate_leakance_gate` rejects at load. + pub fn resolve(&self, epoch: usize) -> f32 { + self.temperature + .range(..=epoch) + .next_back() + .map(|(_, &t)| t) + .unwrap_or_else(|| { + *self.temperature.values().next().expect("leakance_gate.temperature is empty") + }) + } + + /// The last scheduled temperature: what eval, the probe and + /// `dump_parameters` apply. + pub fn final_temperature(&self) -> f32 { + *self.temperature.values().next_back().expect("leakance_gate.temperature is empty") + } +} + /// YAML `kan_head.disaggregation:` block (presence enables the head, unless /// `enabled: false`). The head always consumes `(daily Q', that day's 24h /// precip)` — requires `data_sources.aorc_precip` to be set. See @@ -755,6 +803,10 @@ pub struct Params { /// Stage-dependent Manning roughness, `n(d) = n_0·(d/d_ref)^(−gamma)`. /// Absent ⇒ `gamma = 0`, byte-identical to the historical solver. pub stage_roughness: Option, + /// Temperature-annealed 0/1 gate on the head's `leakance_factor` output. + /// Absent ⇒ the output is used as-is, byte-identical to the historical + /// readers. See `LeakanceGate`. + pub leakance_gate: Option, } impl Params { @@ -795,6 +847,7 @@ impl Default for Params { enforce_positivity: false, subdivision: Subdivision::default(), stage_roughness: None, + leakance_gate: None, } } } @@ -823,6 +876,8 @@ struct ParamsRaw { subdivision: Subdivision, #[serde(default)] stage_roughness: Option, + #[serde(default)] + leakance_gate: Option, } impl From for Params { @@ -920,6 +975,7 @@ impl From for Params { // exactly what `Params::default()` carries. p.subdivision = r.subdivision; p.stage_roughness = r.stage_roughness; + p.leakance_gate = r.leakance_gate; p } } @@ -1076,6 +1132,10 @@ impl Config { path: path.to_path_buf(), source: serde_yaml::Error::custom(msg), })?; + validate_leakance_gate(&cfg).map_err(|msg| DataError::Yaml { + path: path.to_path_buf(), + source: serde_yaml::Error::custom(msg), + })?; validate_stage_roughness(&cfg).map_err(|msg| DataError::Yaml { path: path.to_path_buf(), source: serde_yaml::Error::custom(msg), @@ -1251,6 +1311,36 @@ fn validate_leakance(cfg: &Config) -> std::result::Result<(), String> { Ok(()) } +/// `params.leakance_gate` is only meaningful on a leakance run, and its +/// schedule must hold usable temperatures. +fn validate_leakance_gate(cfg: &Config) -> std::result::Result<(), String> { + let Some(gate) = cfg.params.leakance_gate.as_ref() else { + return Ok(()); + }; + if !cfg.params.use_leakance { + return Err( + "params: `leakance_gate` requires `use_leakance: true` — the gate transforms \ + the head's `leakance_factor` output, which nothing reads while leakance is \ + off, so the block would be silently inert." + .to_string(), + ); + } + if gate.temperature.is_empty() { + return Err( + "params.leakance_gate: `temperature` must map at least one epoch to a \ + temperature (e.g. `{1: 1.0, 10: 0.2}`)." + .to_string(), + ); + } + if let Some((epoch, t)) = gate.temperature.iter().find(|(_, t)| !(t.is_finite() && **t > 0.0)) { + return Err(format!( + "params.leakance_gate: temperature at epoch {epoch} is {t}; every temperature \ + must be positive and finite (1.0 is the exact identity, smaller sharpens)." + )); + } + Ok(()) +} + /// A head that does not emit `q_spatial` holds it at `params.defaults.q_spatial` /// (the n_0 + gamma design). The default must exist and sit inside the box, /// because the readers normalize it into the box and the engine denormalizes @@ -2456,6 +2546,77 @@ params: ); } + fn write_yaml(name: &str, yaml: &str) -> std::path::PathBuf { + let path = std::env::temp_dir().join(name); + std::fs::write(&path, yaml).unwrap(); + path + } + + const GATE_HEAD: &str = r#" +mode: training +geodataset: merit +seed: 1 +np_seed: 1 +"#; + + #[test] + fn leakance_gate_absent_is_none() { + let path = write_yaml("ddrs_leakance_gate_absent.yaml", &format!("{GATE_HEAD}params:\n use_leakance: true\n")); + let cfg = Config::from_yaml_file(&path).unwrap(); + assert!(cfg.params.leakance_gate.is_none()); + assert!(Params::default().leakance_gate.is_none()); + } + + #[test] + fn leakance_gate_rejected_without_use_leakance() { + let path = write_yaml( + "ddrs_leakance_gate_no_leakance.yaml", + &format!("{GATE_HEAD}params:\n leakance_gate:\n temperature: {{1: 1.0, 5: 0.2}}\n"), + ); + let err = Config::from_yaml_file(&path).unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("leakance_gate") && msg.contains("use_leakance"), + "error should name both keys: {msg}" + ); + } + + #[test] + fn leakance_gate_rejects_empty_and_nonpositive_temperatures() { + let path = write_yaml( + "ddrs_leakance_gate_empty.yaml", + &format!("{GATE_HEAD}params:\n use_leakance: true\n leakance_gate:\n temperature: {{}}\n"), + ); + let msg = Config::from_yaml_file(&path).unwrap_err().to_string(); + assert!(msg.contains("temperature"), "{msg}"); + + let path = write_yaml( + "ddrs_leakance_gate_zero.yaml", + &format!("{GATE_HEAD}params:\n use_leakance: true\n leakance_gate:\n temperature: {{1: 1.0, 5: 0.0}}\n"), + ); + let msg = Config::from_yaml_file(&path).unwrap_err().to_string(); + assert!(msg.contains("epoch 5") && msg.contains("positive"), "{msg}"); + } + + #[test] + fn leakance_gate_schedule_resolves_at_epoch_boundaries() { + let path = write_yaml( + "ddrs_leakance_gate_schedule.yaml", + &format!("{GATE_HEAD}params:\n use_leakance: true\n leakance_gate:\n temperature: {{1: 1.0, 3: 0.5, 6: 0.1}}\n"), + ); + let cfg = Config::from_yaml_file(&path).unwrap(); + let gate = cfg.params.leakance_gate.as_ref().expect("block parsed"); + // Same rule as `resolve_lr`: largest key <= epoch, first value before it. + assert_eq!(gate.resolve(0), 1.0); + assert_eq!(gate.resolve(1), 1.0); + assert_eq!(gate.resolve(2), 1.0); + assert_eq!(gate.resolve(3), 0.5); + assert_eq!(gate.resolve(5), 0.5); + assert_eq!(gate.resolve(6), 0.1); + assert_eq!(gate.resolve(100), 0.1); + assert_eq!(gate.final_temperature(), 0.1); + } + #[test] fn leakance_impervious_threshold_parses() { let yaml = r#" diff --git a/src/dump_parameters.rs b/src/dump_parameters.rs index dfeb1836..ddb8a112 100644 --- a/src/dump_parameters.rs +++ b/src/dump_parameters.rs @@ -255,6 +255,12 @@ where // read from the run's config snapshot by the plotting scripts. let dump_gamma = learn_has("gamma"); let mut gamma_phys: Vec = Vec::new(); + // Leakance gate at the FINAL scheduled temperature (mirrors eval): the + // dumped factor is the value routing multiplies, not the raw head output. + let gate_tau = cfg.params.leakance_gate.as_ref().map(|g| g.final_temperature()); + if let Some(tau) = gate_tau { + eprintln!("leakance gate: final temperature tau={tau} (dump applies the last scheduled value)"); + } for start in (0..n_reaches).step_by(batch_size) { let end = (start + batch_size).min(n_reaches); @@ -319,8 +325,13 @@ where let kd_d = denormalize(raw["K_D"].clone(), cfg.params.parameter_ranges.k_d, is_log("K_D")); let dgw_d = denormalize(raw["d_gw"].clone(), cfg.params.parameter_ranges.d_gw, is_log("d_gw")); + let lfac_u = raw["leakance_factor"].clone(); + let lfac_u = match gate_tau { + Some(tau) => crate::training::gate::leakance_gate(lfac_u, tau), + None => lfac_u, + }; let lfac_d = denormalize( - raw["leakance_factor"].clone(), + lfac_u, cfg.params.parameter_ranges.leakance_factor, is_log("leakance_factor"), ); @@ -527,6 +538,12 @@ where // read from the run's config snapshot by the plotting scripts. let dump_gamma = learn_has("gamma"); let mut gamma_phys: Vec = Vec::new(); + // Leakance gate at the FINAL scheduled temperature (mirrors eval): the + // dumped factor is the value routing multiplies, not the raw head output. + let gate_tau = cfg.params.leakance_gate.as_ref().map(|g| g.final_temperature()); + if let Some(tau) = gate_tau { + eprintln!("leakance gate: final temperature tau={tau} (dump applies the last scheduled value)"); + } for start in (0..n_reaches).step_by(batch_size) { let end = (start + batch_size).min(n_reaches); @@ -591,8 +608,13 @@ where let kd_d = denormalize(raw["K_D"].clone(), cfg.params.parameter_ranges.k_d, is_log("K_D")); let dgw_d = denormalize(raw["d_gw"].clone(), cfg.params.parameter_ranges.d_gw, is_log("d_gw")); + let lfac_u = raw["leakance_factor"].clone(); + let lfac_u = match gate_tau { + Some(tau) => crate::training::gate::leakance_gate(lfac_u, tau), + None => lfac_u, + }; let lfac_d = denormalize( - raw["leakance_factor"].clone(), + lfac_u, cfg.params.parameter_ranges.leakance_factor, is_log("leakance_factor"), ); diff --git a/src/training/driver.rs b/src/training/driver.rs index 41676b4c..21c10bf8 100644 --- a/src/training/driver.rs +++ b/src/training/driver.rs @@ -106,6 +106,7 @@ fn run_micro_batch( staids: &[Staid], window: &RhoWindow, mb_label: usize, + gate_tau: Option, ) -> Result>> { let exp = cfg.experiment.as_ref().expect("experiment"); let batch = dataset.collate(staids, window)?; @@ -125,7 +126,7 @@ fn run_micro_batch( let tensors = batch.to_tensors::>(device); let (median_n, n_at_floor) = crate::training::forward::manning_n_stats::(cfg, &tensors, head); - let pred_hourly = forward::(cfg, &tensors, head, device, false); + let pred_hourly = forward::(cfg, &tensors, head, device, false, gate_tau); let daily = tau_trim_and_downsample(pred_hourly, cfg.params.tau); let dims = daily.dims(); let (g, t_days) = (dims[0], dims[1]); @@ -295,7 +296,13 @@ pub fn train( None => sampler.reshuffle(&mut state.rng), } let lr = resolve_lr(&exp.learning_rate, epoch); - eprintln!("epoch {epoch} lr={lr}"); + // Leakance-gate temperature for this epoch, resolved like `lr` and + // stated in the run log so a reader can tell which model trained. + let gate_tau = cfg.params.leakance_gate.as_ref().map(|g| g.resolve(epoch)); + match gate_tau { + Some(tau) => eprintln!("epoch {epoch} lr={lr} leakance_gate_tau={tau}"), + None => eprintln!("epoch {epoch} lr={lr}"), + } let mut mb_done = 0usize; @@ -313,6 +320,7 @@ pub fn train( &staids, &window, state.mini_batch, + gate_tau, )?; let Some(MicroBatchOutcome { loss, @@ -388,6 +396,7 @@ pub fn train( &staids, &window, state.mini_batch, + gate_tau, )?; if let Some(MicroBatchOutcome { loss, diff --git a/src/training/eval.rs b/src/training/eval.rs index f63dd25d..0df70797 100644 --- a/src/training/eval.rs +++ b/src/training/eval.rs @@ -110,6 +110,16 @@ pub fn evaluate( ensure_panic_hook_installed(); take_worker_panicked(); // clear any stale flag from an unrelated prior call + // The gate is scored at its FINAL scheduled temperature + // (`forward_eval_core`); say so in the run log so a reader can tell which + // model these metrics belong to. + if let Some(gate) = cfg.params.leakance_gate.as_ref() { + eprintln!( + "leakance gate: final temperature tau={} (eval applies the last scheduled value)", + gate.final_temperature() + ); + } + let axis = dataset.time_axis().clone(); let n_days_total = axis.num_days; assert!(batch_size_days > 0, "batch_size_days must be positive"); diff --git a/src/training/forward.rs b/src/training/forward.rs index 28cecb8d..eeac17e1 100644 --- a/src/training/forward.rs +++ b/src/training/forward.rs @@ -10,6 +10,7 @@ use crate::config::Config; use crate::data::dataset::RoutingTensors; use crate::routing::mmc::{MuskingumCunge, RoutingInputs, SpatialParameters}; use crate::routing::utils::denormalize; +use crate::training::gate::leakance_gate; /// Gather + grouped sum: `output[g, t] = sum_{k : group_ids[k] == g} runoff[flat_indices[k], t]`. /// @@ -322,13 +323,26 @@ pub fn manning_n_stats( /// with autograd alive on the engine path. /// /// Mirrors `~/projects/ddr/scripts/train.py:67-73` (MLP forward + dmc forward). +/// +/// `gate_tau` — the leakance-gate temperature in force for this forward +/// (`LeakanceGate::resolve(epoch)`, resolved by the driver like the learning +/// rate). Must be `Some` exactly when `params.leakance_gate` is configured and +/// `None` otherwise; the mismatch is asserted so a caller cannot silently +/// train an ungated model against a gated config or vice versa. pub fn forward( cfg: &Config, tensors: &RoutingTensors>, head: &KanHead>, device: &I::Device, carry_state: bool, + gate_tau: Option, ) -> Tensor, 2> { + assert_eq!( + gate_tau.is_some(), + cfg.params.leakance_gate.is_some(), + "forward: `gate_tau` must be Some exactly when `params.leakance_gate` is \ + configured (resolve it with `LeakanceGate::resolve(epoch)`)" + ); let n_active = tensors.adjacency.n; // The head runs at PARENT resolution; expand to the routing's sub-reach // rows before anything denormalizes or slices. No-op when not subdivided. @@ -384,11 +398,15 @@ pub fn forward( ); } } - ( - params_map.get("K_D").cloned(), - params_map.get("d_gw").cloned(), - params_map.get("leakance_factor").cloned(), - ) + // Temperature-annealed 0/1 gate on the NORMALIZED factor, applied to + // the head output before `setup_inputs` denormalizes it. Absent block + // ⇒ `gate_tau` is `None` and the output is untouched (byte-identical). + let factor = params_map.get("leakance_factor").cloned(); + let factor = match gate_tau { + Some(tau) => factor.map(|u| leakance_gate(u, tau)), + None => factor, + }; + (params_map.get("K_D").cloned(), params_map.get("d_gw").cloned(), factor) } else { (None, None, None) }; @@ -717,6 +735,14 @@ fn forward_eval_core( params_map.get("d_gw").expect("checked above").clone(), params_map.get("leakance_factor").expect("checked above").clone(), ); + // Gate at the FINAL scheduled temperature (mirrors `forward`, which + // takes the epoch's value): eval scores the model training ended on. + // Applied to the head output, so an override below still replaces + // the value that reaches denormalization. Pinned by + // `tests/gamma_eval_parity.rs` and `tests/leakance_gate.rs`. + if let Some(gate) = cfg.params.leakance_gate.as_ref() { + factor_t = leakance_gate(factor_t, gate.final_temperature()); + } if let Some(ov) = overrides { assert_eq!( ov.mask.len(), diff --git a/src/training/gate.rs b/src/training/gate.rs new file mode 100644 index 00000000..7fbd9fa2 --- /dev/null +++ b/src/training/gate.rs @@ -0,0 +1,64 @@ +//! Temperature-annealed binary gate for the leakance factor. +//! +//! `zeta = leakance_factor · area_z · K_D · (depth − d_gw)` multiplies +//! `leakance_factor` and `K_D`, so only their product reaches the physics and +//! the two are exactly degenerate (measured `rho = +0.9986` over CONUS on a +//! trained arm). Physically a reach either sits above a losing aquifer or it +//! does not, so `leakance_factor` should be a 0/1 SELECTOR of where leakance +//! acts, not a continuous multiplier. A gate and a conductance are not +//! degenerate with each other. +//! +//! A hard step has zero gradient almost everywhere, so the gate is a +//! sharpening sigmoid applied to the head's NORMALIZED `(0, 1)` output `u`: +//! +//! ```text +//! g = sigmoid( logit(clamp(u, eps, 1 − eps)) / tau ), logit(x) = ln(x / (1 − x)) +//! ``` +//! +//! At `tau = 1` this is `g = u` exactly (sigmoid ∘ logit is the identity), and +//! the code special-cases it to return `u` untouched so the identity is +//! BIT-exact rather than round-tripped through `ln`/`exp`. As `tau` falls the +//! transform sharpens toward a step at `u = 0.5`; `tau` is annealed downward +//! over training by `params.leakance_gate.temperature` (`src/config.rs::LeakanceGate`). +//! +//! The transform is differentiable at every `tau > 0`, so autograd through +//! `clamp → log → div → sigmoid` gives the exact gradient; there is no +//! straight-through estimator. It is applied to the head OUTPUT before the +//! value reaches `setup_inputs`, so `TimestepLeakanceOp` and the sparse +//! backward are untouched. Pinned by `tests/leakance_gate.rs`. + +use burn::tensor::activation::sigmoid; +use burn::tensor::backend::Backend; +use burn::tensor::Tensor; + +/// Clamp margin applied to `u` before the logit. +/// +/// The head ends in a sigmoid, so `u` can be exactly `0.0` or `1.0` in f32 +/// once the pre-activation passes about ±17 (`sigmoid(17) == 1.0_f32`). An +/// unclamped logit there is `±inf`, and its backward `1/(u(1−u))` is `inf`, +/// which turns the gradient into NaN. `1e-6` is chosen because: +/// +/// * `1 − 1e-6` is ~17 ulps below `1.0_f32` (the ulp there is `2^-24 ≈ 6e-8`), +/// so `1 − u` at the clamp is resolved to about 3 % relative accuracy; +/// `1e-7` would sit only 2 ulps away and `1 − u` would carry ~20 % error. +/// * `logit(1e-6) ≈ −13.8`, well inside f32 range, and it only alters values +/// the head has already saturated past `|pre-activation| > 13.8`. +/// * The routing core is f32 throughout (invariant 2); no f64 is introduced. +pub const GATE_EPS: f32 = 1e-6; + +/// Apply the temperature-annealed gate to a normalized head output `u`. +/// +/// `tau == 1.0` returns `u` untouched (bit-exact identity). Any other +/// positive, finite `tau` evaluates the transform above. Generic over the +/// backend so the same function serves the training (`Autodiff`), eval, +/// probe and `dump_parameters` readers. +pub fn leakance_gate(u: Tensor, tau: f32) -> Tensor { + if tau == 1.0 { + return u; + } + debug_assert!(tau.is_finite() && tau > 0.0, "gate temperature must be positive and finite, got {tau}"); + let u = u.clamp(GATE_EPS, 1.0 - GATE_EPS); + // logit(u) = ln(u) − ln(1 − u) + let logit = u.clone().log() - u.neg().add_scalar(1.0).log(); + sigmoid(logit.div_scalar(tau)) +} diff --git a/src/training/mod.rs b/src/training/mod.rs index 6c301ee5..a3ae7cfb 100644 --- a/src/training/mod.rs +++ b/src/training/mod.rs @@ -15,6 +15,7 @@ pub mod checkpoint; pub mod driver; pub mod eval; pub mod forward; +pub mod gate; pub mod loss; pub mod metrics; pub mod optimizer; @@ -37,6 +38,7 @@ pub use loss::{ tau_trim_and_downsample, FilteredPair, }; pub use driver::{train, TrainState}; +pub use gate::{leakance_gate, GATE_EPS}; pub use eval::{evaluate, EvalOutput, EvalParams}; pub use metrics::Metrics; pub use adadelta::{AdaDelta, AdaDeltaConfig}; diff --git a/src/training/probe.rs b/src/training/probe.rs index f635493a..c36c8470 100644 --- a/src/training/probe.rs +++ b/src/training/probe.rs @@ -22,6 +22,7 @@ use crate::nn::kan_head::KanHead; use crate::routing::utils::denormalize; use crate::routing::{MuskingumCunge, RoutingInputs, SpatialParameters}; use crate::training::forward::{gather_params_to_subreaches, scatter_add_by_group}; +use crate::training::gate::leakance_gate; /// Detach `t` from its autograd graph and re-lift it as a `require_grad` /// leaf. Values are bit-identical; only the tape topology changes. @@ -189,6 +190,13 @@ pub fn probe_forward( .find(|(n, _)| n.as_str() == "leakance_factor") .map(|(_, t)| t.clone()) .or_else(|| params_map.get("leakance_factor").cloned()); + // Gate at the FINAL temperature, AFTER lifting: the leaf stays the raw + // head output `u`, so its gradient carries the gate's Jacobian exactly + // as the head's does in training (mirrors `forward`). + let leakance_factor = match cfg.params.leakance_gate.as_ref() { + Some(gate) => leakance_factor.map(|u| leakance_gate(u, gate.final_temperature())), + None => leakance_factor, + }; (k_d, d_gw, leakance_factor) } else { (None, None, None) diff --git a/tests/gamma_eval_parity.rs b/tests/gamma_eval_parity.rs index d630f5e6..9c200350 100644 --- a/tests/gamma_eval_parity.rs +++ b/tests/gamma_eval_parity.rs @@ -26,6 +26,7 @@ use burn::tensor::backend::Backend; use burn::tensor::{Int, Tensor, TensorData}; use chrono::NaiveDate; use common::{mock_config, InnerBackend, TestDevice}; +use ddrs::config::LeakanceGate; use ddrs::data::dataset::RoutingTensors; use ddrs::data::{RhoWindow, Staid}; use ddrs::nn::kan_head::{KanHead, KanHeadConfig}; @@ -33,6 +34,7 @@ use ddrs::sparse::SparseAdjacency; use ddrs::training::forward::{forward, forward_eval}; use ddrs::training::probe::probe_forward; use ndarray::Array2; +use std::collections::BTreeMap; type AB = Autodiff; @@ -113,7 +115,11 @@ fn all_paths(outputs: &[&str], cfg: &ddrs::config::Config) -> (Vec, Vec = head_ad.valid(); let tensors_ad = minimal_routing_tensors::(n, t, f, &device); - let train = forward::(cfg, &tensors_ad, &head_ad, &device, false).inner(); + // Training takes the epoch's gate temperature from the driver; the eval + // and probe readers apply the schedule's FINAL value, so hand training + // that same value here and the three must agree on one model. + let gate_tau = cfg.params.leakance_gate.as_ref().map(|g| g.final_temperature()); + let train = forward::(cfg, &tensors_ad, &head_ad, &device, false, gate_tau).inner(); let eval = forward_eval::( cfg, &minimal_routing_tensors::(n, t, f, &device), @@ -190,7 +196,14 @@ fn every_optional_head_output_reaches_all_three_readers() { leak.params.use_cuda_graphs = false; let mut fixed_q = mock_config(); fixed_q.params.defaults.insert("q_spatial".to_string(), 0.65); - let rows: [(&str, &[&str], &ddrs::config::Config); 8] = [ + // Leakance gate at a sharp FINAL temperature: the schedule's last value + // (0.05) is what eval, the probe and `dump_parameters` must apply, not + // the initial 1.0. `tests/leakance_gate.rs` pins the tau = 1 identity. + let mut leak_gate = leak.clone(); + leak_gate.params.leakance_gate = Some(LeakanceGate { + temperature: [(1usize, 1.0_f32), (2, 0.05)].into_iter().collect::>(), + }); + let rows: [(&str, &[&str], &ddrs::config::Config); 10] = [ ("n only (p, q fixed)", &["n"], &fixed_q), ("n + gamma (p, q fixed)", &["n", "gamma"], &fixed_q), ("n, q only (p fixed)", &["n", "q_spatial"], &base), @@ -207,6 +220,16 @@ fn every_optional_head_output_reaches_all_three_readers() { &["n", "q_spatial", "p_spatial", "x_storage", "gamma", "K_D", "d_gw", "leakance_factor"], &leak, ), + ( + "leakance + gate (final tau 0.05)", + &["n", "q_spatial", "p_spatial", "K_D", "d_gw", "leakance_factor"], + &leak_gate, + ), + ( + "everything + gate (final tau 0.05)", + &["n", "q_spatial", "p_spatial", "x_storage", "gamma", "K_D", "d_gw", "leakance_factor"], + &leak_gate, + ), ]; for (label, outputs, cfg) in rows { let (train, eval, probe) = all_paths(outputs, cfg); diff --git a/tests/leakance_gate.rs b/tests/leakance_gate.rs new file mode 100644 index 00000000..c94c2864 --- /dev/null +++ b/tests/leakance_gate.rs @@ -0,0 +1,398 @@ +//! Temperature-annealed binary gate on the head's `leakance_factor` output +//! (`src/training/gate.rs::leakance_gate`, config `params.leakance_gate`). +//! +//! Three layers, in priority order: +//! +//! 1. `tau = 1` is a BYTE-exact identity, both on the bare function and end +//! to end through all three head readers (`forward`, `forward_eval`, +//! `probe_forward`) against the same config with the block absent. This +//! is the guarantee that lets the block ship absent-by-default. +//! 2. The gate's gradient is the exact autograd gradient of +//! `sigmoid(logit(clamp(u)) / tau)`, checked against central finite +//! differences (of the formula in f64, and of the f32 implementation at +//! its noise floor) at several temperatures, including a sharp one. No +//! straight-through estimator anywhere. +//! 3. Numerical safety: monotone in `u`, the `tau -> 0` limits are 0 below +//! 0.5 and 1 above, and a temperature so small the sigmoid saturates +//! yields finite values and finite gradients even for `u` exactly 0 or 1. + +mod common; + +use burn::backend::{Autodiff, NdArray}; +use burn::module::AutodiffModule; +use burn::tensor::backend::Backend; +use burn::tensor::{Int, Tensor, TensorData}; +use chrono::NaiveDate; +use common::{mock_config, InnerBackend, TestDevice}; +use ddrs::config::{Config, LeakanceGate}; +use ddrs::data::dataset::RoutingTensors; +use ddrs::data::{RhoWindow, Staid}; +use ddrs::nn::kan_head::{KanHead, KanHeadConfig}; +use ddrs::sparse::SparseAdjacency; +use ddrs::training::forward::{forward, forward_eval}; +use ddrs::training::gate::{leakance_gate, GATE_EPS}; +use ddrs::training::probe::probe_forward; +use ndarray::Array2; +use std::collections::BTreeMap; + +type I = NdArray; +type AB = Autodiff; + +// Same acceptance rule as the repo's other gradchecks (`references/testing.md`). +const REL_TOL: f32 = 5e-3; +const ABS_TOL: f32 = 1e-4; + +fn gate_vec(u: &[f32], tau: f32) -> Vec { + let device = Default::default(); + let t = Tensor::::from_data(TensorData::new(u.to_vec(), [u.len()]), &device); + leakance_gate(t, tau).into_data().to_vec::().unwrap() +} + +/// f64 evaluation of the same formula with the implementation's f32 clamp +/// bounds (`1.0_f32 - GATE_EPS` is ~17 ulps below 1, not the real number +/// `1 - 1e-6`). Test-side only: the routing core and the gate stay f32. +fn reference(u: f64, tau: f64) -> f64 { + let lo = GATE_EPS as f64; + let hi = (1.0_f32 - GATE_EPS) as f64; + let u = u.clamp(lo, hi); + let logit = (u / (1.0 - u)).ln(); + 1.0 / (1.0 + (-logit / tau).exp()) +} + +fn sweep() -> Vec { + let mut u: Vec = (0..=1000).map(|k| k as f32 / 1000.0).collect(); + u.extend_from_slice(&[1e-7, 1e-6, 0.5 - 1e-6, 0.5 + 1e-6, 1.0 - 1e-6, 0.999_999_9]); + u +} + +// --------------------------------------------------------------------------- +// 1. Identity at tau = 1 +// --------------------------------------------------------------------------- + +#[test] +fn tau_one_is_bit_exact_identity() { + let u = sweep(); + let g = gate_vec(&u, 1.0); + for (a, b) in u.iter().zip(&g) { + assert_eq!(a.to_bits(), b.to_bits(), "tau=1 must return u untouched: {a} -> {b}"); + } +} + +/// Without the special case, `sigmoid(logit(u))` would NOT be bit-exact in +/// f32: pin that the round trip really does perturb values, so the identity +/// test above is testing something. +#[test] +fn round_trip_without_special_case_is_not_exact() { + let u = sweep(); + // tau = 1 + 1 ulp takes the general branch with a negligible sharpening. + let g = gate_vec(&u, f32::from_bits(1.0_f32.to_bits() + 1)); + let differs = u.iter().zip(&g).filter(|(a, b)| a.to_bits() != b.to_bits()).count(); + assert!(differs > 0, "expected the ln/exp round trip to perturb some values"); +} + +// --------------------------------------------------------------------------- +// 2. Values and gradients +// --------------------------------------------------------------------------- + +#[test] +fn matches_reference_formula() { + let u = sweep(); + for tau in [2.0_f32, 0.5, 0.1, 0.02] { + let g = gate_vec(&u, tau); + for (a, b) in u.iter().zip(&g) { + let r = reference(*a as f64, tau as f64) as f32; + assert!((r - b).abs() < 2e-6, "tau={tau} u={a}: got {b}, reference {r}"); + } + } +} + +/// The autograd gradient of the f32 gate against central finite differences. +/// +/// Two references. (a) The formula evaluated in f64 with a step of 1e-6: the +/// standard gradcheck practice, and the one held to the repo's tolerance +/// (`rel < 5e-3` or `abs < 1e-4`). (b) The f32 implementation itself with a +/// step of 2.5e-4: what the gate really computes, but a difference of two +/// values near 1.0 is quantized at `ulp(1) / 2h ≈ 1.2e-4`, so it is held to an +/// absolute floor of 1e-3 (eight times that noise) where the relative test +/// fails. The step is below the repo's usual 1e-3 because the sharpest +/// temperature has `O((4h/tau)^2 / 12)` truncation at `u = 0.5`. +#[test] +fn gradient_matches_central_differences() { + let device = Default::default(); + // Interior of (0, 1), on both sides of and at the switch point. + let u: Vec = vec![0.2, 0.35, 0.45, 0.5, 0.55, 0.65, 0.8]; + let h32 = 2.5e-4_f32; + let h64 = 1e-6_f64; + const F32_FD_ABS_FLOOR: f32 = 1e-3; + for tau in [1.0_f32, 0.5, 0.2, 0.05, 0.01] { + let leaf = Tensor::, 1>::from_data(TensorData::new(u.clone(), [u.len()]), &device) + .require_grad(); + let loss = leakance_gate(leaf.clone(), tau).sum(); + let grads = loss.backward(); + let analytic: Vec = leaf.grad(&grads).unwrap().into_data().to_vec().unwrap(); + + for (k, &uk) in u.iter().enumerate() { + let a = analytic[k]; + assert!(a.is_finite(), "tau={tau} u={uk}: non-finite gradient"); + + // (a) f64 central difference of the formula. + let fd64 = ((reference(uk as f64 + h64, tau as f64) - reference(uk as f64 - h64, tau as f64)) + / (2.0 * h64)) as f32; + let abs = (a - fd64).abs(); + let rel = abs / fd64.abs().max(1e-12); + assert!( + rel < REL_TOL || abs < ABS_TOL, + "tau={tau} u={uk}: analytic {a:.6e} vs f64 central-difference {fd64:.6e} (rel {rel:.2e}, abs {abs:.2e})" + ); + + // (b) f32 central difference of the implementation itself. + let plus = gate_vec(&[uk + h32], tau)[0]; + let minus = gate_vec(&[uk - h32], tau)[0]; + let fd32 = (plus - minus) / (2.0 * h32); + let abs = (a - fd32).abs(); + let rel = abs / fd32.abs().max(1e-12); + assert!( + rel < REL_TOL || abs < F32_FD_ABS_FLOOR, + "tau={tau} u={uk}: analytic {a:.6e} vs f32 central-difference {fd32:.6e} (rel {rel:.2e}, abs {abs:.2e})" + ); + } + if tau == 1.0 { + // The identity branch: gradient is exactly one everywhere. + assert!(analytic.iter().all(|g| *g == 1.0), "tau=1 gradient must be exactly 1: {analytic:?}"); + } + } +} + +// --------------------------------------------------------------------------- +// 3. Monotonicity, limits, saturation +// --------------------------------------------------------------------------- + +#[test] +fn monotone_in_u_and_fixed_at_half() { + let u: Vec = (1..100).map(|k| k as f32 / 100.0).collect(); + for tau in [1.0_f32, 0.5, 0.1, 0.01, 1e-4] { + let g = gate_vec(&u, tau); + for w in g.windows(2) { + assert!(w[1] >= w[0], "tau={tau}: not monotone at {w:?}"); + } + assert!(g.iter().all(|v| (0.0..=1.0).contains(v)), "tau={tau}: outside [0, 1]"); + // logit(0.5) = ln(0.5) - ln(0.5) is exactly 0 in f32, so 0.5 is a + // fixed point at every temperature. + assert_eq!(gate_vec(&[0.5], tau)[0], 0.5, "tau={tau}"); + } +} + +#[test] +fn small_temperature_limits_are_zero_below_half_and_one_above() { + let below: Vec = vec![0.01, 0.2, 0.4, 0.49]; + let above: Vec = vec![0.51, 0.6, 0.8, 0.99]; + // logit(0.49) = -0.04, so tau = 1e-3 already puts it at sigmoid(-40). + for tau in [1e-3_f32, 1e-4, 1e-5] { + for v in gate_vec(&below, tau) { + assert!(v < 1e-3, "tau={tau}: below-half value {v} did not go to 0"); + } + for v in gate_vec(&above, tau) { + assert!(v > 1.0 - 1e-3, "tau={tau}: above-half value {v} did not go to 1"); + } + } + // Sharpening is monotone in tau: the same u moves further from 0.5 as + // tau falls. + let u = 0.6_f32; + let seq: Vec = [1.0_f32, 0.5, 0.2, 0.1, 0.05].iter().map(|&t| gate_vec(&[u], t)[0]).collect(); + for w in seq.windows(2) { + assert!(w[1] >= w[0], "sharpening not monotone in tau: {seq:?}"); + } +} + +#[test] +fn saturating_temperature_is_finite_forward_and_backward() { + let device = Default::default(); + // Includes exact 0 and 1, which a saturated head sigmoid can emit and + // which are only safe because of the clamp before the logit. + let u: Vec = vec![0.0, 1e-8, 0.3, 0.5, 0.7, 1.0 - 1e-8, 1.0]; + for tau in [1e-6_f32, 1e-12, 1e-30, f32::MIN_POSITIVE] { + let g = gate_vec(&u, tau); + assert!(g.iter().all(|v| v.is_finite()), "tau={tau}: non-finite forward {g:?}"); + assert!(g.iter().all(|v| (0.0..=1.0).contains(v)), "tau={tau}: outside [0, 1] {g:?}"); + assert!(g[0] == 0.0 && g[6] == 1.0 && g[3] == 0.5, "tau={tau}: limits {g:?}"); + + let leaf = Tensor::, 1>::from_data(TensorData::new(u.clone(), [u.len()]), &device) + .require_grad(); + let grads = leakance_gate(leaf.clone(), tau).sum().backward(); + let grad: Vec = leaf.grad(&grads).unwrap().into_data().to_vec().unwrap(); + assert!(grad.iter().all(|v| v.is_finite()), "tau={tau}: non-finite gradient {grad:?}"); + } +} + +// --------------------------------------------------------------------------- +// 4. End to end through the three head readers +// --------------------------------------------------------------------------- + +/// Linear chain of `n` reaches, `t` hourly steps, one gauge at the outlet. +/// Mirrors `tests/gamma_eval_parity.rs::minimal_routing_tensors`. +fn minimal_routing_tensors( + n: usize, + t: usize, + f_attrs: usize, + device: &B::Device, +) -> RoutingTensors { + let adjacency = { + let mut dense = vec![0.0_f32; n * n]; + for i in 0..n - 1 { + dense[(i + 1) * n + i] = 1.0; + } + SparseAdjacency::from_dense(n, &dense, vec![1000.0; n], vec![0.001; n]) + }; + let attrs_vec: Vec = (0..n * f_attrs) + .map(|k| 0.2 + 0.6 * ((k * 7) % 11) as f32 / 10.0) + .collect(); + let spatial_attributes = + Tensor::::from_data(TensorData::new(attrs_vec, [n, f_attrs]), device); + let mut qp = vec![0.0_f32; t * n]; + for ti in 0..t { + let phase = (ti as f32) / (t.max(2) - 1) as f32 * 4.0 * std::f32::consts::PI; + for ri in 0..n { + qp[ti * n + ri] = (5.0 + phase.sin() * 2.0).max(0.1); + } + } + let q_prime = Tensor::::from_data(TensorData::new(qp, [t, n]), device); + let empty = |d: &B::Device| Tensor::::from_data(TensorData::new(Vec::::new(), [0, n]), d); + RoutingTensors { + adjacency, + spatial_attributes, + q_prime, + q_prime_daily: empty(device), + precip_hourly: empty(device), + temp_hourly: empty(device), + observations: Array2::zeros((1, 1)), + flat_indices: Tensor::::from_data(TensorData::from([(n - 1) as i32].as_slice()), device), + group_ids: Tensor::::from_data(TensorData::from([0i32].as_slice()), device), + num_gauges: 1, + gauge_staids: vec![Staid::new("dummy")], + window: RhoWindow { + start_day_idx: 0, + rho_days: 1, + window_start: NaiveDate::from_ymd_opt(1990, 1, 1).unwrap(), + }, + initial_state: None, + impervious_mask: None, + } +} + +const LEAK_OUTPUTS: [&str; 6] = ["n", "q_spatial", "p_spatial", "K_D", "d_gw", "leakance_factor"]; + +fn leak_head(f_attrs: usize, device: &TestDevice) -> KanHead { + KanHeadConfig::new( + (0..f_attrs).map(|i| format!("attr_{i}")).collect(), + LEAK_OUTPUTS.iter().map(|s| s.to_string()).collect(), + 42, + ) + .with_hidden_size(8) + .with_num_hidden_layers(1) + .init::(device) +} + +fn leak_config(schedule: Option<&[(usize, f32)]>) -> Config { + let mut cfg = mock_config(); + cfg.params.use_leakance = true; + cfg.params.use_cuda_graphs = false; + cfg.params.leakance_gate = schedule.map(|s| LeakanceGate { + temperature: s.iter().copied().collect::>(), + }); + cfg +} + +/// `(train, eval, probe)` gauge hydrographs for one config. Training takes +/// the schedule's FINAL temperature so the three readers see the same model. +fn three_readers(cfg: &Config) -> (Vec, Vec, Vec) { + let device = TestDevice::default(); + let (n, t, f) = (5usize, 48usize, 4usize); + let head_ad = leak_head(f, &device); + let head_inner: KanHead = head_ad.valid(); + let tensors_ad = minimal_routing_tensors::(n, t, f, &device); + let gate_tau = cfg.params.leakance_gate.as_ref().map(|g| g.final_temperature()); + + let train = forward::(cfg, &tensors_ad, &head_ad, &device, false, gate_tau).inner(); + let eval = forward_eval::( + cfg, + &minimal_routing_tensors::(n, t, f, &device), + &head_inner, + &device, + false, + None, + None, + None, + ); + let (probe, _leaves) = probe_forward::(cfg, &tensors_ad, &head_ad, &device, &[]); + let v = |t: Tensor| t.into_data().to_vec::().unwrap(); + (v(train), v(eval), v(probe.inner())) +} + +fn assert_bits_equal(label: &str, a: &[f32], b: &[f32]) { + assert_eq!(a.len(), b.len(), "{label}: shape"); + for (k, (x, y)) in a.iter().zip(b).enumerate() { + assert_eq!(x.to_bits(), y.to_bits(), "{label}: bit mismatch at {k}: {x} vs {y}"); + } +} + +/// THE parity test: a config whose gate block schedules `tau = 1` must route +/// byte-identically to the same config without the block, in all three +/// readers. The block absent is the historical code path. +#[test] +fn gate_at_tau_one_is_byte_exact_through_all_readers() { + let (train_off, eval_off, probe_off) = three_readers(&leak_config(None)); + let (train_on, eval_on, probe_on) = three_readers(&leak_config(Some(&[(1, 1.0)]))); + assert!(train_off.iter().all(|v| v.is_finite())); + assert_bits_equal("forward", &train_off, &train_on); + assert_bits_equal("forward_eval", &eval_off, &eval_on); + assert_bits_equal("probe_forward", &probe_off, &probe_on); +} + +/// Teeth for the parity test above: a sharp final temperature must actually +/// move the routed hydrograph, in all three readers, or the identity could be +/// vacuous (the readers silently dropping the gate would pass it). +#[test] +fn gate_at_sharp_tau_changes_output_in_all_readers() { + let (train_off, eval_off, probe_off) = three_readers(&leak_config(None)); + let (train_on, eval_on, probe_on) = three_readers(&leak_config(Some(&[(1, 1.0), (2, 0.05)]))); + let max_diff = |a: &[f32], b: &[f32]| { + a.iter().zip(b).map(|(x, y)| (x - y).abs()).fold(0.0_f32, f32::max) + }; + assert!(train_on.iter().all(|v| v.is_finite()), "gated forward produced non-finite values"); + let d_train = max_diff(&train_off, &train_on); + let d_eval = max_diff(&eval_off, &eval_on); + let d_probe = max_diff(&probe_off, &probe_on); + assert!(d_train > 0.0, "forward: gate at tau=0.05 was a no-op"); + assert!(d_eval > 0.0, "forward_eval: gate at tau=0.05 was a no-op"); + assert!(d_probe > 0.0, "probe_forward: gate at tau=0.05 was a no-op"); + // And the three readers still agree with each other on the gated model. + assert_bits_equal("gated: forward vs probe", &train_on, &probe_on); + let worst = train_on + .iter() + .zip(&eval_on) + .map(|(x, y)| (x - y).abs() / x.abs().max(1e-6)) + .fold(0.0_f32, f32::max); + assert!(worst < 1e-5, "gated: forward vs forward_eval diverge, worst rel diff {worst:.3e}"); +} + +/// `forward` refuses a temperature that contradicts the config, in both +/// directions, so a caller cannot train the wrong model silently. +#[test] +#[should_panic(expected = "gate_tau")] +fn forward_rejects_temperature_without_gate_block() { + let device = TestDevice::default(); + let cfg = leak_config(None); + let head = leak_head(4, &device); + let tensors = minimal_routing_tensors::(5, 48, 4, &device); + let _ = forward::(&cfg, &tensors, &head, &device, false, Some(0.5)); +} + +#[test] +#[should_panic(expected = "gate_tau")] +fn forward_rejects_missing_temperature_with_gate_block() { + let device = TestDevice::default(); + let cfg = leak_config(Some(&[(1, 0.5)])); + let head = leak_head(4, &device); + let tensors = minimal_routing_tensors::(5, 48, 4, &device); + let _ = forward::(&cfg, &tensors, &head, &device, false, None); +} diff --git a/tests/leakance_off_parity.rs b/tests/leakance_off_parity.rs index 328ee6bb..2264d0cd 100644 --- a/tests/leakance_off_parity.rs +++ b/tests/leakance_off_parity.rs @@ -283,6 +283,8 @@ fn head_driven_leakance_changes_output() { &head, &device, false, + // no `params.leakance_gate` in these fixtures, so no gate temperature + None, ); let sum_leak: f32 = out_leak.into_data().to_vec::().unwrap().iter().sum(); @@ -293,6 +295,7 @@ fn head_driven_leakance_changes_output() { &head, &device, false, + None, ); let sum_no_leak: f32 = out_no_leak.into_data().to_vec::().unwrap().iter().sum(); From c8f3c2b4d7ec6c86398320608783329625da5f7f Mon Sep 17 00:00:00 2001 From: Tadd Bindas Date: Wed, 16 Sep 2026 22:43:41 -0400 Subject: [PATCH 2/3] spec: numbered reference of every model equation, as implemented 104 numbered equations over 13 sections, each derived from the code rather than from prior docs, with symbol tables giving units and whether each quantity is learned per reach, set by config, read from data, or derived. Covers the head and its normalization, the new gate, the geometry chain with a catalogue of every clamp and what it does to the gradient, stage-dependent roughness, the Muskingum coefficients, the sparse adjoint, the leakance flux with all six partials, the gradient-path labels, all five losses, and the diagnostics. It found ten code-versus-documentation discrepancies. The load-bearing one: the leakance term computes its width as (p*d)^q while the routing geometry computes the same reach's width as p*d^q, so the area multiplying K_D is smaller by the depth-independent factor p^(q-1) = 0.3445 at the prescribed p = 21, q = 0.65, and a learned K_D absorbs it. The flux and skill results are unaffected because a constant rescaling is absorbable, but K_D is not directly comparable to a measured streambed leakance without dividing it out. Inherited from DDR, so the port is faithful; leakance_reference_match pins that. validate_params.py now divides the factor out and records the caveat. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_0143bx7gVz8ug2oc6DTy5jtz --- experiments/leakance/bias_decomposition.py | 60 + experiments/leakance/validate_params.py | 179 ++ .../2026-09-17-model-equations-reference.md | 1943 +++++++++++++++++ 3 files changed, 2182 insertions(+) create mode 100644 experiments/leakance/bias_decomposition.py create mode 100755 experiments/leakance/validate_params.py create mode 100644 research/specs/2026-09-17-model-equations-reference.md diff --git a/experiments/leakance/bias_decomposition.py b/experiments/leakance/bias_decomposition.py new file mode 100644 index 00000000..70186244 --- /dev/null +++ b/experiments/leakance/bias_decomposition.py @@ -0,0 +1,60 @@ +"""Decompose where each learned term acts. KGE splits skill into three independent +error modes: r (timing and shape), alpha = sd_sim/sd_obs (amplitude), beta = +mean_sim/mean_obs (volume). Roughness conserves mass, so it can move r and alpha +but not beta. Leakance removes mass, so beta is the only mode it can reach that +roughness cannot. Comparing the same three numbers for the unrouted inflow, the +leakance-off arm and the leakance-on arm says which bias each term actually absorbed.""" +import numpy as np, zarr, json +B="/home/tbindas/projects/ddrs/.ddrs/runs" +ON=f"{B}/2026-09-16T04-44-15Z-train-and-test"; OFF=f"{B}/2026-09-16T06-11-35Z-train-and-test" + +def comps(sim, obs): + out=[] + for i in range(sim.shape[0]): + s,o = sim[i], obs[i] + m = np.isfinite(s)&np.isfinite(o)&(o>=0) + if m.sum() < 365: out.append((np.nan,)*4); continue + s,o = s[m], o[m] + so, oo = s.std(), o.std() + if oo <= 0 or so <= 0: out.append((np.nan,)*4); continue + r = np.corrcoef(s,o)[0,1]; a = so/oo; b = s.mean()/o.mean() + kge = 1 - np.sqrt((r-1)**2 + (a-1)**2 + (b-1)**2) + out.append((r,a,b,kge)) + return np.array(out) + +def load_eval(run): + g = zarr.open_group(f"{run}/eval/predictions.zarr", mode="r") + return np.asarray(g["predictions"][:],float), np.asarray(g["observations"][:],float), np.asarray(g["gage_ids"][:]) + +res={} +for lab, run in [("leakance ON", ON), ("leakance OFF", OFF)]: + sim, obs, ids = load_eval(run) + if sim.shape[0] < sim.shape[1] and sim.shape[0] < 5000: pass + else: sim, obs = sim.T, obs.T + res[lab] = (comps(sim, obs), ids) + print(f"{lab}: {sim.shape} gauges x days") + +# unrouted summed inflow, from the shared baseline +bm = json.load(open(f"{ON}/baseline/manifest.json")) +ng = len(bm["gage_ids"]); nd = len(bm["time_range_daily"]) +bp = np.fromfile(f"{ON}/baseline/predictions.f32", dtype=np.float32).reshape(ng, nd).astype(float) +bo = np.fromfile(f"{ON}/baseline/observations.f32", dtype=np.float32).reshape(ng, nd).astype(float) +res["summed inflow (no routing)"] = (comps(bp, bo), np.asarray(bm["gage_ids"])) +print(f"summed inflow: {bp.shape} gauges x days") + +print(f"\n{'arm':28s} {'r':>8s} {'alpha':>8s} {'beta':>8s} {'KGE':>8s} n") +order=["summed inflow (no routing)","leakance OFF","leakance ON"] +med={} +for k in order: + c,_ = res[k]; ok=np.isfinite(c[:,3]) + med[k]=np.nanmedian(c[ok],axis=0) + print(f"{k:28s} " + " ".join(f"{v:8.4f}" for v in med[k]) + f" {ok.sum()}") + +print("\nwhat routing moved (leakance OFF minus unrouted):") +d=med["leakance OFF"]-med["summed inflow (no routing)"] +print(" " + " ".join(f"{n}={v:+.4f}" for n,v in zip(["r","alpha","beta","KGE"],d))) +print("what leakance added on top (ON minus OFF):") +d2=med["leakance ON"]-med["leakance OFF"] +print(" " + " ".join(f"{n}={v:+.4f}" for n,v in zip(["r","alpha","beta","KGE"],d2))) +print("\nvolume bias |beta-1| at the median gauge:") +for k in order: print(f" {k:28s} {abs(med[k][2]-1)*100:6.2f} %") diff --git a/experiments/leakance/validate_params.py b/experiments/leakance/validate_params.py new file mode 100755 index 00000000..9216e642 --- /dev/null +++ b/experiments/leakance/validate_params.py @@ -0,0 +1,179 @@ +#!/usr/bin/env python +"""Validate a leakance arm's learned fields against measured physical ranges, and +map every parameter over the CONUS river network. + +Three validations, each against an external reference rather than internal +consistency: + roughness n -> Chow (1959) tabulated Manning's n for natural channels + leakance K_D -> MODFLOW streambed leakance K/M, against measured streambed Kv + offset d_gw -> depth of the water table below the channel + +usage: validate_params.py [--zeta ] --out-prefix P +""" +import argparse, json +import numpy as np, netCDF4 as nc +import matplotlib; matplotlib.use("Agg") +import matplotlib.pyplot as plt, matplotlib.colors as mcolors +from scipy.stats import spearmanr + +FAB = "/home/tbindas/projects/ddr/data/merit/riv_pfaf_7_MERIT_Hydro_v07_Basins_v01_bugfix1.shp" +ATTR = "/home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc" + +# --- external reference ranges ------------------------------------------------- +# Chow, V.T. (1959) Open-Channel Hydraulics, Table 5-6, natural streams. +CHOW = { + "clean straight, full stage": (0.025, 0.033), + "clean winding, pools and shoals": (0.033, 0.045), + "mountain, gravel and cobbles": (0.030, 0.050), + "mountain, cobbles large boulders": (0.040, 0.070), + "sluggish, weedy, deep pools": (0.050, 0.080), + "very weedy reaches": (0.075, 0.150), +} +# Streambed vertical hydraulic conductivity, measured, four orders of magnitude +# (Nature Scientific Reports 2020, 10:3333): 7.57e-3 .. 1.81 m/day. +KV_MEAS_MDAY = (7.57e-3, 1.81) +BED_THICK_M = 1.0 # leakance = Kv / M; state the assumption rather than hide it + +# The leakance term computes its plan-view width as (p*d)^q, while the routing +# geometry computes the SAME reach's width as p*d^q (src/geometry.rs, Leopold and +# Maddock). Those differ by the depth-independent factor p^(q-1), which is 0.3445 +# at the prescribed p = 21, q = 0.65. So the area multiplying K_D is 0.3445 of the +# routing's plan area, and a learned K_D silently absorbs that factor. To compare +# K_D against a measured streambed leakance we must divide it out. This is +# inherited from DDR (tests/leakance_reference_match.rs matches its formula +# exactly), so it is a property of the formulation, not a port error. The term's +# area is also dimensionally not m^2 unless q = 1, which makes this comparison +# approximate rather than exact. See research/specs/2026-09-17-model-equations-reference.md. +AREA_FORM_FACTOR = 0.3445 + +def band(v, lo, hi): + return float(((v >= lo) & (v <= hi)).mean()) + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("params"); ap.add_argument("--zeta", default=None) + ap.add_argument("--out-prefix", required=True) + ap.add_argument("--fabric", default=FAB) + a = ap.parse_args() + + d = nc.Dataset(a.params) + have = lambda v: v in d.variables + g = lambda v: np.asarray(d[v][:], dtype=float) + cid = np.asarray(d["COMID"][:], dtype=np.int64) + F = {"n": g("n")} + for v in ["K_D", "d_gw", "leakance_factor", "gamma", "q_spatial", "p_spatial", "slope"]: + if have(v): F[v] = g(v) + if "K_D" in F and "leakance_factor" in F: + F["conductance"] = F["leakance_factor"] * F["K_D"] + + report = {"n_reaches": int(cid.size), "source": a.params} + + # ---- 1. roughness against Chow ------------------------------------------- + n = F["n"] + report["roughness"] = { + "p2": float(np.percentile(n, 2)), "median": float(np.median(n)), + "p98": float(np.percentile(n, 98)), + "frac inside Chow natural-channel envelope 0.025-0.150": band(n, 0.025, 0.150), + "frac below 0.025 (smoother than any tabulated natural channel)": float((n < 0.025).mean()), + "frac above 0.150 (rougher than very weedy)": float((n > 0.150).mean()), + "by Chow class": {k: band(n, lo, hi) for k, (lo, hi) in CHOW.items()}, + } + + # ---- 2. leakance against measured streambeds ----------------------------- + if "conductance" in F: + C = F["conductance"] + # divide out the area-form factor so C is comparable to a real leakance + kv_mday = (C / AREA_FORM_FACTOR) * BED_THICK_M * 86400.0 + lo, hi = KV_MEAS_MDAY + report["leakance"] = { + "assumed streambed thickness m": BED_THICK_M, + "area-form factor divided out": AREA_FORM_FACTOR, + "conductance 1/s p10/median/p90": [float(np.percentile(C,10)), float(np.median(C)), float(np.percentile(C,90))], + "implied Kv m/day p10/median/p90": [float(np.percentile(kv_mday,10)), float(np.median(kv_mday)), float(np.percentile(kv_mday,90))], + "measured Kv range m/day": list(KV_MEAS_MDAY), + "frac inside measured range": band(kv_mday, lo, hi), + "frac below measured minimum": float((kv_mday < lo).mean()), + "position in measured range (0=min,1=max, log scale)": + float(np.clip((np.log10(np.median(kv_mday)) - np.log10(lo)) / (np.log10(hi) - np.log10(lo)), 0, 1)), + } + + # ---- 3. the offset as a water table -------------------------------------- + if "d_gw" in F: + dg = F["d_gw"] + rec = {"p10": float(np.percentile(dg,10)), "median": float(np.median(dg)), "p90": float(np.percentile(dg,90)), + "frac positive (offset ABOVE the bed, so the losing-only clamp gates the reach off)": float((dg > 0).mean())} + if a.zeta: + z = nc.Dataset(a.zeta) + dep = np.asarray(z["depth_mean"][:], dtype=float) + # offsets are on the full network, depths on the eval network: join by COMID + zc = np.asarray(z["COMID_eval"][:], dtype=np.int64) + import pandas as pd + pos = pd.Index(cid).get_indexer(zc) + ok = pos >= 0 + rec["frac of eval reaches with offset above mean depth (leakance switched off)"] = \ + float((dg[pos[ok]] > dep[ok]).mean()) + report["offset"] = rec + + # ---- attribute correlations: is any of this reading the landscape? ------- + A = nc.Dataset(ATTR); acid = np.asarray(A["COMID"][:], dtype=np.int64) + idx = {c: i for i, c in enumerate(acid)} + sel = np.array([idx.get(c, -1) for c in cid]); ok = sel >= 0 + attrs = {k: np.asarray(A[k][:], dtype=float)[sel[ok]] + for k in ["aridity", "meanslope", "permeability", "meanP", "meanelevation", "HWSD_sand"]} + corr = {} + for name, v in F.items(): + vv = v[ok]; row = {} + for k, av in attrs.items(): + m = np.isfinite(vv) & np.isfinite(av) + row[k] = round(float(spearmanr(vv[m], av[m]).statistic), 3) + corr[name] = row + report["attribute_correlations"] = corr + + # ---- are the outputs independent? --------------------------------------- + keys = [k for k in ["n", "gamma", "K_D", "d_gw", "leakance_factor"] if k in F] + if len(keys) > 1: + X = np.column_stack([F[k] for k in keys]); X = (X - X.mean(0)) / X.std(0) + s = np.linalg.svd(X, compute_uv=False); p = s**2 / (s**2).sum() + report["output_independence"] = { + "fields": keys, "variance_shares": [round(float(x), 4) for x in p], + "effective_rank": round(float(np.exp(-(p * np.log(p)).sum())), 3), + "pairwise_spearman": {f"{keys[i]}|{keys[j]}": round(float(spearmanr(F[keys[i]], F[keys[j]]).statistic), 4) + for i in range(len(keys)) for j in range(i+1, len(keys))}, + } + + with open(f"{a.out_prefix}_validation.json", "w") as fh: + json.dump(report, fh, indent=2) + print(json.dumps(report, indent=2)) + + # ---- maps ---------------------------------------------------------------- + import geopandas as gpd, pandas as pd, pyogrio + gdf = pyogrio.read_dataframe(a.fabric, columns=["COMID", "uparea"]).set_index("COMID") + gdf = gdf.loc[gdf.index.intersection(cid)] + pos = pd.Index(cid).get_indexer(gdf.index.values) + for k, v in F.items(): gdf[k] = v[pos] + gdf["lw"] = np.clip(0.15 + 0.35*np.log10(np.maximum(gdf["uparea"].values, 1.0))/5.0, 0.15, 1.2) + gdf = gdf.sort_values("uparea") + + panels = [k for k in ["n", "conductance", "d_gw", "leakance_factor", "gamma", "q_spatial"] if k in gdf] + titles = {"n": "Manning's $n$", "conductance": "effective conductance $g\\cdot K_D$ (1/s)", + "d_gw": "groundwater offset $d_{gw}$ (m)", "leakance_factor": "leakance gate $g$", + "gamma": "stage exponent $\\gamma$", "q_spatial": "width exponent $q$"} + ncol = 2; nrow = int(np.ceil(len(panels)/ncol)) + fig, axes = plt.subplots(nrow, ncol, figsize=(9.5*ncol, 5.5*nrow), squeeze=False) + for ax, col in zip(axes.ravel(), panels): + v = gdf[col].values; lo, hi = np.percentile(v, [2, 98]) + logc = col == "conductance" + cmap = {"n": "plasma_r", "conductance": "viridis", "d_gw": "coolwarm", + "leakance_factor": "RdYlBu_r", "gamma": "magma", "q_spatial": "cividis"}.get(col, "viridis") + norm = mcolors.LogNorm(vmin=max(lo,1e-12), vmax=hi) if logc else mcolors.Normalize(vmin=lo, vmax=hi) + gdf.plot(ax=ax, column=col, cmap=cmap, norm=norm, linewidth=gdf["lw"].values, rasterized=True) + ax.set_title(titles.get(col, col), fontsize=12); ax.set_axis_off() + fig.colorbar(plt.cm.ScalarMappable(norm=norm, cmap=cmap), ax=ax, fraction=0.03, pad=0.01) + for ax in axes.ravel()[len(panels):]: ax.set_axis_off() + fig.suptitle(f"Learned parameter fields, CONUS ({cid.size:,} reaches)\n{a.params}", fontsize=12) + fig.tight_layout() + out = f"{a.out_prefix}_maps.png"; fig.savefig(out, dpi=115, bbox_inches="tight") + print("wrote", out) + +if __name__ == "__main__": + main() diff --git a/research/specs/2026-09-17-model-equations-reference.md b/research/specs/2026-09-17-model-equations-reference.md new file mode 100644 index 00000000..b1e5dc34 --- /dev/null +++ b/research/specs/2026-09-17-model-equations-reference.md @@ -0,0 +1,1943 @@ +# ddrs model equations: complete reference as implemented + +**Date:** 2026-09-17 +**Worktree:** `.claude/worktrees/agent-a10b11fb99c152d22`, branch `leakance-gate`, +HEAD `289163b` ("feat(leakance): temperature-annealed gate replaces the continuous factor"). +**Audience:** an analyst judging whether the learned parameters are physically +meaningful and what this work contributes to distributed hydrological modelling. + +## 0. How this document was produced, and how to read it + +Every equation below was transcribed from source read in this worktree, not from +any prior document. The files read in full or in the relevant part were: +`src/nn/kan_head.rs`, `src/routing/utils.rs`, `src/training/gate.rs`, +`src/geometry.rs`, `src/routing/mmc.rs`, `src/routing/mmc_op.rs`, +`src/routing/leakance.rs`, `src/sparse/mod.rs`, `src/training/loss.rs`, +`src/training/forward.rs`, `src/training/metrics.rs`, `src/training/eval.rs`, +`src/training/driver.rs`, `src/dump_parameters.rs`, `src/config.rs`, +`src/data/dataset.rs`, `tests/stage_roughness.rs`, and the pinned +`rskan` v0.1.3 `KanLayer::forward`. Where a docstring and the code disagree the +code is authoritative here and the disagreement is listed in §12. + +**Symbol-status column.** Every symbol table carries a `status` column with one +of five values: + +| status | meaning | +|---|---| +| **learned** | a per-reach KAN output, trained by gradient descent | +| **config** | prescribed by YAML (`params.*`, `experiment.*`) | +| **data** | read from a store (attributes, adjacency, forcing, observations) | +| **derived** | computed from the above inside the forward pass | +| **literal** | a numeric constant hard-coded in Rust, not configurable | + +**Clamps.** Every clamp is named and numbered, because a saturated clamp +sets the incoming gradient to exactly zero for that reach-timestep (the +backward implements each clamp as a boolean mask fill, see §9). A clamp is +therefore both a physical floor and a gradient sink, and a parameter whose +clamp binds on most reach-timesteps is not being trained by that path. + +**Backend precision.** The entire routing core is f32 +(`CLAUDE.md` invariant 2). `Δt = 3600 s` is hard-coded +(`src/routing/mmc.rs::DT_SECONDS`); there is no sub-hourly routing step. + +--- + +## 1. Pipeline overview + +``` + attributes (N x F, z-scored) + | + v §2 KAN head, ends in sigmoid + u in (0,1)^P normalized parameters + | + +---> §3 leakance_factor only: temperature-annealed gate g(u; tau) + | + v §2.3 denormalize (linear or log-space) per parameter box + physical parameters: n_0, q_spatial, p_spatial, [gamma], [K_D, d_gw, f] + | + v §4/§5 per reach, per timestep, from Q_t + depth -> top width -> side slope -> bottom width -> A, P_w, R -> v -> celerity c + | + v §6 + K = L/c ; X (constant, or Cunge-derived) ; c1..c4 + | + +---> §8 zeta = f * area_z * K_D * (d - d_gw) subtracted from the RHS + | + v §7 (I - c1 .* N) Q_{t+1} = c2 (N Q_t) + c3 Q_t + c4 q'_t - zeta + sparse lower-triangular forward substitution + hand-written adjoint + | + v Q_{t+1} = max(x_sol, 1e-4) + gauge scatter-add -> tau shift + daily area pooling -> §10 loss +``` + +--- + +## 2. The parameter head + +### 2.1 Input + +`src/data/dataset.rs::MeritGagesDataset::finalize_attrs` builds the head input. +Missing COMIDs are filled with the attribute's row mean, then each attribute +column is z-scored against precomputed statistics: + +$$ +X_{i,\,\text{fi}} \;=\; \frac{a_{\text{fi},i} - \mu_{\text{fi}}}{\sigma_{\text{fi}}} +\tag{1} +$$ + +and the array is transposed to `(N, F)`. + +| symbol | meaning | units | status | +|---|---|---|---| +| $a_{\text{fi},i}$ | raw attribute `fi` at reach `i` | attribute-specific | data | +| $\mu_{\text{fi}}, \sigma_{\text{fi}}$ | per-attribute mean / std | same | data (statistics JSON) | +| $N$ | reaches in the batch's network (parent resolution) | - | data | +| $F$ | number of attributes, `kan_head.input_var_names` | - | config | + +The shipped $F = 10$ attribute set is `SoilGrids1km_clay, aridity, +meanelevation, meanP, NDVI, meanslope, log10_uparea, SoilGrids1km_sand, +ETPOT_Hargr, Porosity`. **The head sees no discharge, no state, and no time.** +It is a static map from catchment attributes to reach parameters; all +time-dependence in the model enters through the geometry's dependence on $Q_t$. + +### 2.2 Layer stack + +`src/nn/kan_head.rs::KanHead::forward` composes, per parameter group (default: +one group over all outputs): + +$$ +h^{(0)} = \mathrm{Linear}_{F \to H}(X), \qquad +h^{(\ell)} = \mathrm{KanLayer}_{H \to H}\!\left(h^{(\ell-1)}\right), \;\; \ell = 1 \ldots L, +\tag{2} +$$ + +$$ +z = \mathrm{Linear}_{H \to P}\!\left(h^{(L)}\right), \qquad u = \sigma(z) = \frac{1}{1 + e^{-z}} \in (0,1)^{N \times P} +\tag{3} +$$ + +There is **no activation between KAN blocks** (DDR `kan.py` parity, enforced by +`CLAUDE.md` invariant 5). `u` is transposed to `[P, N]` and split row-by-row +into a `HashMap` keyed by `kan_head.learnable_parameters` in order. + +One `rskan::KanLayer` (pinned tag `v0.1.3`, `rskan/src/layer.rs::KanLayer::forward`) +computes, for input $x \in \mathbb{R}^{B \times I}$ and output index $o$: + +$$ +y_{b,o} \;=\; \sum_{i=1}^{I} m_{i,o}\left[\, s^{\mathrm{b}}_{i,o}\,\mathrm{SiLU}(x_{b,i}) +\;+\; s^{\mathrm{sp}}_{i,o} \sum_{j} c_{i,o,j}\, B_{j,k}\!\left(x_{b,i}\right) \right] +\tag{4} +$$ + +$$ +\mathrm{SiLU}(x) = x\,\sigma(x) +\tag{5} +$$ + +with $B_{j,k}$ the order-$k$ B-spline basis on the layer's grid. So every edge +$(i,o)$ carries its own learned univariate function: a fixed SiLU shape scaled +by $s^{\mathrm{b}}$ plus a free spline scaled by $s^{\mathrm{sp}}$. + +| symbol | meaning | units | status | +|---|---|---|---| +| $H$ | `kan_head.hidden_size`, default 21 | - | config | +| $L$ | `kan_head.num_hidden_layers`, default 2 | - | config | +| $P$ | number of emitted parameters | - | config | +| $k$ | B-spline order, default 3 (current arm: 2) | - | config | +| grid | B-spline intervals, default 5 (current arm: 50) | - | config | +| $c_{i,o,j}, s^{\mathrm{b}}, s^{\mathrm{sp}}$ | spline coefficients and scales | - | learned (head weights) | +| $m_{i,o}$ | edge mask | - | learned/fixed (init to 1) | +| $u$ | normalized parameter | dimensionless in $(0,1)$ | derived | + +Initialization (`src/nn/init.rs`, called from `KanHeadConfig::init_trunk`): +input `Linear` weight is Kaiming-normal, $\mathrm{std} = \sqrt{2}/\sqrt{F}$; +output `Linear` weight is Xavier-normal, $\mathrm{std} = 0.1\sqrt{2/(H+P)}$; +both biases zero; **every** inner `KanLayer` receives the SAME seed (a preserved +DDR quirk, invariant 5). + +Three optional architecture switches, all default off and all breaking DDR +parity by construction: +`input_layer_kan` replaces $\mathrm{Linear}_{F\to H}$ with $\mathrm{KanLayer}_{F\to H}$; +`output_layer_kan` replaces $\mathrm{Linear}_{H\to P}$ with $\mathrm{KanLayer}_{H\to P}$; +`parameter_groups` partitions the outputs into independently-weighted trunks, +group 0 seeded with `seed` and group $i>0$ with +`seed + i * 0x9E3779B97F4A7C15` (wrapping). + +**Structural note the analyst needs.** Under the default (all-`Linear` +read-out, one shared trunk), every output is an affine functional of the same +latent vector: $z_{\cdot,o} = w_o \cdot h^{(L)} + b_o$. If the informative part +of $h^{(L)}$ is effectively one direction, then any two outputs are *exactly* +affinely related in logit space. The code's own docstring on +`KanHeadConfig::output_layer_kan` records the measurement: +$\mathrm{logit}(q) = 3.979\,\mathrm{logit}(n) + 2.752$ with $R^2 = 0.987$. +Two "different" learned parameters can therefore be one parameter relabelled. +That is a property of the head topology, not of the physics. + +### 2.3 Normalized to physical: `denormalize` + +`src/routing/utils.rs::denormalize`, with box $[\theta_{\mathrm{lo}}, \theta_{\mathrm{hi}}]$ +from `params.parameter_ranges`: + +Linear branch: +$$ +\theta = u\,(\theta_{\mathrm{hi}} - \theta_{\mathrm{lo}}) + \theta_{\mathrm{lo}} +\tag{6} +$$ + +Log-space branch (parameter name listed in `params.log_space_parameters`): +$$ +\theta = \exp\!\Big[\, u\big(\ln \theta_{\mathrm{hi}} - \Lambda(\theta_{\mathrm{lo}})\big) + \Lambda(\theta_{\mathrm{lo}}) \Big] +\tag{7} +$$ + +where $\Lambda$ is `src/routing/utils.rs::log_space_lower`: + +$$ +\Lambda(\theta_{\mathrm{lo}}) \;=\; +\begin{cases} +\ln \theta_{\mathrm{lo}}, & \theta_{\mathrm{lo}} > 0\\[2pt] +\ln\!\left(\theta_{\mathrm{lo}} + 10^{-6}\right), & \theta_{\mathrm{lo}} \le 0 +\end{cases} +\tag{8} +$$ + +**Why $\Lambda$ exists.** DDR's Python applies `bounds[0] + 1e-6` +unconditionally. For a box whose lower bound is small next to $10^{-6}$ that +inverts the map. Worked example from the docstring, for the historical +$K_D \in [10^{-8}, 10^{-6}]$: unconditional nudging gives +$\ln(1.01\times10^{-6}) = -13.8065$ as the lower end and +$\ln(10^{-6}) = -13.8155$ as the upper, so +$\ln\theta_{\mathrm{hi}} - \Lambda < 0$; the map runs backwards and spans about +1% instead of two decades. The consequence is recorded in the current +experiment config: "before it, $K_D$'s declared box $[10^{-8}, 10^{-6}]$ +collapsed to a 1% band and $K_D$ was frozen, so every leakance run in the July +campaign learned only (`leakance_factor`, `d_gw`)." $\Lambda$ uses +$\ln\theta_{\mathrm{lo}}$ whenever the bound is already positive, and keeps the +old behaviour only for $\theta_{\mathrm{lo}} \le 0$ so that zero-floored boxes +map exactly as before. + +Inverse (`src/training/forward.rs::physical_to_normalized`), sharing the same +$\Lambda$ so the two directions cannot drift: + +$$ +u = \frac{\theta - \theta_{\mathrm{lo}}}{\theta_{\mathrm{hi}} - \theta_{\mathrm{lo}}} +\quad\text{(linear)}, \qquad +u = \frac{\ln \theta - \Lambda(\theta_{\mathrm{lo}})}{\ln \theta_{\mathrm{hi}} - \Lambda(\theta_{\mathrm{lo}})} +\quad\text{(log)} +\tag{9} +$$ + +This inverse is what holds a non-emitted parameter at a prescribed constant: +`src/training/forward.rs::fixed_output_normalized` maps `params.defaults[name]` +through (9) so that (6)/(7) maps it back exactly. + +### 2.4 The parameter boxes + +`src/config.rs::ParameterRanges::default` and +`src/config.rs::Params::default`: + +| parameter | box | log-space by default | units | status | +|---|---|---|---|---| +| `n` (Manning $n_0$) | $[0.015,\,0.25]$ | no | $\mathrm{s\,m^{-1/3}}$ | learned | +| `q_spatial` ($q$, width **exponent**) | $[0,\,1]$ | no | dimensionless | learned or fixed | +| `p_spatial` ($p$, width **coefficient**) | $[1,\,200]$ | **yes** | $\mathrm{m}^{1-q}$ | learned or fixed (21.0) | +| `x_storage` ($X$) | $[0,\,0.5]$ | no | dimensionless | see §12.2 | +| `K_D` | $[10^{-8},\,10^{-6}]$ default; $[10^{-8},\,10^{-4}]$ in the current arm | no by default; **yes** in the current arm | $\mathrm{s^{-1}}$ | learned (leakance only) | +| `d_gw` | $[-2,\,2]$ | no | m | learned (leakance only) | +| `leakance_factor` ($f$) | $[0,\,1]$ | no | dimensionless | learned (leakance only) | +| `gamma` ($\gamma$) | $[0,\,0.5]$ | no | dimensionless | learned, or global config constant | + +`params.log_space_parameters` defaults to `["p_spatial"]` only. Because +$\theta_{\mathrm{lo}} = 1$ there, (7) reduces to $p = 200^{u}$, i.e. $u$ is +$\log_{200} p$. + +`params.defaults` defaults to `{p_spatial: 21.0}`. The current experiment arm +(`config/experiments/sr_n0_leakance_gated.yaml`) learns +`[n, K_D, d_gw, leakance_factor]` and fixes `p_spatial = 21.0`, +`q_spatial = 0.65`, with `log_space_parameters: [p_spatial, K_D]`. + +**What each parameter can and cannot represent.** +- $n_0$ absorbs everything that scales travel time. Since $K = L/c$ and + $c \propto 1/n$, any smooth systematic bias in the celerity (§5, and the + 2.3% beta approximation) is absorbed almost entirely by $n_0$. A learned + $n_0$ is therefore an *effective* travel-time parameter, not a + measurable channel roughness, unless the celerity convention is + independently validated. +- $(p, q)$ are the at-a-station Leopold and Maddock width relation + $T = p\,d^{q}$. They are falsifiable against remotely sensed river width, + which is the stated reason for learning them as a pair. +- $X$ sets the split between translation and attenuation. Under the default + corrected physics it is **not** learned; it is derived from the flow (§6.3). +- $K_D$, $d_gw$, $f$ are the GW-SW exchange triple (§8), with the identifiability + caveat in §8.6. + +--- + +## 3. The leakance gate + +`src/training/gate.rs::leakance_gate`. Applied to the **normalized** head +output $u$ for `leakance_factor` *before* denormalization, at every reader: +training (`src/training/forward.rs::forward`), eval +(`src/training/forward.rs`, `src/training/eval.rs`), the probe +(`src/training/probe.rs`) and `src/dump_parameters.rs`. + +$$ +g \;=\; \sigma\!\left(\frac{\mathrm{logit}\big(\mathrm{clamp}(u,\,\varepsilon_g,\,1-\varepsilon_g)\big)}{\tau}\right), +\qquad \mathrm{logit}(x) = \ln x - \ln(1-x) +\tag{10} +$$ + +with $\varepsilon_g = 10^{-6}$ (`src/training/gate.rs::GATE_EPS`, a Rust +`const`, not configurable). + +**$\tau = 1$ is the identity.** $\sigma \circ \mathrm{logit} = \mathrm{id}$, so +$g = u$ analytically. The code special-cases `tau == 1.0` and returns `u` +untouched, so the identity is *bit*-exact rather than round-tripped through +`ln`/`exp`. This is what makes a $\tau$ schedule starting at 1.0 begin from the +historical ungated model exactly. + +As $\tau \downarrow 0$, (10) sharpens toward a step at $u = 0.5$. It is +differentiable for every $\tau > 0$ (no straight-through estimator): + +$$ +\frac{\partial g}{\partial u} \;=\; \frac{g(1-g)}{\tau\,u\,(1-u)} +\quad\text{inside the clamp}, \qquad 0 \ \text{ outside} +\tag{11} +$$ + +**Epsilon justification, from the `GATE_EPS` docstring.** The head ends in a +sigmoid, so $u$ can be exactly $0.0$ or $1.0$ in f32 once the pre-activation +passes about $\pm 17$ ($\sigma(17) = 1.0$ in f32). An unclamped logit there is +$\pm\infty$ and its backward $1/(u(1-u))$ is $\infty$, giving NaN gradients. +$10^{-6}$ is chosen because $1-10^{-6}$ is about 17 ulps below $1.0$ in f32 +(ulp $= 2^{-24} \approx 6\times10^{-8}$), so $1-u$ resolves to about 3% +relative accuracy; $10^{-7}$ would sit 2 ulps away and carry about 20% error. +$\mathrm{logit}(10^{-6}) \approx -13.8$, well inside f32 range, and the clamp +only alters values the head has already saturated past +$|z| > 13.8$. No f64 is introduced (invariant 2). + +**Temperature schedule.** `src/config.rs::LeakanceGate` holds a +`BTreeMap` keyed by 1-indexed epoch. + +$$ +\tau(\text{epoch}) = \text{value at the largest key} \le \text{epoch}; +\quad\text{if no such key, the first value} +\tag{12} +$$ + +$$ +\tau_{\text{final}} = \text{the last scheduled value} +\tag{13} +$$ + +`LeakanceGate::resolve(epoch)` implements (12) and is what the training driver +passes per epoch (`src/training/driver.rs`, alongside the learning rate). +`LeakanceGate::final_temperature` implements (13) and is what **eval, the +probe, and `dump_parameters` use**, so the scored and reported model is the one +training ended on. Config load (`src/config.rs::validate_leakance_gate`) +requires `use_leakance: true`, a non-empty schedule, and every temperature +positive and finite. + +The current arm anneals `{1: 1.0, 11: 0.5, 21: 0.25, 31: 0.1, 41: 0.05}`. + +**Why the gate exists** (module docstring of `src/training/gate.rs`): +$\zeta$ (equation 45) multiplies $f$ and $K_D$, so only their product reaches +the physics and the two are exactly degenerate; the measured correlation over +CONUS on a trained arm was $\rho = +0.9986$. Physically a reach either sits +above a losing aquifer or it does not, so $f$ should be a $0/1$ **selector** of +*where* leakance acts rather than a continuous multiplier of *how much*. A gate +and a conductance are not degenerate with each other. + +| symbol | meaning | units | status | +|---|---|---|---| +| $u$ | normalized `leakance_factor` head output | dimensionless | learned | +| $g$ | gated factor, still normalized, fed to (6) | dimensionless | derived | +| $\tau$ | gate temperature for this epoch | dimensionless | config schedule | +| $\varepsilon_g$ | clamp margin, $10^{-6}$ | dimensionless | literal | + +--- + +## 4. Channel geometry + +Implemented twice, identically: `src/geometry.rs::compute_trapezoidal_geometry_gamma` +(the reference, autograd-by-BURN) and steps S1-S16 of +`src/routing/mmc_op.rs::forward_chain_inner` (the production path, which runs +at backend-primitive level with the hand-written backward of §9). The equations +below are read from `forward_chain_inner`; `src/geometry.rs` agrees term for +term. + +### 4.1 The width relation, stated unambiguously + +$$ +T \;=\; p \cdot d^{\,q_\varepsilon} +\tag{14} +$$ + +**$p$ is the coefficient. $q$ is the exponent.** $p$ is `p_spatial`, box +$[1, 200]$, log-space; $q$ is `q_spatial`, box $[0, 1]$, linear. Dimensionally +$[p] = \mathrm{m}^{1-q}$, which is why $p$ is not directly comparable across +reaches with different $q$. + +### 4.2 The chain, step by step + +$$ +q_\varepsilon = q + 10^{-6} +\tag{S1, 15} +$$ + +$$ +\mathcal{N} = Q_t \cdot n \cdot (q_\varepsilon + 1) \cdot d_{\mathrm{ref}}^{\,\gamma} +\tag{S2, 16} +$$ + +$$ +\mathcal{D} = p\sqrt{s} + 10^{-8} +\tag{S3, 17} +$$ + +$$ +\mathcal{R} = \mathcal{N} / \mathcal{D} +\tag{S4, 18} +$$ + +$$ +e \;=\; \frac{3}{\,5 + 3 q_\varepsilon + 3\gamma\,} +\tag{S5, 19} +$$ + +$$ +d \;=\; \max\!\left(\mathcal{R}^{\,e},\; d_{\mathrm{lb}}\right) +\tag{S6, 20} +$$ + +$$ +T = p\,d^{\,q_\varepsilon} +\tag{S7, 21} +$$ + +$$ +z_{\mathrm{raw}} = \frac{T\,q_\varepsilon}{2d}, +\qquad +z = \mathrm{clamp}(z_{\mathrm{raw}},\, 0.5,\, 50) +\tag{S8-S9, 22} +$$ + +$$ +b_{\mathrm{raw}} = T - 2 z d, +\qquad +b = \max(b_{\mathrm{raw}},\, b_{\mathrm{lb}}) +\tag{S10-S11, 23} +$$ + +$$ +A = \frac{(T + b)\,d}{2} +\tag{S12, 24} +$$ + +$$ +P_w = b + 2 d\sqrt{1 + z^2} +\tag{S13, 25} +$$ + +$$ +R = A / P_w +\tag{S14, 26} +$$ + +$$ +v_{\mathrm{un}} = \frac{1}{n}\left(\frac{d}{d_{\mathrm{ref}}}\right)^{\!\gamma} R^{2/3}\sqrt{s} +\tag{S15, 27} +$$ + +$$ +v = \mathrm{clamp}(v_{\mathrm{un}},\, v_{\mathrm{lb}},\, 15) +\tag{S16, 28} +$$ + +At $\gamma = 0$, equations 16, 19 and 27 reduce to +$\mathcal{N} = Q_t n (q_\varepsilon+1)$, $e = 3/(5+3q_\varepsilon)$ and +$v_{\mathrm{un}} = n^{-1} R^{2/3}\sqrt{s}$ exactly; the code branches on +`gamma != 0.0` rather than relying on `powf(0.0) == 1.0`, so the historical +path is bit-identical. + +| symbol | meaning | units | status | +|---|---|---|---| +| $Q_t$ | discharge at the start of the step | m³/s | derived (state) | +| $n$ | Manning $n_0$ | s m$^{-1/3}$ | learned | +| $p$ | L&M width coefficient | m$^{1-q}$ | learned or config default | +| $q$ | L&M width exponent | - | learned or config default | +| $s$ | reach slope, floored at `attribute_minimums.slope` | m/m | data (adjacency zarr) | +| $L$ | reach length | m | data (adjacency zarr) | +| $d$ | flow depth | m | derived | +| $T$ | top width | m | derived | +| $z$ | side slope, horizontal per vertical | - | derived | +| $b$ | bottom width | m | derived | +| $A$ | cross-sectional area | m² | derived | +| $P_w$ | wetted perimeter | m | derived | +| $R$ | hydraulic radius | m | derived | +| $v$ | Manning velocity | m/s | derived | +| $\gamma, d_{\mathrm{ref}}$ | stage-roughness pair (§5) | -, m | learned / config | + +### 4.3 Where the depth closed form comes from, and what it assumes + +Equation 20 is *not* an inversion of the trapezoid that equations 24-26 then +build. Substituting $R \approx d$ (wide-channel limit) and the power-law area +$A_{\mathrm{pl}} = \int_0^d T\,\mathrm{d}d' = p\,d^{q+1}/(q+1) = T d/(q+1)$ into +$Q = vA$ with $v = n^{-1}R^{2/3}\sqrt{s}$ gives + +$$ +Q = \frac{p\sqrt{s}}{n\,(q+1)}\, d^{\,q + 5/3} +\;\Longrightarrow\; +d = \left[\frac{Q\,n\,(q+1)}{p\sqrt{s}}\right]^{3/(5+3q)} +\tag{29} +$$ + +which is exactly equations 16-20 at $\gamma = 0$. So the exponent $3/(5+3q)$ +encodes the **power-law section with $R = d$**, while $A$, $P_w$, $R$, $v$ and +the celerity are then evaluated on the **trapezoid**. The two sections do not +agree. Substituting $z = Tq/(2d)$ into equation 23 gives $b = T(1-q)$, hence + +$$ +A_{\mathrm{tz}} = \frac{T(2-q)\,d}{2}, \qquad +\frac{A_{\mathrm{tz}}}{A_{\mathrm{pl}}} = \frac{(2-q)(q+1)}{2} +\tag{30} +$$ + +which equals 1 only at $q = 0$ and $q = 1$, and is 1.114 at $q = 0.65$ (the +value the current arm prescribes). This is the same $(2-q)(q+1)/2$ factor +recorded in `tests/stage_roughness.rs` as the "up to 11%" gap between the true +$\mathrm{d}A/\mathrm{d}d$ and the $T$ the celerity convention uses. It is +inherited from DDR and out of scope of any fix in this tree, but it means the +depth reported by the model is the depth of a power-law channel, not of the +trapezoid whose conveyance is then computed. + +### 4.4 Every clamp, and what it does to the gradient + +Each row names the clamp, its bound, the config key, and the backward label +(§9) that masks it. A saturated clamp zeroes the gradient for that +reach-timestep through that path. + +| # | clamp | bound | source | backward | gradient when saturated | +|---|---|---|---|---|---| +| C1 | slope floor | $s \ge 10^{-3}$ | `attribute_minimums.slope`, applied once in `mmc.rs::setup_inputs` | n/a | $s$ is not a parameter; no gradient path | +| C2 | lateral inflow floor | $q' \ge 10^{-4}$ | `attribute_minimums.discharge`, once per window in `mmc.rs::forward` | n/a | $q'$ is data (a parent only in the adjoint study) | +| C3 | depth floor | $d \ge 10^{-2}$ m | `attribute_minimums.depth` | B6 | zero into $\mathcal{R}$, $e$, hence into $n$, $q$, $p$, $Q_t$ through the depth chain | +| C4 | side-slope band | $z \in [0.5, 50]$ | **literal**, not configurable | B9 | zero into $T$, $q_\varepsilon$, $d$ through $z$ | +| C5 | bottom-width floor | $b \ge 10^{-2}$ m | `attribute_minimums.bottom_width` | B11 | zero into $T$, $z$, $d$ through $b$ | +| C6 | velocity band | $v \in [0.01, 15]$ m/s | lower from `attribute_minimums.velocity`, **upper literal 15** | B16 | zero into all of $R$, $n$, $d$ through the velocity | +| C7 | Muskingum $K$ floor | $K \ge \Delta t (1+\delta)/2$ | `enforce_positivity` only, $\delta = 10^{-2}$ literal | B18' | zero into celerity, hence into all geometry | +| C8 | Cunge $X$ band | $X \in [0, 0.5]$ | literal | B19 | zero into $Q_t$, $T$, $c$ through $X$ | +| C9 | $X$ stability cap | $X \le \min(\tfrac{Cr}{2}, 1-\tfrac{Cr}{2})(1-\delta)$ | `enforce_positivity` only | B19' | routes to one branch only; the Cunge terms are masked off where a cap wins | +| C10 | solve output floor | $Q_{t+1} \ge 10^{-4}$ | `attribute_minimums.discharge` | B28 | zero into everything for that reach-timestep | +| C11 | gate clamp | $u \in [10^{-6}, 1-10^{-6}]$ | literal `GATE_EPS` | autograd | zero into the head's `leakance_factor` logit | +| C12 | losing-only clamp | $\max(0, d - d_{gw})$ | `leakance_losing_only`, default **true** | `zeta_backward` gate | zero into *all six* leakance partials on gaining reaches | +| C13 | KGE per-gauge cap | weighted component sum $\le 10$ | `loss.kge_clamp` | autograd | zero for that gauge's KGE term | + +C10 is the consequential one for mass: see §6.5. + +--- + +## 5. Stage-dependent roughness + +`params.stage_roughness` (a global constant $\gamma$) or `gamma` in +`kan_head.learnable_parameters` (per-reach learned). Config load +(`src/config.rs::validate_stage_roughness`, `::validate_learned_gamma`) rejects +setting both, requires $\gamma \in [0,1]$, $d_{\mathrm{ref}} > 0$, +`ddr_match: false`, `use_cuda_graphs: false`, and, for the learned form, +`use_leakance: false` (the eight-parent leakance op has no $\gamma$ parent, so +$\gamma$ would silently receive no gradient). With a *global* $\gamma$, +$d_{\mathrm{ref}}$ must be exactly 1.0, because $d_{\mathrm{ref}}^{\gamma}$ is +then a constant on the Manning numerator and is fully absorbed by the learned +$n$ field. + +### 5.1 The roughness law + +$$ +n(d) \;=\; n_0 \left(\frac{d}{d_{\mathrm{ref}}}\right)^{\!-\gamma} +\tag{31} +$$ + +$\gamma \ge 0$ means channels get **smoother** as they fill. $n_0$ is the +roughness **at** $d = d_{\mathrm{ref}}$, so changing $d_{\mathrm{ref}}$ +rescales $n_0$ and makes published $n$ values incomparable +(`src/config.rs::StageRoughnessSection` docstring). + +$d_{\mathrm{ref}}$ comes from `params.stage_roughness.d_ref`, default +`default_d_ref()`; it is **required to be 1.0 m** whenever $\gamma$ is global, +and only free when $\gamma$ is a per-reach KAN output. So in every run to date, +$d_{\mathrm{ref}} = 1\ \mathrm{m}$ and $d_{\mathrm{ref}}^{\gamma} = 1$ exactly. + +### 5.2 The modified depth exponent + +Substituting (31) into the Manning inversion of §4.3 moves one power of depth +across, shifting the exponent denominator and putting a constant +$d_{\mathrm{ref}}^{\gamma}$ on the numerator: + +$$ +e \;=\; \frac{3}{\,5 + 3q_\varepsilon + 3\gamma\,}, +\qquad +d = \left[\frac{Q_t\,n_0\,(q_\varepsilon+1)\,d_{\mathrm{ref}}^{\gamma}}{p\sqrt{s} + 10^{-8}}\right]^{e} +\tag{32} +$$ + +The velocity must carry the same factor, because +$1/n(d) = (1/n_0)(d/d_{\mathrm{ref}})^{\gamma}$: + +$$ +v_{\mathrm{un}} = \frac{1}{n_0}\left(\frac{d}{d_{\mathrm{ref}}}\right)^{\!\gamma} R^{2/3}\sqrt{s} +\tag{33} +$$ + +The `src/geometry.rs` comment records why this is called out explicitly: +applying the factor to only the depth inversion and not the velocity "leaves +the velocity exponent at Manning's 2/3 while the depth exponent moves, which is +silently wrong rather than loudly wrong", and `tests/stage_roughness.rs` caught +exactly that on the first implementation. + +### 5.3 The celerity and its correction term + +`ddr_match: true` (deprecated): + +$$ +c = \tfrac{5}{3}\,v +\tag{S17a, 34} +$$ + +This is the wide-rectangular Kleitz-Seddon limit. The in-code comment states it +is WRONG for the trapezoid S7-S13 builds: $\kappa = b/y \approx 0.7$-$1.8$ +there, so the true ratio is $\approx 1.30$-$1.36$, making 5/3 some 22-27% high. + +`ddr_match: false` (the default since 2026-08-19), from +$c = \mathrm{d}Q/\mathrm{d}A$ with $\mathrm{d}A/\mathrm{d}d = T$: + +$$ +c \;=\; v \cdot \beta, +\qquad +\beta \;=\; \underbrace{\frac{5}{3} - \frac{4}{3}\,\frac{A\sqrt{1+z^2}}{T\,P_w}}_{\beta_{\text{trapezoid}}} +\;+\; \underbrace{\gamma\,\frac{A}{T\,d}}_{\text{stage-roughness term}} +\tag{S17b, 35} +$$ + +$\beta_{\text{trapezoid}} \to 5/3$ as $b/y \to \infty$ and $\to 4/3$ as +$b \to 0$. The **trapezoid factor** is $\sqrt{1+z^2}$, which is +$\mathrm{d}P_w/\mathrm{d}d$ divided by 2; it is what makes $\beta$ depend on +the side slope at all. Sanity check recorded in the code: for a pure power-law +section $A/(Td) = 1/(q+1)$, so + +$$ +\beta \;=\; \frac{5 + 3q + 3\gamma}{3(q+1)} +\tag{36} +$$ + +which is $5/3$ at $q = \gamma = 0$, and is the reciprocal of $3e/(q+1)$ with +$e$ from (19). + +### 5.4 The documented approximation in the celerity convention, and its size + +Two approximations, both pinned by tests. + +**(a) The stage-roughness increment is exact.** +`tests/stage_roughness.rs::stage_roughness_celerity_increment_is_exact` +verifies, against central differences at $q \in \{0.084, 0.65\}$, +$d \in \{0.25, 1, 3\}$ m, $\gamma \in \{0.1, 0.183, 0.4\}$: + +$$ +c(d,\gamma) \;=\; \left(\frac{d}{d_{\mathrm{ref}}}\right)^{\!\gamma} c(d, 0) +\;+\; \gamma\, v(d,\gamma)\,\frac{A}{T\,d} +\tag{37} +$$ + +to relative tolerance $10^{-4}$, *under the convention the solver uses*. + +**(b) The convention itself is approximate, by up to 2.3%.** The solver +computes $c = (\mathrm{d}Q/\mathrm{d}d)/T$, i.e. it assumes +$\mathrm{d}A/\mathrm{d}d = T$ and $\mathrm{d}P_w/\mathrm{d}d = 2\sqrt{1+z^2}$, +i.e. a trapezoid of **fixed shape** being filled. This geometry reshapes as it +fills: $b = T(1-q)$ and $z = (pq/2)\,d^{\,q-1}$ both move with depth. The exact +derivatives, from +`tests/stage_roughness.rs::documents_the_preexisting_beta_approximation`: + +$$ +T' = \frac{qT}{d}, \qquad +z' = \frac{z(q-1)}{d}, \qquad +b' = T' - 2\left(z + d z'\right) +\tag{38} +$$ + +$$ +A' = \frac{(T' + b')\,d + (T + b)}{2}, +\qquad +P_w' = b' + 2\sqrt{1+z^2} + 2d\,\frac{z}{\sqrt{1+z^2}}\,z' +\tag{39} +$$ + +$$ +\beta_{\mathrm{exact}} \;=\; \frac{5}{3} - \frac{2}{3}\,\frac{A}{P_w}\,\frac{P_w'}{A'} +\tag{40} +$$ + +(The $z'$ and $b'$ terms are set to zero where the C4 / C5 clamps bind, so the +measurement respects the same clamp masks the code applies.) + +**Measured size, from the test's own assertions and docstring:** the worst +relative gap over the grid +$q \in \{0.05, 0.084, 0.2, 0.35, 0.5, 0.65, 0.85, 1.0\}$, +$d \in \{0.1, 0.25, 1, 3\}$ m was **about 2.3%**, and about **1.2% at the +trained $q \approx 0.084$**. The test asserts `worst < 0.05` (a 5% ceiling, so +a regression is caught) and `worst > 1e-4` (so the test cannot pass vacuously +if $\beta$ is ever made exact). The true $\mathrm{d}A/\mathrm{d}d$ is +$T(2-q)(q+1)/2$, up to **11%** away from the $T$ the convention uses. + +**Why it was documented rather than fixed**, verbatim from the test: the gap +changes sign with depth so it partly averages out, and since $K = L/c$ with +$c \propto 1/n$, a smooth systematic celerity bias is absorbed almost entirely +by the learned roughness. **Only the $q$- and depth-dependent structure of the +error survives training.** This is the single most important caveat for +interpreting a learned $n_0$ as a physical roughness. + +--- + +## 6. Muskingum-Cunge routing + +`src/routing/mmc_op.rs::forward_chain_inner` steps S18-S28. A second, +API-level copy of the coefficients lives at +`src/routing/mmc.rs::MuskingumCunge::calculate_muskingum_coefficients`; it is +term-for-term identical but is **not** on the production per-timestep path. + +### 6.1 Storage time, Courant number, cell Reynolds number + +$$ +K \;=\; \frac{L}{c} +\tag{S18, 41} +$$ + +$$ +Cr \;=\; \frac{\Delta t}{K} \;=\; \frac{c\,\Delta t}{L} +\tag{42} +$$ + +$$ +D_{\mathrm{cell}} \;=\; \frac{Q_t}{T\,s\,c\,L} +\;=\; \frac{Q_t/T}{s\,c\,L} +\tag{43} +$$ + +Equation 42 is the Courant number: cells are reaches, $\Delta x = L$. +Equation 43 is the cell (grid) Reynolds number of the Muskingum-Cunge +literature, $D = q_{\text{unit}}/(s_0 c \Delta x)$ with $q_{\text{unit}} = Q/T$ +the unit-width discharge. In code it appears as the local `w` inside S19, with +a $+10^{-12}$ guard on the denominator: +`w = q_t / (top_width * slope * celerity * length + 1e-12)`. + +| symbol | meaning | units | status | +|---|---|---|---| +| $K$ | Muskingum storage time | s | derived | +| $\Delta t$ | routing step, $3600$ | s | literal (`DT_SECONDS`) | +| $Cr$ | Courant number | - | derived | +| $D_{\mathrm{cell}}$ | cell Reynolds number | - | derived | + +Measured context from `config/experiments/sr_n0_leakance_gated.yaml`: the +median CONUS $Cr$ is **0.226**, i.e. the typical MERIT reach is about 4.4x too +long for an hourly step. + +### 6.2 The four coefficients + +$$ +\mathrm{denom} \;=\; 2K(1-X) + \Delta t +\tag{S20-S23, 44} +$$ + +$$ +c_1 = \frac{-2KX + \Delta t}{\mathrm{denom}}, \quad +c_2 = \frac{2KX + \Delta t}{\mathrm{denom}}, \quad +c_3 = \frac{2K(1-X) - \Delta t}{\mathrm{denom}}, \quad +c_4 = \frac{2\Delta t}{\mathrm{denom}} +\tag{45} +$$ + +All four share the single denominator (44). The identity + +$$ +c_1 + c_2 + c_3 \;=\; \frac{\Delta t + 2K(1-X)}{\mathrm{denom}} \;=\; 1 +\tag{46} +$$ + +holds for **any** $(K, X)$, exactly (up to f32 round-off). $c_2, c_4 > 0$ +unconditionally; $c_1 \ge 0 \iff X \le Cr/2$ and $c_3 \ge 0 \iff X \le 1 - Cr/2$, +which is the window + +$$ +2X \;\le\; Cr \;\le\; 2(1-X) +\tag{47} +$$ + +### 6.3 The storage parameter $X$ + +`ddr_match: true` (deprecated): $X$ is the caller's constant. Every non-learned +path sets it to **0.3** (`src/training/forward.rs`, `Tensor::full([n], 0.3)`). +The in-code comment records that this severs the link between the scheme's +numerical diffusion and the channel's physical hydraulic diffusivity, giving a +median **28x over-diffusion** on CONUS, and is exact only on the measure-zero +locus $D_{\mathrm{cell}} = 0.4$. + +`ddr_match: false` (default), Cunge: + +$$ +X_{\mathrm{cunge}} \;=\; \mathrm{clamp}\!\left(\tfrac{1}{2}\left(1 - D_{\mathrm{cell}}\right),\, 0,\, 0.5\right) +\tag{S19, 48} +$$ + +This is the choice that matches numerical to physical diffusivity: + +$$ +D_{\mathrm{num}} = c L\left(\tfrac12 - X\right) \;=\; \frac{Q_t}{2\,T\,s} \;=\; D_{\mathrm{phys}} +\tag{49} +$$ + +Measured $X \approx 0.49$ on CONUS, which narrows the non-negativity window +(47) from $[0.6, 1.4]$ (at $X = 0.3$) to about $[0.98, 1.02]$. + +### 6.4 `enforce_positivity`: what it does and when it fires + +`params.enforce_positivity`, default **false**, requires `ddr_match: false`. +With $\delta =$ `POSITIVITY_DELTA` $= 10^{-2}$ (a Rust `const` in +`src/routing/mmc_op.rs`): + +$$ +K \;\leftarrow\; \max\!\left(\frac{L}{c},\; \frac{\Delta t\,(1+\delta)}{2}\right) +\tag{S18', 50} +$$ + +$$ +X \;\leftarrow\; \min\!\Big(X_{\mathrm{cunge}},\; +\underbrace{\tfrac{Cr}{2}(1-\delta)}_{\text{hi}_a},\; +\underbrace{\left(1 - \tfrac{Cr}{2}\right)(1-\delta)}_{\text{hi}_b}\Big) +\tag{S19', 51} +$$ + +(50) puts $Cr \in (0,\, 2/(1+\delta)]$, which is what guarantees +$\mathrm{hi}_b > 0$ so no extra `clamp_min` is needed on the three-way `min`. +(51) forces $c_1, c_3 \ge 0$. With $b \ge 0$ and forward substitution +$x_i = b_i + c_{1,i}\sum_{\text{up}} x_j$ in topological order, induction gives +every $x_i \ge 0$, so no negative solve can reach S28. + +**$\delta$ is mandatory, not cosmetic.** At $\delta = 0$ the cap lands exactly +on $c_1 = 0$ / $c_3 = 0$ and f32 round-off crosses it: measured minima +$c_1 = -3.3\times10^{-8}$, $c_3 = -6.8\times10^{-8}$ over a 400k-draw sweep. +$\delta = 10^{-2}$ (about 400x f32 eps) moves those minima to +$+1.8\times10^{-4}$ and $+5.0\times10^{-5}$, costing a 1% tightening of the $X$ +ceiling and a 1% rise in the $K$ floor. + +**When it fires and what it costs.** From the current config's own commentary: +the $X$ cap binds on **95.3%** of reach-timesteps and median $X$ falls +$0.4976 \to 0.0794$ (6.3x). It therefore largely *replaces* the Cunge $X$ +rather than shading it: numerical diffusion becomes stability-set instead of +matched to hydraulic diffusivity. It delivered 0 negative solves out of +197,461,880, but because the cap is applied at runtime to a learned celerity it +made $X \sim Cr \sim 1/n$, handing the optimizer a lever it rode to the +roughness floor: at the same learning rate, `n_at_floor` went +$35\% \to 97.9\%$ with the flag on. **It is off in the current arm.** + +### 6.5 Mass conservation: what is and is not conserved + +**No coefficient is ever clamped.** A grep over `src/routing/` confirms there +is no `c1.clamp`, `c3.clamp`, or equivalent. This is deliberate: because (46) +holds for any $(K, X)$, clamping the **inputs** $K$ and $X$ preserves +$c_1+c_2+c_3 = 1$ exactly, whereas clamping $c_3$ would not. The code comment +on `params.enforce_positivity` states this reasoning explicitly. + +**What does break mass is C10**, the post-solve floor at S28: + +$$ +Q_{t+1} \;=\; \max\!\left(x_{\mathrm{sol}},\; 10^{-4}\right) +\tag{S28, 52} +$$ + +applied in `src/routing/mmc_op.rs::forward_chain_inner`. The +`NEG_SOLVES`/`TOTAL_SOLVES` docstring in that file states the case plainly: +measured on CONUS at mean flow with $X = 0.3$, **69.8%** of reaches sit outside +window (47) (28.4% give $c_1 < 0$ and 41.4% give $c_3 < 0$), so negative +discharge is *expected*, and the clamp "both CREATES MASS and removes the only +symptom". Nothing measured this before 2026-08-02. The counters are diagnostic +only (host readback gated on `track_negative_discharge`, enabled in training, +off in eval, and structurally unavailable on the CUDA-graph path). + +So the accurate statement for the analysis that follows this document is: +*ddrs does not clamp Muskingum coefficients; it clamps $K$ and $X$ (optionally) +and the solve output (always). The first preserves $c_1+c_2+c_3=1$; the second +creates mass whenever the solve goes negative, at a measured rate of +$5.5\times10^{4}$ per $2\times10^{8}$ solves on a trained head, or 69.8% of +reaches by window-violation count at $X = 0.3$.* + +### 6.6 The per-timestep linear system + +$$ +\left(I - c_1 \odot N\right) Q_{t+1} +\;=\; +c_2 \odot (N Q_t) \;+\; c_3 \odot Q_t \;+\; c_4 \odot q'_t \;-\; \zeta_t +\tag{S24-S27, 53} +$$ + +with $\odot$ a row-wise (per-reach) scaling, $N$ the sparse adjacency +(`indices_0`/`indices_1` from the topologically ordered adjacency store, +$N_{ij} = 1$ when $j$ flows directly into $i$), and $\zeta_t$ present only when +leakance is active (§8). Written out: + +$$ +Q_{t+1,i} \;=\; c_{1,i}\!\!\sum_{j \in \mathrm{up}(i)}\!\! Q_{t+1,j} +\;+\; c_{2,i}\!\!\sum_{j \in \mathrm{up}(i)}\!\! Q_{t,j} +\;+\; c_{3,i} Q_{t,i} \;+\; c_{4,i} q'_{t,i} \;-\; \zeta_{t,i} +\tag{54} +$$ + +This is the standard Muskingum update $O_{j+1} = C_1 I_{j+1} + C_2 I_j + C_3 O_j$ +with the inflow $I$ taken as the sum of immediate upstream outflows, which is +what makes the $t+1$ term implicit and requires the solve. + +Intermediates, in code order: +$i_t = N Q_t$ (S24, sparse matrix-vector), +$b = c_2 i_t + c_3 Q_t + c_4 q'_t - \zeta$ (S25), +$A_{\text{values}}$ assembled from $c_1$ (S26), +$x_{\mathrm{sol}}$ from the triangular solve (S27), +$Q_{t+1}$ from (52) (S28). + +### 6.7 Cold start + +`src/routing/utils.rs::compute_hotstart_discharge` / the equivalent CSR path in +`mmc.rs::setup_inputs` with $c \equiv 1$: + +$$ +(I - N)\,Q_0 = q'_0, +\qquad +Q_0 \leftarrow \max(Q_0,\, 10^{-4}) +\tag{55} +$$ + +On a topologically ordered DAG, $I - N$ is lower triangular with unit diagonal, +so this is the exact upstream accumulation of lateral inflow (a cumulative sum +on a linear chain). An externally supplied `initial_state` replaces it when a +state cache is configured. + +### 6.8 Sub-reach subdivision (off by default) + +When `parent_offset` marks a parent reach as split into $m_p$ pieces, the +lateral inflow is divided so total $q'$ per parent is conserved: + +$$ +q'_{\text{row}} = \frac{q'_{\text{parent}}}{m_p} +\tag{56} +$$ + +applied **after** the C2 clamp (`mmc.rs::forward`), because dividing first would +floor each of the $m$ pieces independently and inject $m \cdot 10^{-4}$ instead +of $10^{-4}$. The same divisor is applied to the cold start by default +(`divide_hotstart_by_pieces: true`). Subdivision is a NO-GO and gated off; see +`.claude/REACH-SUBDIVISION.md`. + +--- + +## 7. The sparse solve and its adjoint + +`src/sparse/mod.rs`. `CsrPattern` is built once per network at `setup_inputs` +and reused for every timestep (`Arc`, so per-step cost is a refcount bump). +$A$ is lower triangular because the adjacency is topologically ordered with +$\text{rows}[k] \ge \text{cols}[k]$ (invariant 3). + +### 7.1 Assembly + +`src/sparse/mod.rs::assemble_primitive`, for each non-zero slot $k$ with +row $r(k)$: + +$$ +A_{\text{values}}[k] \;=\; \mathrm{diag}[k] \;-\; c_{1}[r(k)] \cdot \mathrm{adj}[k] +\tag{57} +$$ + +where `diag[k] = 1` at diagonal slots and 0 elsewhere, and `adj[k] = 0` at +diagonal slots. So $A_{ii} = 1$ and $A_{ij} = -c_{1,i} N_{ij}$ for $j \ne i$. + +### 7.2 Forward substitution + +`src/sparse/mod.rs::forward_sub_lower` (a pure-Rust port of +`scipy.sparse.linalg.spsolve_triangular(A, b, lower=True)`): + +$$ +x_i \;=\; \frac{b_i - \sum_{k \in \mathrm{row}(i),\, \mathrm{col}(k) \ne i} A_{\text{values}}[k]\; x_{\mathrm{col}(k)}}{A_{ii}} +\tag{58} +$$ + +evaluated for $i = 0, 1, \ldots, n-1$ in topological order. The CUDA path +(`params.sparse_solver: cuda`) substitutes `cusparseSpSV`; the two are bit-matched +by `tests/cuda_backward_parity.rs`. + +### 7.3 The hand-written adjoint + +`src/sparse/mod.rs::CsrSolveOp impl Backward` with parents +`[a_values, b]`. From $Ax = b$: + +$$ +\mathrm{d}x = A^{-1}\left(\mathrm{d}b - \mathrm{d}A\, x\right) +\tag{59} +$$ + +so with $\bar{x} \equiv \partial L/\partial x$ the incoming gradient: + +$$ +\bar{b} \;=\; A^{-\mathsf{T}}\,\bar{x} +\tag{60} +$$ + +$$ +\bar{A}_{ij} \;=\; -\,\bar{b}_i\, x_j +\qquad\text{i.e.}\qquad +\bar{A}_{\text{values}}[k] \;=\; -\,\bar{b}\big[r(k)\big]\cdot x\big[\mathrm{col}(k)\big] +\tag{61} +$$ + +Equation 60 is evaluated by back-substitution on $A^{\mathsf{T}}$, which is +upper triangular, using the cached transposed pattern +(`src/sparse/mod.rs::back_sub_upper_transposed`; `trans_crow`, `trans_col`, +`trans_to_orig` let it read the original `A_values` without rebuilding +structure): + +$$ +y_i \;=\; \frac{\bar{x}_i - \sum_{k \in \mathrm{row}^{\mathsf{T}}(i),\, \mathrm{col}^{\mathsf{T}}(k) \ne i} A_{\text{values}}\big[\mathrm{trans\_to\_orig}[k]\big]\, y_{\mathrm{col}^{\mathsf{T}}(k)}}{A^{\mathsf{T}}_{ii}}, +\quad i = n-1 \ldots 0 +\tag{62} +$$ + +Equation 61 is evaluated as a per-nnz gather-multiply (host loop on CPU, +`select+mul+neg` on GPU). This is the whole point of invariant 4: **$O(nnz)$ +tape entries per timestep, not $O(n^2)$.** Do not replace it with autograd-tape +unrolling. + +Chaining (61) back through (57): + +$$ +\bar{c}_1[i] \;=\; -\sum_{k:\, r(k) = i} \bar{A}_{\text{values}}[k]\cdot \mathrm{adj}[k] +\tag{63} +$$ + +(`src/sparse/mod.rs::assemble_backward_primitive` / `::cpu_assemble_backward`, +`cusparseSpMV` on the GPU with $\alpha = -1$ folding in the negation), and the +SpMV adjoint: + +$$ +\bar{Q}_t \;=\; N^{\mathsf{T}}\,\bar{i}_t +\qquad\text{i.e.}\qquad +\bar{Q}_t\big[\mathrm{col}(k)\big] \mathrel{+}= \mathrm{adj}[k]\cdot \bar{i}_t\big[r(k)\big] +\tag{64} +$$ + +(`src/sparse/mod.rs::spmv_backward_primitive` / `::cpu_spmv_backward`). + +--- + +## 8. Leakance: the GW-SW exchange term + +`src/routing/leakance.rs`, ported from DDR `_compute_zeta` (`mmc.py`, commit +`c2bd0f9`). `params.use_leakance`, default **false**. Status: CLOSED and NOT +PROMOTABLE (see `CLAUDE.md` and §8.6), but code-complete and gradient-exact, +and enabled in the current experiment arm with the gate of §3. + +### 8.1 The flux + +`src/routing/leakance.rs::zeta_forward`: + +$$ +w_z \;=\; \left(p\,d\right)^{q_\varepsilon} +\tag{65} +$$ + +$$ +\mathcal{A}_z \;=\; w_z \cdot L +\tag{66} +$$ + +$$ +h \;=\; +\begin{cases} +\max\!\left(0,\; d - d_{gw}\right), & \texttt{leakance\_losing\_only = true (default)}\\[2pt] +d - d_{gw}, & \texttt{false (back-compat)} +\end{cases} +\tag{67} +$$ + +$$ +\zeta \;=\; f \cdot \mathcal{A}_z \cdot K_D \cdot h \cdot \mathbb{1}_{\text{mask}} +\tag{68} +$$ + +$$ +b \;\leftarrow\; b - \zeta +\qquad\text{(S25; positive } \zeta \text{ = losing reach)} +\tag{69} +$$ + +$d$ and $q_\varepsilon$ are the **shared** S6 depth and S1 exponent, not +recomputed. `zeta_forward` returns $(w_z, \mathcal{A}_z, \zeta)$; only +$\mathcal{A}_z$ is saved for the backward. + +| symbol | meaning | units | status | +|---|---|---|---| +| $f$ | `leakance_factor`, gated per §3 | dimensionless | learned | +| $K_D$ | leakance = conductivity / streambed thickness | s$^{-1}$ | learned | +| $d_{gw}$ | groundwater depth parameter, box $[-2, 2]$ | m | learned | +| $\mathcal{A}_z$ | plan-view wetted area | see §8.3 | derived | +| $h$ | driving head | m | derived | +| $\zeta$ | exchange flux, positive = losing | m³/s | derived | +| $\mathbb{1}_{\text{mask}}$ | impervious mask, $0$ or $1$ | - | data (constant, no gradient) | + +### 8.2 The impervious mask + +`src/data/dataset.rs::build_mask_from_raw_col`, from the +`corridor_impervious` attribute column with threshold +`params.leakance_impervious_threshold` (default 0.7, i.e. 70% impervious +$\approx$ a concrete-lined channel): + +$$ +\mathbb{1}_{\text{mask}}[i] \;=\; +\begin{cases} +1, & \texttt{corridor\_impervious}[i] \le \text{threshold} \ \ \textbf{or NaN}\\ +0, & \text{otherwise} +\end{cases} +\tag{70} +$$ + +NaN maps to 1 on purpose: no StreamCat coverage means *absence of imperviousness +data*, which is not the same as impervious concrete, so leakance is allowed +where data is missing. The mask is a plain inner-backend tensor, never +autograd-tracked, and it is only built when `use_leakance` is true **and** +`corridor_impervious` is among the configured attributes. Where +$\mathbb{1}_{\text{mask}} = 0$, both $\zeta$ and all six leakance partials are +exactly zero. + +### 8.3 The plan-view area, and a dimensional problem + +Equation 65 is $(p\,d)^{q_\varepsilon}$, **not** $p\,d^{q_\varepsilon}$. It is +therefore not the top width $T$ of equation 21. With $[p] = \mathrm{m}^{1-q}$: + +$$ +[w_z] = \left(\mathrm{m}^{1-q}\cdot\mathrm{m}\right)^{q} = \mathrm{m}^{(2-q)q}, +\qquad +[\mathcal{A}_z] = \mathrm{m}^{(2-q)q + 1} +\tag{71} +$$ + +$[\mathcal{A}_z] = \mathrm{m}^2$ only when $(2-q)q = 1$, i.e. only at $q = 1$. +At the prescribed $q = 0.65$ the exponent is $1.878$, so $\mathcal{A}_z$ has +units $\mathrm{m}^{1.878}$ and $\zeta$ is not in m³/s. The intended quantity +is presumably $T \cdot L = p\,d^{q}L$ (a genuine streambed plan area). This is +inherited from DDR verbatim, is not flagged anywhere in the source, and is +listed in §12 as a discrepancy. **Numerically it means $K_D$ absorbs a +$q$- and $p$-dependent unit conversion, so a learned $K_D$ is not directly +comparable to a measured streambed leakance.** The reported `area_z_mean` +diagnostic is labelled "m^2" in the NetCDF attributes and is not. + +### 8.4 The relation to MODFLOW's river package + +MODFLOW RIV computes $Q_{\text{riv}} = C_{\text{riv}}(h_{\text{riv}} - h_{\text{aq}})$ +with $C_{\text{riv}} = K_v W L / M$ for streambed vertical conductivity $K_v$, +width $W$, length $L$ and thickness $M$. Equation 68 is that form with + +$$ +C \;=\; f \cdot \mathcal{A}_z \cdot K_D, +\qquad +K_D \;\equiv\; \frac{K_v}{M} \quad [\mathrm{s^{-1}}] +\tag{72} +$$ + +so **$K_D$ is leakance proper** (conductivity over streambed thickness, inverse +seconds), $\mathcal{A}_z$ plays the role of $WL$, and $f$ is an extra +dimensionless area-fraction-or-selector that MODFLOW does not have. + +Two departures from MODFLOW worth stating. First, the driving head in +equation 67 is $d - d_{gw}$: a *flow depth* minus a "groundwater depth +parameter" with box $[-2, 2]$ m, not a stage minus an aquifer head. It equals a +true head difference only under the reading that $d_{gw}$ is the water-table +depth below the channel bed. Second, with `leakance_losing_only: true` (the +default) the term is one-directional: gaining reaches produce exactly zero, so +the model cannot represent baseflow gain at all. + +### 8.5 The backward: all six partials + +`src/routing/leakance.rs::zeta_backward`. Because equation 69 subtracts +$\zeta$, the incoming gradient is negated: + +$$ +\bar{\zeta} \;=\; -\bar{b} \cdot \mathbb{1}_{\text{mask}} \cdot \mathbb{1}\!\left[d > d_{gw}\right] +\tag{73} +$$ + +where the $\mathbb{1}[d > d_{gw}]$ gate is applied only when +`losing_only = true`, with the subgradient at the kink $d = d_{gw}$ defined as +**0** (measure-zero, finite-difference safe). Gating $\bar{\zeta}$ once zeroes +all six outputs at a stroke. + +With $m \equiv d - d_{gw}$ (recomputed **unclamped** in the backward, which is +safe precisely because $\bar{\zeta}$ is already gated): + +$$ +\frac{\partial L}{\partial f} \;=\; \bar{\zeta}\,\mathcal{A}_z\,K_D\,m +\tag{74} +$$ + +$$ +\frac{\partial L}{\partial K_D} \;=\; \bar{\zeta}\,f\,\mathcal{A}_z\,m +\tag{75} +$$ + +$$ +\frac{\partial L}{\partial d_{gw}} \;=\; -\,\bar{\zeta}\,f\,\mathcal{A}_z\,K_D +\tag{76} +$$ + +$$ +\frac{\partial L}{\partial p}\bigg|_{\zeta} \;=\; \bar{\zeta}\,f\,K_D\,m\,\mathcal{A}_z\,\frac{q_\varepsilon}{p} +\tag{77} +$$ + +$$ +\frac{\partial L}{\partial q_\varepsilon}\bigg|_{\zeta} \;=\; \bar{\zeta}\,f\,K_D\,m\,\mathcal{A}_z\,\ln\!\left(p\,d\right) +\tag{78} +$$ + +$$ +\frac{\partial L}{\partial d}\bigg|_{\zeta} \;=\; \bar{\zeta}\,f\,K_D\,\mathcal{A}_z +\;+\; \bar{\zeta}\,f\,K_D\,m\,\mathcal{A}_z\,\frac{q_\varepsilon}{d} +\tag{79} +$$ + +Equations 74-76 are registered directly on the three leakance parents. +Equations 77-79 are the *geometry-side* contributions and are folded into the +base operator's existing accumulators (`gp_total`, `gq_spatial`, `gd_total`) +through the `zeta_hook` closure, which is called the moment $\bar{b}$ becomes +available at B27. Verified against central differences by +`src/routing/leakance.rs::grad_tests::zeta_grads_match_central_differences` +(relative tolerance $10^{-3}$ to $10^{-2}$ per partial) and by +`tests/leakance_gradcheck.rs`. + +Derivation check for equation 79: $\zeta = fK_DmL(pd)^{q}$, so +$\partial\zeta/\partial d = fK_DL[(pd)^q + m\,q\,(pd)^{q-1}p] += fK_D[\mathcal{A}_z + m\,\mathcal{A}_z q/d]$, matching term for term. + +### 8.6 The identifiability verdict + +Recorded in `CLAUDE.md` and +`research/findings/2026-07-06-leakance-nogo-scientific-summary.md` §3, and +reproduced here because it directly bears on "are the learned parameters +physically meaningful": **a gauge measures the SUM of $\zeta$ over its upstream +network, and that sum does not determine the per-reach distribution.** Training +therefore constrains the aggregate loss while carrying zero information about +per-reach flux. Every rival explanation (gradient starvation, objective noise, +uninformative inputs, sign ambiguity) was individually refuted. The gate of §3 +is an attempt to change *what* is being identified (a selector rather than a +magnitude), not a refutation of that verdict. + +--- + +## 9. Gradient paths + +`src/routing/mmc_op.rs`. One autograd node per timestep instead of about 33 +BURN ops. The forward chain is labelled S1-S28 (with S18', S19' for the +positivity clamps); the backward mirrors it as B28-B1 (with B18', B19') inside +the single shared function +`src/routing/mmc_op.rs::timestep_backward_core`. + +### 9.1 Backward labels, in execution order + +| label | differentiates | key expression | +|---|---|---| +| **B28** | S28 output clamp (C10) | $\bar{x}_{\mathrm{sol}} = \bar{y}\cdot\mathbb{1}[x_{\mathrm{sol}} > 10^{-4}]$ | +| **B27** | the triangular solve | equations 60-61; calls `dispatch::backward_solve_primitive`, then the per-nnz $-\bar{b}[r]x[c]$ | +| *zeta hook* | equations 73-79 | called with $\bar{b}$ the moment B27 produces it; returns the three geometry grads | +| **B26** | S26 assembly | equation 63, $\bar{c}_1$ from $\bar{A}_{\text{values}}$ | +| **B25** | S25 RHS | $\bar{c}_2 = \bar{b}\,i_t$; $\bar{c}_3 = \bar{b}\,Q_t$; $\bar{c}_4 = \bar{b}\,q'_t$; $\bar{i}_t = c_2\bar{b}$; $\bar{Q}_t \mathrel{+}= c_3\bar{b}$; $\bar{q}'_t = c_4\bar{b}$ | +| **B24** | S24 SpMV | equation 64; **the one $O(nnz)$ piece owned by a single parent**, skipped entirely when $Q_t$ is untracked | +| **B23-B20** | $c_1..c_4$ and the shared denominator | $\overline{\mathrm{denom}}\big|_{c_i} = -\bar{c}_i\,\mathrm{num}_i/\mathrm{denom}^2$; $\overline{\mathrm{num}_i} = \bar{c}_i/\mathrm{denom}$; accumulated into $\overline{2KX}$ and $\overline{2K(1-X)}$ | +| **B19** | the Cunge $X$ (`ddr_match: false` only) | $\bar{X} = 2K(\overline{2KX} - \overline{2K(1-X)})$, then $\partial X/\partial Q_t = -\tfrac{1}{2}/(TscL)$, $\partial X/\partial T = +\tfrac{1}{2}D_{\mathrm{cell}}/T$, $\partial X/\partial c = +\tfrac{1}{2}D_{\mathrm{cell}}/c$, masked by C8 | +| **B19'** | the S19' three-way `min` (`enforce_positivity`) | partitions $\bar{X}$ across branches Cunge > hi$_a$ > hi$_b$ by a `<=` cascade; opens the NEW path $X \to Cr \to K$, accumulated as $\bar{K}\big|_{X\text{-cap}} = -\bar{Cr}\,\Delta t/K^2$ | +| **B18'** | the S18' $K$ floor (C7) | $\bar{K}_{\mathrm{raw}} = \bar{K}\cdot\mathbb{1}[L/c > \Delta t(1+\delta)/2]$ | +| **B18** | $K = L/c$ | $\bar{c} = -\bar{K}_{\mathrm{raw}}L/c^2$, plus B19's $\partial X/\partial c$ term | +| **B17** | the celerity | `ddr_match`: $\bar{v} = \bar{c}\cdot 5/3$. Else with $G = \tfrac43 A\sqrt{1+z^2}/(TP_w)$ and $H = \gamma A/(Td)$: $\bar{v} = \bar{c}\beta$, $\bar\beta = \bar{c}v$, $\partial\beta/\partial A = -G/A + H/A$, $\partial\beta/\partial T = +G/T - H/T$, $\partial\beta/\partial P_w = +G/P_w$, $\partial\beta/\partial z = -Gz/(1+z^2)$, $\partial\beta/\partial d = -H/d$, $\partial c/\partial\gamma = vA/(Td)$ | +| **B16** | S16 velocity clamp (C6) | mask $v_{\mathrm{lb}} < v_{\mathrm{un}} < 15$ | +| **B15** | Manning velocity | $\partial v/\partial n = -v/n$; $\partial v/\partial R = \tfrac23 v/R$; and with stage roughness $\partial v/\partial d\big|_{\text{explicit}} = \gamma v/d$, $\partial v/\partial\gamma = v\ln(d/d_{\mathrm{ref}})$ | +| **B14** | $R = A/P_w$ | $\partial R/\partial A = 1/P_w$; $\partial R/\partial P_w = -R/P_w$ | +| **B13** | wetted perimeter | $\partial P_w/\partial b = 1$; $\partial P_w/\partial d = 2\sqrt{1+z^2}$; $\partial P_w/\partial z = 2dz/\sqrt{1+z^2}$ | +| **B12** | area | $\partial A/\partial T = \partial A/\partial b = d/2$; $\partial A/\partial d = (T+b)/2$ | +| **B11** | bottom-width floor (C5) | mask $b_{\mathrm{raw}} > b_{\mathrm{lb}}$ | +| **B10** | $b_{\mathrm{raw}} = T - 2zd$ | $\partial/\partial T = 1$; $\partial/\partial z = -2d$; $\partial/\partial d = -2z$ | +| **B9** | side-slope clamp (C4) | mask $0.5 < z_{\mathrm{raw}} < 50$ | +| **B8** | $z_{\mathrm{raw}} = Tq_\varepsilon/(2d)$ | $\partial/\partial T = q_\varepsilon/(2d)$; $\partial/\partial q_\varepsilon = T/(2d)$; $\partial/\partial d = -Tq_\varepsilon/(2d^2)$ | +| **B7** | $T = p\,d^{q_\varepsilon}$ | $\partial T/\partial p = T/p$; $\partial T/\partial d = Tq_\varepsilon/d$; $\partial T/\partial q_\varepsilon = T\ln d$ | +| **B6** | depth floor (C3) and $d = \mathcal{R}^e$ | mask $d > d_{\mathrm{lb}}$; $\partial d/\partial\mathcal{R} = e\,d/\mathcal{R}$; $\partial d/\partial e = d\ln\mathcal{R}$ | +| **B5** | the exponent | $\partial e/\partial q_\varepsilon = \partial e/\partial\gamma = -9/(5+3q_\varepsilon+3\gamma)^2$ | +| **B4** | $\mathcal{R} = \mathcal{N}/\mathcal{D}$ | $\partial\mathcal{R}/\partial\mathcal{N} = 1/\mathcal{D}$; $\partial\mathcal{R}/\partial\mathcal{D} = -\mathcal{R}/\mathcal{D}$ | +| **B3** | $\mathcal{D} = p\sqrt{s} + 10^{-8}$ | $\partial\mathcal{D}/\partial p = \sqrt{s}$ (slope is constant) | +| **B2** | the numerator | $\partial\mathcal{N}/\partial Q_t = n(q_\varepsilon+1)d_{\mathrm{ref}}^{\gamma}$; $\partial/\partial n = Q_t(q_\varepsilon+1)d_{\mathrm{ref}}^{\gamma}$; $\partial/\partial q_\varepsilon = Q_t n\,d_{\mathrm{ref}}^{\gamma}$ | +| **B1** | $q_\varepsilon = q + 10^{-6}$ | identity; sums the four $q_\varepsilon$ contributions (B8, B7, B5, B2) plus the zeta one | + +Two implementation facts that matter for reading the code: `state.c1` and +`state.x_storage` are deliberately **never read** in the backward +(`gc1` comes from the assembled A-values at B26, and $X$ comes from the saved +`x_effective`, which is what $c_1..c_4$ were actually built from); and the +$X$-cap term $\bar{K}\big|_{X\text{-cap}}$ must join $\bar{K}$ *after* the +$c_1..c_4$ contribution but *before* the B18' floor mask, because both paths +reach the celerity through the same `clamp_min`. + +### 9.2 How many gradient paths reach each quantity + +- **$Q_t$** has up to **four** paths: the S25 RHS ($c_3 Q_t$), the S24 SpMV + ($N Q_t$), the S2 depth chain, and, under `ddr_match: false`, the Cunge $X$ + at S19. This is why the routing is recurrent in the parameters and not just + in the state. +- **$T$** has up to five contributors before B7 consumes it: B12, B10, B8, + B17's $\partial\beta/\partial T$, and B19's $\partial X/\partial T$. +- **$d$** has up to seven: B13, B12, B10, B8, B7, plus B17's $-H/d$ and B15's + $\gamma v/d$ under stage roughness, plus the zeta term (79). +- **$\gamma$**, when learned, accumulates exactly three: B5 (depth exponent), + B15 (the explicit $(d/d_{\mathrm{ref}})^{\gamma}$ on the velocity), and B17 + (the celerity's $\gamma A/(Td)$). The code asserts all three are `Some` + together or `None` together. + +### 9.3 The three operators and `ParentMask` + +| operator | arity | parents, in fixed order | +|---|---|---| +| `TimestepOp` | `Backward` | `n`, `q_spatial`, `p_spatial`, `q_t`, `q_prime_t` | +| `TimestepGammaOp` | `Backward` | the five above **+ `gamma`** | +| `TimestepLeakanceOp` | `Backward` | the five above **+ `K_D`, `d_gw`, `leakance_factor`** | + +All three share `timestep_backward_core`; only parent bookkeeping differs. +`TimestepGammaOp` exists because a tensor that is not a parent receives no +gradient at all, and widening `TimestepOp` would have given every existing run +a sixth parent it does not use. `TimestepLeakanceOp` has **no `gamma` parent**, +which is why config load rejects learned `gamma` together with `use_leakance`. + +`src/routing/mmc_op.rs::ParentMask` is a six-field `bool` struct built from +`ops.parents` by each `Backward` impl (`ids[i].is_some()`). It gates only: + +1. each parent's **final assembly** (`mask.n.then(|| ...)` and friends), and +2. the one $O(nnz)$ piece that belongs to a single parent, namely B24's + $N^{\mathsf{T}}\bar{i}_t$ for `q_t`. + +The shared chain through the solve and the geometry is needed by every parent +and is never skipped. `register_parent` **panics** if a tracked parent arrives +without a gradient, so a mask that disagrees with `ops.parents` is a loud bug +rather than a silent zero. Typical untracked parents: `p_spatial` when it is a +fixed constant, `q_t` at the first timestep (the hotstart state), and +`q_prime_t` in ordinary training (the lateral inflow is data; only the adjoint +study lifts it to a leaf). + +--- + +## 10. Loss functions + +`src/training/loss.rs`, selected by `experiment.loss.kind` +(`src/config.rs::LossKind`). There are **five** kinds, not four. + +### 10.1 The `tau` output shift and daily pooling + +`src/training/loss.rs::tau_trim_and_downsample`, applied to the hourly routed +prediction $(G, T_{\mathrm{hours}})$ before any loss or metric. + +$$ +\text{slice} \;=\; \big[\,\tau,\; T_{\mathrm{hours}} - (24 - \tau)\,\big), +\qquad +T_{\mathrm{days}} = \frac{T_{\mathrm{hours}} - 24}{24} +\tag{80} +$$ + +$$ +\bar{Q}_{g,i} \;=\; \sum_{j} W_{i,j}\,Q^{\mathrm{hourly}}_{g,\,\tau + j}, +\qquad +W_{i,j} = \frac{\big|\,[\,is,\,(i+1)s\,) \cap [\,j,\,j+1\,)\,\big|}{s}, +\quad s = \frac{L}{M} +\tag{81} +$$ + +Equation 81 is `torch.nn.functional.interpolate(mode="area")` for the 1D case, +implemented as an explicit weight matrix (`::area_pool_weights`); each row sums +to 1, and because the trimmed length is an exact multiple of 24, $s = 24$ and +the pooling is an exact block mean. + +**The sign convention, and the trap.** Under the convention in force since +2026-08-08, pooled day $i$ covers hours $[\tau + 24i,\ \tau + 24(i+1))$ and is +scored against **observation day $i$**. So $\tau$ is the number of hours the +routed output is **advanced** before scoring (a translation-only +inverse-routing shift, same sign and magnitude as dMC-Juniata's tau): +$\tau = 0$ is exactly day-aligned; $\tau = 9$ pairs obs day $i$ with routed +hours $[24i + 9,\ 24i + 33)$. + +The pre-2026-08-08 slice was $[13+\tau,\ -11+\tau)$ with pooled day $i$ scored +against obs day $i+1$, so + +$$ +\tau_{\mathrm{new}} \;=\; \tau_{\mathrm{old}} - 11 +\tag{82} +$$ + +The old shipped value 3 is $\tau_{\mathrm{new}} = -8$; the old measured optimum +20 is $\tau_{\mathrm{new}} = 9$. **`params.tau` is a `u32` asserted to lie in +$[0, 24)$, so $\tau_{\mathrm{new}} = -8$ is unrepresentable.** The default is +still `tau: 3` (`src/config.rs::Params::default`), and under the new +convention that denotes a different shift than the legacy `tau: 3` did. Any config carrying +`tau: 3` inherited from before 2026-08-08 is therefore silently reinterpreted; +`config/experiments/gridded_conus.yaml` carries the migrated `tau: 9`. +DDR-Python's `compute_daily_runoff` still uses the legacy form. Total trim is +24 h under both conventions, so $T_{\mathrm{days}}$ is unchanged and a +$\tau$ sweep scores an identical day sample at every shift. + +### 10.2 Masking rule for missing observations + +Two-stage, and the same rule in both stages. + +**Stage 1, per gauge, in the driver** (`src/training/driver.rs`). After the +post-warmup slice, a gauge is **dropped entirely** if its window contains *any* +NaN: + +$$ +\text{keep}(g) \iff \forall t \in [\text{warmup},\, T_{\mathrm{days}}):\ \neg\,\mathrm{isNaN}\!\left(o_{g,t}\right) +\tag{83} +$$ + +Kept gauges are selected with `Tensor::select` so autograd stays alive on the +predictions. If no gauge survives, the micro-batch is skipped. The driver's +comment records why this matters: without it, a NaN-to-0.0 substitution biases +the head toward predicting near-zero flow and saturates Manning's $n$ at the +lower bound. + +Stage 1 therefore means **every loss below sees NaN-free observations**. The +per-element masking that follows only ever matters for the derivative term. + +**Stage 2, per adjacent pair, inside `obs_diff_term`.** A pair $(j, j+1)$ +counts for gauge $g$ only when both days are finite, so a gap **breaks the +chain**: $N$ valid days with a hole give strictly fewer than $N-1$ pairs. +Implemented as a dense $(G, T-1)$ weight tensor that is zero at invalid pairs, +with the observed-difference tensor also zero there, so no gather is needed, +the term stays differentiable in $p$, and a NaN observation can never reach the +arithmetic. + +`src/training/metrics.rs::Metrics::compute` masks differently again: it drops +individual non-finite *pairs* rather than whole gauges, and emits NaN for a +gauge with no finite pair at all. + +### 10.3 `l1` (the default) + +$$ +L_{\mathrm{L1}} \;=\; \frac{1}{G\,T}\sum_{g=1}^{G}\sum_{t=1}^{T}\left|p_{g,t} - o_{g,t}\right| +\tag{84} +$$ + +A mean over **elements**. `src/training/loss.rs::l1_loss_post_warmup` is the +`ndarray` twin used for logging. + +### 10.4 `nse-batch` (dHBV's `NSELossBatch`) + +$$ +L_{\mathrm{nse}} \;=\; \frac{1}{G\,T}\sum_{g}\sum_{t} +\frac{\left(p_{g,t} - o_{g,t}\right)^2}{\left(\sigma_g + \epsilon\right)^2} +\tag{85} +$$ + +$\epsilon$ is `loss.eps`, default **0.1** (matching DDR `hydrograph_loss`), and +it sits **inside the square, added to $\sigma$**, not added to $\sigma^2$. +$\sigma_g$ is the gauge's observed standard deviation over the **whole training +period**, a fixed vector (`MeritGagesDataset::gauge_obs_std`), *not* recomputed +per window; recomputing per window would make the objective drift between +micro-batches and break accumulation exactness. Pairs with +`experiment.optimizer: adadelta` by design (scale-free), though the current arm +uses Adam. + +### 10.5 `nse-batch-deriv` + +$$ +L \;=\; L_{\mathrm{nse}} \;+\; \lambda \cdot L_{\mathrm{deriv}} +\tag{86} +$$ + +$$ +L_{\mathrm{deriv}} \;=\; \frac{1}{N_{\mathrm{pairs}}}\sum_{(g,j)\ \mathrm{valid}} +\frac{\Big[\left(p_{g,j+1} - p_{g,j}\right) - \left(o_{g,j+1} - o_{g,j}\right)\Big]^2} +{\left(\sigma^{\mathrm{d}}_g + \epsilon\right)^2} +\tag{87} +$$ + +$\lambda$ is `loss.deriv_weight`, default 0.5. $N_{\mathrm{pairs}}$ is the +count of valid adjacent pairs across the whole batch, so (87) is one global +mean, exactly as (85)'s `.mean()` is one global mean. +$\sigma^{\mathrm{d}}_g$ is the gauge's observed **consecutive-day difference** +standard deviation over the training period +(`MeritGagesDataset::gauge_obs_diff_std`, returning 0.0 for a degenerate +gauge, where the $+\epsilon$ is what keeps the term finite). Returns exactly 0 +when $T < 2$ or every pair is broken, so the composite degrades to the level +term rather than to NaN. At $\lambda = 0$ it is numerically identical to (85). + +Rationale from the docstring: the landscape study's curvature probe found the +loss curvature in Manning's $n$ to be governed by the mean square of the +hydrograph's **time derivative**, so a level-only objective leaves the $n$ +valley nearly flat and the aggregate gradient vanishes long before $n$ is +identified. Scoring the derivative directly deepens that valley (predicted +about 3.7x at $\lambda = 0.5$). This mirrors +`experiment::landscape::objective::deriv_window_loss` term for term. + +### 10.6 `nnse-kge` + +Per gauge, over the time axis, with **population** moments (divide by $T$): + +$$ +\mu_p = \overline{p_g}, \quad \mu_o = \overline{o_g}, \quad +\sigma_p = \sqrt{\overline{(p-\mu_p)^2} + \epsilon}, \quad +\sigma_o = \sqrt{\overline{(o-\mu_o)^2} + \epsilon} +\tag{88} +$$ + +$$ +r = \frac{\overline{(p-\mu_p)(o-\mu_o)}}{\sigma_p\,\sigma_o}, +\qquad +\alpha = \frac{\sigma_p}{\sigma_o}, +\qquad +\beta = \frac{\mu_p}{\mu_o + \epsilon} +\tag{89} +$$ + +$$ +1 - \mathrm{KGE} \;=\; \sqrt{(r-1)^2 + (\alpha-1)^2 + (\beta-1)^2} +\tag{90} +$$ + +$$ +\mathrm{NSE} = 1 - \frac{\sum_t (p - o)^2}{\sum_t (o - \mu_o)^2 + \epsilon}, +\qquad +\mathrm{NNSE} = \frac{1}{2 - \mathrm{NSE}}, +\qquad +1 - \mathrm{NNSE} +\tag{91} +$$ + +$$ +L_{\mathrm{nnse\text{-}kge}} \;=\; \lambda_{\mathrm{nnse}}\,\Big\langle 1 - \mathrm{NNSE}\Big\rangle_g +\;+\; \lambda_{\mathrm{kge}}\,\Big\langle 1 - \mathrm{KGE}\Big\rangle_g +\tag{92} +$$ + +$\lambda$'s are `loss.nnse_weight` / `loss.kge_weight`, both default 1.0. +$\langle\cdot\rangle_g$ is a mean over **gauges**, so large basins do not +dominate. **Epsilon placement**, exactly as coded: $+\epsilon$ is added to the +*variance* before the square root in (88) (so it stabilizes $\sigma$, not +$\sigma^2$); $+\epsilon$ is added to $\mu_o$ in the $\beta$ denominator; and +$+\epsilon$ is added to the *sum* of squared observed deviations (SSO) in (91), +not to its mean. The $r$ and $\alpha$ ratios are invariant to the population-vs- +sample choice; NSE's SSE/SSO is too. + +Rationale: L1 and NSE are both maximized at a simulated variance *below* +observed (NSE's optimum sits at $\alpha = r < 1$), so they reward the +Muskingum-Cunge routing for over-attenuating flood peaks. This was the +diagnosed cause of KGE regressing against the summed-Q' baseline: median KGE +$0.723 \to 0.701$ while NSE *improved* $0.639 \to 0.684$, with the whole drop in +$\alpha = \sigma_{\mathrm{sim}}/\sigma_{\mathrm{obs}}$ falling $0.93 \to 0.85$. +KGE's $(\alpha-1)^2$ supplies the restoring gradient; NNSE guards correlation +and volume. + +### 10.7 `kge` (component-weighted) + +$$ +L_{\mathrm{kge}} \;=\; \Big\langle \min\!\big(\,w_r(r-1)^2 + w_\alpha(\alpha-1)^2 + w_\beta(\beta-1)^2,\; \kappa\,\big)\Big\rangle_g +\;+\; \lambda_{\mathrm{nnse}}\Big\langle 1 - \mathrm{NNSE}\Big\rangle_g +\tag{93} +$$ + +with $r, \alpha, \beta$ from (89) and $\epsilon$ placed identically. +$w_r, w_\alpha, w_\beta$ default 1.0; $\kappa$ is `loss.kge_clamp`, default 10 +(clamp C13). When $\lambda_{\mathrm{nnse}} = 0$ the function returns after the +first term, skipping the NNSE computation entirely. + +Two reasons for the squared, unrooted form over (90), from the docstring: +$\sqrt{\cdot}$ has an infinite-slope cusp as the prediction approaches perfect +KGE (the argument goes to 0), whereas the squared form is smooth there; and +$w_\alpha$ can independently up-weight the variance-ratio term, the direct +counter-pressure to MC over-attenuation. $\kappa$ exists because a gauge with +near-constant observed flow has a collapsing $\sigma_o$ / $\mu_o$ denominator +and $(\alpha-1)^2$ / $(\beta-1)^2$ can explode: a single gauge drove batch loss +to about $10^4$ in testing. + +### 10.8 Gradient-accumulation weights + +`src/training/loss.rs::loss_denominator` returns the exact-recombination +weight, which differs by objective because L1 and the NSE forms average over +**elements** while the KGE composites average over **gauges**: + +$$ +n_i \;=\; +\begin{cases} +G^{\mathrm{kept}}_i \cdot T^{\mathrm{post}}_i, & \texttt{l1},\ \texttt{nse-batch},\ \texttt{nse-batch-deriv}\\ +G^{\mathrm{kept}}_i, & \texttt{nnse-kge},\ \texttt{kge} +\end{cases} +\tag{94} +$$ + +$$ +L_{\text{pooled}} \;=\; \frac{\sum_i L_i\, n_i}{\sum_i n_i} +\tag{95} +$$ + +A naive $1/N$ average is wrong whenever micro-batches differ in surviving-gauge +count, which the NaN filter (83) makes the common case. For +`nse-batch-deriv`, recombination is exact only to $O(1/T^{\mathrm{post}})$ +because the derivative term's pair count is $G(T-1)$ rather than $GT$; the +docstring recommends single-batch training with that objective if +micro-batches differ in gauge count. + +| symbol | meaning | units | status | +|---|---|---|---| +| $p_{g,t}$ | daily routed prediction | m³/s | derived | +| $o_{g,t}$ | daily USGS observation | m³/s | data | +| $\sigma_g, \sigma^{\mathrm{d}}_g$ | training-period per-gauge std of obs / obs differences | m³/s | data (fixed per run) | +| $\epsilon$ | `loss.eps`, default 0.1 | m³/s (it is added to a std) | config | +| $\tau$ | output advance in hours, $[0,24)$ | h | config | +| warmup | days dropped from the head of each window | d | config | + +--- + +## 11. Diagnostics used for validation + +### 11.1 The KGE decomposition + +`src/training/metrics.rs::Metrics::compute`, per gauge, on the post-warmup +daily arrays, dropping individual non-finite pairs. **Sample-free population +moments** (divide by $n$, the number of surviving pairs) and **no epsilon**: +the guards are explicit zero tests that emit NaN instead: + +$$ +r \;=\; \frac{\frac1n\sum(p-\mu_p)(o-\mu_o)}{\sigma_p\,\sigma_o} +\quad\text{(NaN if }\sigma_p = 0\text{ or }\sigma_o = 0) +\tag{96} +$$ + +$$ +\alpha \;=\; \frac{\sigma_p}{\sigma_o} +\quad\text{(NaN if }\sigma_o = 0) +\tag{97} +$$ + +$$ +\beta \;=\; \frac{\mu_p}{\mu_o} +\quad\text{(NaN if }\mu_o = 0) +\tag{98} +$$ + +$$ +\mathrm{KGE} \;=\; 1 - \sqrt{(r-1)^2 + (\alpha-1)^2 + (\beta-1)^2} +\tag{99} +$$ + +Precisely: **$r$ is the Pearson correlation of the daily series**, +**$\alpha$ is the variability ratio $\sigma_{\mathrm{sim}}/\sigma_{\mathrm{obs}}$ +of the daily series** (not of the flow-duration curve, and not a +coefficient-of-variation ratio), and **$\beta$ is the bias ratio +$\mu_{\mathrm{sim}}/\mu_{\mathrm{obs}}$**. Note the distinction from the loss: +equation 99 has **no $\epsilon$**, whereas (88)-(89) do, so a reported KGE and +the KGE term inside `nnse-kge` are not the same number on a low-variance gauge. + +The companion metrics, also per gauge: + +$$ +\mathrm{NSE} = 1 - \frac{\sum(p-o)^2}{\sum(o-\mu_o)^2} +\quad\text{(NaN if SSO} = 0), +\qquad +\mathrm{RMSE} = \sqrt{\frac1n\sum(p-o)^2} +\tag{100} +$$ + +$$ +\mathrm{bias} \;=\; \frac1n\sum_t\left(p_t - o_t\right) +\qquad [\mathrm{m^3/s}] +\tag{101} +$$ + +$$ +\mathrm{FHV} = 100\,\frac{\sum_{\mathrm{top\,2\%}}(p^{\uparrow} - o^{\uparrow})}{\sum_{\mathrm{top\,2\%}} o^{\uparrow}}, +\qquad +\mathrm{FLV} = 100\,\frac{\sum_{\mathrm{bottom\,30\%}}(p^{\uparrow} - o^{\uparrow})}{\sum_{\mathrm{bottom\,30\%}} o^{\uparrow}} +\tag{102} +$$ + +$p^{\uparrow}$ and $o^{\uparrow}$ are each series **sorted independently**, so +FHV and FLV are flow-duration-curve volume biases and are **not +timestep-paired**. Indices are `round(0.98 n)` and `round(0.30 n)`. `bias` is a +signed mean difference in m³/s, not a percentage or a ratio. + +### 11.2 The eval-time zeta flux accumulation + +Turned on by `MuskingumCunge::enable_zeta_accumulation` (eval only; the +training path never enables it). Per timestep, `route_timestep` adds to five +inner-backend running sums with no autograd tape +(`src/routing/mmc.rs`, `ZetaSumTensors`): + +$$ +S^{\mathrm{abs}}_i = \sum_{t} \left|\zeta_{t,i}\right|, \quad +S^{\mathrm{net}}_i = \sum_{t} \zeta_{t,i}, \quad +S^{d}_i = \sum_{t} d_{t,i}, \quad +S^{\mathcal{A}}_i = \sum_{t} \mathcal{A}_{z,t,i}, \quad +S^{Q}_i = \sum_{t} Q_{t+1,i} +\tag{103} +$$ + +Across chunked eval calls these merge in `src/training/forward.rs::ZetaSums` +(eval builds a fresh engine per chunk), and `src/training/eval.rs` divides by +the accumulated step count: + +$$ +\overline{X}_i \;=\; \frac{S^{X}_i}{n_{\mathrm{steps}}} +\tag{104} +$$ + +`src/dump_parameters.rs::write_zeta_netcdf` writes these onto a `COMID_eval` +dimension in `/kan_parameters.nc`. **Precisely what each reported +number is:** + +| NetCDF variable | quantity | units | caveat | +|---|---|---|---| +| `zeta` | eval-window mean of $|\zeta|$ per reach | m³/s | **magnitude**, so a reach that alternates sign does not cancel; this is the $|\zeta| > 0.01$ m³/s GO/NO-GO bar | +| `zeta_net` | eval-window mean of signed $\zeta$ | m³/s | positive = losing reach; equals `zeta` exactly under `leakance_losing_only: true` | +| `depth_mean` | eval-window mean routed depth | m | **start-of-step**, computed from $Q_t$ | +| `area_z_mean` | eval-window mean $\mathcal{A}_z$ | labelled m², actually $\mathrm{m}^{(2-q)q+1}$ | see §8.3 | +| `q_mean` | eval-window mean routed discharge | m³/s | **end-of-step** $Q_{t+1}$, so offset one timestep from `depth_mean` | +| `COMID_eval` | reach ids | - | the **eval network** (gauge-subgraph union), NOT full CONUS | + +The step offset between `depth_mean`/`area_z_mean` (start-of-step) and `q_mean` +(end-of-step) is documented in `write_zeta_netcdf`'s docstring and matters if +the analyst ratios them. The `zeta` value written is recomputed from the *saved* +primitives with the same `losing_only` flag and the same impervious mask the +forward used, so it is exactly what was subtracted from $b$, not an +independent estimate. + +Two other diagnostics worth naming because they are easy to misread: + +- **`negative solves`** (`src/routing/mmc_op.rs::negative_solve_stats`): the + count of solve outputs that were negative *before* clamp C10 rewrote them to + $+10^{-4}$, over all reaches and timesteps in one `forward`. A host readback, + gated on `track_negative_discharge` (on in training, off in eval). On the + CUDA-graph path the counters stay at zero because + `forward_chain_inner` is never entered, so `forward` prints + `UNAVAILABLE` rather than a zero, deliberately, so silence cannot be + misread as a measurement. +- **The summed-Q' baseline** (`src/baseline/`): per-gauge sum of upstream + divide $Q_r$ over the eval window against USGS daily observations. No + routing, no learned parameters. It scores the **same** gauge population + training does: `valid_gauges` skips gauges with no subgraph, single-divide + headwater gauges, and gauges with no observation series. The headwater skip + moved the baseline median NSE from 0.142 to 0.315; never compare a trained + median against a baseline computed over a different gauge population. + +--- + +## 12. Code versus documentation discrepancies found + +Ordered by how much they could distort a physical-meaningfulness analysis. + +### 12.1 `area_z` is dimensionally not an area (code-versus-code, undocumented) + +`src/routing/leakance.rs::zeta_forward` computes +$w_z = (p\,d)^{q_\varepsilon}$, whereas the top width everywhere else is +$T = p\,d^{q_\varepsilon}$. Since $[p] = \mathrm{m}^{1-q}$, the plan-view +"area" $\mathcal{A}_z = w_z L$ carries units +$\mathrm{m}^{(2-q)q + 1}$, which is m² only at $q = 1$ (see equation 71); at +the prescribed $q = 0.65$ it is $\mathrm{m}^{1.878}$. Consequently $\zeta$ is +not in m³/s and a learned $K_D$ absorbs a $q$- and $p$-dependent unit +conversion. `src/dump_parameters.rs::write_zeta_netcdf` nevertheless labels +`area_z_mean` with `units = "m^2"`. Nothing in the source flags this; it is +ported verbatim from DDR `_compute_zeta`. **A learned $K_D$ should not be +compared numerically against a measured streambed leakance without redoing +this algebra.** + +### 12.2 A learnable `x_storage` receives no gradient, but the code says it does + +`src/training/forward.rs::forward` denormalizes an `x_storage` head output when +the KAN emits it, with the comment "gradient already flows via the custom +sparse backward in mmc_op.rs". It does not. +`src/routing/mmc_op.rs::timestep_forward` registers exactly five parents +(`n`, `q_spatial`, `p_spatial`, `q_t`, `q_prime_t`), six with `gamma`, eight +with leakance; `x_storage` is **not** among them in any variant. Its node is +extracted (`xst_aut`) and its primitive saved into `TimestepState::x_storage`, +but the backward's own comment states `state.x_storage` "is deliberately never +read here". Doubly inert under the default physics, because with +`ddr_match: false` the S19 Cunge $X$ ignores `xst_in` entirely. Nothing in +`src/config.rs` validates or rejects `x_storage` in +`kan_head.learnable_parameters`, so such a config trains a dead output +silently. Contrast this with the learned-`gamma` case, which is explicitly +rejected alongside `use_leakance` precisely because "gamma would silently +receive no gradient". + +### 12.3 `.claude/ARCHITECTURE.md` describes a solver that no longer exists + +Three claims in the "Per-timestep dataflow" and "Why dense forward +substitution, not sparse CSR + custom autograd" sections are now false: + +1. It states the solve is a **dense** `triangular_solve_lower` with autograd + falling out of BURN. The production path is sparse CSR with a hand-written + `Backward` (`src/sparse/mod.rs::CsrSolveOp`), which `CLAUDE.md` invariant 4 + protects by name. `src/routing/utils.rs::triangular_solve_lower` still + exists but is only reachable from the reference/test path. +2. It states `celerity c = clamp(v, v_lb, 15) * 5/3` unconditionally. That is + the deprecated `ddr_match: true` branch; the default since 2026-08-19 is the + trapezoidal $\beta$ of equation 35. +3. Its dataflow diagram shows `x = 0.3` implicitly via `x_storage`, with no + mention of the Cunge $X$ (equation 48), the positivity clamps (50)-(51), or + stage-dependent roughness (31)-(33). + +The module-map table and the SP-8/9/10 sections are still accurate. + +### 12.4 `docs/book/algorithm.md` presents only the deprecated physics + +Read in this worktree, `docs/book/algorithm.md` states the Leopold and Maddock +coefficient and exponent **correctly** (line 9: "width **exponent** +`q_spatial` and width **coefficient** `p_spatial`", with +`top_width = p * depth^q`), as do `docs/book/architecture.md` and +`docs/book/usage/inputs-formatting.md`. The prompt for this document reported +that page as having them backwards; that is not the case at HEAD `289163b`, so +either it was already fixed or the error is elsewhere. What that page **does** +get wrong is the physics branch: it presents $c = \tfrac53 v$ as *the* celerity +and $X \in [0, 0.5]$ as a supplied constant, with no Cunge derivation, no +positivity clamp, and no stage roughness. Its depth formula (equation 29 here) +and its $c_1..c_4$ block are correct for $\gamma = 0$. + +### 12.5 The prompt's framing of "coefficient clamping" does not match the code + +There is **no** clamping of $c_1, c_2, c_3, c_4$ anywhere in `src/routing/`. +`params.enforce_positivity` clamps the *inputs* $K$ and $X$ (equations 50-51), +and its own config docstring explains why: because $c_1+c_2+c_3 = 1$ holds for +any $(K, X)$, clamping inputs preserves mass exactly while clamping $c_3$ would +not. The clamp that **does** break mass is the post-solve floor +$Q_{t+1} = \max(x_{\mathrm{sol}}, 10^{-4})$ in +`src/routing/mmc_op.rs::forward_chain_inner` (equation 52), which the +`NEG_SOLVES` docstring describes as one that "both CREATES MASS and removes the +only symptom". See §6.5 for the accurate statement and the measured rates. + +### 12.6 `src/geometry.rs` says "invert Manning's equation for a trapezoidal section" + +The closed form it then computes (equations 16-20) is the inversion of the +**power-law** section with $R \approx d$, not of the trapezoid that equations +24-26 subsequently build; the two areas differ by $(2-q)(q+1)/2$, which is +1.114 at $q = 0.65$ (equation 30). `tests/stage_roughness.rs` measures the +related $\mathrm{d}A/\mathrm{d}d$ gap and calls it "up to 11 %", so the fact is +known in the test suite but the geometry docstring still reads as if the +inversion were self-consistent with the trapezoid. + +### 12.7 `params.tau` default is stale relative to its own convention + +`src/config.rs::Params::default` sets `tau: 3`. Under the 2026-08-08 convention +documented on `src/training/loss.rs::tau_trim_and_downsample`, the legacy +shipped value 3 maps to $\tau_{\mathrm{new}} = -8$ (equation 82), which a `u32` +asserted into $[0, 24)$ cannot express. So the default is not the migrated +legacy behaviour, and it is not the measured optimum either (legacy 20, new 9). +Only `config/experiments/gridded_conus.yaml` carries the migrated `tau: 9`; +`config/merit_training.yaml` sets no `tau` and therefore inherits 3. + +### 12.8 `experiment.learning_rate` is documented as inert for AdaDelta but is applied + +Recorded in the current experiment config's own commentary, and worth repeating +because it invalidated a whole run: +`src/config.rs::OptimizerKind::Adadelta` documents the key as ignored +("Ignores the `learning_rate` schedule by design"), but the driver passes `lr` +to `optimizer.step` and `src/training/adadelta.rs` multiplies delta by it. The +2026-07-31 AdaDelta run therefore ran every update about 1000x scaled down and +moved the head by $L_2(\Delta w) = 1.9\times10^{-4}$, i.e. it never left +initialization while reporting plausible losses. dHBV intends AdaDelta at +$lr = 1.0$. **Verified for this document:** +`src/training/driver.rs` passes `lr as f64` into `optimizer.step` on both the +single-batch and accumulation paths, and +`src/training/adadelta.rs::AdaDelta` applies +`tensor - delta.mul_scalar(lr as f32)`. So the schedule is applied, the +docstring is wrong, and an AdaDelta run must set +`learning_rate: {1: 1.0}` explicitly. + +### 12.9 `src/routing/mmc.rs::calculate_muskingum_coefficients` is a second, unused copy + +It is term-for-term identical to S20-S23 in `forward_chain_inner` and is what +`docs/book/algorithm.md` cites, but the production per-timestep path never +calls it (`route_timestep` dispatches straight into `mmc_op`). A reader who +edits one and not the other gets no test failure from the +`compare_ddr_sandbox` gate, which exercises the `mmc_op` path. + +### 12.10 Minor: the loss family has five members, not four + +`src/config.rs::LossKind` is `L1 | NnseKge | Kge | NseBatch | NseBatchDeriv`. +`CLAUDE.md`'s "Training objective" section lists four (`l1`, `nnse-kge`, +`kge`, `nse-batch`) and omits `nse-batch-deriv`, which is the objective the +landscape-curvature work introduced and which carries its own +`loss.deriv_weight` and a second per-gauge statistic (`gauge_obs_diff_std`). + +--- + +## 13. What I could not determine from the code + +Labelled gaps, not guesses. + +1. **Whether the learned $n_0$ is being read at a consistent reference depth + across published runs.** Equation 31 makes $n_0$ the roughness *at* + $d = d_{\mathrm{ref}}$, and `validate_stage_roughness` forces + $d_{\mathrm{ref}} = 1.0$ m for a global $\gamma$. For a *learned* per-reach + $\gamma$ the restriction lifts, and I did not find where (or whether) + `d_ref` is then recorded in the run manifest. If two runs used different + `d_ref`, their `n` fields are not comparable, and I could not confirm the + manifest captures it. + +2. **The `mask` field of `rskan::KanLayer`.** The fixture loader reads + `block_{b}_mask` and the forward multiplies by it (equation 4), but I did not + read `rskan`'s init or its `Module` derive closely enough to say whether + `mask` is a trained `Param` or a frozen buffer initialized to ones. It is + listed as "learned/fixed" in the §2.2 table for that reason. This matters + only if an analyst wants the head's exact trainable-parameter count. + +3. **The B-spline basis and grid update.** Equation 4 cites + $B_{j,k}(x)$ abstractly. I read `rskan/src/layer.rs::forward` and + the call signature of its `coef2curve` helper, but not + `rskan/src/spline.rs`, so I cannot state + the knot placement, the extension rule outside `grid_range`, or whether the + grid is ever re-fit during training. The `kan_grid_range` docstring in + `src/nn/kan_head.rs` says that outside the grid "the spline term flattens + and only the `scale_base * SiLU` path survives", which implies a fixed grid + with flat extrapolation, but I did not verify it in `spline.rs`. + +4. **Which `params.stage_roughness` / learned-`gamma` configuration produced + any specific published number.** The current arm + (`config/experiments/sr_n0_leakance_gated.yaml`) learns + `[n, K_D, d_gw, leakance_factor]` with no `stage_roughness` block, so + $\gamma = 0$ in it despite the filename prefix `sr_`. I did not audit + `.ddrs/runs/*/config.yaml` snapshots, so I cannot say which recorded results + had stage roughness active. + +5. **The disaggregation head's equations.** `src/nn/disagg_head.rs` implements + the precip-conditioned mass-preserving daily-to-hourly forcing head and is a + learned component of the model, but it was outside the ten requested topics + and I did not read it. In the current arm it is present with + `enabled: false`, so the forcing is flat repeat-24; anyone analysing an arm + with `disaggregation.enabled: true` needs that file documented separately. + The one relevant fact I can state from `src/training/forward.rs` is that + when the head is attached it replaces `tensors.q_prime` wholesale and its + output feeds $q'_t$ in equation 53. + +6. **The measured CONUS distributions of $d$, $T$, $z$, $b$, and how often each + clamp in §4.4 binds.** Only three such numbers are recorded in source or + tests: the $X$ cap binds on 95.3% of reach-timesteps under + `enforce_positivity`, 69.8% of reaches violate window (47) at $X = 0.3$, and + `n_at_floor` reached 47.3% at $lr = 0.01$. I found no in-tree measurement of + how often C3 (depth floor), C4 (side-slope band), C5 (bottom-width floor) or + C6 (velocity band) saturate. Since each saturation is a gradient sink, that + is the measurement most worth taking before claiming any learned parameter + is identified. + +7. **Whether `q_prime` is ever a tracked parent outside the adjoint study.** + `ParentMask::q_prime_t` exists and B25 computes $\bar{q}'_t = c_4\bar{b}$ + behind it, and the code comments say "only the adjoint study lifts it" as a + leaf. I did not read `src/experiment/adjoint/` to confirm what it does with + that gradient, so the adjoint influence map's exact definition is outside + this document. + +8. **The DDR-side reference for `_compute_zeta`.** I cite + `src/routing/leakance.rs`'s own header + (`~/projects/ddr/src/ddr/routing/mmc.py:146-197`, commit `c2bd0f9`) for the + claim that equation 65's $(p\,d)^{q}$ form is inherited verbatim. I did not + open the DDR tree in this worktree, so the port fidelity of that one + expression is asserted on the ddrs docstring, not independently checked. From a567f36347d76131dc456a2b9efb59dde8bc9bb8 Mon Sep 17 00:00:00 2001 From: Tadd Bindas Date: Wed, 16 Sep 2026 22:46:18 -0400 Subject: [PATCH 3/3] experiments/leakance: measure geometry-clamp saturation, the missing gradient-sink check The equations reference flags that nothing in the repo measures how often the geometry clamps saturate, which matters because a saturated clamp zeroes the gradient to the parameter behind it, so training cannot identify that parameter at that reach. Measured on the 2026-09-16T04-44-15Z arm over 64,892 eval reaches: the depth floor takes 2.1 percent, and the side-slope bounds, the bottom-width floor and both roughness range bounds take none. The side-slope clamp cannot fire in this configuration: with q fixed at 0.65 its lower bound needs a depth of 1,751 m and its upper bound a depth of 3.2 mm, which the 10 mm depth floor already forbids. So clamps are not the reason the learned outputs collapse onto one direction. That is consistent with the head-topology screen, which found capacity is not the constraint either, leaving the gradient and the objective. Figures are a lower bound: computed from each reach's mean eval depth, so a clamp that fires only at low flow is invisible. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_0143bx7gVz8ug2oc6DTy5jtz --- experiments/leakance/clamp_saturation.py | 66 ++++++++++++++++++++++++ 1 file changed, 66 insertions(+) create mode 100644 experiments/leakance/clamp_saturation.py diff --git a/experiments/leakance/clamp_saturation.py b/experiments/leakance/clamp_saturation.py new file mode 100644 index 00000000..d6933593 --- /dev/null +++ b/experiments/leakance/clamp_saturation.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python +"""How often do the geometry clamps saturate? Each saturated clamp is a gradient +sink: the parameter behind it receives zero gradient at that reach, so no amount +of training can identify it there. The equations reference flags this as the +measurement the repo has never taken. + +Reconstructs the geometry chain from a parameter dump plus the eval diagnostic's +mean depth, and reports the saturated fraction for each clamp in +params.attribute_minimums and for the side-slope bounds. + +Caveat stated up front: this uses each reach's MEAN eval depth, so it is a +lower bound on saturation. A clamp that fires only at low flow is invisible here. + +usage: clamp_saturation.py [--config C] +""" +import argparse, json +import numpy as np, netCDF4 as nc, pandas as pd + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("params"); ap.add_argument("zeta") + ap.add_argument("--depth-lb", type=float, default=0.01) + ap.add_argument("--velocity-lb", type=float, default=0.01) + ap.add_argument("--bottom-width-lb", type=float, default=0.01) + ap.add_argument("--side-slope-bounds", type=float, nargs=2, default=(0.5, 50.0)) + a = ap.parse_args() + + P = nc.Dataset(a.params); Z = nc.Dataset(a.zeta) + pc = np.asarray(P["COMID"][:], dtype=np.int64) + zc = np.asarray(Z["COMID_eval"][:], dtype=np.int64) + pos = pd.Index(pc).get_indexer(zc); ok = pos >= 0 + g = lambda d, v: np.asarray(d[v][:], dtype=float) + depth = g(Z, "depth_mean")[ok] + q = g(P, "q_spatial")[pos[ok]] if "q_spatial" in P.variables else np.full(ok.sum(), 0.65) + p = g(P, "p_spatial")[pos[ok]] if "p_spatial" in P.variables else np.full(ok.sum(), 21.0) + n = g(P, "n")[pos[ok]] + + # geometry chain, mirroring src/geometry.rs + top_width = p * depth**q + side_slope = top_width * q / (2.0 * depth) + side_clamped = np.clip(side_slope, *a.side_slope_bounds) + bottom_width = top_width - 2.0 * side_clamped * depth + + rep = { + "n_eval_reaches": int(ok.sum()), + "caveat": "computed from each reach's MEAN eval depth, so every figure is a LOWER bound on saturation", + "depth at its floor": float((depth <= a.depth_lb * 1.001).mean()), + "side slope at its lower bound 0.5": float((side_slope <= a.side_slope_bounds[0]).mean()), + "side slope at its upper bound 50": float((side_slope >= a.side_slope_bounds[1]).mean()), + "side slope clamped either way": float(((side_slope <= a.side_slope_bounds[0]) | (side_slope >= a.side_slope_bounds[1])).mean()), + "bottom width at or below its floor": float((bottom_width <= a.bottom_width_lb).mean()), + "roughness at its range floor 0.015": float((n <= 0.015 * 1.001).mean()), + "roughness at its range ceiling 0.25": float((n >= 0.25 * 0.999).mean()), + } + # the side-slope clamp is the interesting one: it is a pure function of q + rep["note_side_slope"] = ("side_slope = p*d^q * q / (2d) = (q/2)*p*d^(q-1); with q fixed at 0.65 " + "it saturates the 0.5 lower bound whenever d exceeds (q*p/(2*0.5))^(1/(1-q))") + if np.isfinite(q).all() and np.allclose(q, q[0]): + qq, pp = q[0], p[0] + d_crit = (qq * pp / (2.0 * a.side_slope_bounds[0]))**(1.0 / (1.0 - qq)) + rep["depth above which the side slope pins at 0.5 (m)"] = float(d_crit) + rep["fraction of reaches above that depth"] = float((depth > d_crit).mean()) + print(json.dumps(rep, indent=2)) + +if __name__ == "__main__": + main()