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/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/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() 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. 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();