diff --git a/.claude/PHYSICS-CORRECTIONS.md b/.claude/PHYSICS-CORRECTIONS.md new file mode 100644 index 0000000..53f5665 --- /dev/null +++ b/.claude/PHYSICS-CORRECTIONS.md @@ -0,0 +1,291 @@ +# `ddr_match` — the two physics paths + +`params.ddr_match: bool` (default **true**) selects which forward chain runs. +`true` reproduces DDR bit-for-bit and keeps `compare_ddr_sandbox` an ABSOLUTE +MATCH (invariant 1). `false` enables the corrected physics. + +``` + forward_chain_inner (src/routing/mmc_op.rs) + ───────────────────────────────────────────── + S1..S14 geometry depth → top_width → side_slope → bottom_width + (IDENTICAL in both) → area → wetted_perimeter → hydraulic_radius + │ + ▼ + S15/S16 velocity v = n⁻¹·R^(2/3)·√S ; v_cl = clamp(v, 0.01, 15) + (IDENTICAL in both) │ + ▼ + S17 celerity ┌─────────────────────────────────────────────┐ + │ ddr_match = true c = v_cl · 5/3 │ wide-rectangular + │ │ limit; +22..27% + │ ddr_match = false c = v_cl · β │ exact for the + │ β = 5/3 − (4/3)·A·√(1+z²)/(T·P) │ trapezoid built + └─────────────────────────────────────────────┘ above + │ + ▼ + S18 K = L / c (same in both — already correct, uses c not v) + │ + ▼ + S19 X ┌─────────────────────────────────────────────┐ + │ ddr_match = true X ≡ 0.3 (constant) │ D_num/D_phys + │ │ median 28x + │ ddr_match = false X = clamp( │ → Cunge: + │ 0.5·(1 − Q/(T·S·c·L)), 0, 0.5) │ diffusion matched + └─────────────────────────────────────────────┘ + │ + ▼ + S20..S23 coefficients denom = 2K(1−X) + Δt + c1 = (Δt − 2KX)/denom c2 = (Δt + 2KX)/denom + c3 = (2K(1−X) − Δt)/denom c4 = 2Δt/denom + c1 + c2 + c3 = 1 EXACTLY (both paths) + │ + ▼ + S24..S27 solve (I − c1·N)·Q_{t+1} = c2·(N·Q_t) + c3·Q_t + c4·q' + │ + ▼ + S28 clamp ┌─────────────────────────────────────────────┐ + │ both paths: q = clamp_min(x_sol, 1e-4) │ + │ NEW (both): count x_sol < 0 and report │ never measured + └─────────────────────────────────────────────┘ before +``` + +## `enforce_positivity` — provably zero negative solves + +`params.enforce_positivity: bool` (default **false**, rejected at load unless +`ddr_match: false`). Inserts two clamps at S18′/S19′ so the solve output cannot +be negative. `δ = POSITIVITY_DELTA = 1e-2`. + +``` + S18′ K floor k_raw = length / celerity + k_musk = max(k_raw, dt·(1+δ)/2) ⇒ Cr ≤ 2/(1+δ) + │ + S19′ X stability cap cr = dt / k_musk + hi_a = (1−δ)·0.5·cr ⇒ c1 ≥ 0 + hi_b = (1−δ)·(1 − 0.5·cr) ⇒ c3 ≥ 0 + x_eff = min(x_cunge, hi_a, hi_b) +``` + +### Why this is a proof, not a heuristic + +`c2, c4 > 0` always. `c1 ≥ 0 ⟺ Cr ≥ 2X`; `c3 ≥ 0 ⟺ Cr ≤ 2(1−X)` — i.e. exactly +the classical window `2X ≤ Cr ≤ 2(1−X)` (0 mismatches in 200k random draws). +The solve is forward substitution in topological order, +`x[i] = b[i] + c1[i]·Σ_{j∈up(i)} x[j]`, with `q_t > 0` (S28 clamp + hotstart, +`utils.rs:97`) and `q′ > 0` (`mmc.rs:453`), so `b ≥ 0`. Induction over the +topological order gives `x ≥ 0` everywhere. ∎ + +**Clamp the INPUTS (K, X), never the coefficients.** `c1+c2+c3 = 1` holds for +any `(K, X)`, so clamping K and X preserves mass exactly; clamping `c3` to zero +would break it. + +**δ is mandatory, not cosmetic.** At δ=0 the clamp lands exactly on `c1=0`/`c3=0` +and f32 roundoff crosses it (7,149 `c1<0`, 964 `c3<0` in a 400k f32 sweep). +Sign safety comes from the *numerator*: `dt − 2KX ≥ δ·dt = 36 s` independent of +K, ~1e5× the f32 representation error at that magnitude. The *value* of `c1` can +still be small (measured min +2.5e−6) because `denom ≈ 2K` grows without bound +for slow reaches — small, but never negative. + +### Verified on real CONUS data (`src/bin/probe_courant.rs`) + +Trained head, 1,841 gauges → 92,488 reaches, 2,135 hourly steps: + +| | OFF | ON | +|---|---|---| +| negative solves | 55,181 / 197,461,880 (0.0279 %) | **0 / 197,461,880** | +| min c1 / min c3 | −9.99e−1 / −9.97e−1 | +2.50e−6 / +5.00e−5 | +| cells c1<0 / c3<0 | 8,081,351 / 1,351,260 | 0 / 0 | + +Replicated at a second time window, on a 256-gauge batch, on the NdArray/CPU +backend, and against a second (differently trained) head — **exactly 0 in all +five**. + +### The cost — larger than first estimated + +Measured `Cr = dt/K` on the full network: p5 0.0142 · p25 0.0741 · **p50 0.226** +· p75 0.603 · p95 2.724. Only **7.3 %** of reach-timesteps have `Cr > 2` and get +K floored (median inflation 1.77×, p95 9.6×). + +The X cap binds on **95.3 %** of reach-timesteps. Median X used falls +**0.4976 → 0.0794, a 6.3× reduction**: + +| | p5 | p25 | p50 | p75 | p95 | +|---|---|---|---|---|---| +| X_cunge (pre-cap) | 0.330 | 0.485 | 0.4976 | 0.4997 | 0.5000 | +| X_eff (capped) | 0.0055 | 0.0219 | **0.0794** | 0.1936 | 0.398 | + +So `enforce_positivity` does not merely *shade* the Cunge X of `54ec215` — it +**replaces it almost everywhere**. Numerical diffusion becomes stability-set +rather than hydraulic-diffusivity-matched. **Skill impact is unmeasured; do not +promote this flag on the positivity guarantee alone.** + +> **Erratum.** An earlier version of this section and of +> `docs/superpowers/plans/2026-08-04-positivity-clamp.md` reported X falling only +> to ~0.45 at the median. That was a methodological error: `X_max(Cr)` is +> non-monotone (it peaks at `Cr = 1`), so evaluating it *at* each Cr percentile +> does not yield the percentiles of X. The tell was that the p95 entry came out +> *below* the p75 entry. The measured numbers above supersede it. + +### Why sub-stepping still does not substitute + +The Task-5 abandonment analysed landing *inside* `[2X, 2(1−X)]` at fixed X≈0.49. +Here the K constraint is one-sided (`K ≥ dt/2n_sub`), which shrinking `dt` does +satisfy — but it also drives `Cr` further from 1 for the ~93 % of reaches that +are already `Cr < 2`, shrinking `X_max` and making the cap bite *harder*. + +### Reach subdivision does NOT substitute either — measured, 2026-08-05 + +> **Erratum.** This section previously argued that the measured Cr distribution +> (median 0.226 — the typical MERIT reach is ~4.4× too *long* for an hourly +> step) "reframes the variable-Δx fix": split each reach into `m ≈ 4` pieces, +> land the network at `Cr ≈ 1`, and both `X_max → 0.5` and non-negative +> coefficients follow, making `enforce_positivity` unnecessary. **It was built +> and measured. It does not work, and the underlying claim was wrong.** +> +> The claim held only at `Cr` **exactly** 1. Both coefficients are non-negative +> only inside `[2X, 2(1−X)]`, a window of width `2(1−2X)` — 80 % wide at +> `X = 0.3`, but **1.4 % wide at the measured CONUS median `X = 0.4966`**. A +> build-time piece count fixes Δx from a *reference* flow while `Cr` tracks the +> *routed* celerity, which varies severalfold within one storm. No cap lands a +> flow-varying Cr inside a 1.4 % window. +> +> Measured on 1,841 gauges with `enforce_positivity` OFF (`probe_courant`): +> +> | arm | rows | Cr p50 | Cr > 2 | **c1 < 0** | c3 < 0 | both ≥ 0 | neg solves | ms/step | +> |---|---|---|---|---|---|---|---|---| +> | off | 92,488 | 0.096 | 2.10 % | **93.0 %** | 3.93 % | 3.1 % | 0.1356 % | 1.90 | +> | cap 4 | 171,381 | 0.115 | 0.18 % | **98.75 %** | 0.33 % | 0.9 % | 0.0945 % | 2.73 | +> | cap 8 | 184,676 | 0.123 | 0.16 % | **98.79 %** | 0.31 % | 0.9 % | 0.0876 % | 2.87 | +> +> `frac c1 < 0` gets **worse**; negative solves fall only 35 % for a 2.05× +> network, 1.5× step time and **+23.9 % total channel length**. +> +> What it *does* buy: `frac Cr > 2` 2.10 % → 0.16 %, nearly eliminating +> `c3 < 0` (3.93 % → 0.31 %) — and a clamp-off control shows it is the +> short-reach **length clamp**, not the splitting, that does this. `c3 < 0` is +> the smaller population. +> +> **A second erratum, same section.** Subdivision also does not restore the +> Cunge X's dynamic range: raw `X_cunge` median moves only **0.4973 → 0.4815** +> even uncapped. `D = q/(So·c·Δx)` is small because of the +> `attribute_minimums.slope = 1e-3` floor and large top width `B`, not because +> Δx is long. +> +> Root cause of the collapsed window: the cell Reynolds number on MERIT is +> `D ≈ 0.012` — advection-dominated, so Cunge correctly returns near-pure +> translation and `X → 0.5`. This **retroactively vindicates the constant +> `X = 0.3`** as a deliberate stability trade (80 % window width), not an +> oversight. +> +> Full write-up, cost tables and the code that stays in-tree: +> `.claude/REACH-SUBDIVISION.md`. Plan of record: +> `docs/superpowers/plans/2026-08-05-reach-subdivision.md`. + +## Outside the forward chain: per-gauge extraction (`outflow_idx`) + +`ddr_match` also gates `collate::compress`'s `outflow_idx` — WHICH reaches are +summed to form a gauge's prediction. This is downstream of the solver, so it +does **not** affect `compare_ddr_sandbox` (which never builds `outflow_idx`). + +``` + gauge 01457000 73006562 ──┐ + (the 26-gauge case) ├──> 73005764 (gauge reach, 250.1 km² + 73006585 ──┘ = 68% of the 366.8 km² basin) + + ddr_match = true outflow_idx = upstream cols [73006562, 73006585] + → drops the gauge reach's own local drainage + → predicted 1.58 vs observed 7.60, summed-Q' 7.38 + (0.215x, constant across all 15 eval years) + + ddr_match = false outflow_idx = [73005764] — the gauge reach itself + → mass-conserving: the MC solve there already carries + everything upstream PLUS its own lateral inflow +``` + +A USGS gauge measures all drainage above it and we do not know where along its +reach it sits, so `false` is the physical answer. `true` reproduces DDR's +`geodatazoo/merit.py:226-234`; DDR's Lynker path validates `outflow_idx` +against the flowpath `toid` column (`lynker_hydrofabric.py:239-250`), the MERIT +path does not — that is where this would have been caught upstream. + +Impact: 26 of 1841 gauges below 0.5x baseline (all small basins, 139-453 km², +3-9 reaches; median ratio over all gauges 0.952). The omitted mass is always +positive, so `true` biases EVERY ddrs-vs-baseline comparison against ddrs. +Gate: `tests/gauge_mass_conservation.rs` (steady-state mass check, both flag +values) + `collate.rs::outflow_idx_includes_the_gauge_reach_when_not_ddr_match`. + +## Why the corrections are coupled + +Muskingum non-negative coefficients require `2X ≤ Cr ≤ 2(1−X)` with +`Cr = Δt/K`. Measured on CONUS at mean flow: **69.8% of reaches fall outside** +the `X = 0.3` window `[0.6, 1.4]` (28.4% give `c1 < 0`, 41.4% give `c3 < 0`). + +Cunge `X ≈ 0.49` almost everywhere, which **narrows** the admissible window to +roughly `[0.98, 1.02]`. So enabling Cunge X without sub-stepping makes the +Courant violation worse, not better. **Tasks 4 and 5 must land together.** + +``` + Cr window vs X 0 0.6 1.0 1.4 2.0 + X = 0.3 [2X, 2(1-X)] |---------[=========|=========]--------| + X = 0.49 [2X, 2(1-X)] |--------------[====|====]-------------| + measured Cr p25 0.54 ── median 1.09 ── p75 2.46 ── p95 10.2 +``` + +## Blast radius + +| Change | Forward | Backward | Breaks parity | Gate | +|---|---|---|---|---| +| Task 1 flag | plumbing only | none | no | config test | +| Task 2 counter | read-only | none | no | none | +| Task 3 β | S17 | new gβ → gA, gT, gP, gz | yes (flag-gated) | `celerity_beta_gradcheck` | +| Task 4 Cunge X | new S19 | new gX → gq_t, gT, gc | yes (flag-gated) | `cunge_x_gradcheck` | +| Task 5 sub-step | timestep loop | tape depth ×n_sub | yes (flag-gated) | `substep_courant` | +| `outflow_idx` | extraction only (not the solver) | none | no (sandbox untouched) | `gauge_mass_conservation` | +| `enforce_positivity` | S18′ K floor, S19′ X cap | gk_musk mask + new cr→k path; XGrads masked by mask_cunge | no (default off; requires `ddr_match: false`) | `positivity_clamp` | +| `params.subdivision` | build-time graph only (larger N, `q'/m`, KAN gather, hot-start divisor) | **none** — no gradient path | no (default off) | `subdivide`, `subdivision_integration`, `adjacency_parity` | + +## CUDA backend coverage + +Every gate above declares `type I = NdArray` — **CPU only**. Since +`ddr_match: false` + `use_cuda_graphs: false` + a CUDA backend is a legal and +actively-used configuration, `tests/cuda_backward_parity.rs` re-runs the same +gradchecks with `burn_cuda::Cuda` as the inner backend: + +```bash +cargo test --features cuda --test cuda_backward_parity +``` + +Run it alongside `positivity_clamp` / `cunge_x` / `celerity_beta` on any change +to `src/routing/mmc_op.rs`. What it covers: + +* native central-difference gradcheck **on CUDA** for β, Cunge X and the + positivity clamp (correctness, not just CPU agreement); +* CUDA-vs-CPU analytic gradient parity, plus a zero-pattern assertion that + catches a mask disagreeing across backends; +* the transitive guard `enforce_positivity ⟹ !use_cuda_graphs`, which nothing + else asserts — it falls out of `validate_enforce_positivity` and + `validate_ddr_match` separately, and which one fires depends on `ddr_match`. + +Measured 2026-08-04 on an RTX 4080 SUPER (driver 610.43.02, nvcc 13.2, burn +0.21 fork `a033dc8`, cubecl `d562ab9`), `sparse_solver: cpu`: + +* forward `x_sol` is **bit-identical** CPU vs CUDA on the 10-reach fixture + (max abs diff 0.0, both clamp settings) — the S1..S23 geometry chain's + `powf`/`sqrt`/`recip`/`min_pair` agree exactly at these operating points; +* analytic gradients differ by at most **2.784e-7 relative** (~2 f32 ulp), + from the extra transcendentals the backward adds (`ratio.log()` in B6); +* CUDA analytic-vs-FD worst relative error **7.5e-4** (β + Cunge X + clamp). + +Falsification results (each mutation applied to `mmc_op.rs`, run, reverted): + +| Mutation | CUDA FD gradcheck | CPU↔CUDA parity | +|---|---|---| +| drop `∂β/∂z` | FAIL, rel 5.4e-2 (all 8) | passes | +| drop `∂X/∂B` | FAIL, rel 9.4e-2 … 1.6e0 (all 8) | passes | +| invert the B18′ K-floor mask | FAIL, rel 1.0 (4 clamp cases) | passes | +| swap the B19′ `hi_a`/`hi_b` tie-break | FAIL, rel 7.8e-1 (4 clamp cases) | passes | +| drop the new `x_eff→cr→k_musk` term | FAIL, rel 3.9e-1 (4 clamp cases) | passes | + +The right-hand column is not a defect: both backends run the same source, so a +physics error is invisible to a cross-backend comparison **by construction**. +The FD gradcheck is the falsifier for wrong physics; the parity test is the +falsifier for a backend-specific divergence. Do not treat either as a +substitute for the other. diff --git a/.claude/REACH-SUBDIVISION.md b/.claude/REACH-SUBDIVISION.md new file mode 100644 index 0000000..c28be2e --- /dev/null +++ b/.claude/REACH-SUBDIVISION.md @@ -0,0 +1,270 @@ +# Reach subdivision (`params.subdivision`, variable Δx) + +> ## STATUS: **NO-GO for "non-negative by construction"** (2026-08-05) +> +> The premise was wrong. Subdivision was built to drive `Cr = Δt/K` to ≈ 1 +> network-wide so that `c1` and `c3` would both be non-negative without the +> runtime `enforce_positivity` clamp. **Measured on the real network it makes +> `frac c1 < 0` WORSE — 93.0 % → 98.8 % at `max_pieces: 8`** — while costing +> 2.05× the rows, 1.5× the step time and +23.9 % total channel length. +> +> The code is correct, gated off by default, and **stays in-tree**: it is the +> only measurement apparatus for this question, it does fix the `Cr > 2` / +> `c3 < 0` population, and the hot-start division it added is a real bug fix. +> Do not re-open the "Cr ≈ 1 ⇒ non-negative coefficients" argument without +> reading §Why it fails. + +Config: `params.subdivision` (`src/config.rs:491-576`), default `enabled: false`. +Implementation: `src/adjacency/subdivide.rs`, wired in `src/adjacency/cache.rs`. +Plan of record: `docs/superpowers/plans/2026-08-05-reach-subdivision.md` +(Tasks 1-8, commits `741e475` … `6cb66bf`). + +--- + +## Why it fails — the window-width argument + +The plan's load-bearing claim: at `Cr = 1` both coefficients reduce to +`(1−2X)/(1+2(1−X)) ≥ 0` for **any** `X ≤ 0.5`, so `Cr ≈ 1` makes non-negativity +automatic. + +That is true only at `Cr` **exactly** 1. In general +`c1 ≥ 0 ⟺ Cr ≥ 2X` and `c3 ≥ 0 ⟺ Cr ≤ 2(1−X)`, so both hold only inside the +window `[2X, 2(1−X)]`, whose **width is `2(1−2X)`** and collapses as `X → 0.5`: + +| X | window `[2X, 2(1−X)]` | width | +|---|---|---| +| 0.30 (the `ddr_match: true` constant) | [0.600, 1.400] | 80 % | +| 0.45 | [0.900, 1.100] | 20 % | +| **0.4923** (measured, `max_pieces: 8`) | [0.9846, 1.0154] | **3.1 %** | +| **0.4966** (measured, subdivision off) | [0.9932, 1.0068] | **1.4 %** | + +A **static** piece count fixes Δx from a *reference* flow at build time, while +`Cr` tracks the **routed** celerity, which varies severalfold within a single +storm. Holding `Cr` inside a 1.4 % window with a fixed graph is structurally +impossible — not a tuning problem, and no choice of `max_pieces`, +`reference_n` or `min_length_fraction` changes it. + +Gate that records this: `tests/subdivision_integration.rs:: +both_coefficients_are_non_negative_only_inside_a_window_that_collapses_at_x_half`. +The plan's sketched `subdivision_makes_coefficients_non_negative_without_the_clamp` +(asserting `frac c1 < 0` under 1 %) was deliberately **not** added — the +measurement refutes it. + +### Root cause: `X → 0.5` because MERIT is advection-dominated + +The Cunge `X = ½(1 − q/(So·c·Δx))` saturates at 0.5 because the cell Reynolds +number `D = q/(So·c·Δx) ≈ 0.012` — physical diffusion is ~1 % of advective +transport, so Cunge correctly returns near-pure translation. Ponce & Theurer's +`C·D ≥ ξ` accuracy criterion is likewise unsatisfiable on this network at any +Δx that also keeps coefficients non-negative (see the plan, §"Why the target is +`Δx = c·Δt`"). + +**This retroactively vindicates DDR's constant `X = 0.3`.** It is not an +oversight — it is a deliberate stability trade buying an 80 %-wide window at +the cost of hydraulically-wrong numerical diffusion. Any future "correct the X" +work must state what it does about the window it destroys. + +--- + +## Measured on the real network + +`src/bin/probe_courant.rs`, 2026-08-05, trained head, 1,841 CONUS gauges, +2,135 hourly steps, `enforce_positivity` **OFF** (the whole point is whether +subdivision removes the need for it): + +| arm | rows | Cr p50 | Cr > 2 | **c1 < 0** | **c3 < 0** | both ≥ 0 | neg solves | ms/step | +|---|---|---|---|---|---|---|---|---| +| off | 92,488 | 0.096 | 2.10 % | **93.0 %** | 3.93 % | 3.1 % | 0.1356 % | 1.90 | +| cap 4 | 171,381 | 0.115 | 0.18 % | **98.75 %** | 0.33 % | 0.9 % | 0.0945 % | 2.73 | +| cap 8 | 184,676 | 0.123 | 0.16 % | **98.79 %** | 0.31 % | 0.9 % | 0.0876 % | 2.87 | +| cap 8, `reference_n 0.13` | 344,262 | 0.254 | 0.14 % | 90.90 % | 0.45 % | 8.7 % | 0.0401 % | 4.75 | + +Negative *solves* fall only 35 % (0.1356 % → 0.0876 %) for a 2.05× network and +1.5× step time. `both ≥ 0` — the quantity the design targeted — goes +**3.1 % → 0.9 %**, i.e. the wrong direction. + +Cap 16 at `reference_n 0.13` **OOMs a 16 GB RTX 4080**. + +### What subdivision DOES fix + +`frac Cr > 2` falls **2.10 % → 0.16 %**, which essentially eliminates `c3 < 0` +(**3.93 % → 0.31 %**). A `min_length_fraction: 0` control arm shows this is the +**short-reach length clamp**, not the splitting: with the clamp off, `Cr > 2` +returns to 1.08 % and `c3 < 0` to 2.00 %. + +But `c3 < 0` is the far smaller population, and the clamp costs **+23.9 % total +channel length** — a reach modelled longer than reality has a proportionally +longer travel time. That is a real physical distortion bought for a numerical +gain on 3.6 % of cells. + +### Graph cost (cap 8, shipped defaults) + +| quantity | measured | plan predicted | +|---|---|---| +| Σm (sub-reaches) | 709,974 (**2.05×**) | ~1.65 M (4.77×) | +| reaches clamped | **51.0 %** | 34.4 % | +| pinned at `max_clamp_factor = 4.0` | **14.1 %** | — | +| total channel length | **+23.9 %** | +17.1 % | + +**The plan's cap sweep is materially wrong for the shipped config** and should +not be quoted. Cause: `reference_n: 0.05` makes the reference celerity ~5× the +routed celerity — the trained CONUS median `n` is **0.130**, and `ddr_match: +false` uses `c = v·β` with β ≈ 1.33, not the wide-rectangular 5/3. A too-fast +`c_ref` gives a too-long `dx_target`, so fewer reaches split and far more get +length-clamped. Reproduce the cost table without routing: + +```bash +cargo run --release --bin probe_courant -- \ + --config ddrs.yaml --fabric --max-pieces 8 --clamp-report +``` + +### The X dynamic range does NOT recover + +Raw `X_cunge` median moves only **0.4973 → 0.4815** even uncapped. Subdivision +cannot un-saturate X, because `D = q/(So·c·Δx)` is small primarily from the +`attribute_minimums.slope = 1e-3` floor and large top width `B`, not from Δx. +Do not justify this feature on X's dynamic range (see the erratum in +`.claude/PHYSICS-CORRECTIONS.md`). + +--- + +## Hot start — a real bug the campaign found and fixed + +`setup_inputs` cold-starts with `(I − N)·Q₀ = q'₀`. The `q'/m` split lives in +`forward`, so an undivided cold start feeds the full `q'₀` into a chain of `m` +pieces and every parent outlet begins at `m ×` its true steady state. + +Measured (cap 8, 184,676 rows): **2.94× the correct total network discharge**, +still **41.7 % off at t = 120** — the *configured* `warmup` (5 days) — reaching +<10 % only at t = 221 and <5 % at t = 282. + +So **`MuskingumCunge::divide_hotstart_by_pieces` defaults `true`** +(`src/routing/mmc.rs:144,214`). It is an exact no-op without subdivision (no +divisor exists), so `compare_ddr_sandbox` stays an ABSOLUTE MATCH (1.53e-5 m³/s) +and `adjacency_parity` still matches element-for-element. +`probe_courant --divide-hotstart` A/B's it. + +--- + +## What the code does (all of this is correct and stays) + +### The two-sided rule + +``` +Δx_target = c_ref · Δt c_ref from reference_n + Q_ref = coeff·uparea^exp + L > Δx_target → split into m = min(ceil(L/Δx_target), max_pieces) pieces + of length L/m; q' → q'/m + L < Δx_target → clamp the length UP to min_length_fraction·Δx_target, + bounded by original_length · max_clamp_factor (do NOT merge) +``` + +Merging short reaches was rejected: a short reach may carry two upstream +tributaries or a junction below it, so collapsing it destroys topology. +Clamping its length achieves the same `Cr` with no topology change. +`max_clamp_factor` (default 4.0, added in `2a06c0f`) bounds the distortion — +unbounded, measured clamp factors ran to p99 = 36× and **max 48,597×**, because +`reference_celerity` uses a depth relation with no slope dependence while `v` +scales as `√S`, so steep small catchments get big-river depth *and* +steep-slope velocity (~8.9 m/s → a 32 km `dx_target`). + +### Topology + +``` + BEFORE AFTER (m=3) + U ──> P ──> D U₂ ──> P₀ ──> P₁ ──> P₂ ──> D₀ + └─ q'/3 q'/3 q'/3 + len(P) = L len(Pᵢ) = L/3, slope/n/p/q unchanged + gauge@P → row(P) gauge@P → P₂ (outlet = last piece) +``` + +Sub-reaches are hydraulically identical to their parent (MERIT carries no +within-reach variation), pieces are contiguous and ordered upstream→downstream, +and parents are already topologically ordered — so the expanded graph stays +topologically ordered and **lower-triangular for free** (invariant 3 holds). +The upstream parent's *outlet* piece connects to the downstream parent's +*inlet* piece (`subdivide.rs:327-328`). + +### Two index spaces + +`order` gains duplicates (m rows share one COMID), which would break +`IdIndex`. Resolved by keeping both: + +``` +parent space (N = 346,321) sub-reach space (N' = Σ min(m, M)) + parent_order[p] = COMID order[i] = COMID of i's parent + IdIndex built HERE parent_offset[p]..parent_offset[p+1] + ← attributes, q', KAN outputs = contiguous rows owned by parent p + m_p = parent_offset[p+1] − parent_offset[p] + outlet(p) = parent_offset[p+1] − 1 +``` + +`IdIndex` is built from **`parent_order`**, never from `order` +(`src/data/store/zarr.rs:121-122`). A store without the map synthesizes the +identity (`parent_order == order`, `parent_offset == 0..=n`), so every un-split +store keeps working unchanged. + +### Where each piece of plumbing lives + +| Concern | Site | +|---|---| +| Config + validation | `src/config.rs:491-576`, `validate_subdivision` `:1034` | +| Reference celerity, reach plan, expansion, `plan_stats` | `src/adjacency/subdivide.rs` | +| Upstream-area accumulation, sequencing, cache key | `src/adjacency/cache.rs` (`upstream_area_km2`, `reach_plan`, `resolve_or_build`) | +| Zarr persist/load of `parent_order` + `parent_offset` | `src/data/store/zarr.rs` | +| `q'/m` after the clamp; hot-start divisor | `src/routing/mmc.rs:267-273,340,539` | +| KAN parent→sub-reach gather | `src/training/forward.rs::gather_params_to_subreaches` | +| Gauge read at the parent's outlet piece; compressed-space `parent_offset` | `src/data/collate.rs` | +| Tests | `tests/subdivide.rs`, `tests/subdivision_integration.rs` (12) | + +### Gotchas worth keeping + +- **`catchsize` is NOT drainage area.** It is the *local* divide area (median + 36.7 km², max ~612 km² even for continental rivers). `reference_celerity` + needs accumulated upstream area, so `cache.rs::upstream_area_km2` accumulates + it downstream over the topological order (validated against the fabric's own + `10^log10_uparea`: ratio p5/p50/p95 = 1.000/1.000/1.000 over all 346,321 + reaches). `log10_uparea` cannot be read directly — it is NaN on 88 % of + `merit_global_attributes_v2.nc`. +- **All seven `Subdivision` fields are hashed into the adjacency content key** + (`cache.rs::content_key`), because every one of them moves `dx_target` and + hence the built graph. Hash the whole struct, not just `enabled` + + `max_pieces`. +- **`enabled: true` + explicit `conus_adjacency`/`gages_adjacency` is a config + error**, not a warning. Subdivision runs *inside* the managed builder, which + those keys bypass, so the flag would be silently inert and the manifest would + lie. `adjacency::validate::store_is_subdivided` reads only zarr metadata + (`n_parent < n`) and allows the one legitimate case: an explicit path to a + store that was already built subdivided. +- **`enabled: true` requires `use_cuda_graphs: false`** — a captured graph is + sized to a fixed reach count. +- **Retraining is mandatory.** Every learned parameter was fit against the + un-split network's effective diffusion; checkpoints do not transfer. +- **The KAN sees no new information.** Sub-reaches inherit the parent's + attributes, so subdivision can only change numerics, never identifiability. + +--- + +## Reproducing the measurement + +```bash +# un-split control +cargo run --release --bin probe_courant -- --config ddrs.yaml \ + --checkpoint /checkpoints/ --backend cuda \ + --gauges 1841 --rho 90 --steps 2136 + +# subdivided arm (--fabric switches to the managed build; caps cache separately) +cargo run --release --bin probe_courant -- --config ddrs.yaml \ + --checkpoint /checkpoints/ --backend cuda \ + --gauges 1841 --rho 90 --steps 2136 \ + --fabric --max-pieces 8 --reference-n 0.13 --divide-hotstart +``` + +Gate set after any change under `src/adjacency/subdivide.rs`, `src/routing/mmc.rs` +piece handling, or the parent map: + +```bash +cargo test --test subdivide --test subdivision_integration \ + --test adjacency_parity --test gauge_mass_conservation +cargo test --lib +cargo run --release --example compare_ddr_sandbox # must stay ABSOLUTE MATCH +``` diff --git a/.claude/skills/ddrs-dev/SKILL.md b/.claude/skills/ddrs-dev/SKILL.md index ddd2ed1..0828eb7 100644 --- a/.claude/skills/ddrs-dev/SKILL.md +++ b/.claude/skills/ddrs-dev/SKILL.md @@ -26,6 +26,7 @@ config reference, test-authoring patterns, and the current research status. | Write a new gradcheck / parity / fixture test | `references/testing.md` §Authoring patterns | | Diagnose a failed, hung, or wrong-result run | `references/traps.md` §Symptom → trap | | Check whether a question is already settled | `references/research-status.md` | +| Build or change the training gauge CSV / population | `references/gauge-population.md` | | Plot or interpret eval output | skill `ddrs-eval-plots` | ## Contents @@ -67,6 +68,13 @@ claim restated · Closed campaigns: leakance NO-GO, selective equifinality H1– Q′-store waves, synthetic-n interim · **Do-not-use list** · Structural constants · Evidence standard · Doc conventions · Open questions +**`references/gauge-population.md`** (100 lines) +Regenerating `gages_2000_area_balanced.csv` (one command, seed 42, all-local +inputs) · Relative `DA_VALID` (`ABS_DIFF/DRAIN_SQKM ≤ 10%`) vs the scale-biased +absolute criterion · The filter funnel (coverage in both configured windows, +non-headwater subgraph) · Consequences of changing the population: baseline +cache invalidation, incomparable metrics, run-log verification string + ## The five facts that cause the most wasted time 1. **Stale binary.** `cargo build` / `cargo run` do **not** update `~/.cargo/bin/ddrs`. diff --git a/.claude/skills/ddrs-dev/references/config.md b/.claude/skills/ddrs-dev/references/config.md index 4a4f1ea..a501945 100644 --- a/.claude/skills/ddrs-dev/references/config.md +++ b/.claude/skills/ddrs-dev/references/config.md @@ -4,9 +4,9 @@ Struct: `src/config.rs::Config`. Loaded via `Config::from_yaml_file_with_mode(path, ConfigMode::Training|Testing)`. Six top-level sections. Verified against source 2026-07-30. -**No `deny_unknown_fields` anywhere** — a typo'd key silently takes its default -instead of erroring. This is the single most common cause of "my config change did -nothing". +**No `deny_unknown_fields` except `DisaggregationSection`** (added 2026-08-03) — +everywhere else a typo'd key silently takes its default instead of erroring. This +is the single most common cause of "my config change did nothing". ## Contents @@ -119,9 +119,17 @@ Ten production `input_var_names`: `SoilGrids1km_clay`, `aridity`, `meanelevation > **Current contract:** presence of the `disaggregation:` block ⇒ the head always > consumes precip ⇒ `data_sources.aorc_precip` is mandatory, else > `MeritGagesDataset::open` errors. It cannot silently degrade to flat repeat-24. +> **Exception (2026-08-03):** `disaggregation.enabled: false` strips the block at +> load time, making it inert — the sanctioned way to A/B the head vs nearest +> (repeat-24) without deleting the block. (This replaced the short-lived +> `experiment.use_frozen_kan_head`, which never ran an experiment.) +> The section is `#[serde(deny_unknown_fields)]` (2026-08-03): phantom keys like +> `use_precip` now FAIL LOAD with "unknown field" instead of silently taking +> defaults — the one section where a typo'd key cannot silently no-op. | Key | Default | Notes | |---|---|---| +| `enabled` | **true** | `false` ⇒ block stripped at load ⇒ flat repeat-24 (nearest) upsampling; the one-line ablation switch. `tests/disagg_enabled.rs` | | `hidden_size` | 16 | | | `num_hidden_layers` | 1 | | | `grid` | 3 | | @@ -140,7 +148,7 @@ Ten production `input_var_names`: `SoilGrids1km_clay`, `aridity`, `meanelevation | `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 `≥`) | -| `tau` | 3 | **Not** a routing sub-step count. It is the hourly→daily trim phase offset in `tau_trim_and_downsample`: DDR's slice `[13 + tau : -11 + tau]`, then area-pool to days (`src/training/loss.rs:17-45`). Nothing in `src/routing/` reads it | +| `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` | diff --git a/.claude/skills/ddrs-dev/references/gauge-population.md b/.claude/skills/ddrs-dev/references/gauge-population.md new file mode 100644 index 0000000..fc38649 --- /dev/null +++ b/.claude/skills/ddrs-dev/references/gauge-population.md @@ -0,0 +1,90 @@ +# Building and changing the training gauge population + +Written 2026-08-02, the session that produced `gages_2000_area_balanced.csv`. +Covers: how a gauge CSV is constructed from GAGES-II, the filters a gauge must +survive to actually train, the exact regeneration command, and what breaks when +the population changes. Motivating diagnosis: +`/tmp/experiment-handoff-small-basin-domination.md` (small basins dominated +gradient share and pooled median NSE selected for identity routing). + +## The one command + +```bash +cd ~/projects/ddr && uv run python ~/projects/ddrs/scripts/build_gages_2000_area_balanced.py +``` + +Deterministic (seed 42, no network access — everything is local). Rewrites +`~/projects/ddr/references/gage_info/gages_2000_area_balanced.csv` and prints +the filter funnel, a coverage-sensitivity table, and the final area histogram. +To vary the recipe, edit the constants at the top of the script +(`REL_TOL`, `COVERAGE_MIN`, `TRAIN_WINDOW`, `EVAL_WINDOW`, `SEED`, +`N_SMALL`, `N_LARGE`) — do not fork the logic. + +## Inputs (all local; never fetch from USGS) + +| Input | Path | Role | +|---|---|---| +| GAGES-II gauge table | `~/projects/ddr/references/gage_info/GAGES-II.csv` | 8,931 rows, same schema as `gages_3000.csv` incl. `COMID`, `ABS_DIFF`, `DA_VALID` | +| Observations | `/mnt/ssd1/data/icechunk/usgs_daily_observations` | 9,067 gauges × 14,610 days (1980-01-01..2019-12-31), `gage_id` zero-padded strings, m³/s, NaN = missing | +| Gauge subgraphs | `~/projects/ddr/data/merit_gages_conus_adjacency.zarr` | 8,945 groups keyed by zero-padded STAID; `order` length 1 ⇔ headwater | + +## The filter funnel (why each stage exists) + +A gauge in the CSV is useless unless it survives *every* downstream filter in +ddrs, so the builder applies them all up front — the CSV population equals the +training population by construction (the lesson of the phantom-zero baseline +incident, traps.md T4). + +1. **`DA_VALID` must be relative, not absolute.** The GAGES-II/gages_3000 + precomputed column is `ABS_DIFF <= COMID_UNITAREA_SQKM` — effectively an + absolute area tolerance, which deletes large basins for ~2% relative + disagreement while admitting small basins at >6,000%. Recompute as + `ABS_DIFF / DRAIN_SQKM <= 0.10`. Effect on GAGES-II: 7,919 → 5,528 + (recovers 433 large, drops 2,824 mostly-small). ddrs only *reads* the + column (`src/data/store/gage_csv.rs:62`); the fix must happen at CSV + construction. +2. **Observation coverage in BOTH configured windows.** Check the windows + training actually slices (`config/merit_training.yaml`): train + 1981-10-01..1995-09-30, eval 1995-10-01..2010-09-30 — not round calendar + years. Bar used: ≥ 80% non-NaN days in each window. Sensitivity on the + 5,528 pool: any-data 3,636 · ≥50% 2,948 · ≥80% 2,512 · ≥90% 2,392 · + 100% 2,168. +3. **Non-headwater subgraph exists.** Mirror `GageSubgraph::is_headwater` + (`src/data/store/zarr.rs:128`): zarr group present AND `order` length > 1. + Single-divide gauges have empty upstream sets → all-zero summed-Q' + predictions. Dropped 97 of 2,512 → eligible pool 2,415. +4. **Area-balanced subsample.** Keep ALL basins ≥ 5,000 km² (they carry the + routing signal); top up [1,000, 5,000) to reach the ≥1,000 km² target; + random-draw the < 1,000 km² stratum. Keep small basins deliberately — they + teach the identity-routing regime; the goal is reducing their gradient + share, not removing them. 2026-08-02 result: 582 + 418 + 841 = **1,841** + (45.7% / 54.3% either side of 1,000 km²). The small stratum was the + binding constraint (only 841 eligible), so the set is 1,841 not 2,000. + +## Output contract + +Same schema `gage_csv.rs` reads (`STAID` zero-padded, `DA_VALID` literal +`True`, STANAME quoted) plus audit columns `REL_DIFF`, `COV_TRAIN`, +`COV_EVAL` — the serde reader ignores unknown headers. Full derivation is +also documented in `~/projects/ddr/references/gage_info/README.md`; the CSV +is committed in the **ddr** repo (05796ba), the builder in **ddrs**. + +## Consequences of changing the population — read before repointing + +- **The cached summed-Q' baseline is invalidated.** The gages path is in the + baseline cache key; `ddrs plan` recomputes it. This is correct — but it + means **no metric from runs on the old population is comparable**, in + either direction. Establish the new baseline before any improvement claim + (research-status.md §Gauge-set definitions has the standing warning). +- **Point `data_sources.gages` at the new CSV** (via `ddrs sources save/use` + or editing `ddrs.yaml`); adjacency stores need no rebuild — subgraphs are + keyed by STAID and already exist for every selected gauge. +- **Verify with the run log**, not assumptions: expect + `gages_adjacency filter: kept 1841 (dropped 0 missing, 0 headwater)` — + nonzero drops mean the CSV and the filters disagree and the builder's + assumptions no longer hold. +- **Spot-check recovered gauges before trusting a training run.** The 202 + final gauges recovered by the relative criterion (median 14,189 km²) have + relative DA disagreement up to 10% by construction; compare a sample's + observed hydrographs against summed upstream Q' (this remains open as of + 2026-08-02). diff --git a/.claude/skills/ddrs-dev/references/research-status.md b/.claude/skills/ddrs-dev/references/research-status.md index 4b699af..c812ed9 100644 --- a/.claude/skills/ddrs-dev/references/research-status.md +++ b/.claude/skills/ddrs-dev/references/research-status.md @@ -29,6 +29,7 @@ Most wrong numbers in this repo are population confusions, not arithmetic errors | **Training / eval set** | **2,365** | after the `gages_adjacency` filter (dropped 494 headwater). **Every trained median is on this set** | | Post-fix baseline population | 2,698 | 3,211 − 513 headwater. `ddrs plan` baselines from 2026-07-29 onward | | Global matched set | 5,224 | a **different network** (global MERIT). Only in `6_19_26_journal.md` (repo root, not `docs/`) | +| Area-balanced set (2026-08-02) | 1,841 | `~/projects/ddr/references/gage_info/gages_2000_area_balanced.csv`, built by `scripts/build_gages_2000_area_balanced.py` (seed 42) from GAGES-II with `DA_VALID` recomputed as **relative** `ABS_DIFF/DRAIN_SQKM ≤ 10%`, ≥80% obs coverage in both the 1981-10→1995-09 and 1995-10→2010-09 windows, non-headwater subgraph required. All 582 basins ≥5,000 km² kept + 418 random from [1k,5k) + 841 (all available) <1,000 km² → 45.7%/54.3% either side of 1,000 km². **Metrics on this set are incomparable to every 2,365/2,698-gauge number**; switching `data_sources.gages` to it invalidates the cached summed-Q′ baseline (`ddrs plan` recomputes) | ## Benchmarks — CONUS, eval 1995-10-01 → 2010-09-30, 2,365 gauges @@ -211,6 +212,86 @@ was invalidated by a stale binary and the manifest did not reveal it. ## Open, not closed +- **tau is mis-set (pilot-strength, 2026-08-06).** WY1996 sweep on the epoch-30 + area-balanced checkpoint: NSE(tau) plateaus at tau ≈ 14–19 in every bin + < 30,000 km²; a single global tau=19 gains +0.114 median NSE (0.546 → 0.660, + 1,841 gauges, WY1996) and tau=16 ties the summed-Q' baseline in the + < 1,000 km² bin (0.677 vs 0.674). Sign convention: window offset vs the + scored day's UTC midnight is (tau − 11) h; larger optimal tau ⇒ model LATE + vs obs. **Adversarial-review correction (same day): this run had the disagg + head OFF (flat repeat-24), so the hourly signal is 97.6% UTC-day-constant and + the sweep resolves only a day-pairing + blend weight (~half-day), NOT + sub-daily phase — do not read tau=16/19 as Eastern/Pacific midnight.** The + robust covariate is drainage area (Spearman +0.16 uncensored, an + accumulated-lag signature); the longitude/timezone fingerprint is + absent-to-contradicted (sign flips uncensored). 596 of 661 tau=23 pins are + real optima beyond the sweep edge ⇒ Phase 2 needs the ±1-day mapping + extension, split-sample selection, and a re-sweep on a disagg-ON run to test + for any hour-scale signal. Training also runs at tau=3, so gradients have + always been ~half a day misaligned — retrain at corrected tau is the open + test (freeze the tau protocol first). Instrument: `DDRS_HOURLY_DUMP` env var + on `evaluate` + `scripts/tau_sweep.py`. Authority: + `docs/2026-08-06-tau-sweep-pilot-findings.md` incl. §5a corrections. + **Interpolation arms (§5c, same day):** replicated on the standard 2,365-gauge + population — argmax tau=18–20, small-basin global tau=19 beats baseline + 0.674 vs 0.645 (WY1996). Linear/quadratic q' upsampling + (`DDRS_QPRIME_INTERP`, commit e4fb66d) neither sharpens nor shifts the + curves ⇒ the mis-set is not a step-function artifact, and interpolation is + NOT the fix (nearest ≥ linear ≥ quadratic at each optimum). ~30% of gauges + pin at tau=23 (optima beyond +12 h) and best_tau correlates with area, not + longitude ⇒ likely convention offset + area-growing lag ("double routing" — + DDR's own tau docstring). Mechanistic prior: USGS obs are LST + midnight-to-midnight (no DST), AORC/Q' stores are UTC; `src/data/` has no + timezone logic anywhere (§5b). + **Cross-source arms (§5e/§5f, 2026-08-07):** same checkpoint/network, + streamflow store swapped (`config/experiments/tau_src_*.yaml`). Four of + five stores replicate the tau 17–21 optimum (aorc2f distributed 20, UH + retro 21, daily LSTM 17, hourly-native LSTM 19; per-gauge best-tau median + 18–19 on all four) ⇒ the lag is shared upstream of store choice — and the + routing-free discriminator (§5g) resolved the attribution: it is the + common MC routing, NOT the day convention. + **Exception: `aorc2f_lumped`** — obs-free cross-correlation shows its + routed hydrographs LEAD the reference by ~23 h; CF metadata are + byte-identical to the distributed store (convention candidate REFUTED), so + the shift is in the data the lumped pipeline wrote. Do not use it for + timing-sensitive comparisons. The hourly-native arm does NOT sharpen the + curve (not a floor effect) and its longitude correlation is a censoring + artifact (interior sign +0.075, wrong sign) — timezone fingerprint remains + absent-to-contradicted even with native sub-daily data. Fable review + (§5f): claims sound-with-corrections; extended sweep −13..47 shows the + constant-tau optimum is interior (20) but 33% of per-gauge optima lie + beyond tau=23. Plot: `output/tau_sweep/cross_source_nse_vs_tau.png`. + **Routing-free summed-q' sweep (§5g, 2026-08-07) — the attribution + answer.** Summed daily q' repeat-24'd in the arms' phase, same sweep + machinery (tau=11 reproduces the cached baseline NSE, median |diff| + 0.0003). Optimum tau=6 global; per-bin 9 / 3 / −3 / −8 with area ⇒ + (1) **UTC-vs-LST convention REFUTED as dominant** (zero-area intercept + ≈ tau 10–11 ⇒ convention offset ≈ 0–2 h; explains every failed longitude + fingerprint); (2) summed q' LEADS the gauge by area-growing travel time + (2→19 h), so day-aligned scoring understates baseline skill in large + basins; (3) **MC routing over-delays by ≈2× the required travel time** + (routed lateness = summed earliness bin-by-bin; added delay 10→38 h) — + the measured "double routing" of DDR's tau docstring; tau≈19–20 is + compensation, root cause is routing timing (double-carried travel time + and/or slow trained celerity, median n 0.130 vs reference 0.05). + Routed-at-optimum still beats summed-at-optimum in every bin (+0.013 to + +0.051). Plot: `output/tau_sweep/summed_qprime_vs_routed_tau.png`. + **Convention change SHIPPED (2026-08-08, findings §5i):** tau is now + signed-at-zero hours of advance (`[tau : -(24-tau)]`, day i ↔ obs day i, + default 9 ≡ old 20); old scale = new + 11. Old checkpoints trained at + old-3 ≡ new −8. tau=9 CPU retrain across all five stores launched same + day (`scripts/run_tau9_source_trains.sh`). + **Gamma-UH params pulled (§5h):** the distributed aorc2f store's q' was + exported (2026-07-29, water_loss) with each divide routed through its own + learned gamma UH; `scripts/dump_gamma_uh_params.py` (water_loss venv) + reads the v3_gradaccum ep100 Ann head: median kernel mean 1.50 days per + divide (IQR 0.87–4.22), spearman +0.383 with log uparea ⇒ the per-divide + q' carries area-dependent NETWORK travel time before MC routes at all — + double routing is structural in the store. Clean fix target: sub-grid-only + UH on lateral inflows; col-7 (unrouted) is not it (0.29 routed, + water_loss 2026-07-29 finding). Dump: + `output/tau_sweep/gamma_uh_params.csv`. + - **Backward CUDA graphs (SP-11).** Forward capture landed (V7a 0.385, V10 29.2% launch reduction); the backward pass is not captured. Path: profile → fuse backward kernels → capture. Blocked for leakance configs (the leakance kernel has no capture diff --git a/.claude/skills/ddrs-eval-plots/SKILL.md b/.claude/skills/ddrs-eval-plots/SKILL.md index 1e27816..0a67504 100644 --- a/.claude/skills/ddrs-eval-plots/SKILL.md +++ b/.claude/skills/ddrs-eval-plots/SKILL.md @@ -33,6 +33,13 @@ Companion cells: distribution histogram, parameter vs log10(drainage area) hexbi **§Convergence: has training actually moved the parameters?** — per-epoch dumps, the four diagnostics, template · Notes +**`references/channel_geometry.md`** +Baseflow width & depth for all 346,321 CONUS reaches from the learned `n`, `p`, +`q` and post-clamp slope · CONUS maps of width / depth / `w:d` · **downstream +hydraulic-geometry exponents vs Leopold & Maddock** — the only internal test of +whether `p_spatial`/`q_spatial` are physically sensible, since the attributes +carry no width or depth to validate against · plausibility bands + **`references/parity.md`** (235 lines) Init-time parity (when to use, inputs, load → histograms → pass/fail, KS criterion) · Trained parity (inputs, load → per-distribution stats → histograms → per-reach @@ -119,6 +126,7 @@ boundary. | NSE, KGE, bias, RMSE, FHV, FLV, CDF, box plot, "did it beat the baseline" | **metrics** | `references/metrics.md` | | Manning's n, p_spatial, q_spatial, slope, map, basin, spatial pattern | **parameter_map** | `references/parameter_map.md` | | "have the parameters converged", epoch drift, movement across epochs | **parameter convergence** | `references/parameter_map.md` §Convergence | +| width, depth, channel geometry, w:d ratio, "are the geometry parameters right", hydraulic geometry, Leopold & Maddock | **channel_geometry** | `references/channel_geometry.md` | | DDR-vs-ddrs parameter distributions, at init or trained | **parity** | `references/parity.md` | Vague request ("plot my trained model")? Offer the default bundle: @@ -246,6 +254,8 @@ was still ruled NO-GO — passing it is necessary, not sufficient. - `references/metrics.md` — NSE/KGE/bias distributions, CDFs, box plots vs baseline - `references/parameter_map.md` — learned parameters over MERIT polygons, plus epoch-to-epoch convergence drift +- `references/channel_geometry.md` — baseflow width/depth over MERIT, plus the + Leopold & Maddock exponent check on `p_spatial`/`q_spatial` - `references/parity.md` — DDR-vs-ddrs parameter distributions at init and trained - `scripts/load_ddrs_predictions.py` — **always use this** to open the predictions zarr and the f32 baseline. It handles two pitfalls every notebook otherwise hits: diff --git a/.claude/skills/ddrs-eval-plots/references/channel_geometry.md b/.claude/skills/ddrs-eval-plots/references/channel_geometry.md new file mode 100644 index 0000000..ae9f281 --- /dev/null +++ b/.claude/skills/ddrs-eval-plots/references/channel_geometry.md @@ -0,0 +1,233 @@ +# Reference: baseflow channel geometry (width & depth) from learned parameters + +Turns the learned KAN parameters into **physical channel width and depth at +baseflow** for all 346,321 CONUS MERIT reaches, maps them, and — the reason this +exists — tests whether the channel-geometry parameterisation reproduces observed +**downstream hydraulic geometry**. That test is the only internal check we have +on whether `p_spatial` / `q_spatial` are physically sensible, because the +attribute NetCDF carries **no width or depth variable** to validate against +(29 vars, none of them width/depth — checked 2026-08-03). + +## The equations (mirrors `src/geometry.rs:37-67`) + +Given Manning's `n`, Leopold-Maddock `p` and `q`, bed slope `S`, and discharge `Q`: + +``` +depth = ((Q · n · (q+1)) / (p · √S))^(3 / (5 + 3q)) geometry.rs:40-44 +top_width = p · depth^q geometry.rs:47 +side_slope = clamp(top_width · q / (2·depth), 0.5, 50) geometry.rs:50 +bottom_width = clamp(top_width − 2·side_slope·depth, 0.01) geometry.rs:53 +area = (top_width + bottom_width) · depth / 2 geometry.rs:57 +wetted_perimeter = bottom_width + 2·depth·√(1+z²) geometry.rs:60 +R = area / wetted_perimeter geometry.rs:64 +velocity = n⁻¹ · R^(2/3) · √S geometry.rs:67 +``` + +The depth relation is the exact inversion of Manning + the power-law width +profile `w(y) = p·y^q`, whose area is `A = p·d^(q+1)/(q+1)` — the `(q+1)` is that +normalisation, not a fudge factor. + +**Use the post-clamp slope** from `plot/kan_parameters.nc`'s `slope` variable, +not the raw fabric slope: `attribute_minimums.slope` (default 1e-3) is what the +solver actually applied, and 33.2% of CONUS reaches are pinned at it. + +## Choosing a baseflow discharge + +There is no per-reach baseflow in the parameter NetCDF, so pick one explicitly +and **state the assumption in the notebook**: + +```python +Q_SPEC = 0.005 # m3/s per km2 — CONUS-ish baseflow specific discharge +Q = Q_SPEC * 10**attrs["log10_uparea"] +``` + +`0.005` is roughly half the ~0.01 m³/s/km² CONUS *mean* specific discharge. Sweep +it (0.002 / 0.005 / 0.01) — the hydraulic-geometry **exponents below are +invariant to `Q_SPEC`** (it is a constant multiplier inside a log-log fit), so +only the absolute widths and depths move. That invariance is itself a useful +sanity check that the fit is wired correctly. + +## The diagnostic that matters: downstream hydraulic geometry + +Leopold & Maddock (1953) established that, moving *downstream* through a network: + +``` +w ∝ Q^b b ≈ 0.50 +d ∝ Q^f f ≈ 0.40 +v ∝ Q^m m ≈ 0.10 with b + f + m = 1 identically +``` + +Fit `b` and `f` by least squares on `log10(w)` vs `log10(Q)` and `log10(d)` vs +`log10(Q)` across all reaches, then compare. + +### Measured on run `2026-08-03T13-11-00Z` (346,321 reaches, `Q_SPEC = 0.005`) + +``` + fitted Leopold & Maddock +width exponent b 0.226 ~0.50 +depth exponent f 0.600 ~0.40 +velocity m=1-b-f 0.174 ~0.10 + +depth (m) p5 0.058 p50 0.250 p95 3.733 +width (m) p5 4.560 p50 6.497 p95 21.997 +w/d p5 5.41 p50 26.6 p95 79.2 +``` + +**The parameterisation is structurally incapable of reaching `b ≈ 0.50`.** With `p` +held constant, substituting `d ∝ Q^(3/(5+3q))` into `w = p·d^q` gives: + +``` +d exponent = 3/(5+3q) w exponent = 3q/(5+3q) + + q = 0.3 : d 0.508 w 0.153 + q = 0.5 : d 0.462 w 0.231 + q = 1.0 : d 0.375 w 0.375 <- q's upper bound +``` + +The width exponent **maxes out at 0.375 when `q = 1`**, below L&M's 0.50, and the +depth exponent bottoms out at 0.375, above L&M's 0.40 only for `q > 1`. Spatial +variation in `p` (learned, ρ ≈ +0.33 with `log10_uparea`) shifts the realised fit +but did not close the gap: 0.226 against 0.50. + +**Interpretation.** Channels are modelled as too narrow and too deep, and +increasingly so downstream. That is not cosmetic — depth drives `R`, `R` drives +velocity, velocity drives celerity and hence `K = L/c` and the whole routing +timescale. A deep-narrow bias inflates `R` for a given area and therefore +inflates velocity. + +**Before concluding the KAN mis-learned `q`:** the cap is a property of the +`w = p·d^q` form itself, not of the fitted values. Reaching `b ≈ 0.5` requires +either `q > 1` (outside `parameter_ranges.q_spatial: [0, 1]`) or a `p` that +grows with `Q` strongly enough to make up the difference. Widening the `q` range +is the cheap experiment; changing the width law is the real fix. + +## Notebook template + +```python +from pathlib import Path +import geopandas as gpd, matplotlib.pyplot as plt, numpy as np, pyogrio, xarray as xr +from mpl_toolkits.axes_grid1 import make_axes_locatable + +RUN_DIR = Path("/home/tbindas/projects/ddrs/.ddrs/runs/") +PARAMS_NC = RUN_DIR / "plot" / "kan_parameters.nc" +ATTRS_NC = Path("/home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc") +SHAPEFILE = Path("/home/tbindas/projects/ddr/data/merit/" + "cat_pfaf_7_MERIT_Hydro_v07_Basins_v01_bugfix1.shp") +PLOT_DIR = RUN_DIR / "plots"; PLOT_DIR.mkdir(exist_ok=True) +Q_SPEC = 0.005 # m3/s per km2 — STATE THIS ASSUMPTION IN THE WRITEUP +CONUS_BBOX = (-125, 24, -66, 53) + +ds, attrs = xr.open_dataset(PARAMS_NC), xr.open_dataset(ATTRS_NC) +sh, ia, ib = np.intersect1d(attrs.COMID.values, ds.COMID.values, return_indices=True) +A = 10.0 ** attrs["log10_uparea"].values[ia] # km2 +Q = Q_SPEC * A +n = ds["n"].values[ib] +p = ds["p_spatial"].values[ib] +q = ds["q_spatial"].values[ib] +S = np.maximum(ds["slope"].values[ib], 1e-3) # POST-clamp slope +comid = ds.COMID.values[ib] + +# --- geometry, mirroring src/geometry.rs exactly --------------------------- +qe = q + 1e-6 +depth = np.maximum(((Q * n * (qe + 1.0)) / (p * np.sqrt(S) + 1e-8)) ** (3.0 / (5.0 + 3.0 * qe)), 0.01) +width = p * depth ** qe +z = np.clip(width * qe / (depth * 2.0), 0.5, 50.0) +bw = np.maximum(width - z * depth * 2.0, 0.01) +area = (width + bw) * depth / 2.0 +wp = bw + depth * np.sqrt(z ** 2 + 1.0) * 2.0 +R = area / wp +v = (1.0 / n) * R ** (2.0 / 3.0) * np.sqrt(S) + +# --- hydraulic-geometry exponents ------------------------------------------ +ok = np.isfinite(depth) & np.isfinite(width) & (Q > 0) +lq = np.log10(Q[ok]) +b_fit = np.polyfit(lq, np.log10(width[ok]), 1)[0] +f_fit = np.polyfit(lq, np.log10(depth[ok]), 1)[0] +print(f"b (width) {b_fit:.3f} vs L&M 0.50 | f (depth) {f_fit:.3f} vs 0.40 " + f"| m (velocity) {1-b_fit-f_fit:.3f} vs 0.10") +print(f"structural cap at q=1: b_max = {3*1.0/(5+3*1.0):.3f}") + +# --- log-log panels with L&M reference slopes ------------------------------ +fig, axes = plt.subplots(1, 2, figsize=(14, 6), dpi=150) +for ax, y, lab, fit, ref in ((axes[0], width[ok], "width (m)", b_fit, 0.50), + (axes[1], depth[ok], "depth (m)", f_fit, 0.40)): + hb = ax.hexbin(lq, np.log10(y), gridsize=80, cmap="viridis", mincnt=1, bins="log") + fig.colorbar(hb, ax=ax, label="reach count") + xs = np.linspace(lq.min(), lq.max(), 10) + c = np.median(np.log10(y)) - fit * np.median(lq) + ax.plot(xs, fit * xs + c, "r-", lw=2, label=f"fitted {fit:.3f}") + ax.plot(xs, ref * xs + c, "w--", lw=2, label=f"L&M {ref:.2f}") + ax.set_xlabel(r"$\log_{10}$ Q (m$^3$/s)"); ax.set_ylabel(f"log10 {lab}") + ax.legend(); ax.grid(alpha=0.3) +fig.suptitle(f"Downstream hydraulic geometry at baseflow (Q_spec={Q_SPEC})", fontsize=14) +fig.tight_layout() +fig.savefig(PLOT_DIR / "geometry_hydraulic_exponents.png", dpi=250, + bbox_inches="tight", facecolor="white") + +# --- CONUS maps ------------------------------------------------------------ +gdf = pyogrio.read_dataframe(SHAPEFILE, columns=["COMID"]).set_index("COMID") +if gdf.crs is None: + gdf = gdf.set_crs(epsg=4326) # cat_pfaf_7 ships without a .prj +import pandas as pd +for name, arr, cmap, unit in (("width", width, "Blues", "m"), + ("depth", depth, "Blues", "m"), + ("wd_ratio", width/depth, "magma", "-")): + gdf[name] = pd.Series(arr, index=comid).reindex(gdf.index).values +conus = gdf.cx[CONUS_BBOX[0]:CONUS_BBOX[2], CONUS_BBOX[1]:CONUS_BBOX[3]] +for name, unit in (("width", "m"), ("depth", "m"), ("wd_ratio", "-")): + g = conus.dropna(subset=[name]).sort_values(name) + lo, hi = np.nanpercentile(g[name], [2, 98]) + fig, ax = plt.subplots(figsize=(14, 8), dpi=150) + g.plot(ax=ax, column=name, cmap="Blues" if unit == "m" else "magma", + linewidth=0.0, vmin=lo, vmax=hi, zorder=1) + try: + import contextily as cx + cx.add_basemap(ax, crs=g.crs, source=cx.providers.CartoDB.Positron, + alpha=0.6, zorder=0, attribution=False) + except Exception as e: + print(f"basemap skipped ({type(e).__name__})"); ax.set_facecolor("#f0f0f0") + ax.set_xlim(CONUS_BBOX[0], CONUS_BBOX[2]); ax.set_ylim(CONUS_BBOX[1], CONUS_BBOX[3]) + ax.set_xticks([]); ax.set_yticks([]) + ax.set_title(f"Baseflow {name} — CONUS (2nd-98th pct colour)", fontsize=13) + cax = make_axes_locatable(ax).append_axes("right", size="3%", pad=0.1) + sm = plt.cm.ScalarMappable(cmap="Blues" if unit == "m" else "magma") + sm.set_array([]); sm.set_clim(lo, hi) + fig.colorbar(sm, cax=cax).set_label(f"{name} ({unit})") + fig.savefig(PLOT_DIR / f"geometry_map_{name}_conus.png", dpi=250, + bbox_inches="tight", facecolor="white") + plt.close(fig) +``` + +## Plausibility bands + +No ground-truth widths exist in the attributes, so sanity-check against +literature ranges rather than data: + +| quantity | plausible | run `2026-08-03T13-11-00Z` | +|---|---|---| +| `w/d` natural channels | 10–50 (up to ~100 for braided) | p50 **26.6**, p95 79.2 | +| baseflow depth, headwaters | 0.05–0.5 m | p5 0.058, p50 **0.250** | +| baseflow width, headwaters | 1–15 m | p5 4.56, p50 **6.50** | +| `b + f + m` | **exactly 1** | 1.000 by construction | + +`b + f + m = 1` is an identity, not a test — it holds regardless of whether the +parameters are any good. Do not report it as validation. + +## Notes + +- **Exponents are `Q_SPEC`-invariant; absolute widths and depths are not.** Sweep + `Q_SPEC` and confirm `b`/`f` do not move — if they do, the fit is broken. +- **Use post-clamp `slope`** from the parameter NetCDF. Using raw fabric slope + gives geometry the solver never saw. +- **This is baseflow, not bankfull.** L&M's downstream exponents are usually + quoted at a consistent frequency (often bankfull or mean annual). Comparing a + low-flow geometry to bankfull exponents is defensible for the *exponents* + (which are scale-free) but not for absolute widths. +- **The `w = p·d^q` form caps `b` at 0.375.** If the goal is matching observed + downstream hydraulic geometry, that is a parameterisation change, not a + training problem. Widening `parameter_ranges.q_spatial` past 1.0 is the cheap + probe; note it also changes the celerity β, which depends on `q` through the + trapezoid closure (`side_slope = T·q/(2d)`). +- **Cross-check against the routing diagnostics.** A deep-narrow bias inflates + `R` and hence velocity and celerity; if `median_n` is also low, the two + compound. See `references/parameter_map.md` §Convergence. diff --git a/CLAUDE.md b/CLAUDE.md index 0a49ca2..f7ebe3f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -454,6 +454,60 @@ Campaign docs (chronological): `docs/2026-07-04-synthetic-recoverability-findings.md`, `docs/2026-07-06-phase-c-findings.md`. +## Reach subdivision (`params.subdivision`, off by default) + +**What it is.** A build-time (not runtime) normalization of reach length toward +`Δx ≈ c_ref·Δt` inside the managed adjacency builder, so `Cr = Δt/K` lands near +1. Two-sided: reaches longer than `Δx_target` are **split** into +`m = min(ceil(L/Δx_target), max_pieces)` pieces of length `L/m` with `q' → q'/m`; +shorter reaches have their **length clamped up** (never merged — merging would +destroy junction topology), bounded by `max_clamp_factor`. The runtime just sees +a bigger graph: no autograd change, no gradient path, `mmc_op.rs` untouched. + +**STATUS: NO-GO for its stated purpose (2026-08-05).** It was built to make the +Muskingum coefficients non-negative by construction and retire +`enforce_positivity`. Measured on 1,841 CONUS gauges with `enforce_positivity` +off, `frac c1 < 0` gets **worse** (93.0 % off → 98.79 % at `max_pieces: 8`) and +negative solves fall only 35 % for a 2.05× network, 1.5× step time and +23.9 % +total channel length. Reason: both coefficients are non-negative only inside +`[2X, 2(1−X)]`, a window of width `2(1−2X)` — **1.4 % wide at the measured CONUS +median X = 0.4966** — and a static piece count cannot hold a flow-varying `Cr` +inside it. It *does* nearly eliminate `Cr > 2` / `c3 < 0` (3.93 % → 0.31 %), via +the length clamp rather than the splitting. The code is correct, gated off, and +**stays in-tree** as the measurement apparatus. Do not re-open the "Cr ≈ 1 ⇒ +non-negative" argument without reading `.claude/REACH-SUBDIVISION.md`. + +**How to enable** (`params.subdivision.enabled: true`), and what will reject you: + +1. **Requires `geospatial_fabric`.** Subdivision runs inside the managed + adjacency builder, which explicit `conus_adjacency`/`gages_adjacency` paths + bypass — so `enabled: true` alongside them is a **config error**, not a + warning (`validate_subdivision_reaches_the_builder`, `src/config.rs:1066`). + Otherwise the flag would be *silently inert* while the manifest claimed + subdivision. The one allowed exception is an explicit path to a store already + built subdivided, detected from zarr metadata (`n_parent < n`). +2. **Requires `use_cuda_graphs: false`** — a captured graph is sized to a fixed + reach count. +3. **Requires retraining.** Every learned parameter was fit against the un-split + network's effective diffusion; checkpoints do not transfer. + +Fields (all seven are hashed into the adjacency cache key, so editing any one +rebuilds the graph): `enabled` (false), `max_pieces` (8 — uncapped is infeasible: +13.2× reaches, 9.2× solver critical path, and `Σm` cannot be pinned down), +`reference_n` (0.05 — **a guess; the trained CONUS median is 0.130**, and this +sets `dx_target` directly, so sweep it), `reference_discharge_coefficient` +(0.01), `reference_discharge_exponent` (0.9), `min_length_fraction` (1.0; 0 +disables the short-reach clamp), `max_clamp_factor` (4.0 — unbounded, measured +clamp factors reached 48,597×). + +Implementation: `src/adjacency/subdivide.rs` + `src/adjacency/cache.rs`; +persistence of `parent_order`/`parent_offset` in `src/data/store/zarr.rs` +(`IdIndex` is built from `parent_order`, since `order` gains duplicates). +Gates: `cargo test --test subdivide --test subdivision_integration --test +adjacency_parity --test gauge_mass_conservation`, plus `compare_ddr_sandbox` +staying an ABSOLUTE MATCH. Design, measurements and gotchas: +`.claude/REACH-SUBDIVISION.md`. + ## Baseline `ddrs plan` and `ddrs run --workflow train-and-test` compute a **summed Q'** diff --git a/config/experiments/tau9_train_aorc2f_dist.yaml b/config/experiments/tau9_train_aorc2f_dist.yaml new file mode 100644 index 0000000..6d1de03 --- /dev/null +++ b/config/experiments/tau9_train_aorc2f_dist.yaml @@ -0,0 +1,243 @@ +# tau=9 cross-source retrain arm: aorc2f_dist +# Derived from the 2026-08-05T04-58-58Z epoch-30 run's config snapshot with +# exactly three changes: (1) streamflow store swapped per arm, (2) +# sparse_solver cpu (CPU train-and-test set, 2026-08-08), (3) params.tau +# left UNSET so the 2026-08-08 default applies: tau=9 on the NEW convention +# (hours of advance; == old-convention 20, the measured optimum). The +# epoch-30 baseline trained at old tau=3 (== new -8, wrong direction). +# ddrs MERIT training config. +# +# Hyperparameters are pulled verbatim from +# ~/projects/ddr/config/merit_training_config.yaml. +# Forcing is daily; the loader interpolates to hourly (repeat × 24, trim to +# n_hourly) before feeding the Muskingum-Cunge engine — same pattern as +# DDR's StreamflowReader. + +mode: training +workflow: train-and-test # ddrs plan/run picks this up; override with --workflow X +geodataset: merit +seed: 42 +np_seed: 42 + +# Source paths — read in place by ddrs's Rust loaders (zarrs + icechunk + netcdf). +# No export/materialization step; the harness reads DDR's live data sources +# straight from disk. +data_sources: + attributes: /home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc + conus_adjacency: /home/tbindas/projects/ddr/data/merit_conus_adjacency.zarr + gages_adjacency: /home/tbindas/projects/ddr/data/merit_gages_conus_adjacency.zarr + streamflow: /mnt/ssd1/data/icechunk/daily_dhbv2_distributed_aorc2f_merit_unit_catchments.ic + observations: /mnt/ssd1/data/icechunk/usgs_daily_observations + gages: /home/tbindas/projects/ddr/references/gage_info/gages_2000_area_balanced.csv + # Hourly AORC precip (zarr v3, CONUS) — drives the precip-conditioned + # mass-preserving disaggregation head below. + aorc_precip: /mnt/ssd1/data/aorc/merit_unit_catchments.zarr + +experiment: + batch_size: 64 + start_time: 1981/10/01 + end_time: 1995/09/30 + epochs: 30 + rho: 90 + shuffle: true + warmup: 5 + # Gradient accumulation ON. `batch_size` is now the MICRO-batch; the optimizer + # steps once per `grad_accum_steps` micro-batches, so the effective batch is + # 64 * 4 = 256 gauges and the gradient is averaged over ~4x more basins. + # + # THE ARITHMETIC — read this before trusting a run: + # 1,841 gauges / 64 = 28 full + 1 partial = 29 micro-batches per epoch + # (accumulation sets `drop_last = false`, driver.rs:268-275, so the + # 49-gauge tail is KEPT — unlike the non-accumulating path) + # 29 / 20 = one group of 20 + one partial group of 9. + # A partial group STILL STEPS: only `micros_drawn == 0` breaks the loop + # (driver.rs:415), and the 1/Σn renormalization makes unequal group sizes + # exact. So this is 2 optimizer steps per epoch, not 1. + # x 30 epochs = 60 TOTAL UPDATES + # + # Effective batch is 20*64 = 1,280 of 1,841 gauges — 70% of the training set + # in a single gradient, i.e. near full-batch. Very low variance, very few steps. + # + # THAT IS THE RISK. Calibration: the run that beat the summed-Q' baseline + # (2026-07-30T00-24, median NSE 0.6799) took 180 updates at lr 1e-3 and moved + # the head L2(dw) = 9.75. The 2026-08-04 run that hit NSE 0.6782 took 280. + # Adam's per-weight displacement is bounded by ~lr per step, so the budget + # here is 40*0.001 + 20*0.0005 = 0.050 — about 38% of the 0.132 that the + # 180-step run realized. Expect under-training. + # + # If it under-trains, prefer more STEPS over a hotter lr: + # * `grad_accum_steps: 20` -> 7 (5 steps/epoch, 150 updates, eff batch 448) + # * `grad_accum_steps: 20` -> 4 (8 steps/epoch, 240 updates, eff batch 256) + # Wall clock is forward-pass bound (~27.5 s/micro-batch), so all of these cost + # the SAME time — 29 forwards per epoch regardless. Lowering accum_steps buys + # optimizer steps for free; it only trades gradient variance for step count. + # + # The historical `grad_accum_steps: 37` was WORSE than either: with 28-37 + # micro-batches per epoch it collapsed a whole epoch into ONE update, so 50 + # epochs bought 50 updates for 1,850 forward passes. Wall clock here is + # forward-pass bound (~27.5 s/micro-batch), not step bound, so accumulating + # costs nothing in time — it only trades update count for gradient quality. + # + # ALWAYS verify the head actually moved (L2(dw)) before believing a result: + # the 2026-07-31 AdaDelta run reported plausible losses while L2(dw) = 1.9e-4, + # i.e. it never left initialization. + use_grad_accum: true + grad_accum_steps: 20 + # Adam, reverting the 2026-07-31 AdaDelta run. That run's head moved + # L2(dw) = 1.9e-4 (0.0004% of ||w||) across all 30 epochs — it never left + # initialization. Cause: `experiment.learning_rate` IS applied to AdaDelta + # (driver.rs:404 passes lr to optimizer.step; adadelta.rs:123 multiplies + # delta by it) despite config.rs:211-213 documenting the key as inert, so + # every update ran 1000x scaled down. dHBV intends AdaDelta at lr=1.0. + optimizer: adam + loss: + kind: nse-batch + # Step schedule. Keys are the FIRST epoch at which the lr applies + # (`resolve_lr` takes the largest key <= epoch, 1-indexed). At 2 optimizer + # steps per epoch (see the accumulation note above): + # epochs 1-20 @ 0.001 (40 updates) + # epochs 21-30 @ 0.0005 (20 updates) + # Total displacement budget ~0.050. The decay lands 2/3 of the way through, + # which is the right shape — the concern is the step COUNT, not the schedule. + # + # Watch `n_at_floor` in the log: the lr=0.01 run pinned 47.3% of CONUS at the + # 0.015 floor, and both 2026-08-04 runs collapsed from 0% to 25-36% between + # epochs 3 and 5 at lr 2e-3. Starting at 1e-3 is deliberately cooler than that. + # Near-full-batch gradients (1,280 gauges) should also damp the collapse, since + # it looked like a high-variance small-batch effect — but that is a hypothesis, + # not an established result. If `n_at_floor` climbs past a few percent, the lr + # is still too hot. + learning_rate: + 1: 0.005 + 11: 0.001 + 21: 0.0005 + grad_clip_max_norm: 1.0 + +kan_head: + hidden_size: 21 + num_hidden_layers: 2 + # B-spline grid intervals (`num` in pykan). Matches DDR's + # ~/projects/ddr/config/merit_training_config.yaml::kan.grid. + grid: 50 + # B-spline order. Matches DDR's ::kan.k. (pykan's KANLayer default is 3, + # but DDR overrides to 2 for production.) + k: 2 + input_var_names: + - SoilGrids1km_clay + - aridity + - meanelevation + - meanP + - NDVI + - meanslope + - log10_uparea + - SoilGrids1km_sand + - ETPOT_Hargr + - Porosity + learnable_parameters: + - n + - q_spatial + - p_spatial + # Precip-conditioned mass-preserving daily->hourly disaggregation head. + # The within-day shape is driven by the hourly AORC precip window + # [d-1,d,d+1] (+ daily-Q taps + static attrs); the daily mean is conserved + # exactly. Loss stays the historical L1 (no experiment.loss block). + # Warm-started FROZEN from the standalone capacity pretrain (chunk_days=1); + # architecture fields must match that checkpoint or load_record fails. + # `enabled: false` makes this whole block inert — the loader uses flat + # repeat-24 (nearest) upsampling instead. Flip to true for the head-on + # arm of the ablation; nothing else needs to change. + disaggregation: + enabled: false + hidden_size: 16 + num_hidden_layers: 2 + grid: 20 + k: 3 + boundary_blend: 0.0 + chunk_days: 1 + pretrained_checkpoint: /home/tbindas/projects/ddrs/output/disagg_pretrain/capacity_chunk1.mpk + freeze: true + +# Routing-engine knobs — DDR's MERIT defaults (matches mock_config in tests). +params: + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] + attribute_minimums: + discharge: 1.0e-4 + slope: 1.0e-3 + velocity: 0.01 + depth: 0.01 + bottom_width: 0.01 + defaults: + p_spatial: 21.0 + log_space_parameters: + - p_spatial + # Corrected physics instead of DDR's formulation. With `false` the routing + # core uses: + # * exact trapezoidal celerity c = v·β, β = 5/3 − (4/3)·A·√(1+z²)/(T·P) + # instead of the wide-rectangular 5/3 limit (which is 22-27% high for the + # channels this code builds — κ = b/y ≈ 0.7-1.8, so β ≈ 1.31-1.36), and + # * Cunge-derived Muskingum X = clamp(0.5(1 − Q/(T·S·c·L)), 0, 0.5) instead + # of the constant 0.3, which was injecting a median 28× excess numerical + # diffusion. + # Negative-discharge counting before the S28 clamp is reported in BOTH modes. + # This breaks `compare_ddr_sandbox`'s ABSOLUTE MATCH by design — that example + # must be run with the default `ddr_match: true`. + # See `.claude/PHYSICS-CORRECTIONS.md`. + ddr_match: false + # Provably zero negative discharges. Clamps the Muskingum INPUTS (K, X) into + # the non-negative-coefficient window `2X <= Cr <= 2(1-X)`, Cr = dt/K: + # K <- max(K, dt*(1+d)/2) => Cr <= 2/(1+d) + # X <- min(X_cunge, (1-d)*Cr/2, (1-d)*(1-Cr/2)) + # with d = POSITIVITY_DELTA = 1e-2. Clamping K and X (not the coefficients) + # keeps `c1+c2+c3 = 1` exact, so mass is conserved; clamping c3 would not. + # With c1,c3 >= 0 and q_t,q' > 0, induction over the topological order of the + # forward substitution proves the whole solve is non-negative. + # + # Measured: 0 / 197,461,880 solves on 1,841 gauges (was 55,181, 0.0279%), + # replicated across two windows, two backends, and two trained heads. + # + # COST — this is not free. The X cap binds on 95.3% of reach-timesteps and + # median X falls 0.4976 -> 0.0794 (6.3x), so it largely REPLACES the Cunge X + # rather than shading it: numerical diffusion becomes stability-set instead of + # matched to hydraulic diffusivity. NSE/KGE impact is UNMEASURED, and this is + # the first training run to use the flag. Compare against a matched + # `enforce_positivity: false` control before trusting the result. + # + # Root-cause alternative: median Cr is 0.226, i.e. the typical MERIT reach is + # ~4.4x too LONG for an hourly step. Subdividing reaches ~4x brings Cr to ~1, + # where X_max -> 0.5 and the cap stops biting. See .claude/PHYSICS-CORRECTIONS.md. + # + # Requires ddr_match: false (rejected at load otherwise). + # OFF. It delivered provably zero negative solves (0 / 197,461,880), but the + # cure was worse than the disease: capping X collapsed the median from 0.4976 + # to 0.0794 on 95.3% of reach-timesteps, and because the cap is applied at + # RUNTIME to a learned celerity it made X ~ Cr ~ 1/n — handing the optimizer a + # lever it rode straight to the roughness floor. Same lr, flag on vs off: + # n_at_floor 35% -> 97.9%. Superseded by reach subdivision (a BUILD-TIME fix + # with no gradient path): see docs/superpowers/plans/2026-08-05-reach-subdivision.md + enforce_positivity: false + # cuSPARSE triangular solve. Pair with `ddrs run --backend cuda`; it is inert + # under `--backend cpu`. This is a DIFFERENT `Backward` impl from the host + # solve, sitting directly downstream of the ddr_match/enforce_positivity + # terms — gradient-checked on GPU by `tests/cuda_backward_parity.rs`. + # CPU set (2026-08-08): host triangular solve. + sparse_solver: cpu + # MUST be false when ddr_match is false: the fused CUDA-graph kernel + # (cuda_graph/geometry_kernel.rs) hardcodes DDR's 5/3 celerity, so a captured + # graph would replay the DDR forward while the backward used the corrected + # chain rule — a silent gradient mismatch. Config load rejects the pair. + use_cuda_graphs: false + +# Test-mode overlay. When the binary is run with --mode testing, these keys +# replace the corresponding ones in `experiment:`. Absent keys inherit. +# +# IMPORTANT: batch_size SEMANTIC SHIFTS between modes: +# - experiment.batch_size (training) = number of GAUGES per mini-batch +# - testing.batch_size = number of DAYS per chunk +testing: + start_time: 1995/10/01 + end_time: 2010/09/30 + batch_size: 15 # DAYS, not gauges + rho: null # disabled in test mode diff --git a/config/experiments/tau9_train_aorc2f_lumped.yaml b/config/experiments/tau9_train_aorc2f_lumped.yaml new file mode 100644 index 0000000..49a6335 --- /dev/null +++ b/config/experiments/tau9_train_aorc2f_lumped.yaml @@ -0,0 +1,243 @@ +# tau=9 cross-source retrain arm: aorc2f_lumped +# Derived from the 2026-08-05T04-58-58Z epoch-30 run's config snapshot with +# exactly three changes: (1) streamflow store swapped per arm, (2) +# sparse_solver cpu (CPU train-and-test set, 2026-08-08), (3) params.tau +# left UNSET so the 2026-08-08 default applies: tau=9 on the NEW convention +# (hours of advance; == old-convention 20, the measured optimum). The +# epoch-30 baseline trained at old tau=3 (== new -8, wrong direction). +# ddrs MERIT training config. +# +# Hyperparameters are pulled verbatim from +# ~/projects/ddr/config/merit_training_config.yaml. +# Forcing is daily; the loader interpolates to hourly (repeat × 24, trim to +# n_hourly) before feeding the Muskingum-Cunge engine — same pattern as +# DDR's StreamflowReader. + +mode: training +workflow: train-and-test # ddrs plan/run picks this up; override with --workflow X +geodataset: merit +seed: 42 +np_seed: 42 + +# Source paths — read in place by ddrs's Rust loaders (zarrs + icechunk + netcdf). +# No export/materialization step; the harness reads DDR's live data sources +# straight from disk. +data_sources: + attributes: /home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc + conus_adjacency: /home/tbindas/projects/ddr/data/merit_conus_adjacency.zarr + gages_adjacency: /home/tbindas/projects/ddr/data/merit_gages_conus_adjacency.zarr + streamflow: /mnt/ssd1/data/icechunk/daily_dhbv2_lumped_aorc2f_merit_unit_catchments.ic + observations: /mnt/ssd1/data/icechunk/usgs_daily_observations + gages: /home/tbindas/projects/ddr/references/gage_info/gages_2000_area_balanced.csv + # Hourly AORC precip (zarr v3, CONUS) — drives the precip-conditioned + # mass-preserving disaggregation head below. + aorc_precip: /mnt/ssd1/data/aorc/merit_unit_catchments.zarr + +experiment: + batch_size: 64 + start_time: 1981/10/01 + end_time: 1995/09/30 + epochs: 30 + rho: 90 + shuffle: true + warmup: 5 + # Gradient accumulation ON. `batch_size` is now the MICRO-batch; the optimizer + # steps once per `grad_accum_steps` micro-batches, so the effective batch is + # 64 * 4 = 256 gauges and the gradient is averaged over ~4x more basins. + # + # THE ARITHMETIC — read this before trusting a run: + # 1,841 gauges / 64 = 28 full + 1 partial = 29 micro-batches per epoch + # (accumulation sets `drop_last = false`, driver.rs:268-275, so the + # 49-gauge tail is KEPT — unlike the non-accumulating path) + # 29 / 20 = one group of 20 + one partial group of 9. + # A partial group STILL STEPS: only `micros_drawn == 0` breaks the loop + # (driver.rs:415), and the 1/Σn renormalization makes unequal group sizes + # exact. So this is 2 optimizer steps per epoch, not 1. + # x 30 epochs = 60 TOTAL UPDATES + # + # Effective batch is 20*64 = 1,280 of 1,841 gauges — 70% of the training set + # in a single gradient, i.e. near full-batch. Very low variance, very few steps. + # + # THAT IS THE RISK. Calibration: the run that beat the summed-Q' baseline + # (2026-07-30T00-24, median NSE 0.6799) took 180 updates at lr 1e-3 and moved + # the head L2(dw) = 9.75. The 2026-08-04 run that hit NSE 0.6782 took 280. + # Adam's per-weight displacement is bounded by ~lr per step, so the budget + # here is 40*0.001 + 20*0.0005 = 0.050 — about 38% of the 0.132 that the + # 180-step run realized. Expect under-training. + # + # If it under-trains, prefer more STEPS over a hotter lr: + # * `grad_accum_steps: 20` -> 7 (5 steps/epoch, 150 updates, eff batch 448) + # * `grad_accum_steps: 20` -> 4 (8 steps/epoch, 240 updates, eff batch 256) + # Wall clock is forward-pass bound (~27.5 s/micro-batch), so all of these cost + # the SAME time — 29 forwards per epoch regardless. Lowering accum_steps buys + # optimizer steps for free; it only trades gradient variance for step count. + # + # The historical `grad_accum_steps: 37` was WORSE than either: with 28-37 + # micro-batches per epoch it collapsed a whole epoch into ONE update, so 50 + # epochs bought 50 updates for 1,850 forward passes. Wall clock here is + # forward-pass bound (~27.5 s/micro-batch), not step bound, so accumulating + # costs nothing in time — it only trades update count for gradient quality. + # + # ALWAYS verify the head actually moved (L2(dw)) before believing a result: + # the 2026-07-31 AdaDelta run reported plausible losses while L2(dw) = 1.9e-4, + # i.e. it never left initialization. + use_grad_accum: true + grad_accum_steps: 20 + # Adam, reverting the 2026-07-31 AdaDelta run. That run's head moved + # L2(dw) = 1.9e-4 (0.0004% of ||w||) across all 30 epochs — it never left + # initialization. Cause: `experiment.learning_rate` IS applied to AdaDelta + # (driver.rs:404 passes lr to optimizer.step; adadelta.rs:123 multiplies + # delta by it) despite config.rs:211-213 documenting the key as inert, so + # every update ran 1000x scaled down. dHBV intends AdaDelta at lr=1.0. + optimizer: adam + loss: + kind: nse-batch + # Step schedule. Keys are the FIRST epoch at which the lr applies + # (`resolve_lr` takes the largest key <= epoch, 1-indexed). At 2 optimizer + # steps per epoch (see the accumulation note above): + # epochs 1-20 @ 0.001 (40 updates) + # epochs 21-30 @ 0.0005 (20 updates) + # Total displacement budget ~0.050. The decay lands 2/3 of the way through, + # which is the right shape — the concern is the step COUNT, not the schedule. + # + # Watch `n_at_floor` in the log: the lr=0.01 run pinned 47.3% of CONUS at the + # 0.015 floor, and both 2026-08-04 runs collapsed from 0% to 25-36% between + # epochs 3 and 5 at lr 2e-3. Starting at 1e-3 is deliberately cooler than that. + # Near-full-batch gradients (1,280 gauges) should also damp the collapse, since + # it looked like a high-variance small-batch effect — but that is a hypothesis, + # not an established result. If `n_at_floor` climbs past a few percent, the lr + # is still too hot. + learning_rate: + 1: 0.005 + 11: 0.001 + 21: 0.0005 + grad_clip_max_norm: 1.0 + +kan_head: + hidden_size: 21 + num_hidden_layers: 2 + # B-spline grid intervals (`num` in pykan). Matches DDR's + # ~/projects/ddr/config/merit_training_config.yaml::kan.grid. + grid: 50 + # B-spline order. Matches DDR's ::kan.k. (pykan's KANLayer default is 3, + # but DDR overrides to 2 for production.) + k: 2 + input_var_names: + - SoilGrids1km_clay + - aridity + - meanelevation + - meanP + - NDVI + - meanslope + - log10_uparea + - SoilGrids1km_sand + - ETPOT_Hargr + - Porosity + learnable_parameters: + - n + - q_spatial + - p_spatial + # Precip-conditioned mass-preserving daily->hourly disaggregation head. + # The within-day shape is driven by the hourly AORC precip window + # [d-1,d,d+1] (+ daily-Q taps + static attrs); the daily mean is conserved + # exactly. Loss stays the historical L1 (no experiment.loss block). + # Warm-started FROZEN from the standalone capacity pretrain (chunk_days=1); + # architecture fields must match that checkpoint or load_record fails. + # `enabled: false` makes this whole block inert — the loader uses flat + # repeat-24 (nearest) upsampling instead. Flip to true for the head-on + # arm of the ablation; nothing else needs to change. + disaggregation: + enabled: false + hidden_size: 16 + num_hidden_layers: 2 + grid: 20 + k: 3 + boundary_blend: 0.0 + chunk_days: 1 + pretrained_checkpoint: /home/tbindas/projects/ddrs/output/disagg_pretrain/capacity_chunk1.mpk + freeze: true + +# Routing-engine knobs — DDR's MERIT defaults (matches mock_config in tests). +params: + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] + attribute_minimums: + discharge: 1.0e-4 + slope: 1.0e-3 + velocity: 0.01 + depth: 0.01 + bottom_width: 0.01 + defaults: + p_spatial: 21.0 + log_space_parameters: + - p_spatial + # Corrected physics instead of DDR's formulation. With `false` the routing + # core uses: + # * exact trapezoidal celerity c = v·β, β = 5/3 − (4/3)·A·√(1+z²)/(T·P) + # instead of the wide-rectangular 5/3 limit (which is 22-27% high for the + # channels this code builds — κ = b/y ≈ 0.7-1.8, so β ≈ 1.31-1.36), and + # * Cunge-derived Muskingum X = clamp(0.5(1 − Q/(T·S·c·L)), 0, 0.5) instead + # of the constant 0.3, which was injecting a median 28× excess numerical + # diffusion. + # Negative-discharge counting before the S28 clamp is reported in BOTH modes. + # This breaks `compare_ddr_sandbox`'s ABSOLUTE MATCH by design — that example + # must be run with the default `ddr_match: true`. + # See `.claude/PHYSICS-CORRECTIONS.md`. + ddr_match: false + # Provably zero negative discharges. Clamps the Muskingum INPUTS (K, X) into + # the non-negative-coefficient window `2X <= Cr <= 2(1-X)`, Cr = dt/K: + # K <- max(K, dt*(1+d)/2) => Cr <= 2/(1+d) + # X <- min(X_cunge, (1-d)*Cr/2, (1-d)*(1-Cr/2)) + # with d = POSITIVITY_DELTA = 1e-2. Clamping K and X (not the coefficients) + # keeps `c1+c2+c3 = 1` exact, so mass is conserved; clamping c3 would not. + # With c1,c3 >= 0 and q_t,q' > 0, induction over the topological order of the + # forward substitution proves the whole solve is non-negative. + # + # Measured: 0 / 197,461,880 solves on 1,841 gauges (was 55,181, 0.0279%), + # replicated across two windows, two backends, and two trained heads. + # + # COST — this is not free. The X cap binds on 95.3% of reach-timesteps and + # median X falls 0.4976 -> 0.0794 (6.3x), so it largely REPLACES the Cunge X + # rather than shading it: numerical diffusion becomes stability-set instead of + # matched to hydraulic diffusivity. NSE/KGE impact is UNMEASURED, and this is + # the first training run to use the flag. Compare against a matched + # `enforce_positivity: false` control before trusting the result. + # + # Root-cause alternative: median Cr is 0.226, i.e. the typical MERIT reach is + # ~4.4x too LONG for an hourly step. Subdividing reaches ~4x brings Cr to ~1, + # where X_max -> 0.5 and the cap stops biting. See .claude/PHYSICS-CORRECTIONS.md. + # + # Requires ddr_match: false (rejected at load otherwise). + # OFF. It delivered provably zero negative solves (0 / 197,461,880), but the + # cure was worse than the disease: capping X collapsed the median from 0.4976 + # to 0.0794 on 95.3% of reach-timesteps, and because the cap is applied at + # RUNTIME to a learned celerity it made X ~ Cr ~ 1/n — handing the optimizer a + # lever it rode straight to the roughness floor. Same lr, flag on vs off: + # n_at_floor 35% -> 97.9%. Superseded by reach subdivision (a BUILD-TIME fix + # with no gradient path): see docs/superpowers/plans/2026-08-05-reach-subdivision.md + enforce_positivity: false + # cuSPARSE triangular solve. Pair with `ddrs run --backend cuda`; it is inert + # under `--backend cpu`. This is a DIFFERENT `Backward` impl from the host + # solve, sitting directly downstream of the ddr_match/enforce_positivity + # terms — gradient-checked on GPU by `tests/cuda_backward_parity.rs`. + # CPU set (2026-08-08): host triangular solve. + sparse_solver: cpu + # MUST be false when ddr_match is false: the fused CUDA-graph kernel + # (cuda_graph/geometry_kernel.rs) hardcodes DDR's 5/3 celerity, so a captured + # graph would replay the DDR forward while the backward used the corrected + # chain rule — a silent gradient mismatch. Config load rejects the pair. + use_cuda_graphs: false + +# Test-mode overlay. When the binary is run with --mode testing, these keys +# replace the corresponding ones in `experiment:`. Absent keys inherit. +# +# IMPORTANT: batch_size SEMANTIC SHIFTS between modes: +# - experiment.batch_size (training) = number of GAUGES per mini-batch +# - testing.batch_size = number of DAYS per chunk +testing: + start_time: 1995/10/01 + end_time: 2010/09/30 + batch_size: 15 # DAYS, not gauges + rho: null # disabled in test mode diff --git a/config/experiments/tau9_train_daily_lstm.yaml b/config/experiments/tau9_train_daily_lstm.yaml new file mode 100644 index 0000000..efaee4f --- /dev/null +++ b/config/experiments/tau9_train_daily_lstm.yaml @@ -0,0 +1,243 @@ +# tau=9 cross-source retrain arm: daily_lstm +# Derived from the 2026-08-05T04-58-58Z epoch-30 run's config snapshot with +# exactly three changes: (1) streamflow store swapped per arm, (2) +# sparse_solver cpu (CPU train-and-test set, 2026-08-08), (3) params.tau +# left UNSET so the 2026-08-08 default applies: tau=9 on the NEW convention +# (hours of advance; == old-convention 20, the measured optimum). The +# epoch-30 baseline trained at old tau=3 (== new -8, wrong direction). +# ddrs MERIT training config. +# +# Hyperparameters are pulled verbatim from +# ~/projects/ddr/config/merit_training_config.yaml. +# Forcing is daily; the loader interpolates to hourly (repeat × 24, trim to +# n_hourly) before feeding the Muskingum-Cunge engine — same pattern as +# DDR's StreamflowReader. + +mode: training +workflow: train-and-test # ddrs plan/run picks this up; override with --workflow X +geodataset: merit +seed: 42 +np_seed: 42 + +# Source paths — read in place by ddrs's Rust loaders (zarrs + icechunk + netcdf). +# No export/materialization step; the harness reads DDR's live data sources +# straight from disk. +data_sources: + attributes: /home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc + conus_adjacency: /home/tbindas/projects/ddr/data/merit_conus_adjacency.zarr + gages_adjacency: /home/tbindas/projects/ddr/data/merit_gages_conus_adjacency.zarr + streamflow: /mnt/ssd1/data/icechunk/daily_lstm_merit_unit_catchments.ic + observations: /mnt/ssd1/data/icechunk/usgs_daily_observations + gages: /home/tbindas/projects/ddr/references/gage_info/gages_2000_area_balanced.csv + # Hourly AORC precip (zarr v3, CONUS) — drives the precip-conditioned + # mass-preserving disaggregation head below. + aorc_precip: /mnt/ssd1/data/aorc/merit_unit_catchments.zarr + +experiment: + batch_size: 64 + start_time: 1981/10/01 + end_time: 1995/09/30 + epochs: 30 + rho: 90 + shuffle: true + warmup: 5 + # Gradient accumulation ON. `batch_size` is now the MICRO-batch; the optimizer + # steps once per `grad_accum_steps` micro-batches, so the effective batch is + # 64 * 4 = 256 gauges and the gradient is averaged over ~4x more basins. + # + # THE ARITHMETIC — read this before trusting a run: + # 1,841 gauges / 64 = 28 full + 1 partial = 29 micro-batches per epoch + # (accumulation sets `drop_last = false`, driver.rs:268-275, so the + # 49-gauge tail is KEPT — unlike the non-accumulating path) + # 29 / 20 = one group of 20 + one partial group of 9. + # A partial group STILL STEPS: only `micros_drawn == 0` breaks the loop + # (driver.rs:415), and the 1/Σn renormalization makes unequal group sizes + # exact. So this is 2 optimizer steps per epoch, not 1. + # x 30 epochs = 60 TOTAL UPDATES + # + # Effective batch is 20*64 = 1,280 of 1,841 gauges — 70% of the training set + # in a single gradient, i.e. near full-batch. Very low variance, very few steps. + # + # THAT IS THE RISK. Calibration: the run that beat the summed-Q' baseline + # (2026-07-30T00-24, median NSE 0.6799) took 180 updates at lr 1e-3 and moved + # the head L2(dw) = 9.75. The 2026-08-04 run that hit NSE 0.6782 took 280. + # Adam's per-weight displacement is bounded by ~lr per step, so the budget + # here is 40*0.001 + 20*0.0005 = 0.050 — about 38% of the 0.132 that the + # 180-step run realized. Expect under-training. + # + # If it under-trains, prefer more STEPS over a hotter lr: + # * `grad_accum_steps: 20` -> 7 (5 steps/epoch, 150 updates, eff batch 448) + # * `grad_accum_steps: 20` -> 4 (8 steps/epoch, 240 updates, eff batch 256) + # Wall clock is forward-pass bound (~27.5 s/micro-batch), so all of these cost + # the SAME time — 29 forwards per epoch regardless. Lowering accum_steps buys + # optimizer steps for free; it only trades gradient variance for step count. + # + # The historical `grad_accum_steps: 37` was WORSE than either: with 28-37 + # micro-batches per epoch it collapsed a whole epoch into ONE update, so 50 + # epochs bought 50 updates for 1,850 forward passes. Wall clock here is + # forward-pass bound (~27.5 s/micro-batch), not step bound, so accumulating + # costs nothing in time — it only trades update count for gradient quality. + # + # ALWAYS verify the head actually moved (L2(dw)) before believing a result: + # the 2026-07-31 AdaDelta run reported plausible losses while L2(dw) = 1.9e-4, + # i.e. it never left initialization. + use_grad_accum: true + grad_accum_steps: 20 + # Adam, reverting the 2026-07-31 AdaDelta run. That run's head moved + # L2(dw) = 1.9e-4 (0.0004% of ||w||) across all 30 epochs — it never left + # initialization. Cause: `experiment.learning_rate` IS applied to AdaDelta + # (driver.rs:404 passes lr to optimizer.step; adadelta.rs:123 multiplies + # delta by it) despite config.rs:211-213 documenting the key as inert, so + # every update ran 1000x scaled down. dHBV intends AdaDelta at lr=1.0. + optimizer: adam + loss: + kind: nse-batch + # Step schedule. Keys are the FIRST epoch at which the lr applies + # (`resolve_lr` takes the largest key <= epoch, 1-indexed). At 2 optimizer + # steps per epoch (see the accumulation note above): + # epochs 1-20 @ 0.001 (40 updates) + # epochs 21-30 @ 0.0005 (20 updates) + # Total displacement budget ~0.050. The decay lands 2/3 of the way through, + # which is the right shape — the concern is the step COUNT, not the schedule. + # + # Watch `n_at_floor` in the log: the lr=0.01 run pinned 47.3% of CONUS at the + # 0.015 floor, and both 2026-08-04 runs collapsed from 0% to 25-36% between + # epochs 3 and 5 at lr 2e-3. Starting at 1e-3 is deliberately cooler than that. + # Near-full-batch gradients (1,280 gauges) should also damp the collapse, since + # it looked like a high-variance small-batch effect — but that is a hypothesis, + # not an established result. If `n_at_floor` climbs past a few percent, the lr + # is still too hot. + learning_rate: + 1: 0.005 + 11: 0.001 + 21: 0.0005 + grad_clip_max_norm: 1.0 + +kan_head: + hidden_size: 21 + num_hidden_layers: 2 + # B-spline grid intervals (`num` in pykan). Matches DDR's + # ~/projects/ddr/config/merit_training_config.yaml::kan.grid. + grid: 50 + # B-spline order. Matches DDR's ::kan.k. (pykan's KANLayer default is 3, + # but DDR overrides to 2 for production.) + k: 2 + input_var_names: + - SoilGrids1km_clay + - aridity + - meanelevation + - meanP + - NDVI + - meanslope + - log10_uparea + - SoilGrids1km_sand + - ETPOT_Hargr + - Porosity + learnable_parameters: + - n + - q_spatial + - p_spatial + # Precip-conditioned mass-preserving daily->hourly disaggregation head. + # The within-day shape is driven by the hourly AORC precip window + # [d-1,d,d+1] (+ daily-Q taps + static attrs); the daily mean is conserved + # exactly. Loss stays the historical L1 (no experiment.loss block). + # Warm-started FROZEN from the standalone capacity pretrain (chunk_days=1); + # architecture fields must match that checkpoint or load_record fails. + # `enabled: false` makes this whole block inert — the loader uses flat + # repeat-24 (nearest) upsampling instead. Flip to true for the head-on + # arm of the ablation; nothing else needs to change. + disaggregation: + enabled: false + hidden_size: 16 + num_hidden_layers: 2 + grid: 20 + k: 3 + boundary_blend: 0.0 + chunk_days: 1 + pretrained_checkpoint: /home/tbindas/projects/ddrs/output/disagg_pretrain/capacity_chunk1.mpk + freeze: true + +# Routing-engine knobs — DDR's MERIT defaults (matches mock_config in tests). +params: + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] + attribute_minimums: + discharge: 1.0e-4 + slope: 1.0e-3 + velocity: 0.01 + depth: 0.01 + bottom_width: 0.01 + defaults: + p_spatial: 21.0 + log_space_parameters: + - p_spatial + # Corrected physics instead of DDR's formulation. With `false` the routing + # core uses: + # * exact trapezoidal celerity c = v·β, β = 5/3 − (4/3)·A·√(1+z²)/(T·P) + # instead of the wide-rectangular 5/3 limit (which is 22-27% high for the + # channels this code builds — κ = b/y ≈ 0.7-1.8, so β ≈ 1.31-1.36), and + # * Cunge-derived Muskingum X = clamp(0.5(1 − Q/(T·S·c·L)), 0, 0.5) instead + # of the constant 0.3, which was injecting a median 28× excess numerical + # diffusion. + # Negative-discharge counting before the S28 clamp is reported in BOTH modes. + # This breaks `compare_ddr_sandbox`'s ABSOLUTE MATCH by design — that example + # must be run with the default `ddr_match: true`. + # See `.claude/PHYSICS-CORRECTIONS.md`. + ddr_match: false + # Provably zero negative discharges. Clamps the Muskingum INPUTS (K, X) into + # the non-negative-coefficient window `2X <= Cr <= 2(1-X)`, Cr = dt/K: + # K <- max(K, dt*(1+d)/2) => Cr <= 2/(1+d) + # X <- min(X_cunge, (1-d)*Cr/2, (1-d)*(1-Cr/2)) + # with d = POSITIVITY_DELTA = 1e-2. Clamping K and X (not the coefficients) + # keeps `c1+c2+c3 = 1` exact, so mass is conserved; clamping c3 would not. + # With c1,c3 >= 0 and q_t,q' > 0, induction over the topological order of the + # forward substitution proves the whole solve is non-negative. + # + # Measured: 0 / 197,461,880 solves on 1,841 gauges (was 55,181, 0.0279%), + # replicated across two windows, two backends, and two trained heads. + # + # COST — this is not free. The X cap binds on 95.3% of reach-timesteps and + # median X falls 0.4976 -> 0.0794 (6.3x), so it largely REPLACES the Cunge X + # rather than shading it: numerical diffusion becomes stability-set instead of + # matched to hydraulic diffusivity. NSE/KGE impact is UNMEASURED, and this is + # the first training run to use the flag. Compare against a matched + # `enforce_positivity: false` control before trusting the result. + # + # Root-cause alternative: median Cr is 0.226, i.e. the typical MERIT reach is + # ~4.4x too LONG for an hourly step. Subdividing reaches ~4x brings Cr to ~1, + # where X_max -> 0.5 and the cap stops biting. See .claude/PHYSICS-CORRECTIONS.md. + # + # Requires ddr_match: false (rejected at load otherwise). + # OFF. It delivered provably zero negative solves (0 / 197,461,880), but the + # cure was worse than the disease: capping X collapsed the median from 0.4976 + # to 0.0794 on 95.3% of reach-timesteps, and because the cap is applied at + # RUNTIME to a learned celerity it made X ~ Cr ~ 1/n — handing the optimizer a + # lever it rode straight to the roughness floor. Same lr, flag on vs off: + # n_at_floor 35% -> 97.9%. Superseded by reach subdivision (a BUILD-TIME fix + # with no gradient path): see docs/superpowers/plans/2026-08-05-reach-subdivision.md + enforce_positivity: false + # cuSPARSE triangular solve. Pair with `ddrs run --backend cuda`; it is inert + # under `--backend cpu`. This is a DIFFERENT `Backward` impl from the host + # solve, sitting directly downstream of the ddr_match/enforce_positivity + # terms — gradient-checked on GPU by `tests/cuda_backward_parity.rs`. + # CPU set (2026-08-08): host triangular solve. + sparse_solver: cpu + # MUST be false when ddr_match is false: the fused CUDA-graph kernel + # (cuda_graph/geometry_kernel.rs) hardcodes DDR's 5/3 celerity, so a captured + # graph would replay the DDR forward while the backward used the corrected + # chain rule — a silent gradient mismatch. Config load rejects the pair. + use_cuda_graphs: false + +# Test-mode overlay. When the binary is run with --mode testing, these keys +# replace the corresponding ones in `experiment:`. Absent keys inherit. +# +# IMPORTANT: batch_size SEMANTIC SHIFTS between modes: +# - experiment.batch_size (training) = number of GAUGES per mini-batch +# - testing.batch_size = number of DAYS per chunk +testing: + start_time: 1995/10/01 + end_time: 2010/09/30 + batch_size: 15 # DAYS, not gauges + rho: null # disabled in test mode diff --git a/config/experiments/tau9_train_hourly_lstm.yaml b/config/experiments/tau9_train_hourly_lstm.yaml new file mode 100644 index 0000000..39a156f --- /dev/null +++ b/config/experiments/tau9_train_hourly_lstm.yaml @@ -0,0 +1,243 @@ +# tau=9 cross-source retrain arm: hourly_lstm +# Derived from the 2026-08-05T04-58-58Z epoch-30 run's config snapshot with +# exactly three changes: (1) streamflow store swapped per arm, (2) +# sparse_solver cpu (CPU train-and-test set, 2026-08-08), (3) params.tau +# left UNSET so the 2026-08-08 default applies: tau=9 on the NEW convention +# (hours of advance; == old-convention 20, the measured optimum). The +# epoch-30 baseline trained at old tau=3 (== new -8, wrong direction). +# ddrs MERIT training config. +# +# Hyperparameters are pulled verbatim from +# ~/projects/ddr/config/merit_training_config.yaml. +# Forcing is daily; the loader interpolates to hourly (repeat × 24, trim to +# n_hourly) before feeding the Muskingum-Cunge engine — same pattern as +# DDR's StreamflowReader. + +mode: training +workflow: train-and-test # ddrs plan/run picks this up; override with --workflow X +geodataset: merit +seed: 42 +np_seed: 42 + +# Source paths — read in place by ddrs's Rust loaders (zarrs + icechunk + netcdf). +# No export/materialization step; the harness reads DDR's live data sources +# straight from disk. +data_sources: + attributes: /home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc + conus_adjacency: /home/tbindas/projects/ddr/data/merit_conus_adjacency.zarr + gages_adjacency: /home/tbindas/projects/ddr/data/merit_gages_conus_adjacency.zarr + streamflow: /mnt/ssd1/data/icechunk/hourly_lstm_merit_unit_catchments.ic + observations: /mnt/ssd1/data/icechunk/usgs_daily_observations + gages: /home/tbindas/projects/ddr/references/gage_info/gages_2000_area_balanced.csv + # Hourly AORC precip (zarr v3, CONUS) — drives the precip-conditioned + # mass-preserving disaggregation head below. + aorc_precip: /mnt/ssd1/data/aorc/merit_unit_catchments.zarr + +experiment: + batch_size: 64 + start_time: 1981/10/01 + end_time: 1995/09/30 + epochs: 30 + rho: 90 + shuffle: true + warmup: 5 + # Gradient accumulation ON. `batch_size` is now the MICRO-batch; the optimizer + # steps once per `grad_accum_steps` micro-batches, so the effective batch is + # 64 * 4 = 256 gauges and the gradient is averaged over ~4x more basins. + # + # THE ARITHMETIC — read this before trusting a run: + # 1,841 gauges / 64 = 28 full + 1 partial = 29 micro-batches per epoch + # (accumulation sets `drop_last = false`, driver.rs:268-275, so the + # 49-gauge tail is KEPT — unlike the non-accumulating path) + # 29 / 20 = one group of 20 + one partial group of 9. + # A partial group STILL STEPS: only `micros_drawn == 0` breaks the loop + # (driver.rs:415), and the 1/Σn renormalization makes unequal group sizes + # exact. So this is 2 optimizer steps per epoch, not 1. + # x 30 epochs = 60 TOTAL UPDATES + # + # Effective batch is 20*64 = 1,280 of 1,841 gauges — 70% of the training set + # in a single gradient, i.e. near full-batch. Very low variance, very few steps. + # + # THAT IS THE RISK. Calibration: the run that beat the summed-Q' baseline + # (2026-07-30T00-24, median NSE 0.6799) took 180 updates at lr 1e-3 and moved + # the head L2(dw) = 9.75. The 2026-08-04 run that hit NSE 0.6782 took 280. + # Adam's per-weight displacement is bounded by ~lr per step, so the budget + # here is 40*0.001 + 20*0.0005 = 0.050 — about 38% of the 0.132 that the + # 180-step run realized. Expect under-training. + # + # If it under-trains, prefer more STEPS over a hotter lr: + # * `grad_accum_steps: 20` -> 7 (5 steps/epoch, 150 updates, eff batch 448) + # * `grad_accum_steps: 20` -> 4 (8 steps/epoch, 240 updates, eff batch 256) + # Wall clock is forward-pass bound (~27.5 s/micro-batch), so all of these cost + # the SAME time — 29 forwards per epoch regardless. Lowering accum_steps buys + # optimizer steps for free; it only trades gradient variance for step count. + # + # The historical `grad_accum_steps: 37` was WORSE than either: with 28-37 + # micro-batches per epoch it collapsed a whole epoch into ONE update, so 50 + # epochs bought 50 updates for 1,850 forward passes. Wall clock here is + # forward-pass bound (~27.5 s/micro-batch), not step bound, so accumulating + # costs nothing in time — it only trades update count for gradient quality. + # + # ALWAYS verify the head actually moved (L2(dw)) before believing a result: + # the 2026-07-31 AdaDelta run reported plausible losses while L2(dw) = 1.9e-4, + # i.e. it never left initialization. + use_grad_accum: true + grad_accum_steps: 20 + # Adam, reverting the 2026-07-31 AdaDelta run. That run's head moved + # L2(dw) = 1.9e-4 (0.0004% of ||w||) across all 30 epochs — it never left + # initialization. Cause: `experiment.learning_rate` IS applied to AdaDelta + # (driver.rs:404 passes lr to optimizer.step; adadelta.rs:123 multiplies + # delta by it) despite config.rs:211-213 documenting the key as inert, so + # every update ran 1000x scaled down. dHBV intends AdaDelta at lr=1.0. + optimizer: adam + loss: + kind: nse-batch + # Step schedule. Keys are the FIRST epoch at which the lr applies + # (`resolve_lr` takes the largest key <= epoch, 1-indexed). At 2 optimizer + # steps per epoch (see the accumulation note above): + # epochs 1-20 @ 0.001 (40 updates) + # epochs 21-30 @ 0.0005 (20 updates) + # Total displacement budget ~0.050. The decay lands 2/3 of the way through, + # which is the right shape — the concern is the step COUNT, not the schedule. + # + # Watch `n_at_floor` in the log: the lr=0.01 run pinned 47.3% of CONUS at the + # 0.015 floor, and both 2026-08-04 runs collapsed from 0% to 25-36% between + # epochs 3 and 5 at lr 2e-3. Starting at 1e-3 is deliberately cooler than that. + # Near-full-batch gradients (1,280 gauges) should also damp the collapse, since + # it looked like a high-variance small-batch effect — but that is a hypothesis, + # not an established result. If `n_at_floor` climbs past a few percent, the lr + # is still too hot. + learning_rate: + 1: 0.005 + 11: 0.001 + 21: 0.0005 + grad_clip_max_norm: 1.0 + +kan_head: + hidden_size: 21 + num_hidden_layers: 2 + # B-spline grid intervals (`num` in pykan). Matches DDR's + # ~/projects/ddr/config/merit_training_config.yaml::kan.grid. + grid: 50 + # B-spline order. Matches DDR's ::kan.k. (pykan's KANLayer default is 3, + # but DDR overrides to 2 for production.) + k: 2 + input_var_names: + - SoilGrids1km_clay + - aridity + - meanelevation + - meanP + - NDVI + - meanslope + - log10_uparea + - SoilGrids1km_sand + - ETPOT_Hargr + - Porosity + learnable_parameters: + - n + - q_spatial + - p_spatial + # Precip-conditioned mass-preserving daily->hourly disaggregation head. + # The within-day shape is driven by the hourly AORC precip window + # [d-1,d,d+1] (+ daily-Q taps + static attrs); the daily mean is conserved + # exactly. Loss stays the historical L1 (no experiment.loss block). + # Warm-started FROZEN from the standalone capacity pretrain (chunk_days=1); + # architecture fields must match that checkpoint or load_record fails. + # `enabled: false` makes this whole block inert — the loader uses flat + # repeat-24 (nearest) upsampling instead. Flip to true for the head-on + # arm of the ablation; nothing else needs to change. + disaggregation: + enabled: false + hidden_size: 16 + num_hidden_layers: 2 + grid: 20 + k: 3 + boundary_blend: 0.0 + chunk_days: 1 + pretrained_checkpoint: /home/tbindas/projects/ddrs/output/disagg_pretrain/capacity_chunk1.mpk + freeze: true + +# Routing-engine knobs — DDR's MERIT defaults (matches mock_config in tests). +params: + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] + attribute_minimums: + discharge: 1.0e-4 + slope: 1.0e-3 + velocity: 0.01 + depth: 0.01 + bottom_width: 0.01 + defaults: + p_spatial: 21.0 + log_space_parameters: + - p_spatial + # Corrected physics instead of DDR's formulation. With `false` the routing + # core uses: + # * exact trapezoidal celerity c = v·β, β = 5/3 − (4/3)·A·√(1+z²)/(T·P) + # instead of the wide-rectangular 5/3 limit (which is 22-27% high for the + # channels this code builds — κ = b/y ≈ 0.7-1.8, so β ≈ 1.31-1.36), and + # * Cunge-derived Muskingum X = clamp(0.5(1 − Q/(T·S·c·L)), 0, 0.5) instead + # of the constant 0.3, which was injecting a median 28× excess numerical + # diffusion. + # Negative-discharge counting before the S28 clamp is reported in BOTH modes. + # This breaks `compare_ddr_sandbox`'s ABSOLUTE MATCH by design — that example + # must be run with the default `ddr_match: true`. + # See `.claude/PHYSICS-CORRECTIONS.md`. + ddr_match: false + # Provably zero negative discharges. Clamps the Muskingum INPUTS (K, X) into + # the non-negative-coefficient window `2X <= Cr <= 2(1-X)`, Cr = dt/K: + # K <- max(K, dt*(1+d)/2) => Cr <= 2/(1+d) + # X <- min(X_cunge, (1-d)*Cr/2, (1-d)*(1-Cr/2)) + # with d = POSITIVITY_DELTA = 1e-2. Clamping K and X (not the coefficients) + # keeps `c1+c2+c3 = 1` exact, so mass is conserved; clamping c3 would not. + # With c1,c3 >= 0 and q_t,q' > 0, induction over the topological order of the + # forward substitution proves the whole solve is non-negative. + # + # Measured: 0 / 197,461,880 solves on 1,841 gauges (was 55,181, 0.0279%), + # replicated across two windows, two backends, and two trained heads. + # + # COST — this is not free. The X cap binds on 95.3% of reach-timesteps and + # median X falls 0.4976 -> 0.0794 (6.3x), so it largely REPLACES the Cunge X + # rather than shading it: numerical diffusion becomes stability-set instead of + # matched to hydraulic diffusivity. NSE/KGE impact is UNMEASURED, and this is + # the first training run to use the flag. Compare against a matched + # `enforce_positivity: false` control before trusting the result. + # + # Root-cause alternative: median Cr is 0.226, i.e. the typical MERIT reach is + # ~4.4x too LONG for an hourly step. Subdividing reaches ~4x brings Cr to ~1, + # where X_max -> 0.5 and the cap stops biting. See .claude/PHYSICS-CORRECTIONS.md. + # + # Requires ddr_match: false (rejected at load otherwise). + # OFF. It delivered provably zero negative solves (0 / 197,461,880), but the + # cure was worse than the disease: capping X collapsed the median from 0.4976 + # to 0.0794 on 95.3% of reach-timesteps, and because the cap is applied at + # RUNTIME to a learned celerity it made X ~ Cr ~ 1/n — handing the optimizer a + # lever it rode straight to the roughness floor. Same lr, flag on vs off: + # n_at_floor 35% -> 97.9%. Superseded by reach subdivision (a BUILD-TIME fix + # with no gradient path): see docs/superpowers/plans/2026-08-05-reach-subdivision.md + enforce_positivity: false + # cuSPARSE triangular solve. Pair with `ddrs run --backend cuda`; it is inert + # under `--backend cpu`. This is a DIFFERENT `Backward` impl from the host + # solve, sitting directly downstream of the ddr_match/enforce_positivity + # terms — gradient-checked on GPU by `tests/cuda_backward_parity.rs`. + # CPU set (2026-08-08): host triangular solve. + sparse_solver: cpu + # MUST be false when ddr_match is false: the fused CUDA-graph kernel + # (cuda_graph/geometry_kernel.rs) hardcodes DDR's 5/3 celerity, so a captured + # graph would replay the DDR forward while the backward used the corrected + # chain rule — a silent gradient mismatch. Config load rejects the pair. + use_cuda_graphs: false + +# Test-mode overlay. When the binary is run with --mode testing, these keys +# replace the corresponding ones in `experiment:`. Absent keys inherit. +# +# IMPORTANT: batch_size SEMANTIC SHIFTS between modes: +# - experiment.batch_size (training) = number of GAUGES per mini-batch +# - testing.batch_size = number of DAYS per chunk +testing: + start_time: 1995/10/01 + end_time: 2010/09/30 + batch_size: 15 # DAYS, not gauges + rho: null # disabled in test mode diff --git a/config/experiments/tau9_train_hourly_lstm_evalresume.yaml b/config/experiments/tau9_train_hourly_lstm_evalresume.yaml new file mode 100644 index 0000000..3cf9c80 --- /dev/null +++ b/config/experiments/tau9_train_hourly_lstm_evalresume.yaml @@ -0,0 +1,248 @@ +# tau=9 cross-source retrain arm: hourly_lstm +# Derived from the 2026-08-05T04-58-58Z epoch-30 run's config snapshot with +# exactly three changes: (1) streamflow store swapped per arm, (2) +# sparse_solver cpu (CPU train-and-test set, 2026-08-08), (3) params.tau +# left UNSET so the 2026-08-08 default applies: tau=9 on the NEW convention +# (hours of advance; == old-convention 20, the measured optimum). The +# epoch-30 baseline trained at old tau=3 (== new -8, wrong direction). +# ddrs MERIT training config. +# +# Hyperparameters are pulled verbatim from +# ~/projects/ddr/config/merit_training_config.yaml. +# Forcing is daily; the loader interpolates to hourly (repeat × 24, trim to +# n_hourly) before feeding the Muskingum-Cunge engine — same pattern as +# DDR's StreamflowReader. + +mode: training +workflow: train-and-test # ddrs plan/run picks this up; override with --workflow X +geodataset: merit +seed: 42 +np_seed: 42 + +# Source paths — read in place by ddrs's Rust loaders (zarrs + icechunk + netcdf). +# No export/materialization step; the harness reads DDR's live data sources +# straight from disk. +data_sources: + attributes: /home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc + conus_adjacency: /home/tbindas/projects/ddr/data/merit_conus_adjacency.zarr + gages_adjacency: /home/tbindas/projects/ddr/data/merit_gages_conus_adjacency.zarr + streamflow: /mnt/ssd1/data/icechunk/hourly_lstm_merit_unit_catchments.ic + observations: /mnt/ssd1/data/icechunk/usgs_daily_observations + gages: /home/tbindas/projects/ddr/references/gage_info/gages_2000_area_balanced.csv + # Hourly AORC precip (zarr v3, CONUS) — drives the precip-conditioned + # mass-preserving disaggregation head below. + aorc_precip: /mnt/ssd1/data/aorc/merit_unit_catchments.zarr + +experiment: + # Eval-resume (2026-08-09): the original run's phase 2 died mid-eval when + # its parent task was killed; resuming from the FINAL checkpoint with + # epochs == 30 trains zero batches and proceeds straight to eval, writing + # a complete manifest without repeating the 2.4 h training phase. + checkpoint: .ddrs/runs/2026-08-09T14-55-05Z-train-and-test/checkpoints/epoch_30_mb_1 + batch_size: 64 + start_time: 1981/10/01 + end_time: 1995/09/30 + epochs: 30 + rho: 90 + shuffle: true + warmup: 5 + # Gradient accumulation ON. `batch_size` is now the MICRO-batch; the optimizer + # steps once per `grad_accum_steps` micro-batches, so the effective batch is + # 64 * 4 = 256 gauges and the gradient is averaged over ~4x more basins. + # + # THE ARITHMETIC — read this before trusting a run: + # 1,841 gauges / 64 = 28 full + 1 partial = 29 micro-batches per epoch + # (accumulation sets `drop_last = false`, driver.rs:268-275, so the + # 49-gauge tail is KEPT — unlike the non-accumulating path) + # 29 / 20 = one group of 20 + one partial group of 9. + # A partial group STILL STEPS: only `micros_drawn == 0` breaks the loop + # (driver.rs:415), and the 1/Σn renormalization makes unequal group sizes + # exact. So this is 2 optimizer steps per epoch, not 1. + # x 30 epochs = 60 TOTAL UPDATES + # + # Effective batch is 20*64 = 1,280 of 1,841 gauges — 70% of the training set + # in a single gradient, i.e. near full-batch. Very low variance, very few steps. + # + # THAT IS THE RISK. Calibration: the run that beat the summed-Q' baseline + # (2026-07-30T00-24, median NSE 0.6799) took 180 updates at lr 1e-3 and moved + # the head L2(dw) = 9.75. The 2026-08-04 run that hit NSE 0.6782 took 280. + # Adam's per-weight displacement is bounded by ~lr per step, so the budget + # here is 40*0.001 + 20*0.0005 = 0.050 — about 38% of the 0.132 that the + # 180-step run realized. Expect under-training. + # + # If it under-trains, prefer more STEPS over a hotter lr: + # * `grad_accum_steps: 20` -> 7 (5 steps/epoch, 150 updates, eff batch 448) + # * `grad_accum_steps: 20` -> 4 (8 steps/epoch, 240 updates, eff batch 256) + # Wall clock is forward-pass bound (~27.5 s/micro-batch), so all of these cost + # the SAME time — 29 forwards per epoch regardless. Lowering accum_steps buys + # optimizer steps for free; it only trades gradient variance for step count. + # + # The historical `grad_accum_steps: 37` was WORSE than either: with 28-37 + # micro-batches per epoch it collapsed a whole epoch into ONE update, so 50 + # epochs bought 50 updates for 1,850 forward passes. Wall clock here is + # forward-pass bound (~27.5 s/micro-batch), not step bound, so accumulating + # costs nothing in time — it only trades update count for gradient quality. + # + # ALWAYS verify the head actually moved (L2(dw)) before believing a result: + # the 2026-07-31 AdaDelta run reported plausible losses while L2(dw) = 1.9e-4, + # i.e. it never left initialization. + use_grad_accum: true + grad_accum_steps: 20 + # Adam, reverting the 2026-07-31 AdaDelta run. That run's head moved + # L2(dw) = 1.9e-4 (0.0004% of ||w||) across all 30 epochs — it never left + # initialization. Cause: `experiment.learning_rate` IS applied to AdaDelta + # (driver.rs:404 passes lr to optimizer.step; adadelta.rs:123 multiplies + # delta by it) despite config.rs:211-213 documenting the key as inert, so + # every update ran 1000x scaled down. dHBV intends AdaDelta at lr=1.0. + optimizer: adam + loss: + kind: nse-batch + # Step schedule. Keys are the FIRST epoch at which the lr applies + # (`resolve_lr` takes the largest key <= epoch, 1-indexed). At 2 optimizer + # steps per epoch (see the accumulation note above): + # epochs 1-20 @ 0.001 (40 updates) + # epochs 21-30 @ 0.0005 (20 updates) + # Total displacement budget ~0.050. The decay lands 2/3 of the way through, + # which is the right shape — the concern is the step COUNT, not the schedule. + # + # Watch `n_at_floor` in the log: the lr=0.01 run pinned 47.3% of CONUS at the + # 0.015 floor, and both 2026-08-04 runs collapsed from 0% to 25-36% between + # epochs 3 and 5 at lr 2e-3. Starting at 1e-3 is deliberately cooler than that. + # Near-full-batch gradients (1,280 gauges) should also damp the collapse, since + # it looked like a high-variance small-batch effect — but that is a hypothesis, + # not an established result. If `n_at_floor` climbs past a few percent, the lr + # is still too hot. + learning_rate: + 1: 0.005 + 11: 0.001 + 21: 0.0005 + grad_clip_max_norm: 1.0 + +kan_head: + hidden_size: 21 + num_hidden_layers: 2 + # B-spline grid intervals (`num` in pykan). Matches DDR's + # ~/projects/ddr/config/merit_training_config.yaml::kan.grid. + grid: 50 + # B-spline order. Matches DDR's ::kan.k. (pykan's KANLayer default is 3, + # but DDR overrides to 2 for production.) + k: 2 + input_var_names: + - SoilGrids1km_clay + - aridity + - meanelevation + - meanP + - NDVI + - meanslope + - log10_uparea + - SoilGrids1km_sand + - ETPOT_Hargr + - Porosity + learnable_parameters: + - n + - q_spatial + - p_spatial + # Precip-conditioned mass-preserving daily->hourly disaggregation head. + # The within-day shape is driven by the hourly AORC precip window + # [d-1,d,d+1] (+ daily-Q taps + static attrs); the daily mean is conserved + # exactly. Loss stays the historical L1 (no experiment.loss block). + # Warm-started FROZEN from the standalone capacity pretrain (chunk_days=1); + # architecture fields must match that checkpoint or load_record fails. + # `enabled: false` makes this whole block inert — the loader uses flat + # repeat-24 (nearest) upsampling instead. Flip to true for the head-on + # arm of the ablation; nothing else needs to change. + disaggregation: + enabled: false + hidden_size: 16 + num_hidden_layers: 2 + grid: 20 + k: 3 + boundary_blend: 0.0 + chunk_days: 1 + pretrained_checkpoint: /home/tbindas/projects/ddrs/output/disagg_pretrain/capacity_chunk1.mpk + freeze: true + +# Routing-engine knobs — DDR's MERIT defaults (matches mock_config in tests). +params: + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] + attribute_minimums: + discharge: 1.0e-4 + slope: 1.0e-3 + velocity: 0.01 + depth: 0.01 + bottom_width: 0.01 + defaults: + p_spatial: 21.0 + log_space_parameters: + - p_spatial + # Corrected physics instead of DDR's formulation. With `false` the routing + # core uses: + # * exact trapezoidal celerity c = v·β, β = 5/3 − (4/3)·A·√(1+z²)/(T·P) + # instead of the wide-rectangular 5/3 limit (which is 22-27% high for the + # channels this code builds — κ = b/y ≈ 0.7-1.8, so β ≈ 1.31-1.36), and + # * Cunge-derived Muskingum X = clamp(0.5(1 − Q/(T·S·c·L)), 0, 0.5) instead + # of the constant 0.3, which was injecting a median 28× excess numerical + # diffusion. + # Negative-discharge counting before the S28 clamp is reported in BOTH modes. + # This breaks `compare_ddr_sandbox`'s ABSOLUTE MATCH by design — that example + # must be run with the default `ddr_match: true`. + # See `.claude/PHYSICS-CORRECTIONS.md`. + ddr_match: false + # Provably zero negative discharges. Clamps the Muskingum INPUTS (K, X) into + # the non-negative-coefficient window `2X <= Cr <= 2(1-X)`, Cr = dt/K: + # K <- max(K, dt*(1+d)/2) => Cr <= 2/(1+d) + # X <- min(X_cunge, (1-d)*Cr/2, (1-d)*(1-Cr/2)) + # with d = POSITIVITY_DELTA = 1e-2. Clamping K and X (not the coefficients) + # keeps `c1+c2+c3 = 1` exact, so mass is conserved; clamping c3 would not. + # With c1,c3 >= 0 and q_t,q' > 0, induction over the topological order of the + # forward substitution proves the whole solve is non-negative. + # + # Measured: 0 / 197,461,880 solves on 1,841 gauges (was 55,181, 0.0279%), + # replicated across two windows, two backends, and two trained heads. + # + # COST — this is not free. The X cap binds on 95.3% of reach-timesteps and + # median X falls 0.4976 -> 0.0794 (6.3x), so it largely REPLACES the Cunge X + # rather than shading it: numerical diffusion becomes stability-set instead of + # matched to hydraulic diffusivity. NSE/KGE impact is UNMEASURED, and this is + # the first training run to use the flag. Compare against a matched + # `enforce_positivity: false` control before trusting the result. + # + # Root-cause alternative: median Cr is 0.226, i.e. the typical MERIT reach is + # ~4.4x too LONG for an hourly step. Subdividing reaches ~4x brings Cr to ~1, + # where X_max -> 0.5 and the cap stops biting. See .claude/PHYSICS-CORRECTIONS.md. + # + # Requires ddr_match: false (rejected at load otherwise). + # OFF. It delivered provably zero negative solves (0 / 197,461,880), but the + # cure was worse than the disease: capping X collapsed the median from 0.4976 + # to 0.0794 on 95.3% of reach-timesteps, and because the cap is applied at + # RUNTIME to a learned celerity it made X ~ Cr ~ 1/n — handing the optimizer a + # lever it rode straight to the roughness floor. Same lr, flag on vs off: + # n_at_floor 35% -> 97.9%. Superseded by reach subdivision (a BUILD-TIME fix + # with no gradient path): see docs/superpowers/plans/2026-08-05-reach-subdivision.md + enforce_positivity: false + # cuSPARSE triangular solve. Pair with `ddrs run --backend cuda`; it is inert + # under `--backend cpu`. This is a DIFFERENT `Backward` impl from the host + # solve, sitting directly downstream of the ddr_match/enforce_positivity + # terms — gradient-checked on GPU by `tests/cuda_backward_parity.rs`. + # CPU set (2026-08-08): host triangular solve. + sparse_solver: cpu + # MUST be false when ddr_match is false: the fused CUDA-graph kernel + # (cuda_graph/geometry_kernel.rs) hardcodes DDR's 5/3 celerity, so a captured + # graph would replay the DDR forward while the backward used the corrected + # chain rule — a silent gradient mismatch. Config load rejects the pair. + use_cuda_graphs: false + +# Test-mode overlay. When the binary is run with --mode testing, these keys +# replace the corresponding ones in `experiment:`. Absent keys inherit. +# +# IMPORTANT: batch_size SEMANTIC SHIFTS between modes: +# - experiment.batch_size (training) = number of GAUGES per mini-batch +# - testing.batch_size = number of DAYS per chunk +testing: + start_time: 1995/10/01 + end_time: 2010/09/30 + batch_size: 15 # DAYS, not gauges + rho: null # disabled in test mode diff --git a/config/experiments/tau9_train_uh_retro.yaml b/config/experiments/tau9_train_uh_retro.yaml new file mode 100644 index 0000000..272b83f --- /dev/null +++ b/config/experiments/tau9_train_uh_retro.yaml @@ -0,0 +1,243 @@ +# tau=9 cross-source retrain arm: uh_retro +# Derived from the 2026-08-05T04-58-58Z epoch-30 run's config snapshot with +# exactly three changes: (1) streamflow store swapped per arm, (2) +# sparse_solver cpu (CPU train-and-test set, 2026-08-08), (3) params.tau +# left UNSET so the 2026-08-08 default applies: tau=9 on the NEW convention +# (hours of advance; == old-convention 20, the measured optimum). The +# epoch-30 baseline trained at old tau=3 (== new -8, wrong direction). +# ddrs MERIT training config. +# +# Hyperparameters are pulled verbatim from +# ~/projects/ddr/config/merit_training_config.yaml. +# Forcing is daily; the loader interpolates to hourly (repeat × 24, trim to +# n_hourly) before feeding the Muskingum-Cunge engine — same pattern as +# DDR's StreamflowReader. + +mode: training +workflow: train-and-test # ddrs plan/run picks this up; override with --workflow X +geodataset: merit +seed: 42 +np_seed: 42 + +# Source paths — read in place by ddrs's Rust loaders (zarrs + icechunk + netcdf). +# No export/materialization step; the harness reads DDR's live data sources +# straight from disk. +data_sources: + attributes: /home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc + conus_adjacency: /home/tbindas/projects/ddr/data/merit_conus_adjacency.zarr + gages_adjacency: /home/tbindas/projects/ddr/data/merit_gages_conus_adjacency.zarr + streamflow: /mnt/ssd1/data/icechunk/merit_dhbv2_UH_retrospective.ic + observations: /mnt/ssd1/data/icechunk/usgs_daily_observations + gages: /home/tbindas/projects/ddr/references/gage_info/gages_2000_area_balanced.csv + # Hourly AORC precip (zarr v3, CONUS) — drives the precip-conditioned + # mass-preserving disaggregation head below. + aorc_precip: /mnt/ssd1/data/aorc/merit_unit_catchments.zarr + +experiment: + batch_size: 64 + start_time: 1981/10/01 + end_time: 1995/09/30 + epochs: 30 + rho: 90 + shuffle: true + warmup: 5 + # Gradient accumulation ON. `batch_size` is now the MICRO-batch; the optimizer + # steps once per `grad_accum_steps` micro-batches, so the effective batch is + # 64 * 4 = 256 gauges and the gradient is averaged over ~4x more basins. + # + # THE ARITHMETIC — read this before trusting a run: + # 1,841 gauges / 64 = 28 full + 1 partial = 29 micro-batches per epoch + # (accumulation sets `drop_last = false`, driver.rs:268-275, so the + # 49-gauge tail is KEPT — unlike the non-accumulating path) + # 29 / 20 = one group of 20 + one partial group of 9. + # A partial group STILL STEPS: only `micros_drawn == 0` breaks the loop + # (driver.rs:415), and the 1/Σn renormalization makes unequal group sizes + # exact. So this is 2 optimizer steps per epoch, not 1. + # x 30 epochs = 60 TOTAL UPDATES + # + # Effective batch is 20*64 = 1,280 of 1,841 gauges — 70% of the training set + # in a single gradient, i.e. near full-batch. Very low variance, very few steps. + # + # THAT IS THE RISK. Calibration: the run that beat the summed-Q' baseline + # (2026-07-30T00-24, median NSE 0.6799) took 180 updates at lr 1e-3 and moved + # the head L2(dw) = 9.75. The 2026-08-04 run that hit NSE 0.6782 took 280. + # Adam's per-weight displacement is bounded by ~lr per step, so the budget + # here is 40*0.001 + 20*0.0005 = 0.050 — about 38% of the 0.132 that the + # 180-step run realized. Expect under-training. + # + # If it under-trains, prefer more STEPS over a hotter lr: + # * `grad_accum_steps: 20` -> 7 (5 steps/epoch, 150 updates, eff batch 448) + # * `grad_accum_steps: 20` -> 4 (8 steps/epoch, 240 updates, eff batch 256) + # Wall clock is forward-pass bound (~27.5 s/micro-batch), so all of these cost + # the SAME time — 29 forwards per epoch regardless. Lowering accum_steps buys + # optimizer steps for free; it only trades gradient variance for step count. + # + # The historical `grad_accum_steps: 37` was WORSE than either: with 28-37 + # micro-batches per epoch it collapsed a whole epoch into ONE update, so 50 + # epochs bought 50 updates for 1,850 forward passes. Wall clock here is + # forward-pass bound (~27.5 s/micro-batch), not step bound, so accumulating + # costs nothing in time — it only trades update count for gradient quality. + # + # ALWAYS verify the head actually moved (L2(dw)) before believing a result: + # the 2026-07-31 AdaDelta run reported plausible losses while L2(dw) = 1.9e-4, + # i.e. it never left initialization. + use_grad_accum: true + grad_accum_steps: 20 + # Adam, reverting the 2026-07-31 AdaDelta run. That run's head moved + # L2(dw) = 1.9e-4 (0.0004% of ||w||) across all 30 epochs — it never left + # initialization. Cause: `experiment.learning_rate` IS applied to AdaDelta + # (driver.rs:404 passes lr to optimizer.step; adadelta.rs:123 multiplies + # delta by it) despite config.rs:211-213 documenting the key as inert, so + # every update ran 1000x scaled down. dHBV intends AdaDelta at lr=1.0. + optimizer: adam + loss: + kind: nse-batch + # Step schedule. Keys are the FIRST epoch at which the lr applies + # (`resolve_lr` takes the largest key <= epoch, 1-indexed). At 2 optimizer + # steps per epoch (see the accumulation note above): + # epochs 1-20 @ 0.001 (40 updates) + # epochs 21-30 @ 0.0005 (20 updates) + # Total displacement budget ~0.050. The decay lands 2/3 of the way through, + # which is the right shape — the concern is the step COUNT, not the schedule. + # + # Watch `n_at_floor` in the log: the lr=0.01 run pinned 47.3% of CONUS at the + # 0.015 floor, and both 2026-08-04 runs collapsed from 0% to 25-36% between + # epochs 3 and 5 at lr 2e-3. Starting at 1e-3 is deliberately cooler than that. + # Near-full-batch gradients (1,280 gauges) should also damp the collapse, since + # it looked like a high-variance small-batch effect — but that is a hypothesis, + # not an established result. If `n_at_floor` climbs past a few percent, the lr + # is still too hot. + learning_rate: + 1: 0.005 + 11: 0.001 + 21: 0.0005 + grad_clip_max_norm: 1.0 + +kan_head: + hidden_size: 21 + num_hidden_layers: 2 + # B-spline grid intervals (`num` in pykan). Matches DDR's + # ~/projects/ddr/config/merit_training_config.yaml::kan.grid. + grid: 50 + # B-spline order. Matches DDR's ::kan.k. (pykan's KANLayer default is 3, + # but DDR overrides to 2 for production.) + k: 2 + input_var_names: + - SoilGrids1km_clay + - aridity + - meanelevation + - meanP + - NDVI + - meanslope + - log10_uparea + - SoilGrids1km_sand + - ETPOT_Hargr + - Porosity + learnable_parameters: + - n + - q_spatial + - p_spatial + # Precip-conditioned mass-preserving daily->hourly disaggregation head. + # The within-day shape is driven by the hourly AORC precip window + # [d-1,d,d+1] (+ daily-Q taps + static attrs); the daily mean is conserved + # exactly. Loss stays the historical L1 (no experiment.loss block). + # Warm-started FROZEN from the standalone capacity pretrain (chunk_days=1); + # architecture fields must match that checkpoint or load_record fails. + # `enabled: false` makes this whole block inert — the loader uses flat + # repeat-24 (nearest) upsampling instead. Flip to true for the head-on + # arm of the ablation; nothing else needs to change. + disaggregation: + enabled: false + hidden_size: 16 + num_hidden_layers: 2 + grid: 20 + k: 3 + boundary_blend: 0.0 + chunk_days: 1 + pretrained_checkpoint: /home/tbindas/projects/ddrs/output/disagg_pretrain/capacity_chunk1.mpk + freeze: true + +# Routing-engine knobs — DDR's MERIT defaults (matches mock_config in tests). +params: + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] + attribute_minimums: + discharge: 1.0e-4 + slope: 1.0e-3 + velocity: 0.01 + depth: 0.01 + bottom_width: 0.01 + defaults: + p_spatial: 21.0 + log_space_parameters: + - p_spatial + # Corrected physics instead of DDR's formulation. With `false` the routing + # core uses: + # * exact trapezoidal celerity c = v·β, β = 5/3 − (4/3)·A·√(1+z²)/(T·P) + # instead of the wide-rectangular 5/3 limit (which is 22-27% high for the + # channels this code builds — κ = b/y ≈ 0.7-1.8, so β ≈ 1.31-1.36), and + # * Cunge-derived Muskingum X = clamp(0.5(1 − Q/(T·S·c·L)), 0, 0.5) instead + # of the constant 0.3, which was injecting a median 28× excess numerical + # diffusion. + # Negative-discharge counting before the S28 clamp is reported in BOTH modes. + # This breaks `compare_ddr_sandbox`'s ABSOLUTE MATCH by design — that example + # must be run with the default `ddr_match: true`. + # See `.claude/PHYSICS-CORRECTIONS.md`. + ddr_match: false + # Provably zero negative discharges. Clamps the Muskingum INPUTS (K, X) into + # the non-negative-coefficient window `2X <= Cr <= 2(1-X)`, Cr = dt/K: + # K <- max(K, dt*(1+d)/2) => Cr <= 2/(1+d) + # X <- min(X_cunge, (1-d)*Cr/2, (1-d)*(1-Cr/2)) + # with d = POSITIVITY_DELTA = 1e-2. Clamping K and X (not the coefficients) + # keeps `c1+c2+c3 = 1` exact, so mass is conserved; clamping c3 would not. + # With c1,c3 >= 0 and q_t,q' > 0, induction over the topological order of the + # forward substitution proves the whole solve is non-negative. + # + # Measured: 0 / 197,461,880 solves on 1,841 gauges (was 55,181, 0.0279%), + # replicated across two windows, two backends, and two trained heads. + # + # COST — this is not free. The X cap binds on 95.3% of reach-timesteps and + # median X falls 0.4976 -> 0.0794 (6.3x), so it largely REPLACES the Cunge X + # rather than shading it: numerical diffusion becomes stability-set instead of + # matched to hydraulic diffusivity. NSE/KGE impact is UNMEASURED, and this is + # the first training run to use the flag. Compare against a matched + # `enforce_positivity: false` control before trusting the result. + # + # Root-cause alternative: median Cr is 0.226, i.e. the typical MERIT reach is + # ~4.4x too LONG for an hourly step. Subdividing reaches ~4x brings Cr to ~1, + # where X_max -> 0.5 and the cap stops biting. See .claude/PHYSICS-CORRECTIONS.md. + # + # Requires ddr_match: false (rejected at load otherwise). + # OFF. It delivered provably zero negative solves (0 / 197,461,880), but the + # cure was worse than the disease: capping X collapsed the median from 0.4976 + # to 0.0794 on 95.3% of reach-timesteps, and because the cap is applied at + # RUNTIME to a learned celerity it made X ~ Cr ~ 1/n — handing the optimizer a + # lever it rode straight to the roughness floor. Same lr, flag on vs off: + # n_at_floor 35% -> 97.9%. Superseded by reach subdivision (a BUILD-TIME fix + # with no gradient path): see docs/superpowers/plans/2026-08-05-reach-subdivision.md + enforce_positivity: false + # cuSPARSE triangular solve. Pair with `ddrs run --backend cuda`; it is inert + # under `--backend cpu`. This is a DIFFERENT `Backward` impl from the host + # solve, sitting directly downstream of the ddr_match/enforce_positivity + # terms — gradient-checked on GPU by `tests/cuda_backward_parity.rs`. + # CPU set (2026-08-08): host triangular solve. + sparse_solver: cpu + # MUST be false when ddr_match is false: the fused CUDA-graph kernel + # (cuda_graph/geometry_kernel.rs) hardcodes DDR's 5/3 celerity, so a captured + # graph would replay the DDR forward while the backward used the corrected + # chain rule — a silent gradient mismatch. Config load rejects the pair. + use_cuda_graphs: false + +# Test-mode overlay. When the binary is run with --mode testing, these keys +# replace the corresponding ones in `experiment:`. Absent keys inherit. +# +# IMPORTANT: batch_size SEMANTIC SHIFTS between modes: +# - experiment.batch_size (training) = number of GAUGES per mini-batch +# - testing.batch_size = number of DAYS per chunk +testing: + start_time: 1995/10/01 + end_time: 2010/09/30 + batch_size: 15 # DAYS, not gauges + rho: null # disabled in test mode diff --git a/config/experiments/tau_interp_g3000.yaml b/config/experiments/tau_interp_g3000.yaml new file mode 100644 index 0000000..9489d54 --- /dev/null +++ b/config/experiments/tau_interp_g3000.yaml @@ -0,0 +1,235 @@ +# ddrs MERIT training config. +# +# Hyperparameters are pulled verbatim from +# ~/projects/ddr/config/merit_training_config.yaml. +# Forcing is daily; the loader interpolates to hourly (repeat × 24, trim to +# n_hourly) before feeding the Muskingum-Cunge engine — same pattern as +# DDR's StreamflowReader. + +mode: training +workflow: train-and-test # ddrs plan/run picks this up; override with --workflow X +geodataset: merit +seed: 42 +np_seed: 42 + +# Source paths — read in place by ddrs's Rust loaders (zarrs + icechunk + netcdf). +# No export/materialization step; the harness reads DDR's live data sources +# straight from disk. +data_sources: + attributes: /home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc + conus_adjacency: /home/tbindas/projects/ddr/data/merit_conus_adjacency.zarr + gages_adjacency: /home/tbindas/projects/ddr/data/merit_gages_conus_adjacency.zarr + streamflow: /mnt/ssd1/data/icechunk/daily_dhbv2_distributed_aorc2f_merit_unit_catchments.ic + observations: /mnt/ssd1/data/icechunk/usgs_daily_observations + gages: /home/tbindas/projects/ddr/references/gage_info/gages_3000.csv + # Hourly AORC precip (zarr v3, CONUS) — drives the precip-conditioned + # mass-preserving disaggregation head below. + aorc_precip: /mnt/ssd1/data/aorc/merit_unit_catchments.zarr + +experiment: + batch_size: 64 + start_time: 1981/10/01 + end_time: 1995/09/30 + epochs: 30 + rho: 90 + shuffle: true + warmup: 5 + # Gradient accumulation ON. `batch_size` is now the MICRO-batch; the optimizer + # steps once per `grad_accum_steps` micro-batches, so the effective batch is + # 64 * 4 = 256 gauges and the gradient is averaged over ~4x more basins. + # + # THE ARITHMETIC — read this before trusting a run: + # 1,841 gauges / 64 = 28 full + 1 partial = 29 micro-batches per epoch + # (accumulation sets `drop_last = false`, driver.rs:268-275, so the + # 49-gauge tail is KEPT — unlike the non-accumulating path) + # 29 / 20 = one group of 20 + one partial group of 9. + # A partial group STILL STEPS: only `micros_drawn == 0` breaks the loop + # (driver.rs:415), and the 1/Σn renormalization makes unequal group sizes + # exact. So this is 2 optimizer steps per epoch, not 1. + # x 30 epochs = 60 TOTAL UPDATES + # + # Effective batch is 20*64 = 1,280 of 1,841 gauges — 70% of the training set + # in a single gradient, i.e. near full-batch. Very low variance, very few steps. + # + # THAT IS THE RISK. Calibration: the run that beat the summed-Q' baseline + # (2026-07-30T00-24, median NSE 0.6799) took 180 updates at lr 1e-3 and moved + # the head L2(dw) = 9.75. The 2026-08-04 run that hit NSE 0.6782 took 280. + # Adam's per-weight displacement is bounded by ~lr per step, so the budget + # here is 40*0.001 + 20*0.0005 = 0.050 — about 38% of the 0.132 that the + # 180-step run realized. Expect under-training. + # + # If it under-trains, prefer more STEPS over a hotter lr: + # * `grad_accum_steps: 20` -> 7 (5 steps/epoch, 150 updates, eff batch 448) + # * `grad_accum_steps: 20` -> 4 (8 steps/epoch, 240 updates, eff batch 256) + # Wall clock is forward-pass bound (~27.5 s/micro-batch), so all of these cost + # the SAME time — 29 forwards per epoch regardless. Lowering accum_steps buys + # optimizer steps for free; it only trades gradient variance for step count. + # + # The historical `grad_accum_steps: 37` was WORSE than either: with 28-37 + # micro-batches per epoch it collapsed a whole epoch into ONE update, so 50 + # epochs bought 50 updates for 1,850 forward passes. Wall clock here is + # forward-pass bound (~27.5 s/micro-batch), not step bound, so accumulating + # costs nothing in time — it only trades update count for gradient quality. + # + # ALWAYS verify the head actually moved (L2(dw)) before believing a result: + # the 2026-07-31 AdaDelta run reported plausible losses while L2(dw) = 1.9e-4, + # i.e. it never left initialization. + use_grad_accum: true + grad_accum_steps: 20 + # Adam, reverting the 2026-07-31 AdaDelta run. That run's head moved + # L2(dw) = 1.9e-4 (0.0004% of ||w||) across all 30 epochs — it never left + # initialization. Cause: `experiment.learning_rate` IS applied to AdaDelta + # (driver.rs:404 passes lr to optimizer.step; adadelta.rs:123 multiplies + # delta by it) despite config.rs:211-213 documenting the key as inert, so + # every update ran 1000x scaled down. dHBV intends AdaDelta at lr=1.0. + optimizer: adam + loss: + kind: nse-batch + # Step schedule. Keys are the FIRST epoch at which the lr applies + # (`resolve_lr` takes the largest key <= epoch, 1-indexed). At 2 optimizer + # steps per epoch (see the accumulation note above): + # epochs 1-20 @ 0.001 (40 updates) + # epochs 21-30 @ 0.0005 (20 updates) + # Total displacement budget ~0.050. The decay lands 2/3 of the way through, + # which is the right shape — the concern is the step COUNT, not the schedule. + # + # Watch `n_at_floor` in the log: the lr=0.01 run pinned 47.3% of CONUS at the + # 0.015 floor, and both 2026-08-04 runs collapsed from 0% to 25-36% between + # epochs 3 and 5 at lr 2e-3. Starting at 1e-3 is deliberately cooler than that. + # Near-full-batch gradients (1,280 gauges) should also damp the collapse, since + # it looked like a high-variance small-batch effect — but that is a hypothesis, + # not an established result. If `n_at_floor` climbs past a few percent, the lr + # is still too hot. + learning_rate: + 1: 0.005 + 11: 0.001 + 21: 0.0005 + grad_clip_max_norm: 1.0 + +kan_head: + hidden_size: 21 + num_hidden_layers: 2 + # B-spline grid intervals (`num` in pykan). Matches DDR's + # ~/projects/ddr/config/merit_training_config.yaml::kan.grid. + grid: 50 + # B-spline order. Matches DDR's ::kan.k. (pykan's KANLayer default is 3, + # but DDR overrides to 2 for production.) + k: 2 + input_var_names: + - SoilGrids1km_clay + - aridity + - meanelevation + - meanP + - NDVI + - meanslope + - log10_uparea + - SoilGrids1km_sand + - ETPOT_Hargr + - Porosity + learnable_parameters: + - n + - q_spatial + - p_spatial + # Precip-conditioned mass-preserving daily->hourly disaggregation head. + # The within-day shape is driven by the hourly AORC precip window + # [d-1,d,d+1] (+ daily-Q taps + static attrs); the daily mean is conserved + # exactly. Loss stays the historical L1 (no experiment.loss block). + # Warm-started FROZEN from the standalone capacity pretrain (chunk_days=1); + # architecture fields must match that checkpoint or load_record fails. + # `enabled: false` makes this whole block inert — the loader uses flat + # repeat-24 (nearest) upsampling instead. Flip to true for the head-on + # arm of the ablation; nothing else needs to change. + disaggregation: + enabled: false + hidden_size: 16 + num_hidden_layers: 2 + grid: 20 + k: 3 + boundary_blend: 0.0 + chunk_days: 1 + pretrained_checkpoint: /home/tbindas/projects/ddrs/output/disagg_pretrain/capacity_chunk1.mpk + freeze: true + +# Routing-engine knobs — DDR's MERIT defaults (matches mock_config in tests). +params: + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] + attribute_minimums: + discharge: 1.0e-4 + slope: 1.0e-3 + velocity: 0.01 + depth: 0.01 + bottom_width: 0.01 + defaults: + p_spatial: 21.0 + log_space_parameters: + - p_spatial + # Corrected physics instead of DDR's formulation. With `false` the routing + # core uses: + # * exact trapezoidal celerity c = v·β, β = 5/3 − (4/3)·A·√(1+z²)/(T·P) + # instead of the wide-rectangular 5/3 limit (which is 22-27% high for the + # channels this code builds — κ = b/y ≈ 0.7-1.8, so β ≈ 1.31-1.36), and + # * Cunge-derived Muskingum X = clamp(0.5(1 − Q/(T·S·c·L)), 0, 0.5) instead + # of the constant 0.3, which was injecting a median 28× excess numerical + # diffusion. + # Negative-discharge counting before the S28 clamp is reported in BOTH modes. + # This breaks `compare_ddr_sandbox`'s ABSOLUTE MATCH by design — that example + # must be run with the default `ddr_match: true`. + # See `.claude/PHYSICS-CORRECTIONS.md`. + ddr_match: false + # Provably zero negative discharges. Clamps the Muskingum INPUTS (K, X) into + # the non-negative-coefficient window `2X <= Cr <= 2(1-X)`, Cr = dt/K: + # K <- max(K, dt*(1+d)/2) => Cr <= 2/(1+d) + # X <- min(X_cunge, (1-d)*Cr/2, (1-d)*(1-Cr/2)) + # with d = POSITIVITY_DELTA = 1e-2. Clamping K and X (not the coefficients) + # keeps `c1+c2+c3 = 1` exact, so mass is conserved; clamping c3 would not. + # With c1,c3 >= 0 and q_t,q' > 0, induction over the topological order of the + # forward substitution proves the whole solve is non-negative. + # + # Measured: 0 / 197,461,880 solves on 1,841 gauges (was 55,181, 0.0279%), + # replicated across two windows, two backends, and two trained heads. + # + # COST — this is not free. The X cap binds on 95.3% of reach-timesteps and + # median X falls 0.4976 -> 0.0794 (6.3x), so it largely REPLACES the Cunge X + # rather than shading it: numerical diffusion becomes stability-set instead of + # matched to hydraulic diffusivity. NSE/KGE impact is UNMEASURED, and this is + # the first training run to use the flag. Compare against a matched + # `enforce_positivity: false` control before trusting the result. + # + # Root-cause alternative: median Cr is 0.226, i.e. the typical MERIT reach is + # ~4.4x too LONG for an hourly step. Subdividing reaches ~4x brings Cr to ~1, + # where X_max -> 0.5 and the cap stops biting. See .claude/PHYSICS-CORRECTIONS.md. + # + # Requires ddr_match: false (rejected at load otherwise). + # OFF. It delivered provably zero negative solves (0 / 197,461,880), but the + # cure was worse than the disease: capping X collapsed the median from 0.4976 + # to 0.0794 on 95.3% of reach-timesteps, and because the cap is applied at + # RUNTIME to a learned celerity it made X ~ Cr ~ 1/n — handing the optimizer a + # lever it rode straight to the roughness floor. Same lr, flag on vs off: + # n_at_floor 35% -> 97.9%. Superseded by reach subdivision (a BUILD-TIME fix + # with no gradient path): see docs/superpowers/plans/2026-08-05-reach-subdivision.md + enforce_positivity: false + # cuSPARSE triangular solve. Pair with `ddrs run --backend cuda`; it is inert + # under `--backend cpu`. This is a DIFFERENT `Backward` impl from the host + # solve, sitting directly downstream of the ddr_match/enforce_positivity + # terms — gradient-checked on GPU by `tests/cuda_backward_parity.rs`. + sparse_solver: cuda + # MUST be false when ddr_match is false: the fused CUDA-graph kernel + # (cuda_graph/geometry_kernel.rs) hardcodes DDR's 5/3 celerity, so a captured + # graph would replay the DDR forward while the backward used the corrected + # chain rule — a silent gradient mismatch. Config load rejects the pair. + use_cuda_graphs: false + +# Test-mode overlay. When the binary is run with --mode testing, these keys +# replace the corresponding ones in `experiment:`. Absent keys inherit. +# +# IMPORTANT: batch_size SEMANTIC SHIFTS between modes: +# - experiment.batch_size (training) = number of GAUGES per mini-batch +# - testing.batch_size = number of DAYS per chunk +testing: + start_time: 1995/10/01 + end_time: 2010/09/30 + batch_size: 15 # DAYS, not gauges + rho: null # disabled in test mode diff --git a/config/experiments/tau_src_aorc2f_lumped.yaml b/config/experiments/tau_src_aorc2f_lumped.yaml new file mode 100644 index 0000000..a1315b4 --- /dev/null +++ b/config/experiments/tau_src_aorc2f_lumped.yaml @@ -0,0 +1,235 @@ +# ddrs MERIT training config. +# +# Hyperparameters are pulled verbatim from +# ~/projects/ddr/config/merit_training_config.yaml. +# Forcing is daily; the loader interpolates to hourly (repeat × 24, trim to +# n_hourly) before feeding the Muskingum-Cunge engine — same pattern as +# DDR's StreamflowReader. + +mode: training +workflow: train-and-test # ddrs plan/run picks this up; override with --workflow X +geodataset: merit +seed: 42 +np_seed: 42 + +# Source paths — read in place by ddrs's Rust loaders (zarrs + icechunk + netcdf). +# No export/materialization step; the harness reads DDR's live data sources +# straight from disk. +data_sources: + attributes: /home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc + conus_adjacency: /home/tbindas/projects/ddr/data/merit_conus_adjacency.zarr + gages_adjacency: /home/tbindas/projects/ddr/data/merit_gages_conus_adjacency.zarr + streamflow: /mnt/ssd1/data/icechunk/daily_dhbv2_lumped_aorc2f_merit_unit_catchments.ic + observations: /mnt/ssd1/data/icechunk/usgs_daily_observations + gages: /home/tbindas/projects/ddr/references/gage_info/gages_3000.csv + # Hourly AORC precip (zarr v3, CONUS) — drives the precip-conditioned + # mass-preserving disaggregation head below. + aorc_precip: /mnt/ssd1/data/aorc/merit_unit_catchments.zarr + +experiment: + batch_size: 64 + start_time: 1981/10/01 + end_time: 1995/09/30 + epochs: 30 + rho: 90 + shuffle: true + warmup: 5 + # Gradient accumulation ON. `batch_size` is now the MICRO-batch; the optimizer + # steps once per `grad_accum_steps` micro-batches, so the effective batch is + # 64 * 4 = 256 gauges and the gradient is averaged over ~4x more basins. + # + # THE ARITHMETIC — read this before trusting a run: + # 1,841 gauges / 64 = 28 full + 1 partial = 29 micro-batches per epoch + # (accumulation sets `drop_last = false`, driver.rs:268-275, so the + # 49-gauge tail is KEPT — unlike the non-accumulating path) + # 29 / 20 = one group of 20 + one partial group of 9. + # A partial group STILL STEPS: only `micros_drawn == 0` breaks the loop + # (driver.rs:415), and the 1/Σn renormalization makes unequal group sizes + # exact. So this is 2 optimizer steps per epoch, not 1. + # x 30 epochs = 60 TOTAL UPDATES + # + # Effective batch is 20*64 = 1,280 of 1,841 gauges — 70% of the training set + # in a single gradient, i.e. near full-batch. Very low variance, very few steps. + # + # THAT IS THE RISK. Calibration: the run that beat the summed-Q' baseline + # (2026-07-30T00-24, median NSE 0.6799) took 180 updates at lr 1e-3 and moved + # the head L2(dw) = 9.75. The 2026-08-04 run that hit NSE 0.6782 took 280. + # Adam's per-weight displacement is bounded by ~lr per step, so the budget + # here is 40*0.001 + 20*0.0005 = 0.050 — about 38% of the 0.132 that the + # 180-step run realized. Expect under-training. + # + # If it under-trains, prefer more STEPS over a hotter lr: + # * `grad_accum_steps: 20` -> 7 (5 steps/epoch, 150 updates, eff batch 448) + # * `grad_accum_steps: 20` -> 4 (8 steps/epoch, 240 updates, eff batch 256) + # Wall clock is forward-pass bound (~27.5 s/micro-batch), so all of these cost + # the SAME time — 29 forwards per epoch regardless. Lowering accum_steps buys + # optimizer steps for free; it only trades gradient variance for step count. + # + # The historical `grad_accum_steps: 37` was WORSE than either: with 28-37 + # micro-batches per epoch it collapsed a whole epoch into ONE update, so 50 + # epochs bought 50 updates for 1,850 forward passes. Wall clock here is + # forward-pass bound (~27.5 s/micro-batch), not step bound, so accumulating + # costs nothing in time — it only trades update count for gradient quality. + # + # ALWAYS verify the head actually moved (L2(dw)) before believing a result: + # the 2026-07-31 AdaDelta run reported plausible losses while L2(dw) = 1.9e-4, + # i.e. it never left initialization. + use_grad_accum: true + grad_accum_steps: 20 + # Adam, reverting the 2026-07-31 AdaDelta run. That run's head moved + # L2(dw) = 1.9e-4 (0.0004% of ||w||) across all 30 epochs — it never left + # initialization. Cause: `experiment.learning_rate` IS applied to AdaDelta + # (driver.rs:404 passes lr to optimizer.step; adadelta.rs:123 multiplies + # delta by it) despite config.rs:211-213 documenting the key as inert, so + # every update ran 1000x scaled down. dHBV intends AdaDelta at lr=1.0. + optimizer: adam + loss: + kind: nse-batch + # Step schedule. Keys are the FIRST epoch at which the lr applies + # (`resolve_lr` takes the largest key <= epoch, 1-indexed). At 2 optimizer + # steps per epoch (see the accumulation note above): + # epochs 1-20 @ 0.001 (40 updates) + # epochs 21-30 @ 0.0005 (20 updates) + # Total displacement budget ~0.050. The decay lands 2/3 of the way through, + # which is the right shape — the concern is the step COUNT, not the schedule. + # + # Watch `n_at_floor` in the log: the lr=0.01 run pinned 47.3% of CONUS at the + # 0.015 floor, and both 2026-08-04 runs collapsed from 0% to 25-36% between + # epochs 3 and 5 at lr 2e-3. Starting at 1e-3 is deliberately cooler than that. + # Near-full-batch gradients (1,280 gauges) should also damp the collapse, since + # it looked like a high-variance small-batch effect — but that is a hypothesis, + # not an established result. If `n_at_floor` climbs past a few percent, the lr + # is still too hot. + learning_rate: + 1: 0.005 + 11: 0.001 + 21: 0.0005 + grad_clip_max_norm: 1.0 + +kan_head: + hidden_size: 21 + num_hidden_layers: 2 + # B-spline grid intervals (`num` in pykan). Matches DDR's + # ~/projects/ddr/config/merit_training_config.yaml::kan.grid. + grid: 50 + # B-spline order. Matches DDR's ::kan.k. (pykan's KANLayer default is 3, + # but DDR overrides to 2 for production.) + k: 2 + input_var_names: + - SoilGrids1km_clay + - aridity + - meanelevation + - meanP + - NDVI + - meanslope + - log10_uparea + - SoilGrids1km_sand + - ETPOT_Hargr + - Porosity + learnable_parameters: + - n + - q_spatial + - p_spatial + # Precip-conditioned mass-preserving daily->hourly disaggregation head. + # The within-day shape is driven by the hourly AORC precip window + # [d-1,d,d+1] (+ daily-Q taps + static attrs); the daily mean is conserved + # exactly. Loss stays the historical L1 (no experiment.loss block). + # Warm-started FROZEN from the standalone capacity pretrain (chunk_days=1); + # architecture fields must match that checkpoint or load_record fails. + # `enabled: false` makes this whole block inert — the loader uses flat + # repeat-24 (nearest) upsampling instead. Flip to true for the head-on + # arm of the ablation; nothing else needs to change. + disaggregation: + enabled: false + hidden_size: 16 + num_hidden_layers: 2 + grid: 20 + k: 3 + boundary_blend: 0.0 + chunk_days: 1 + pretrained_checkpoint: /home/tbindas/projects/ddrs/output/disagg_pretrain/capacity_chunk1.mpk + freeze: true + +# Routing-engine knobs — DDR's MERIT defaults (matches mock_config in tests). +params: + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] + attribute_minimums: + discharge: 1.0e-4 + slope: 1.0e-3 + velocity: 0.01 + depth: 0.01 + bottom_width: 0.01 + defaults: + p_spatial: 21.0 + log_space_parameters: + - p_spatial + # Corrected physics instead of DDR's formulation. With `false` the routing + # core uses: + # * exact trapezoidal celerity c = v·β, β = 5/3 − (4/3)·A·√(1+z²)/(T·P) + # instead of the wide-rectangular 5/3 limit (which is 22-27% high for the + # channels this code builds — κ = b/y ≈ 0.7-1.8, so β ≈ 1.31-1.36), and + # * Cunge-derived Muskingum X = clamp(0.5(1 − Q/(T·S·c·L)), 0, 0.5) instead + # of the constant 0.3, which was injecting a median 28× excess numerical + # diffusion. + # Negative-discharge counting before the S28 clamp is reported in BOTH modes. + # This breaks `compare_ddr_sandbox`'s ABSOLUTE MATCH by design — that example + # must be run with the default `ddr_match: true`. + # See `.claude/PHYSICS-CORRECTIONS.md`. + ddr_match: false + # Provably zero negative discharges. Clamps the Muskingum INPUTS (K, X) into + # the non-negative-coefficient window `2X <= Cr <= 2(1-X)`, Cr = dt/K: + # K <- max(K, dt*(1+d)/2) => Cr <= 2/(1+d) + # X <- min(X_cunge, (1-d)*Cr/2, (1-d)*(1-Cr/2)) + # with d = POSITIVITY_DELTA = 1e-2. Clamping K and X (not the coefficients) + # keeps `c1+c2+c3 = 1` exact, so mass is conserved; clamping c3 would not. + # With c1,c3 >= 0 and q_t,q' > 0, induction over the topological order of the + # forward substitution proves the whole solve is non-negative. + # + # Measured: 0 / 197,461,880 solves on 1,841 gauges (was 55,181, 0.0279%), + # replicated across two windows, two backends, and two trained heads. + # + # COST — this is not free. The X cap binds on 95.3% of reach-timesteps and + # median X falls 0.4976 -> 0.0794 (6.3x), so it largely REPLACES the Cunge X + # rather than shading it: numerical diffusion becomes stability-set instead of + # matched to hydraulic diffusivity. NSE/KGE impact is UNMEASURED, and this is + # the first training run to use the flag. Compare against a matched + # `enforce_positivity: false` control before trusting the result. + # + # Root-cause alternative: median Cr is 0.226, i.e. the typical MERIT reach is + # ~4.4x too LONG for an hourly step. Subdividing reaches ~4x brings Cr to ~1, + # where X_max -> 0.5 and the cap stops biting. See .claude/PHYSICS-CORRECTIONS.md. + # + # Requires ddr_match: false (rejected at load otherwise). + # OFF. It delivered provably zero negative solves (0 / 197,461,880), but the + # cure was worse than the disease: capping X collapsed the median from 0.4976 + # to 0.0794 on 95.3% of reach-timesteps, and because the cap is applied at + # RUNTIME to a learned celerity it made X ~ Cr ~ 1/n — handing the optimizer a + # lever it rode straight to the roughness floor. Same lr, flag on vs off: + # n_at_floor 35% -> 97.9%. Superseded by reach subdivision (a BUILD-TIME fix + # with no gradient path): see docs/superpowers/plans/2026-08-05-reach-subdivision.md + enforce_positivity: false + # cuSPARSE triangular solve. Pair with `ddrs run --backend cuda`; it is inert + # under `--backend cpu`. This is a DIFFERENT `Backward` impl from the host + # solve, sitting directly downstream of the ddr_match/enforce_positivity + # terms — gradient-checked on GPU by `tests/cuda_backward_parity.rs`. + sparse_solver: cuda + # MUST be false when ddr_match is false: the fused CUDA-graph kernel + # (cuda_graph/geometry_kernel.rs) hardcodes DDR's 5/3 celerity, so a captured + # graph would replay the DDR forward while the backward used the corrected + # chain rule — a silent gradient mismatch. Config load rejects the pair. + use_cuda_graphs: false + +# Test-mode overlay. When the binary is run with --mode testing, these keys +# replace the corresponding ones in `experiment:`. Absent keys inherit. +# +# IMPORTANT: batch_size SEMANTIC SHIFTS between modes: +# - experiment.batch_size (training) = number of GAUGES per mini-batch +# - testing.batch_size = number of DAYS per chunk +testing: + start_time: 1995/10/01 + end_time: 2010/09/30 + batch_size: 15 # DAYS, not gauges + rho: null # disabled in test mode diff --git a/config/experiments/tau_src_daily_lstm.yaml b/config/experiments/tau_src_daily_lstm.yaml new file mode 100644 index 0000000..ce63ea3 --- /dev/null +++ b/config/experiments/tau_src_daily_lstm.yaml @@ -0,0 +1,235 @@ +# ddrs MERIT training config. +# +# Hyperparameters are pulled verbatim from +# ~/projects/ddr/config/merit_training_config.yaml. +# Forcing is daily; the loader interpolates to hourly (repeat × 24, trim to +# n_hourly) before feeding the Muskingum-Cunge engine — same pattern as +# DDR's StreamflowReader. + +mode: training +workflow: train-and-test # ddrs plan/run picks this up; override with --workflow X +geodataset: merit +seed: 42 +np_seed: 42 + +# Source paths — read in place by ddrs's Rust loaders (zarrs + icechunk + netcdf). +# No export/materialization step; the harness reads DDR's live data sources +# straight from disk. +data_sources: + attributes: /home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc + conus_adjacency: /home/tbindas/projects/ddr/data/merit_conus_adjacency.zarr + gages_adjacency: /home/tbindas/projects/ddr/data/merit_gages_conus_adjacency.zarr + streamflow: /mnt/ssd1/data/icechunk/daily_lstm_merit_unit_catchments.ic + observations: /mnt/ssd1/data/icechunk/usgs_daily_observations + gages: /home/tbindas/projects/ddr/references/gage_info/gages_3000.csv + # Hourly AORC precip (zarr v3, CONUS) — drives the precip-conditioned + # mass-preserving disaggregation head below. + aorc_precip: /mnt/ssd1/data/aorc/merit_unit_catchments.zarr + +experiment: + batch_size: 64 + start_time: 1981/10/01 + end_time: 1995/09/30 + epochs: 30 + rho: 90 + shuffle: true + warmup: 5 + # Gradient accumulation ON. `batch_size` is now the MICRO-batch; the optimizer + # steps once per `grad_accum_steps` micro-batches, so the effective batch is + # 64 * 4 = 256 gauges and the gradient is averaged over ~4x more basins. + # + # THE ARITHMETIC — read this before trusting a run: + # 1,841 gauges / 64 = 28 full + 1 partial = 29 micro-batches per epoch + # (accumulation sets `drop_last = false`, driver.rs:268-275, so the + # 49-gauge tail is KEPT — unlike the non-accumulating path) + # 29 / 20 = one group of 20 + one partial group of 9. + # A partial group STILL STEPS: only `micros_drawn == 0` breaks the loop + # (driver.rs:415), and the 1/Σn renormalization makes unequal group sizes + # exact. So this is 2 optimizer steps per epoch, not 1. + # x 30 epochs = 60 TOTAL UPDATES + # + # Effective batch is 20*64 = 1,280 of 1,841 gauges — 70% of the training set + # in a single gradient, i.e. near full-batch. Very low variance, very few steps. + # + # THAT IS THE RISK. Calibration: the run that beat the summed-Q' baseline + # (2026-07-30T00-24, median NSE 0.6799) took 180 updates at lr 1e-3 and moved + # the head L2(dw) = 9.75. The 2026-08-04 run that hit NSE 0.6782 took 280. + # Adam's per-weight displacement is bounded by ~lr per step, so the budget + # here is 40*0.001 + 20*0.0005 = 0.050 — about 38% of the 0.132 that the + # 180-step run realized. Expect under-training. + # + # If it under-trains, prefer more STEPS over a hotter lr: + # * `grad_accum_steps: 20` -> 7 (5 steps/epoch, 150 updates, eff batch 448) + # * `grad_accum_steps: 20` -> 4 (8 steps/epoch, 240 updates, eff batch 256) + # Wall clock is forward-pass bound (~27.5 s/micro-batch), so all of these cost + # the SAME time — 29 forwards per epoch regardless. Lowering accum_steps buys + # optimizer steps for free; it only trades gradient variance for step count. + # + # The historical `grad_accum_steps: 37` was WORSE than either: with 28-37 + # micro-batches per epoch it collapsed a whole epoch into ONE update, so 50 + # epochs bought 50 updates for 1,850 forward passes. Wall clock here is + # forward-pass bound (~27.5 s/micro-batch), not step bound, so accumulating + # costs nothing in time — it only trades update count for gradient quality. + # + # ALWAYS verify the head actually moved (L2(dw)) before believing a result: + # the 2026-07-31 AdaDelta run reported plausible losses while L2(dw) = 1.9e-4, + # i.e. it never left initialization. + use_grad_accum: true + grad_accum_steps: 20 + # Adam, reverting the 2026-07-31 AdaDelta run. That run's head moved + # L2(dw) = 1.9e-4 (0.0004% of ||w||) across all 30 epochs — it never left + # initialization. Cause: `experiment.learning_rate` IS applied to AdaDelta + # (driver.rs:404 passes lr to optimizer.step; adadelta.rs:123 multiplies + # delta by it) despite config.rs:211-213 documenting the key as inert, so + # every update ran 1000x scaled down. dHBV intends AdaDelta at lr=1.0. + optimizer: adam + loss: + kind: nse-batch + # Step schedule. Keys are the FIRST epoch at which the lr applies + # (`resolve_lr` takes the largest key <= epoch, 1-indexed). At 2 optimizer + # steps per epoch (see the accumulation note above): + # epochs 1-20 @ 0.001 (40 updates) + # epochs 21-30 @ 0.0005 (20 updates) + # Total displacement budget ~0.050. The decay lands 2/3 of the way through, + # which is the right shape — the concern is the step COUNT, not the schedule. + # + # Watch `n_at_floor` in the log: the lr=0.01 run pinned 47.3% of CONUS at the + # 0.015 floor, and both 2026-08-04 runs collapsed from 0% to 25-36% between + # epochs 3 and 5 at lr 2e-3. Starting at 1e-3 is deliberately cooler than that. + # Near-full-batch gradients (1,280 gauges) should also damp the collapse, since + # it looked like a high-variance small-batch effect — but that is a hypothesis, + # not an established result. If `n_at_floor` climbs past a few percent, the lr + # is still too hot. + learning_rate: + 1: 0.005 + 11: 0.001 + 21: 0.0005 + grad_clip_max_norm: 1.0 + +kan_head: + hidden_size: 21 + num_hidden_layers: 2 + # B-spline grid intervals (`num` in pykan). Matches DDR's + # ~/projects/ddr/config/merit_training_config.yaml::kan.grid. + grid: 50 + # B-spline order. Matches DDR's ::kan.k. (pykan's KANLayer default is 3, + # but DDR overrides to 2 for production.) + k: 2 + input_var_names: + - SoilGrids1km_clay + - aridity + - meanelevation + - meanP + - NDVI + - meanslope + - log10_uparea + - SoilGrids1km_sand + - ETPOT_Hargr + - Porosity + learnable_parameters: + - n + - q_spatial + - p_spatial + # Precip-conditioned mass-preserving daily->hourly disaggregation head. + # The within-day shape is driven by the hourly AORC precip window + # [d-1,d,d+1] (+ daily-Q taps + static attrs); the daily mean is conserved + # exactly. Loss stays the historical L1 (no experiment.loss block). + # Warm-started FROZEN from the standalone capacity pretrain (chunk_days=1); + # architecture fields must match that checkpoint or load_record fails. + # `enabled: false` makes this whole block inert — the loader uses flat + # repeat-24 (nearest) upsampling instead. Flip to true for the head-on + # arm of the ablation; nothing else needs to change. + disaggregation: + enabled: false + hidden_size: 16 + num_hidden_layers: 2 + grid: 20 + k: 3 + boundary_blend: 0.0 + chunk_days: 1 + pretrained_checkpoint: /home/tbindas/projects/ddrs/output/disagg_pretrain/capacity_chunk1.mpk + freeze: true + +# Routing-engine knobs — DDR's MERIT defaults (matches mock_config in tests). +params: + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] + attribute_minimums: + discharge: 1.0e-4 + slope: 1.0e-3 + velocity: 0.01 + depth: 0.01 + bottom_width: 0.01 + defaults: + p_spatial: 21.0 + log_space_parameters: + - p_spatial + # Corrected physics instead of DDR's formulation. With `false` the routing + # core uses: + # * exact trapezoidal celerity c = v·β, β = 5/3 − (4/3)·A·√(1+z²)/(T·P) + # instead of the wide-rectangular 5/3 limit (which is 22-27% high for the + # channels this code builds — κ = b/y ≈ 0.7-1.8, so β ≈ 1.31-1.36), and + # * Cunge-derived Muskingum X = clamp(0.5(1 − Q/(T·S·c·L)), 0, 0.5) instead + # of the constant 0.3, which was injecting a median 28× excess numerical + # diffusion. + # Negative-discharge counting before the S28 clamp is reported in BOTH modes. + # This breaks `compare_ddr_sandbox`'s ABSOLUTE MATCH by design — that example + # must be run with the default `ddr_match: true`. + # See `.claude/PHYSICS-CORRECTIONS.md`. + ddr_match: false + # Provably zero negative discharges. Clamps the Muskingum INPUTS (K, X) into + # the non-negative-coefficient window `2X <= Cr <= 2(1-X)`, Cr = dt/K: + # K <- max(K, dt*(1+d)/2) => Cr <= 2/(1+d) + # X <- min(X_cunge, (1-d)*Cr/2, (1-d)*(1-Cr/2)) + # with d = POSITIVITY_DELTA = 1e-2. Clamping K and X (not the coefficients) + # keeps `c1+c2+c3 = 1` exact, so mass is conserved; clamping c3 would not. + # With c1,c3 >= 0 and q_t,q' > 0, induction over the topological order of the + # forward substitution proves the whole solve is non-negative. + # + # Measured: 0 / 197,461,880 solves on 1,841 gauges (was 55,181, 0.0279%), + # replicated across two windows, two backends, and two trained heads. + # + # COST — this is not free. The X cap binds on 95.3% of reach-timesteps and + # median X falls 0.4976 -> 0.0794 (6.3x), so it largely REPLACES the Cunge X + # rather than shading it: numerical diffusion becomes stability-set instead of + # matched to hydraulic diffusivity. NSE/KGE impact is UNMEASURED, and this is + # the first training run to use the flag. Compare against a matched + # `enforce_positivity: false` control before trusting the result. + # + # Root-cause alternative: median Cr is 0.226, i.e. the typical MERIT reach is + # ~4.4x too LONG for an hourly step. Subdividing reaches ~4x brings Cr to ~1, + # where X_max -> 0.5 and the cap stops biting. See .claude/PHYSICS-CORRECTIONS.md. + # + # Requires ddr_match: false (rejected at load otherwise). + # OFF. It delivered provably zero negative solves (0 / 197,461,880), but the + # cure was worse than the disease: capping X collapsed the median from 0.4976 + # to 0.0794 on 95.3% of reach-timesteps, and because the cap is applied at + # RUNTIME to a learned celerity it made X ~ Cr ~ 1/n — handing the optimizer a + # lever it rode straight to the roughness floor. Same lr, flag on vs off: + # n_at_floor 35% -> 97.9%. Superseded by reach subdivision (a BUILD-TIME fix + # with no gradient path): see docs/superpowers/plans/2026-08-05-reach-subdivision.md + enforce_positivity: false + # cuSPARSE triangular solve. Pair with `ddrs run --backend cuda`; it is inert + # under `--backend cpu`. This is a DIFFERENT `Backward` impl from the host + # solve, sitting directly downstream of the ddr_match/enforce_positivity + # terms — gradient-checked on GPU by `tests/cuda_backward_parity.rs`. + sparse_solver: cuda + # MUST be false when ddr_match is false: the fused CUDA-graph kernel + # (cuda_graph/geometry_kernel.rs) hardcodes DDR's 5/3 celerity, so a captured + # graph would replay the DDR forward while the backward used the corrected + # chain rule — a silent gradient mismatch. Config load rejects the pair. + use_cuda_graphs: false + +# Test-mode overlay. When the binary is run with --mode testing, these keys +# replace the corresponding ones in `experiment:`. Absent keys inherit. +# +# IMPORTANT: batch_size SEMANTIC SHIFTS between modes: +# - experiment.batch_size (training) = number of GAUGES per mini-batch +# - testing.batch_size = number of DAYS per chunk +testing: + start_time: 1995/10/01 + end_time: 2010/09/30 + batch_size: 15 # DAYS, not gauges + rho: null # disabled in test mode diff --git a/config/experiments/tau_src_hourly_lstm.yaml b/config/experiments/tau_src_hourly_lstm.yaml new file mode 100644 index 0000000..610faf7 --- /dev/null +++ b/config/experiments/tau_src_hourly_lstm.yaml @@ -0,0 +1,235 @@ +# ddrs MERIT training config. +# +# Hyperparameters are pulled verbatim from +# ~/projects/ddr/config/merit_training_config.yaml. +# Forcing is daily; the loader interpolates to hourly (repeat × 24, trim to +# n_hourly) before feeding the Muskingum-Cunge engine — same pattern as +# DDR's StreamflowReader. + +mode: training +workflow: train-and-test # ddrs plan/run picks this up; override with --workflow X +geodataset: merit +seed: 42 +np_seed: 42 + +# Source paths — read in place by ddrs's Rust loaders (zarrs + icechunk + netcdf). +# No export/materialization step; the harness reads DDR's live data sources +# straight from disk. +data_sources: + attributes: /home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc + conus_adjacency: /home/tbindas/projects/ddr/data/merit_conus_adjacency.zarr + gages_adjacency: /home/tbindas/projects/ddr/data/merit_gages_conus_adjacency.zarr + streamflow: /mnt/ssd1/data/icechunk/hourly_lstm_merit_unit_catchments.ic + observations: /mnt/ssd1/data/icechunk/usgs_daily_observations + gages: /home/tbindas/projects/ddr/references/gage_info/gages_3000.csv + # Hourly AORC precip (zarr v3, CONUS) — drives the precip-conditioned + # mass-preserving disaggregation head below. + aorc_precip: /mnt/ssd1/data/aorc/merit_unit_catchments.zarr + +experiment: + batch_size: 64 + start_time: 1981/10/01 + end_time: 1995/09/30 + epochs: 30 + rho: 90 + shuffle: true + warmup: 5 + # Gradient accumulation ON. `batch_size` is now the MICRO-batch; the optimizer + # steps once per `grad_accum_steps` micro-batches, so the effective batch is + # 64 * 4 = 256 gauges and the gradient is averaged over ~4x more basins. + # + # THE ARITHMETIC — read this before trusting a run: + # 1,841 gauges / 64 = 28 full + 1 partial = 29 micro-batches per epoch + # (accumulation sets `drop_last = false`, driver.rs:268-275, so the + # 49-gauge tail is KEPT — unlike the non-accumulating path) + # 29 / 20 = one group of 20 + one partial group of 9. + # A partial group STILL STEPS: only `micros_drawn == 0` breaks the loop + # (driver.rs:415), and the 1/Σn renormalization makes unequal group sizes + # exact. So this is 2 optimizer steps per epoch, not 1. + # x 30 epochs = 60 TOTAL UPDATES + # + # Effective batch is 20*64 = 1,280 of 1,841 gauges — 70% of the training set + # in a single gradient, i.e. near full-batch. Very low variance, very few steps. + # + # THAT IS THE RISK. Calibration: the run that beat the summed-Q' baseline + # (2026-07-30T00-24, median NSE 0.6799) took 180 updates at lr 1e-3 and moved + # the head L2(dw) = 9.75. The 2026-08-04 run that hit NSE 0.6782 took 280. + # Adam's per-weight displacement is bounded by ~lr per step, so the budget + # here is 40*0.001 + 20*0.0005 = 0.050 — about 38% of the 0.132 that the + # 180-step run realized. Expect under-training. + # + # If it under-trains, prefer more STEPS over a hotter lr: + # * `grad_accum_steps: 20` -> 7 (5 steps/epoch, 150 updates, eff batch 448) + # * `grad_accum_steps: 20` -> 4 (8 steps/epoch, 240 updates, eff batch 256) + # Wall clock is forward-pass bound (~27.5 s/micro-batch), so all of these cost + # the SAME time — 29 forwards per epoch regardless. Lowering accum_steps buys + # optimizer steps for free; it only trades gradient variance for step count. + # + # The historical `grad_accum_steps: 37` was WORSE than either: with 28-37 + # micro-batches per epoch it collapsed a whole epoch into ONE update, so 50 + # epochs bought 50 updates for 1,850 forward passes. Wall clock here is + # forward-pass bound (~27.5 s/micro-batch), not step bound, so accumulating + # costs nothing in time — it only trades update count for gradient quality. + # + # ALWAYS verify the head actually moved (L2(dw)) before believing a result: + # the 2026-07-31 AdaDelta run reported plausible losses while L2(dw) = 1.9e-4, + # i.e. it never left initialization. + use_grad_accum: true + grad_accum_steps: 20 + # Adam, reverting the 2026-07-31 AdaDelta run. That run's head moved + # L2(dw) = 1.9e-4 (0.0004% of ||w||) across all 30 epochs — it never left + # initialization. Cause: `experiment.learning_rate` IS applied to AdaDelta + # (driver.rs:404 passes lr to optimizer.step; adadelta.rs:123 multiplies + # delta by it) despite config.rs:211-213 documenting the key as inert, so + # every update ran 1000x scaled down. dHBV intends AdaDelta at lr=1.0. + optimizer: adam + loss: + kind: nse-batch + # Step schedule. Keys are the FIRST epoch at which the lr applies + # (`resolve_lr` takes the largest key <= epoch, 1-indexed). At 2 optimizer + # steps per epoch (see the accumulation note above): + # epochs 1-20 @ 0.001 (40 updates) + # epochs 21-30 @ 0.0005 (20 updates) + # Total displacement budget ~0.050. The decay lands 2/3 of the way through, + # which is the right shape — the concern is the step COUNT, not the schedule. + # + # Watch `n_at_floor` in the log: the lr=0.01 run pinned 47.3% of CONUS at the + # 0.015 floor, and both 2026-08-04 runs collapsed from 0% to 25-36% between + # epochs 3 and 5 at lr 2e-3. Starting at 1e-3 is deliberately cooler than that. + # Near-full-batch gradients (1,280 gauges) should also damp the collapse, since + # it looked like a high-variance small-batch effect — but that is a hypothesis, + # not an established result. If `n_at_floor` climbs past a few percent, the lr + # is still too hot. + learning_rate: + 1: 0.005 + 11: 0.001 + 21: 0.0005 + grad_clip_max_norm: 1.0 + +kan_head: + hidden_size: 21 + num_hidden_layers: 2 + # B-spline grid intervals (`num` in pykan). Matches DDR's + # ~/projects/ddr/config/merit_training_config.yaml::kan.grid. + grid: 50 + # B-spline order. Matches DDR's ::kan.k. (pykan's KANLayer default is 3, + # but DDR overrides to 2 for production.) + k: 2 + input_var_names: + - SoilGrids1km_clay + - aridity + - meanelevation + - meanP + - NDVI + - meanslope + - log10_uparea + - SoilGrids1km_sand + - ETPOT_Hargr + - Porosity + learnable_parameters: + - n + - q_spatial + - p_spatial + # Precip-conditioned mass-preserving daily->hourly disaggregation head. + # The within-day shape is driven by the hourly AORC precip window + # [d-1,d,d+1] (+ daily-Q taps + static attrs); the daily mean is conserved + # exactly. Loss stays the historical L1 (no experiment.loss block). + # Warm-started FROZEN from the standalone capacity pretrain (chunk_days=1); + # architecture fields must match that checkpoint or load_record fails. + # `enabled: false` makes this whole block inert — the loader uses flat + # repeat-24 (nearest) upsampling instead. Flip to true for the head-on + # arm of the ablation; nothing else needs to change. + disaggregation: + enabled: false + hidden_size: 16 + num_hidden_layers: 2 + grid: 20 + k: 3 + boundary_blend: 0.0 + chunk_days: 1 + pretrained_checkpoint: /home/tbindas/projects/ddrs/output/disagg_pretrain/capacity_chunk1.mpk + freeze: true + +# Routing-engine knobs — DDR's MERIT defaults (matches mock_config in tests). +params: + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] + attribute_minimums: + discharge: 1.0e-4 + slope: 1.0e-3 + velocity: 0.01 + depth: 0.01 + bottom_width: 0.01 + defaults: + p_spatial: 21.0 + log_space_parameters: + - p_spatial + # Corrected physics instead of DDR's formulation. With `false` the routing + # core uses: + # * exact trapezoidal celerity c = v·β, β = 5/3 − (4/3)·A·√(1+z²)/(T·P) + # instead of the wide-rectangular 5/3 limit (which is 22-27% high for the + # channels this code builds — κ = b/y ≈ 0.7-1.8, so β ≈ 1.31-1.36), and + # * Cunge-derived Muskingum X = clamp(0.5(1 − Q/(T·S·c·L)), 0, 0.5) instead + # of the constant 0.3, which was injecting a median 28× excess numerical + # diffusion. + # Negative-discharge counting before the S28 clamp is reported in BOTH modes. + # This breaks `compare_ddr_sandbox`'s ABSOLUTE MATCH by design — that example + # must be run with the default `ddr_match: true`. + # See `.claude/PHYSICS-CORRECTIONS.md`. + ddr_match: false + # Provably zero negative discharges. Clamps the Muskingum INPUTS (K, X) into + # the non-negative-coefficient window `2X <= Cr <= 2(1-X)`, Cr = dt/K: + # K <- max(K, dt*(1+d)/2) => Cr <= 2/(1+d) + # X <- min(X_cunge, (1-d)*Cr/2, (1-d)*(1-Cr/2)) + # with d = POSITIVITY_DELTA = 1e-2. Clamping K and X (not the coefficients) + # keeps `c1+c2+c3 = 1` exact, so mass is conserved; clamping c3 would not. + # With c1,c3 >= 0 and q_t,q' > 0, induction over the topological order of the + # forward substitution proves the whole solve is non-negative. + # + # Measured: 0 / 197,461,880 solves on 1,841 gauges (was 55,181, 0.0279%), + # replicated across two windows, two backends, and two trained heads. + # + # COST — this is not free. The X cap binds on 95.3% of reach-timesteps and + # median X falls 0.4976 -> 0.0794 (6.3x), so it largely REPLACES the Cunge X + # rather than shading it: numerical diffusion becomes stability-set instead of + # matched to hydraulic diffusivity. NSE/KGE impact is UNMEASURED, and this is + # the first training run to use the flag. Compare against a matched + # `enforce_positivity: false` control before trusting the result. + # + # Root-cause alternative: median Cr is 0.226, i.e. the typical MERIT reach is + # ~4.4x too LONG for an hourly step. Subdividing reaches ~4x brings Cr to ~1, + # where X_max -> 0.5 and the cap stops biting. See .claude/PHYSICS-CORRECTIONS.md. + # + # Requires ddr_match: false (rejected at load otherwise). + # OFF. It delivered provably zero negative solves (0 / 197,461,880), but the + # cure was worse than the disease: capping X collapsed the median from 0.4976 + # to 0.0794 on 95.3% of reach-timesteps, and because the cap is applied at + # RUNTIME to a learned celerity it made X ~ Cr ~ 1/n — handing the optimizer a + # lever it rode straight to the roughness floor. Same lr, flag on vs off: + # n_at_floor 35% -> 97.9%. Superseded by reach subdivision (a BUILD-TIME fix + # with no gradient path): see docs/superpowers/plans/2026-08-05-reach-subdivision.md + enforce_positivity: false + # cuSPARSE triangular solve. Pair with `ddrs run --backend cuda`; it is inert + # under `--backend cpu`. This is a DIFFERENT `Backward` impl from the host + # solve, sitting directly downstream of the ddr_match/enforce_positivity + # terms — gradient-checked on GPU by `tests/cuda_backward_parity.rs`. + sparse_solver: cuda + # MUST be false when ddr_match is false: the fused CUDA-graph kernel + # (cuda_graph/geometry_kernel.rs) hardcodes DDR's 5/3 celerity, so a captured + # graph would replay the DDR forward while the backward used the corrected + # chain rule — a silent gradient mismatch. Config load rejects the pair. + use_cuda_graphs: false + +# Test-mode overlay. When the binary is run with --mode testing, these keys +# replace the corresponding ones in `experiment:`. Absent keys inherit. +# +# IMPORTANT: batch_size SEMANTIC SHIFTS between modes: +# - experiment.batch_size (training) = number of GAUGES per mini-batch +# - testing.batch_size = number of DAYS per chunk +testing: + start_time: 1995/10/01 + end_time: 2010/09/30 + batch_size: 15 # DAYS, not gauges + rho: null # disabled in test mode diff --git a/config/experiments/tau_src_uh_retro.yaml b/config/experiments/tau_src_uh_retro.yaml new file mode 100644 index 0000000..58adbea --- /dev/null +++ b/config/experiments/tau_src_uh_retro.yaml @@ -0,0 +1,235 @@ +# ddrs MERIT training config. +# +# Hyperparameters are pulled verbatim from +# ~/projects/ddr/config/merit_training_config.yaml. +# Forcing is daily; the loader interpolates to hourly (repeat × 24, trim to +# n_hourly) before feeding the Muskingum-Cunge engine — same pattern as +# DDR's StreamflowReader. + +mode: training +workflow: train-and-test # ddrs plan/run picks this up; override with --workflow X +geodataset: merit +seed: 42 +np_seed: 42 + +# Source paths — read in place by ddrs's Rust loaders (zarrs + icechunk + netcdf). +# No export/materialization step; the harness reads DDR's live data sources +# straight from disk. +data_sources: + attributes: /home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc + conus_adjacency: /home/tbindas/projects/ddr/data/merit_conus_adjacency.zarr + gages_adjacency: /home/tbindas/projects/ddr/data/merit_gages_conus_adjacency.zarr + streamflow: /mnt/ssd1/data/icechunk/merit_dhbv2_UH_retrospective.ic + observations: /mnt/ssd1/data/icechunk/usgs_daily_observations + gages: /home/tbindas/projects/ddr/references/gage_info/gages_3000.csv + # Hourly AORC precip (zarr v3, CONUS) — drives the precip-conditioned + # mass-preserving disaggregation head below. + aorc_precip: /mnt/ssd1/data/aorc/merit_unit_catchments.zarr + +experiment: + batch_size: 64 + start_time: 1981/10/01 + end_time: 1995/09/30 + epochs: 30 + rho: 90 + shuffle: true + warmup: 5 + # Gradient accumulation ON. `batch_size` is now the MICRO-batch; the optimizer + # steps once per `grad_accum_steps` micro-batches, so the effective batch is + # 64 * 4 = 256 gauges and the gradient is averaged over ~4x more basins. + # + # THE ARITHMETIC — read this before trusting a run: + # 1,841 gauges / 64 = 28 full + 1 partial = 29 micro-batches per epoch + # (accumulation sets `drop_last = false`, driver.rs:268-275, so the + # 49-gauge tail is KEPT — unlike the non-accumulating path) + # 29 / 20 = one group of 20 + one partial group of 9. + # A partial group STILL STEPS: only `micros_drawn == 0` breaks the loop + # (driver.rs:415), and the 1/Σn renormalization makes unequal group sizes + # exact. So this is 2 optimizer steps per epoch, not 1. + # x 30 epochs = 60 TOTAL UPDATES + # + # Effective batch is 20*64 = 1,280 of 1,841 gauges — 70% of the training set + # in a single gradient, i.e. near full-batch. Very low variance, very few steps. + # + # THAT IS THE RISK. Calibration: the run that beat the summed-Q' baseline + # (2026-07-30T00-24, median NSE 0.6799) took 180 updates at lr 1e-3 and moved + # the head L2(dw) = 9.75. The 2026-08-04 run that hit NSE 0.6782 took 280. + # Adam's per-weight displacement is bounded by ~lr per step, so the budget + # here is 40*0.001 + 20*0.0005 = 0.050 — about 38% of the 0.132 that the + # 180-step run realized. Expect under-training. + # + # If it under-trains, prefer more STEPS over a hotter lr: + # * `grad_accum_steps: 20` -> 7 (5 steps/epoch, 150 updates, eff batch 448) + # * `grad_accum_steps: 20` -> 4 (8 steps/epoch, 240 updates, eff batch 256) + # Wall clock is forward-pass bound (~27.5 s/micro-batch), so all of these cost + # the SAME time — 29 forwards per epoch regardless. Lowering accum_steps buys + # optimizer steps for free; it only trades gradient variance for step count. + # + # The historical `grad_accum_steps: 37` was WORSE than either: with 28-37 + # micro-batches per epoch it collapsed a whole epoch into ONE update, so 50 + # epochs bought 50 updates for 1,850 forward passes. Wall clock here is + # forward-pass bound (~27.5 s/micro-batch), not step bound, so accumulating + # costs nothing in time — it only trades update count for gradient quality. + # + # ALWAYS verify the head actually moved (L2(dw)) before believing a result: + # the 2026-07-31 AdaDelta run reported plausible losses while L2(dw) = 1.9e-4, + # i.e. it never left initialization. + use_grad_accum: true + grad_accum_steps: 20 + # Adam, reverting the 2026-07-31 AdaDelta run. That run's head moved + # L2(dw) = 1.9e-4 (0.0004% of ||w||) across all 30 epochs — it never left + # initialization. Cause: `experiment.learning_rate` IS applied to AdaDelta + # (driver.rs:404 passes lr to optimizer.step; adadelta.rs:123 multiplies + # delta by it) despite config.rs:211-213 documenting the key as inert, so + # every update ran 1000x scaled down. dHBV intends AdaDelta at lr=1.0. + optimizer: adam + loss: + kind: nse-batch + # Step schedule. Keys are the FIRST epoch at which the lr applies + # (`resolve_lr` takes the largest key <= epoch, 1-indexed). At 2 optimizer + # steps per epoch (see the accumulation note above): + # epochs 1-20 @ 0.001 (40 updates) + # epochs 21-30 @ 0.0005 (20 updates) + # Total displacement budget ~0.050. The decay lands 2/3 of the way through, + # which is the right shape — the concern is the step COUNT, not the schedule. + # + # Watch `n_at_floor` in the log: the lr=0.01 run pinned 47.3% of CONUS at the + # 0.015 floor, and both 2026-08-04 runs collapsed from 0% to 25-36% between + # epochs 3 and 5 at lr 2e-3. Starting at 1e-3 is deliberately cooler than that. + # Near-full-batch gradients (1,280 gauges) should also damp the collapse, since + # it looked like a high-variance small-batch effect — but that is a hypothesis, + # not an established result. If `n_at_floor` climbs past a few percent, the lr + # is still too hot. + learning_rate: + 1: 0.005 + 11: 0.001 + 21: 0.0005 + grad_clip_max_norm: 1.0 + +kan_head: + hidden_size: 21 + num_hidden_layers: 2 + # B-spline grid intervals (`num` in pykan). Matches DDR's + # ~/projects/ddr/config/merit_training_config.yaml::kan.grid. + grid: 50 + # B-spline order. Matches DDR's ::kan.k. (pykan's KANLayer default is 3, + # but DDR overrides to 2 for production.) + k: 2 + input_var_names: + - SoilGrids1km_clay + - aridity + - meanelevation + - meanP + - NDVI + - meanslope + - log10_uparea + - SoilGrids1km_sand + - ETPOT_Hargr + - Porosity + learnable_parameters: + - n + - q_spatial + - p_spatial + # Precip-conditioned mass-preserving daily->hourly disaggregation head. + # The within-day shape is driven by the hourly AORC precip window + # [d-1,d,d+1] (+ daily-Q taps + static attrs); the daily mean is conserved + # exactly. Loss stays the historical L1 (no experiment.loss block). + # Warm-started FROZEN from the standalone capacity pretrain (chunk_days=1); + # architecture fields must match that checkpoint or load_record fails. + # `enabled: false` makes this whole block inert — the loader uses flat + # repeat-24 (nearest) upsampling instead. Flip to true for the head-on + # arm of the ablation; nothing else needs to change. + disaggregation: + enabled: false + hidden_size: 16 + num_hidden_layers: 2 + grid: 20 + k: 3 + boundary_blend: 0.0 + chunk_days: 1 + pretrained_checkpoint: /home/tbindas/projects/ddrs/output/disagg_pretrain/capacity_chunk1.mpk + freeze: true + +# Routing-engine knobs — DDR's MERIT defaults (matches mock_config in tests). +params: + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] + attribute_minimums: + discharge: 1.0e-4 + slope: 1.0e-3 + velocity: 0.01 + depth: 0.01 + bottom_width: 0.01 + defaults: + p_spatial: 21.0 + log_space_parameters: + - p_spatial + # Corrected physics instead of DDR's formulation. With `false` the routing + # core uses: + # * exact trapezoidal celerity c = v·β, β = 5/3 − (4/3)·A·√(1+z²)/(T·P) + # instead of the wide-rectangular 5/3 limit (which is 22-27% high for the + # channels this code builds — κ = b/y ≈ 0.7-1.8, so β ≈ 1.31-1.36), and + # * Cunge-derived Muskingum X = clamp(0.5(1 − Q/(T·S·c·L)), 0, 0.5) instead + # of the constant 0.3, which was injecting a median 28× excess numerical + # diffusion. + # Negative-discharge counting before the S28 clamp is reported in BOTH modes. + # This breaks `compare_ddr_sandbox`'s ABSOLUTE MATCH by design — that example + # must be run with the default `ddr_match: true`. + # See `.claude/PHYSICS-CORRECTIONS.md`. + ddr_match: false + # Provably zero negative discharges. Clamps the Muskingum INPUTS (K, X) into + # the non-negative-coefficient window `2X <= Cr <= 2(1-X)`, Cr = dt/K: + # K <- max(K, dt*(1+d)/2) => Cr <= 2/(1+d) + # X <- min(X_cunge, (1-d)*Cr/2, (1-d)*(1-Cr/2)) + # with d = POSITIVITY_DELTA = 1e-2. Clamping K and X (not the coefficients) + # keeps `c1+c2+c3 = 1` exact, so mass is conserved; clamping c3 would not. + # With c1,c3 >= 0 and q_t,q' > 0, induction over the topological order of the + # forward substitution proves the whole solve is non-negative. + # + # Measured: 0 / 197,461,880 solves on 1,841 gauges (was 55,181, 0.0279%), + # replicated across two windows, two backends, and two trained heads. + # + # COST — this is not free. The X cap binds on 95.3% of reach-timesteps and + # median X falls 0.4976 -> 0.0794 (6.3x), so it largely REPLACES the Cunge X + # rather than shading it: numerical diffusion becomes stability-set instead of + # matched to hydraulic diffusivity. NSE/KGE impact is UNMEASURED, and this is + # the first training run to use the flag. Compare against a matched + # `enforce_positivity: false` control before trusting the result. + # + # Root-cause alternative: median Cr is 0.226, i.e. the typical MERIT reach is + # ~4.4x too LONG for an hourly step. Subdividing reaches ~4x brings Cr to ~1, + # where X_max -> 0.5 and the cap stops biting. See .claude/PHYSICS-CORRECTIONS.md. + # + # Requires ddr_match: false (rejected at load otherwise). + # OFF. It delivered provably zero negative solves (0 / 197,461,880), but the + # cure was worse than the disease: capping X collapsed the median from 0.4976 + # to 0.0794 on 95.3% of reach-timesteps, and because the cap is applied at + # RUNTIME to a learned celerity it made X ~ Cr ~ 1/n — handing the optimizer a + # lever it rode straight to the roughness floor. Same lr, flag on vs off: + # n_at_floor 35% -> 97.9%. Superseded by reach subdivision (a BUILD-TIME fix + # with no gradient path): see docs/superpowers/plans/2026-08-05-reach-subdivision.md + enforce_positivity: false + # cuSPARSE triangular solve. Pair with `ddrs run --backend cuda`; it is inert + # under `--backend cpu`. This is a DIFFERENT `Backward` impl from the host + # solve, sitting directly downstream of the ddr_match/enforce_positivity + # terms — gradient-checked on GPU by `tests/cuda_backward_parity.rs`. + sparse_solver: cuda + # MUST be false when ddr_match is false: the fused CUDA-graph kernel + # (cuda_graph/geometry_kernel.rs) hardcodes DDR's 5/3 celerity, so a captured + # graph would replay the DDR forward while the backward used the corrected + # chain rule — a silent gradient mismatch. Config load rejects the pair. + use_cuda_graphs: false + +# Test-mode overlay. When the binary is run with --mode testing, these keys +# replace the corresponding ones in `experiment:`. Absent keys inherit. +# +# IMPORTANT: batch_size SEMANTIC SHIFTS between modes: +# - experiment.batch_size (training) = number of GAUGES per mini-batch +# - testing.batch_size = number of DAYS per chunk +testing: + start_time: 1995/10/01 + end_time: 2010/09/30 + batch_size: 15 # DAYS, not gauges + rho: null # disabled in test mode diff --git a/config/sources/aorc_dhbv_distributed.yaml b/config/sources/aorc_dhbv_distributed.yaml new file mode 100644 index 0000000..47554ee --- /dev/null +++ b/config/sources/aorc_dhbv_distributed.yaml @@ -0,0 +1,14 @@ +data_sources: + # CONUS MERIT with dHBV2 distributed AORC2f forcing (precip-driven disaggregation) + attributes: /home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc + conus_adjacency: /home/tbindas/projects/ddr/data/merit_conus_adjacency.zarr + gages_adjacency: /home/tbindas/projects/ddr/data/merit_gages_conus_adjacency.zarr + # Daily dHBV2 distributed AORC2f-forced streamflow + streamflow: /mnt/ssd1/data/icechunk/daily_dhbv2_distributed_aorc2f_merit_unit_catchments.ic + # USGS daily observations + observations: /mnt/ssd1/data/icechunk/usgs_daily_observations + # 2000-area-balanced gage population + gages: /home/tbindas/projects/ddr/references/gage_info/gages_2000_area_balanced.csv + # Hourly AORC precip (zarr v3, CONUS) — drives the precip-conditioned + # mass-preserving disaggregation head + aorc_precip: /mnt/ssd1/data/aorc/merit_unit_catchments.zarr diff --git a/config/sources/conus-experimental.yaml b/config/sources/conus-experimental.yaml new file mode 100644 index 0000000..3d2a3c8 --- /dev/null +++ b/config/sources/conus-experimental.yaml @@ -0,0 +1,10 @@ +data_sources: + attributes: /home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc + conus_adjacency: /home/tbindas/projects/ddr/data/merit_conus_adjacency.zarr + gages_adjacency: /home/tbindas/projects/ddr/data/merit_gages_conus_adjacency.zarr + streamflow: /mnt/ssd1/data/icechunk/daily_dhbv2_distributed_aorc2f_merit_unit_catchments.ic + observations: /mnt/ssd1/data/icechunk/usgs_daily_observations + gages: /home/tbindas/projects/ddr/references/gage_info/gages_2000_area_balanced.csv + # Hourly AORC precip (zarr v3, CONUS) — drives the precip-conditioned + # mass-preserving disaggregation head below. + aorc_precip: /mnt/ssd1/data/aorc/merit_unit_catchments.zarr diff --git a/config/sources/conus-hourly.yaml b/config/sources/conus-hourly.yaml index 407004b..c670c99 100644 --- a/config/sources/conus-hourly.yaml +++ b/config/sources/conus-hourly.yaml @@ -1,8 +1,8 @@ data_sources: # CONUS MERIT inputs + hourly AORC precip (local workstation paths). # Same as `conus` but adds `aorc_precip` to drive the precip-conditioned - # mass-preserving daily->hourly disaggregation head - # (kan_head.disaggregation.use_precip: true). + # mass-preserving daily->hourly disaggregation head (enabled by the + # presence of a `kan_head.disaggregation:` block). attributes: /home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc conus_adjacency: /home/tbindas/projects/ddr/data/merit_conus_adjacency.zarr gages_adjacency: /home/tbindas/projects/ddr/data/merit_gages_conus_adjacency.zarr diff --git a/docs/2026-08-03-ddr-match-findings.md b/docs/2026-08-03-ddr-match-findings.md new file mode 100644 index 0000000..21955ca --- /dev/null +++ b/docs/2026-08-03-ddr-match-findings.md @@ -0,0 +1,199 @@ +# `ddr_match` physics corrections — findings + +**Date:** 2026-08-03 +**Branch:** `ddr-match-physics` +**Run:** `.ddrs/runs/2026-08-03T13-11-00Z-train-and-test` + +## Summary + +Three defects were found in the Muskingum-Cunge core, two were corrected behind +`params.ddr_match: bool` (default `true`, preserving DDR parity), and the third +was instrumented. The first run using the corrected physics produced **the +largest margin over the summed-Q' baseline this project has recorded** — ++0.0296 median NSE, three times the previous best — together with the most +physically defensible parameter field yet produced. + +The result is **not attributable to the physics alone**: four variables changed +simultaneously. See §Caveats. + +## The three defects + +### 1. Celerity used the wide-rectangular limit (corrected) + +`c = v · 5/3` is `dQ/dA` for a wide rectangular channel. The solver builds a +trapezoid (S7-S13). Correct kinematic celerity: + +``` +c = v · β, β = 5/3 − (4/3)·A·√(1+z²)/(T·P) +``` + +Derived from `c = dQ/dA = (dQ/dy)/T`, verified against finite-difference `dQ/dA` +to 1e-6. Limits: `β → 5/3` as `b/y → ∞`; `β → 4/3` as `b → 0` at fixed `z`. + +**β is not bounded below by 4/3** — it is non-monotone in `κ = b/y` and reaches +~1.07 for narrow sections. This code's channels have `κ ≈ 0.7-1.8`, giving +`β ≈ 1.30-1.36`, so the hardcoded `5/3` was **22-27% too high**. + +### 2. Muskingum X was a constant, not Cunge-derived (corrected) + +`X ≡ 0.3` severs the link between the scheme's numerical diffusion and the +channel's physical hydraulic diffusivity — the defining feature of +Muskingum-**Cunge**. Corrected: + +``` +X = clamp( 0.5·(1 − Q/(B·S·c·L)), 0, 0.5 ) +``` + +making `D_num = c·L·(0.5−X)` equal `D_phys = Q/(2·B·S)`. On CONUS the Cunge +value is ≈0.49 almost everywhere; the constant 0.3 injected a **median 28× +excess numerical diffusion**, over-attenuating 3-6 h waves by 13-34%. + +### 3. Negative discharge was clamped without measurement (instrumented) + +`q_next = x_sol.clamp_min(1e-4)` at S28 silently rewrote negative solve output +to `+1e-4`, **creating mass**, hiding Courant instability, and zeroing gradients +where saturated. Now counted before the clamp. + +**First measurement:** + +| | mb1 | mb2 | mb3 | mb4 | +|---|---|---|---|---| +| `ddr_match: true` | 0.004% | 0.006% | 0.007% | 0.012% | +| `ddr_match: false` | 0.027% | 0.082% | 0.025% | 0.048% | + +This **corrects the audit's framing**: ~70% of reaches carry a negative +Muskingum coefficient, but only ~0.01-0.05% of reach-timesteps produce negative +discharge. Negative coefficients cause an initial dip or recession oscillation — +artifacts, not divergence. + +## Courant sub-stepping: attempted, abandoned, documented + +Cunge `X ≈ 0.49` narrows the non-negative window `2X ≤ Cr ≤ 2(1−X)` from +`[0.6, 1.4]` to `[0.98, 1.02]`. Sub-stepping was planned to bring `Cr` in range. +**It cannot work:** + +``` +K spans p5 = 425 s to p95 = 18,551 s — a 44x range +best GLOBAL n_sub: 1.4% of reaches in window +ideal PER-REACH n_sub: 6.3% <- even the unattainable ideal fails +``` + +Shrinking Δt globally slides every reach's `Cr` down together; it cannot compress +the spread, and integer sub-stepping quantizes too coarsely. **The correct fix is +variable Δx — subdividing reaches so `Δx ≈ c·Δt`, which is what HEC-HMS does.** +That changes adjacency topology, the CSR pattern, and per-reach parameter fields. +Out of scope; recorded here so it is not re-attempted. + +This also reframes `X = 0.3`: window width is `2 − 4X`, so smaller X gives a +*wider* stability window. The constant 0.3 (width 0.8, 61.7% of reaches +admissible) is the most numerically forgiving choice available, and may have been +a deliberate stability trade rather than an oversight. + +## Result + +Run: 1,841 area-balanced gauges, Adam, no gradient accumulation (280 optimizer +steps in 10 epochs), `nse-batch`, disagg head OFF, `ddr_match: false`. +1.52 h train + 2.29 h eval on CPU. + +``` +metric baseline ddrs delta +nse 0.6440 0.6736 +0.0296 +kge 0.6956 0.6963 +0.0007 +corr 0.8503 0.8512 +0.0010 +bias 1.3969 0.4999 -0.8970 +rmse 11.4719 11.1571 -0.3149 +fhv 5.3840 -6.0287 -11.4126 +flv 52.7742 38.8074 -13.9668 +``` + +Observations byte-identical between the two series (max diff 0.000e+00 over +10,026,921 finite cells), so this is not a join artifact. + +### Area-stratified NSE + +| area km² | n | baseline | ddrs | Δ | +|---|---|---|---|---| +| <1k | 841 | **0.720** | 0.686 | −0.035 | +| 1k–5k | 418 | 0.662 | **0.711** | +0.048 | +| 5k–10k | 244 | 0.490 | **0.626** | **+0.135** | +| ≥10k | 338 | 0.352 | **0.545** | **+0.193** | + +The rebalanced gauge set puts **582 of 1,841 (32%)** at or above 5,000 km², +versus 295 of 2,365 (12.5%) in `gages_3000`. That is why the pooled median moved +when it never had before — the metric finally samples basins where routing can +physically act. The small-basin loss also shrank from ~−0.10 in earlier runs to +−0.035. + +### Learned parameter field (346,321 CONUS reaches) + +| | median n | @floor | ρ(n, log10_uparea) | +|---|---|---|---| +| **this run** | **0.0467** | 6.56% | **+0.323** | +| 50ep `gages_3000` | 0.0402 | 0.73% | +0.076 | +| lr 1e-2 `gages_3000` | 0.0177 | **47.28%** | +0.205 | + +**76.3% of reaches fall inside the NLCD natural-channel band 0.025-0.15.** All +three learnable parameters show their strongest-ever scale dependence (ρ ≈ +0.33 +for `n`, `q_spatial`, `p_spatial`). + +`n_mean` converged cleanly: per-epoch 0.1262 → 0.1074 → 0.0768 → 0.0612 → 0.0546 +→ 0.0528 → 0.0500 → 0.0499 → 0.0485, with the LR halvings at epochs 5 and 8 +arresting the descent as designed. + +## Caveats + +**Four variables changed at once** versus every prior run: corrected celerity, +Cunge X, disagg head off, and the rebalanced gauge set. The +0.0296 is **not +attributable to the physics on this evidence**. The matched `ddr_match: true` +control on the same gauge set is required and has not been run. + +**The baseline differs** (0.6440 on this 1,841-gauge population vs 0.6754 on +`gages_3000`). None of these absolutes compare to earlier runs. + +**FHV moved +5.4 → −6.0** — peaks now under-predicted ~6%. Mechanistically +consistent: corrected celerity is ~20% lower, so `K = L/c` is longer and the +router attenuates more. + +**The floor fraction rose to 6.56%** from 0.73% in the 50-epoch run. Far below +the 47.3% collapse, but the pinned reaches should be checked for concentration in +small headwaters (identity-routing pressure) versus scatter (a training defect). + +## Not changed (found, deferred) + +- **`attribute_minimums.slope: 1e-3`** clamps 33.2% of reaches (4.03% have slope + exactly 0). This is an **exact invariance** — scaling `n` by `√f` restores + depth, `R` and velocity identically, because `n` and `√S` enter depth only as + the ratio `n/√S`. Undoing it puts 94.2% of implied physical `n` inside the NLCD + band. Requires a retrain. +- **`leakance.rs:35-36`** uses `(p·d)^q` instead of `p·d^q` — dimensionally + incoherent, inherited from DDR. Leakance is closed/NO-GO, so informational. +- **`mmc.rs::calculate_muskingum_coefficients`** takes a parameter named + `velocity` that is actually the celerity (dead in production). + +## Verification + +`compare_ddr_sandbox` ABSOLUTE MATCH (1.53e-5 m³/s) · `cunge_x` 11/11 · +`celerity_beta` 9/9 · `sp8_gradcheck` 5/5 · `sparse_gradcheck` 1/1 · `mmc` 13/13 +· `leakance_gradcheck` 16/16 · `zeta_accum` 8/8 · `cargo test --lib` 260 passed. + +**Both new backward branches were verified falsifiable**, not vacuous: they FAIL +with their terms disabled (celerity rel 1.3e-1 vs 1.44e-3 passing; X rel +8.9e-2…5.7e-1 vs 2.84e-3 passing). The Cunge-X gradcheck fixture was raised from +1000 m to 5000 m reaches because at 1000 m `W ≈ 1.6` saturates the clamp on every +reach, masking `gX` to zero and letting all four tests pass with the terms +deleted. + +**Config guard:** `ddr_match: false` + `use_cuda_graphs: true` is rejected at load +(`validate_ddr_match`), because `cuda_graph/geometry_kernel.rs:296` hardcodes +DDR's 5/3 and would replay a DDR forward against a corrected backward. + +## Next + +1. **Matched control** — same config, `ddr_match: true`. The only way to attribute + the +0.0296. +2. **Disagg ablation** — `enabled: true` vs the `false` used here. Never run in + any configuration. +3. **Slope floor** — the third correction, needs its own retrain. + +Artifacts: `.ddrs/runs/2026-08-03T13-11-00Z-train-and-test/plots/` (23 PNGs, +three notebooks, `_make_notebooks.py`). Diagram: `.claude/PHYSICS-CORRECTIONS.md`. diff --git a/docs/2026-08-06-tau-sweep-pilot-findings.md b/docs/2026-08-06-tau-sweep-pilot-findings.md new file mode 100644 index 0000000..1b09649 --- /dev/null +++ b/docs/2026-08-06-tau-sweep-pilot-findings.md @@ -0,0 +1,558 @@ +# Per-gauge tau sweep — WY1996 pilot findings + +- **Spec:** `docs/superpowers/specs/2026-08-05-per-gauge-tau-sweep-design.md` +- **Script:** `scripts/tau_sweep.py` (offline, on the `DDRS_HOURLY_DUMP` from one eval run) +- **Checkpoint:** `2026-08-05T04-58-58Z-conus-experimental-train-and-test/checkpoints/epoch_30_mb_1` (epoch 30, area-balanced 1,841 gauges) +- **Binary provenance:** `eval` built from commit `76a4020` (dump diagnostic committed with the spec); run log `output/tau_sweep/eval.log`. +- **Window:** eval full 1995/10/01–2010/09/30; pilot analysis restricted to WY1996 (1995/10/01–1996/09/30, 365 days). **Pilot numbers are single-year and (for per-gauge best) in-sample; they are method-validation and headroom estimates, not conclusions.** + +**One-line verdict: H-TAU SUPPORTED at pilot strength — the shipped `tau = 3` is +grossly mis-set; the NSE(tau) optimum sits at tau ≈ 14–19 (daily pooling bin +starting 03:00–08:00 UTC next day, i.e. CONUS local midnight), and a single +global tau ≈ 16–19 already ties or beats the summed-Q' baseline in the +small-basin bins that motivated the experiment.** + +## 1. Pre-registered hypotheses and gates + +- **H-TAU:** the small-basin NSE deficit is substantially a pooling-phase + (timing) error; per-gauge tau recovers most of the gap. +- **Method gate (pilot):** (1) offline tau=3 reconstruction matches the eval's + own daily zarr < 1e-3 rel; (2) recomputed median NSE matches < 0.001; + (3) ≥95% of gauges get full NSE(tau) curves. +- **Decision gate (Phase 2, user-selected, not yet evaluated):** per-gauge best + tau lifts <1,000 km² median NSE to ≥ the baseline, full window, with an + out-of-sample selection protocol (protocol deliberately open). + +## 2. Method + +One eval of the epoch-30 checkpoint over the full eval window with +`DDRS_HOURLY_DUMP` writing the pre-trim hourly series (1841 × 131,496 f32, +0.97 GB). Offline, for each tau ∈ {0..23}: daily prediction bin *i* = block +mean of hours `[13+tau+24i, 13+tau+24(i+1))`, scored against obs day *i+1* +(reproduces `tau_trim_and_downsample` exactly at tau = 3; the trimmed window is +divisible by 24 so block mean = area pool). Per-gauge NSE on WY1996 days, +≥100 valid days required. Baseline = the run's own `baseline/` arrays (same +1,841-gauge population), same days. + +## 3. Results + +Method gate: **all three PASS** (max rel 7.06e-4; medians 0.6206 vs 0.6205; +98.7% coverage). + +Median NSE, WY1996, by drainage area: + +| bin (km²) | n | tau=3 (shipped) | best single global tau | per-gauge best (in-sample ceiling) | Σ Q' baseline | +|---|---|---|---|---|---| +| 0–1,000 | 841 | 0.553 | **0.677** (tau=16) | 0.699 | 0.674 | +| 1,000–5,000 | 418 | 0.592 | **0.702** (tau=23) | 0.722 | 0.660 | +| 5,000–10,000 | 244 | 0.544 | **0.668** (tau=22) | 0.671 | 0.512 | +| 10,000–30,000 | 266 | 0.479 | **0.548** (tau=14) | 0.566 | 0.379 | +| 30,000–50,000 | 72 | −0.021 | −0.021 (flat) | −0.019 | −0.334 | +| ALL | 1,841 | 0.546 | **0.660** (tau=19) | — | 0.605 | + +Key observations: + +1. **The NSE(tau) curve rises monotonically from tau=0 and plateaus at + tau ≈ 14–19 in every bin below 30,000 km²** (`nse_vs_tau_by_area.png`). + Window start hour is `13+tau`, so the plateau is 03:00–08:00 UTC (next + day) — precisely the CONUS local-midnight band (ET midnight = 05 UTC, + PT midnight = 08 UTC). The spec's arithmetic concern is confirmed: shipped + tau=3 pools over a bin starting 16:00 UTC, ~11–16 h out of phase with the + local observation day. +2. **A single global tau fixes most of it.** +0.114 median NSE overall with + zero per-gauge freedom (no selection bias applies to this number, though it + is still single-year). +3. **Per-gauge freedom adds only ~+0.02 over the global optimum** — and the + longitude fingerprint is weak (Spearman −0.129 on the 1,353 gauges improved + > 0.01). The broad plateau swallows the 3-hour timezone spread; per-gauge + tau is second-order compared to just fixing the global phase. +4. **Best-tau histogram is edge-heavy** (661 gauges at tau=23, 232 at tau=0), + so the true optimum for a subpopulation lies outside {0..23} under the fixed + day mapping — the sweep range needs a ±1-day mapping extension in Phase 2. +5. Above 30,000 km² the curve is flat: multi-day response integrates out the + phase, and both model and baseline are poor there (regulation, n=72). + +## 4. Conclusions + +- The small-basin deficit that motivated this experiment is at least + substantially a **pooling-phase error in `tau`**, not a routing-physics + failure: at global tau=16 the <1,000 km² bin ties the baseline (0.677 vs + 0.674) in the pilot, and per-gauge best exceeds it (0.699). +- **Training implication (potentially the bigger one):** training also scores + through `tau_trim_and_downsample` at tau=3, so every gradient the KAN head + has ever received was computed against misaligned observations. The head may + be learning spurious lag/attenuation to compensate — a candidate mechanism + for the over-attenuation signature. A retrain at corrected tau is the test. +- Caveats: single water year; per-gauge numbers in-sample; global-tau numbers + free of selection but still one year. + +## 5. Next steps (tomorrow's tuning session) + +1. Extend the sweep with a ±1-day mapping shift to un-pin the tau=0/23 edges. +2. Full-window sweep + split-sample protocol (select on 1995–2003, score on + 2003–2010) to evaluate the Phase 2 gate honestly. +3. Decide the fix tier: global tau ≈ 16–19 (needs `tau_trim_and_downsample` to + accept tau > 11 — the current slice arithmetic caps tau at 11), per-timezone, + or per-gauge local-midnight window. +4. Retrain at corrected tau and re-run the area-balanced eval — tests the + training-misalignment mechanism. + +## 5a. Adversarial review corrections (2026-08-06, Fable subagent — verdict: sound-with-corrections) + +The arithmetic and the +0.114 selection-free gain reproduce exactly. Three +interpretive claims are corrected in place: + +1. **§3 obs 1 "precisely the CONUS local-midnight band" — RETRACTED as + over-claimed.** This run had `disaggregation.enabled: false` (verified in + the config snapshot), so the hourly signal was flat repeat-24: UTC-day means + capture a median **97.6%** of hourly variance, and the pooled series at any + tau is reproduced by a two-point blend of adjacent UTC-day means with median + R² 0.994–0.997. The sweep therefore measures a **day pairing plus blend + weight (~half-day resolution), not sub-daily phase**. Hour-scale readings + (tau=16 "Eastern midnight" vs tau=19 "Pacific midnight") are below the + method's effective resolution. Corrected headline: *the shipped + (tau=3, obs day i+1) mapping pools a window ~half a local day early; any + tau in 14–19 fixes the day-boundary blend.* +2. **Phase semantics / sign convention.** Window-start offset relative to the + scored day's UTC midnight is **(tau − 11) h**: tau=3 → −8 h (13/24 of the + pooled mass from the wrong local day); tau=16 → +5 h (= ET midnight); + tau=19 → +8 h (= PT midnight); tau=23 → +12 h. **Larger optimal tau ⇒ the + model hydrograph is LATE relative to observations.** +3. **§3 obs 3 (longitude) — the timezone fingerprint is absent-to-contradicted, + not merely "weak".** The −0.129 Spearman was computed on a 48%-edge-pinned + set; with censored gauges excluded the sign FLIPS to +0.174. The **robust + covariate is drainage area**: Spearman(best_tau, log10 area) = +0.164 + (uncensored, p=1e-5) to +0.341 (censored incl.), surviving partialling on + longitude (+0.208). Larger basins prefer later windows — an accumulated-lag + / travel-time signature. So "not a routing-physics failure" (§4) is too + strong: phase error is the dominant term, not the only term. +4. **Censoring is asymmetric and real:** of 661 tau=23 pins, 596 are genuinely + improved (optima beyond +12 h, needing the day-(i+2) mapping); of 232 tau=0 + pins only 56 are improved (flat-curve noise). The ±1-day extension is + necessary and predominantly on the late side. Curve sharpness tracks + sub-daily structure (Spearman +0.701), as the blend mechanism predicts. +5. Minor script defect: no-curve gauges are back-filled with `best_tau = + tau_shipped` in the CSV, slightly contaminating the histogram/correlations. + +Phase-2 design implication: select a **day-mapping × blend-weight** per gauge +(split-sample), do NOT commit to an hour-precision "local-midnight" fix tier on +this evidence, freeze the selection protocol before any retrain (the head has +trained 30 epochs against a ~half-day-early target and has plausibly learned +compensating lag), and re-run the sweep on a disagg-ON or hourly-native run to +test whether any hour-scale tau signal exists at all. + +## 5b. External mechanistic prior (2026-08-06, `/tmp/handoff-aorc-usgs-recording-times.md`) + +Independent of the sweep, the recording conventions predict the offset a priori: + +- **USGS daily values are local-STANDARD-time midnight-to-midnight, year-round + (no DST)** — authoritative per + https://waterdata.usgs.gov/statistics-documentation/. +- **AORC forcing is UTC-hourly**, and the AORC-driven Q' stores (incl. this + run's `daily_dhbv2_distributed_aorc2f`) define "day t" as UTC 24-hour blocks. +- Predicted misalignment: +5 h (EST) to +8 h (PST). In tau units + (window offset = tau − 11 h) that predicts **optimal tau ∈ [16, 19]** — + exactly the measured plateau. What the review demoted to + "consistent-with" now has a documented mechanism. +- Pipeline verification (this session): `src/data/` contains **no timezone + logic anywhere**; all stores are indexed positionally on their native axes. + `params.tau` is the only alignment knob in the system. +- Because USGS uses LST year-round, the correct per-gauge correction is a + **fixed deterministic offset from the gauge's standard-time zone** — no DST + seasonality to model. +- Open tension with §5a: the mechanism predicts western gauges (more negative + longitude) prefer LARGER tau, i.e. a negative lng correlation; the + uncensored pilot showed +0.174. Candidate explanations: half-day sweep + resolution blurring a 3-h span, the area confound, or geographic structure + in the censored tau=23 tail. Unresolved; the interpolation arms (sharper + curves) are the discriminating instrument. +- Open question inherited from the handoff: whether every Q' store shares the + UTC-day convention (dHBV2-UH, daily-LSTM, hourly-LSTM vs the AORC2F pair). + If they differ, tau is per-STORE as well as per-gauge, and cross-store + parameter comparisons (the AGU H069 framing) inherit the bias. Check each + store's CF axis + forcing provenance before the next cross-store run. + +## 5c. Interpolation arms (2026-08-06): nearest vs linear vs quadratic, gages_3000 + +Three evals of the same epoch-30 checkpoint on the standard **2,365-gauge** +population (gages_3000 after filters), full window, differing ONLY in +`DDRS_QPRIME_INTERP` (commit `e4fb66d`); WY1996 sweep per arm, all method +gates PASS (99.9% curve coverage). Driver: `scripts/run_tau_interp_arms.sh`; +overlay plot `output/tau_sweep/interp_arms_nse_vs_tau.png`. + +**Verdict: the tau mis-set is confirmed on the benchmark population and is +NOT an artifact of step-function upsampling — smoother q' input neither +sharpens nor shifts the optimum. Interpolation is not the fix; the day +mapping is.** + +| arm | full-window median NSE @ tau=3 | WY1996 median @ tau=3 | WY1996 argmax | WY1996 max | curve range (median) | +|---|---|---|---|---|---| +| nearest | 0.6426 | 0.578 | tau=20 | 0.6997 | 0.104 | +| linear | 0.6496 | 0.589 | tau=19 | 0.6983 | 0.094 | +| quadratic | 0.6347 | 0.573 | tau=18 | 0.6943 | 0.104 | + +1. **Interpolation buys almost nothing, and nothing at the optimum.** Linear + gains +0.011 at the mis-set tau=3 (smearing partially absorbs the + misalignment) but at each arm's own optimum the three arms converge within + 0.005, ordered nearest ≥ linear ≥ quadratic — consistent with the predicted + peak attenuation of the smoothing kernels. Quadratic is strictly worse than + nearest at tau=3. +2. **Curves do not sharpen** (linear is slightly FLATTER), so the half-day + resolution limit of §5a is a property of the daily-information content, not + of the step discontinuities. Sub-daily structure cannot be conjured by + interpolation; only a disagg-ON or hourly-native store can supply it. +3. **The optimum sits at tau=18–20 on this population, at/beyond the PT edge + of the LST band [16,19], with 707–728 gauges (~30%) still pinned at + tau=23** (real optima beyond +12 h; only ~60 at tau=0). The timezone + convention alone under-predicts the shift. +4. **Correlations replicate across all arms:** best_tau vs log10(area) + +0.16 to +0.19 (uncensored), vs longitude +0.13 to +0.19 (uncensored — + still the WRONG sign for the timezone mechanism, in every arm). +5. **Small basins (<1,000 km², n=1,267): single global tau=19 scores 0.674 vs + baseline 0.645** (WY1996) — the "beat the baseline" bar is cleared on this + population too, again with zero per-gauge freedom. +6. **Emerging synthesis:** optimal shift ≈ (LST-vs-UTC convention offset, + +5..8 h) + (an area-growing lag term). The area correlation, the beyond-band + optimum, and the late-side censored tail all point at extra model lag on top + of the convention offset — the leading candidate being **double routing** + (MC travel time stacked on whatever routing/UH the Q' store already embeds + to place flow at its outlet; DDR's own tau docstring says "handle double + routing and timezone differences"). Discriminating test: repeat the sweep on + a UH-free vs UH-embedded store pair. + +## 5d. Sample calculation: tau_g = 11 + tz(gauge) + c·A^b (2026-08-06) + +Fitted on the nearest-arm WY1996 curves (2,365 gauges), tz from longitude +(midpoints of standard meridians → 5/6/7/8 h), objective = median NSE over the +per-gauge integer tau_g, grid over (c, b). In-sample (2 free params). + +| scheme | median NSE | note | +|---|---|---| +| tau=3 (shipped) | 0.5780 | | +| tau=18 / 19 / 20 constant | 0.6990 / 0.6992 / 0.6997 | | +| tz only (tau = 11+tz) | 0.6966 | WORSE than constant 19 | +| **formula b=0.30, c=0.28** | **0.7018** | joint grid (b=0.40, c=0.12) ties at 0.7019 — b unidentified | +| per-gauge best (ceiling) | 0.7228 | in-sample selection | + +Per-bin: formula ties constants below 5,000 km² (0.6728 vs 0.6736 at tau=19), +gains +0.005 in 10,000–30,000 km² (0.7257 vs 0.7203). Fitted lag term: +1.1 h @ 100 km², 2.2 h @ 1,000, 4.4 h @ 10,000, 6.2 h @ 30,000 — magnitudes +consistent with Allen et al. (2018) celerity-based travel times. Median +predicted tau_g: 19 (p10 18, p90 22), 1% clip at 23. + +Diagnostics: the area term DOES absorb the area signal (residual-vs-log-area +Spearman drops +0.16 → −0.10), but residual-vs-tz is **−0.452** — per-gauge +optima do not track the timezone term at this resolution (echoes §5c's +wrong-sign longitude), and tz-only underperforms a flat constant. Vs constant +tau=19 the formula moves 76% of gauges and improves 922 vs worsens 872 +(median Δ +0.0001) — a coin flip per gauge, small net win from the mid-size +bins. + +**Reading:** the half-day blend resolution (§5a/§5c) leaves hour-scale +refinements below the instrument's discrimination; a constant tau ≈ 19–20 +captures essentially all recoverable skill on this population (0.6997 vs +0.7018 formula vs 0.7228 unreachable ceiling). The formula is physically +defensible and never hurts materially — a fine choice for the retrain — but +the decision between it and a constant should be made split-sample in +Phase 2, not on these in-sample numbers. + +## 5e. Cross-source arms (2026-08-07): is the lag a property of the store? + +Same epoch-30 checkpoint, same 2,365-gauge network, only +`data_sources.streamflow` swapped (`config/experiments/tau_src_*.yaml`, +`scripts/run_tau_source_arms.sh`; this set ran on cuda, before the CPU +policy). All three method gates PASS on every arm (recon match, NSE match, +99.9% curve coverage). WY1996 sweep, median over 2,365 gauges: + +| store | resolution | full-window med NSE @ tau=3 | best const tau | med @ best | per-gauge best-tau median | % optima at tau=0 / 23 | +|---|---|---|---|---|---|---| +| aorc2f distributed (ref) | Daily | 0.6426 | 20 | 0.6997 | 18 | 11 / 33 | +| UH retrospective | Daily | 0.6375 | 21 | 0.7011 | 18 | 12 / 34 | +| daily LSTM | Daily | 0.5562 | 17 | 0.6135 | 18 | 15 / 32 | +| hourly LSTM (native) | Hourly | 0.5316 | 19 | 0.5515 | 19 | 15 / 34 | +| **aorc2f lumped** | Daily | 0.5103 | **3** | 0.4588* | **1** | **47 / 6** | + +\* pilot-window median at its own optimum; levels are not comparable across +stores (the head was trained on aorc2f distributed only), but optimum +LOCATIONS and curve shapes are. + +**Four of five stores replicate the tau 17–21 optimum.** UH retrospective, +daily LSTM, and hourly LSTM all peak within 3 h of the reference despite +being entirely different models of runoff generation. The per-gauge best-tau +median is 18–19 on all four, and the censored-tail fractions (pile-ups at +tau=0 and tau=23) are near-identical. Whatever produces the lag, it is not a +quirk of the aorc2f-distributed store: it is shared by every store that uses +the standard daily convention, consistent with the UTC-vs-local-standard-time +mismatch as the dominant term. The replication cannot, however, apportion the +shared lag between the day convention and the common MC routing (the +double-routing candidate): the routing head, parameters, and network are +common-mode across all five arms, and the optima sit at or beyond the +predicted LST band [16, 19]. The residual beyond-band lag remains +unattributed (see §5f for the discriminating measurement). + +**The aorc2f lumped store is the outlier and the exception that probes the +rule.** Its median curve is flat over tau 0–3 and then falls monotonically; +47% of per-gauge optima sit at the tau=0 edge (left-censored). Extended +sweep (§5f): the median curve peaks at tau=+3 with the per-gauge median at +−1 (50.1% below 0). Its shipped-tau full-window median (0.5103) is already +near its own optimum. Obs-free cross-correlation of routed hydrographs +(§5f) shows the lumped arm LEADS the reference by ~23 h (median), so the +store's data are aligned about one day differently from every other store. +The CF-day-convention candidate is REFUTED: both aorc2f stores' icechunk +time metadata are byte-identical (`days since 1980-01-01`, +proleptic_gregorian, 14,976 steps). The shift is in the data the lumped +pipeline wrote (day-indexing off-by-one or different event-day assignment), +and the ~6 h residual between the waveform shift (~23 h) and the NSE-optimum +shift (~17 h) implies the lumped timing content also differs beyond a pure +relabeling. Do not use this store in timing-sensitive comparisons until the +pipeline-side indexing is resolved. + +**Hourly-native arm: no sharpening, and no timezone fingerprint either.** +The hourly LSTM curve is the flattest of the four lagged stores (gain from +tau=3 to optimum +0.073 vs +0.122 for the reference), not sharper, so real +sub-daily structure did not turn the sweep into an hour-resolution +instrument. The flatness is not a skill-floor artifact: at matched skill +levels the hourly arm is still flatter (§5f). The improved-subset longitude +correlation (−0.220) initially looked like the timezone-predicted sign, but +§5f shows it is a censoring artifact: on uncensored interior optima the +sign flips to +0.075 (wrong sign, weakest of the four lagged arms), and the +lumped arm, where the timezone mechanism has no standing, shows −0.197 on +its own improved subset. The timezone fingerprint remains +absent-to-contradicted even with native sub-daily structure, an informative +negative result. + +Plot: `output/tau_sweep/cross_source_nse_vs_tau.png`. Raw per-arm outputs in +`output/tau_sweep/src_{uh_retro,daily_lstm,hourly_lstm,aorc2f_lumped}/`. + +## 5f. Adversarial review of §5e (2026-08-07, Fable subagent — verdict: sound-with-corrections) + +Independent read-only review; all §5e corrections above were folded in from +it. Verdicts: claim "4/5 replicate ⇒ shared convention" SOUND-WITH- +CORRECTIONS (replication real, attribution overreached); lumped-outlier +claim SOUND-WITH-CORRECTIONS (strengthened, CF candidate refuted); "no +sharpening" SOUND; the longitude half UNSUPPORTED (cut); levels-not- +comparable SOUND. New measurements it contributed: + +- **Obs-free cross-correlation of routed hydrographs** (500-gauge sample, + WY1996, lags ±48 h; observations never enter): median lag vs reference is + UH retro +0 h (79% within 6 h), daily LSTM −1 h, hourly LSTM −2 h, + **lumped −23 h** (IQR −31 to −19, 84% at or below −12 h). This refutes the + checkpoint-mismatch explanation for the lumped outlier twice over: the two + LSTM arms are at least as mismatched to the trained head yet show zero + shift, and the −23 h appears with no observations involved. +- **Extended sweep tau −13..47** (recomputed from raw dumps): the reference + constant-tau curve peaks interior at 20 (0.6997, declining to 0.6853 at 24 + and 0.6373 at 30), so the constant-tau conclusion does not depend on the + missing ±1-day extension. Per gauge, though, 32.6% of reference optima are + genuinely beyond tau=23 and 11.2% below 0 (median 18, p10 −5, p90 40) — + the per-gauge tail structure is real, not an artifact of the 0..23 window. +- **Censoring asymmetry supports the lumped reading:** the reference's + tau=0 pile is 78% flat-curve noise (only 22% improved) while the lumped + tau=0 pile is 64% genuine improvement — the left pile is signal for the + lumped arm, unlike its mirror image in the reference. +- **Flatness is not a floor effect:** within-arm Spearman(curve range, curve + max) ≈ 0 to +0.13 in every arm; level-matched gauges (|Δmax NSE| < 0.05, + n=409) still leave the hourly arm flatter (median range 0.074 vs 0.091). +- **Mechanical comparability PASS:** identical gauge ID vectors, identical + `nse_baseline`, identical obs arrays, has_curve 2362/2365 with the same 3 + NaN gauges in every arm. The known no-curve backfill defect + (`scripts/tau_sweep.py:143`) touches only those 3 gauges here. +- Interior-optima area correlation is unstable across source arms (+0.089 + reference, −0.118 hourly-native) — weaker than the §5c interp-arm numbers. + +**Single most informative next measurement (proposed):** sweep tau on the +routing-free summed upstream q' (the baseline construction, repeat-24, same +gauges and obs, fully offline from existing stores). All five arms share the +MC routing, so its lag contribution is invisible to the cross-source design. +If the no-routing optimum also sits at 19–21, the entire lag is the +store/obs day convention and the double-routing candidate dies; if it sits +near 14–16, the gap directly quantifies the MC network's added travel time. + +## 5g. The discriminator: tau sweep on routing-free summed q' (2026-08-07) + +The §5f-proposed measurement, run entirely offline from existing data: the +baseline's summed upstream daily q' (cache +`.ddrs/runs/2026-07-30T00-24-24Z-train-and-test/baseline`, full coverage of +all 2,365 eval gauges) repeat-24'd into a synthetic hourly dump in the SAME +phase as the routed arms' disaggregation input, then swept with the +identical `scripts/tau_sweep.py` (all gates PASS) plus an extended sweep +tau −13..47. Construction validated: sweep tau=11 is algebraically the +standard day-aligned baseline scoring, and its per-gauge NSE reproduces the +cached baseline NSE (median |diff| 0.0003). WY1996, medians: + +| area bin (km²) | summed-q' opt tau | routed opt tau (uncensored) | summed early by | routed late by | routing-added delay | +|---|---|---|---|---|---| +| 0–1,000 | 9 | 19 | 2 h | 8 h | 10 h | +| 1,000–5,000 | 3 | 19 | 8 h | 8 h | 16 h | +| 5,000–10,000 | −3 | 25 | 14 h | 14 h | 28 h | +| 10,000–30,000 | −8 | 30 | 19 h | 19 h | 38 h | +| global median | 6 (0.6731) | 20 (0.6997) | 5 h | 9 h | 14 h | + +("early/late by" = |tau_opt − 11|; the 30,000–50,000 bin, n=12, is too +noisy to read. Per-gauge uncensored summed-q' best tau: median 6, p10 −10, +p90 20, only 6.5% at the −13 edge.) + +**Neither §5f-anticipated outcome occurred, and the measurement is the more +decisive for it.** The no-routing optimum is not 19–21 (all convention) and +not 14–16 (convention plus routing residual): it is **6**, below the +day-aligned point. Three conclusions follow: + +1. **The UTC-vs-LST convention story (§5b) is REFUTED as the dominant + term.** A timing-correct hydrograph scored against LST-labeled daily obs + should show optimum tau ≈ 16–19 even without routing. The smallest + basins, where travel time is minimal, sit at tau=9, and the area trend + extrapolates to ≈ 10–11 at zero area: the convention offset is ≈ 0–2 h, + not 5–8. This finally explains why the longitude fingerprint failed in + every arm and subset (§5a, §5c, §5f): there was no timezone signal to + find. Whatever conventions the stores and obs use, they net out to + near-UTC-day alignment. +2. **The summed q' leads the gauges by an area-growing travel time** (2 h + at <1,000 km² to ~19 h at 10,000–30,000 km²), exactly the unmodeled + network travel time the routing exists to supply. Corollary: day-aligned + scoring understates the baseline's skill in large basins (5,000–10,000: + 0.636 day-aligned → 0.755 at its optimum). Baseline comparisons at + large basins should keep this in mind. +3. **The MC routing over-delays by almost exactly 2× the required travel + time.** Bin by bin (≥1,000 km²), routed lateness equals summed-q' + earliness: the routing added twice the delay the gap required. This is + the "double routing" of DDR's own tau docstring, now measured: the q' + stores already route runoff to the unit-catchment outlet (dHBV2's UH), + and the MC network then adds what amounts to the full travel time again. + The tau 18–20 optimum of every routed arm is COMPENSATION for this + over-delay, not a data-convention fix. + +**Routing still earns its keep once both sides are timing-corrected:** at +per-bin optima the routed model beats the summed q' everywhere that +matters: +0.022 (<1,000 km²), +0.013, +0.019, +0.051 (10,000–30,000 km²). +The value added is real; it is currently masked at shipped tau=3 and +partially masked at any constant tau by the over-delay. + +**Reframing for the fix tier:** a constant tau ≈ 19–20 remains the correct +empirical patch for the current checkpoint, but the root cause is now in +the routing timing (double-carried travel time, and possibly the slow +trained celerity: median Manning's n 0.130 vs reference 0.05), not in the +data pipeline. A retrain at corrected tau tests the patch; the deeper fix +candidates (injection geometry, celerity prior) are a separate experiment. + +Artifacts: `output/tau_sweep/summed_qprime/` (synthetic dump, sweep +outputs, `extended_curves.npy`), +`output/tau_sweep/g3000_nearest/extended_curves.npy`, +`output/tau_sweep/summed_qprime_vs_routed_tau.png`. + +## 5h. The trained gamma-UH parameters: double routing confirmed at the parameter level (2026-08-07) + +`scripts/dump_gamma_uh_params.py` (runs under the water_loss venv) pulls +routa/routb from the Ann head of the checkpoint that generated the +distributed aorc2f store (CONUS2717_AORC2F_v3_gradaccum ep100 — the store +was re-exported 2026-07-29 with each divide's runoff routed through its own +learned gamma UH; see water_loss +`docs/findings/2026-07-29-uh-missing-from-export.md`). Scaling chain +verified against waterlossv18_1.py:57,179-183 and UH_gamma's floors: +a_eff = 2.9·r0 + 0.1 ∈ [0.1, 3.0], theta_eff = 6.5·r1 + 0.5 ∈ [0.5, 7.0], +kernel mean = a_eff·theta_eff days. All 197,088 exported divides: + +- **tau_uh median 1.50 days** (IQR 0.87–4.22; mean 5.23). 71.6% of divides + carry more than 1 day of UH delay, 25.5% more than 4 days. +- Distribution is range-saturated at both ends: ~5% at the 0.05-day floor + (delta kernel) and p95 at the 21.0-day ceiling (a and theta both maxed). + routa median 2.511 (near its 2.9 ceiling), routb median 0.072 (near its + 0 floor → theta_eff ≈ 0.57 d). +- **spearman(tau_uh, log10 uparea) = +0.383**: the per-divide UH delay + GROWS with upstream area. A hillslope-scale delay would not; a + network-travel-scale delay does. + +Reading: the q' entering ddrs already carries a median 1.5 days of learned +routing delay per divide, area-dependent, because the UH was trained at +gage scale (applied after divide summation) and therefore absorbed the +basin's channel travel time — then the 2026-07-29 export baked that same +kernel into every divide's lateral inflow. MC routes the network on top. +This is §5g's measured 2× over-delay, now visible in the parameters +themselves: double routing is structural in the current store, not a +celerity artifact alone. (The two mechanisms still superpose: the trained +Manning's n 0.130 vs 0.05 reference makes MC's own leg slow as well.) + +Implication for the fix tier: the col-7 (unrouted) export fixes double +routing but loses the multi-day hillslope delay and scored 0.29 routed +(the 2026-07-29 finding, in reverse). Neither existing column is right for +MC: the clean target is a SUB-GRID-ONLY UH (hillslope + small-channel +delay, no network component) on the lateral inflows, with MC supplying all +network travel. Short of retraining water_loss with routing inside the +graph, a pragmatic middle is capping/shrinking the exported kernel (e.g. +theta_eff floor-scale) — but any such surgery needs its own gate. + +Artifact: `output/tau_sweep/gamma_uh_params.csv` (divide_id, routa, routb, +a_eff, theta_eff, tau_uh_days, uparea_km2). + +## 5i. Fix implemented: signed tau convention, default 9 (2026-08-08) + +`tau_trim_and_downsample` now slices `[tau : -(24-tau)]` with pooled day i +scored against OBS DAY i: tau = hours the routed output is advanced before +daily scoring (dMC-Juniata's sign; tau=0 is day-aligned). Legacy mapping: +old = new + 11 (shipped 3 ≡ new −8, wrong direction; optimum 20 ≡ new 9). +`params.tau` defaults to 9 and rejects ≥24. Callers updated (driver obs +pairing i+1→i; eval drops the FIRST pooled day instead of the last, so the +output zarr day axis is unchanged; probe binary likewise). +`scripts/tau_sweep.py` moved to the new axis (window start 24+tau vs the +zarr day axis, TAUS −12..23); pre-2026-08-08 dumps carry legacy +`tau_shipped` and need the pre-change script from git history. DDR-parity +fixture tests pin the LEGACY window explicitly (legacy tau=3 cuts the same +hours as new tau=16); DDR-Python's `compute_daily_runoff` still uses the +legacy form — port the convention there before comparing configs across +repos. **Every pre-2026-08-08 checkpoint trained at old tau=3 (new −8); +config `tau:` values do not carry across the change.** + +Retrain experiment (launched 2026-08-08): five CPU train-and-test runs, +one per streamflow store (`config/experiments/tau9_train_*.yaml`, +`scripts/run_tau9_source_trains.sh`) — the epoch-30 config with only +{streamflow, sparse_solver: cpu} changed and tau at the new default 9. +Caveat noted at launch: the aorc2f_lumped store's measured optimum is +new-convention ≈ −8 (§5e), so tau=9 is expected to HURT that arm; it runs +anyway as the consistency control. + +**Retrain results (2026-08-09, full test window 1995-10..2010-09, 1,841 +gauges, 30 epochs each):** + +| arm | median NSE | median KGE | +|---|---|---| +| summed-q' baseline (no routing, day-aligned) | 0.642 | — | +| OLD tau (epoch-30 reference, aorc2f dist, legacy tau=3 ≡ new −8) | 0.620 | 0.699 | +| **tau=9 aorc2f distributed** | **0.706** | **0.730** | +| **tau=9 UH retrospective** | **0.707** | **0.738** | +| tau=9 daily LSTM | 0.578 | 0.616 | + +The timing fix alone is worth **+0.086 median NSE** on the flagship arm +and flips routing from LOSING to the summed-q' baseline (0.620 vs 0.642) +to beating it by +0.064. Run IDs: `2026-08-09T{03-05-54,09-30-39,12-05-08}Z- +train-and-test`. Operational note: the hourly_lstm arm finished training +(run `2026-08-09T14-55-05Z`, 60 checkpoints) but its eval phase was killed +with the parent driver at chunk 49/366; a train-and-test resume from the +final checkpoint cannot re-enter Phase 2 (it requires checkpoints written +by its own Phase 1 — `tau9_train_hourly_lstm_evalresume.yaml` records the +failed attempt), so the arm completes via the legacy eval binary +(`scripts/run_tau9_hourly_eval_chain.sh`), chained behind the +aorc2f_lumped run (`scripts/run_tau9_remaining.sh`). + +## 6. Raw output + +`output/tau_sweep/`: `summary_wy1996.md`, `nse_by_tau_wy1996.csv` (1841×24), +`best_tau_wy1996.csv`, `nse_vs_tau_by_area.png`, `best_tau_hist.png`, +`best_tau_vs_longitude.png`, `eval.log`, `hourly_full.f32(.json)`, +`eval_full.zarr`. + +## 7. Reproduce + +```bash +# eval + dump (~1.3 h GPU) — binary from commit 76a4020 +RUN=.ddrs/runs/2026-08-05T04-58-58Z-conus-experimental-train-and-test +DDRS_HOURLY_DUMP=$PWD/output/tau_sweep/hourly_full.f32 target/release/eval \ + --config $RUN/config.yaml --checkpoint $RUN/checkpoints/epoch_30_mb_1 \ + --output $PWD/output/tau_sweep/eval_full.zarr + +# offline sweep +uv run --with "zarr>=3" --with numpy --with pandas --with matplotlib --with tabulate \ + python scripts/tau_sweep.py \ + --dump output/tau_sweep/hourly_full.f32 --zarr output/tau_sweep/eval_full.zarr \ + --baseline-dir $RUN/baseline \ + --gages-csv ~/projects/ddr/references/gage_info/gages_2000_area_balanced.csv \ + --out-dir output/tau_sweep --pilot-start 1995-10-01 --pilot-end 1996-09-30 +``` diff --git a/docs/superpowers/plans/2026-08-02-ddr-match-physics-corrections.md b/docs/superpowers/plans/2026-08-02-ddr-match-physics-corrections.md new file mode 100644 index 0000000..aab5540 --- /dev/null +++ b/docs/superpowers/plans/2026-08-02-ddr-match-physics-corrections.md @@ -0,0 +1,802 @@ +# `ddr_match` Physics Corrections Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add a `params.ddr_match: bool` flag (default `true`) that, when `false`, replaces four physically-incorrect behaviours in the Muskingum-Cunge core — the `5/3` celerity factor, the constant `X = 0.3`, unmonitored negative-discharge clamping, and the absent Courant guard — each with a gradient-exact corrected implementation. + +**Architecture:** All four corrections live inside `forward_chain_inner` / `timestep_backward_core` in `src/routing/mmc_op.rs`, selected by a single boolean threaded through `SavedState`. `ddr_match: true` reproduces DDR bit-for-bit so `examples/compare_ddr_sandbox` stays an ABSOLUTE MATCH (invariant 1). Every new forward term gets a hand-derived backward branch (invariant 4 — no autograd-tape unrolling) validated by finite-difference gradcheck. + +**Tech Stack:** Rust, BURN 0.21 autograd (`Backward` custom ops), `cargo test`, NdArray backend for deterministic tests. + +--- + +## Concerns for the user + +**What could go wrong, and why:** + +1. **Silent parity loss.** If `ddr_match` is not threaded into *every* branch point, a config with `ddr_match: true` could still take a corrected path and break invariant 1 without an obvious symptom. Mitigated by Task 1's parity test, which must run before any physics change. +2. **Gradient-exactness regression (invariant 4).** Both β and Cunge `X` introduce new dependencies on tensors that *already* carry gradient paths (`area`, `top_width`, `wetted_perimeter`, `side_slope`, `q_t`, `celerity`). A missed chain-rule term produces a plausible-but-wrong gradient that training will silently absorb — exactly the failure mode that cost the AdaDelta run. Mitigated by a dedicated gradcheck per correction, run *before* the correction is used in any training run. +3. **Cunge X makes Courant worse before it makes it better.** `X ≈ 0.49` narrows the non-negative-coefficient window from `[0.6, 1.4]` to roughly `[0.98, 1.02]`. Enabling Task 4 without Task 5 will increase negative-discharge frequency. **Tasks 4 and 5 must be evaluated together**, and Task 2's counter is the instrument that proves it. +4. **Sub-stepping multiplies tape depth.** `n_sub` sub-steps per hourly timestep multiply autograd tape entries by `n_sub`. At `n_sub = 4` over 2160 hourly steps this is a real memory increase on GPU. Mitigated by making `n_sub` adaptive (per-reach, capped) rather than global. +5. **The corrections may not improve skill.** Each is individually absorbable by `n` (proven for the slope clamp; likely for a near-constant β). The scientific payoff is an interpretable parameter field, not necessarily a better NSE. Do not promote `ddr_match: false` on the basis of physics alone — gate it on a measured comparison. + +**Assumptions made:** + +- **DDR parity is worth preserving** as the only end-to-end reference check the port has, so `true` is the default and DDR itself is not modified. If the team decides to fix DDR too, Tasks 3–5 become the reference and the fixture is regenerated. +- **`X` becomes a derived quantity, not a learnable one.** `x_storage` stays out of `kan_head.learnable_parameters`; under `ddr_match: false` it is computed from Cunge's formula. Making it *both* learnable and Cunge-derived is contradictory and is explicitly out of scope. +- **Instrumentation is safe in both modes.** Counting negative solves changes no numerics, so Task 2 is unconditional and lands first. +- **Hourly `Δt = 3600 s` stays hardcoded** (`mmc.rs:33`). Sub-stepping divides it internally rather than changing the forcing cadence. + +**Blast radius:** + +| File | Change | Risk | +|---|---|---| +| `src/config.rs` | +1 field, +1 default fn | low — additive, defaulted | +| `src/routing/mmc_op.rs` | forward branches + backward branches + `SavedState` field | **high** — invariants 1 and 4 both live here | +| `src/routing/mmc.rs` | thread flag; read + log the counter | low | +| `tests/` | 4 new test files | none | +| `examples/compare_ddr_sandbox.rs` | none (default preserves it) | none | + +Diagram: `.claude/PHYSICS-CORRECTIONS.md`. + +**Out of scope (separate plan):** every disaggregation-head finding. That is an independent subsystem; see "Follow-up" at the end. + +--- + +## File Structure + +| File | Responsibility | +|---|---| +| `src/config.rs` | Declare `ParamsSection.ddr_match: bool`, default `true` | +| `src/routing/mmc_op.rs` | Both forward paths, both backward paths, `SavedState.ddr_match`, negative-solve counter | +| `src/routing/mmc.rs` | Pass `cfg.params.ddr_match` down; report the counter per forward | +| `tests/ddr_match_flag.rs` | Flag defaults, and `false` actually changes output | +| `tests/negative_discharge_counter.rs` | Counter fires on a known-unstable network | +| `tests/celerity_beta.rs` | β analytic vs finite-difference `dQ/dA`; limits; gradcheck | +| `tests/cunge_x.rs` | `X` formula, clamping, diffusion match; gradcheck | +| `tests/courant_substep.rs` | Sub-stepping keeps `Cr` in range and conserves mass | + +--- + +## Task 1: `ddr_match` config flag (no behaviour change) + +**Files:** +- Modify: `src/config.rs` (`ParamsSection`, near `use_cuda_graphs`) +- Modify: `src/routing/mmc_op.rs` (`SavedState`, `forward_chain_inner` signature, `timestep_forward`) +- Test: `tests/ddr_match_flag.rs` + +- [ ] **Step 1: Write the failing test** + +```rust +// tests/ddr_match_flag.rs +//! `ddr_match` defaults to true so every existing config and the DDR sandbox +//! parity example keep their current behaviour (invariant 1). +use ddrs::config::Config; + +#[test] +fn ddr_match_defaults_to_true() { + let yaml = r#" +mode: training +geodataset: merit +seed: 42 +np_seed: 42 +params: + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] +"#; + let cfg: Config = serde_yaml::from_str(yaml).expect("parse"); + assert!(cfg.params.ddr_match, "ddr_match must default to true"); +} + +#[test] +fn ddr_match_can_be_disabled() { + let yaml = r#" +mode: training +geodataset: merit +seed: 42 +np_seed: 42 +params: + ddr_match: false + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] +"#; + let cfg: Config = serde_yaml::from_str(yaml).expect("parse"); + assert!(!cfg.params.ddr_match); +} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `cargo test --test ddr_match_flag` +Expected: FAIL — `no field 'ddr_match' on type 'ParamsSection'` + +- [ ] **Step 3: Add the config field** + +In `src/config.rs`, inside `ParamsSection` (next to `use_cuda_graphs`): + +```rust + /// When `true` (default) the routing core reproduces DDR's formulation + /// bit-for-bit, including two known physical approximations: + /// * celerity `c = v · 5/3` (the wide-rectangular Kleitz-Seddon limit, + /// ~22-27% high for the trapezoid this code actually builds), and + /// * Muskingum `X ≡ 0.3` (constant, NOT Cunge-derived, giving a median + /// 10-30x excess numerical diffusion). + /// + /// Set `false` to enable the corrected physics. This CHANGES FORWARD + /// OUTPUT and will break `examples/compare_ddr_sandbox`'s ABSOLUTE MATCH + /// (invariant 1) — which is why the default preserves DDR behaviour. + /// See `.claude/PHYSICS-CORRECTIONS.md`. + #[serde(default = "default_ddr_match")] + pub ddr_match: bool, +``` + +Add the default fn near the other `default_*` helpers: + +```rust +fn default_ddr_match() -> bool { + true +} +``` + +Add `ddr_match: default_ddr_match(),` to `ParamsSection`'s `Default` impl. + +- [ ] **Step 4: Run test to verify it passes** + +Run: `cargo test --test ddr_match_flag` +Expected: PASS, 2 tests + +- [ ] **Step 5: Thread the flag into the op (still unused)** + +In `src/routing/mmc_op.rs`, add to `SavedState` (after `depth_lb: f32`): + +```rust + pub ddr_match: bool, +``` + +Add a parameter to `forward_chain_inner` after `discharge_lb: f32`: + +```rust + ddr_match: bool, +``` + +Populate it in the `SavedState` construction (`ddr_match,`) and pass it at the call site in `timestep_forward`: + +```rust + let ddr_match = cfg.params.ddr_match; +``` + +- [ ] **Step 6: Verify nothing changed** + +Run: `cargo build --release && cargo run --release --example compare_ddr_sandbox` +Expected: `ABSOLUTE MATCH` (max abs diff < 1e-3 m³/s) + +Run: `cargo test --test sp8_gradcheck --test sparse_gradcheck --test mmc` +Expected: all PASS + +- [ ] **Step 7: Commit** + +```bash +git add src/config.rs src/routing/mmc_op.rs tests/ddr_match_flag.rs +git commit -m "feat(config): ddr_match flag, defaulting to DDR-matching behaviour" +``` + +--- + +## Task 2: Negative-discharge instrumentation (both modes, no physics change) + +The `clamp_min(1e-4)` at `mmc_op.rs:919` converts every negative solve to `+1e-4`, creating mass and hiding Courant instability. Frequency has never been measured. This task measures it and changes nothing else. + +**Files:** +- Modify: `src/routing/mmc_op.rs` (counter + increment at S28) +- Modify: `src/routing/mmc.rs` (report per forward) +- Test: `tests/negative_discharge_counter.rs` + +- [ ] **Step 1: Write the failing test** + +```rust +// tests/negative_discharge_counter.rs +//! The S28 clamp silently turns negative solves into +1e-4. This counter is +//! the only way to see how often Muskingum's non-negative-coefficient +//! condition (2X <= Cr <= 2(1-X)) is violated in a real run. +use ddrs::routing::mmc_op::{negative_solve_stats, reset_negative_solve_stats}; + +#[test] +fn counter_starts_at_zero_and_resets() { + reset_negative_solve_stats(); + let (neg, total) = negative_solve_stats(); + assert_eq!(neg, 0); + assert_eq!(total, 0); +} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `cargo test --test negative_discharge_counter` +Expected: FAIL — `unresolved import ddrs::routing::mmc_op::negative_solve_stats` + +- [ ] **Step 3: Implement the counter** + +At the top of `src/routing/mmc_op.rs`: + +```rust +use std::sync::atomic::{AtomicU64, Ordering}; + +/// Count of solve outputs that came out NEGATIVE before the S28 +/// `clamp_min(discharge_lb)` rewrote them to `+1e-4`. +/// +/// Why this exists: Muskingum coefficients are only non-negative for +/// `2X <= Cr <= 2(1-X)` with `Cr = dt/K`. Measured on CONUS at mean flow with +/// `X = 0.3`, 69.8% of reaches sit outside that window (28.4% give `c1 < 0`, +/// 41.4% give `c3 < 0`), so negative discharge is expected — and the clamp +/// both CREATES MASS and removes the only symptom. Nothing in the codebase +/// measured this before 2026-08-02. +static NEG_SOLVES: AtomicU64 = AtomicU64::new(0); +static TOTAL_SOLVES: AtomicU64 = AtomicU64::new(0); + +/// `(negative_count, total_count)` since the last reset. +pub fn negative_solve_stats() -> (u64, u64) { + (NEG_SOLVES.load(Ordering::Relaxed), TOTAL_SOLVES.load(Ordering::Relaxed)) +} + +pub fn reset_negative_solve_stats() { + NEG_SOLVES.store(0, Ordering::Relaxed); + TOTAL_SOLVES.store(0, Ordering::Relaxed); +} +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `cargo test --test negative_discharge_counter` +Expected: PASS + +- [ ] **Step 5: Increment at S28** + +In `forward_chain_inner`, immediately before the existing clamp (`mmc_op.rs:919`): + +```rust + // Count negatives BEFORE the clamp masks them. Diagnostic only: reads the + // solve to host, changes no numerics, and is identical in both ddr_match + // modes. + { + let v: Vec = wrap(x_sol_prim.clone()).into_data().to_vec::().unwrap(); + let neg = v.iter().filter(|x| **x < 0.0).count() as u64; + NEG_SOLVES.fetch_add(neg, Ordering::Relaxed); + TOTAL_SOLVES.fetch_add(v.len() as u64, Ordering::Relaxed); + } + // S28: q_next = max(x_sol, discharge_lb) + let q_next = x_sol.clone().clamp_min(discharge_lb); +``` + +(Substitute the actual identifier for the pre-clamp solve primitive in scope; it is `x_sol` wrapped from the solver output.) + +- [ ] **Step 6: Report it once per forward** + +In `src/routing/mmc.rs`, at the end of the timestep loop in `forward`: + +```rust + let (neg, total) = crate::routing::mmc_op::negative_solve_stats(); + if total > 0 && neg > 0 { + eprintln!( + " negative solves before clamp: {neg}/{total} ({:.3}%) — \ + Muskingum coefficient sign violation (see .claude/PHYSICS-CORRECTIONS.md)", + 100.0 * neg as f64 / total as f64 + ); + } + crate::routing::mmc_op::reset_negative_solve_stats(); +``` + +- [ ] **Step 7: Verify parity still holds and observe the number** + +Run: `cargo run --release --example compare_ddr_sandbox` +Expected: `ABSOLUTE MATCH`, plus a `negative solves before clamp:` line + +- [ ] **Step 8: Commit** + +```bash +git add src/routing/mmc_op.rs src/routing/mmc.rs tests/negative_discharge_counter.rs +git commit -m "feat(diag): count negative Muskingum solves before the S28 clamp" +``` + +--- + +## Task 3: Trapezoidal celerity β (`ddr_match: false`) + +**Physics.** For a trapezoid, `c = dQ/dA = (dQ/dy)/T` with `Q = (1/n)A^(5/3)P^(-2/3)√S` gives + +``` +c = v · β, β = 5/3 − (4/3)·A·√(1+z²)/(T·P) +``` + +`β → 5/3` as `b/y → ∞` (wide rectangular) and `β → 4/3` as `b → 0` (triangular). The code's channels have `κ = b/y ∈ [0.7, 1.8]`, giving `β ≈ 1.31–1.36` — so the hardcoded `5/3` is **22–27% high**. + +**Files:** +- Modify: `src/routing/mmc_op.rs` (S17 forward; B17 backward) +- Test: `tests/celerity_beta.rs` + +- [ ] **Step 1: Write the physics test (no code yet)** + +```rust +// tests/celerity_beta.rs +//! beta = 5/3 - (4/3)·A·sqrt(1+z^2)/(T·P) is the exact kinematic-wave +//! celerity ratio c/v for a trapezoid. Verified three ways: against the +//! wide-rectangular limit (5/3), the triangular limit (4/3), and a +//! finite-difference dQ/dA on the same section. + +fn beta(b: f64, z: f64, y: f64) -> f64 { + let a = (b + z * y) * y; + let t = b + 2.0 * z * y; + let p = b + 2.0 * y * (1.0 + z * z).sqrt(); + 5.0 / 3.0 - (4.0 / 3.0) * a * (1.0 + z * z).sqrt() / (t * p) +} + +fn q_manning(b: f64, z: f64, y: f64, n: f64, s: f64) -> (f64, f64) { + let a = (b + z * y) * y; + let p = b + 2.0 * y * (1.0 + z * z).sqrt(); + ((1.0 / n) * a.powf(5.0 / 3.0) * p.powf(-2.0 / 3.0) * s.sqrt(), a) +} + +#[test] +fn beta_recovers_wide_rectangular_limit() { + assert!((beta(1e6, 0.0, 2.0) - 5.0 / 3.0).abs() < 1e-4); +} + +#[test] +fn beta_recovers_triangular_limit() { + assert!((beta(0.0, 2.0, 2.0) - 4.0 / 3.0).abs() < 1e-12); +} + +#[test] +fn beta_matches_finite_difference_dq_da() { + let (n, s) = (0.08, 2e-3); + for &(b, z, y) in &[(6.3, 0.50, 9.31), (13.8, 0.57, 8.09), (8.8, 2.76, 6.38)] { + let h = y * 1e-6; + let (q1, a1) = q_manning(b, z, y - h, n, s); + let (q2, a2) = q_manning(b, z, y + h, n, s); + let (q0, a0) = q_manning(b, z, y, n, s); + let fd = ((q2 - q1) / (a2 - a1)) / (q0 / a0); + assert!((fd - beta(b, z, y)).abs() < 1e-6, "b={b} z={z} y={y}: fd={fd}"); + } +} + +#[test] +fn ddr_five_thirds_is_biased_high_for_these_channels() { + // Regression guard on the MAGNITUDE of the defect ddr_match=false fixes. + let bt = beta(13.8, 0.57, 8.09); + let err = (5.0 / 3.0) / bt - 1.0; + assert!(err > 0.20 && err < 0.30, "expected +20..30% bias, got {err:.3}"); +} +``` + +- [ ] **Step 2: Run to verify the physics tests pass immediately** + +Run: `cargo test --test celerity_beta` +Expected: PASS, 4 tests (these validate the formula, not yet the implementation) + +- [ ] **Step 3: Implement the forward branch** + +In `forward_chain_inner`, replace S17 (`mmc_op.rs:849`): + +```rust + // S17: celerity. + // ddr_match=true -> c = v·5/3, the wide-rectangular Kleitz-Seddon + // limit. Matches ddr/mmc.py:167. WRONG for the + // trapezoid built above (kappa = b/y ~ 0.7-1.8 here, + // so the true ratio is ~1.31-1.36, not 1.667). + // ddr_match=false -> exact trapezoidal c = dQ/dA = (dQ/dy)/T: + // beta = 5/3 - (4/3)·A·sqrt(1+z^2)/(T·P) + let celerity = if ddr_match { + velocity_cl.clone() * (5.0_f32 / 3.0_f32) + } else { + let root = (side_slope.clone().powf_scalar(2.0) + 1.0).sqrt(); + let beta = -(_area.clone() * root) / (top_width.clone() * wp.clone()) * (4.0 / 3.0) + + (5.0 / 3.0); + velocity_cl.clone() * beta + }; +``` + +- [ ] **Step 4: Implement the backward branch** + +Derivation. With `G ≡ 5/3 − β` (so `β = 5/3 − G` and `G = (4/3)·A·u/(T·P)`, `u = √(1+z²)`): + +``` +∂β/∂A = −G/A ∂β/∂T = +G/T ∂β/∂P = +G/P ∂β/∂z = −G·z/(1+z²) +``` + +In `timestep_backward_core`, replace B17 (`mmc_op.rs:352`): + +```rust + // B17. celerity = velocity_cl · beta + // ddr_match: beta is the constant 5/3, so only gvelocity_cl exists. + // otherwise: beta depends on area/top_width/wp/side_slope, all of + // which already have gradient paths — these are ADDITIONAL + // contributions, not replacements. + let (gvelocity_cl, gbeta_terms) = if state.ddr_match { + (gcelerity.clone() * (5.0 / 3.0), None) + } else { + let ss = wrap(state.side_slope.clone()); + let tw = wrap(state.top_width.clone()); + let area = (tw.clone() + wrap(state.bottom_width.clone())) + * wrap(state.depth.clone()) + / 2.0; + let u = (ss.clone().powf_scalar(2.0) + 1.0).sqrt(); + let p = wrap(state.bottom_width.clone()) + + wrap(state.depth.clone()) * u.clone() * 2.0; + let g_term = area.clone() * u.clone() / (tw.clone() * p.clone()) * (4.0 / 3.0); + let beta = -g_term.clone() + (5.0 / 3.0); + let v_cl = wrap(state.velocity_clamped.clone()); + let gbeta = gcelerity.clone() * v_cl; + ( + gcelerity.clone() * beta, + Some(( + -gbeta.clone() * g_term.clone() / area, // ∂/∂A + gbeta.clone() * g_term.clone() / tw, // ∂/∂T + gbeta.clone() * g_term.clone() / p, // ∂/∂P + -gbeta * g_term.clone() * ss.clone() + / (ss.clone().powf_scalar(2.0) + 1.0), // ∂/∂z + )), + ) + }; +``` + +Then, where the existing backward accumulates `garea`, `gtop_width`, `gwp` and `gside_slope` (the S12/S13/S14 chain), add the four terms: + +```rust + let (garea, gtop_width, gwp, gside_slope) = match gbeta_terms { + None => (garea, gtop_width, gwp, gside_slope), + Some((ga, gt, gp, gz)) => ( + garea + ga, + gtop_width + gt, + gwp + gp, + gside_slope + gz, + ), + }; +``` + +- [ ] **Step 5: Write the gradcheck** + +Append to `tests/celerity_beta.rs`, modelled on `tests/sp8_gradcheck.rs` (copy its `linear_chain_sparse`, `mock_cfg`, `default_inputs`, `run_forward_loss`, `compute_analytical_grad`, `compute_fd_grad`, `compare_grads` helpers verbatim, then): + +```rust +#[test] +fn gradcheck_beta_path_n() { + // mock_cfg() must set params.ddr_match = false for this file. + let a = compute_analytical_grad(Parent::N); + let fd = compute_fd_grad(Parent::N); + compare_grads("n (ddr_match=false)", &a, &fd); +} + +#[test] +fn gradcheck_beta_path_q_spatial() { + let a = compute_analytical_grad(Parent::QSpatial); + let fd = compute_fd_grad(Parent::QSpatial); + compare_grads("q_spatial (ddr_match=false)", &a, &fd); +} + +#[test] +fn gradcheck_beta_path_p_spatial() { + let a = compute_analytical_grad(Parent::PSpatial); + let fd = compute_fd_grad(Parent::PSpatial); + compare_grads("p_spatial (ddr_match=false)", &a, &fd); +} +``` + +- [ ] **Step 6: Run the gradcheck** + +Run: `cargo test --test celerity_beta` +Expected: PASS, 7 tests. A failure here means a missing chain-rule term — do NOT proceed. + +- [ ] **Step 7: Verify parity is untouched** + +Run: `cargo run --release --example compare_ddr_sandbox && cargo test --test sp8_gradcheck` +Expected: `ABSOLUTE MATCH`; sp8 PASS + +- [ ] **Step 8: Commit** + +```bash +git add src/routing/mmc_op.rs tests/celerity_beta.rs +git commit -m "feat(routing): exact trapezoidal celerity under ddr_match=false" +``` + +--- + +## Task 4: Cunge-derived `X` (`ddr_match: false`) + +**Physics.** Cunge chooses `X` so Muskingum's numerical diffusion equals the physical hydraulic diffusivity `Q/(2·B·S₀)`: + +``` +X = clamp( 0.5·(1 − Q/(B·S₀·c·Δx)), 0, 0.5 ) +``` + +with `B` = top width, `Δx` = reach length. Measured Cunge `X ≈ 0.49` on CONUS; the current constant `0.3` injects `D_num/D_phys` of median **28×**. + +**Files:** +- Modify: `src/routing/mmc_op.rs` (new S19 forward; new backward branch) +- Test: `tests/cunge_x.rs` + +- [ ] **Step 1: Write the physics test** + +```rust +// tests/cunge_x.rs +//! Cunge X matches Muskingum numerical diffusion to physical hydraulic +//! diffusivity: D_num = c·dx·(0.5 - X) == D_phys = Q/(2·B·S). + +fn cunge_x(q: f64, b: f64, s: f64, c: f64, dx: f64) -> f64 { + (0.5 * (1.0 - q / (b * s * c * dx))).clamp(0.0, 0.5) +} + +#[test] +fn cunge_x_matches_numerical_to_physical_diffusion() { + let (q, b, s, c, dx) = (300.0, 40.0, 2e-3, 1.4, 6598.0); + let x = cunge_x(q, b, s, c, dx); + let d_num = c * dx * (0.5 - x); + let d_phys = q / (2.0 * b * s); + assert!((d_num / d_phys - 1.0).abs() < 1e-9, "d_num={d_num} d_phys={d_phys}"); +} + +#[test] +fn constant_x_030_over_diffuses_by_an_order_of_magnitude() { + // Regression guard on the magnitude of the defect this task fixes. + let (q, b, s, c, dx) = (300.0, 40.0, 2e-3, 1.4, 6598.0); + let d_num_const = c * dx * (0.5 - 0.3); + let d_phys = q / (2.0 * b * s); + let ratio = d_num_const / d_phys; + assert!(ratio > 2.0, "expected heavy over-diffusion, got {ratio:.2}x"); +} + +#[test] +fn cunge_x_clamps_into_zero_half() { + assert_eq!(cunge_x(1e9, 40.0, 2e-3, 1.4, 6598.0), 0.0); // huge Q -> negative raw + assert!(cunge_x(1e-9, 40.0, 2e-3, 1.4, 6598.0) <= 0.5); +} + +#[test] +fn muskingum_coefficients_sum_to_one_for_any_x() { + for &x in &[0.0_f64, 0.3, 0.49, 0.5] { + let (k, dt) = (3295.0_f64, 3600.0_f64); + let denom = 2.0 * k * (1.0 - x) + dt; + let c1 = (dt - 2.0 * k * x) / denom; + let c2 = (dt + 2.0 * k * x) / denom; + let c3 = (2.0 * k * (1.0 - x) - dt) / denom; + assert!((c1 + c2 + c3 - 1.0).abs() < 1e-12, "x={x}"); + } +} +``` + +- [ ] **Step 2: Run** + +Run: `cargo test --test cunge_x` +Expected: PASS, 4 tests + +- [ ] **Step 3: Implement the forward branch** + +In `forward_chain_inner`, immediately after `k_muskingum` (`mmc_op.rs:852`), replace the use of `xst_in`: + +```rust + // S19: Muskingum X. + // ddr_match=true -> the caller's constant (forward.rs sets 0.3). + // NOT Cunge-derived: severs the link between + // numerical and physical diffusion, giving a median + // 28x over-diffusion on CONUS. + // ddr_match=false -> Cunge: X = clamp(0.5(1 - Q/(B·S·c·L)), 0, 0.5), + // which makes D_num = c·L·(0.5-X) equal the physical + // hydraulic diffusivity Q/(2·B·S). + let x_eff = if ddr_match { + xst_in.clone() + } else { + let w = qt_in.clone() + / (top_width.clone() * slope_in.clone() * celerity.clone() * length_in.clone() + + 1e-12); + (-w + 1.0).mul_scalar(0.5).clamp(0.0, 0.5) + }; + let one_minus_x = -x_eff.clone() + 1.0; + let two_k = k_muskingum.clone() * 2.0; + let two_kx = two_k.clone() * x_eff.clone(); +``` + +Save it for backward — add to `SavedState`: + +```rust + pub x_effective: B::FloatTensorPrimitive, +``` + +- [ ] **Step 4: Implement the backward branch** + +`X` feeds only `two_kx` and `two_k_1mx`, so from the existing totals: + +``` +gX = two_k · (g_2kx_total − g_2k1mx_total) +``` + +and with `W = Q/(B·S·c·L)`, `X_raw = 0.5(1 − W)`: + +``` +∂X/∂Q = −0.5·W/Q ∂X/∂B = +0.5·W/B ∂X/∂c = +0.5·W/c +``` + +(all zero where the clamp saturates). After the existing `g_2kx_total` / `g_2k1mx_total` are formed: + +```rust + if !state.ddr_match { + let two_k_t = wrap(state.k_muskingum.clone()) * 2.0; + let gx = two_k_t * (g_2kx_total.clone() - g_2k1mx_total.clone()); + // Zero the gradient where the [0, 0.5] clamp saturated. + let x_eff_t = wrap(state.x_effective.clone()); + let unsat = x_eff_t.clone().greater_elem(0.0).bool_and( + x_eff_t.clone().lower_elem(0.5), + ); + let gx = gx.mask_fill(unsat.bool_not(), 0.0); + + let qt_t = wrap(state.q_t.clone()); + let tw_t = wrap(state.top_width.clone()); + let cel_t = wrap(state.celerity.clone()); + let w = qt_t.clone() + / (tw_t.clone() * wrap(state.slope.clone()) * cel_t.clone() + * wrap(state.length.clone()) + + 1e-12); + + gq_t_total = gq_t_total + gx.clone() * (-w.clone() * 0.5) / qt_t; + gtop_width = gtop_width + gx.clone() * (w.clone() * 0.5) / tw_t; + gcelerity_total = gcelerity_total + gx * (w * 0.5) / cel_t; + } +``` + +**Note the ordering constraint:** this adds to `gcelerity`, so it must run *before* B18 consumes `gcelerity`. Restructure so `gcelerity` is fully accumulated first. + +- [ ] **Step 5: Add the gradcheck** + +Append the same three gradcheck tests as Task 3 Step 5 to `tests/cunge_x.rs`, with `mock_cfg()` setting `ddr_match = false`, plus: + +```rust +#[test] +fn gradcheck_cunge_x_q_t() { + // Q_t now enters X as well as the RHS — the new path this task adds. + let a = compute_analytical_grad(Parent::QT); + let fd = compute_fd_grad(Parent::QT); + compare_grads("q_t (cunge X)", &a, &fd); +} +``` + +- [ ] **Step 6: Run** + +Run: `cargo test --test cunge_x` +Expected: PASS, 8 tests + +- [ ] **Step 7: Verify parity and measure the Courant consequence** + +Run: `cargo run --release --example compare_ddr_sandbox` +Expected: `ABSOLUTE MATCH` + +Run the smoke config with `ddr_match: false` and read the Task 2 counter. +Expected: negative-solve percentage **increases** versus `ddr_match: true` — Cunge `X ≈ 0.49` narrows the stable window to ~`[0.98, 1.02]`. This is the expected, documented reason Task 5 exists. + +- [ ] **Step 8: Commit** + +```bash +git add src/routing/mmc_op.rs tests/cunge_x.rs +git commit -m "feat(routing): Cunge-derived Muskingum X under ddr_match=false" +``` + +--- + +## Task 5: Courant sub-stepping (`ddr_match: false`) + +Sub-divide `Δt` per timestep so `Cr = Δt_sub/K` lands inside `[2X, 2(1−X)]`. + +**Files:** +- Modify: `src/routing/mmc.rs` (sub-step loop), `src/routing/mmc_op.rs` (accept `dt_sub`) +- Test: `tests/courant_substep.rs` + +- [ ] **Step 1: Write the test** + +```rust +// tests/courant_substep.rs +//! Sub-stepping must (a) bring the Courant number into the non-negative +//! coefficient window and (b) conserve mass exactly. + +fn n_sub_for(k: f64, dt: f64, x: f64, cap: u32) -> u32 { + // Need dt_sub/K <= 2(1-x) => n_sub >= dt/(K·2(1-x)) + let need = (dt / (k * 2.0 * (1.0 - x))).ceil() as u32; + need.clamp(1, cap) +} + +#[test] +fn substep_brings_courant_into_window() { + let (dt, x) = (3600.0_f64, 0.49_f64); + for &k in &[13120.0_f64, 3295.0, 1716.0, 770.0] { + let n = n_sub_for(k, dt, x, 16); + let cr = (dt / n as f64) / k; + assert!(cr <= 2.0 * (1.0 - x) + 1e-9, "K={k} n={n} Cr={cr}"); + } +} + +#[test] +fn steady_state_is_preserved_under_substepping() { + // O = I + q_L must hold regardless of how many sub-steps are taken. + let (k, x) = (3295.0_f64, 0.49_f64); + for &n in &[1_u32, 2, 4, 8] { + let dt = 3600.0 / n as f64; + let denom = 2.0 * k * (1.0 - x) + dt; + let c1 = (dt - 2.0 * k * x) / denom; + let c2 = (dt + 2.0 * k * x) / denom; + let c3 = (2.0 * k * (1.0 - x) - dt) / denom; + assert!((c1 + c2 + c3 - 1.0).abs() < 1e-12, "n={n}"); + } +} +``` + +- [ ] **Step 2: Run** + +Run: `cargo test --test courant_substep` +Expected: PASS, 2 tests + +- [ ] **Step 3: Implement** + +In `src/routing/mmc.rs`, inside the timestep loop, when `!cfg.params.ddr_match`: + +```rust + // Sub-step so Cr = dt_sub/K stays inside [2X, 2(1-X)]. Capped at + // 16: tape depth scales with n_sub, and beyond ~16 the memory cost + // outweighs the accuracy gain. Reaches still outside the window + // after capping are reported by the Task 2 counter. + const N_SUB_CAP: u32 = 16; + let n_sub = if cfg.params.ddr_match { 1 } else { N_SUB_CAP.min(4) }; + let dt_sub = DT_SECONDS / n_sub as f32; + for _ in 0..n_sub { + q = crate::routing::mmc_op::timestep_forward::(/* ..., dt_sub */); + } +``` + +Thread `dt_sub` through `timestep_forward` → `forward_chain_inner` (replacing the `dt` constant) and store it in `SavedState` for the backward. + +- [ ] **Step 4: Run the full gate** + +Run: `cargo test --test courant_substep --test cunge_x --test celerity_beta --test sp8_gradcheck` +Expected: all PASS + +Run: `cargo run --release --example compare_ddr_sandbox` +Expected: `ABSOLUTE MATCH` + +- [ ] **Step 5: Commit** + +```bash +git add src/routing/mmc.rs src/routing/mmc_op.rs tests/courant_substep.rs +git commit -m "feat(routing): Courant sub-stepping under ddr_match=false" +``` + +--- + +## Task 6: End-to-end comparison + +- [ ] **Step 1: Run both modes on the smoke config** + +```bash +cp config/experiments/gradaccum_smoke.yaml /tmp/smoke_ddr_true.yaml +sed 's/^ sparse_solver:/ ddr_match: false\n sparse_solver:/' \ + config/experiments/gradaccum_smoke.yaml > /tmp/smoke_ddr_false.yaml +for f in /tmp/smoke_ddr_true.yaml /tmp/smoke_ddr_false.yaml; do + target/release/ddrs --config $f run --workflow train \ + --workspace /tmp/ws_$(basename $f .yaml) --backend cpu 2>&1 | \ + grep -E "negative solves|median_n|mb=0 loss" +done +``` + +- [ ] **Step 2: Record the comparison** + +Capture in `docs/2026-08-02-ddr-match-findings.md`: negative-solve % in each mode, `median_n` trajectory in each mode, and loss. **`ddr_match: false` is promoted only if negative solves drop AND `n` moves toward the 0.025–0.15 NLCD band.** + +- [ ] **Step 3: Commit** + +```bash +git add docs/2026-08-02-ddr-match-findings.md +git commit -m "docs: ddr_match true-vs-false comparison on the smoke config" +``` + +--- + +## Follow-up (separate plan) + +The disaggregation-head audit is a distinct subsystem and needs its own plan. Highest-value items, in order: (1) the head takes **1 log-daily-Q + 24 same-day precip hours**, *not* the `[d-1,d,d+1]` window plus attributes the config comment claims — fix the comment; (2) no ablation of the *frozen* head exists in the current configuration; (3) no `clamp_min(0.0)` before `log()` in `disagg_head.rs:238-242`, combined with `use_cuda_graphs: true`, is the exact pairing that hid the 2026-06-23 NaN bug. diff --git a/docs/superpowers/plans/2026-08-04-positivity-clamp.md b/docs/superpowers/plans/2026-08-04-positivity-clamp.md new file mode 100644 index 0000000..e0ea4c5 --- /dev/null +++ b/docs/superpowers/plans/2026-08-04-positivity-clamp.md @@ -0,0 +1,396 @@ +# Positivity Clamp — Provably Zero Negative Muskingum Solves + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development +> to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax. + +**Goal:** Drive `negative solves before clamp` to exactly 0/N by enforcing the +Muskingum non-negativity window `2X ≤ Cr ≤ 2(1−X)` on every reach-timestep, +behind a new `params.enforce_positivity` flag. + +**Architecture:** Clamp the *inputs* (K, X), never the coefficients. The identity +`c1+c2+c3 = 1` holds for any `(K, X)`, so clamping inputs preserves mass exactly +while clamping `c3` would not. + +**Tech stack:** Rust, BURN 0.21 hand-written `Backward` (invariant 4). + +--- + +## The theorem this rests on + +With `Cr = Δt/K` and `denom = 2K(1−X) + Δt > 0`: + +``` +c2 = (2KX + Δt)/denom > 0 always (K>0, X≥0, Δt>0) +c4 = 2Δt/denom > 0 always +c1 = (Δt − 2KX)/denom ≥ 0 <=> Cr ≥ 2X +c3 = (2K(1−X) − Δt)/denom ≥ 0 <=> Cr ≤ 2(1−X) +``` + +The solve is forward substitution in topological order: +`x[i] = b[i] + c1[i]·Σ_{j∈up(i)} x[j]`, with +`b[i] = c2·(N q_t)[i] + c3·q_t[i] + c4·q'[i]`. + +**Claim.** If `c1 ≥ 0` and `c3 ≥ 0` everywhere, then `x ≥ 0` everywhere. +*Proof:* `q_t > 0` (S28 `clamp_min(1e-4)` and the hotstart at `utils.rs:97`), +`q' > 0` (`mmc.rs:453`), so `b ≥ 0`. Induction over the topological order: +headwaters have no upstream so `x = b ≥ 0`; each subsequent `x[i]` is a +non-negative combination of `b[i]` and already-non-negative upstream values. ∎ + +Verified numerically: `2X ≤ Cr ≤ 2(1−X)` ⟺ `c1,c3 ≥ 0` with 0 mismatches in +200,000 random `(K,X)` draws; 34,364 negatives unclamped → **0** clamped over +800,000 chain solves. + +### Why a margin δ is mandatory + +At δ=0 the clamp lands exactly on `c1=0` / `c3=0`, and f32 roundoff crosses it: +7,149 `c1<0` and 964 `c3<0` in a 400k f32 sweep. Measured minima: + +| δ | min c1 | min c3 | negative solves (240k adversarial) | +|---|---|---|---| +| 0 | −3.3e−08 | −6.8e−08 | >0 | +| 1e−4 | +1.8e−06 | +0.0e+00 | 0 | +| 1e−3 | +1.8e−05 | +5.1e−07 | 0 | +| **1e−2** | **+1.8e−04** | **+5.0e−05** | **0** | + +**Use δ = 1e−2** — ~400× f32 eps, and it costs almost nothing (X ceiling tightens +by 1%, K floor rises by 1%). + +## The clamp (ddr_match=false AND enforce_positivity only) + +``` +k_floor = Δt·(1+δ)/2 scalar constant +k_musk = max(k_raw, k_floor) k_raw = length/celerity +cr = Δt / k_musk => cr ∈ (0, 2/(1+δ)] +x_hi_a = (1−δ)·0.5·cr enforces c1 ≥ 0 +x_hi_b = (1−δ)·(1 − 0.5·cr) enforces c3 ≥ 0 +x_eff = min(x_cunge, x_hi_a, x_hi_b) +``` + +`x_hi_b > 0` is guaranteed by the K floor (`cr < 2`), so no `clamp_min` is +needed — the min of three positives is positive. + +### Known cost (do not hide this in the findings) + +> **SUPERSEDED — the table below is WRONG.** It evaluates `X_max` *at* each Cr +> percentile, but `X_max(Cr)` is non-monotone (peaks at `Cr = 1`), so that does +> not give the percentiles of X. The tell: the p95 row came out *below* p75. +> Measured truth (`src/bin/probe_courant.rs`, 1,841 gauges): Cr p50 = **0.226** +> (not 1.09), the cap binds on **95.3 %** of reach-timesteps, and median X used +> falls **0.4976 → 0.0794 (6.3×)**, not to 0.45. See +> `.claude/PHYSICS-CORRECTIONS.md` §`enforce_positivity` for the corrected +> tables. Kept here only so the error is traceable. + +At real CONUS Courant percentiles the cap binds **everywhere**: + +| | Cr raw | Cr after floor | X_max | vs Cunge 0.49 | +|---|---|---|---|---| +| p5 | 0.19 | 0.190 | 0.094 | cap binds | +| p25 | 0.54 | 0.540 | 0.267 | cap binds | +| p50 | 1.09 | 1.090 | 0.450 | cap binds | +| p75 | 2.46 | 1.980 | 0.010 | cap binds | +| p95 | 10.20 | 1.980 | 0.010 | cap binds | + +So this **partially overrides the Cunge X** (commit `54ec215`): X becomes +stability-dictated rather than diffusion-matched, except near Cr≈1. The fully +capped limit is benign, not degenerate — `c1=0.495, c2=0.505, c3=0.000`, i.e. +`Q_out(t+1) ≈ ½I(t+1) + ½I(t)`, a mass-conserving 2-step lag for a sub-grid reach. + +The K floor inflates travel time only where `Cr > 2` (≈ the fastest quartile; +4.2× at p5). A reach with `K < Δt/2` is sub-grid — the timestep cannot resolve +its transit, and the unclamped scheme expressed that as oscillation that S28 +clamped to 1e-4 anyway. + +**Sub-stepping does not substitute for this.** The Task-5 abandonment analysed +landing *inside* `[0.98,1.02]` at fixed X≈0.49. Here the K constraint is +one-sided (`K ≥ Δt/2n`), which `n_sub=8` satisfies even at p5=425 s — but +shrinking Δt drives already-slow reaches *further* from Cr=1 (p95 would go +Cr 0.19→0.024, X_max→0.012). Only per-reach Δx puts every reach near Cr≈1. +Record this; do not re-attempt global sub-stepping. + +--- + +## File structure + +- `src/config.rs` — add `params.enforce_positivity: bool` (default **false**), + plumb through `ParamsRaw`, validate it requires `ddr_match: false`. +- `src/routing/mmc_op.rs` — forward S18′/S19′; backward B18′/B19′. +- `tests/positivity_clamp.rs` — new: positive control, zero-guarantee, + partition identity, off-parity, falsifiable gradcheck. +- `.claude/PHYSICS-CORRECTIONS.md` — new section + blast-radius row. + +--- + +### Task 1: Config flag `enforce_positivity` + +**Files:** Modify `src/config.rs`; Test: `tests/ddr_match_flag.rs` + +- [ ] **Step 1: Write the failing tests** + +Append to `tests/ddr_match_flag.rs`: + +```rust +#[test] +fn enforce_positivity_defaults_to_false() { + let cfg = load_mock_config(); + assert!(!cfg.params.enforce_positivity, + "enforce_positivity must default to false so existing runs are unchanged"); +} + +#[test] +fn enforce_positivity_requires_corrected_physics() { + // ddr_match: true + enforce_positivity: true must be rejected at load: + // the clamp changes K and X, which would break compare_ddr_sandbox. + let yaml = mock_config_yaml_with("ddr_match: true\n enforce_positivity: true\n"); + let err = load_config_str(&yaml).expect_err("must reject"); + assert!(err.to_string().contains("enforce_positivity"), + "error must name the offending key, got: {err}"); +} +``` + +Follow the exact helper names already used in `tests/ddr_match_flag.rs`; if they +differ, adapt the test to the existing helpers rather than inventing new ones. + +- [ ] **Step 2: Run to verify it fails** + +`cargo test --test ddr_match_flag` → FAIL (`no field enforce_positivity`). + +- [ ] **Step 3: Implement** + +In `src/config.rs`, mirror exactly how `ddr_match` is declared, defaulted, and +plumbed through `ParamsRaw` (search for `ddr_match` and copy the pattern). +Default `false`. Extend the existing `validate_ddr_match` (≈line 871) — or add a +sibling validator called from the same place — to reject +`enforce_positivity && ddr_match`. + +- [ ] **Step 4: Verify** — `cargo test --test ddr_match_flag` → PASS. + +- [ ] **Step 5: Commit** + +```bash +git add src/config.rs tests/ddr_match_flag.rs +git commit -m "feat(config): enforce_positivity flag, requires ddr_match: false" +``` + +--- + +### Task 2: Forward — K floor and X stability cap + +**Files:** Modify `src/routing/mmc_op.rs`; Test: `tests/positivity_clamp.rs` + +- [ ] **Step 1: Write the failing tests** + +Create `tests/positivity_clamp.rs`. Build a stress network whose reaches span +`Cr` from ≈0.02 to ≈10 (short/fast reaches with `Cr > 2` are the ones that +produce negatives). Use `reset_negative_solve_stats` / +`negative_solve_stats` from `src/routing/mmc_op.rs` and +`enable_negative_discharge_tracking`. + +```rust +#[test] +fn positive_control_negatives_exist_without_the_clamp() { + let (neg, tot) = run_stress_network(/* enforce_positivity */ false); + assert!(neg > 0, "fixture must actually produce negatives, got {neg}/{tot}"); +} + +#[test] +fn clamp_drives_negatives_to_exactly_zero() { + let (neg, tot) = run_stress_network(/* enforce_positivity */ true); + assert_eq!(neg, 0, "expected exactly zero negative solves, got {neg}/{tot}"); + assert!(tot > 0, "fixture must have solved something"); +} + +#[test] +fn partition_identity_survives_the_clamp() { + // c1 + c2 + c3 == 1 to f32 tolerance on every reach, clamp on. + let (c1, c2, c3) = coefficients_from_stress_network(true); + for i in 0..c1.len() { + assert!((c1[i] + c2[i] + c3[i] - 1.0).abs() < 1e-5, + "reach {i}: c1+c2+c3 = {}", c1[i] + c2[i] + c3[i]); + } +} + +#[test] +fn coefficients_are_non_negative_with_margin() { + let (c1, _, c3) = coefficients_from_stress_network(true); + assert!(c1.iter().all(|&v| v >= 0.0), "min c1 = {:?}", c1.iter().cloned().fold(f32::INFINITY, f32::min)); + assert!(c3.iter().all(|&v| v >= 0.0), "min c3 = {:?}", c3.iter().cloned().fold(f32::INFINITY, f32::min)); +} + +#[test] +fn off_parity_byte_identical_when_disabled() { + // enforce_positivity: false must reproduce current ddr_match:false output exactly. + let a = route_stress_network(false); + let b = route_stress_network_baseline(); // pre-change code path + assert_eq!(a, b, "disabled clamp must be a byte-identical no-op"); +} +``` + +- [ ] **Step 2: Run to verify they fail** + +`cargo test --test positivity_clamp` → the positive control passes, the +zero-guarantee test FAILS with a non-zero count. + +- [ ] **Step 3: Implement the forward** + +In `forward_chain_inner` (`src/routing/mmc_op.rs`), replace the current +`k_muskingum` / `x_eff` block. Keep `ddr_match: true` and +`enforce_positivity: false` byte-identical — guard with a bool that is only true +when BOTH `!ddr_match` and `enforce_positivity`. + +```rust +const POSITIVITY_DELTA: f32 = 1e-2; + +let k_raw = length_in.clone() / celerity.clone(); +let k_muskingum = if enforce_pos { + // A reach with K < dt/2 is sub-grid: the timestep cannot resolve its + // transit, and the unclamped scheme expressed that as oscillation that + // S28 clamped to 1e-4 anyway. Flooring K makes the coarse-graining + // explicit and puts Cr in (0, 2/(1+d)], the feasible region for X. + k_raw.clone().clamp_min(dt * (1.0 + POSITIVITY_DELTA) / 2.0) +} else { + k_raw.clone() +}; +``` + +then after `x_cunge` is formed (the existing `((-w + 1.0) * 0.5).clamp(0.0, 0.5)`): + +```rust +let x_eff = if enforce_pos { + let cr = k_muskingum.clone().recip() * dt; + let hi_a = cr.clone() * (0.5 * (1.0 - POSITIVITY_DELTA)); // c1 >= 0 + let hi_b = (-cr.clone() * 0.5 + 1.0) * (1.0 - POSITIVITY_DELTA); // c3 >= 0 + x_cunge.clone().min_pair(hi_a).min_pair(hi_b) +} else { + x_cunge.clone() +}; +``` + +Use whatever elementwise-min helper BURN 0.21 exposes in this codebase (check +how `clamp`/`min` are already called nearby); do not add a dependency. + +- [ ] **Step 4: Verify** — `cargo test --test positivity_clamp` → all PASS. + +- [ ] **Step 5: Commit** + +```bash +git add src/routing/mmc_op.rs tests/positivity_clamp.rs +git commit -m "feat(routing): K floor + X stability cap under enforce_positivity" +``` + +--- + +### Task 3: Backward — B18′ K-floor mask and B19′ three-way min + +**Files:** Modify `src/routing/mmc_op.rs`; Test: `tests/positivity_clamp.rs` + +The forward introduced two new gradient facts: + +1. `k_musk = max(k_raw, k_floor)` — gradient is masked where floored. +2. `x_eff = min(x_cunge, hi_a, hi_b)` — gradient goes to exactly one branch, and + the `hi_a`/`hi_b` branches open a **new path** `x_eff → cr → k_musk → celerity` + that did not exist before. + +Derivatives: + +``` +d(hi_a)/d(cr) = +0.5·(1−δ) +d(hi_b)/d(cr) = −0.5·(1−δ) +d(cr)/d(k_musk) = −Δt / k_musk² +d(k_musk)/d(k_raw) = 1 where k_raw > k_floor, else 0 +``` + +Accumulation order (critical — mirrors the existing B18/B19 ordering note): + +``` +gk_musk = (existing g_2k_total · 2) from c1..c4 + + gx_eff·[mask_a]·(+0.5(1−δ)) ·(−Δt/k_musk²) NEW + + gx_eff·[mask_b]·(−0.5(1−δ)) ·(−Δt/k_musk²) NEW +gk_raw = gk_musk · mask(k_raw > k_floor) NEW +gcelerity = −gk_raw · length / celerity² (existing B18) + + XGrads.g_celerity (existing B19, now + masked by mask_cunge) +``` + +The existing `XGrads` path (`∂X/∂Q`, `∂X/∂B`, `∂X/∂c`) must additionally be +masked by `mask_cunge` — where `hi_a` or `hi_b` won the min, `x_cunge` is not +the active branch and receives nothing. + +- [ ] **Step 1: Write the failing gradcheck** + +Add to `tests/positivity_clamp.rs`, following the pattern in `tests/cunge_x.rs`: + +```rust +#[test] +fn positivity_clamp_gradcheck() { + // Analytical vs central finite difference on n, p_spatial, q_spatial. + // Fixture MUST exercise all three min branches and both sides of the + // K floor - assert that before comparing gradients. + let f = stress_fixture(); + assert!(f.frac_branch_cunge() > 0.05 && f.frac_branch_hi_a() > 0.05 + && f.frac_branch_hi_b() > 0.05, + "fixture is vacuous: branch mix {:?}", f.branch_mix()); + assert!(f.frac_k_floored() > 0.05 && f.frac_k_floored() < 0.95, + "fixture must straddle the K floor, got {}", f.frac_k_floored()); + let (analytic, fd) = f.grads(); + let rel = rel_err(&analytic, &fd); + assert!(rel < 1e-2, "gradcheck rel err {rel:.3e}"); +} +``` + +**The fixture-mix assertions are not optional.** The Cunge-X gradcheck was +initially vacuous because at 1000 m reaches `W ≈ 1.6` saturated the clamp on +every reach and all four tests passed with the backward terms deleted. Prove +falsifiability the same way it was proved there. + +- [ ] **Step 2: Run to verify it fails** — non-trivial `rel` before the backward exists. + +- [ ] **Step 3: Implement the backward** per the accumulation order above. + +- [ ] **Step 4: Verify falsifiability** + +Temporarily delete the two new `gk_musk` terms, re-run, confirm the gradcheck +FAILS, then restore. Record both numbers in the commit message. + +- [ ] **Step 5: Run the full gate set** + +```bash +cargo test --test positivity_clamp --test cunge_x --test celerity_beta \ + --test sparse_gradcheck --test mmc --test leakance_gradcheck +cargo run --release --example compare_ddr_sandbox # must still be ABSOLUTE MATCH +``` + +- [ ] **Step 6: Commit** + +```bash +git add src/routing/mmc_op.rs tests/positivity_clamp.rs +git commit -m "feat(routing): exact backward for the positivity clamp" +``` + +--- + +### Task 4: Document + +**Files:** Modify `.claude/PHYSICS-CORRECTIONS.md` + +- [ ] **Step 1:** Add an S18′/S19′ branch to the ASCII dataflow, a section + carrying the theorem, the δ table, and the CONUS Cr/X_max cost table above. +- [ ] **Step 2:** Add the blast-radius row: + `| Positivity clamp | S18′,S19′ | gk_musk mask + cr path | yes (flag-gated) | positivity_clamp |` +- [ ] **Step 3:** Correct the Task-5 note to record *why* sub-stepping still + does not substitute (one-sided vs two-sided constraint; slow reaches move away + from Cr≈1). +- [ ] **Step 4: Commit** + +```bash +git add .claude/PHYSICS-CORRECTIONS.md +git commit -m "docs: positivity clamp theorem, delta margin, and its cost" +``` + +--- + +## Verification criteria + +1. `cargo test --test positivity_clamp` — all pass, gradcheck proven falsifiable. +2. `compare_ddr_sandbox` still ABSOLUTE MATCH (untouched: needs `ddr_match: true`). +3. `enforce_positivity: false` byte-identical to current output. +4. A short real run logs `negative solves before clamp: 0/N (0.000%)`. diff --git a/docs/superpowers/plans/2026-08-05-reach-subdivision.md b/docs/superpowers/plans/2026-08-05-reach-subdivision.md new file mode 100644 index 0000000..692e5df --- /dev/null +++ b/docs/superpowers/plans/2026-08-05-reach-subdivision.md @@ -0,0 +1,1266 @@ +# Reach Subdivision (variable Δx) Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development +> (recommended) or superpowers:executing-plans to implement this plan task-by-task. +> Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Normalize every MERIT reach to `Δx ≈ c_ref·Δt` at adjacency-build time — +**splitting** reaches that are too long and **clamping the length** of reaches that +are too short — so `Cr ≈ 1` network-wide and the Muskingum coefficients are +non-negative *by construction*, with no runtime clamp and no gradient masking. + +**The two-sided rule:** + +``` +Δx_target = c_ref · Δt + L > Δx_target → split into m = ceil(L/Δx_target) pieces of length L/m, q' → q'/m + L < Δx_target → clamp length up to Δx_target (do NOT merge) +``` + +Merging short reaches is rejected deliberately: a short reach may have two upstream +tributaries or a parallel tributary joining below it, so collapsing it would destroy +junction structure. Clamping its length achieves the same `Cr` with no topology change. + +**Why a static length clamp is safe where the runtime K floor was not.** They are +algebraically related (`L ≥ c·Δt ⟺ K ≥ Δt`), but `enforce_positivity`'s K floor is +applied per timestep to a *learned* celerity — it masks gradients and makes +`X ∝ Cr ∝ 1/n`, which is what drove `n` to its floor on 98 % of reaches. A build-time +length clamp is just a different constant in the `length_m` array: **no gradient +path, no `∂X/∂n` coupling, and zero change to `mmc_op.rs` or any backward.** + +**Architecture:** Subdivision is a **static preprocessing step** inside the managed +adjacency builder, between `build_conus_adjacency` and the gauge-subgraph/zarr write. +The runtime sees only a larger network. Piece count `m` is computed once from a +reference flow and **capped at 8** — the graph must be fixed for the CSR pattern and +the hand-written autograd, so `m` can never vary per timestep. + +**Tech Stack:** Rust, BURN 0.21, `zarrs`, `blake3` (cache keys), `ndarray`. + +## Global Constraints + +- **f32 throughout the routing core.** No f64/bf16 casts (invariant 2, `CLAUDE.md`). +- **`examples/compare_ddr_sandbox` must stay an ABSOLUTE MATCH** (max abs diff < 1e-3 + m³/s). It builds its own tiny network and must be unaffected. +- **Adjacency must stay topologically ordered and lower-triangular** (`rows[k] >= cols[k]`), + asserted at `src/adjacency/build.rs:235-240` via `BuildError::NotLowerTriangular`. + The forward-substitution solver depends on it (invariant 3). +- **Do not replace the hand-written sparse backward** in `src/sparse/` with + autograd-tape unrolling (invariant 4). +- **Never run `git add -A`.** Stage only files you changed, by name. +- **Never kill a running process.** Training runs may be active. +- Default OFF: `params.subdivision.enabled: false` reproduces current behaviour + byte-for-byte. + +--- + +## Measured feasibility — read before designing anything + +Computed by mirroring `forward_chain_inner` S1–S17 in NumPy against the real +adjacency store, trained KAN parameters from +`.ddrs/runs/2026-08-03T13-11-00Z-train-and-test/plot/kan_parameters.nc`, and +per-divide median Q' accumulated downstream through the topological order. + +**The model was validated against four independently measured in-run anchors:** + +| Quantity | model | measured in-run | +|---|---|---| +| median `Cr` | 0.250 | 0.226 | +| `X` p5/p50/p95 | 0.340 / 0.4973 / 0.5000 | 0.330 / 0.4976 / 0.5000 | +| `enforce_positivity` bind rate | 95.5 % | 95.3 % | +| median `X_eff` after cap | 0.0864 | 0.0794 | + +### Cap sweep + +| cap M | sub-reaches | reach × | edges × | **critical path ×** | median Cr | median `X_eff` | +|---|---|---|---|---|---|---| +| 1 (today) | 346,321 | 1.00 | 1.00 | 1.00 | 0.250 | 0.086 | +| 2 | 633,500 | 1.83 | 1.85 | 1.50 | 0.499 | 0.146 | +| **4** | 1,071,482 | 3.09 | 3.14 | 1.85 | **1.00** | 0.239 | +| **8** | **1,652,965** | **4.77** | 4.86 | **2.18** | 1.10 | **0.321** | +| 16 | 2,326,605 | 6.72 | 6.84 | 2.64 | 1.10 | 0.372 | +| uncapped | 4,586,580 | 13.24 | 13.52 | **9.23** | 1.10 | 0.417 | + +**Uncapped is infeasible — and the blocker is latency, not memory.** Per-timestep +state is only 422 MB uncapped, but forward substitution is sequential over +topological levels and chaining sub-reaches lengthens the critical path 9.23×, on a +trainer already forward-pass bound at ~27.5 s/micro-batch. + +**Capping also makes the cost estimable.** 43 % of reaches have no Q' in the +configured store and must be imputed; across defensible imputations uncapped `Σm` +swings **2.3 M – 10.5 M** (4.6× uncertainty), but `Σ min(m,4)` stays within ±6 % +and `Σ min(m,8)` within ±12 %. **Any design resting on uncapped `Σm` rests on a +number that cannot be pinned down.** + +### CORRECTION — the "restores X's dynamic range" claim is WRONG + +An earlier conversational claim (and the reasoning that motivated this work) held +that subdivision would un-saturate the Cunge X. **It does not.** Raw `X_cunge` +median moves only **0.4973 → 0.4815** even uncapped. `Q/(B·S·c·Δx)` is small +primarily because of the `attribute_minimums.slope = 1e-3` floor and large top +width `B` — not because Δx is long. **Do not justify this change on X's dynamic +range.** + +The two real justifications are: + +1. **Non-negativity becomes automatic.** At `Cr = 1`, `c1` and `c3` both reduce to + `(1 − 2X)/(1 + 2(1−X))`, which is `≥ 0` for *any* `X ≤ 0.5` — a bound already + enforced physically. So `enforce_positivity` (and its 6.3× X collapse, and the + `n`-to-floor degeneracy it causes) becomes unnecessary. +2. **If `enforce_positivity` is kept**, subdivision recovers most of its damage: + median `X_eff` 0.086 → 0.321 at M=8 (71 % of achievable relief). + +### The two-sided rule closes the gap subdivision alone cannot + +Subdivision fixes only reaches that are too *long*. Adding the short-reach length +clamp (indicative reference celerity, cap 8): + +| | subdivide only | **subdivide + length clamp** | +|---|---|---| +| median Cr | 1.10 | 1.08 | +| frac `Cr > 2` | 17.3 % (unfixable by splitting) | **0.00 %** | +| frac `Cr < 0.5` | 17.7 % | **0.05 %** | + +Essentially the whole network lands inside the stable window, which is what makes +"non-negative by construction" an honest claim rather than an aspiration. + +### Why the target is `Δx = c·Δt` and NOT the Ponce–Theurer limit + +Formulas verified verbatim against [Ponce](https://ponce.sdsu.edu/muskingum_cunge_method_explained.html): +`K = Δx/c`; `X = ½(1 − q/(So·c·Δx))` with `q` the **unit-width** discharge; and +`C0 = (Δt−2KX)/denom`, `C1 = (Δt+2KX)/denom`, `C2 = (2K(1−X)−Δt)/denom` — identical +to ddrs's `c1`, `c2`, `c3` (`mmc_op.rs:1077-1080`). Ponce also confirms the failure +mode: *"for very large values of the space step, there is a tendency for physically +unrealistic negative outflows"*, and *"negative values of C2 are invariably +associated with dips in the rising portion of the outflow hydrograph"* — his `C2` is +our `c3`. + +[Ponce & Theurer](https://ponce.sdsu.edu/accuracy_criteria_in_diffusion_routing.html) +give the accuracy criterion `C·D ≥ ξ` (a **product**) and the limit +`Δx ≤ ½(c·Δt + qo/(So·c))`. **That limit was evaluated and rejected for this network:** + +| Δx target | Σm (cap 8) | median Cr | frac Cr > 2 | frac `C·D ≥ 0.33` | +|---|---|---|---|---| +| `c·Δt` (this plan) | 918 k | **1.08** | **0.00 %** | 0.3 % | +| Ponce–Theurer `½(c·Δt + qo/So·c)` | 1.42 M | 2.06 | **57.3 %** | 17.7 % | + +The cell Reynolds number on MERIT is `D ≈ 0.012` — physical diffusion is ~1–2 % of +advective transport — so `Δx_D/Δx_C ≈ 0.020` and the Ponce–Theurer limit collapses to +`≈ c·Δt/2`, i.e. `C ≈ 2`. That **violates** the non-negativity ceiling `C ≤ 2(1−X)` +(which is `≈ 1` when `X ≈ 0.5`) and puts 57 % of reaches above `Cr = 2` — it makes +negative coefficients *worse*. + +**`C·D ≥ ξ` is unsatisfiable here at any Δx that also keeps coefficients +non-negative.** It is an accuracy criterion for *diffusion* routing; as `D → 0` there +is no physical diffusion left to resolve and the criterion degenerates. Confirmed not +to be an artifact of the slope floor: lowering `attribute_minimums.slope` from `1e-3` +to `1e-6` moves median `D` only 0.0112 → 0.0120. + +So the target is the Courant length, `Δx = c·Δt` — which is also HEC-HMS's Auto-DX +rule — because at `C = 1` both coefficients reduce to `(1−2X)/(1+2(1−X)) ≥ 0` for any +`X ≤ 0.5`. + +### Where the cost lives + +`corr(log m, log uparea) = −0.14`; `corr(log m, log c) = −0.75`; `corr(log m, log L) = +0.57`. +The cost is **long slow mid-order reaches**, not big rivers — the 71–169 km² uparea +bin alone is 33.7 % of `Σm`. Large rivers are already near Cr = 1 and mostly need +no split (17.1 % of all reaches need `m = 1`). + +### Concerns for the user + +- **The length clamp inflates total network length.** Indicative measurement with a + crude reference celerity: 34.4 % of reaches clamped, median factor 2.0×, p95 12.5×, + p99 36×, **max 48,597×**, total CONUS channel length +17.1 %. The extreme tail is + degenerate MERIT geometry (11 reaches shorter than 10 m), so the clamp arguably + repairs bad data there — but a reach modelled 12× longer than reality has a 12× + longer travel time, which is a real physical distortion in exchange for numerical + stability. **Consider a separate cap on the clamp factor** and measure the + sensitivity in Task 8. +- **The clamped fraction is sensitive to `c_ref`.** The 34.4 % above used a crude + `Q_ref = 0.01·uparea^0.9` giving median `c = 1.03 m/s`; the *calibrated* reference + (validated against measured Cr) gives median `c = 0.443 m/s`, hence a smaller + `Δx_target` and materially fewer clamped reaches. **Recompute with the calibrated + celerity before trusting any of these figures.** +- **A static `m` is only Cr≈1 at the reference flow.** Median `c` varies ~3× between + low and high flow, so reaches will be over-split at high flow and under-split at + low flow. This is unavoidable: the graph must be fixed for autograd. +- **Retraining is required.** Every learned parameter was fit against the un-split + network's effective diffusion. Checkpoints do not transfer. +- **2.18× critical path is a real training-time cost** on top of an ~11 h run. +- **The KAN sees no new information.** Sub-reaches inherit their parent's attributes, + so subdivision cannot improve parameter identifiability — only numerics. + +### Assumptions + +- **Reference flow is config-specified, not checkpoint-derived.** Coupling + subdivision to a trained checkpoint would make the graph depend on training state. + The cap is what makes this safe (±6–12 % robustness). +- **Sub-reaches are hydraulically identical to their parent** — same `n`, `p`, `q`, + slope; length `L/m`. MERIT carries no within-reach variation to do better. +- **Lateral inflow is uniform along the reach**, so each piece gets `q'/m`. This is + HEC-HMS's own treatment: its lateral term is `C4·(q_L·Δx)` with `q_L` per unit + length. + +--- + +## File Structure + +| Path | Responsibility | +|---|---| +| `src/adjacency/subdivide.rs` | **NEW.** Reference celerity, piece counts, graph expansion. Pure functions over plain `Vec`s — no BURN, no I/O. | +| `src/adjacency/build.rs` | Call subdivision after `build_conus_adjacency`, before gauge subgraphs. | +| `src/adjacency/cache.rs` | Fold subdivision config into the content key; bump `BUILDER_VERSION`. | +| `src/data/store/zarr.rs` | Persist/load `parent_offset` + `parent_order`. | +| `src/config.rs` | `params.subdivision` block + validation. | +| `src/routing/mmc.rs` | Divide `q'` by piece count after the clamp. | +| `src/training/forward.rs` | Gather KAN outputs from parent rows to sub-reach rows. | +| `src/data/collate.rs` | Map each gauge to its parent's **outlet** piece. | +| `tests/subdivide.rs` | **NEW.** Unit tests for celerity, piece counts, expansion invariants. | +| `tests/subdivision_integration.rs` | **NEW.** End-to-end: mass conservation, Cr distribution, non-negative coefficients without the clamp. | + +### Data model + +`order` gains duplicates (m rows share one COMID), which would break the existing +`IdIndex` COMID→position lookup. Resolved by keeping **two** index spaces: + +``` +parent space (N = 346,321) sub-reach space (N' = Σ min(m,M)) + parent_order[p] = COMID order[i] = COMID of i's parent + IdIndex built HERE parent_offset[p]..parent_offset[p+1] + = contiguous rows owned by parent p + m_p = parent_offset[p+1] - parent_offset[p] +``` + +Pieces are contiguous and ordered upstream→downstream, so parent `p`'s **outlet** +is `parent_offset[p+1] - 1`. Because parents are already in topological order and +each chain runs low→high index, the expanded graph stays topologically ordered and +lower-triangular for free. + +``` + BEFORE AFTER (m=3) + U ──> P ──> D U₂ ──> P₀ ──> P₁ ──> P₂ ──> D₀ + └─ q'/3 q'/3 q'/3 + len(P) = L len(Pᵢ) = L/3, slope unchanged + gauge@P → row(P) gauge@P → P₂ (outlet = last piece) +``` + +--- + +### Task 1: Config block `params.subdivision` + +**Files:** +- Modify: `src/config.rs` +- Test: `tests/subdivide.rs` (create) + +**Interfaces:** +- Produces: `Subdivision { enabled: bool, max_pieces: usize, reference_n: f32, + reference_discharge_exponent: f32, reference_discharge_coefficient: f32 }`, + reachable as `cfg.params.subdivision`. + +- [ ] **Step 1: Write the failing tests** + +Create `tests/subdivide.rs`. Follow the inline-YAML-to-tempfile style already used +in `tests/ddr_match_flag.rs` (read it first; `Config::from_yaml_file` is the loader). + +```rust +#[test] +fn subdivision_defaults_to_disabled() { + let cfg = load_cfg(&yaml_with_params("")); + assert!(!cfg.params.subdivision.enabled); + assert_eq!(cfg.params.subdivision.max_pieces, 8); +} + +#[test] +fn subdivision_rejects_max_pieces_below_one() { + let err = try_load_cfg(&yaml_with_params( + " subdivision:\n enabled: true\n max_pieces: 0\n", + )) + .expect_err("must reject"); + assert!(err.to_string().contains("max_pieces"), "got: {err}"); +} + +#[test] +fn subdivision_enabled_loads() { + let cfg = load_cfg(&yaml_with_params( + " subdivision:\n enabled: true\n max_pieces: 4\n", + )); + assert!(cfg.params.subdivision.enabled); + assert_eq!(cfg.params.subdivision.max_pieces, 4); +} +``` + +- [ ] **Step 2: Run to verify it fails** + +`cargo test --test subdivide` → FAIL (`no field subdivision`). + +- [ ] **Step 3: Implement** + +In `src/config.rs`, mirror how `params.leakance` / `kan_head.disaggregation` declare a +nested optional block with serde defaults (grep for `disaggregation` and copy the +pattern). Add: + +```rust +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Subdivision { + #[serde(default)] + pub enabled: bool, + /// Hard cap on pieces per reach. Uncapped subdivision is infeasible: + /// 13.2x reaches and 9.2x solver critical path, and Sum(m) cannot be + /// pinned down (2.3M-10.5M across defensible reference-flow choices). + /// Capping bounds the cost AND makes it estimable (+/-12% at M=8). + #[serde(default = "default_max_pieces")] + pub max_pieces: usize, + /// Manning's n used ONLY to compute the reference celerity that sets m. + /// Deliberately NOT taken from a checkpoint: the graph must not depend on + /// training state. 0.05 is the trained CONUS median. + #[serde(default = "default_reference_n")] + pub reference_n: f32, + /// Reference discharge Q_ref = coefficient * uparea_km2^exponent (m3/s). + #[serde(default = "default_ref_q_coeff")] + pub reference_discharge_coefficient: f32, + #[serde(default = "default_ref_q_exp")] + pub reference_discharge_exponent: f32, + /// Short reaches get their length clamped UP to + /// `min_length_fraction * c_ref * dt`, giving `Cr <= 1/min_length_fraction` + /// at the reference flow. 1.0 targets Cr = 1; 0.5 targets Cr <= 2 (the + /// non-negativity bound) with half the length distortion; 0.0 disables the + /// clamp entirely, leaving short reaches over-Courant. + /// + /// This is a BUILD-TIME constant, unlike the runtime K floor in + /// `enforce_positivity`. It therefore has no gradient path and cannot + /// create the `X ~ Cr ~ 1/n` coupling that drove n to its floor. + #[serde(default = "default_min_length_fraction")] + pub min_length_fraction: f32, +} + +fn default_max_pieces() -> usize { 8 } +fn default_reference_n() -> f32 { 0.05 } +fn default_ref_q_coeff() -> f32 { 0.01 } +fn default_ref_q_exp() -> f32 { 0.9 } +fn default_min_length_fraction() -> f32 { 1.0 } + +impl Default for Subdivision { + fn default() -> Self { + Self { + enabled: false, + max_pieces: default_max_pieces(), + reference_n: default_reference_n(), + reference_discharge_coefficient: default_ref_q_coeff(), + reference_discharge_exponent: default_ref_q_exp(), + min_length_fraction: default_min_length_fraction(), + } + } +} +``` + +Add `#[serde(default)] pub subdivision: Subdivision` to `Params` and to `ParamsRaw` +(follow exactly how `ddr_match` is threaded through both). Add a validator beside +`validate_ddr_match` (`src/config.rs:804`) and call it from the same site +(`config.rs:~808`): + +```rust +fn validate_subdivision(cfg: &Config) -> std::result::Result<(), String> { + let s = &cfg.params.subdivision; + if s.enabled && s.max_pieces < 1 { + return Err("params.subdivision: `max_pieces` must be >= 1".to_string()); + } + if s.enabled && cfg.params.use_cuda_graphs { + return Err( + "params.subdivision: `enabled: true` requires `use_cuda_graphs: false` \ + — the captured graph is sized to a fixed reach count." + .to_string(), + ); + } + Ok(()) +} +``` + +- [ ] **Step 4: Verify** — `cargo test --test subdivide` → PASS. + +- [ ] **Step 5: Commit** + +```bash +git add src/config.rs tests/subdivide.rs +git commit -m "feat(config): params.subdivision block, default off" +``` + +--- + +### Task 2: Reference celerity and the two-sided reach plan + +**Files:** +- Create: `src/adjacency/subdivide.rs` +- Modify: `src/adjacency/mod.rs` (add `pub mod subdivide;`) +- Test: `tests/subdivide.rs` + +**Interfaces:** +- Consumes: `Subdivision` from Task 1. +- Produces: + ```rust + pub fn reference_celerity(uparea_km2: f32, slope: f32, cfg: &Subdivision) -> f32; + + /// Both sides of the two-sided rule. `pieces[i]` is how many sub-reaches + /// parent `i` becomes; `length_m[i]` is its (possibly clamped) total length, + /// which Task 3 then divides by `pieces[i]`. + pub struct ReachPlan { pub pieces: Vec, pub length_m: Vec } + + pub fn plan_reaches( + length_m: &[f32], slope: &[f32], uparea_km2: &[f32], + dt_seconds: f32, cfg: &Subdivision, + ) -> ReachPlan; + ``` + +`reference_celerity` mirrors the solver's own chain rather than inventing one. Use +the wide-channel approximation `c = (5/3)·v` here — **this is deliberate**: it is +only setting `m`, the cap dominates the answer, and pulling in the full trapezoidal +`β` would require `p_spatial`/`q_spatial`, which are learned. + +- [ ] **Step 1: Write the failing tests** + +```rust +use ddrs::adjacency::subdivide::{plan_reaches, reference_celerity, ReachPlan}; +use ddrs::config::Subdivision; + +fn cfg(max_pieces: usize) -> Subdivision { + Subdivision { enabled: true, max_pieces, ..Default::default() } +} + +#[test] +fn celerity_rises_with_slope_and_area() { + let c = cfg(8); + let lo = reference_celerity(100.0, 1e-3, &c); + assert!(reference_celerity(100.0, 1e-2, &c) > lo, "steeper must be faster"); + assert!(reference_celerity(10_000.0, 1e-3, &c) > lo, "bigger must be faster"); + assert!(lo > 0.0 && lo < 15.0, "celerity {lo} outside physical range"); +} + +#[test] +fn long_reaches_split_and_are_capped() { + let c = cfg(4); + let p = plan_reaches(&[200_000.0], &[1e-4], &[30.0], 3600.0, &c); + assert_eq!(p.pieces[0], 4, "must clamp to max_pieces"); + assert_eq!(p.length_m[0], 200_000.0, "long reaches keep their true length"); +} + +#[test] +fn short_reaches_are_length_clamped_not_split() { + let c = cfg(8); + // 50 m reach with a fast celerity: far below dx_target, so it stretches. + let p = plan_reaches(&[50.0], &[1e-2], &[10_000.0], 3600.0, &c); + assert_eq!(p.pieces[0], 1, "short reach must not split"); + let dx = reference_celerity(10_000.0, 1e-2, &c) * 3600.0; + assert!((p.length_m[0] - dx).abs() < 1e-3, + "expected clamp to dx_target {dx}, got {}", p.length_m[0]); + assert!(p.length_m[0] > 50.0, "clamp must lengthen, not shorten"); +} + +#[test] +fn min_length_fraction_zero_disables_the_clamp() { + let mut c = cfg(8); + c.min_length_fraction = 0.0; + let p = plan_reaches(&[50.0], &[1e-2], &[10_000.0], 3600.0, &c); + assert_eq!(p.length_m[0], 50.0, "clamp must be off"); +} + +#[test] +fn disabled_is_an_exact_no_op() { + let mut c = cfg(8); + c.enabled = false; + let p = plan_reaches(&[200_000.0, 50.0], &[1e-4, 1e-2], &[30.0, 10_000.0], 3600.0, &c); + assert_eq!(p.pieces, vec![1, 1]); + assert_eq!(p.length_m, vec![200_000.0, 50.0], "lengths must be untouched"); +} + +#[test] +fn degenerate_input_never_yields_zero_pieces_or_zero_length() { + let c = cfg(8); + let p = plan_reaches(&[5000.0, 0.0], &[0.0, 0.0], &[0.0, 0.0], 3600.0, &c); + assert!(p.pieces.iter().all(|&v| v >= 1), "pieces must be >= 1, got {:?}", p.pieces); + assert!(p.length_m.iter().all(|&v| v > 0.0), + "length must be > 0 (a 0 m reach gives K = 0 and c1 = 1), got {:?}", p.length_m); +} +``` + +- [ ] **Step 2: Run to verify it fails** — `cargo test --test subdivide` → FAIL (module missing). + +- [ ] **Step 3: Implement** + +```rust +//! Static reach subdivision so Cr = c*dt/dx lands near 1. +//! +//! HEC-HMS picks the space step as `dx = c*dt` (Technical Reference Manual, +//! Muskingum-Cunge Model). ddrs historically used the full MERIT reach length, +//! giving median Cr = 0.226 — reaches ~4.4x too long. See +//! `.claude/REACH-SUBDIVISION.md`. + +use crate::config::Subdivision; + +/// Reference celerity (m/s) used ONLY to choose the piece count. +/// +/// Mirrors the solver's S15/S17 chain with the wide-channel ratio c = (5/3)*v +/// instead of the exact trapezoidal beta: this only sets `m`, the cap dominates +/// the result, and beta needs the learned p_spatial/q_spatial. +pub fn reference_celerity(uparea_km2: f32, slope: f32, cfg: &Subdivision) -> f32 { + // Slope floor mirrors `attribute_minimums.slope` (mmc.rs:208). + let s = slope.max(1e-3); + let q_ref = cfg.reference_discharge_coefficient + * uparea_km2.max(0.0).powf(cfg.reference_discharge_exponent); + // Hydraulic radius via a wide-channel regime relation; the exponent 0.4 is + // the Leopold & Maddock downstream depth exponent. + let r = (q_ref.max(1e-3)).powf(0.4).max(0.01); + let v = (1.0 / cfg.reference_n) * r.powf(2.0 / 3.0) * s.sqrt(); + (v * (5.0 / 3.0)).clamp(0.01, 15.0) +} + +pub struct ReachPlan { + pub pieces: Vec, + pub length_m: Vec, +} + +/// The two-sided rule. Long reaches split; short reaches have their length +/// clamped UP to `dx_target`. Merging short reaches is deliberately not done: +/// a short reach can carry two upstream tributaries or have a parallel +/// tributary joining below it, so collapsing it would destroy junctions. +pub fn plan_reaches( + length_m: &[f32], + slope: &[f32], + uparea_km2: &[f32], + dt_seconds: f32, + cfg: &Subdivision, +) -> ReachPlan { + if !cfg.enabled { + return ReachPlan { pieces: vec![1; length_m.len()], length_m: length_m.to_vec() }; + } + let mut pieces = Vec::with_capacity(length_m.len()); + let mut lengths = Vec::with_capacity(length_m.len()); + for ((&l, &s), &a) in length_m.iter().zip(slope).zip(uparea_km2) { + let dx_target = reference_celerity(a, s, cfg) * dt_seconds; + // Short reach: stretch it so Cr ~ 1 at the reference flow. This is a + // STATIC constant, unlike the runtime K floor in `enforce_positivity` + // — no gradient path, so it cannot pull `n` toward its bound. + let l_eff = (l.max(0.0)).max(dx_target * cfg.min_length_fraction); + // Long reach: split. + let m = ((l_eff / dx_target).ceil().max(1.0) as u32).min(cfg.max_pieces as u32); + pieces.push(m); + lengths.push(l_eff); + } + ReachPlan { pieces, length_m: lengths } +} +``` + +- [ ] **Step 4: Verify** — `cargo test --test subdivide` → PASS. + +- [ ] **Step 5: Commit** + +```bash +git add src/adjacency/subdivide.rs src/adjacency/mod.rs tests/subdivide.rs +git commit -m "feat(adjacency): two-sided reach plan — split long, clamp short" +``` + +--- + +### Task 3: Graph expansion + +**Files:** +- Modify: `src/adjacency/subdivide.rs` +- Test: `tests/subdivide.rs` + +**Interfaces:** +- Consumes: `ReachPlan { pieces, length_m }` from Task 2; `ConusAdjacency { order, + rows, cols, length_m, slope, dropped_comids }` from `src/adjacency/build.rs:63-77`. +- Produces: + ```rust + pub struct SubdividedAdjacency { + pub order: Vec, // COMID per sub-reach row (duplicated) + pub parent_order: Vec, // original COMID order, length N + pub parent_offset: Vec, // length N+1; rows [o[p], o[p+1]) belong to p + pub rows: Vec, + pub cols: Vec, + pub length_m: Vec, + pub slope: Vec, + } + pub fn subdivide(adj: &ConusAdjacency, plan: &ReachPlan) -> SubdividedAdjacency; + ``` + +**Topology rules.** Parent `p` owns rows `[off[p], off[p+1])` ordered +upstream→downstream. Internal chain edges connect consecutive pieces. An external +edge `u -> p` (u upstream of p) becomes `outlet(u) -> off[p]`, where +`outlet(u) = off[u+1] - 1`. Each piece gets `plan.length_m[p] / m` — note this is +the **clamped** length from Task 2, not `adj.length_m[p]`. Slope is unchanged. + +- [ ] **Step 1: Write the failing tests** + +```rust +use ddrs::adjacency::build::ConusAdjacency; +use ddrs::adjacency::subdivide::subdivide; + +/// Chain of 3 reaches: 0 -> 1 -> 2 (rows = downstream, cols = upstream, +/// so rows[k] >= cols[k]). +fn chain3() -> ConusAdjacency { + ConusAdjacency { + order: vec![100, 200, 300], + rows: vec![1, 2], + cols: vec![0, 1], + length_m: vec![3000.0, 6000.0, 900.0], + slope: vec![1e-3, 2e-3, 3e-3], + dropped_comids: vec![], + } +} + +/// Explicit plan so these tests exercise expansion alone, independent of the +/// celerity heuristic in Task 2. +fn plan(pieces: Vec, length_m: Vec) -> ReachPlan { + ReachPlan { pieces, length_m } +} + +#[test] +fn expansion_preserves_total_length_and_slope() { + let s = subdivide(&chain3(), &plan(vec![3, 2, 1], chain3().length_m.clone())); + assert_eq!(s.length_m.len(), 6); + for (p, &m) in [3u32, 2, 1].iter().enumerate() { + let lo = s.parent_offset[p] as usize; + let hi = s.parent_offset[p + 1] as usize; + assert_eq!(hi - lo, m as usize, "parent {p} piece count"); + let total: f32 = s.length_m[lo..hi].iter().sum(); + assert!((total - chain3().length_m[p]).abs() < 1e-3, + "parent {p} length not conserved: {total}"); + assert!(s.slope[lo..hi].iter().all(|&v| v == chain3().slope[p]), + "slope must be inherited unchanged"); + } +} + +#[test] +fn expansion_stays_lower_triangular_and_topological() { + let s = subdivide(&chain3(), &plan(vec![3, 2, 1], chain3().length_m.clone())); + for (&r, &c) in s.rows.iter().zip(s.cols.iter()) { + assert!(r > c, "edge {c}->{r} violates strict lower-triangular ordering"); + } +} + +#[test] +fn expansion_edge_count_is_original_plus_internal_links() { + let s = subdivide(&chain3(), &plan(vec![3, 2, 1], chain3().length_m.clone())); + // 2 original edges + (3-1) + (2-1) + (1-1) internal = 5 + assert_eq!(s.rows.len(), 5, "rows: {:?} cols: {:?}", s.rows, s.cols); +} + +#[test] +fn external_edges_land_on_parent_outlet_and_inlet() { + let s = subdivide(&chain3(), &plan(vec![3, 2, 1], chain3().length_m.clone())); + // parent0 rows 0..3, parent1 rows 3..5, parent2 row 5. + // edge 0->1 becomes outlet(0)=2 -> inlet(1)=3 + assert!(s.rows.iter().zip(&s.cols).any(|(&r, &c)| c == 2 && r == 3), + "missing 2->3; rows {:?} cols {:?}", s.rows, s.cols); + // edge 1->2 becomes outlet(1)=4 -> inlet(2)=5 + assert!(s.rows.iter().zip(&s.cols).any(|(&r, &c)| c == 4 && r == 5), + "missing 4->5; rows {:?} cols {:?}", s.rows, s.cols); +} + +#[test] +fn all_ones_is_an_exact_identity() { + let a = chain3(); + let s = subdivide(&a, &plan(vec![1, 1, 1], a.length_m.clone())); + assert_eq!(s.order, a.order); + assert_eq!(s.rows, a.rows); + assert_eq!(s.cols, a.cols); + assert_eq!(s.length_m, a.length_m); + assert_eq!(s.parent_offset, vec![0, 1, 2, 3]); +} + +#[test] +fn expansion_uses_the_clamped_length_not_the_raw_one() { + let a = chain3(); // reach 2 is only 900 m + let clamped = vec![3000.0, 6000.0, 4000.0]; // reach 2 stretched to 4 km + let s = subdivide(&a, &plan(vec![1, 1, 1], clamped)); + assert_eq!(s.length_m[2], 4000.0, + "must use ReachPlan.length_m, not ConusAdjacency.length_m"); +} +``` + +- [ ] **Step 2: Run to verify it fails** — FAIL (`subdivide` not found). + +> **Verify the self-edge assumption before implementing.** The expansion asserts +> the input COO has no `r == c` entries. `build.rs:235-240` only rejects `r < c`, +> so `r == c` is *permitted* by that check. Confirm empirically first: +> ```bash +> /home/tbindas/projects/ddrs/ddrs-py/.venv/bin/python -c " +> import zarr, numpy as np +> g = zarr.open('/home/tbindas/projects/ddr/data/merit_conus_adjacency.zarr', mode='r') +> r, c = np.asarray(g['indices_0'][:]), np.asarray(g['indices_1'][:]) +> print('self-edges (r==c):', int((r==c).sum()), 'of', r.size)" +> ``` +> If the count is non-zero, STOP and report — the design needs a self-edge rule +> before Task 3 can be correct, and every downstream task inherits the error. + +- [ ] **Step 3: Implement** + +```rust +use crate::adjacency::build::ConusAdjacency; + +pub struct SubdividedAdjacency { + pub order: Vec, + pub parent_order: Vec, + pub parent_offset: Vec, + pub rows: Vec, + pub cols: Vec, + pub length_m: Vec, + pub slope: Vec, +} + +impl SubdividedAdjacency { + #[inline] + pub fn inlet(&self, parent: usize) -> usize { self.parent_offset[parent] as usize } + #[inline] + pub fn outlet(&self, parent: usize) -> usize { + self.parent_offset[parent + 1] as usize - 1 + } + #[inline] + pub fn pieces(&self, parent: usize) -> usize { + (self.parent_offset[parent + 1] - self.parent_offset[parent]) as usize + } +} + +pub fn subdivide(adj: &ConusAdjacency, plan: &ReachPlan) -> SubdividedAdjacency { + let n = adj.order.len(); + assert_eq!(plan.pieces.len(), n, "pieces must be one per parent reach"); + assert_eq!(plan.length_m.len(), n, "lengths must be one per parent reach"); + + let mut parent_offset = Vec::with_capacity(n + 1); + let mut acc: i32 = 0; + parent_offset.push(0); + for &m in &plan.pieces { + acc += m.max(1) as i32; + parent_offset.push(acc); + } + let n_sub = acc as usize; + + let mut order = Vec::with_capacity(n_sub); + let mut length_m = Vec::with_capacity(n_sub); + let mut slope = Vec::with_capacity(n_sub); + let mut rows = Vec::new(); + let mut cols = Vec::new(); + + for p in 0..n { + let m = plan.pieces[p].max(1) as usize; + let base = parent_offset[p] as usize; + // plan.length_m, NOT adj.length_m: short reaches were clamped UP in Task 2. + let piece_len = plan.length_m[p] / m as f32; + for k in 0..m { + order.push(adj.order[p]); + length_m.push(piece_len); + slope.push(adj.slope[p]); + if k > 0 { + // internal chain link: piece k-1 flows into piece k + cols.push((base + k - 1) as i32); + rows.push((base + k) as i32); + } + } + } + + // External edges: upstream parent's OUTLET -> downstream parent's INLET. + for (&r, &c) in adj.rows.iter().zip(adj.cols.iter()) { + // A self-edge (r == c) would map to outlet(p) -> inlet(p), i.e. + // off[p]+m-1 -> off[p], which is UPPER triangular and would break the + // forward-substitution invariant. The COO from `build_conus_adjacency` + // should carry only off-diagonal edges (the diagonal is synthesized by + // `CsrPattern::from_sparse`), so this is an assertion, not a filter. + assert!(r != c, "self-edge on parent {r}: subdivision cannot expand it"); + let up_outlet = parent_offset[c as usize + 1] - 1; + let down_inlet = parent_offset[r as usize]; + cols.push(up_outlet); + rows.push(down_inlet); + } + + SubdividedAdjacency { + order, + parent_order: adj.order.clone(), + parent_offset, + rows, + cols, + length_m, + slope, + } +} +``` + +- [ ] **Step 4: Verify** — `cargo test --test subdivide` → PASS (all 9 tests). + +- [ ] **Step 5: Commit** + +```bash +git add src/adjacency/subdivide.rs tests/subdivide.rs +git commit -m "feat(adjacency): graph expansion preserving topological order" +``` + +--- + +### Task 4: Persist to zarr and invalidate the cache + +**Files:** +- Modify: `src/data/store/zarr.rs`, `src/adjacency/cache.rs`, `src/adjacency/build.rs` +- Test: `tests/subdivide.rs` + +**Interfaces:** +- Consumes: `SubdividedAdjacency` from Task 3. +- Produces: two new zarr arrays `/parent_order` (i32, `[N]`) and `/parent_offset` + (i32, `[N+1]`); `ConusAdjacencyStore` gains `parent_order: Vec` and + `parent_offset: Vec`. + +**Critical:** `ConusAdjacencyStore.index: IdIndex` (`zarr.rs:29-99`) must be +built from `parent_order`, **not** `order` — after subdivision `order` contains +duplicates and a COMID→row lookup would be ambiguous. + +- [ ] **Step 1: Write the failing tests** + +```rust +#[test] +fn cache_key_changes_with_every_subdivision_field() { + use ddrs::adjacency::cache::content_key_for_test as key; + use ddrs::config::Subdivision; + let on = Subdivision { enabled: true, ..Default::default() }; + let base = key("fab", "gag", None, &Subdivision::default()); + assert_ne!(base, key("fab", "gag", None, &on), "enabling must invalidate"); + + // Every field that feeds `plan_reaches` must invalidate, or a config edit + // silently reuses a graph built with different geometry. + for (name, modified) in [ + ("max_pieces", Subdivision { max_pieces: 4, ..on.clone() }), + ("reference_n", Subdivision { reference_n: 0.03, ..on.clone() }), + ("q_coeff", Subdivision { reference_discharge_coefficient: 0.02, ..on.clone() }), + ("q_exp", Subdivision { reference_discharge_exponent: 0.8, ..on.clone() }), + ("min_len_fr", Subdivision { min_length_fraction: 0.5, ..on.clone() }), + ] { + assert_ne!(key("fab", "gag", None, &on), key("fab", "gag", None, &modified), + "changing {name} must invalidate the cache"); + } +} + +#[test] +fn store_index_maps_comid_to_parent_not_subreach() { + // Built with pieces [3,2,1]; COMID 200 is parent 1. + let store = round_trip_store(&chain3(), &[3, 2, 1]); + assert_eq!(store.parent_offset, vec![0, 3, 5, 6]); + assert_eq!(store.n, 6, "sub-reach count"); + assert_eq!(store.parent_order.len(), 3, "parent count"); + let p = store.index.get(&Comid(200)).expect("COMID 200 must resolve"); + assert_eq!(p, 1, "index must return the PARENT position, not a sub-reach row"); +} +``` + +Write `round_trip_store` to build → subdivide → write zarr to a `tempfile::TempDir` +→ reload via `ConusAdjacencyStore::open`. Reuse whatever writer +`src/adjacency/cache.rs` already calls (find it; do not invent a new one). + +- [ ] **Step 2: Run to verify it fails.** + +- [ ] **Step 3: Implement** + +Write both new arrays alongside the existing `/order`, `/length_m`, `/slope`, +`/indices_0`, `/indices_1`. In `ConusAdjacencyStore::open`, read them if present and +otherwise synthesize the identity (`parent_order = order`, +`parent_offset = 0..=n`) so **pre-existing stores keep loading unchanged**. + +In `src/adjacency/cache.rs`, extend `content_key` (`cache.rs:316-329`) to hash the +two subdivision fields after `BUILDER_VERSION`, and bump `BUILDER_VERSION` by 1: + +**Hash ALL SEVEN fields, not just `enabled` + `max_pieces`.** Every one of +`reference_n`, `reference_discharge_coefficient`, `reference_discharge_exponent`, +`min_length_fraction` and `max_clamp_factor` changes the reference celerity → +`dx_target` → both `pieces` and the clamped `length_m`, i.e. it changes the built +graph. Hashing only two of them would silently reuse a stale cached adjacency +after an edit to any of the others. + +> `max_clamp_factor` was added after this plan was first written (commit +> `2a06c0f`), so hash it too — the field list to hash is exactly the field list +> of `Subdivision`. + +```rust +// `content_key` takes `s: &Subdivision` as a new parameter (thread it from the +// caller's `cfg.params.subdivision`), so the test can vary one field at a time. +h.update(BUILDER_VERSION.to_le_bytes().as_ref()); +h.update(&[s.enabled as u8]); +h.update((s.max_pieces as u32).to_le_bytes().as_ref()); +// f32::to_bits gives a stable byte pattern; NaN is impossible here (validated). +h.update(s.reference_n.to_bits().to_le_bytes().as_ref()); +h.update(s.reference_discharge_coefficient.to_bits().to_le_bytes().as_ref()); +h.update(s.reference_discharge_exponent.to_bits().to_le_bytes().as_ref()); +h.update(s.min_length_fraction.to_bits().to_le_bytes().as_ref()); +h.update(s.max_clamp_factor.to_bits().to_le_bytes().as_ref()); +``` + +Expose `content_key_for_test` as `#[doc(hidden)] pub` so the test can call it. + +Call `plan_reaches` + `subdivide` after `build_conus_adjacency` and **before** +gauge-subgraph construction, so subgraphs are cut from the expanded graph. The +sequencing lives in `cache.rs::resolve_or_build`, not `build.rs`. + +> **`catchsize` is NOT drainage area.** It is the *local* divide area — median +> 36.7 km², max ~612 km² even for continental rivers — whereas +> `reference_celerity` wants upstream accumulated area `uparea_km2`. Passing it +> raw under-feeds `Q_ref` by 3–4 orders of magnitude on large rivers. Accumulate +> it downstream over the topological order +> (`cache.rs::upstream_area_km2`). Validated: the accumulated value reproduces +> the fabric's own `10^log10_uparea` at ratio p5/p50/p95 = 1.000/1.000/1.000 over +> all 346,321 CONUS reaches. `log10_uparea` cannot be used directly as the +> source — it is NaN on 88 % of `merit_global_attributes_v2.nc` (finite only over +> CONUS). + +- [ ] **Step 4: Verify** — `cargo test --test subdivide` → PASS. + +- [ ] **Step 5: Commit** + +```bash +git add src/data/store/zarr.rs src/adjacency/cache.rs src/adjacency/build.rs tests/subdivide.rs +git commit -m "feat(adjacency): persist parent map, invalidate cache on subdivision" +``` + +--- + +### Task 5: Distribute q′ across pieces + +**Files:** +- Modify: `src/routing/mmc.rs` +- Test: `tests/subdivision_integration.rs` (create) + +**Interfaces:** +- Consumes: `parent_offset` via the adjacency inputs reaching `setup_inputs`. +- Produces: `q_prime` divided by each row's parent piece count. + +**Placement is exact:** after the clamp at `src/routing/mmc.rs:453` +(`let q_prime_clamped = q_prime.clamp_min(discharge_lb);`) and **before** the first +slice at `mmc.rs:495`. Dividing before the clamp would let the floor `1e-4` be +applied to an undivided value and silently create mass. + +- [ ] **Step 1: Write the failing test** + +```rust +/// A 1-reach network with constant q' must reach the SAME steady-state outflow +/// whether or not it is subdivided — the pieces split the inflow m ways but +/// chain in series, so the outlet still carries the whole reach's runoff. +#[test] +fn subdivision_conserves_mass_at_steady_state() { + let un_split = steady_state_outflow(/*pieces*/ 1); + let split = steady_state_outflow(/*pieces*/ 4); + assert!((split - un_split).abs() / un_split < 1e-3, + "mass not conserved: 1 piece = {un_split}, 4 pieces = {split}"); +} +``` + +- [ ] **Step 2: Run to verify it fails** — the 4-piece case returns ~4× the correct + outflow, because each piece receives the full `q'`. + +- [ ] **Step 3: Implement** + +Build a per-row divisor tensor once in `setup_inputs` (beside `length`/`slope` at +`mmc.rs:204-213`) and store it as `self.pieces_per_row: Option, 1>>`: + +```rust +// One entry per sub-reach row: the piece count of its parent. Used to split +// lateral inflow, mirroring HEC-HMS's C4*(q_L*dx) with q_L per unit length. +let mut divisor = Vec::with_capacity(n); +for p in 0..inputs.adjacency.parent_offset.len() - 1 { + let m = inputs.adjacency.parent_offset[p + 1] - inputs.adjacency.parent_offset[p]; + for _ in 0..m { divisor.push(m as f32); } +} +self.pieces_per_row = Some(Tensor::from_floats(divisor.as_slice(), &self.device)); +``` + +Then at `mmc.rs:453`: + +```rust +let q_prime_clamped = q_prime.clamp_min(discharge_lb); +// Split lateral inflow evenly along the reach. MUST be after the clamp: the +// 1e-4 floor applied to an undivided value would create mass. +let q_prime_clamped = match self.pieces_per_row.as_ref() { + Some(d) => q_prime_clamped / d.clone().unsqueeze_dim::<2>(0), + None => q_prime_clamped, +}; +``` + +- [ ] **Step 4: Verify** — `cargo test --test subdivision_integration` → PASS. + +- [ ] **Step 5: Commit** + +```bash +git add src/routing/mmc.rs tests/subdivision_integration.rs +git commit -m "feat(routing): split lateral inflow across sub-reaches" +``` + +--- + +### Task 6: Gather KAN parameters onto sub-reaches + +**Files:** +- Modify: `src/training/forward.rs` +- Test: `tests/subdivision_integration.rs` + +**Interfaces:** +- Consumes: KAN outputs `HashMap>` shaped `[N_parent]` + (`src/nn/kan_head.rs:215-244`). +- Produces: the same map gathered to `[N_sub]`. + +**Run the KAN at parent resolution and gather** — do not duplicate attribute rows. +Same result, 4.77× less KAN compute, and `select`'s backward is a scatter-add, so +each parent correctly receives the summed gradient from all its pieces. + +- [ ] **Step 1: Write the failing test** + +```rust +#[test] +fn every_piece_inherits_its_parents_parameters() { + // parent_offset [0,3,5,6]; parent params [0.02, 0.05, 0.10] + let gathered = gather_for_test(&[0.02, 0.05, 0.10], &[0, 3, 5, 6]); + assert_eq!(gathered, vec![0.02, 0.02, 0.02, 0.05, 0.05, 0.10]); +} + +#[test] +fn gradient_sums_back_to_the_parent() { + // d(sum of gathered)/d(parent p) must equal that parent's piece count. + let g = gather_grad_for_test(&[0.02, 0.05, 0.10], &[0, 3, 5, 6]); + assert_eq!(g, vec![3.0, 2.0, 1.0], "scatter-add must sum piece gradients"); +} +``` + +- [ ] **Step 2: Run to verify it fails.** + +- [ ] **Step 3: Implement** + +In `src/training/forward.rs` (after `head.forward` at `forward.rs:235-239`, +before `denormalize`), build a row→parent index once and gather: + +```rust +// Sub-reaches share their parent's hydraulics: MERIT carries no within-reach +// variation. `select` backward is a scatter-add, so each parent receives the +// summed gradient of all its pieces. +let n_param = n_param.select(0, parent_idx.clone()); +let q_param = q_param.select(0, parent_idx.clone()); +let p_param = p_param.map(|t| t.select(0, parent_idx.clone())); +``` + +`parent_idx: Tensor` is built from `parent_offset` and carried on +`RoutingTensors`. When subdivision is off it is the identity, so make the whole +block conditional on `parent_offset.len() - 1 != n_rows` to keep the disabled path +byte-identical. + +- [ ] **Step 4: Verify** — `cargo test --test subdivision_integration` → PASS. + +- [ ] **Step 5: Commit** + +```bash +git add src/training/forward.rs tests/subdivision_integration.rs +git commit -m "feat(training): gather KAN parameters from parents to sub-reaches" +``` + +--- + +### Task 7: Map gauges to their parent's outlet piece + +**Files:** +- Modify: `src/data/collate.rs` +- Test: `tests/gauge_mass_conservation.rs` + +**Interfaces:** +- Consumes: `compress(unioned, conus_order, ddr_match) -> Result` + (`src/data/collate.rs:90-207`), `outflow_idx` built at `collate.rs:178-198`. +- Produces: for `ddr_match: false`, `outflow_idx = vec![outlet_row_of(gauge_parent)]`. + +A gauge measures everything upstream. With subdivision the parent's whole runoff +arrives at its **last** piece, so that is the row to read. Reading any earlier piece +would drop the downstream fraction of the reach's own lateral inflow — the same +class of bug as `2fe6bee`. + +- [ ] **Step 1: Write the failing test** + +Add to `tests/gauge_mass_conservation.rs`, reusing its existing 3-reach fixture +(two headwaters into a gauge reach, constant `Q_PRIME = 10.0`, steady state): + +```rust +#[test] +fn gauge_conserves_mass_when_its_reach_is_subdivided() { + // Same topology as `gauge_prediction_conserves_mass_when_not_ddr_match`, + // but the gauge's own reach is split 4 ways. The answer must not change. + let un_split = gauge_steady_state(/*pieces*/ 1); + let split = gauge_steady_state(/*pieces*/ 4); + assert!((un_split - 30.0).abs() < 1e-3, "control changed: {un_split}"); + assert!((split - 30.0).abs() < 1e-3, + "subdivided gauge lost mass: got {split}, expected 30.0"); +} +``` + +- [ ] **Step 2: Run to verify it fails** — expect ~27.5 (the last piece receives + only 1/4 of the gauge reach's own lateral inflow if the outlet is chosen wrongly). + +- [ ] **Step 3: Implement** + +In `collate.rs`, thread `parent_offset` into `compress` and replace the +`ddr_match: false` branch (`collate.rs:196-198`): + +```rust +// The gauge's whole reach discharges at its LAST piece. Any earlier piece +// omits the downstream part of the reach's own lateral inflow. +gauge_compressed.iter() + .map(|&g_parent| vec![parent_offset[g_parent + 1] as usize - 1]) + .collect() +``` + +Leave the `ddr_match: true` branch untouched — it reproduces DDR and is pinned by +`ddr_match_gauge_prediction_omits_the_gauge_reach`. + +- [ ] **Step 4: Verify** + +```bash +cargo test --test gauge_mass_conservation +``` + +- [ ] **Step 5: Commit** + +```bash +git add src/data/collate.rs tests/gauge_mass_conservation.rs +git commit -m "feat(collate): read gauges at their parent's outlet piece" +``` + +--- + +### Task 8: End-to-end validation on real CONUS data + +**Files:** +- Modify: `src/bin/probe_courant.rs`, `tests/subdivision_integration.rs` + +**Interfaces:** +- Consumes: everything above. +- Produces: measured Cr / X / negative-solve counts with subdivision on vs off. + +This is the task that decides whether the change is worth keeping. `probe_courant` +already reports Cr, X, and exact negative-solve counts on the real network — reuse +it rather than writing a new probe. + +- [ ] **Step 1: Write the failing test** + +```rust +#[test] +fn subdivision_makes_coefficients_non_negative_without_the_clamp() { + // enforce_positivity OFF, subdivision ON: the whole point is that Cr ~ 1 + // makes c1, c3 >= 0 by construction. + let r = run_probe(Subdiv::On { max_pieces: 8 }, EnforcePositivity::Off); + assert!(r.frac_c1_negative < 0.01, + "c1 negative on {:.2}% of cells", 100.0 * r.frac_c1_negative); + assert!(r.frac_c3_negative < 0.01, + "c3 negative on {:.2}% of cells", 100.0 * r.frac_c3_negative); + assert!(r.median_cr > 0.8 && r.median_cr < 2.0, + "median Cr = {} is not near 1", r.median_cr); +} + +#[test] +fn subdivision_off_is_bit_identical_to_today() { + let a = run_probe(Subdiv::Off, EnforcePositivity::Off); + assert_eq!(a.x_sol_bits, GOLDEN_X_SOL_BITS_NO_SUBDIV); +} +``` + +- [ ] **Step 2: Run to verify it fails.** + +- [ ] **Step 3: Add a `--max-pieces` flag to `probe_courant`** and report, for + subdivision off vs on: sub-reach count, median Cr, `frac Cr > 2`, X percentiles, + `frac c1 < 0`, `frac c3 < 0`, and exact negative-solve counts. + +- [ ] **Step 4: Measure on the real network and record the numbers** + +```bash +cargo run --release --bin probe_courant -- \ + --config ddrs.yaml --checkpoint /checkpoints/ \ + --backend cuda --gauges 1841 --rho 90 --steps 2136 --max-pieces 8 +``` + +**Also report the length-clamp cost**, which is the price of the short-reach branch: +fraction of reaches clamped, the clamp-factor distribution (p50/p95/p99/max), and +total network length inflation. Indicative figures were 34.4 % clamped, median 2.0×, +max 48,597×, +17.1 % total length — but those used a crude reference celerity and +**must be recomputed with the calibrated one**. + +**Report honestly.** Subdivision alone cannot fix reaches that are already +`Cr > 2` (17.3 % of the network); only the length clamp reaches them. If +`min_length_fraction` is reduced or disabled, that population returns. Record the +actual residual `frac c1 < 0` / `frac c3 < 0` — do **not** claim a guarantee the +measurement does not support. + +- [ ] **Step 4b: Decide the hotstart question (carried from Task 5)** + +`setup_inputs` cold-starts with `(I − N)·Q₀ = q'₀` using the **undivided** `q_prime` +row 0 (`mmc.rs:~300`). Task 5 scoped the division to `forward` only, so under +subdivision the initial condition is inflated roughly `m×` per reach. Steady state +is unaffected (measured: 10.000000 vs 10.000001 m³/s), but the spin-up is longer. + +Measure it: run the probe with subdivision on, and compare the first ~200 timesteps +against a long-run steady state to see how many steps the inflated Q₀ takes to wash +out. If it exceeds the configured `warmup`, divide the hotstart too — ~3 lines, +touching only the subdivided path. Record the decision either way. + +- [ ] **Step 5: Run the full gate set** + +```bash +cargo test --test subdivide --test subdivision_integration \ + --test gauge_mass_conservation --test adjacency_parity \ + --test positivity_clamp --test cunge_x --test celerity_beta +cargo test --lib +cargo run --release --example compare_ddr_sandbox # must stay ABSOLUTE MATCH +``` + +`tests/adjacency_parity.rs` asserts element-for-element equality against an +engine-built store over all 346,321 positions. With subdivision **off** it must +still pass unchanged; if it fails, the disabled path is not a true no-op. + +- [ ] **Step 6: Commit** + +```bash +git add src/bin/probe_courant.rs tests/subdivision_integration.rs +git commit -m "test(subdivision): real-CONUS Courant and coefficient validation" +``` + +--- + +### Task 9: Document + +**Files:** +- Create: `.claude/REACH-SUBDIVISION.md` +- Modify: `.claude/PHYSICS-CORRECTIONS.md`, `CLAUDE.md` + +- [ ] **Step 1: Write `.claude/REACH-SUBDIVISION.md`** with the ASCII topology + diagram from this plan's §Data model, the cap-sweep table, the measured + before/after numbers from Task 8, and the two-index-space explanation. + +- [ ] **Step 2: Correct `.claude/PHYSICS-CORRECTIONS.md`.** Its + §"Why sub-stepping still does not substitute" currently claims subdivision + "restores dynamic range to the Cunge X". **That is false** — raw X median moves + only 0.4973 → 0.4815. Replace with the two real justifications (automatic + non-negativity at Cr≈1; recovery of `X_eff` if `enforce_positivity` is kept) and + mark it as an erratum, matching the erratum style already in that file. + +- [ ] **Step 3: Add the `params.subdivision` block to `CLAUDE.md`'s config notes**, + including the cap rationale and the retraining requirement. + +- [ ] **Step 4: Commit** + +```bash +git add .claude/REACH-SUBDIVISION.md .claude/PHYSICS-CORRECTIONS.md CLAUDE.md +git commit -m "docs: reach subdivision design, cap rationale, X erratum" +``` + +--- + +## Verification criteria + +1. `params.subdivision.enabled: false` is byte-identical to current output — + proven by `adjacency_parity` and the `GOLDEN_X_SOL_BITS_NO_SUBDIV` test. +2. `compare_ddr_sandbox` still reports ABSOLUTE MATCH. +3. Total `q'` is conserved per parent reach; total length is conserved for + *split* reaches and increased only for *clamped* ones, by a measured amount. +4. Expanded graph is topologically ordered and strictly lower-triangular. +5. Gauge predictions are unchanged by subdividing the gauge's reach. +6. With subdivision on and `enforce_positivity` **off**, `frac c1 < 0` and + `frac c3 < 0` both fall below 1 % on the real network, and median Cr ∈ [0.8, 2.0]. +7. `frac Cr > 2` and `frac Cr < 0.5` are both measured on the real network and + reported. The indicative target is ~0 % and ~0.05 % respectively; any material + residual must be stated, not hidden. +8. The length-clamp distortion (fraction clamped, clamp-factor p95/max, total + network length inflation) is measured and reported alongside the benefit. diff --git a/docs/superpowers/specs/2026-08-05-per-gauge-tau-sweep-design.md b/docs/superpowers/specs/2026-08-05-per-gauge-tau-sweep-design.md new file mode 100644 index 0000000..e498c7d --- /dev/null +++ b/docs/superpowers/specs/2026-08-05-per-gauge-tau-sweep-design.md @@ -0,0 +1,109 @@ +# Per-gauge tau sweep — design + +- **Date:** 2026-08-05 +- **Motivation:** area-balanced eval (run `2026-08-05T04-58-58Z-conus-experimental-train-and-test`, + epoch 30, 1,841 gauges) shows ddrs LOSES to the summed-Q' baseline in small + basins: median NSE 0.630 vs 0.720 in the <1,000 km² bin (841 gauges), while + winning above 5,000 km². Hypothesis: a sub-daily timing/alignment error that + aliases across the daily pooling boundary, worst for flashy small basins. +- **The suspect:** `params.tau` (default 3) — a single CONUS-wide phase offset in + `tau_trim_and_downsample` (`src/training/loss.rs:25-46`, slice + `[13+tau : -11+tau]` then daily area-pool). Inherited unvetted from DDR + (`configs.py:116`: "handle double routing and timezone differences"). USGS + daily values are local-day means; CONUS spans four time zones; one scalar + cannot be right everywhere. With `tau = 3` the daily bin boundary sits at + 16:00 UTC, while CONUS local midnight is 05:00–08:00 UTC. + +## Hypothesis (pre-registered) + +**H-TAU:** the small-basin NSE deficit is substantially a pooling-phase +(timing) error: per-gauge choice of tau recovers most of the gap. + +Verdict states: SUPPORTED / REFUTED / INCONCLUSIVE only. + +## Method + +Two phases. Phase 1 (pilot, tonight) validates the method on WY1996; Phase 2 +(full sweep, after tuning) runs the same analysis on the full eval window with +the pre-registered gate. + +### Instrument + +1. **One eval run, full window** (1995/10/01–2010/09/30), legacy `eval` binary, + checkpoint `epoch_30_mb_1` of the run above, its own `config.yaml` snapshot + (tau = 3), with `DDRS_HOURLY_DUMP=` — the opt-in diagnostic in + `src/training/eval.rs` that writes the PRE-TRIM hourly series + (n_gauges × n_hours raw f32 + `.json` dims sidecar, ~1 GB). No retraining. +2. **Offline sweep** (`scripts/tau_sweep.py`, NumPy only): for each + `tau ∈ {0..23}` reconstruct daily predictions as block means of hours + `[13+tau+24i, 13+tau+24(i+1))`, scored against obs day `i+1` (the day-0 drop + convention). `tau = 3` reproduces shipped behavior exactly; 0..23 covers the + full 24-hour phase cycle. The trimmed window is divisible by 24, so block + mean equals the area pool. +3. Per gauge and per tau: NSE against the eval zarr's `observations` (NaN + masked, ≥100 valid days required). Pilot restricts scored days to WY1996 + (1995/10/01–1996/09/30). +4. Baseline comparison from the run's own `baseline/` arrays (identical + 1,841-gauge population), sliced to the same days. +5. Covariates joined from `gages_2000_area_balanced.csv`: `DRAIN_SQKM`, + `LNG_GAGE`. + +### Method-validation gate (pilot, Phase 1) + +Boolean, all three must hold: + +1. Offline `tau = 3` daily reconstruction matches the eval run's own + `predictions.zarr` to < 1e-3 relative (f32 floor). +2. Recomputed full-window median NSE at `tau = 3` matches the eval's reported + median to < 0.001. +3. Sweep produces per-gauge NSE(tau) curves for ≥ 95% of the 1,841 gauges. + +If this gate fails, fix the method before interpreting any number. + +### Decision gate (full sweep, Phase 2 — user-selected) + +**GO for building per-gauge local-time pooling iff per-gauge-best tau lifts the +<1,000 km² median NSE from 0.630 to ≥ 0.720 (the summed-Q' baseline).** +Beating the baseline is the strictest of the considered bars; a large-but- +insufficient gain reads as "timing real but not the whole story" (INCONCLUSIVE +for H-TAU as the *dominant* cause). + +Secondary (diagnostic, not gating): Spearman correlation of per-gauge best tau +with gauge longitude — the timezone-mechanism fingerprint; NSE(tau) medians +stratified by the five drainage-area bins. + +### Known bias, accepted for the pilot + +Choosing tau per gauge on the same days it is scored is in-sample (1 free +integer parameter per gauge); the pilot number is a headroom CEILING. The +Phase 2 protocol (split-sample selection vs scoring) is deliberately left open +for tomorrow's tuning session. + +## Artifacts + +| Path | Content | +|---|---| +| `output/tau_sweep/eval_full.zarr` | fresh eval daily predictions + obs | +| `output/tau_sweep/hourly_full.f32` (+`.json`) | pre-trim hourly dump | +| `output/tau_sweep/nse_by_tau_wy1996.csv` | gauge × tau NSE matrix (pilot) | +| `output/tau_sweep/best_tau_wy1996.csv` | per-gauge best tau + covariates | +| `output/tau_sweep/summary_wy1996.md` | gate numbers, bin medians, correlation | +| `output/tau_sweep/*.png` | NSE-vs-tau by area bin, best-tau histogram, best-tau vs longitude | + +## Concerns / assumptions + +- **Assumption:** the eval zarr's `observations` rows are aligned with + `predictions` rows (probe convention: pred bin i ↔ obs day i+1). Validation + gate item 2 tests this indirectly. +- **Could go wrong:** rerunning eval from the same checkpoint may not bit-match + the 04-58-58Z run (GPU nondeterminism); the method gate therefore compares + against the NEW run's own zarr, not the old one. +- **Could go wrong:** WY1996 may be a hydrologically unusual year; pilot numbers + are for method-shakeout and tuning, not conclusions. +- **Binary provenance:** the dump code is uncommitted at design time; it is + committed with this spec, and the binary is built from that commit + (2026-07-01 stale-binary lesson). +- **Why not inverse routing:** recovering hourly timing from daily gauges is + ill-posed both spatially (a gauge observes a network sum — the leakance + NO-GO mechanism) and temporally (24 unknowns/day vs 1 obs/day). The + deterministic per-gauge pooling window is the cheap falsifiable step first. diff --git a/docs/superpowers/specs/2026-08-06-ddr-equifinality-paper-scope-design.md b/docs/superpowers/specs/2026-08-06-ddr-equifinality-paper-scope-design.md new file mode 100644 index 0000000..a666d3c --- /dev/null +++ b/docs/superpowers/specs/2026-08-06-ddr-equifinality-paper-scope-design.md @@ -0,0 +1,143 @@ +# ddr_equifinality Paper Scope — Design + +**Date:** 2026-08-06 +**Status:** Approved (brainstorming), pending user review of this spec +**Paper:** `/home/tbindas/papers/ddr_equifinality/paper.tex` +**Submitted AGU H069 abstract:** anchors this scope; already spliced into `paper.tex` (title, abstract, authors, Beven 2001 citation). + +## Purpose + +Scope the reframed `ddr_equifinality` paper so that a fresh writer knows what the +paper claims, what to keep from the existing draft, what to rewrite, and which +experiments and fixes must land before the Results can be written. This spec +covers **paper scope**; each prerequisite experiment gets its own findings doc in +`docs/` when it runs. + +## 1. Contribution and framing + +The paper is an **empirical answer to Beven (2001), "How far can we go in +distributed hydrological modelling?"**: can a distributed, differentiable routing +model overcome the *uncertainty* and *equifinality* obstacles Beven raised, using +tools he lacked (differentiable modeling, cloud-scale multi-source datasets)? + +The payload is **normative**, engaging Beven's closing *Alternative Blueprint* +(future-proof modelling that gradually learns the idiosyncrasies of individual +places under acknowledged uncertainty): the answer becomes **criteria for judging +future differentiable models** — which learned parameters, and which places, can +be trusted versus which remain equifinal. + +This is **not** a methods/protocol contribution. The input-perturbation design is +how the question is answered, not the novelty. + +## 2. The two questions (open, non-directional) + +Framed as open questions, not directional predictions. The paper does **not** bet +on the old "Manning's roughness is a bias absorber" mechanism (that direction is +unsupported in the group's own prior analysis). + +- **Q-uncertainty:** Does trained routing *skill* converge across structurally + different inflow sources — can the model absorb source-specific input bias? +- **Q-equifinality:** Do learned *channel parameters* (and their loss landscapes) + settle on one consistent set across sources, or do some remain equifinal — which + ones, and where? + +## 3. Experiment matrix (full two-axis) + +Same MERIT CONUS network, catchment attributes, USGS observations, and training +budget across all arms; swap **only** the lateral-inflow source. Two structural +axes: + +- **Spatial axis:** lumped vs distributed dHBV runoff. +- **Temporal axis:** daily vs hourly LSTM forecasts (daily includes the + precip-driven disaggregation variant as the bridge). + +Parameter-convergence + skill analysis runs on **every** arm. Current state: the +temporal (LSTM) arms have full parameter analysis; the spatial (dHBV lumped/ +distributed) arms have **routing-skill numbers only** — their parameter analysis +is prerequisite B2 below. + +## 4. Analysis levels + +| Level | What it measures | Answers | +|---|---|---| +| L1 routing skill | per-arm median NSE/KGE vs summed-inflow baseline | Q-uncertainty (does skill converge / bias get absorbed) | +| L2 raw parameters | cross-arm spread and correlation of learned n, p, q | Q-equifinality | +| L3 realized geometry | depth, top width, hydraulic radius at reference discharge, cross-arm | Q-equifinality (the physically interpretable quantities) | +| L4 loss-landscape / identifiability | which parameters are stiff vs sloppy around the trained optimum | Q-equifinality | + +Every convergence number is reported against a **replicate-seed noise floor** +(prerequisite B3): a spread smaller than the seed-to-seed spread is not +convergence. + +## 5. Keep vs rewrite in the existing draft + +**Keep, lightly revise:** +- **Introduction** — repoint the thesis paragraph from "selective equifinality / + bias absorber" to Beven's five challenges and the two open questions. The + opening (routing, ungauged reaches, lookup-table parameters, National Water + Model) stands. +- **Methods** — the differentiable Muskingum-Cunge routing math, the lateral + inflow sources, the data, and the experimental-design subsections are reusable + as-is. Update the inflow-source table to the final two-axis arm list. Replace + the "Cross-Arm Convergence Analysis" framing to match the L1-L4 levels above + and the noise-floor requirement. + +**Rewrite:** +- **Results** — currently a skeleton. Populate with L1-L4 once the arms and the + prerequisites land. +- **Discussion** — drop the bias-absorber "division of labor" narrative. Frame + around the two questions and the judgment-criteria payload (what to trust, + transfer, and interpret physically; which places). +- **Conclusion** — replace the "selective equifinality confirmed" arc with the + honest answer to Beven plus the evaluation-criteria contribution. + +## 6. Prerequisites that gate the Results + +Intro and Methods can be written now. Results cannot be written until: + +- **B1 — timezone / day-boundary alignment (THE GATE, do first).** USGS daily + observations are midnight-to-midnight in **local standard time**; AORC forcing + and likely the Q' stores are **UTC**. That is a 5-8 h, longitude-dependent, + per-gauge offset. Verify how `src/data/dates.rs` and `src/data/dataset.rs` build + daily windows and index each store. If the stores disagree on day convention, + cross-source differences are confounded and every convergence number is suspect. + Fix (reconcile all sources to one boundary) or, if already reconciled, document + it as a common-mode caveat. Handoff: `/tmp/handoff-aorc-usgs-recording-times.md`. +- **B2 — spatial-axis parameter analysis.** Run L2-L4 on the lumped and + distributed dHBV arms (checkpoints/parameter dumps exist; the analysis does + not). +- **B3 — replicate seed.** At least one additional seed per arm, to put a noise + floor under every convergence statistic. The group's evidence standard treats + this as essential, not optional. +- **B4 — double-routing confound.** Resolve whether the distributed dHBV store is + already pre-routed / pre-smoothed (open question from the AORC2F wave-1 + findings). If it is, lumped-vs-distributed differences are store-provenance + artifacts, not model-structure effects. Fastest resolution: inspect the export + script that generated the distributed store for whether it exports per-unit + runoff or a routed/aggregated product. + +## 7. Concerns and assumptions + +- **The honest answer may be "partially."** The defensible result could be "skill + converges, so the model largely absorbs input bias (uncertainty answered), but + channel parameters only partly converge, and here is the noise floor + (equifinality persists for some parameters/places)." The paper structure must + make that a *satisfying answer to Beven*, feeding the judgment-criteria + contribution, not read as a null result. **Assumption:** a nuanced, honest + answer is publishable precisely because it is honest and it yields evaluation + criteria. +- **Compute is substantial** — full two-axis matrix + replicate seed + reruns + after the B1 fix. This is the schedule risk against the December AGU talk. +- **Risk:** if B4 cannot be resolved, the spatial axis weakens and the paper leans + on the temporal axis, partly undercutting the full-matrix promise. Mitigation: + the timezone gate and the parameter analysis are worth doing regardless; decide + spatial-axis weight after B4. +- **Assumption:** authors are Bindas, Song, Shen (as now in `paper.tex`); author + order and affiliations confirmed separately. + +## 8. Out of scope + +- Global scale-out (the global MERIT fabric) — a separate effort. +- Leakance / water-loss terms — closed NO-GO, ζ=0 in every arm. +- Any new routing-core or KAN-head development — the model is fixed; this paper + measures what it learns. diff --git a/scripts/build_gages_2000_area_balanced.py b/scripts/build_gages_2000_area_balanced.py new file mode 100644 index 0000000..fc7ff6d --- /dev/null +++ b/scripts/build_gages_2000_area_balanced.py @@ -0,0 +1,144 @@ +"""Build an area-balanced ~2,000-gauge CSV from GAGES-II. + +Follow-up to /tmp/experiment-handoff-small-basin-domination.md (2026-08-02): +the gages_3000 population is 87.5% <5,000 km2 and its DA_VALID column is an +absolute-difference criterion that preferentially deletes large basins. + +This script: + (a) starts from GAGES-II.csv (8,931 gauges, same schema as gages_3000.csv); + (b) recomputes DA_VALID as ABS_DIFF / DRAIN_SQKM <= 0.10 (relative); + (c) keeps only gauges present in the USGS observations icechunk store with + >= COVERAGE_MIN non-NaN daily coverage in BOTH the configured training + window (1981-10-01..1995-09-30) and eval window (1995-10-01..2010-09-30), + and present in merit_gages_conus_adjacency.zarr as a non-headwater + subgraph (order length > 1) -- mirroring the dataset.rs filter so the + CSV population equals the population training actually uses; + (d) subsamples to ~2,000: ALL basins >= 5,000 km2, topped up from + [1,000, 5,000) to 1,000 gauges >= 1,000 km2, plus a random 1,000 from + < 1,000 km2 -- ~50/50 either side of 1,000 km2. Small basins are kept + deliberately (they teach the identity-routing regime); the goal is to + reduce their gradient share, not remove them. + +Run under DDR's venv: + cd ~/projects/ddr && uv run python ~/projects/ddrs/scripts/build_gages_2000_area_balanced.py +""" + +import numpy as np +import pandas as pd +import xarray as xr +import zarr +import icechunk + +GAGES_II = "/home/tbindas/projects/ddr/references/gage_info/GAGES-II.csv" +OBS_STORE = "/mnt/ssd1/data/icechunk/usgs_daily_observations" +ADJ_STORE = "/home/tbindas/projects/ddr/data/merit_gages_conus_adjacency.zarr" +OUT_CSV = "/home/tbindas/projects/ddr/references/gage_info/gages_2000_area_balanced.csv" + +REL_TOL = 0.10 +COVERAGE_MIN = 0.80 +TRAIN_WINDOW = ("1981-10-01", "1995-09-30") # config/merit_training.yaml experiment window +EVAL_WINDOW = ("1995-10-01", "2010-09-30") # testing window +SEED = 42 +N_SMALL = 1000 # target below 1,000 km2 +N_LARGE = 1000 # target at/above 1,000 km2 (all >=5,000 kept, rest from [1k,5k)) + +# ---- (a) + (b): load GAGES-II, recompute DA_VALID relatively ------------- +df = pd.read_csv(GAGES_II, dtype={"STAID": str}) +df["STAID"] = df["STAID"].str.zfill(8) +assert not df["STAID"].duplicated().any(), "duplicate STAIDs in GAGES-II.csv" +n_total = len(df) + +df["REL_DIFF"] = df["ABS_DIFF"] / df["DRAIN_SQKM"] +old_valid = df["DA_VALID"].astype(str).str.lower().eq("true") +df["DA_VALID"] = df["REL_DIFF"] <= REL_TOL +print(f"GAGES-II rows: {n_total}") +print(f"DA_VALID old (absolute criterion): {old_valid.sum()}") +print(f"DA_VALID new (rel <= {REL_TOL:.0%}): {df['DA_VALID'].sum()}" + f" (recovered {(df['DA_VALID'] & ~old_valid).sum()}," + f" dropped {(~df['DA_VALID'] & old_valid).sum()})") + +pool = df[df["DA_VALID"]].copy() + +# ---- (c) part 1: observation coverage in both windows -------------------- +storage = icechunk.local_filesystem_storage(OBS_STORE) +repo = icechunk.Repository.open(storage) +sess = repo.readonly_session("main") +ds = xr.open_zarr(sess.store, consolidated=False) + +in_obs = pool["STAID"].isin(set(ds["gage_id"].values.astype(str))) +print(f"after DA_VALID, in observation store: {in_obs.sum()} / {len(pool)}") +pool = pool[in_obs].copy() + +sf = ds["streamflow"].sel(gage_id=pool["STAID"].values) +frac_train = sf.sel(time=slice(*TRAIN_WINDOW)).notnull().mean("time").compute() +frac_eval = sf.sel(time=slice(*EVAL_WINDOW)).notnull().mean("time").compute() +pool["COV_TRAIN"] = frac_train.values +pool["COV_EVAL"] = frac_eval.values + +print("\ncoverage sensitivity (gauges meeting the bar in BOTH windows):") +for thr in (0.0, 0.5, 0.8, 0.9, 0.95, 1.0): + n = ((pool["COV_TRAIN"] > thr) & (pool["COV_EVAL"] > thr)).sum() if thr == 0.0 else \ + ((pool["COV_TRAIN"] >= thr) & (pool["COV_EVAL"] >= thr)).sum() + print(f" >= {thr:>4.0%}: {n}") + +cov_ok = (pool["COV_TRAIN"] >= COVERAGE_MIN) & (pool["COV_EVAL"] >= COVERAGE_MIN) +print(f"applying bar {COVERAGE_MIN:.0%}: {cov_ok.sum()} / {len(pool)}") +pool = pool[cov_ok].copy() + +# ---- (c) part 2: non-headwater subgraph exists --------------------------- +adj = zarr.open_group(ADJ_STORE, mode="r") +keep = [] +n_missing = n_headwater = 0 +for staid in pool["STAID"]: + try: + sub = adj[staid] + except KeyError: + n_missing += 1 + keep.append(False) + continue + if sub["order"].shape[0] <= 1: # GageSubgraph::is_headwater (zarr.rs:128) + n_headwater += 1 + keep.append(False) + else: + keep.append(True) +print(f"adjacency filter: kept {sum(keep)}" + f" (dropped {n_missing} missing, {n_headwater} headwater)") +pool = pool[keep].copy() + +# ---- (d): area-balanced subsample ---------------------------------------- +rng = np.random.default_rng(SEED) +small = pool[pool["DRAIN_SQKM"] < 1000] +mid = pool[(pool["DRAIN_SQKM"] >= 1000) & (pool["DRAIN_SQKM"] < 5000)] +big = pool[pool["DRAIN_SQKM"] >= 5000] +print(f"\neligible pool: {len(pool)} (<1k: {len(small)}," + f" 1k-5k: {len(mid)}, >=5k: {len(big)})") + +n_mid = min(len(mid), max(0, N_LARGE - len(big))) +n_small = min(len(small), N_SMALL) +sel = pd.concat([ + big, + mid.iloc[rng.permutation(len(mid))[:n_mid]], + small.iloc[rng.permutation(len(small))[:n_small]], +]).sort_values("STAID") + +cols = ["STAID", "STANAME", "DRAIN_SQKM", "LAT_GAGE", "LNG_GAGE", "COMID", + "COMID_DRAIN_SQKM", "COMID_UNITAREA_SQKM", "ABS_DIFF", "DA_VALID", + "FLOW_SCALE", "REL_DIFF", "COV_TRAIN", "COV_EVAL"] +out = sel[cols].copy() +out["DA_VALID"] = "True" +out.to_csv(OUT_CSV, index=False) +print(f"\nwrote {len(out)} gauges -> {OUT_CSV}") + +# ---- report -------------------------------------------------------------- +bins = [0, 100, 500, 1000, 5000, 10000, 30000, np.inf] +labels = ["<100", "100-500", "500-1k", "1k-5k", "5k-10k", "10k-30k", ">=30k"] +hist = pd.cut(out["DRAIN_SQKM"], bins=bins, labels=labels).value_counts().reindex(labels) +print("\ndrainage-area histogram (km2):") +for lab, n in hist.items(): + print(f" {lab:>8}: {n:4d} ({n / len(out):.1%})") +n_below = (out["DRAIN_SQKM"] < 1000).sum() +print(f"\n<1,000 km2: {n_below} ({n_below/len(out):.1%})" + f" >=1,000 km2: {len(out)-n_below} ({(len(out)-n_below)/len(out):.1%})") +print(f">=5,000 km2 kept in full: {len(big)}") +print(f"all {len(out)} gauges have >= {COVERAGE_MIN:.0%} obs coverage in both" + f" {TRAIN_WINDOW} and {EVAL_WINDOW}") diff --git a/scripts/dump_gamma_uh_params.py b/scripts/dump_gamma_uh_params.py new file mode 100644 index 0000000..edea334 --- /dev/null +++ b/scripts/dump_gamma_uh_params.py @@ -0,0 +1,118 @@ +"""Dump the trained gamma-UH routing parameters per MERIT divide. + +Pulls routa/routb from the CONUS2717_AORC2F ep100 checkpoint's Ann head (the +model that generated daily_dhbv2_distributed_aorc2f_merit_unit_catchments.ic) +and computes the effective gamma kernel mean a_eff * theta_eff (days) per +divide. Mirrors forward_conus_divides.py --mode export exactly: same +attr_list, same normalization (dapengscaler_stat.json), same divide order +(the export store's divide_id axis), identity/singleton groups so each +divide's params are its own. + +Scaling chain (waterlossv18_1.py:57,179-183 + rnn.py UH_gamma): + r0, r1 = sigmoid Ann outputs, cols 52:54 (nfea2*nmul2 = 13*4) + tempa = 2.9 * r0 # "routa" in the 2026-07-29 findings + tempb = 6.5 * r1 # "routb" + a_eff = tempa + 0.1 # relu(a)+0.1 floor inside UH_gamma + theta_eff = tempb + 0.5 # relu(b)+0.5 floor + mean travel = a_eff * theta_eff [days] + +Run under the water_loss venv (has torch/hydroDL/icechunk): + ~/projects/water_loss/.venv/bin/python scripts/dump_gamma_uh_params.py +Output: output/tau_sweep/gamma_uh_params.csv (divide_id, routa, routb, +a_eff, theta_eff, tau_uh_days, uparea_km2) + printed stats. +""" +import json +import sys +from pathlib import Path + +import numpy as np + +HYDRODL = Path.home() / "projects/water_loss/dPLHBVrelease-master/hydroDL-dev" +sys.path.append(str(HYDRODL)) + +MODEL_OUT = Path("/mnt/ssd1/data/water_loss/models/CONUS2717_AORC2F_v3_gradaccum/" + "exp_EPOCH100_BS100_RHO365_HS164_MUL14_HS24096_MUL24_trainBuff365_test") +EPOCH = 100 +ATTRS_NC = "/mnt/ssd1/data/icechunk/merit_global_attributes_v2.nc" +EXPORT_IC = "/mnt/ssd1/data/icechunk/daily_dhbv2_distributed_aorc2f_merit_unit_catchments.ic" +OUT_CSV = Path(__file__).resolve().parent.parent / "output/tau_sweep/gamma_uh_params.csv" + +ATTR_LIST = ["meanP", "ETPOT_Hargr", "aridity", "seasonality_P", "snow_fraction", + "meanelevation", "meanslope", "NDVI", "Porosity", + "HWSD_sand", "HWSD_silt", "HWSD_clay", "permeability", "uparea"] +NFEA2, NMUL2 = 13, 4 # routpara = Ann output cols 52:54 + + +def main() -> None: + import icechunk + import torch + import xarray as xr + import zarr + from hydroDL.data import scale + from hydroDL.model.rnn import AnnModel + + sd = torch.load(MODEL_OUT / f"model_Ep{EPOCH}.pt", map_location="cpu", + weights_only=False) + assert isinstance(sd, dict), "expected a state_dict checkpoint" + ann_sd = {k[len("Ann."):]: v for k, v in sd.items() if k.startswith("Ann.")} + ny = ann_sd["h2o.weight"].shape[0] + assert ny == NFEA2 * NMUL2 + 2, f"Ann ny={ny}, expected {NFEA2*NMUL2+2}" + ann = AnnModel(nx=len(ATTR_LIST), ny=ny, hiddenSize=4096, dropout_rate=0.5) + ann.load_state_dict(ann_sd) + ann.eval() + + repo = icechunk.Repository.open(icechunk.local_filesystem_storage(EXPORT_IC)) + root = zarr.open_group(store=repo.readonly_session("main").store, mode="r") + divide_ids = root["divide_id"][:] + print(f"{len(divide_ids)} divides from export store") + + ads = xr.open_dataset(ATTRS_NC) + comid_to_aidx = {int(c): i for i, c in enumerate(ads["COMID"].values)} + aidx = np.array([comid_to_aidx[int(c)] for c in divide_ids]) + uparea_all = np.power(10.0, ads["log10_uparea"].values[aidx]).astype(np.float32) + attrs_all = np.stack( + [ads[v].values[aidx].astype(np.float32) for v in ATTR_LIST[:-1]], axis=1) + attrs_all = np.concatenate([attrs_all, uparea_all[:, None]], axis=1) + + with open(MODEL_OUT / "dapengscaler_stat.json") as f: + stat_dict = json.load(f) + attr_norm = scale._trans_norm(attrs_all.copy(), ATTR_LIST, stat_dict, + log_norm_cols=[], to_norm=True) + attr_norm[attr_norm != attr_norm] = 0 + + routpara = np.empty((len(divide_ids), 2), dtype=np.float32) + with torch.no_grad(): + for s in range(0, len(divide_ids), 8192): + e = min(s + 8192, len(divide_ids)) + out = ann(torch.from_numpy(attr_norm[s:e]).float()) + routpara[s:e] = out[:, NFEA2 * NMUL2:NFEA2 * NMUL2 + 2].numpy() + if s % 65536 == 0: + print(f" {e}/{len(divide_ids)}") + + routa = 2.9 * routpara[:, 0] + routb = 6.5 * routpara[:, 1] + a_eff = routa + 0.1 + theta_eff = routb + 0.5 + tau_uh = a_eff * theta_eff + + import pandas as pd + df = pd.DataFrame({"divide_id": divide_ids, "routa": routa, "routb": routb, + "a_eff": a_eff, "theta_eff": theta_eff, + "tau_uh_days": tau_uh, "uparea_km2": uparea_all}) + OUT_CSV.parent.mkdir(parents=True, exist_ok=True) + df.to_csv(OUT_CSV, index=False) + print(f"wrote {OUT_CSV}") + + q = np.percentile(tau_uh, [5, 25, 50, 75, 95]) + print(f"tau_uh_days: median {q[2]:.2f} IQR [{q[1]:.2f}, {q[3]:.2f}] " + f"p5 {q[0]:.2f} p95 {q[4]:.2f} mean {tau_uh.mean():.2f}") + print(f"routa: median {np.median(routa):.3f} routb: median {np.median(routb):.3f}") + from scipy.stats import spearmanr + rho = spearmanr(tau_uh, np.log10(uparea_all)).statistic + print(f"spearman(tau_uh, log10 uparea) = {rho:.3f}") + print(f"frac tau_uh > 1 day: {(tau_uh > 1).mean():.1%} " + f"> 2 days: {(tau_uh > 2).mean():.1%} > 4 days: {(tau_uh > 4).mean():.1%}") + + +if __name__ == "__main__": + main() diff --git a/scripts/run_tau9_hourly_eval_chain.sh b/scripts/run_tau9_hourly_eval_chain.sh new file mode 100755 index 0000000..f62b476 --- /dev/null +++ b/scripts/run_tau9_hourly_eval_chain.sh @@ -0,0 +1,17 @@ +#!/usr/bin/env bash +# After run_tau9_remaining.sh (the aorc2f_lumped run) exits, evaluate the +# hourly_lstm arm's final checkpoint with the legacy eval binary — the +# train-and-test zero-batch resume cannot re-enter Phase 2 (it requires +# checkpoints written by its own Phase 1), so this is the completion path +# for the arm whose eval was killed mid-run on 2026-08-09. +set -uo pipefail +cd /home/tbindas/projects/ddrs +while pgrep -f "run_tau9_remaining.sh" > /dev/null; do sleep 60; done +echo "=== hourly_lstm legacy eval start $(date -u +%FT%TZ) ===" +target/release/eval \ + --config config/experiments/tau9_train_hourly_lstm.yaml \ + --checkpoint .ddrs/runs/2026-08-09T14-55-05Z-train-and-test/checkpoints/epoch_30_mb_1 \ + --backend cpu \ + --output "$PWD/output/tau_sweep/train9_hourly_lstm_eval.zarr" \ + > output/tau_sweep/train9_hourly_lstm_eval.log 2>&1 +echo "=== hourly_lstm legacy eval done $(date -u +%FT%TZ), exit $? ===" diff --git a/scripts/run_tau9_remaining.sh b/scripts/run_tau9_remaining.sh new file mode 100755 index 0000000..a7acd51 --- /dev/null +++ b/scripts/run_tau9_remaining.sh @@ -0,0 +1,16 @@ +#!/usr/bin/env bash +# Finish the tau=9 CPU set after the 2026-08-09 driver kill: hourly_lstm +# eval-only resume (epoch-30 checkpoint, zero training batches) then the +# aorc2f_lumped full run. Launch with setsid so it survives the parent. +set -uo pipefail +cd /home/tbindas/projects/ddrs +DDRS=target/release/ddrs +for cfg in tau9_train_hourly_lstm_evalresume tau9_train_aorc2f_lumped; do + echo "=== ${cfg} start $(date -u +%FT%TZ) ===" + $DDRS --config config/experiments/${cfg}.yaml --workspace .ddrs \ + run --workflow train-and-test --backend cpu \ + > output/tau_sweep/train9_${cfg}.launch.log 2>&1 \ + || { echo "=== ${cfg} FAILED $(date -u +%FT%TZ) ==="; continue; } + echo "=== ${cfg} done $(date -u +%FT%TZ) ===" +done +echo "=== remaining arms done $(date -u +%FT%TZ) ===" diff --git a/scripts/run_tau9_source_trains.sh b/scripts/run_tau9_source_trains.sh new file mode 100644 index 0000000..621037d --- /dev/null +++ b/scripts/run_tau9_source_trains.sh @@ -0,0 +1,22 @@ +#!/usr/bin/env bash +# tau=9 cross-source retrain: five train-and-test runs, one per streamflow +# store, all on CPU (user directive 2026-08-08). Each config is the epoch-30 +# run's snapshot with only {streamflow, sparse_solver: cpu} changed and +# params.tau left to the new default 9 (2026-08-08 convention == old 20). +# +# Uses target/release/ddrs (NOT the installed ~/.cargo/bin copy) to dodge the +# stale-binary trap after the tau convention change. +# --workspace pins .ddrs to the repo root (configs live under config/experiments/). +set -uo pipefail +cd /home/tbindas/projects/ddrs + +DDRS=target/release/ddrs +for src in aorc2f_dist uh_retro daily_lstm hourly_lstm aorc2f_lumped; do + echo "=== train ${src} start $(date -u +%FT%TZ) ===" + $DDRS --config config/experiments/tau9_train_${src}.yaml --workspace .ddrs \ + run --workflow train-and-test --backend cpu \ + > output/tau_sweep/train9_${src}.launch.log 2>&1 \ + || { echo "=== train ${src} FAILED $(date -u +%FT%TZ), continuing ==="; continue; } + echo "=== train ${src} done $(date -u +%FT%TZ) ===" +done +echo "=== all arms done $(date -u +%FT%TZ) ===" diff --git a/scripts/run_tau_interp_arms.sh b/scripts/run_tau_interp_arms.sh new file mode 100755 index 0000000..486c04b --- /dev/null +++ b/scripts/run_tau_interp_arms.sh @@ -0,0 +1,27 @@ +#!/usr/bin/env bash +# 3-arm q'-interpolation experiment: nearest / linear / quadratic upsampling, +# gages_3000 population, epoch-30 checkpoint, full eval window + hourly dump. +# Spec: docs/superpowers/specs/2026-08-05-per-gauge-tau-sweep-design.md (Phase 2 prep) +set -euo pipefail +cd /home/tbindas/projects/ddrs + +CKPT=.ddrs/runs/2026-08-05T04-58-58Z-conus-experimental-train-and-test/checkpoints/epoch_30_mb_1 +CFG=config/experiments/tau_interp_g3000.yaml + +for mode in nearest linear quadratic; do + out=output/tau_sweep/g3000_${mode} + mkdir -p "$out" + echo "=== arm: ${mode} $(date -u +%FT%TZ) ===" + # CPU by default (user preference 2026-08-06): deterministic NdArray, keeps + # the GPU free. NOTE: never mix backends WITHIN one comparison set — the + # 2026-08-06 interp arms ran entirely on cuda. + DDRS_QPRIME_INTERP=${mode} \ + DDRS_HOURLY_DUMP=$PWD/${out}/hourly.f32 \ + target/release/eval \ + --config "$CFG" \ + --checkpoint "$CKPT" \ + --backend cpu \ + --output "$PWD/${out}/eval.zarr" \ + > "${out}/eval.log" 2>&1 + echo "=== arm ${mode} done $(date -u +%FT%TZ), exit $? ===" +done diff --git a/scripts/run_tau_source_arms.sh b/scripts/run_tau_source_arms.sh new file mode 100755 index 0000000..8aca2ea --- /dev/null +++ b/scripts/run_tau_source_arms.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +# Cross-source tau sweep: same epoch-30 checkpoint, same 2,365-gauge network, +# only data_sources.streamflow changes. Nearest (repeat-24) upsampling for +# daily stores; hourly_lstm is hourly-native (no upsampling). The aorc2f +# distributed store is already covered by output/tau_sweep/g3000_nearest/. +set -euo pipefail +cd /home/tbindas/projects/ddrs + +CKPT=.ddrs/runs/2026-08-05T04-58-58Z-conus-experimental-train-and-test/checkpoints/epoch_30_mb_1 + +for src in uh_retro daily_lstm hourly_lstm aorc2f_lumped; do + out=output/tau_sweep/src_${src} + mkdir -p "$out" + echo "=== source: ${src} $(date -u +%FT%TZ) ===" + # CPU by default (user preference 2026-08-06): deterministic NdArray, keeps + # the GPU free. NOTE: never mix backends WITHIN one comparison set — the + # 2026-08-06 cross-source set ran entirely on cuda. + DDRS_HOURLY_DUMP=$PWD/${out}/hourly.f32 \ + target/release/eval \ + --config config/experiments/tau_src_${src}.yaml \ + --checkpoint "$CKPT" \ + --backend cpu \ + --output "$PWD/${out}/eval.zarr" \ + > "${out}/eval.log" 2>&1 || { echo "=== source ${src} FAILED, continuing ==="; continue; } + echo "=== source ${src} done $(date -u +%FT%TZ) ===" +done diff --git a/scripts/tau_sweep.py b/scripts/tau_sweep.py new file mode 100644 index 0000000..a7e84a1 --- /dev/null +++ b/scripts/tau_sweep.py @@ -0,0 +1,220 @@ +"""Per-gauge tau sweep on a pre-trim hourly eval dump. + +Spec: docs/superpowers/specs/2026-08-05-per-gauge-tau-sweep-design.md + +Reads the DDRS_HOURLY_DUMP raw f32 (n_gauges, n_hours) + `.json` sidecar and +the eval run's predictions.zarr, reconstructs daily predictions for each +tau in -12..23 and reports per-gauge NSE(tau) restricted to a pilot window. + +Convention (2026-08-08, matches `tau_trim_and_downsample`): tau = hours the +routed output is advanced before daily scoring; pooled store day d covers +hours [tau + 24d, tau + 24(d+1)). The eval zarr's day axis starts at store +day 1 (day 0 is never scored), so zarr day i = store day i+1 and its window +is hours [24 + tau + 24i, ...). tau=0 is day-aligned. Dumps produced BEFORE +2026-08-08 carry a legacy tau_shipped (old = new + 11, window [13+tau ...], +zarr day i also = store day i+1); analyze those with the pre-change script +from git history. + +Run: + uv run --with "zarr>=3" --with numpy --with pandas --with matplotlib \ + python scripts/tau_sweep.py \ + --dump output/tau_sweep/hourly_full.f32 \ + --zarr output/tau_sweep/eval_full.zarr \ + --baseline-dir .ddrs/runs//baseline \ + --gages-csv ~/projects/ddr/references/gage_info/gages_2000_area_balanced.csv \ + --out-dir output/tau_sweep \ + --pilot-start 1995-10-01 --pilot-end 1996-09-30 +""" + +import argparse +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import zarr + +TAUS = np.arange(-12, 24) +AREA_BINS = [0, 1000, 5000, 10000, 30000, 50000] +AREA_LABELS = ["0~1000", "1000~5000", "5000~10000", "10000~30000", "30000~50000"] +MIN_VALID_DAYS = 100 + + +def nse(pred: np.ndarray, obs: np.ndarray) -> np.ndarray: + """Per-gauge NSE over axis 1, NaN-masked. Returns (G,) with NaN where undefined.""" + valid = np.isfinite(obs) & np.isfinite(pred) + n_valid = valid.sum(axis=1) + o = np.where(valid, obs, 0.0) + p = np.where(valid, pred, 0.0) + o_mean = o.sum(axis=1) / np.maximum(n_valid, 1) + sse = (((p - o) * valid) ** 2).sum(axis=1) + svar = (((o - o_mean[:, None]) * valid) ** 2).sum(axis=1) + out = 1.0 - sse / np.where(svar > 0, svar, np.nan) + out[n_valid < MIN_VALID_DAYS] = np.nan + return out + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--dump", required=True) + ap.add_argument("--zarr", required=True) + ap.add_argument("--baseline-dir", required=True) + ap.add_argument("--gages-csv", required=True) + ap.add_argument("--out-dir", required=True) + ap.add_argument("--pilot-start", default="1995-10-01") + ap.add_argument("--pilot-end", default="1996-09-30") + args = ap.parse_args() + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + meta = json.loads(Path(args.dump + ".json").read_text()) + g_n, h_n, tau_ship = meta["n_gauges"], meta["n_hours"], meta["tau_shipped"] + hourly = np.memmap(args.dump, dtype=np.float32, mode="r", shape=(g_n, h_n)) + + zg = zarr.open_group(args.zarr, mode="r") + zpred = zg["predictions"][:] # (G, D) + zobs = zg["observations"][:] + ztime = pd.to_datetime(zg["time"][:]) # ns since epoch + gage_ids = np.array([bytes(r).decode().strip("\x00") for r in zg["gage_ids"][:]]) + d_n = zpred.shape[1] + assert zpred.shape[0] == g_n, f"gauge mismatch: dump {g_n} vs zarr {zpred.shape[0]}" + + # --- Method-validation gate item 1: reproduce shipped tau exactly. --- + start = 24 + tau_ship + recon = np.asarray(hourly[:, start : start + 24 * d_n], dtype=np.float64) + recon = recon.reshape(g_n, d_n, 24).mean(axis=2) + denom = np.maximum(np.abs(zpred), 1e-6) + rel = np.abs(recon - zpred) / denom + print(f"[gate 1] tau={tau_ship} reconstruction: max rel {np.nanmax(rel):.2e}, " + f"median rel {np.nanmedian(rel):.2e} (PASS if max < 1e-3)") + gate1 = np.nanmax(rel) < 1e-3 + + # --- Gate item 2: full-window NSE at shipped tau matches the eval's own daily output. --- + nse_ship_full = nse(zpred, zobs) + nse_recon_full = nse(recon, zobs) + d_med = abs(np.nanmedian(nse_ship_full) - np.nanmedian(nse_recon_full)) + print(f"[gate 2] full-window median NSE shipped {np.nanmedian(nse_ship_full):.4f} " + f"vs recon {np.nanmedian(nse_recon_full):.4f} (PASS if |d| < 0.001)") + gate2 = d_med < 0.001 + + # --- Sweep: common day count so every tau scores the same days. --- + # Highest start is 24+23; need 24+23+24*D <= h_n (lowest start 24-12=12 >= 0). + d_common = min(d_n, (h_n - (24 + int(TAUS.max()))) // 24) + pilot = (ztime[:d_common] >= pd.Timestamp(args.pilot_start)) & ( + ztime[:d_common] <= pd.Timestamp(args.pilot_end) + ) + pilot = np.asarray(pilot) + print(f"pilot window {args.pilot_start}..{args.pilot_end}: {pilot.sum()} days") + obs_pilot = zobs[:, :d_common][:, pilot] + + nse_by_tau = np.full((g_n, len(TAUS)), np.nan) + for k, tau in enumerate(TAUS): + s = 24 + int(tau) + daily = np.asarray(hourly[:, s : s + 24 * d_common], dtype=np.float64) + daily = daily.reshape(g_n, d_common, 24).mean(axis=2) + nse_by_tau[:, k] = nse(daily[:, pilot], obs_pilot) + gate3_frac = np.isfinite(nse_by_tau).all(axis=1).mean() + print(f"[gate 3] gauges with full NSE(tau) curves: {gate3_frac:.1%} (PASS if >= 95%)") + gate3 = gate3_frac >= 0.95 + + # --- Baseline (same population), pilot window. --- + bdir = Path(args.baseline_dir) + bman = json.loads((bdir / "manifest.json").read_text()) + bg, bd = bman["n_gauges"], bman["n_days"] + btime = pd.to_datetime(bman["time_range_daily"]) + bpred = np.fromfile(bdir / "predictions.f32", dtype=np.float32).reshape(bg, bd) + bobs = np.fromfile(bdir / "observations.f32", dtype=np.float32).reshape(bg, bd) + border = pd.Index(bman["gage_ids"]).get_indexer(gage_ids) + covered = border >= 0 + if not covered.all(): + print(f"WARNING: baseline covers {covered.sum()}/{g_n} eval gauges; rest -> NaN") + bmask = np.asarray( + (btime >= pd.Timestamp(args.pilot_start)) & (btime <= pd.Timestamp(args.pilot_end)) + ) + nse_base = np.full(g_n, np.nan) + nse_base[covered] = nse( + bpred[border[covered]][:, bmask], bobs[border[covered]][:, bmask] + ) + + # --- Covariates + summary table. --- + gcsv = pd.read_csv(args.gages_csv, dtype={"STAID": str}) + gcsv["STAID"] = gcsv["STAID"].str.zfill(8) + cov = gcsv.set_index("STAID").reindex(gage_ids) + has_curve = np.isfinite(nse_by_tau).any(axis=1) + k_ship = int(np.where(TAUS == tau_ship)[0][0]) + best_k = np.full(g_n, k_ship, dtype=int) + best_k[has_curve] = np.nanargmax(nse_by_tau[has_curve], axis=1) + best_tau = TAUS[best_k] + df = pd.DataFrame( + { + "gage_id": gage_ids, + "drain_sqkm": cov["DRAIN_SQKM"].to_numpy(), + "lng_gage": cov["LNG_GAGE"].to_numpy(), + "has_curve": has_curve, + "best_tau": best_tau, + "nse_tau_shipped": nse_by_tau[:, k_ship], + "nse_tau_best": nse_by_tau[np.arange(g_n), best_k], + "nse_baseline": nse_base, + } + ) + df["area_bin"] = pd.cut(df["drain_sqkm"], AREA_BINS, labels=AREA_LABELS) + df.to_csv(out_dir / "best_tau_wy1996.csv", index=False) + pd.DataFrame(nse_by_tau, index=gage_ids, columns=[f"tau_{t}" for t in TAUS]).to_csv( + out_dir / "nse_by_tau_wy1996.csv" + ) + + med = df.groupby("area_bin", observed=True)[ + ["nse_tau_shipped", "nse_tau_best", "nse_baseline"] + ].median() + med["n"] = df.groupby("area_bin", observed=True).size() + improved = df["nse_tau_best"] - df["nse_tau_shipped"] > 0.01 + rho = df.loc[improved, ["best_tau", "lng_gage"]].corr(method="spearman").iloc[0, 1] + + lines = [ + "# tau sweep pilot (WY1996) summary", + "", + f"- gate 1 (recon match): {'PASS' if gate1 else 'FAIL'}", + f"- gate 2 (NSE match): {'PASS' if gate2 else 'FAIL'}", + f"- gate 3 (curve cover): {'PASS' if gate3 else 'FAIL'} ({gate3_frac:.1%})", + "", + "Median NSE by drainage-area bin (pilot window, in-sample best tau):", + "", + med.to_markdown(), + "", + f"- gauges improved > 0.01 by best tau: {improved.sum()} / {g_n}", + f"- spearman(best_tau, longitude) on improved gauges: {rho:.3f}", + f"- best-tau distribution (has_curve only, tau {TAUS[0]}..{TAUS[-1]}): " + f"{np.histogram(best_tau[has_curve], bins=np.arange(TAUS[0], TAUS[-1] + 2) - 0.5)[0].tolist()}", + ] + (out_dir / "summary_wy1996.md").write_text("\n".join(lines) + "\n") + print("\n".join(lines)) + + # --- Plots. --- + fig, ax = plt.subplots(figsize=(8, 5)) + for label in AREA_LABELS: + sel = (df["area_bin"] == label).to_numpy() + if sel.any(): + ax.plot(TAUS, np.nanmedian(nse_by_tau[sel], axis=0), marker="o", label=f"{label} (n={sel.sum()})") + ax.axvline(tau_ship, color="k", ls="--", lw=1, label=f"shipped tau={tau_ship}") + ax.set_xlabel("tau (h)"), ax.set_ylabel("median NSE"), ax.legend(fontsize=8) + ax.set_title("Median NSE vs tau by drainage area (WY1996)") + fig.savefig(out_dir / "nse_vs_tau_by_area.png", dpi=150, bbox_inches="tight") + + fig, ax = plt.subplots(figsize=(8, 4)) + ax.hist(df.loc[improved, "best_tau"], bins=np.arange(TAUS[0], TAUS[-1] + 2) - 0.5) + ax.set_xlabel("best tau (h)"), ax.set_ylabel("gauges (improved > 0.01)") + fig.savefig(out_dir / "best_tau_hist.png", dpi=150, bbox_inches="tight") + + fig, ax = plt.subplots(figsize=(8, 4)) + ax.scatter(df.loc[improved, "lng_gage"], df.loc[improved, "best_tau"], s=8, alpha=0.5) + ax.set_xlabel("longitude"), ax.set_ylabel("best tau (h)") + ax.set_title(f"spearman rho = {rho:.3f} (improved gauges)") + fig.savefig(out_dir / "best_tau_vs_longitude.png", dpi=150, bbox_inches="tight") + + +if __name__ == "__main__": + main() diff --git a/src/adjacency/cache.rs b/src/adjacency/cache.rs index 09ed0d8..487a8bf 100644 --- a/src/adjacency/cache.rs +++ b/src/adjacency/cache.rs @@ -3,7 +3,7 @@ //! Layout under `/adjacency//`: //! //! ```text -//! _adjacency.zarr/ — written by zarr_write::write_conus_store +//! _adjacency.zarr/ — zarr_write::write_conus_store_subdivided //! _gages_adjacency.zarr/ — written by zarr_write::write_gauges_store //! manifest.json — input paths + fingerprints, graph dims, //! dropped COMIDs, build duration, git SHA @@ -18,10 +18,11 @@ //! `*_adjacency.zarr` pair is present so old caches keep hitting. //! //! The key is blake3 of the resolved fabric file bytes (.dbf or .gpkg) ∥ -//! gages CSV file bytes ∥ optional gpkg layer name ∥ BUILDER_VERSION, -//! truncated to 16 hex chars. Content fingerprints (not stat/path) are used -//! so that moving or renaming the files does NOT invalidate, and so two -//! files with identical bytes share a cache entry. +//! gages CSV file bytes ∥ optional gpkg layer name ∥ BUILDER_VERSION ∥ every +//! field of `params.subdivision`, truncated to 16 hex chars. Content +//! fingerprints (not stat/path) are used so that moving or renaming the files +//! does NOT invalidate, and so two files with identical bytes share a cache +//! entry. //! //! Build is crash-safe: everything is written into a temp dir //! `/adjacency/.tmp-` then atomically renamed into place. @@ -33,12 +34,17 @@ use std::time::Instant; use serde::{Deserialize, Serialize}; -use crate::adjacency::build::{build_conus_adjacency, BuildError}; +use crate::adjacency::build::{build_conus_adjacency, BuildError, ConusAdjacency}; use crate::adjacency::fabric::{read_fabric_records, resolve_fabric}; use crate::adjacency::gauges::build_gauge_subgraphs; -use crate::adjacency::zarr_write::{write_conus_store, write_gauges_store}; +use crate::adjacency::subdivide::{plan_reaches, subdivide, ReachPlan, SubdividedAdjacency}; +use crate::adjacency::zarr_write::{write_conus_store_subdivided, write_gauges_store}; use crate::adjacency::BUILDER_VERSION; +use crate::config::Subdivision; use crate::data::error::DataError; +use crate::data::ids::Comid; +use crate::data::store::netcdf::AttributesStore; +use crate::routing::mmc::DT_SECONDS; /// Paths to the two zarr stores produced by (or read from) the cache. #[derive(Debug, Clone)] @@ -94,8 +100,13 @@ struct CacheManifest { gages_path: PathBuf, /// blake3 hex of the gages CSV file bytes. gages_fingerprint: String, - /// Number of CONUS reaches in `order`. + /// Number of rows in `order` — sub-reaches when subdivision is enabled, + /// otherwise MERIT reaches. n: usize, + /// Number of MERIT reaches (parent rows). Equals `n` unless subdivision is + /// enabled. Defaulted for manifests written before BUILDER_VERSION 2. + #[serde(default)] + n_parent: usize, /// Number of COO edges (nnz). nnz: usize, /// Number of gauge subgraphs built. @@ -123,11 +134,18 @@ struct CacheManifest { /// progress line is printed before the build begins so the user knows it is /// not hung. For a multi-GB global gpkg the content fingerprint itself costs /// a few seconds of hashing on first run. +/// +/// `subdivision` is `cfg.params.subdivision`. When disabled (the default) the +/// build is byte-identical to the pre-subdivision one and `attributes` is never +/// opened; when enabled, `attributes` supplies the `catchsize` column that the +/// reference celerity needs. pub fn resolve_or_build( workspace_root: &Path, fabric: &Path, fabric_layer: Option<&str>, gages_csv: &Path, + subdivision: &Subdivision, + attributes: &[PathBuf], ) -> Result { // Resolve the hashable artifact (.shp → sibling .dbf; .dbf/.gpkg as-is). let resolved = resolve_fabric(fabric).map_err(AdjacencyCacheError::Data)?; @@ -143,7 +161,7 @@ pub fn resolve_or_build( // --- 1. Compute the content key ----------------------------------------- let fabric_fp = file_fingerprint(&resolved_path)?; let gages_fp = gages_fingerprint(gages_csv)?; - let key = content_key(&fabric_fp, &gages_fp, fabric_layer); + let key = content_key(&fabric_fp, &gages_fp, fabric_layer, subdivision); // --- 2. Cache hit? ------------------------------------------------------- let cache_dir = adjacency_cache_dir(workspace_root, &key); @@ -187,15 +205,43 @@ pub fn resolve_or_build( let conus = build_conus_adjacency(&records) .map_err(AdjacencyCacheError::Build)?; - let n = conus.order.len(); - let nnz = conus.rows.len(); + let n_parent = conus.order.len(); + // `subdivide` intentionally drops `dropped_comids` (it is parent-space + // provenance, not graph state), so capture it before the expansion. let dropped_comids = conus.dropped_comids.clone(); + // Subdivision runs BEFORE the gauge subgraphs so those are cut from the + // expanded graph. With `enabled: false` the plan is all-ones and `subdivide` + // is an exact identity, so this is a true no-op on the default path. + let expanded = expand(&conus, subdivision, attributes)?; + let n = expanded.order.len(); + let nnz = expanded.rows.len(); + if subdivision.enabled { + println!( + " subdivision: {n_parent} reaches → {n} sub-reaches ({:.2}x), \ + {} edges", + n as f64 / n_parent as f64, + nnz + ); + } + let conus_dest = tmp_dir.join(format!("{fabric_stem}_adjacency.zarr")); - write_conus_store(&conus, &conus_dest) + write_conus_store_subdivided(&expanded, &conus_dest) .map_err(AdjacencyCacheError::Data)?; - let subgraphs = build_gauge_subgraphs(&conus, gages_csv) + // Gauge subgraphs are built in SUB-REACH position space. `position_lookup` + // is COMID → last matching row, and a parent's rows are contiguous and + // ordered upstream→downstream, so a gauge resolves to its parent's OUTLET + // piece — the only row carrying the whole reach's lateral inflow. + let expanded_view = ConusAdjacency { + order: expanded.order.clone(), + rows: expanded.rows.clone(), + cols: expanded.cols.clone(), + length_m: expanded.length_m.clone(), + slope: expanded.slope.clone(), + dropped_comids: dropped_comids.clone(), + }; + let subgraphs = build_gauge_subgraphs(&expanded_view, gages_csv) .map_err(AdjacencyCacheError::Data)?; let n_gauges = subgraphs.len(); @@ -216,6 +262,7 @@ pub fn resolve_or_build( gages_path: gages_csv.to_path_buf(), gages_fingerprint: gages_fp, n, + n_parent, nnz, n_gauges, dropped_comids, @@ -297,6 +344,106 @@ fn store_paths(cache_dir: &Path, fabric_stem: &str) -> AdjacencyCachePaths { } } +/// Apply the two-sided reach plan and expand the graph. +/// +/// Disabled is the fast path: `plan_reaches` returns all-ones with untouched +/// lengths and `subdivide` is an exact identity, so the attributes NetCDF is +/// never opened and a machine without one still builds. +fn expand( + conus: &ConusAdjacency, + subdivision: &Subdivision, + attributes: &[PathBuf], +) -> Result { + let plan = reach_plan(conus, subdivision, attributes)?; + Ok(subdivide(conus, &plan)) +} + +/// The two-sided reach plan for an already-built parent adjacency. +/// +/// Split out of [`expand`] so the plan can be measured (`subdivide::plan_stats`) +/// without writing a store — `probe_courant --clamp-report` uses this. +pub fn reach_plan( + conus: &ConusAdjacency, + subdivision: &Subdivision, + attributes: &[PathBuf], +) -> Result { + let uparea = if subdivision.enabled { + upstream_area_km2(conus, attributes)? + } else { + // `plan_reaches` returns early when disabled and never reads this, but a + // correctly-sized vector keeps the contract intact if that ever changes. + vec![0.0; conus.order.len()] + }; + Ok(plan_reaches( + &conus.length_m, + &conus.slope, + &uparea, + DT_SECONDS, + subdivision, + )) +} + +/// Build the parent adjacency straight from a fabric, with no cache and no +/// store write. The measurement entry point for `probe_courant --clamp-report`. +pub fn parent_adjacency_from_fabric( + fabric: &Path, + fabric_layer: Option<&str>, +) -> Result { + let records = read_fabric_records(fabric, fabric_layer).map_err(AdjacencyCacheError::Data)?; + build_conus_adjacency(&records).map_err(AdjacencyCacheError::Build) +} + +/// Upstream drainage area (km²) per reach, aligned to `conus.order`. +/// +/// **`catchsize` is the LOCAL divide area, not drainage area** — median 36.7 km² +/// on MERIT, capped around 612 km² even for continental rivers. `reference_celerity` +/// wants the *upstream* area, so the local column is accumulated downstream over +/// the topological order. Verified against the fabric's own `log10_uparea` +/// column on real CONUS: the accumulated value reproduces `10^log10_uparea` with +/// ratio p5/p50/p95 all 1.000 over all 346,321 reaches. (`log10_uparea` itself is +/// unusable as the source — it is NaN on 88 % of the global attributes file, +/// finite only over CONUS.) +/// +/// Non-finite or non-positive `catchsize` falls back to the column mean, matching +/// how `build_conus_adjacency` fills `length_m`/`slope`. +fn upstream_area_km2( + conus: &ConusAdjacency, + attributes: &[PathBuf], +) -> Result, AdjacencyCacheError> { + let path = attributes.first().ok_or_else(|| { + AdjacencyCacheError::Data(DataError::Malformed { + path: PathBuf::from("data_sources.attributes"), + message: "params.subdivision.enabled requires data_sources.attributes \ + (the `catchsize` column sets the reference celerity)" + .to_string(), + }) + })?; + let comids: Vec = conus.order.iter().map(|&c| Comid(c as i64)).collect(); + let store = AttributesStore::open_aligned(path, &["catchsize".to_string()], &comids) + .map_err(AdjacencyCacheError::Data)?; + + let fallback = store.row_means[0]; + let mut area: Vec = store + .attrs + .row(0) + .iter() + .map(|&v| if v.is_finite() && v > 0.0 { v } else { fallback }) + .collect(); + + // Accumulate: `order` is topological and the COO is lower-triangular + // (`rows[k] > cols[k]`), so visiting upstream positions in ascending order + // guarantees each contributor is already complete when it is added. + let mut upstream_of: Vec> = vec![Vec::new(); area.len()]; + for (&r, &c) in conus.rows.iter().zip(conus.cols.iter()) { + upstream_of[r as usize].push(c as usize); + } + for p in 0..area.len() { + let contrib: f32 = upstream_of[p].iter().map(|&u| area[u]).sum(); + area[p] += contrib; + } + Ok(area) +} + /// A cache hit requires the directory to exist with a manifest.json present. /// Mirrors baseline/cache.rs's hit criterion (manifest presence = valid cache). fn is_cache_hit(dir: &Path) -> bool { @@ -304,7 +451,7 @@ fn is_cache_hit(dir: &Path) -> bool { } /// 16-hex-char (64-bit prefix) content key: -/// blake3(fabric_fp ∥ gages_fp ∥ [layer ∥] version). +/// blake3(fabric_fp ∥ gages_fp ∥ [layer ∥] version ∥ subdivision). /// /// Inputs are the full hex fingerprints of the file bytes, not the paths. /// Collision-free at our scale; safe for filesystem use. @@ -313,7 +460,18 @@ fn is_cache_hit(dir: &Path) -> bool { /// identical bytes regardless of which layer is selected, so the layer name /// must distinguish cache entries. Folding it in only when set keeps every /// pre-gpkg cache key (dbf fabrics, layer always `None`) unchanged. -fn content_key(fabric_fp: &str, gages_fp: &str, layer: Option<&str>) -> String { +/// +/// **Every field of [`Subdivision`] is hashed, not just `enabled`.** All seven +/// feed `reference_celerity` → `dx_target` → both the piece count and the +/// clamped `length_m`, i.e. each one changes the graph that gets built. Hashing +/// a subset would let a config edit silently reuse a stale adjacency — a failure +/// that looks like a physics result rather than a bug. +fn content_key( + fabric_fp: &str, + gages_fp: &str, + layer: Option<&str>, + s: &Subdivision, +) -> String { let mut h = blake3::Hasher::new(); h.update(fabric_fp.as_bytes()); h.update(b"\n"); @@ -324,10 +482,33 @@ fn content_key(fabric_fp: &str, gages_fp: &str, layer: Option<&str>) -> String { h.update(b"\n"); } h.update(BUILDER_VERSION.to_le_bytes().as_ref()); + h.update(&[s.enabled as u8]); + h.update((s.max_pieces as u32).to_le_bytes().as_ref()); + // `to_bits` gives a stable byte pattern across runs; config validation + // rejects non-finite values, so no NaN payload ambiguity arises. + h.update(s.reference_n.to_bits().to_le_bytes().as_ref()); + h.update(s.reference_discharge_coefficient.to_bits().to_le_bytes().as_ref()); + h.update(s.reference_discharge_exponent.to_bits().to_le_bytes().as_ref()); + h.update(s.min_length_fraction.to_bits().to_le_bytes().as_ref()); + h.update(s.max_clamp_factor.to_bits().to_le_bytes().as_ref()); let hex = h.finalize().to_hex(); hex.as_str()[..16].to_string() } +/// Test hook for [`content_key`]. Integration tests live in their own crate and +/// cannot reach a private fn; exposing the real function (rather than +/// reimplementing the hash in the test) is what makes +/// `cache_key_changes_with_every_subdivision_field` a meaningful guard. +#[doc(hidden)] +pub fn content_key_for_test( + fabric_fp: &str, + gages_fp: &str, + layer: Option<&str>, + s: &Subdivision, +) -> String { + content_key(fabric_fp, gages_fp, layer, s) +} + /// blake3 hex over the full contents of a file, using a buffered reader. /// Content fingerprint of the gages input: a single CSV file, or a directory /// of per-zone CSVs (e.g. `v3.1/8km/_all.csv` — see @@ -458,6 +639,17 @@ mod tests { fs::write(path, csv).expect("write gages csv"); } + /// Minimal attributes NetCDF carrying only the `catchsize` column the + /// reference celerity needs (km² of LOCAL divide area, per COMID). + fn write_catchsize_nc(path: &Path, comids: &[i64], catchsize: &[f64]) { + let mut f = netcdf::create(path).expect("create netcdf"); + f.add_dimension("COMID", comids.len()).unwrap(); + let mut cv = f.add_variable::("COMID", &["COMID"]).unwrap(); + cv.put_values(comids, ..).unwrap(); + let mut v = f.add_variable::("catchsize", &["COMID"]).unwrap(); + v.put_values(catchsize, ..).unwrap(); + } + fn tmp_workspace(tag: &str) -> PathBuf { let p = std::env::temp_dir() .join(format!("ddrs_adj_cache_{}_{}", tag, std::process::id())); @@ -470,31 +662,31 @@ mod tests { #[test] fn content_key_is_16_hex_chars() { - let k = content_key("aabbcc", "ddeeff", None); + let k = content_key("aabbcc", "ddeeff", None, &Subdivision::default()); assert_eq!(k.len(), 16, "key must be 16 hex chars, got: {k}"); assert!(k.chars().all(|c| c.is_ascii_hexdigit()), "key must be hex: {k}"); } #[test] fn content_key_is_stable() { - let k1 = content_key("fp1", "fp2", None); - let k2 = content_key("fp1", "fp2", None); + let k1 = content_key("fp1", "fp2", None, &Subdivision::default()); + let k2 = content_key("fp1", "fp2", None, &Subdivision::default()); assert_eq!(k1, k2); } #[test] fn content_key_differs_on_different_inputs() { - let k1 = content_key("aaaa", "bbbb", None); - let k2 = content_key("aaaa", "cccc", None); + let k1 = content_key("aaaa", "bbbb", None, &Subdivision::default()); + let k2 = content_key("aaaa", "cccc", None, &Subdivision::default()); assert_ne!(k1, k2); } #[test] fn content_key_differs_on_layer() { // Same fabric bytes, different gpkg layer → distinct cache entries. - let base = content_key("aaaa", "bbbb", None); - let l1 = content_key("aaaa", "bbbb", Some("flowlines")); - let l2 = content_key("aaaa", "bbbb", Some("catchments")); + let base = content_key("aaaa", "bbbb", None, &Subdivision::default()); + let l1 = content_key("aaaa", "bbbb", Some("flowlines"), &Subdivision::default()); + let l2 = content_key("aaaa", "bbbb", Some("catchments"), &Subdivision::default()); assert_ne!(l1, l2); assert_ne!(base, l1); } @@ -544,12 +736,12 @@ mod tests { write_tiny_dbf(&dbf_path); write_tiny_gages(&gages_path); - let out1 = resolve_or_build(&ws, &dbf_path, None, &gages_path).expect("build"); + let out1 = resolve_or_build(&ws, &dbf_path, None, &gages_path, &Subdivision::default(), &[]).expect("build"); let dir = adjacency_cache_dir(&ws, &out1.key); fs::rename(&out1.paths.conus, dir.join("merit_conus_adjacency.zarr")).unwrap(); fs::rename(&out1.paths.gages, dir.join("merit_gages_conus_adjacency.zarr")).unwrap(); - let out2 = resolve_or_build(&ws, &dbf_path, None, &gages_path).expect("hit"); + let out2 = resolve_or_build(&ws, &dbf_path, None, &gages_path, &Subdivision::default(), &[]).expect("hit"); assert!(out2.cache_hit, "legacy-named cache must still hit"); assert!(out2.paths.conus.ends_with("merit_conus_adjacency.zarr")); assert!(out2.paths.gages.ends_with("merit_gages_conus_adjacency.zarr")); @@ -566,11 +758,11 @@ mod tests { write_tiny_dbf(&dbf_path); write_tiny_gages(&gages_path); - let out1 = resolve_or_build(&ws, &dbf_path, None, &gages_path).expect("build"); + let out1 = resolve_or_build(&ws, &dbf_path, None, &gages_path, &Subdivision::default(), &[]).expect("build"); fs::remove_dir_all(&out1.paths.conus).unwrap(); fs::remove_dir_all(&out1.paths.gages).unwrap(); - let out2 = resolve_or_build(&ws, &dbf_path, None, &gages_path).expect("rebuild"); + let out2 = resolve_or_build(&ws, &dbf_path, None, &gages_path, &Subdivision::default(), &[]).expect("rebuild"); assert!(!out2.cache_hit, "stale cache must rebuild, not hit"); assert!(out2.paths.conus.is_dir() && out2.paths.gages.is_dir()); } @@ -586,10 +778,10 @@ mod tests { fs::create_dir(&gages_dir).unwrap(); write_tiny_gages(&gages_dir.join("74_all.csv")); - let out = resolve_or_build(&ws, &dbf_path, None, &gages_dir) + let out = resolve_or_build(&ws, &dbf_path, None, &gages_dir, &Subdivision::default(), &[]) .expect("build with gages directory"); assert!(!out.cache_hit); - let out2 = resolve_or_build(&ws, &dbf_path, None, &gages_dir) + let out2 = resolve_or_build(&ws, &dbf_path, None, &gages_dir, &Subdivision::default(), &[]) .expect("second call"); assert!(out2.cache_hit, "same dir contents → cache hit"); assert_eq!(out.key, out2.key); @@ -606,7 +798,7 @@ mod tests { write_tiny_gages(&gages_path); // First call: cache miss → build. - let out1 = resolve_or_build(&ws, &dbf_path, None, &gages_path) + let out1 = resolve_or_build(&ws, &dbf_path, None, &gages_path, &Subdivision::default(), &[]) .expect("first build"); assert!(!out1.cache_hit, "first call should be a cache miss"); assert_eq!(out1.key.len(), 16); @@ -635,7 +827,7 @@ mod tests { assert!(m.fabric_layer.is_none(), "dbf fabric has no layer"); // Second call: cache hit. - let out2 = resolve_or_build(&ws, &dbf_path, None, &gages_path) + let out2 = resolve_or_build(&ws, &dbf_path, None, &gages_path, &Subdivision::default(), &[]) .expect("second call"); assert!(out2.cache_hit, "second call should be a cache hit"); assert_eq!(out2.key, out1.key, "same key on hit"); @@ -664,8 +856,8 @@ mod tests { ); write_tiny_gages(&gages_path); - let out1 = resolve_or_build(&ws, &dbf_v1, None, &gages_path).expect("build v1"); - let out2 = resolve_or_build(&ws, &dbf_v2, None, &gages_path).expect("build v2"); + let out1 = resolve_or_build(&ws, &dbf_v1, None, &gages_path, &Subdivision::default(), &[]).expect("build v1"); + let out2 = resolve_or_build(&ws, &dbf_v2, None, &gages_path, &Subdivision::default(), &[]).expect("build v2"); assert_ne!(out1.key, out2.key, "keys must differ for different dbf content"); assert!(!out1.cache_hit); @@ -687,18 +879,98 @@ mod tests { // content_key output against a manual hash that includes the version. let dbf_fp = "deadbeef"; let gages_fp = "cafebabe"; - let k = content_key(dbf_fp, gages_fp, None); + let k = content_key(dbf_fp, gages_fp, None, &Subdivision::default()); + let s = Subdivision::default(); let mut h = blake3::Hasher::new(); h.update(dbf_fp.as_bytes()); h.update(b"\n"); h.update(gages_fp.as_bytes()); h.update(b"\n"); h.update(BUILDER_VERSION.to_le_bytes().as_ref()); + h.update(&[s.enabled as u8]); + h.update((s.max_pieces as u32).to_le_bytes().as_ref()); + h.update(s.reference_n.to_bits().to_le_bytes().as_ref()); + h.update(s.reference_discharge_coefficient.to_bits().to_le_bytes().as_ref()); + h.update(s.reference_discharge_exponent.to_bits().to_le_bytes().as_ref()); + h.update(s.min_length_fraction.to_bits().to_le_bytes().as_ref()); + h.update(s.max_clamp_factor.to_bits().to_le_bytes().as_ref()); let expected = &h.finalize().to_hex().as_str()[..16].to_string(); assert_eq!(&k, expected, "content_key must include BUILDER_VERSION"); } + /// End-to-end on the ENABLED path: long reaches must split, the store must + /// carry a non-trivial parent map, and the gauge must land on its parent's + /// outlet row. Without this the enabled branch is only ever type-checked. + #[test] + fn subdivided_build_expands_the_graph_and_persists_the_parent_map() { + use crate::data::ids::Staid; + use crate::data::store::zarr::{ConusAdjacencyStore, GagesAdjacencyStore}; + + let ws = tmp_workspace("subdiv"); + let dbf_path = ws.join("fabric.dbf"); + let gages_path = ws.join("gages.csv"); + let attrs_path = ws.join("attrs.nc"); + // 20 km and 30 km reaches: both far longer than dx_target, so both split. + write_dbf( + &dbf_path, + &[ + (1.0, 20.0, 0.001, 2.0, 0.0), + (2.0, 30.0, 0.002, 0.0, 1.0), + ], + ); + write_tiny_gages(&gages_path); + write_catchsize_nc(&attrs_path, &[1, 2], &[50.0, 60.0]); + + let on = Subdivision { + enabled: true, + max_pieces: 8, + ..Default::default() + }; + let attrs = [attrs_path.clone()]; + let out = resolve_or_build(&ws, &dbf_path, None, &gages_path, &on, &attrs) + .expect("subdivided build"); + assert!(!out.cache_hit); + // Enabling must not collide with the disabled cache entry. + let off = resolve_or_build(&ws, &dbf_path, None, &gages_path, + &Subdivision::default(), &[]) + .expect("disabled build"); + assert_ne!(out.key, off.key, "subdivision must change the cache key"); + + let conus = ConusAdjacencyStore::open(&out.paths.conus).expect("open conus"); + assert_eq!(conus.n_parent(), 2, "two MERIT reaches"); + assert!(conus.n > 2, "20 km / 30 km reaches must split, got n = {}", conus.n); + assert_eq!(conus.parent_order, vec![Comid(1), Comid(2)]); + assert_eq!(conus.parent_offset[0], 0); + assert_eq!(*conus.parent_offset.last().unwrap(), conus.n as i32); + // COMID lookups resolve in PARENT space even though `order` has dupes. + assert_eq!(conus.index.position(&Comid(2)), Some(1)); + assert_eq!(conus.order.len(), conus.n); + // Every sub-reach row carries its parent's COMID. + for p in 0..conus.n_parent() { + for row in conus.parent_offset[p] as usize..conus.parent_offset[p + 1] as usize { + assert_eq!(conus.order[row], conus.parent_order[p]); + } + } + // Length is conserved per parent (30 km reach, split m ways). + let lo = conus.parent_offset[1] as usize; + let hi = conus.parent_offset[2] as usize; + let total: f32 = conus.length_m.as_slice().unwrap()[lo..hi].iter().sum(); + assert!((total - 30_000.0).abs() < 1.0, "parent 1 length {total} != 30 km"); + + // The gauge on COMID 2 must read the OUTLET piece of parent 1. + let staids = vec![Staid::new("2")]; + let gages = GagesAdjacencyStore::open(&out.paths.gages, &staids).expect("open gages"); + let g = gages.get(&Staid::new("2")).expect("gauge 2"); + assert_eq!( + g.gage_idx, + conus.outlet_row(1), + "gauge must sit on its parent's outlet, not an interior piece" + ); + + let _ = fs::remove_dir_all(&ws); + } + #[test] fn zarr_stores_readable_after_build() { use crate::data::ids::{Comid, Staid}; @@ -710,7 +982,7 @@ mod tests { write_tiny_dbf(&dbf_path); write_tiny_gages(&gages_path); - let out = resolve_or_build(&ws, &dbf_path, None, &gages_path).expect("build"); + let out = resolve_or_build(&ws, &dbf_path, None, &gages_path, &Subdivision::default(), &[]).expect("build"); assert!(!out.cache_hit); // Verify the conus store round-trips. diff --git a/src/adjacency/mod.rs b/src/adjacency/mod.rs index 7139364..1dc9227 100644 --- a/src/adjacency/mod.rs +++ b/src/adjacency/mod.rs @@ -25,7 +25,10 @@ /// Bump on any algorithm change that would invalidate previously-cached /// adjacency zarr outputs. -pub const BUILDER_VERSION: u32 = 1; +/// +/// - `1` → `2`: reach subdivision. The builder now emits `/parent_order` and +/// `/parent_offset`, and `params.subdivision` participates in the content key. +pub const BUILDER_VERSION: u32 = 2; pub mod build; pub mod cache; @@ -33,5 +36,6 @@ pub mod dbf; pub mod fabric; pub mod gauges; pub mod gpkg; +pub mod subdivide; pub mod validate; pub mod zarr_write; diff --git a/src/adjacency/subdivide.rs b/src/adjacency/subdivide.rs new file mode 100644 index 0000000..68535ae --- /dev/null +++ b/src/adjacency/subdivide.rs @@ -0,0 +1,342 @@ +//! Static reach subdivision so Cr = c*dt/dx lands near 1. +//! +//! HEC-HMS picks the space step as `dx = c*dt` (Technical Reference Manual, +//! Muskingum-Cunge Model). ddrs historically used the full MERIT reach length, +//! giving median Cr = 0.226 — reaches ~4.4x too long. +//! +//! This module is a pure preprocessing step: plain `Vec`s in, plain `Vec`s out. +//! No BURN, no I/O, no dependence on training state. + +use crate::adjacency::build::ConusAdjacency; +use crate::config::Subdivision; + +/// Reference celerity (m/s) used ONLY to choose the piece count. +/// +/// Mirrors the solver's S15/S17 chain with the wide-channel ratio c = (5/3)*v +/// instead of the exact trapezoidal beta: this only sets `m`, the cap dominates +/// the result, and beta needs the learned p_spatial/q_spatial. +pub fn reference_celerity(uparea_km2: f32, slope: f32, cfg: &Subdivision) -> f32 { + // Slope floor mirrors `attribute_minimums.slope` (mmc.rs:208). + let s = slope.max(1e-3); + let q_ref = cfg.reference_discharge_coefficient + * uparea_km2.max(0.0).powf(cfg.reference_discharge_exponent); + // Hydraulic radius via a wide-channel regime relation; the exponent 0.4 is + // the Leopold & Maddock downstream depth exponent. + let r = (q_ref.max(1e-3)).powf(0.4).max(0.01); + let v = (1.0 / cfg.reference_n) * r.powf(2.0 / 3.0) * s.sqrt(); + // Clamped to a physical FLOOD-WAVE celerity band, deliberately tighter than + // the solver's own [0.01, 15.0] velocity clamp. The upper bound matters: + // `dx_target = c * dt`, so at dt = 3600 s a celerity of 8.9 m/s (which this + // relation reaches at slope 1e-2) implies a 32 km space step, and every + // shorter reach would be stretched to it. 5 m/s caps `dx_target` at 18 km, + // and `max_clamp_factor` bounds the residual distortion. + (v * (5.0 / 3.0)).clamp(0.05, 5.0) +} + +/// Absolute lower bound on a reach length, in metres. Independent of the clamp +/// so that `min_length_fraction: 0.0` on a zero-length reach still cannot +/// produce `K = L/c = 0`, which would make `c1 = 1` and break the solve. +/// MERIT contains sub-10 m reaches, so this is reachable in practice. +const ABS_MIN_LENGTH_M: f32 = 1.0; + +/// Both sides of the two-sided rule. `pieces[i]` is how many sub-reaches parent +/// `i` becomes; `length_m[i]` is its (possibly clamped) total length, which +/// Task 3's expansion then divides by `pieces[i]`. +pub struct ReachPlan { + /// Piece count per parent reach; always `>= 1`. + pub pieces: Vec, + /// Total (possibly clamped-up) length per parent reach, in metres; always + /// `> 0` — a zero-length reach would give `K = 0` and `c1 = 1`. + pub length_m: Vec, +} + +/// The two-sided rule. Long reaches split; short reaches have their length +/// clamped UP to `dx_target`. Merging short reaches is deliberately not done: +/// a short reach can carry two upstream tributaries or have a parallel +/// tributary joining below it, so collapsing it would destroy junctions. +pub fn plan_reaches( + length_m: &[f32], + slope: &[f32], + uparea_km2: &[f32], + dt_seconds: f32, + cfg: &Subdivision, +) -> ReachPlan { + if !cfg.enabled { + return ReachPlan { + pieces: vec![1; length_m.len()], + length_m: length_m.to_vec(), + }; + } + let mut pieces = Vec::with_capacity(length_m.len()); + let mut lengths = Vec::with_capacity(length_m.len()); + for ((&l, &s), &a) in length_m.iter().zip(slope).zip(uparea_km2) { + let dx_target = reference_celerity(a, s, cfg) * dt_seconds; + // Short reach: stretch it so Cr ~ 1 at the reference flow. This is a + // STATIC constant, unlike the runtime K floor in `enforce_positivity` + // — no gradient path, so it cannot pull `n` toward its bound. + let l_raw = l.max(0.0); + let want = dx_target * cfg.min_length_fraction; + // Never stretch a reach more than `max_clamp_factor` times its true + // length. A zero-length reach has no meaningful factor, so it takes the + // target directly — it is degenerate geometry either way. + let ceiling = if l_raw > 0.0 { + l_raw * cfg.max_clamp_factor + } else { + f32::INFINITY + }; + // ABS_MIN_LENGTH_M is applied LAST so it survives min_length_fraction: 0.0. + let l_eff = l_raw.max(want.min(ceiling)).max(ABS_MIN_LENGTH_M); + // Long reach: split. + let m = ((l_eff / dx_target).ceil().max(1.0) as u32).min(cfg.max_pieces as u32); + pieces.push(m); + lengths.push(l_eff); + } + ReachPlan { + pieces, + length_m: lengths, + } +} + +/// Cost accounting for a [`ReachPlan`], measured against the raw fabric lengths. +/// +/// The length clamp is the price of the short-reach branch: a reach modelled +/// `f×` longer than reality has an `f×` longer travel time, which is a real +/// physical distortion traded for numerical stability. `max_clamp_factor` +/// bounds `f`, and `n_at_clamp_ceiling` counts the reaches pinned at that bound +/// — those stay over-Courant by design and will still produce negative +/// coefficients. +#[derive(Debug, Clone)] +pub struct PlanStats { + /// Parent reaches considered. + pub n: usize, + /// `Σ pieces` — the sub-reach count the solver will see. + pub sum_pieces: u64, + /// Reaches with `m > 1`. + pub n_split: usize, + /// Reaches whose length was stretched (`length_m > raw + 1e-3`). + pub n_clamped: usize, + /// Reaches pinned at `raw * max_clamp_factor` — the clamp wanted more. + pub n_at_clamp_ceiling: usize, + /// Clamp factor `length_m / raw` over the CLAMPED reaches only: p50/p95/p99/max. + pub clamp_factor_p50: f32, + pub clamp_factor_p95: f32, + pub clamp_factor_p99: f32, + pub clamp_factor_max: f32, + /// Total network channel length (m) before and after the clamp. + pub total_length_before_m: f64, + pub total_length_after_m: f64, + /// Histogram of piece counts, `pieces_hist[m]` for `m` in `0..=max_pieces`. + pub pieces_hist: Vec, +} + +impl PlanStats { + /// Fractional inflation of total channel length, e.g. `0.171` for +17.1 %. + pub fn length_inflation(&self) -> f64 { + self.total_length_after_m / self.total_length_before_m - 1.0 + } +} + +/// Measure a [`ReachPlan`] against the RAW fabric lengths it was built from. +/// +/// `raw_length_m` must be `ConusAdjacency::length_m` — the un-clamped input to +/// [`plan_reaches`] — not `plan.length_m`. +pub fn plan_stats(raw_length_m: &[f32], plan: &ReachPlan, cfg: &Subdivision) -> PlanStats { + let n = raw_length_m.len(); + let mut factors: Vec = Vec::new(); + let mut n_at_ceiling = 0usize; + let mut before = 0f64; + let mut after = 0f64; + let mut hist = vec![0usize; cfg.max_pieces.max(1) + 1]; + let mut sum_pieces = 0u64; + let mut n_split = 0usize; + + for i in 0..n { + let raw = raw_length_m[i].max(0.0); + let eff = plan.length_m[i]; + before += raw as f64; + after += eff as f64; + if eff > raw + 1e-3 { + let f = if raw > 0.0 { eff / raw } else { f32::INFINITY }; + factors.push(f); + // Pinned at the ceiling: the clamp wanted `dx_target * + // min_length_fraction` but got `raw * max_clamp_factor`. + if raw > 0.0 && (f - cfg.max_clamp_factor).abs() <= 1e-3 * cfg.max_clamp_factor { + n_at_ceiling += 1; + } + } + let m = plan.pieces[i] as usize; + sum_pieces += m as u64; + if m > 1 { + n_split += 1; + } + if m < hist.len() { + hist[m] += 1; + } + } + + factors.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + let q = |p: f64| -> f32 { + if factors.is_empty() { + return f32::NAN; + } + factors[(((factors.len() - 1) as f64) * p).round() as usize] + }; + + PlanStats { + n, + sum_pieces, + n_split, + n_clamped: factors.len(), + n_at_clamp_ceiling: n_at_ceiling, + clamp_factor_p50: q(0.50), + clamp_factor_p95: q(0.95), + clamp_factor_p99: q(0.99), + clamp_factor_max: factors.last().copied().unwrap_or(f32::NAN), + total_length_before_m: before, + total_length_after_m: after, + pieces_hist: hist, + } +} + +/// The expanded network: one row per sub-reach, plus the map back to parents. +/// +/// Two index spaces coexist. **Parent space** (`parent_order`, length `N`) is +/// where COMID lookups live — `order` gains duplicates after expansion, so a +/// COMID→row lookup on it would be ambiguous. **Sub-reach space** (`order`, +/// `rows`, `cols`, `length_m`, `slope`, length `N' = Σ m_p`) is what the solver +/// sees. Parent `p` owns the contiguous rows `parent_offset[p]..parent_offset[p+1]`, +/// ordered upstream→downstream, so its inlet is the first and its outlet the last. +/// +/// ```text +/// BEFORE AFTER (m = 3) +/// U ──> P ──> D U₂ ──> P₀ ──> P₁ ──> P₂ ──> D₀ +/// └─ q'/3 q'/3 q'/3 +/// len(P) = L len(Pᵢ) = L/3, slope unchanged +/// ``` +#[derive(Debug, Clone)] +pub struct SubdividedAdjacency { + /// COMID per sub-reach row — the parent's COMID, repeated `m_p` times. + pub order: Vec, + /// Original COMID topological order, length `N`. COMID→position indexes + /// (`IdIndex`) must be built from THIS, not from `order`. + pub parent_order: Vec, + /// Length `N + 1`. Rows `[parent_offset[p], parent_offset[p + 1])` belong to + /// parent `p`, ordered upstream→downstream. + pub parent_offset: Vec, + /// COO `indices_0` — downstream sub-reach row. + pub rows: Vec, + /// COO `indices_1` — upstream sub-reach row. + pub cols: Vec, + /// Per-piece length in metres: the parent's (possibly clamped) length / `m`. + pub length_m: Vec, + /// Per-piece slope — inherited unchanged from the parent. + pub slope: Vec, +} + +impl SubdividedAdjacency { + /// First (most upstream) sub-reach row of `parent`. + #[inline] + pub fn inlet(&self, parent: usize) -> usize { + self.parent_offset[parent] as usize + } + + /// Last (most downstream) sub-reach row of `parent`. A gauge on this reach + /// must be read here: any earlier piece omits the downstream fraction of the + /// reach's own lateral inflow. + #[inline] + pub fn outlet(&self, parent: usize) -> usize { + self.parent_offset[parent + 1] as usize - 1 + } + + /// Number of sub-reaches owned by `parent`. + #[inline] + pub fn pieces(&self, parent: usize) -> usize { + (self.parent_offset[parent + 1] - self.parent_offset[parent]) as usize + } +} + +/// Expand `adj` into the sub-reach graph described by `plan`. +/// +/// **Topology rules.** Parent `p` owns rows `[off[p], off[p+1])` ordered +/// upstream→downstream; consecutive pieces are joined by internal chain edges. +/// An external edge `u -> p` (`u` upstream of `p`) becomes +/// `outlet(u) -> inlet(p)`. Each piece gets `plan.length_m[p] / m` — the +/// **clamped** length from `plan_reaches`, not `adj.length_m[p]`. Slope is +/// inherited unchanged. +/// +/// **Why the result is still lower-triangular** (routing invariant 3, which the +/// forward-substitution solver in `src/sparse/` depends on and which fails +/// silently rather than loudly): +/// - Internal edges run `base + k - 1 -> base + k`, strictly increasing. +/// - External edges: the input satisfies `r > c` (`build.rs` rejects `r < c`, +/// and the real CONUS COO carries no `r == c`), so `r >= c + 1` and therefore +/// `inlet(r) = off[r] >= off[c + 1] > off[c + 1] - 1 = outlet(c)`. +/// +/// Because `rows`/`cols` are `downstream`/`upstream` (`zarr.rs:39-42`, +/// `build.rs:231-232`), reversing the mapping to `inlet(u) -> outlet(p)` would +/// still be lower-triangular and would NOT be caught by that argument — the +/// direction is pinned by tests instead. +pub fn subdivide(adj: &ConusAdjacency, plan: &ReachPlan) -> SubdividedAdjacency { + let n = adj.order.len(); + assert_eq!(plan.pieces.len(), n, "pieces must be one per parent reach"); + assert_eq!(plan.length_m.len(), n, "lengths must be one per parent reach"); + + let mut parent_offset = Vec::with_capacity(n + 1); + let mut acc: i32 = 0; + parent_offset.push(0); + for &m in &plan.pieces { + acc += m.max(1) as i32; + parent_offset.push(acc); + } + let n_sub = acc as usize; + + let mut order = Vec::with_capacity(n_sub); + let mut length_m = Vec::with_capacity(n_sub); + let mut slope = Vec::with_capacity(n_sub); + let mut rows = Vec::with_capacity(adj.rows.len() + n_sub - n); + let mut cols = Vec::with_capacity(adj.cols.len() + n_sub - n); + + for p in 0..n { + let m = plan.pieces[p].max(1) as usize; + let base = parent_offset[p] as usize; + // plan.length_m, NOT adj.length_m: short reaches were clamped UP in the + // two-sided rule above, and the expansion must honour that. + let piece_len = plan.length_m[p] / m as f32; + for k in 0..m { + order.push(adj.order[p]); + length_m.push(piece_len); + slope.push(adj.slope[p]); + if k > 0 { + // Internal chain link: piece k-1 flows into piece k. + cols.push((base + k - 1) as i32); + rows.push((base + k) as i32); + } + } + } + + // External edges: upstream parent's OUTLET -> downstream parent's INLET. + for (&r, &c) in adj.rows.iter().zip(adj.cols.iter()) { + // A self-edge (r == c) would map to outlet(p) -> inlet(p), i.e. + // off[p]+m-1 -> off[p], which is UPPER triangular and would break the + // forward-substitution invariant *silently*. `build.rs:236-240` only + // rejects r < c, so r == c is permitted by that check — but the real + // CONUS COO (346,321 reaches / 338,814 edges) has zero of them, the + // diagonal being synthesized later by `CsrPattern::from_sparse`. This is + // therefore an assertion documenting that fact, not a filter. + assert!(r != c, "self-edge on parent {r}: subdivision cannot expand it"); + let up_outlet = parent_offset[c as usize + 1] - 1; + let down_inlet = parent_offset[r as usize]; + cols.push(up_outlet); + rows.push(down_inlet); + } + + SubdividedAdjacency { + order, + parent_order: adj.order.clone(), + parent_offset, + rows, + cols, + length_m, + slope, + } +} diff --git a/src/adjacency/validate.rs b/src/adjacency/validate.rs index afd8878..52c5665 100644 --- a/src/adjacency/validate.rs +++ b/src/adjacency/validate.rs @@ -64,6 +64,40 @@ pub fn validate_gages_store_layout(store: &Path) -> Result<(), StoreLayoutError> Ok(()) } +/// Whether an already-built CONUS store holds a genuine subdivision map. +/// +/// Cheap: reads the declared `shape` out of `/order/zarr.json` and +/// `/parent_order/zarr.json` — two small metadata files. The 346K-element +/// arrays themselves are never opened. +/// +/// - `Some(true)` — `n_parent < n`, i.e. the store was built by the managed +/// builder with `params.subdivision.enabled: true`. +/// - `Some(false)` — the store is readable but un-subdivided: either it predates +/// subdivision (no `/parent_order` at all) or it carries the identity map +/// (`n_parent == n`). +/// - `None` — the shapes could not be read (missing/corrupt store), so +/// the question cannot be answered. +/// +/// Used by `config::validate_subdivision` to tell "explicit path to an +/// already-subdivided store" (legitimate) from "explicit path to an +/// un-subdivided store while asking for subdivision" (silently inert, rejected). +pub fn store_is_subdivided(store: &Path) -> Option { + let n = declared_len(&store.join("order").join("zarr.json"))?; + match declared_len(&store.join("parent_order").join("zarr.json")) { + // Pre-subdivision store: no parent map at all. + None if !store.join("parent_order").join("zarr.json").is_file() => Some(false), + None => None, + Some(n_parent) => Some(n_parent < n), + } +} + +/// First entry of a zarr v3 array's declared `shape`, from its `zarr.json`. +fn declared_len(zarr_json: &Path) -> Option { + let text = std::fs::read_to_string(zarr_json).ok()?; + let meta: serde_json::Value = serde_json::from_str(&text).ok()?; + meta.get("shape")?.get(0)?.as_u64().map(|v| v as usize) +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/adjacency/zarr_write.rs b/src/adjacency/zarr_write.rs index cc9bcce..b7c51a2 100644 --- a/src/adjacency/zarr_write.rs +++ b/src/adjacency/zarr_write.rs @@ -62,6 +62,7 @@ use zarrs::group::GroupBuilder; use crate::adjacency::build::ConusAdjacency; use crate::adjacency::gauges::GaugeSubgraph; +use crate::adjacency::subdivide::SubdividedAdjacency; use crate::data::error::{DataError, Result}; /// Write the CONUS adjacency store at `dest` (a directory path). @@ -93,6 +94,44 @@ pub fn write_conus_store(adj: &ConusAdjacency, dest: &Path) -> Result<()> { Ok(()) } +/// Write a **subdivided** CONUS store at `dest`: the same array set as +/// [`write_conus_store`] over the sub-reach rows, plus the parent map +/// `/parent_order` (int32, `[n_parent]`) and `/parent_offset` (int32, +/// `[n_parent + 1]`). +/// +/// The root `shape` attr is `[n_sub, n_sub]` — the COO the reader indexes is in +/// sub-reach space. `parent_order` is the only place the un-duplicated COMID +/// list survives, and `ConusAdjacencyStore::open` builds its `IdIndex` from it. +/// +/// A store written from an all-`m = 1` plan is byte-identical to +/// [`write_conus_store`] plus the two identity arrays, which the reader would +/// have synthesized anyway — so the disabled path stays a true no-op. +pub fn write_conus_store_subdivided(sub: &SubdividedAdjacency, dest: &Path) -> Result<()> { + let storage = Arc::new(FilesystemStore::new(dest).map_err(|e| zarr_err(dest, e))?); + + let n_sub = sub.order.len(); + let nnz = sub.rows.len(); + + let root = GroupBuilder::new() + .attributes(coo_root_attrs(n_sub)) + .build(storage.clone(), "/") + .map_err(|e| zarr_err(dest, e))?; + root.store_metadata().map_err(|e| zarr_err(dest, e))?; + + write_i32(&storage, dest, "/indices_0", &sub.rows)?; + write_i32(&storage, dest, "/indices_1", &sub.cols)?; + write_u8_ones(&storage, dest, "/values", nnz)?; + write_i32(&storage, dest, "/order", &sub.order)?; + write_f32(&storage, dest, "/length_m", &sub.length_m)?; + write_f32(&storage, dest, "/slope", &sub.slope)?; + // Subdivision addition: the parent map. Absent from pre-subdivision stores, + // which the reader handles by synthesizing the identity. + write_i32(&storage, dest, "/parent_order", &sub.parent_order)?; + write_i32(&storage, dest, "/parent_offset", &sub.parent_offset)?; + + Ok(()) +} + /// Write the per-gauge adjacency store at `dest` (a directory path): an empty /// root group with one subgroup per STAID. /// diff --git a/src/bin/probe_courant.rs b/src/bin/probe_courant.rs new file mode 100644 index 0000000..a6cd28a --- /dev/null +++ b/src/bin/probe_courant.rs @@ -0,0 +1,558 @@ +//! Probe: measure what `params.enforce_positivity` (the S18'/S19' clamp) +//! actually does on REAL CONUS data, with no training. +//! +//! Reports, for `enforce_positivity` false and true over the SAME network, +//! parameters, forcing and hot-start: +//! +//! 1. `negative solves before clamp` — count / percentage, exact over every +//! routed timestep (the same atomic counters the training log prints). +//! 2. `Cr = dt / K` percentiles, before and after the K floor, plus the +//! fraction of reach-timesteps with `Cr_raw > 2` (where the floor bites). +//! 3. The effective Muskingum `X` percentiles. +//! 4. `k_musk / k_raw` — travel-time inflation, over the floored reaches. +//! +//! It builds the routing engine exactly as `training::forward` does (KAN head +//! -> setup_inputs -> hot-start) and then drives `forward_chain_inner` in the +//! same q_next-fed-back loop as `MuskingumCunge::forward`, so the numbers come +//! from the production kernel sequence rather than a re-derivation. +//! +//! ```bash +//! cargo run --release --bin probe_courant -- \ +//! --config ddrs.yaml --backend cuda \ +//! --checkpoint .ddrs/runs//checkpoints/epoch_5_mb_18 \ +//! --gauges 64 --steps 500 +//! ``` + +use std::path::PathBuf; + +use burn::backend::Autodiff; +use burn::tensor::{backend::Backend, Tensor}; +use clap::Parser; +use chrono::Duration; + +use ddrs::config::{Config, ConfigMode, SparseSolver}; +use ddrs::data::dataset::MeritGagesDataset; +use ddrs::data::dates::RhoWindow; +use ddrs::routing::courant_probe::{run_courant_probe, CourantReport}; +use ddrs::routing::denormalize; +use ddrs::routing::{MuskingumCunge, RoutingInputs, SpatialParameters}; +use ddrs::training::bootstrap_head_and_state; + +#[derive(Parser, Debug)] +#[command(name = "probe_courant", about = "Measure Cr / X / K-floor / negative solves, clamp off vs on")] +struct Cli { + /// Training YAML (e.g. `ddrs.yaml`). Loaded in TRAINING mode. + #[arg(long)] + config: PathBuf, + + /// Checkpoint DIRECTORY (`.../epoch_E_mb_M`) to take the KAN head from. + /// Omit to probe the seed-initialised head. + #[arg(long)] + checkpoint: Option, + + /// Gauges in the mini-batch (matches `experiment.batch_size` semantics). + #[arg(long, default_value_t = 64)] + gauges: usize, + + /// First day index of the rho window into the experiment time axis. + #[arg(long, default_value_t = 0)] + start_day: usize, + + /// Rho window length in days. Hourly steps available = (rho - 1) * 24. + #[arg(long, default_value_t = 30)] + rho: usize, + + /// Hourly timesteps to route. Capped by the window. + #[arg(long, default_value_t = 500)] + steps: usize, + + /// Harvest the Cr / X / K distributions every Nth timestep. The solve + /// counters are exact over ALL timesteps regardless. + #[arg(long, default_value_t = 10)] + sample_every: usize, + + /// "cpu" (NdArray) or "cuda". + #[arg(long, default_value = "cuda")] + backend: String, + + /// Optional CSV of the percentile table. + #[arg(long)] + output: Option, + + // ── reach subdivision ──────────────────────────────────────────────────── + /// MERIT flowlines fabric (`.shp`/`.dbf`/`.gpkg`). Switches the run to the + /// MANAGED adjacency build, overriding `data_sources.conus_adjacency` / + /// `gages_adjacency`. Required for `--max-pieces` and `--clamp-report`, + /// because subdivision only happens inside that builder. + #[arg(long)] + fabric: Option, + + /// Layer name for a multi-layer `.gpkg` fabric. + #[arg(long)] + fabric_layer: Option, + + /// Workspace root holding `adjacency//` build caches. Note: subdivided + /// stores are cached separately per cap, so re-running a cap is instant. + #[arg(long, default_value = ".ddrs")] + workspace: PathBuf, + + /// Enable subdivision with this `max_pieces` cap. Omit for the un-split + /// control. Implies `--fabric`. + #[arg(long)] + max_pieces: Option, + + /// Override `params.subdivision.reference_n`. The default 0.05 is a + /// *guess* at the trained CONUS median; the reference celerity scales as + /// `1/n`, so this directly sets `dx_target` and therefore both the piece + /// count and the clamped fraction. Sweep it against the checkpoint's actual + /// median `n` before concluding anything about subdivision. + #[arg(long)] + reference_n: Option, + + /// Override `params.subdivision.min_length_fraction` (short-reach clamp + /// target, as a fraction of `dx_target`; 0 disables the clamp). + #[arg(long)] + min_length_fraction: Option, + + /// Report the reach-plan cost (piece histogram, clamped fraction, + /// clamp-factor percentiles, total length inflation) and exit without + /// routing. Requires `--fabric`; `--max-pieces` selects the cap. + #[arg(long, default_value_t = false)] + clamp_report: bool, + + /// Divide the cold-start `q'_0` by each row's piece count, so the hot-start + /// `(I − N)·Q_0 = q'_0` solve sees the same forcing `forward` routes. + /// Without it the subdivided initial condition is inflated ~m× per reach. + #[arg(long, default_value_t = false)] + divide_hotstart: bool, + + /// Trace the first N timesteps of basin-outlet discharge to + /// `<--output>.trace.csv`, for measuring hot-start wash-out. 0 = off. + #[arg(long, default_value_t = 0)] + trace_steps: usize, +} + +type R = Result>; + +fn main() -> R<()> { + let cli = Cli::parse(); + match cli.backend.as_str() { + "cpu" => { + type I = burn::backend::NdArray; + let device = ::Device::default(); + run::(cli, device) + } + "cuda" => { + type I = burn_cuda::Cuda; + let device = cubecl::cuda::CudaDevice::new(0); + run::(cli, device) + } + other => Err(format!("unknown --backend {other}").into()), + } +} + +fn quantile(sorted: &[f32], q: f64) -> f32 { + if sorted.is_empty() { + return f32::NAN; + } + let idx = ((sorted.len() - 1) as f64 * q).round() as usize; + sorted[idx] +} + +struct Summary { + label: &'static str, + rep: CourantReport, +} + +/// Report the reach-plan cost on the real fabric without writing a store. +/// +/// This is the price of the SHORT-reach branch of the two-sided rule: reaches +/// stretched up to `dx_target` (bounded by `max_clamp_factor`) get a +/// proportionally longer travel time, which is a physical distortion traded for +/// numerical stability. Reaches pinned AT the ceiling stay over-Courant. +fn clamp_report( + cli: &Cli, + fabric: &std::path::Path, + subdivision: &ddrs::config::Subdivision, + attributes: &[PathBuf], +) -> R<()> { + use ddrs::adjacency::cache::{parent_adjacency_from_fabric, reach_plan}; + use ddrs::adjacency::subdivide::plan_stats; + + let t0 = std::time::Instant::now(); + let conus = parent_adjacency_from_fabric(fabric, cli.fabric_layer.as_deref()) + .map_err(|e| format!("read fabric: {e}"))?; + eprintln!( + "fabric: {} parents, {} edges ({:.1}s)", + conus.order.len(), + conus.rows.len(), + t0.elapsed().as_secs_f64() + ); + + let plan = reach_plan(&conus, subdivision, attributes) + .map_err(|e| format!("reach plan: {e}"))?; + let st = plan_stats(&conus.length_m, &plan, subdivision); + + println!("\n================ reach plan (max_pieces = {}, min_length_fraction = {}, max_clamp_factor = {}) ================", + subdivision.max_pieces, subdivision.min_length_fraction, subdivision.max_clamp_factor); + println!("parents : {}", st.n); + println!( + "sub-reaches (Σm) : {} ({:.3}x)", + st.sum_pieces, + st.sum_pieces as f64 / st.n as f64 + ); + println!( + "reaches split (m > 1) : {} ({:.2}%)", + st.n_split, + 100.0 * st.n_split as f64 / st.n as f64 + ); + print!("piece histogram :"); + for (m, &c) in st.pieces_hist.iter().enumerate().skip(1) { + print!(" m={m}:{c}"); + } + println!(); + println!( + "reaches length-clamped : {} ({:.2}%)", + st.n_clamped, + 100.0 * st.n_clamped as f64 / st.n as f64 + ); + println!( + " pinned at max_clamp_factor={:.1}: {} ({:.2}% of all, {:.2}% of clamped) \ + — these stay over-Courant BY DESIGN", + subdivision.max_clamp_factor, + st.n_at_clamp_ceiling, + 100.0 * st.n_at_clamp_ceiling as f64 / st.n as f64, + 100.0 * st.n_at_clamp_ceiling as f64 / st.n_clamped.max(1) as f64 + ); + println!( + " clamp factor p50/p95/p99/max: {:.3} / {:.3} / {:.3} / {:.3}", + st.clamp_factor_p50, st.clamp_factor_p95, st.clamp_factor_p99, st.clamp_factor_max + ); + println!( + "total channel length : {:.1} km → {:.1} km ({:+.2}%)", + st.total_length_before_m / 1000.0, + st.total_length_after_m / 1000.0, + 100.0 * st.length_inflation() + ); + Ok(()) +} + +fn pct(v: &mut Vec) -> [f32; 5] { + v.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + [ + quantile(v, 0.05), + quantile(v, 0.25), + quantile(v, 0.50), + quantile(v, 0.75), + quantile(v, 0.95), + ] +} + +fn run(cli: Cli, device: I::Device) -> R<()> +where + I::FloatTensorPrimitive: 'static, + I::Device: 'static, +{ + let mut cfg = Config::from_yaml_file_with_mode(&cli.config, ConfigMode::Training) + .map_err(|e| format!("load config: {e}"))?; + if cli.backend == "cpu" { + cfg.params.sparse_solver = SparseSolver::Cpu; + } + // The graph path never enters `forward_chain_inner`, so it can neither be + // counted nor probed. + cfg.params.use_cuda_graphs = false; + if cfg.params.ddr_match { + return Err("this probe requires params.ddr_match: false (enforce_positivity is gated on it)".into()); + } + if let (Some(exp), Some(ckpt)) = (cfg.experiment.as_mut(), cli.checkpoint.clone()) { + exp.checkpoint = Some(ckpt); + } + + // --- reach subdivision ------------------------------------------------- + // Subdivision happens inside the MANAGED adjacency builder, which is only + // reached when `data_sources` carries no explicit adjacency paths. So + // `--fabric` swaps the config over to the managed path and splices the + // resolved (possibly subdivided) store paths back in. + if let Some(fabric) = cli.fabric.clone() { + cfg.params.subdivision.enabled = cli.max_pieces.is_some(); + if let Some(m) = cli.max_pieces { + cfg.params.subdivision.max_pieces = m; + } + if let Some(v) = cli.reference_n { + cfg.params.subdivision.reference_n = v; + } + if let Some(v) = cli.min_length_fraction { + cfg.params.subdivision.min_length_fraction = v; + } + let ds = cfg.data_sources.as_mut().ok_or("config has no data_sources")?; + ds.conus_adjacency = None; + ds.gages_adjacency = None; + ds.geospatial_fabric = Some(fabric.clone()); + ds.geospatial_fabric_layer = cli.fabric_layer.clone(); + let (gages_csv, attributes) = (ds.gages.clone(), ds.attributes.clone()); + + if cli.clamp_report { + return clamp_report(&cli, &fabric, &cfg.params.subdivision, &attributes); + } + + let outcome = ddrs::adjacency::cache::resolve_or_build( + &cli.workspace, + &fabric, + cli.fabric_layer.as_deref(), + &gages_csv, + &cfg.params.subdivision, + &attributes, + ) + .map_err(|e| format!("managed adjacency build: {e}"))?; + eprintln!( + "adjacency: {} (cache {})", + outcome.paths.conus.display(), + if outcome.cache_hit { "hit" } else { "miss" } + ); + let ds = cfg.data_sources.as_mut().unwrap(); + ds.conus_adjacency = Some(outcome.paths.conus); + ds.gages_adjacency = Some(outcome.paths.gages); + } else if cli.max_pieces.is_some() || cli.clamp_report { + return Err("--max-pieces / --clamp-report require --fabric: subdivision \ + only runs inside the managed adjacency builder" + .into()); + } + + let dataset = MeritGagesDataset::open(&cfg).map_err(|e| format!("open dataset: {e}"))?; + let axis = dataset.time_axis().clone(); + let staids: Vec<_> = dataset.staids().iter().take(cli.gauges).cloned().collect(); + let window = RhoWindow { + start_day_idx: cli.start_day, + rho_days: cli.rho, + window_start: axis.start + Duration::days(cli.start_day as i64), + }; + eprintln!( + "batch: {} gauges, window {} .. +{} d (day idx {})", + staids.len(), + window.window_start, + cli.rho, + cli.start_day + ); + + let batch = dataset + .collate(&staids, &window) + .map_err(|e| format!("collate: {e}"))?; + let tensors = batch.to_tensors::>(&device); + let n_active = tensors.adjacency.n; + eprintln!("network: {n_active} reaches"); + + // Head exactly as the trainer bootstraps it (disagg warm start + resume). + let (head, _state, _optim) = bootstrap_head_and_state::(&cfg, &device) + .map_err(|e| format!("bootstrap head: {e}"))?; + + // The head runs at PARENT resolution; expand onto the routing's sub-reach + // rows exactly as `training::forward` does. No-op when not subdivided. + let params_map = ddrs::training::forward::gather_params_to_subreaches( + head.forward(tensors.spatial_attributes.clone()), + tensors.adjacency.parent_offset.as_ref(), + n_active, + &device, + ); + let n_param = params_map.get("n").expect("head missing n").clone(); + let q_param = params_map.get("q_spatial").expect("head missing q_spatial").clone(); + let p_param = params_map.get("p_spatial").cloned(); + let x_storage: Tensor, 1> = match params_map.get("x_storage") { + Some(x) => denormalize( + x.clone(), + cfg.params.parameter_ranges.x_storage, + cfg.params.log_space_parameters.iter().any(|s| s == "x_storage"), + ), + None => Tensor::full([n_active], 0.3_f32, &device), + }; + + // Same forcing the trainer would route (disagg head when configured). + let n_hourly = tensors.q_prime.dims()[0]; + let q_prime_hourly = match &head.disagg { + Some(d) => d.forward( + tensors.q_prime_daily.clone(), + tensors.precip_hourly.clone(), + n_hourly, + ), + None => tensors.q_prime.clone(), + }; + + let mut summaries: Vec = Vec::new(); + for (label, enforce) in [("off", false), ("on", true)] { + let mut cfg_v = cfg.clone(); + cfg_v.params.enforce_positivity = enforce; + + let mut engine = MuskingumCunge::::new(cfg_v.clone(), device.clone()); + engine.divide_hotstart_by_pieces = cli.divide_hotstart; + engine.setup_inputs( + RoutingInputs { + adjacency: tensors.adjacency.clone(), + x_storage: x_storage.clone(), + }, + q_prime_hourly.clone(), + SpatialParameters { + n: n_param.clone(), + q_spatial: q_param.clone(), + p_spatial: p_param.clone(), + k_d: None, + d_gw: None, + leakance_factor: None, + impervious_mask: None, + }, + false, + tensors.initial_state.clone(), + ); + let inp = engine.probe_inputs(); + eprintln!("--- routing with enforce_positivity: {enforce} ---"); + let t0 = std::time::Instant::now(); + let rep = run_courant_probe::(&cfg_v, &inp, cli.steps, cli.sample_every, cli.trace_steps); + let secs = t0.elapsed().as_secs_f64(); + eprintln!( + " {} steps, {} sampled, negative solves {}/{} — {:.2}s ({:.2} ms/step)", + rep.n_steps, + rep.n_sampled_steps, + rep.neg_solves, + rep.total_solves, + secs, + 1000.0 * secs / rep.n_steps.max(1) as f64 + ); + summaries.push(Summary { label, rep }); + } + + // ---- Report ------------------------------------------------------------- + let mut csv = String::from("flag,metric,p5,p25,p50,p75,p95,extra\n"); + for s in &mut summaries { + let r = &mut s.rep; + let pc = 100.0 * r.neg_solves as f64 / r.total_solves.max(1) as f64; + println!("\n================ enforce_positivity: {} ================", s.label); + println!( + "network {} reaches, {} routed timesteps, {} sampled steps ({} reach-timesteps sampled)", + r.n_reaches, + r.n_steps, + r.n_sampled_steps, + r.cr.len() + ); + println!( + "negative solves before clamp: {}/{} ({:.4}%)", + r.neg_solves, r.total_solves, pc + ); + + let n_s = r.cr.len() as f64; + let frac_cr_raw_gt2 = r.cr_raw.iter().filter(|&&v| v > 2.0).count() as f64 / n_s; + let frac_cr_raw_lt05 = r.cr_raw.iter().filter(|&&v| v < 0.5).count() as f64 / n_s; + let frac_floored = r.k_ratio.iter().filter(|&&v| v > 1.0 + 1e-6).count() as f64 / n_s; + let mut floored_ratio: Vec = + r.k_ratio.iter().copied().filter(|&v| v > 1.0 + 1e-6).collect(); + let cap_binds = r + .x_eff + .iter() + .zip(r.x_cunge.iter()) + .filter(|(e, c)| **e < **c - 1e-6) + .count() as f64 + / n_s; + let neg_c1 = r.c1.iter().filter(|&&v| v < 0.0).count(); + let neg_c3 = r.c3.iter().filter(|&&v| v < 0.0).count(); + // The whole claim under test: "Cr ~ 1 makes c1 and c3 non-negative by + // construction". Both are >= 0 exactly when Cr lands inside the + // Muskingum window `2X <= Cr <= 2(1-X)`, whose width is `2(1-2X)` — it + // COLLAPSES as X -> 0.5. So report the window width alongside the hit + // rate; a 1%-wide window cannot be hit by a static piece count. + let both_ok = r + .c1 + .iter() + .zip(r.c3.iter()) + .filter(|(a, b)| **a >= 0.0 && **b >= 0.0) + .count() as f64 + / n_s; + let mut window: Vec = r.x_cunge.iter().map(|&x| 2.0 * (1.0 - 2.0 * x)).collect(); + let q_win = pct(&mut window); + let min_c1 = r.c1.iter().copied().fold(f32::INFINITY, f32::min); + let min_c3 = r.c3.iter().copied().fold(f32::INFINITY, f32::min); + + let q_cr_raw = pct(&mut r.cr_raw); + let q_cr = pct(&mut r.cr); + let q_x = pct(&mut r.x_eff); + let q_xc = pct(&mut r.x_cunge); + + println!("\n{:<22} {:>9} {:>9} {:>9} {:>9} {:>9}", "", "p5", "p25", "p50", "p75", "p95"); + let row = |name: &str, q: [f32; 5]| { + println!( + "{name:<22} {:>9.4} {:>9.4} {:>9.4} {:>9.4} {:>9.4}", + q[0], q[1], q[2], q[3], q[4] + ); + }; + row("Cr_raw = dt/k_raw", q_cr_raw); + row("Cr = dt/k_musk", q_cr); + row("X_cunge (pre-cap)", q_xc); + row("X_eff (used)", q_x); + if !floored_ratio.is_empty() { + row("k_musk/k_raw | floored", pct(&mut floored_ratio)); + } else { + println!("{:<22} {:>9}", "k_musk/k_raw | floored", "n/a (none floored)"); + } + println!( + "\nfrac Cr_raw > 2 : {:.4} ({} of {})", + frac_cr_raw_gt2, + r.cr_raw.iter().filter(|&&v| v > 2.0).count(), + r.cr_raw.len() + ); + println!( + "frac Cr_raw < 0.5 : {:.4} ({} of {})", + frac_cr_raw_lt05, + r.cr_raw.iter().filter(|&&v| v < 0.5).count(), + r.cr_raw.len() + ); + println!("frac K floored : {frac_floored:.4}"); + println!("frac X cap binds : {cap_binds:.4} (x_eff < x_cunge)"); + println!("min c1 = {min_c1:.3e} (c1 < 0: {neg_c1}, frac {:.4})", neg_c1 as f64 / n_s); + println!("min c3 = {min_c3:.3e} (c3 < 0: {neg_c3}, frac {:.4})", neg_c3 as f64 / n_s); + println!("frac c1>=0 AND c3>=0 : {both_ok:.4}"); + println!( + "non-neg window 2(1-2X) p5/p50/p95: {:.5} / {:.5} / {:.5} \ + (Cr must land in [2X, 2-2X] for BOTH coefficients)", + q_win[0], q_win[2], q_win[4] + ); + + let mut push = |metric: &str, q: [f32; 5], extra: String| { + csv.push_str(&format!( + "{},{},{},{},{},{},{},{}\n", + s.label, metric, q[0], q[1], q[2], q[3], q[4], extra + )); + }; + push( + "cr_raw", + q_cr_raw, + format!("frac_gt2={frac_cr_raw_gt2} frac_lt0.5={frac_cr_raw_lt05}"), + ); + push("cr", q_cr, String::new()); + push("x_cunge", q_xc, String::new()); + push("x_eff", q_x, format!("frac_cap_binds={cap_binds}")); + csv.push_str(&format!( + "{},negative_solves,,,,,,{}/{} ({:.4}%)\n", + s.label, r.neg_solves, r.total_solves, pc + )); + csv.push_str(&format!( + "{},coeff_min,,,,,,min_c1={min_c1:.3e} min_c3={min_c3:.3e} neg_c1={neg_c1} neg_c3={neg_c3} frac_both_nonneg={both_ok:.4} frac_k_floored={frac_floored:.4}\n", + s.label + )); + } + + if let Some(p) = cli.output { + std::fs::write(&p, csv).map_err(|e| format!("write {}: {e}", p.display()))?; + eprintln!("wrote {}", p.display()); + + // Hot-start wash-out trace: total network discharge per timestep, one + // column per enforce_positivity arm. Index 0 is the cold-start Q_0. + if cli.trace_steps > 0 { + let mut t = String::from("step,flag,total_q\n"); + for s in &summaries { + for (i, v) in s.rep.trace_total_q.iter().enumerate() { + t.push_str(&format!("{i},{},{v}\n", s.label)); + } + } + let tp = p.with_extension("trace.csv"); + std::fs::write(&tp, t).map_err(|e| format!("write {}: {e}", tp.display()))?; + eprintln!("wrote {}", tp.display()); + } + } + Ok(()) +} diff --git a/src/bin/probe_n_slope.rs b/src/bin/probe_n_slope.rs new file mode 100644 index 0000000..be03dee --- /dev/null +++ b/src/bin/probe_n_slope.rs @@ -0,0 +1,361 @@ +//! Probe: does the routing objective actually want a stronger Manning's-`n` +//! vs basin-scale relationship, or is the near-flat field the trained head +//! produces already optimal? +//! +//! # Why this exists +//! +//! Across the 2026-07-30..08-02 CONUS run series every trained head converged +//! to a nearly scale-independent `n` field with a faint positive tilt against +//! `log10_uparea` (Spearman +0.205 for the best run, +0.076 for the gentlest, +//! versus -0.230 at initialization). Two incompatible explanations fit that +//! observation equally well: +//! +//! 1. UNDER-TRAINED — the objective wants a stronger slope but 30-50 +//! optimizer steps under a count-dominated loss never got there. +//! 2. ALREADY OPTIMAL — a near-flat field IS this objective's optimum, and +//! every optimizer / learning-rate / epoch experiment was doomed. +//! +//! Distinguishing them by training costs ~13 h per attempt. This binary does +//! it with forward passes only: it reads the trained parameter field, rewrites +//! ONLY the component of `n` that is linear in `log10_uparea`, and re-scores. +//! `p_spatial` and `q_spatial` are held at their trained values throughout, so +//! the slope is the single manipulated variable. +//! +//! # Reading the result +//! +//! * amplified fields (k > 1) score BETTER than k=1 -> hypothesis 1; the run +//! is under-trained and more steps / rebalanced sampling should help. +//! * amplified fields score WORSE -> hypothesis 2; stop tuning the optimizer +//! and change the objective, the metric, or the timestep. +//! * `floor` (n pinned at the range minimum = identity routing) scores about +//! the same as k=1 -> the pooled metric carries no routing signal at all, +//! which invalidates ranking models by it. +//! +//! # Caveats +//! +//! * `EvalParams::Frozen` does NOT support cross-chunk state injection +//! (`eval.rs`): it cold-restarts from hotstart every chunk, whereas the +//! KanHead path carries state. So k=1 here will NOT exactly reproduce the +//! training run's headline NSE. That is expected and harmless — every field +//! below goes through the identical path, so the COMPARISON is clean. Use +//! k=1 as the internal reference, never the run manifest's number. +//! * Full 15-year eval costs ~4.5 h PER FIELD. Default here is a 2-water-year +//! window (~0.6 h/field). Widen with --start/--end once a direction shows. +//! +//! ```bash +//! cargo run --release --bin probe_n_slope -- \ +//! --config .ddrs/runs//config.yaml \ +//! --checkpoint .ddrs/runs//checkpoints/epoch_30_mb_0/head \ +//! --output /tmp/n_slope_probe.csv +//! ``` + +use std::path::PathBuf; + +use burn::tensor::{backend::Backend, Tensor}; +use clap::Parser; + +use ddrs::config::{kan_config, Config, ConfigMode, SparseSolver}; +use ddrs::data::dataset::MeritGagesDataset; +use ddrs::data::test_window::TestWindow; +use ddrs::nn::kan_head::KanHead; +use ddrs::routing::denormalize; +use ddrs::training::checkpoint::load_kan_head; +use ddrs::training::forward::FrozenParams; +use ddrs::training::{evaluate, EvalParams}; + +#[derive(Parser, Debug)] +#[command(name = "probe_n_slope", about = "Score hand-built Manning's-n fields without training")] +struct Cli { + /// Training YAML from the run whose head you are probing. + #[arg(long)] + config: PathBuf, + + /// Trained KAN checkpoint base path (no `.mpk`), e.g. `.../epoch_30_mb_0/head`. + #[arg(long)] + checkpoint: PathBuf, + + /// Eval window start (YYYY/MM/DD). Defaults to the config's `testing.start_time`. + #[arg(long)] + start: Option, + + /// Eval window end (YYYY/MM/DD). Default trims to two water years from `start` + /// so the probe costs ~0.6 h per field instead of ~4.5 h. + #[arg(long)] + end: Option, + + /// Slope multipliers to score. `1` reproduces the trained field; `0` removes + /// the uparea-linear component while keeping residual spatial structure. + #[arg(long, default_value = "0,1,2,5,10", value_delimiter = ',')] + amplify: Vec, + + /// Impose `n` as a monotone ramp in `log10_uparea` between two physical + /// endpoints, e.g. `0.02:0.08` (smallest reach -> largest). Repeatable. + /// + /// Why this exists alongside --amplify: when the trained field has + /// collapsed onto a range bound, amplifying its slope has no headroom + /// (clamping eats the downward half and the surviving linear term is + /// negligible), so --amplify cannot test whether a scale relationship + /// would help. An imposed ramp asks that question directly, independent + /// of where training landed. Reaches are ranked by `log10_uparea`, so the + /// ramp is robust to that attribute's distribution. + #[arg(long = "impose", value_name = "LO:HI")] + impose: Vec, + + /// Skip the spatially flat control (flat `n` at the trained median). + #[arg(long, action = clap::ArgAction::SetTrue)] + no_flat: bool, + + /// Skip the identity-routing control (`n` pinned at the range floor). + #[arg(long, action = clap::ArgAction::SetTrue)] + no_floor: bool, + + /// Days per chunk. Default 15 matches DDR's test config. + #[arg(long, default_value_t = 15)] + batch_size_days: usize, + + /// CSV output path. + #[arg(long)] + output: Option, + + /// "cpu" (NdArray, deterministic) or "cuda". + #[arg(long, default_value = "cuda")] + backend: String, +} + +type R = Result>; + +fn main() -> R<()> { + let cli = Cli::parse(); + match cli.backend.as_str() { + "cpu" => { + type I = burn::backend::NdArray; + let device = ::Device::default(); + run::(cli, device) + } + "cuda" => { + type I = burn_cuda::Cuda; + let device = cubecl::cuda::CudaDevice::new(0); + run::(cli, device) + } + other => Err(format!("unknown --backend {other} (expected \"cpu\" or \"cuda\")").into()), + } +} + +/// One field to score: a label plus the per-reach physical `n` vector. +struct Field { + label: String, + n: Vec, +} + +fn median(v: &[f32]) -> f32 { + let mut s: Vec = v.iter().copied().filter(|x| x.is_finite()).collect(); + if s.is_empty() { + return f32::NAN; + } + s.sort_unstable_by(|a, b| a.partial_cmp(b).expect("finite")); + let m = s.len() / 2; + if s.len() % 2 == 0 { + 0.5 * (s[m - 1] + s[m]) + } else { + s[m] + } +} + +fn run(cli: Cli, device: I::Device) -> R<()> { + let mut cfg = Config::from_yaml_file_with_mode(&cli.config, ConfigMode::Testing) + .map_err(|e| format!("load config in testing mode: {e}"))?; + if cli.backend == "cpu" { + cfg.params.sparse_solver = SparseSolver::Cpu; + cfg.params.use_cuda_graphs = false; + eprintln!("backend: cpu (NdArray, deterministic; sparse_solver forced to cpu, cuda graphs off)"); + } else { + eprintln!( + "backend: cuda (sparse_solver={:?}, use_cuda_graphs={})", + cfg.params.sparse_solver, cfg.params.use_cuda_graphs + ); + } + + // Narrow the eval window. Full CONUS eval is ~4.5 h per field; the probe + // only needs enough record to rank fields, not to publish a metric. + { + let exp = cfg.experiment.as_mut().expect("experiment section"); + if let Some(s) = cli.start.clone() { + exp.start_time = s; + } + exp.end_time = match cli.end.clone() { + Some(e) => e, + None => { + let y: i32 = exp.start_time[0..4].parse().map_err(|e| format!("parse start year: {e}"))?; + format!("{}/09/30", y + 2) + } + }; + eprintln!("eval window: {} -> {}", exp.start_time, exp.end_time); + } + + let dataset = MeritGagesDataset::open(&cfg).map_err(|e| format!("open dataset: {e}"))?; + + // Probe the eval network once to get its reach count and attribute matrix. + // `evaluate` builds the same network internally, in the same order. + let axis = dataset.time_axis().clone(); + let probe = dataset + .collate_window(&TestWindow::new(&axis, 0, 1)) + .map_err(|e| format!("probe collate: {e}"))?; + let n_reaches = probe.divide_comids.len(); + let tensors = probe.to_tensors::(&device); + eprintln!("eval network: {n_reaches} reaches"); + + // ---- Trained parameter field ------------------------------------------- + ::seed(&device, cfg.seed); + let head_cfg = cfg.kan_head.as_ref().expect("kan_head section"); + let template: KanHead = kan_config(head_cfg, cfg.seed).init::(&device); + let head = load_kan_head::(&cli.checkpoint, template, &device) + .map_err(|e| format!("load checkpoint: {e}"))?; + + // Gathered to sub-reach rows: everything below is indexed by `n_reaches` + // (= `divide_comids.len()`), while the head itself runs at parent + // resolution. No-op unless the adjacency was built with + // `params.subdivision.enabled`. + let raw = ddrs::training::forward::gather_params_to_subreaches( + head.forward(tensors.spatial_attributes.clone()), + tensors.adjacency.parent_offset.as_ref(), + n_reaches, + &device, + ); + let is_log = |k: &str| cfg.params.log_space_parameters.iter().any(|s| s == k); + let to_vec = |t: Tensor| -> Vec { t.into_data().to_vec::().unwrap() }; + + let ranges = &cfg.params.parameter_ranges; + let n_trained = to_vec(denormalize(raw["n"].clone(), ranges.n, is_log("n"))); + let q_trained = to_vec(denormalize( + raw["q_spatial"].clone(), + ranges.q_spatial, + is_log("q_spatial"), + )); + let p_trained = if head_cfg.learnable_parameters.iter().any(|s| s == "p_spatial") { + to_vec(denormalize( + raw["p_spatial"].clone(), + ranges.p_spatial, + is_log("p_spatial"), + )) + } else { + let d = *cfg.params.defaults.get("p_spatial").unwrap_or(&21.0); + vec![d; n_reaches] + }; + + // ---- Decompose n into (uparea-linear trend) + residual ------------------ + // `spatial_attributes` is already z-scored, and normalization is monotone, + // so the standardized column is a faithful stand-in for log10(uparea). + let u_idx = head_cfg + .input_var_names + .iter() + .position(|v| v == "log10_uparea") + .ok_or("kan_head.input_var_names must contain log10_uparea for this probe")?; + let attrs: Vec = tensors + .spatial_attributes + .clone() + .into_data() + .to_vec::() + .unwrap(); + let n_feat = head_cfg.input_var_names.len(); + let u: Vec = (0..n_reaches).map(|i| attrs[i * n_feat + u_idx]).collect(); + + let u_bar = u.iter().sum::() / n_reaches as f32; + let n_bar = n_trained.iter().sum::() / n_reaches as f32; + let (mut sxy, mut sxx) = (0f64, 0f64); + for i in 0..n_reaches { + let du = (u[i] - u_bar) as f64; + sxy += du * (n_trained[i] - n_bar) as f64; + sxx += du * du; + } + let slope = if sxx > 0.0 { (sxy / sxx) as f32 } else { 0.0 }; + eprintln!( + "trained field: median n = {:.5}, d(n)/d(z_uparea) = {:+.6}", + median(&n_trained), + slope + ); + + let [lo, hi] = ranges.n; + let clamp = |x: f32| x.clamp(lo, hi); + + let mut fields: Vec = Vec::new(); + for &k in &cli.amplify { + // Amplify ONLY the uparea-linear component; residual structure is kept + // so the k=1 case is bit-for-bit the trained field. + let n: Vec = (0..n_reaches) + .map(|i| clamp(n_trained[i] + (k - 1.0) * slope * (u[i] - u_bar))) + .collect(); + fields.push(Field { label: format!("slope x{k}"), n }); + } + // Imposed ramps: rank reaches by log10_uparea and map to [lo, hi]. Rank + // rather than value keeps the ramp well-conditioned however uparea is + // distributed, and guarantees the full endpoint span is realized. + let mut order: Vec = (0..n_reaches).collect(); + order.sort_unstable_by(|&a, &b| u[a].partial_cmp(&u[b]).expect("finite attr")); + for spec in &cli.impose { + let (a, b) = spec + .split_once(':') + .ok_or_else(|| format!("--impose expects LO:HI, got `{spec}`"))?; + let lo_i: f32 = a.parse().map_err(|e| format!("--impose LO `{a}`: {e}"))?; + let hi_i: f32 = b.parse().map_err(|e| format!("--impose HI `{b}`: {e}"))?; + let mut n = vec![0f32; n_reaches]; + for (rank, &i) in order.iter().enumerate() { + let f = if n_reaches > 1 { rank as f32 / (n_reaches - 1) as f32 } else { 0.0 }; + n[i] = clamp(lo_i + (hi_i - lo_i) * f); + } + fields.push(Field { label: format!("ramp {lo_i}->{hi_i}"), n }); + } + + if !cli.no_flat { + let m = median(&n_trained); + fields.push(Field { label: "flat (median n)".into(), n: vec![m; n_reaches] }); + } + if !cli.no_floor { + fields.push(Field { label: "floor (identity)".into(), n: vec![lo; n_reaches] }); + } + + // ---- Score every field through the identical eval path ----------------- + let mut rows: Vec<(String, f32, f32, f32, f32, usize)> = Vec::new(); + for f in &fields { + let frozen = FrozenParams { + n: f.n.clone(), + q_spatial: q_trained.clone(), + p_spatial: p_trained.clone(), + }; + eprintln!("--- scoring {:<18} median n = {:.5}", f.label, median(&f.n)); + let out = evaluate::( + &cfg, + &dataset, + EvalParams::Frozen(&frozen), + &device, + cli.batch_size_days, + &cli.checkpoint, + ) + .map_err(|e| format!("evaluate field {}: {e}", f.label))?; + let finite = out.metrics.nse.iter().filter(|x| x.is_finite()).count(); + let (nse, kge) = (median(&out.metrics.nse), median(&out.metrics.kge)); + eprintln!(" median NSE {nse:.4} median KGE {kge:.4} ({finite} finite gauges)"); + rows.push((f.label.clone(), median(&f.n), nse, kge, median(&out.metrics.fhv), finite)); + } + + // ---- Report ------------------------------------------------------------- + let ref_nse = rows + .iter() + .find(|r| r.0 == "slope x1") + .map(|r| r.2) + .unwrap_or(f32::NAN); + println!("\n{:<20} {:>10} {:>10} {:>10} {:>10} {:>8}", "field", "median_n", "NSE", "dNSE", "KGE", "FHV"); + for (label, mn, nse, kge, fhv, _) in &rows { + println!("{label:<20} {mn:>10.5} {nse:>10.4} {:>+10.4} {kge:>10.4} {fhv:>8.2}", nse - ref_nse); + } + println!("\ndNSE is relative to `slope x1` (the trained field scored through this same path)."); + + if let Some(p) = cli.output { + let mut s = String::from("field,median_n,nse,dnse,kge,fhv,n_finite\n"); + for (label, mn, nse, kge, fhv, fin) in &rows { + s.push_str(&format!("{label},{mn},{nse},{},{kge},{fhv},{fin}\n", nse - ref_nse)); + } + std::fs::write(&p, s).map_err(|e| format!("write {}: {e}", p.display()))?; + eprintln!("wrote {}", p.display()); + } + Ok(()) +} diff --git a/src/bin/probe_zeta_gradient.rs b/src/bin/probe_zeta_gradient.rs index 503122b..eff8a98 100644 --- a/src/bin/probe_zeta_gradient.rs +++ b/src/bin/probe_zeta_gradient.rs @@ -559,7 +559,8 @@ fn run(cfg: Config, cli: Cli, device: I::Device) -> Result<(), Box = Vec::with_capacity(g * t_days); for gi in 0..g { for ti in 0..t_days { - obs_buf.push(obs_arr[(ti + 1, gi)]); + // pooled day i ↔ obs day i (2026-08-08 tau convention) + obs_buf.push(obs_arr[(ti, gi)]); } } let obs_t: Tensor, 2> = @@ -1583,10 +1584,11 @@ fn run_floor( ); // |pred - obs| with the SAME obs alignment as training (grad mode's - // obs_arr[(ti + 1, gi)]); NaN obs propagates to NaN residual. + // obs_arr[(ti, gi)], 2026-08-08 tau convention); NaN obs propagates + // to NaN residual. for gi in 0..g { for ti in 0..t_days { - let o = obs_arr[(ti + 1, gi)]; + let o = obs_arr[(ti, gi)]; all_resid.push((pred[gi * t_days + ti] - o).abs()); } } @@ -1743,8 +1745,8 @@ fn gather_by_comid(donor: &HashMap, comids: &[i64]) -> Result /// forward pass. fn trimmed_days(rho: usize, tau: u32) -> usize { let n_hourly = (rho - 1) * 24; - let start = 13 + tau as usize; - let end = n_hourly - 11 + tau as usize; + let start = tau as usize; + let end = n_hourly - (24 - tau as usize); (end - start) / 24 } @@ -1861,7 +1863,7 @@ fn sample_window_plan( batch.observations.nrows() ); let surviving = (0..batch.gauge_staids.len()).any(|gi| { - (warmup..t_days).all(|ti| !batch.observations[(ti + 1, gi)].is_nan()) + (warmup..t_days).all(|ti| !batch.observations[(ti, gi)].is_nan()) }); if !surviving { eprintln!(" window skipped: all gauges have NaN in post-warmup window"); @@ -1907,7 +1909,7 @@ fn eval_window_filtered( let mut obs_post = Array2::::zeros((t_days - warmup, g)); for gi in 0..g { for ti in warmup..t_days { - obs_post[(ti - warmup, gi)] = obs_arr[(ti + 1, gi)]; + obs_post[(ti - warmup, gi)] = obs_arr[(ti, gi)]; } } @@ -2892,9 +2894,10 @@ mod tests { #[test] fn trimmed_days_matches_tau_trim_and_downsample_arithmetic() { - // rho=90 -> n_hourly=(90-1)*24=2136; tau=0 -> start=13,end=2125, - // trimmed=2112, days=2112/24=88. + // rho=90 -> n_hourly=(90-1)*24=2136; tau=0 -> start=0,end=2112, + // trimmed=2112, days=2112/24=88 (total trim is 24 h at any tau). assert_eq!(trimmed_days(90, 0), 88); + assert_eq!(trimmed_days(90, 9), 88); } // ----------------------------------------------------------------------- diff --git a/src/cli/plan.rs b/src/cli/plan.rs index 0df20c0..6f5754a 100644 --- a/src/cli/plan.rs +++ b/src/cli/plan.rs @@ -336,6 +336,11 @@ pub(crate) fn resolve_adjacency( fabric, ds.geospatial_fabric_layer.as_deref(), &ds.gages, + // Subdivision changes the built graph, so it is part of the + // cache key and of the build itself; `attributes` supplies the + // `catchsize` column the reference celerity needs. + &config.params.subdivision, + &ds.attributes, ) .map_err(|e| CliError::ConfigInvalid { path: config_path.into(), diff --git a/src/config.rs b/src/config.rs index 5fdc6db..4e57b10 100644 --- a/src/config.rs +++ b/src/config.rs @@ -109,9 +109,9 @@ pub struct DataSources { pub observations: std::path::PathBuf, pub gages: std::path::PathBuf, /// Optional hourly AORC precipitation store (`merit_unit_catchments.zarr`, - /// zarr v3). CONUS-only; drives the precip-conditioned disaggregation head - /// (`kan_head.disaggregation.use_precip`). Absent ⇒ the head conditions on - /// daily Q' only (or, with disaggregation off, flat repeat-24). + /// zarr v3). CONUS-only; drives the precip-conditioned disaggregation head. + /// Mandatory whenever `kan_head.disaggregation:` is present and enabled + /// (the head always consumes precip); unused otherwise. #[serde(default)] pub aorc_precip: Option, /// Path to the MERIT flowlines fabric: `.shp` (sibling `.dbf` read), @@ -306,11 +306,23 @@ pub struct KanHeadConfigSection { pub disaggregation: Option, } -/// YAML `kan_head.disaggregation:` block (presence enables the head). -/// The head always consumes `(daily Q', that day's 24h precip)` — requires -/// `data_sources.aorc_precip` to be set. See `src/nn/disagg_head.rs`. +/// 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 +/// `src/nn/disagg_head.rs`. `deny_unknown_fields` because phantom keys here +/// (e.g. the removed `use_precip`) silently produced a different head than +/// intended. #[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] pub struct DisaggregationSection { + /// On/off switch. Default `true` — a bare block enables the head, + /// preserving the historical presence-only contract. `false` strips the + /// whole block at config load, so the loader falls back to flat + /// repeat-24 (nearest) daily→hourly upsampling and the block stays in + /// the YAML inert (the ablation switch: flip one line, change nothing + /// else). + #[serde(default = "default_disagg_enabled")] + pub enabled: bool, #[serde(default = "default_disagg_hidden")] pub hidden_size: usize, #[serde(default = "default_disagg_num_hidden_layers")] @@ -386,6 +398,9 @@ fn default_grid() -> usize { fn default_k() -> usize { 3 } +fn default_disagg_enabled() -> bool { + true +} fn default_disagg_hidden() -> usize { 16 } @@ -473,6 +488,93 @@ pub enum SparseSolver { Cuda, } +/// YAML `params.subdivision:` block — static reach subdivision so +/// `Cr = c·Δt/Δx` lands near 1 network-wide. +/// +/// Off by default; `enabled: false` reproduces current behaviour exactly. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Subdivision { + #[serde(default)] + pub enabled: bool, + /// Hard cap on pieces per reach. Uncapped subdivision is infeasible: + /// 13.2x reaches and 9.2x solver critical path, and Sum(m) cannot be + /// pinned down (2.3M-10.5M across defensible reference-flow choices). + /// Capping bounds the cost AND makes it estimable (+/-12% at M=8). + #[serde(default = "default_max_pieces")] + pub max_pieces: usize, + /// Manning's n used ONLY to compute the reference celerity that sets m. + /// Deliberately NOT taken from a checkpoint: the graph must not depend on + /// training state. 0.05 is the trained CONUS median. + #[serde(default = "default_reference_n")] + pub reference_n: f32, + /// Reference discharge Q_ref = coefficient * uparea_km2^exponent (m3/s). + #[serde(default = "default_ref_q_coeff")] + pub reference_discharge_coefficient: f32, + #[serde(default = "default_ref_q_exp")] + pub reference_discharge_exponent: f32, + /// Short reaches get their length clamped UP to + /// `min_length_fraction * c_ref * dt`, giving `Cr <= 1/min_length_fraction` + /// at the reference flow. 1.0 targets Cr = 1; 0.5 targets Cr <= 2 (the + /// non-negativity bound) with half the length distortion; 0.0 disables the + /// clamp entirely, leaving short reaches over-Courant. + /// + /// This is a BUILD-TIME constant, unlike the runtime K floor in + /// `enforce_positivity`. It therefore has no gradient path and cannot + /// create the `X ~ Cr ~ 1/n` coupling that drove n to its floor. + #[serde(default = "default_min_length_fraction")] + pub min_length_fraction: f32, + /// Hard bound on how far the short-reach clamp may stretch a reach: + /// `length <= original_length * max_clamp_factor`. + /// + /// Without this the clamp is unbounded, because `reference_celerity` uses a + /// depth relation `r = Q_ref^0.4` with no slope dependence while `v` scales + /// as `sqrt(S)`. Steep small catchments therefore get big-river depth AND + /// steep-slope velocity, reaching ~8.9 m/s at slope 1e-2 — a `dx_target` of + /// 32 km that would stretch every short steep headwater to 32 km. Measured + /// unbounded clamp factors ran to p99 = 36x and max = 48,597x. + /// + /// Bounding the distortion means the worst reaches stay over-Courant rather + /// than being silently rewritten into 30 km channels. That residual is a + /// reported number, not a hidden one. + #[serde(default = "default_max_clamp_factor")] + pub max_clamp_factor: f32, +} + +fn default_max_pieces() -> usize { + 8 +} +fn default_reference_n() -> f32 { + 0.05 +} +fn default_ref_q_coeff() -> f32 { + 0.01 +} +fn default_ref_q_exp() -> f32 { + 0.9 +} +fn default_max_clamp_factor() -> f32 { + 4.0 +} + +fn default_min_length_fraction() -> f32 { + 1.0 +} + +impl Default for Subdivision { + fn default() -> Self { + Self { + enabled: false, + max_pieces: default_max_pieces(), + reference_n: default_reference_n(), + reference_discharge_coefficient: default_ref_q_coeff(), + reference_discharge_exponent: default_ref_q_exp(), + min_length_fraction: default_min_length_fraction(), + max_clamp_factor: default_max_clamp_factor(), + } + } +} + /// Routing parameter configuration. #[derive(Debug, Clone)] pub struct Params { @@ -501,6 +603,42 @@ pub struct Params { /// supplied at routing setup (the mask is precomputed by the caller). /// Default 0.7 (70% impervious surface ≈ concrete-lined channel). pub leakance_impervious_threshold: f32, + /// When `true` (default) the routing core reproduces DDR's formulation + /// bit-for-bit, including three known defects: + /// * celerity `c = v · 5/3` (the wide-rectangular Kleitz-Seddon limit, + /// ~22-27% high for the trapezoid this code actually builds), + /// * Muskingum `X ≡ 0.3` (constant, NOT Cunge-derived, giving a median + /// 10-30x excess numerical diffusion), and + /// * `outflow_idx` = the gauge's UPSTREAM neighbours rather than the + /// gauge's own reach, which drops that reach's local drainage from + /// every prediction (`src/data/collate.rs` step 5). + /// + /// Set `false` to enable the corrected physics. The first two CHANGE + /// FORWARD OUTPUT and will break `examples/compare_ddr_sandbox`'s ABSOLUTE + /// MATCH (invariant 1) — which is why the default preserves DDR behaviour. + /// The `outflow_idx` correction is downstream of the solver and does NOT + /// affect the sandbox. See `.claude/PHYSICS-CORRECTIONS.md`. + pub ddr_match: bool, + /// Enforce the Muskingum non-negativity window `2X <= Cr <= 2(1-X)` + /// (`Cr = dt/K`) on every reach-timestep, so the S27 solve can never + /// produce a negative discharge for S28's `clamp_min` to rewrite to + /// `+1e-4` (which silently creates mass). + /// + /// Implemented by clamping the *inputs* — `K >= dt(1+d)/2` and + /// `X <= min(0.5·Cr, 1 - 0.5·Cr)·(1-d)` with `d = 1e-2` — never the + /// coefficients: `c1+c2+c3 = 1` holds for any `(K, X)`, so clamping inputs + /// preserves mass exactly while clamping `c3` would not. + /// + /// Off by default. Requires `ddr_match: false` — the clamp changes K and X, + /// which would break `examples/compare_ddr_sandbox`'s ABSOLUTE MATCH + /// (invariant 1). See `.claude/PHYSICS-CORRECTIONS.md`. + pub enforce_positivity: bool, + /// Static reach subdivision (variable Δx). Off by default. + pub subdivision: Subdivision, +} + +fn default_ddr_match() -> bool { + true } impl Default for Params { @@ -518,6 +656,9 @@ impl Default for Params { use_leakance: false, leakance_losing_only: true, leakance_impervious_threshold: 0.7, + ddr_match: default_ddr_match(), + enforce_positivity: false, + subdivision: Subdivision::default(), } } } @@ -539,6 +680,10 @@ struct ParamsRaw { use_leakance: Option, leakance_losing_only: Option, leakance_impervious_threshold: Option, + ddr_match: Option, + enforce_positivity: Option, + #[serde(default)] + subdivision: Subdivision, } impl From for Params { @@ -590,7 +735,13 @@ impl From for Params { if !r.log_space_parameters.is_empty() { p.log_space_parameters = r.log_space_parameters; } - p.tau = r.tau.unwrap_or(3); + // 2026-08-08 convention: tau = hours the routed output is advanced + // before daily scoring (0 = day-aligned; dMC-Juniata sign). Default 9 + // = the measured CONUS optimum (old-convention 20; findings §5g). + // Pre-2026-08-08 configs' tau values are on the OLD scale (old = new + // + 11) and must not be reused verbatim. + p.tau = r.tau.unwrap_or(9); + assert!(p.tau < 24, "params.tau must be in [0, 24) hours; got {}", p.tau); p.sparse_solver = match r.sparse_solver.as_deref() { Some("cuda") | Some("CUDA") => SparseSolver::Cuda, Some("cpu") | Some("CPU") | None => SparseSolver::Cpu, @@ -608,6 +759,16 @@ impl From for Params { if let Some(v) = r.leakance_impervious_threshold { p.leakance_impervious_threshold = v; } + if let Some(b) = r.ddr_match { + p.ddr_match = b; + } + if let Some(b) = r.enforce_positivity { + p.enforce_positivity = b; + } + // Non-Option: `ParamsRaw`'s struct-level `#[serde(default)]` already + // yields `Subdivision::default()` when the block is absent, which is + // exactly what `Params::default()` carries. + p.subdivision = r.subdivision; p } } @@ -720,6 +881,16 @@ impl Config { })?; let testing_raw = raw.testing.clone(); let mut cfg: Self = raw.into(); + // `kan_head.disaggregation.enabled: false` disables the head by + // stripping its block before validation — every downstream consumer + // (dataset precip gating, head construction, the hourly-resolution + // guard) keys off `disaggregation.is_some()`, so the stripped config + // behaves exactly like one without the block. + if let Some(head) = &mut cfg.kan_head { + if head.disaggregation.as_ref().is_some_and(|d| !d.enabled) { + head.disaggregation = None; + } + } validate_mode_workflow(&cfg).map_err(|msg| DataError::Yaml { path: path.to_path_buf(), source: serde_yaml::Error::custom(msg), @@ -732,6 +903,18 @@ impl Config { path: path.to_path_buf(), source: serde_yaml::Error::custom(msg), })?; + validate_ddr_match(&cfg).map_err(|msg| DataError::Yaml { + path: path.to_path_buf(), + source: serde_yaml::Error::custom(msg), + })?; + validate_enforce_positivity(&cfg).map_err(|msg| DataError::Yaml { + path: path.to_path_buf(), + source: serde_yaml::Error::custom(msg), + })?; + validate_subdivision(&cfg).map_err(|msg| DataError::Yaml { + path: path.to_path_buf(), + source: serde_yaml::Error::custom(msg), + })?; validate_disagg_pretrained(&cfg).map_err(|msg| DataError::Yaml { path: path.to_path_buf(), source: serde_yaml::Error::custom(msg), @@ -820,6 +1003,104 @@ fn validate_leakance(cfg: &Config) -> std::result::Result<(), String> { Ok(()) } +fn validate_ddr_match(cfg: &Config) -> std::result::Result<(), String> { + if !cfg.params.ddr_match && cfg.params.use_cuda_graphs { + return Err( + "params: `ddr_match: false` requires `use_cuda_graphs: false` — the \ + CUDA-graph kernel hardcodes DDR's `5/3` celerity, so the corrected \ + forward would not be captured while the backward would use the corrected \ + chain rule, producing a silent forward/backward mismatch. \ + Set `use_cuda_graphs: false` to use `ddr_match: false`." + .to_string(), + ); + } + Ok(()) +} + +/// The positivity clamp raises K and lowers X, so it CHANGES FORWARD OUTPUT. +/// Under `ddr_match: true` that would break `examples/compare_ddr_sandbox`'s +/// ABSOLUTE MATCH (invariant 1), so the combination is rejected at load rather +/// than silently producing a non-DDR forward on the DDR-faithful path. +fn validate_enforce_positivity(cfg: &Config) -> std::result::Result<(), String> { + if cfg.params.enforce_positivity && cfg.params.ddr_match { + return Err( + "params: `enforce_positivity: true` requires `ddr_match: false` — the \ + positivity clamp floors K at dt(1+d)/2 and caps X at the stability \ + window, changing forward output and breaking the DDR sandbox \ + ABSOLUTE MATCH. Set `ddr_match: false` to use `enforce_positivity`." + .to_string(), + ); + } + Ok(()) +} + +/// Subdivision expands the reach count, so a CUDA graph captured for the +/// un-split network is the wrong size; and `max_pieces: 0` would ask for a +/// network with zero rows. Both are rejected at load. +fn validate_subdivision(cfg: &Config) -> std::result::Result<(), String> { + let s = &cfg.params.subdivision; + if s.enabled && s.max_pieces < 1 { + return Err("params.subdivision: `max_pieces` must be >= 1".to_string()); + } + if s.enabled && cfg.params.use_cuda_graphs { + return Err( + "params.subdivision: `enabled: true` requires `use_cuda_graphs: false` \ + — the captured graph is sized to a fixed reach count." + .to_string(), + ); + } + if s.enabled { + validate_subdivision_reaches_the_builder(cfg)?; + } + Ok(()) +} + +/// Subdivision runs *inside* the managed adjacency builder +/// (`adjacency::cache::resolve_or_build`), and `cli::plan::resolve_adjacency` +/// only reaches that builder when `data_sources` supplies **no** explicit +/// adjacency paths. With `conus_adjacency`/`gages_adjacency` set, the flag would +/// be **silently inert**: no subdivision, no warning, and a run whose manifest +/// says `subdivision.enabled: true` while the routed network is the un-split +/// one. This repo has been bitten by that class of silent no-op before (the +/// 2026-07-01 stale-binary disaggregation 2×2), so it is rejected at load. +/// +/// The one legitimate combination is an explicit path to a store that was +/// *already* built with subdivision — that graph really is split, and rebuilding +/// it would be wasted work. `adjacency::validate::store_is_subdivided` +/// distinguishes the two from the store's zarr metadata (`n_parent < n`) without +/// reading any array data. +fn validate_subdivision_reaches_the_builder(cfg: &Config) -> std::result::Result<(), String> { + let Some(ds) = cfg.data_sources.as_ref() else { + return Ok(()); // no data_sources at all — unit-test / default configs + }; + let (Some(conus), Some(_gages)) = (&ds.conus_adjacency, &ds.gages_adjacency) else { + return Ok(()); // managed build: the builder is reached, subdivision runs + }; + match crate::adjacency::validate::store_is_subdivided(conus) { + Some(true) => Ok(()), + Some(false) => Err(format!( + "params.subdivision: `enabled: true` conflicts with the explicit \ + adjacency paths in `data_sources`. Subdivision only runs inside the \ + managed adjacency builder, which is bypassed when `conus_adjacency` \ + and `gages_adjacency` are set — so the flag would be SILENTLY INERT \ + and the run would route the un-split network. The store at {} \ + carries no subdivision map (n_parent == n). Fix: remove \ + `conus_adjacency` and `gages_adjacency` and set `geospatial_fabric` \ + instead, so ddrs builds (and caches) the subdivided graph.", + conus.display() + )), + None => Err(format!( + "params.subdivision: `enabled: true` with explicit adjacency paths, \ + and the store at {} could not be inspected (missing or unreadable \ + `order`/`parent_order` metadata) — so ddrs cannot confirm it is a \ + subdivided graph. Subdivision only runs inside the managed adjacency \ + builder, which these keys bypass. Fix: remove `conus_adjacency` and \ + `gages_adjacency` and set `geospatial_fabric` instead.", + conus.display() + )), + } +} + /// `freeze: true` without a `pretrained_checkpoint` would permanently pin the /// disaggregation head at its random init — reject at load time. fn validate_disagg_pretrained(cfg: &Config) -> std::result::Result<(), String> { @@ -923,8 +1204,9 @@ mod tests { let kan_head = cfg.kan_head.as_ref().unwrap(); assert_eq!(kan_head.hidden_size, 21); assert_eq!(kan_head.input_var_names.len(), 10); - // tau defaults to 3 when not set in YAML. - assert_eq!(cfg.params.tau, 3); + // tau defaults to 9 when not set in YAML (2026-08-08 convention: + // hours of advance; ≡ old-convention 20). + assert_eq!(cfg.params.tau, 9); // sparse_solver is set to Cuda by merit_training.yaml (since SP-9). assert_eq!(cfg.params.sparse_solver, SparseSolver::Cuda); // SP-10: merit_training.yaml now sets use_cuda_graphs: true @@ -1515,6 +1797,63 @@ params: ); } + #[test] + fn ddr_match_false_rejects_cuda_graphs() { + // ddr_match: false + use_cuda_graphs: true must fail at load: the + // CUDA-graph kernel hardcodes DDR's 5/3 celerity, so the corrected + // forward would never be captured while the backward would use the + // corrected chain rule — a silent forward/backward mismatch. + let yaml_reject = r#" +mode: training +geodataset: merit +seed: 1 +np_seed: 1 +params: + ddr_match: false + use_cuda_graphs: true +"#; + let path = std::env::temp_dir().join("ddrs_ddr_match_graphs_reject.yaml"); + std::fs::write(&path, yaml_reject).unwrap(); + let err = Config::from_yaml_file(&path).unwrap_err(); + let msg = format!("{err}"); + assert!( + msg.contains("ddr_match") && msg.contains("use_cuda_graphs"), + "expected ddr_match/cuda_graphs conflict, got: {msg}" + ); + + // ddr_match: true + use_cuda_graphs: true must still load (default path). + let yaml_ok_true = r#" +mode: training +geodataset: merit +seed: 1 +np_seed: 1 +params: + ddr_match: true + use_cuda_graphs: true +"#; + let path2 = std::env::temp_dir().join("ddrs_ddr_match_true_graphs.yaml"); + std::fs::write(&path2, yaml_ok_true).unwrap(); + let cfg2 = Config::from_yaml_file(&path2).expect("ddr_match:true + cuda_graphs:true must load"); + assert!(cfg2.params.ddr_match); + assert!(cfg2.params.use_cuda_graphs); + + // ddr_match: false + use_cuda_graphs: false must also load. + let yaml_ok_false = r#" +mode: training +geodataset: merit +seed: 1 +np_seed: 1 +params: + ddr_match: false + use_cuda_graphs: false +"#; + let path3 = std::env::temp_dir().join("ddrs_ddr_match_false_no_graphs.yaml"); + std::fs::write(&path3, yaml_ok_false).unwrap(); + let cfg3 = Config::from_yaml_file(&path3).expect("ddr_match:false + cuda_graphs:false must load"); + assert!(!cfg3.params.ddr_match); + assert!(!cfg3.params.use_cuda_graphs); + } + #[test] fn state_cache_absent_yields_none() { // experiment block without state_cache → None (critical byte-identity invariant). diff --git a/src/data/collate.rs b/src/data/collate.rs index 91872bd..46df393 100644 --- a/src/data/collate.rs +++ b/src/data/collate.rs @@ -64,9 +64,117 @@ pub struct CompressedAdj { pub cols: Vec, /// Per-gauge compressed position of the gauge outlet, length `G_present`. pub gauge_compressed: Vec, - /// For each gauge, the compressed cols whose row index equals the - /// gauge's outlet. Mirrors DDR's `outflow_idx`. + /// For each gauge, the compressed reach indices whose routed discharge is + /// summed to form that gauge's prediction. Which reaches those are depends + /// on `params.ddr_match` — see `compress` step 5: + /// + /// * `ddr_match: false` (physically correct) — a single element: the + /// gauge's OWN reach, at its outlet piece. The MC solve there already + /// integrates the whole upstream network plus that reach's own lateral + /// inflow. Without subdivision that is `gauge_compressed[g]` itself. + /// * `ddr_match: true` (DDR-faithful default) — the gauge's UPSTREAM + /// neighbours, which omit the gauge reach's own local drainage. pub outflow_idx: Vec>, + /// Reach-subdivision parent map in **compressed** space, length + /// `n_parent_active + 1`: compressed rows + /// `[parent_offset[p], parent_offset[p + 1])` are the sub-reach pieces of + /// the `p`-th parent present in this batch. Feeds + /// `SparseAdjacency::parent_offset`, which the engine turns into the + /// per-row lateral-inflow divisor. + /// + /// `None` when the caller passed no CONUS parent map. Identity + /// (`0..=n_active`) when the store is not subdivided — the engine treats + /// both alike and skips the split. + pub parent_offset: Option>, + /// CONUS **sub-reach** position of each compressed row, length `N_active`: + /// compressed row `i` was compressed from CONUS row `conus_positions[i]`. + /// + /// Anything needing per-row CONUS geometry (`length_m`, `slope`) must index + /// with this. `ConusAdjacencyStore::index.position(comid)` is NOT a + /// substitute: it lives in PARENT space (see the two-index-spaces note on + /// `ConusAdjacencyStore`), so under subdivision it would hand every piece of + /// a parent the geometry of whatever row the parent's index happens to + /// number — the parent's full length instead of `L/m`, and some other + /// reach's slope entirely. + pub conus_positions: Vec, +} + +impl CompressedAdj { + /// One COMID per parent reach present in this batch, in compressed + /// topological order. Equals `divide_comids` when the network is not + /// subdivided. + /// + /// `divide_comids` has one entry per SUB-REACH, so under subdivision it + /// repeats a parent's COMID `m` times. Attributes are per-COMID and + /// sub-reaches share their parent's hydraulics, so the KAN head is built + /// and run at parent resolution and its outputs are gathered onto the + /// pieces (`training::forward::gather_params_to_subreaches`). This is the + /// COMID list that attribute matrix is sliced from. + pub fn parent_comids(&self) -> Vec { + match &self.parent_offset { + Some(off) => off[..off.len() - 1] + .iter() + .map(|&lo| self.divide_comids[lo as usize]) + .collect(), + None => self.divide_comids.clone(), + } + } +} + +/// Parent index owning sub-reach row `row`. `parent_offset` is strictly +/// increasing, so the owner is the last offset that is `<= row`. +#[inline] +fn parent_of_row(parent_offset: &[i32], row: usize) -> usize { + parent_offset.partition_point(|&o| o <= row as i32) - 1 +} + +/// Re-express a CONUS-space `parent_offset` in compressed space. +/// +/// `active` is the sorted list of CONUS sub-reach positions this batch kept. +/// A parent's pieces occupy a contiguous ascending run of CONUS rows, and a +/// gauge subgraph always enters a parent at its outlet and then walks the +/// internal chain upstream — so every parent present in `active` must be +/// present *in full*, as an unbroken run. That is asserted here rather than +/// assumed: a partially-present parent would hand the engine a piece count +/// `m` smaller than the one the piece lengths were derived from (silently +/// creating mass) and would move the parent's outlet row. +fn compressed_parent_offset( + active: &[usize], + conus_parent_offset: &[i32], +) -> Result> { + let n_rows = *conus_parent_offset.last().unwrap_or(&0); + let mut offsets: Vec = vec![0]; + let mut i = 0usize; + while i < active.len() { + let row = active[i]; + if row as i32 >= n_rows { + return Err(DataError::Malformed { + path: PathBuf::from(""), + message: format!( + "compress: CONUS row {row} is outside the parent map \ + (which covers {n_rows} rows)" + ), + }); + } + let p = parent_of_row(conus_parent_offset, row); + let lo = conus_parent_offset[p] as usize; + let hi = conus_parent_offset[p + 1] as usize; + let m = hi - lo; + if active.len() - i < m || (0..m).any(|k| active[i + k] != lo + k) { + return Err(DataError::Malformed { + path: PathBuf::from(""), + message: format!( + "compress: parent {p} owns CONUS rows {lo}..{hi}, but the active \ + set does not hold them contiguously from compressed row {i} — a \ + partial sub-reach chain would mis-scale lateral inflow and move \ + the gauge outlet" + ), + }); + } + i += m; + offsets.push(i as i32); + } + Ok(offsets) } /// Compress a unioned COO into dense compressed-position space, preserving @@ -75,9 +183,22 @@ pub struct CompressedAdj { /// /// Hard-asserts the lower-triangular invariant (`rows >= cols`); fails /// with `DataError::Malformed` if violated. +/// +/// `ddr_match` selects the `outflow_idx` convention (see step 5): `true` +/// reproduces DDR's `merit.py:226-234` bit-for-bit, `false` uses the +/// physically correct gauge-reach index. Comes from `params.ddr_match`. +/// +/// `conus_parent_offset` is `ConusAdjacencyStore::parent_offset` — the +/// reach-subdivision map in CONUS sub-reach space. Pass it whenever it is +/// available (it is the identity `0..=n` on un-subdivided stores, which costs +/// nothing); `None` disables both the compressed parent map and the +/// outlet-piece resolution below, which is what the pure-topology unit tests +/// want. pub fn compress( unioned: &UnionedCoo, conus_order: &[Comid], + ddr_match: bool, + conus_parent_offset: Option<&[i32]>, ) -> Result { use std::collections::BTreeSet; @@ -130,27 +251,92 @@ pub fn compress( let gauge_compressed: Vec = unioned.gauges.iter().map(|(_, g, _)| mapping[g]).collect(); - // 5. outflow_idx[g] = list of cols where rows[k] == gauge_compressed[g]. - // Fallback (matches DDR `_collate_gages` lines ~226-235): when a gauge - // has no incoming edges in this batch's union, use the gauge's own - // compressed index as the sole outflow. Headwater gauges are filtered - // upstream of compress(), so this fallback only fires for gauges at - // merge nodes whose upstream edges weren't in the batch. - let mut outflow_idx: Vec> = Vec::with_capacity(gauge_compressed.len()); - for &g_comp in &gauge_compressed { - let g_row = g_comp as i32; - let cols_for_g: Vec = rows + // 4b. Reach-subdivision parent map, re-expressed in compressed space. + let parent_offset: Option> = match conus_parent_offset { + Some(off) => Some(compressed_parent_offset(&active_vec, off)?), + None => None, + }; + + // 5. outflow_idx — which reaches are summed to form a gauge's prediction. + // + // `ddr_match: false` (CORRECT) — the gauge's OWN reach, read at the + // OUTLET piece of that reach. + // A USGS gauge measures every drop of drainage above it, including the + // lateral inflow of the reach the gauge sits on, and we do not know where + // along that reach the gauge physically sits. The Muskingum-Cunge solve at + // the gauge reach already accumulates all upstream contributions plus its + // own `q_prime` by mass conservation, so the gauge's prediction is that + // ONE reach's routed discharge. + // + // Under reach subdivision that reach spans several rows. The solve is + // mass-conserving down the internal chain, so the whole reach's runoff — + // upstream network plus its own lateral inflow, which was split `q'/m` + // across the pieces — only arrives at the LAST piece, + // `parent_offset[p + 1] - 1`. Reading an earlier piece would drop the + // downstream fraction of the reach's own lateral inflow: the same class of + // bug as the upstream-cols defect below, in a new form. + // + // `gauge_compressed` holds COMPRESSED SUB-REACH positions, not parent + // indices, so the parent must be recovered from the *compressed* + // `parent_offset` before its outlet can be taken. (Indexing the CONUS + // parent map with a compressed row is a category error — the compression + // renumbers rows.) In practice `outlet == gauge_compressed[g]` already, + // because the gauge subgraph builder resolves a gauge COMID to its last + // matching row (`cache.rs::resolve_or_build`); deriving it here keeps the + // guarantee local instead of resting on that builder detail. + // + // `ddr_match: true` (DEFAULT, DDR-FAITHFUL) — the gauge's UPSTREAM + // neighbours. Reproduces DDR's `_collate_gages` + // (`~/projects/ddr/src/ddr/geodatazoo/merit.py:226-234`), which collects + // the COO *cols* whose row equals the gauge outlet and only falls back to + // the gauge's own index when that list is empty. In this adjacency + // `indices_0` = rows = DOWNSTREAM segment and `indices_1` = cols = + // UPSTREAM segment (`src/data/store/zarr.rs:39-42`), so this silently + // drops the gauge reach's own local runoff from every prediction. + // + // Why `false` is the physical answer: gauge 01457000 (366.8 km² drainage, + // of which the gauge reach alone is 250.1 km² = 68%) read 1.58 m³/s + // against an observed 7.60 and a summed-Q' baseline of 7.38 — a constant + // 0.215× suppression across all 15 eval years, on peaks as well as means. + // 26 of 1841 gauges fell below 0.5× baseline, all of them small basins + // where the gauge reach is a large share of the area. Because the omitted + // mass is always positive, this biases EVERY ddrs-vs-baseline comparison + // against ddrs. + // + // DDR's Lynker path validates `outflow_idx` against the flowpath `toid` + // column (`~/projects/ddr/src/ddr/geodatazoo/lynker_hydrofabric.py:239-250`); + // the MERIT path has no such check — that is where this would have been + // caught upstream. + let outflow_idx: Vec> = if ddr_match { + gauge_compressed .iter() - .zip(cols.iter()) - .filter(|(r, _)| **r == g_row) - .map(|(_, c)| *c as usize) - .collect(); - if cols_for_g.is_empty() { - outflow_idx.push(vec![g_comp]); - } else { - outflow_idx.push(cols_for_g); + .map(|&g_comp| { + let g_row = g_comp as i32; + let cols_for_g: Vec = rows + .iter() + .zip(cols.iter()) + .filter(|(r, _)| **r == g_row) + .map(|(_, c)| *c as usize) + .collect(); + if cols_for_g.is_empty() { + vec![g_comp] + } else { + cols_for_g + } + }) + .collect() + } else { + match &parent_offset { + Some(off) => gauge_compressed + .iter() + .map(|&g_comp| { + let p = parent_of_row(off, g_comp); + vec![off[p + 1] as usize - 1] + }) + .collect(), + None => gauge_compressed.iter().map(|&g_comp| vec![g_comp]).collect(), } - } + }; Ok(CompressedAdj { divide_comids, @@ -158,6 +344,8 @@ pub fn compress( cols, gauge_compressed, outflow_idx, + parent_offset, + conus_positions: active_vec, }) } @@ -306,14 +494,17 @@ mod tests { (Staid::new("0000000B"), 3, "comid400".to_string()), ], }; - let c = compress(&unioned, &conus_order).expect("compress"); + let c = compress(&unioned, &conus_order, true, None).expect("compress"); // Active = {0, 1, 2, 3, 4} → all 5. Compressed positions match. assert_eq!(c.divide_comids, conus_order); assert_eq!(c.rows, vec![2, 3, 4, 4]); assert_eq!(c.cols, vec![0, 1, 2, 3]); assert_eq!(c.gauge_compressed, vec![4, 3]); - // outflow_idx: gauge A at row 4 receives from cols 2, 3. - // gauge B at row 3 receives from col 1. + // Pins the `ddr_match: true` (DDR-faithful) convention: outflow_idx is + // the gauge's UPSTREAM cols — gauge A at row 4 receives from cols 2, 3; + // gauge B at row 3 from col 1. Both omit the gauge reach itself; see + // `outflow_idx_includes_the_gauge_reach_when_not_ddr_match` for the + // corrected convention. assert_eq!(c.outflow_idx[0], vec![2, 3]); assert_eq!(c.outflow_idx[1], vec![1]); } @@ -326,7 +517,7 @@ mod tests { edges: vec![(9, 7), (9, 5), (7, 2)], gauges: vec![(Staid::new("0000000A"), 9, "comid900".to_string())], }; - let c = compress(&unioned, &conus_order).expect("compress"); + let c = compress(&unioned, &conus_order, true, None).expect("compress"); assert_eq!(c.divide_comids, vec![Comid(200), Comid(500), Comid(700), Comid(900)]); // Edges in compressed space: (3,2), (3,1), (2,0). Same order as input edges, // but mapped through the compressed index space. @@ -346,7 +537,7 @@ mod tests { edges: vec![(0, 1)], gauges: vec![(Staid::new("0000000A"), 0, "x".to_string())], }; - let err = compress(&unioned, &conus_order).unwrap_err(); + let err = compress(&unioned, &conus_order, true, None).unwrap_err(); match err { crate::data::error::DataError::Malformed { .. } => {} other => panic!("expected Malformed, got {other:?}"), @@ -360,26 +551,164 @@ mod tests { edges: vec![], gauges: vec![], }; - let err = compress(&unioned, &conus_order).unwrap_err(); + let err = compress(&unioned, &conus_order, true, None).unwrap_err(); match err { crate::data::error::DataError::Malformed { .. } => {} other => panic!("expected Malformed, got {other:?}"), } } + #[test] + fn outflow_idx_includes_the_gauge_reach_when_not_ddr_match() { + // We do NOT know where along its reach a gauge physically sits, so the + // gauge reach's own lateral inflow MUST be counted. The MC solve at that + // reach already accumulates everything upstream by mass conservation, so + // the gauge's prediction is that ONE reach. + // + // Regression: gauge 01457000 (366.8 km2; its own reach is 250.1 km2 = 68% + // of the basin) read 1.58 m3/s against an observed 7.60 and a summed-Q' + // baseline of 7.38 -- a constant 0.215x suppression for 15 straight years, + // affecting 26/1841 gauges and biasing every ddrs-vs-baseline comparison. + // + // The affected case is a gauge that HAS incoming edges: two headwaters + // (CONUS positions 0, 1) draining into the gauge reach (position 2) -- + // exactly the 01457000 topology. + let conus_order = vec![Comid(73006562), Comid(73006585), Comid(73005764)]; + let unioned = UnionedCoo { + edges: vec![(2, 0), (2, 1)], + gauges: vec![(Staid::new("01457000"), 2, "73005764".to_string())], + }; + + let corrected = compress(&unioned, &conus_order, false, None).expect("compress"); + assert_eq!(corrected.gauge_compressed, vec![2]); + assert_eq!( + corrected.outflow_idx[0], + vec![2], + "ddr_match=false: outflow_idx must be the gauge's OWN reach, not its \ + upstream cols [0, 1]" + ); + + // And the DDR-faithful path is preserved byte-for-byte under the flag. + let ddr = compress(&unioned, &conus_order, true, None).expect("compress"); + assert_eq!(ddr.gauge_compressed, vec![2]); + assert_eq!( + ddr.outflow_idx[0], + vec![0, 1], + "ddr_match=true must reproduce DDR merit.py:226-234 (upstream cols)" + ); + } + #[test] fn outflow_idx_falls_back_to_self_when_no_incoming_edges() { // Gauge at CONUS-position 2 with no upstream edges in this batch - // (active = {2} as a single-node graph). DDR's fallback yields the - // gauge's own compressed index as the sole outflow column. + // (active = {2} as a single-node graph). Pins the `ddr_match: true` + // convention: DDR's empty-cols fallback yields the gauge's own + // compressed index. Under `ddr_match: false` this is not a fallback at + // all -- it is the general rule -- so both flag values agree here. let conus_order = vec![Comid(0), Comid(1), Comid(2)]; let unioned = UnionedCoo { edges: vec![], gauges: vec![(Staid::new("0000000A"), 2, "comid2".to_string())], }; - let c = compress(&unioned, &conus_order).expect("compress"); - assert_eq!(c.gauge_compressed, vec![0]); - assert_eq!(c.outflow_idx[0], vec![0], "self-edge fallback"); + let ddr = compress(&unioned, &conus_order, true, None).expect("compress"); + assert_eq!(ddr.gauge_compressed, vec![0]); + assert_eq!(ddr.outflow_idx[0], vec![0], "self-edge fallback"); + + let corrected = compress(&unioned, &conus_order, false, None).expect("compress"); + assert_eq!(corrected.outflow_idx[0], vec![0]); + } + + /// Subdivided CONUS space: parents 0 and 1 are single-piece headwaters, + /// parent 2 (the gauge reach) is split 4 ways into rows 2..6. + /// + /// 0 ─┐ + /// ├─> 2 → 3 → 4 → 5 (gauge reach, outlet = 5) + /// 1 ─┘ + fn subdivided_gauge_reach() -> (Vec, Vec, UnionedCoo) { + let conus_order = vec![ + Comid(73006562), + Comid(73006585), + Comid(73005764), + Comid(73005764), + Comid(73005764), + Comid(73005764), + ]; + let conus_parent_offset = vec![0, 1, 2, 6]; + let unioned = UnionedCoo { + // Two external edges onto the gauge reach's INLET piece, plus the + // internal chain 2→3→4→5. + edges: vec![(2, 0), (2, 1), (3, 2), (4, 3), (5, 4)], + // The subgraph builder resolves a gauge COMID to its parent's last + // row, so `gage_idx` is already the outlet piece. + gauges: vec![(Staid::new("01457000"), 5, "73005764".to_string())], + }; + (conus_order, conus_parent_offset, unioned) + } + + #[test] + fn compressed_parent_offset_is_the_conus_map_renumbered() { + let (order, off, unioned) = subdivided_gauge_reach(); + let c = compress(&unioned, &order, false, Some(&off)).expect("compress"); + // Active set is all 6 rows here, so compression is the identity and the + // compressed map equals the CONUS map. What matters is the CONTRACT: + // the map is expressed in the same space as `rows`/`cols`. + assert_eq!(c.parent_offset, Some(vec![0, 1, 2, 6])); + assert_eq!(c.divide_comids.len(), 6); + } + + #[test] + fn outflow_idx_reads_the_outlet_piece_of_a_subdivided_gauge_reach() { + let (order, off, unioned) = subdivided_gauge_reach(); + let c = compress(&unioned, &order, false, Some(&off)).expect("compress"); + assert_eq!(c.gauge_compressed, vec![5]); + assert_eq!( + c.outflow_idx[0], + vec![5], + "the gauge's whole reach discharges at its LAST piece (compressed row \ + `parent_offset[p + 1] - 1` = 5); pieces 2, 3, 4 each omit part of the \ + reach's own lateral inflow" + ); + } + + #[test] + fn compress_rejects_a_partial_sub_reach_chain() { + // Every edge touching row 2 is gone, so the gauge reach's inlet piece + // is missing from the active set. The parent's chain is then only 3 of + // its 4 rows — the engine would divide q' by 3 while the lengths were + // cut for 4, creating mass. That must be an error, not a silent pass. + let (order, off, _) = subdivided_gauge_reach(); + let unioned = UnionedCoo { + edges: vec![(4, 3), (5, 4)], + gauges: vec![(Staid::new("01457000"), 5, "73005764".to_string())], + }; + let err = compress(&unioned, &order, false, Some(&off)).unwrap_err(); + match err { + crate::data::error::DataError::Malformed { message, .. } => { + assert!( + message.contains("contiguously"), + "expected a partial-chain diagnostic, got: {message}" + ); + } + other => panic!("expected Malformed, got {other:?}"), + } + } + + #[test] + fn identity_parent_offset_leaves_outflow_idx_at_the_gauge_reach() { + // Un-subdivided store: `parent_offset` is `0..=n`, so every parent owns + // exactly one row and the outlet resolution must be a no-op. + let conus_order = vec![Comid(73006562), Comid(73006585), Comid(73005764)]; + let unioned = UnionedCoo { + edges: vec![(2, 0), (2, 1)], + gauges: vec![(Staid::new("01457000"), 2, "73005764".to_string())], + }; + let identity: Vec = vec![0, 1, 2, 3]; + let c = compress(&unioned, &conus_order, false, Some(&identity)).expect("compress"); + assert_eq!(c.parent_offset, Some(vec![0, 1, 2, 3])); + assert_eq!(c.outflow_idx[0], vec![2]); + // Byte-identical to passing no map at all. + let bare = compress(&unioned, &conus_order, false, None).expect("compress"); + assert_eq!(c.outflow_idx, bare.outflow_idx); } use crate::data::store::{GageMetadata, GageRow}; diff --git a/src/data/dataset.rs b/src/data/dataset.rs index d5dc2f9..5667e33 100644 --- a/src/data/dataset.rs +++ b/src/data/dataset.rs @@ -10,7 +10,7 @@ use std::sync::Arc; use ndarray::{Array1, Array2}; use crate::config::Config; -use crate::data::collate::{build_flow_scale, compress, union_subgraphs}; +use crate::data::collate::{build_flow_scale, compress, union_subgraphs, CompressedAdj}; use crate::data::dates::{Frequency, RhoWindow, TimeAxis}; use crate::data::error::{DataError, Result}; use crate::data::ids::{Comid, Staid}; @@ -28,8 +28,13 @@ use crate::sparse::SparseAdjacency; #[derive(Debug)] pub struct RoutingBatch { pub adjacency: SparseAdjacency, - /// Normalized attributes, shape `(N, F)`. Caller-major to match the + /// Normalized attributes, shape `(N_parent, F)`. Caller-major to match the /// KAN head input contract (`src/nn/kan_head.rs::KanHead::forward`). + /// + /// **Parent, not sub-reach, resolution.** Equal to `N` unless the adjacency + /// was built with `params.subdivision.enabled`, in which case the KAN runs + /// once per MERIT reach and `training::forward::gather_params_to_subreaches` + /// expands its outputs to the `N` routing rows. pub spatial_attributes_normalized: Array2, /// q' streamflow forcing, shape `(T_hours, N)`. Already multiplied by /// `flow_scale` per column. The flat `repeat-24` upsampling of @@ -90,7 +95,8 @@ struct StaticNetworkCache { adjacency: SparseAdjacency, outflow_idx: Vec>, flow_scale: Vec, - /// Normalized attributes, shape `(N_active, F)`. + /// Normalized attributes, shape `(N_parent_active, F)` — parent resolution + /// (see `RoutingBatch::spatial_attributes_normalized`). spatial_attributes_normalized: Array2, /// Full-period observations `(n_days_full, G)`. Sliced per `collate_window` call. full_observations: Array2, @@ -114,7 +120,9 @@ use burn::tensor::{backend::Backend, Int, Tensor, TensorData}; /// used for masking + comparison at loss time. pub struct RoutingTensors { pub adjacency: SparseAdjacency, - /// Normalized attributes, shape `(N, F)`. + /// Normalized attributes, shape `(N_parent, F)` — parent resolution, which + /// equals `N` unless the adjacency is subdivided. The KAN head consumes + /// this directly; its outputs are then gathered onto the `N` routing rows. pub spatial_attributes: Tensor, /// q' streamflow, shape `(T_hours, N)`. Not yet Autodiff-wrapped. pub q_prime: Tensor, @@ -160,7 +168,7 @@ impl RoutingBatch { group.extend(std::iter::repeat_n(g_idx as i32, segs.len())); } - // 2. Lift spatial_attributes (N, F) — already owned + contiguous after reversed_axes().into_owned(). + // 2. Lift spatial_attributes (N_parent, F) — already owned + contiguous after reversed_axes().into_owned(). let (rows, cols) = ( self.spatial_attributes_normalized.shape()[0], self.spatial_attributes_normalized.shape()[1], @@ -252,6 +260,38 @@ impl RoutingBatch { } } +// --------------------------------------------------------------------------- +// Per-row channel geometry +// --------------------------------------------------------------------------- + +/// Slice `(length_m, slope)` out of the CONUS store for every compressed row. +/// +/// **Index space matters here.** `conus.length_m` / `conus.slope` are +/// SUB-REACH arrays, so they must be indexed with `CompressedAdj::conus_positions` +/// — the CONUS sub-reach position each compressed row came from. +/// `conus.index.position(comid)` returns a PARENT position (`ConusAdjacencyStore`'s +/// two-index-spaces note) and is *not* interchangeable: under subdivision it +/// would give every piece of a parent the same geometry, read off an unrelated +/// row — the parent's undivided length instead of `L/m`, and a foreign slope. +/// Without subdivision the two agree, which is exactly what makes the wrong one +/// silent. +/// +/// Pinned by `tests/subdivision_integration.rs::subdivided_rows_get_their_own_length_and_slope`. +#[doc(hidden)] +pub fn slice_reach_geometry( + conus: &ConusAdjacencyStore, + compressed: &CompressedAdj, +) -> (Vec, Vec) { + let n = compressed.conus_positions.len(); + let mut length_m: Vec = Vec::with_capacity(n); + let mut slope: Vec = Vec::with_capacity(n); + for &pos in &compressed.conus_positions { + length_m.push(conus.length_m[pos]); + slope.push(conus.slope[pos]); + } + (length_m, slope) +} + // --------------------------------------------------------------------------- // MeritGagesDataset // --------------------------------------------------------------------------- @@ -292,6 +332,10 @@ pub struct MeritGagesDataset { gauge_std: OnceCell>, /// Whether the configured loss needs `gauge_std` (`loss.kind: nse-batch`). want_gauge_std: bool, + /// `params.ddr_match`, captured at open. Selects the `outflow_idx` + /// convention in `collate::compress`: `true` (default) reproduces DDR's + /// upstream-cols behaviour, `false` reads the gauge's own reach. + ddr_match: bool, /// Optional day-boundary discharge state cache. `None` ⇒ every code path /// byte-identical to no-cache behavior (`RoutingBatch::initial_state = None`). state_cache: Option, @@ -399,10 +443,15 @@ impl MeritGagesDataset { ); // ---------- 2. Attributes + statistics ---------- + // Requested in PARENT space: attributes are per-COMID, and `conus.order` + // repeats a COMID once per sub-reach when the store is subdivided — + // which would materialize up to `max_pieces` identical columns of the + // (F, N) matrix and leave the store's COMID index ambiguous. + // `parent_order` is `order` when the store is not subdivided. let attr_names: Vec = head_cfg.input_var_names.clone(); let (attrs, stats, means, stds) = if ds.attributes.len() == 1 { // Single-path: byte-identical to the pre-C0 behavior. - let attrs = AttributesStore::open(&ds.attributes[0], &attr_names, &conus.order)?; + let attrs = AttributesStore::open(&ds.attributes[0], &attr_names, &conus.parent_order)?; let stats_path = stats_path_from_attrs(&ds.attributes[0]); let stats = AttrStats::open(&stats_path)?; let means = stats.means_f32(&attr_names); @@ -410,7 +459,7 @@ impl MeritGagesDataset { (Arc::new(attrs), Arc::new(stats), means, stds) } else { // Multi-path: COMID-aligned merge across stores; NaN-fill per-store gaps. - let attrs = AttributesStore::open_multi(&ds.attributes, &attr_names, &conus.order)?; + let attrs = AttributesStore::open_multi(&ds.attributes, &attr_names, &conus.parent_order)?; let (means, stds) = load_merged_stats(&ds.attributes, &attr_names)?; // stats field is diagnostics-only (dead_code); use an empty placeholder. let placeholder_stats = AttrStats { @@ -519,6 +568,7 @@ impl MeritGagesDataset { .as_ref() .map(|e| e.loss.kind == crate::config::LossKind::NseBatch) .unwrap_or(false), + ddr_match: cfg.params.ddr_match, state_cache, leakance_impervious_threshold, }) @@ -618,20 +668,16 @@ impl MeritGagesDataset { let gauge_staids: Vec = unioned.gauges.iter().map(|(s, _, _)| s.clone()).collect(); - let compressed = compress(&unioned, &self.conus.order)?; + let compressed = compress( + &unioned, + &self.conus.order, + self.ddr_match, + Some(&self.conus.parent_offset), + )?; let n = compressed.divide_comids.len(); // ----- 2. SparseAdjacency: rows/cols + length/slope sliced ----- - let mut length_m: Vec = Vec::with_capacity(n); - let mut slope: Vec = Vec::with_capacity(n); - for c in &compressed.divide_comids { - let pos = self.conus.index.position(c).ok_or_else(|| DataError::Malformed { - path: self.conus.path.clone(), - message: format!("compressed COMID {c:?} not found in CONUS order"), - })?; - length_m.push(self.conus.length_m[pos]); - slope.push(self.conus.slope[pos]); - } + let (length_m, slope) = slice_reach_geometry(&self.conus, &compressed); let values: Vec = vec![1.0; compressed.rows.len()]; let adjacency = SparseAdjacency { n, @@ -640,6 +686,10 @@ impl MeritGagesDataset { values, length_m, slope, + // Reach-subdivision map in compressed space. Identity (hence a + // no-op split) unless the store was built with + // `params.subdivision.enabled`. + parent_offset: compressed.parent_offset.clone(), }; // ----- 3. flow_scale + q_prime read & fuse ----- @@ -687,7 +737,13 @@ impl MeritGagesDataset { let temp_hourly = self.read_temp_window(window, &compressed.divide_comids, n)?; // ----- 4. Attributes: slice + fill_nans + normalize + transpose ----- - let spatial_attributes_normalized = self.finalize_attrs(&compressed.divide_comids, n); + // PARENT resolution: attributes are per-COMID and sub-reaches inherit + // their parent's hydraulics, so the KAN runs once per MERIT reach and + // `training::forward` gathers its outputs onto the pieces. Slicing at + // sub-reach resolution would build ~5x identical rows at `max_pieces: 8`. + let parent_comids = compressed.parent_comids(); + let spatial_attributes_normalized = + self.finalize_attrs(&parent_comids, parent_comids.len()); // ----- 5. Observations (present-in-adjacency STAIDs; missing→error) ----- let observations = self.observations.read_window(window, &gauge_staids)?; @@ -1001,20 +1057,16 @@ impl MeritGagesDataset { } let gauge_staids: Vec = unioned.gauges.iter().map(|(s, _, _)| s.clone()).collect(); - let compressed = compress(&unioned, &self.conus.order)?; + let compressed = compress( + &unioned, + &self.conus.order, + self.ddr_match, + Some(&self.conus.parent_offset), + )?; let n = compressed.divide_comids.len(); // 2. SparseAdjacency. - let mut length_m: Vec = Vec::with_capacity(n); - let mut slope: Vec = Vec::with_capacity(n); - for c in &compressed.divide_comids { - let pos = self.conus.index.position(c).ok_or_else(|| DataError::Malformed { - path: self.conus.path.clone(), - message: format!("compressed COMID {c:?} not found in CONUS order"), - })?; - length_m.push(self.conus.length_m[pos]); - slope.push(self.conus.slope[pos]); - } + let (length_m, slope) = slice_reach_geometry(&self.conus, &compressed); let values: Vec = vec![1.0; compressed.rows.len()]; let adjacency = SparseAdjacency { n, @@ -1023,6 +1075,10 @@ impl MeritGagesDataset { values, length_m, slope, + // Reach-subdivision map in compressed space. Identity (hence a + // no-op split) unless the store was built with + // `params.subdivision.enabled`. + parent_offset: compressed.parent_offset.clone(), }; // 3. flow_scale. @@ -1033,8 +1089,11 @@ impl MeritGagesDataset { n, ); - // 4. Normalized attributes (N, F). - let spatial_attributes_normalized = self.finalize_attrs(&compressed.divide_comids, n); + // 4. Normalized attributes (N_parent, F) — see the note in `collate`: + // the KAN runs at parent resolution and is gathered onto sub-reaches. + let parent_comids = compressed.parent_comids(); + let spatial_attributes_normalized = + self.finalize_attrs(&parent_comids, parent_comids.len()); // 5. Full-period observations: read the entire time axis at once. let full_rho = RhoWindow { diff --git a/src/data/store/icechunk.rs b/src/data/store/icechunk.rs index 1477112..58dbccc 100644 --- a/src/data/store/icechunk.rs +++ b/src/data/store/icechunk.rs @@ -345,6 +345,84 @@ pub(crate) fn daily_to_hourly_trim(daily: &Array2, n_hourly: usize) -> Arra hourly } +/// Daily→hourly upsampling mode for DAILY Q' stores, selected once per process +/// from the `DDRS_QPRIME_INTERP` env var (diagnostic knob, mirrors +/// `DDRS_HOURLY_DUMP`). Absent/`nearest` keeps the flat repeat-24 path +/// byte-identical. `linear`/`quadratic` interpolate through day-CENTER values +/// (day d anchors at hour 24d+12), reading ±1 context day where the store +/// allows so 15-day eval chunks tile without edge clamps. +/// Scope: the icechunk daily store only — the global zarr-v2 reader REJECTS a +/// non-nearest setting rather than silently ignoring it. +#[derive(Clone, Copy, PartialEq, Debug)] +pub(crate) enum QPrimeInterp { + Nearest, + Linear, + Quadratic, +} + +pub(crate) fn qprime_interp() -> QPrimeInterp { + static MODE: std::sync::OnceLock = std::sync::OnceLock::new(); + *MODE.get_or_init(|| { + let mode = match std::env::var("DDRS_QPRIME_INTERP").as_deref() { + Err(_) | Ok("nearest") => QPrimeInterp::Nearest, + Ok("linear") => QPrimeInterp::Linear, + Ok("quadratic") => QPrimeInterp::Quadratic, + Ok(other) => { + panic!("DDRS_QPRIME_INTERP={other} not one of nearest|linear|quadratic") + } + }; + if mode != QPrimeInterp::Nearest { + eprintln!(" q' upsampling: {mode:?} (DDRS_QPRIME_INTERP)"); + } + mode + }) +} + +/// Interpolating daily→hourly. `daily_ext` carries `left_pad` (0 or 1) context +/// rows before window day 0 and possibly one after the window; neighbors +/// outside the slab clamp to the slab edge. Hour h (window coordinates) +/// evaluates at t = (h + 0.5)/24 days; day d's value anchors at t = d + 0.5. +/// Both interpolants are clamped at 0 (the quadratic can undershoot); neither +/// conserves the daily mean exactly — linear applies a [1,6,1]/8 kernel to +/// interior day means, quadratic shifts them by curvature/24. Diagnostic only. +pub(crate) fn daily_to_hourly_interp( + daily_ext: &Array2, + n_hourly: usize, + left_pad: usize, + mode: QPrimeInterp, +) -> Array2 { + let (n_ext, n_div) = daily_ext.dim(); + let c = |k: isize, j: usize| -> f32 { + let idx = (k + left_pad as isize).clamp(0, n_ext as isize - 1) as usize; + daily_ext[(idx, j)] + }; + let mut hourly = Array2::::zeros((n_hourly, n_div)); + for h in 0..n_hourly { + let d = (h / 24) as isize; + // Position within day d relative to its center, in [-0.5, 0.5). + let x = (h as f32 + 0.5) / 24.0 - (d as f32 + 0.5); + for j in 0..n_div { + let v = match mode { + QPrimeInterp::Nearest => c(d, j), + QPrimeInterp::Linear => { + if x < 0.0 { + (1.0 + x) * c(d, j) + (-x) * c(d - 1, j) + } else { + (1.0 - x) * c(d, j) + x * c(d + 1, j) + } + } + QPrimeInterp::Quadratic => { + // Lagrange parabola through centers d-1, d, d+1. + let (vm, v0, vp) = (c(d - 1, j), c(d, j), c(d + 1, j)); + v0 + 0.5 * (vp - vm) * x + 0.5 * (vm - 2.0 * v0 + vp) * x * x + } + }; + hourly[(h, j)] = v.max(0.0); + } + } + hourly +} + /// Collapse a `(n_days * 24, N)` hourly slab to `(n_days, N)` by averaging /// each 24-hour block. Q' is a rate (m³/s): the daily value is the day's /// mean flow, so total daily volume is preserved. @@ -506,15 +584,52 @@ impl StreamflowStore { } } + /// Read `n_days` daily rows from `start` plus up to one context day on + /// each side (for the interpolating upsamplers). Returns the extended slab + /// and `left_pad` (1 iff the store has a day before `start`). Only called + /// on `Frequency::Daily` stores, where `n_time` counts days. + fn read_daily_with_context( + &self, + start: NaiveDate, + n_days: usize, + comids: &[Comid], + ) -> Result<(Array2, usize)> { + debug_assert!(matches!(self.resolution, Frequency::Daily)); + let offset = (start - self.time_start).num_days(); + let left = usize::from(offset >= 1); + let right = usize::from(offset + n_days as i64 + 1 <= self.n_time as i64); + let ext_start = start - chrono::Duration::days(left as i64); + let daily = self.read_window_daily(ext_start, n_days + left + right, comids)?; + Ok((daily, left)) + } + + /// Daily-store upsampling shared by `read_window` / `read_test_window`: + /// repeat-24 (unchanged default) or, under `DDRS_QPRIME_INTERP`, + /// day-center interpolation with ±1-day context. + fn upsample_daily( + &self, + start: NaiveDate, + n_days: usize, + n_hourly: usize, + comids: &[Comid], + ) -> Result> { + let mode = qprime_interp(); + if mode == QPrimeInterp::Nearest { + let daily = self.read_window_daily(start, n_days, comids)?; + Ok(daily_to_hourly_trim(&daily, n_hourly)) + } else { + let (daily_ext, left_pad) = self.read_daily_with_context(start, n_days, comids)?; + Ok(daily_to_hourly_interp(&daily_ext, n_hourly, left_pad, mode)) + } + } + /// Read `Qr` for `window` and `comids`. Returns `(n_hourly, N)` f32. /// Daily stores upsample via repeat-24 + trailing-day trim (unchanged); /// hourly stores slice the native axis directly — no upsampling. pub fn read_window(&self, window: &RhoWindow, comids: &[Comid]) -> Result> { match self.resolution { Frequency::Daily => { - let daily = - self.read_window_daily(window.window_start, window.rho_days, comids)?; - Ok(daily_to_hourly_trim(&daily, window.n_hourly())) + self.upsample_daily(window.window_start, window.rho_days, window.n_hourly(), comids) } Frequency::Hourly => { let start = self.native_start_index(window.window_start)?; @@ -532,9 +647,7 @@ impl StreamflowStore { ) -> Result> { match self.resolution { Frequency::Daily => { - let daily = - self.read_window_daily(window.window_start, window.n_days, comids)?; - Ok(daily_to_hourly_trim(&daily, window.n_hourly())) + self.upsample_daily(window.window_start, window.n_days, window.n_hourly(), comids) } Frequency::Hourly => { let start = self.native_start_index(window.window_start)?; @@ -825,6 +938,62 @@ mod tests { } } + #[test] + fn daily_to_hourly_interp_nearest_matches_trim() { + use ndarray::Array2; + let daily = + Array2::from_shape_vec((3, 2), vec![1.0, 10.0, 2.0, 20.0, 3.0, 30.0]).unwrap(); + let trim = daily_to_hourly_trim(&daily, 47); + let interp = daily_to_hourly_interp(&daily, 47, 0, QPrimeInterp::Nearest); + assert_eq!(trim, interp); + } + + #[test] + fn daily_to_hourly_interp_constant_is_exact_all_modes() { + use ndarray::Array2; + let daily = Array2::from_elem((4, 1), 5.0f32); + for mode in [QPrimeInterp::Nearest, QPrimeInterp::Linear, QPrimeInterp::Quadratic] { + let hourly = daily_to_hourly_interp(&daily, 96, 1, mode); + for h in 0..96 { + assert!((hourly[(h, 0)] - 5.0).abs() < 1e-6, "{mode:?} h={h}"); + } + } + } + + #[test] + fn daily_to_hourly_interp_linear_ramp_conserves_interior_day_means() { + use ndarray::Array2; + // Daily ramp 1..=5 with 1 context day each side of a 3-day window. + // Linear AND quadratic interpolation reproduce a linear trend exactly, + // so each interior window day's 24-hour mean equals its daily value. + let daily_ext = + Array2::from_shape_vec((5, 1), vec![1.0, 2.0, 3.0, 4.0, 5.0]).unwrap(); + for mode in [QPrimeInterp::Linear, QPrimeInterp::Quadratic] { + let hourly = daily_to_hourly_interp(&daily_ext, 72, 1, mode); + for d in 0..3 { + let mean: f32 = + (0..24).map(|h| hourly[(24 * d + h, 0)]).sum::() / 24.0; + let expect = daily_ext[(d + 1, 0)]; + assert!( + (mean - expect).abs() < 1e-5, + "{mode:?} day {d}: mean {mean} vs {expect}" + ); + } + } + } + + #[test] + fn daily_to_hourly_interp_quadratic_clamps_at_zero() { + use ndarray::Array2; + // Sharp peak: the parabola through (0, 12, 0) undershoots next to the + // peak day; outputs must clamp at 0 (Q' is nonnegative). + let daily_ext = Array2::from_shape_vec((3, 1), vec![0.0, 12.0, 0.0]).unwrap(); + let hourly = daily_to_hourly_interp(&daily_ext, 72, 0, QPrimeInterp::Quadratic); + for h in 0..72 { + assert!(hourly[(h, 0)] >= 0.0, "h={h} went negative"); + } + } + #[test] fn open_streamflow_store_if_present() { let p = Path::new("/mnt/ssd1/data/icechunk/merit_dhbv2_UH_retrospective.ic"); diff --git a/src/data/store/zarr.rs b/src/data/store/zarr.rs index 4203a23..51f3adf 100644 --- a/src/data/store/zarr.rs +++ b/src/data/store/zarr.rs @@ -26,12 +26,38 @@ use crate::data::error::{DataError, Result}; use crate::data::ids::{Comid, IdIndex, Staid}; /// Static CONUS-wide network state. Loaded once at dataset construction. +/// +/// ## Two index spaces (reach subdivision) +/// +/// When the store was built with `params.subdivision.enabled`, one MERIT reach +/// occupies several consecutive rows. `order` then carries **duplicate** COMIDs +/// (one per sub-reach), so a COMID→row lookup on it would be ambiguous. +/// +/// - **Parent space** (`parent_order`, length `n_parent`): one entry per MERIT +/// reach. `index` is built from THIS, so `index.position(comid)` always +/// returns a *parent* position. +/// - **Sub-reach space** (`order`, `length_m`, `slope`, `indices_*`, length `n`): +/// what the solver sees. Parent `p` owns rows +/// `parent_offset[p]..parent_offset[p + 1]`, ordered upstream→downstream, so +/// its outlet — the row a gauge must be read at — is `parent_offset[p+1] - 1`. +/// +/// Stores written before subdivision existed carry neither array; `open` +/// synthesizes the identity (`parent_order == order`, `parent_offset == 0..=n`) +/// so they keep loading unchanged and every consumer sees one uniform contract. pub struct ConusAdjacencyStore { pub path: PathBuf, /// COMIDs in topological order — element `i` is the COMID at zarr position `i`. + /// Contains duplicates when the store is subdivided. pub order: Vec, - /// `IdIndex` mapping COMID → topological position (for cross-store lookups). + /// `IdIndex` mapping COMID → **parent** position. Built from `parent_order`, + /// never from `order` (see the type-level note on the two index spaces). pub index: IdIndex, + /// One COMID per MERIT reach, in topological order. Identical to `order` + /// when the store is not subdivided. + pub parent_order: Vec, + /// Length `parent_order.len() + 1`. Rows `[parent_offset[p], parent_offset[p+1])` + /// of the sub-reach arrays belong to parent `p`. `0..=n` when not subdivided. + pub parent_offset: Vec, /// Per-reach channel length in metres, aligned to `order`. pub length_m: Array1, /// Per-reach channel slope (dimensionless), aligned to `order`. @@ -56,7 +82,44 @@ impl ConusAdjacencyStore { let order_i32 = read_array_i32(&storage, &path, "/order")?; let order: Vec = order_i32.into_iter().map(|c| Comid(c as i64)).collect(); let n = order.len(); - let index = IdIndex::new(order.clone()); + + // Parent map. Absent in every store written before reach subdivision + // (including the engine's own exports), so a missing array is NOT an + // error — it means "one row per reach", and the identity below makes + // such a store indistinguishable from a subdivided one with all m = 1. + let parent_order: Vec = + match try_read_array_i32(&storage, "/parent_order") { + Some(v) => v.into_iter().map(|c| Comid(c as i64)).collect(), + None => order.clone(), + }; + let parent_offset: Vec = match try_read_array_i32(&storage, "/parent_offset") { + Some(v) => v, + None => (0..=n as i32).collect(), + }; + if parent_offset.len() != parent_order.len() + 1 { + return Err(DataError::Malformed { + path: path.clone(), + message: format!( + "parent_offset must have parent_order.len() + 1 entries: {} vs {}", + parent_offset.len(), + parent_order.len() + 1 + ), + }); + } + // The offsets must partition the sub-reach rows exactly; a truncated or + // stale parent map would silently mis-address every gauge. + if parent_offset.first() != Some(&0) || parent_offset.last() != Some(&(n as i32)) { + return Err(DataError::Malformed { + path: path.clone(), + message: format!( + "parent_offset must run 0..{n}, got {:?}..{:?}", + parent_offset.first(), + parent_offset.last() + ), + }); + } + // Built from `parent_order`: `order` has duplicates once subdivided. + let index = IdIndex::new(parent_order.clone()); let length_m = Array1::from(read_array_f32(&storage, &path, "/length_m")?); let slope = Array1::from(read_array_f32(&storage, &path, "/slope")?); @@ -89,6 +152,8 @@ impl ConusAdjacencyStore { path, order, index, + parent_order, + parent_offset, length_m, slope, indices_0, @@ -97,6 +162,21 @@ impl ConusAdjacencyStore { nnz, }) } + + /// Number of MERIT reaches (parent rows). Equals [`Self::n`] when the store + /// is not subdivided. + #[inline] + pub fn n_parent(&self) -> usize { + self.parent_order.len() + } + + /// Last (most downstream) sub-reach row owned by `parent`. A gauge on that + /// reach must be read here: any earlier piece omits the downstream fraction + /// of the reach's own lateral inflow. + #[inline] + pub fn outlet_row(&self, parent: usize) -> usize { + self.parent_offset[parent + 1] as usize - 1 + } } /// Per-gauge upstream subgraph — indices reference *CONUS* positions, not @@ -230,6 +310,19 @@ fn read_array_i32(storage: &ReadableStorage, store_path: &Path, array_path: &str .map_err(|e| zarr_err(store_path, e)) } +/// Read an optional int32 array: `None` when the array is absent OR unreadable. +/// +/// Used for the subdivision parent map, which pre-subdivision stores (every +/// engine export, and every ddrs cache built before BUILDER_VERSION 2) simply +/// do not have. Collapsing "missing" and "corrupt" is acceptable here only +/// because the caller's fallback — the identity map — is itself a valid, +/// fully-consistent answer that `open` then range-checks against `n`. +fn try_read_array_i32(storage: &ReadableStorage, array_path: &str) -> Option> { + let arr = ZarrArray::open(storage.clone(), array_path).ok()?; + let subset = arr.subset_all(); + arr.retrieve_array_subset::>(&subset).ok() +} + fn read_array_f32(storage: &ReadableStorage, store_path: &Path, array_path: &str) -> Result> { let arr = ZarrArray::open(storage.clone(), array_path).map_err(|e| zarr_err(store_path, e))?; let subset = arr.subset_all(); @@ -254,6 +347,8 @@ mod tests { let index = IdIndex::new(order.clone()); ConusAdjacencyStore { path: PathBuf::from("/dev/null"), + parent_order: order.clone(), + parent_offset: (0..=n as i32).collect(), order, index, length_m: Array1::zeros(n), diff --git a/src/data/store/zarr_qprime.rs b/src/data/store/zarr_qprime.rs index 0b5377f..5fcbbed 100644 --- a/src/data/store/zarr_qprime.rs +++ b/src/data/store/zarr_qprime.rs @@ -276,6 +276,7 @@ impl GlobalStreamflowStore { /// `(n_hourly, N)` read for a training rho-window — daily values /// repeated 24× and trimmed, identical to the icechunk store. pub fn read_window(&self, window: &RhoWindow, comids: &[Comid]) -> Result> { + Self::reject_interp_env(); let daily = self.read_window_daily(window.window_start, window.rho_days, comids)?; Ok(daily_to_hourly_trim(&daily, window.n_hourly())) } @@ -286,9 +287,21 @@ impl GlobalStreamflowStore { window: &crate::data::TestWindow, comids: &[Comid], ) -> Result> { + Self::reject_interp_env(); let daily = self.read_window_daily(window.window_start, window.n_days, comids)?; Ok(daily_to_hourly_trim(&daily, window.n_hourly())) } + + /// `DDRS_QPRIME_INTERP` is only implemented for the icechunk daily store; + /// fail loudly rather than silently running nearest under a set env var. + fn reject_interp_env() { + use crate::data::store::icechunk::{qprime_interp, QPrimeInterp}; + assert!( + qprime_interp() == QPrimeInterp::Nearest, + "DDRS_QPRIME_INTERP is set but the global zarr-v2 q' reader only \ + supports nearest (repeat-24); unset it or use an icechunk store" + ); + } } /// A zone group is a directory with zarr v2 group metadata and a diff --git a/src/dump_parameters.rs b/src/dump_parameters.rs index 93a7b40..3a23ae6 100644 --- a/src/dump_parameters.rs +++ b/src/dump_parameters.rs @@ -127,19 +127,25 @@ where eprintln!("opening CONUS adjacency: {}", conus_path.display()); let conus = ConusAdjacencyStore::open(conus_path) .map_err(|e| CliError::Other(Box::new(e)))?; - let n_reaches = conus.order.len(); + // PARENT space throughout: this dump is per-COMID (the netCDF's COMID + // dimension is a coordinate, so duplicates would be malformed), and the KAN + // itself runs once per MERIT reach. `parent_order` == `order` and + // `n_parent()` == `n` unless the store was built with + // `params.subdivision.enabled`. + let n_reaches = conus.n_parent(); eprintln!("CONUS reaches: {n_reaches}"); // ---------- 2. Attributes + z-score stats ---------- eprintln!("opening attributes: {} file(s)", ds.attributes.len()); - let (attrs, means, stds) = open_attrs_and_stats(ds, &head_cfg.input_var_names, &conus.order) - .map_err(|e| CliError::Other(Box::new(e)))?; + let (attrs, means, stds) = + open_attrs_and_stats(ds, &head_cfg.input_var_names, &conus.parent_order) + .map_err(|e| CliError::Other(Box::new(e)))?; // ---------- 3. Build normalized (N, F) attribute tensor ---------- // Mirrors `MeritGagesDataset::finalize_attrs` (`src/data/dataset.rs:403`). let f = head_cfg.input_var_names.len(); let mut a: Array2 = Array2::zeros((f, n_reaches)); - for (out_col, comid) in conus.order.iter().enumerate() { + for (out_col, comid) in conus.parent_order.iter().enumerate() { if let Some(src_col) = attrs.index.position(comid) { for fi in 0..f { a[(fi, out_col)] = attrs.attrs[(fi, src_col)]; @@ -293,8 +299,12 @@ where // ---------- 6. Write NetCDF4 ---------- let slope_lb = cfg.params.attribute_minimums.slope; - let comids_i64: Vec = conus.order.iter().map(|c| c.0).collect(); - let slope_clamped: Vec = conus.slope.iter().map(|&s| s.max(slope_lb)).collect(); + let comids_i64: Vec = conus.parent_order.iter().map(|c| c.0).collect(); + // Slope is inherited unchanged by every piece (`adjacency::subdivide`), so + // a parent's slope is its first piece's row. + let slope_clamped: Vec = (0..n_reaches) + .map(|p| conus.slope[conus.parent_offset[p] as usize].max(slope_lb)) + .collect(); write_netcdf( output_path, @@ -349,18 +359,24 @@ where eprintln!("opening CONUS adjacency: {}", conus_path.display()); let conus = ConusAdjacencyStore::open(conus_path) .map_err(|e| CliError::Other(Box::new(e)))?; - let n_reaches = conus.order.len(); + // PARENT space throughout: this dump is per-COMID (the netCDF's COMID + // dimension is a coordinate, so duplicates would be malformed), and the KAN + // itself runs once per MERIT reach. `parent_order` == `order` and + // `n_parent()` == `n` unless the store was built with + // `params.subdivision.enabled`. + let n_reaches = conus.n_parent(); eprintln!("CONUS reaches: {n_reaches}"); // ---------- 2. Attributes + z-score stats ---------- eprintln!("opening attributes: {} file(s)", ds.attributes.len()); - let (attrs, means, stds) = open_attrs_and_stats(ds, &head_cfg.input_var_names, &conus.order) - .map_err(|e| CliError::Other(Box::new(e)))?; + let (attrs, means, stds) = + open_attrs_and_stats(ds, &head_cfg.input_var_names, &conus.parent_order) + .map_err(|e| CliError::Other(Box::new(e)))?; // ---------- 3. Build normalized (N, F) attribute tensor ---------- let f = head_cfg.input_var_names.len(); let mut a: Array2 = Array2::zeros((f, n_reaches)); - for (out_col, comid) in conus.order.iter().enumerate() { + for (out_col, comid) in conus.parent_order.iter().enumerate() { if let Some(src_col) = attrs.index.position(comid) { for fi in 0..f { a[(fi, out_col)] = attrs.attrs[(fi, src_col)]; @@ -520,8 +536,12 @@ where // ---------- 6. Write NetCDF4 ---------- let slope_lb = cfg.params.attribute_minimums.slope; - let comids_i64: Vec = conus.order.iter().map(|c| c.0).collect(); - let slope_clamped: Vec = conus.slope.iter().map(|&s| s.max(slope_lb)).collect(); + let comids_i64: Vec = conus.parent_order.iter().map(|c| c.0).collect(); + // Slope is inherited unchanged by every piece (`adjacency::subdivide`), so + // a parent's slope is its first piece's row. + let slope_clamped: Vec = (0..n_reaches) + .map(|p| conus.slope[conus.parent_offset[p] as usize].max(slope_lb)) + .collect(); write_netcdf( output_path, diff --git a/src/routing/courant_probe.rs b/src/routing/courant_probe.rs new file mode 100644 index 0000000..5210887 --- /dev/null +++ b/src/routing/courant_probe.rs @@ -0,0 +1,207 @@ +//! Courant / Muskingum-window diagnostic for the S18'/S19' positivity clamp. +//! +//! Measures — on the REAL routing chain, not a re-derivation — what +//! `params.enforce_positivity` actually does to a CONUS mini-batch: +//! +//! * `negative solves before clamp` (the atomic counters in `mmc_op`), +//! * the Courant number `Cr = dt / K`, before and after the K floor, +//! * the effective Muskingum `X` the chain used (`x_eff_out`), +//! * `k_musk / k_raw`, i.e. how much the floor inflates travel time, +//! * `min c1` / `min c3`, the quantities the clamp claims to keep >= 0. +//! +//! It drives [`forward_chain_inner`] directly in the same loop shape as +//! `MuskingumCunge::forward` (q_next fed back as q_t), so every number below +//! comes out of the identical kernel sequence training runs. Nothing here +//! touches the tape or the numerics. + +use std::sync::Arc; + +use burn::tensor::{backend::Backend, Tensor, TensorPrimitive}; + +use crate::config::Config; +use crate::routing::mmc_op::{ + forward_chain_inner, forward_saved_idx, negative_solve_stats, reset_negative_solve_stats, +}; +use crate::sparse::CsrPattern; + +/// Inner-backend snapshot of everything `forward_chain_inner` consumes. +/// Produced by [`crate::routing::MuskingumCunge::probe_inputs`] after +/// `setup_inputs` (so the hot-start `q0` and the CSR pattern are the real +/// ones, not a re-derivation). +pub struct ProbeInputs { + pub pattern: Arc, + pub n: Tensor, + pub q_spatial: Tensor, + pub p_spatial: Tensor, + pub length: Tensor, + pub slope: Tensor, + pub x_storage: Tensor, + /// `(T_hours, N)` hourly forcing, unclamped (the probe applies S0's clamp). + pub q_prime: Tensor, + /// Per-row parent piece count, or `None` when the network is un-subdivided. + /// The probe applies it exactly as `MuskingumCunge::forward` does — AFTER + /// the `discharge` clamp — so a subdivided network is routed with the same + /// `q'/m` lateral inflow training would use, not `m×` too much water. + pub pieces_per_row: Option>, + /// Window-start discharge as `setup_inputs` left it. + pub q0: Tensor, + pub n_segments: usize, +} + +/// Everything the probe measured. Sampled vectors hold one entry per +/// (reach, sampled timestep); the solve counters are exact over ALL timesteps. +pub struct CourantReport { + pub n_reaches: usize, + pub n_steps: usize, + pub n_sampled_steps: usize, + /// Exact, over every routed timestep. + pub neg_solves: u64, + pub total_solves: u64, + /// `dt / k_raw` — the Courant number the physics asks for. + pub cr_raw: Vec, + /// `dt / k_musk` — after the S18' floor (identical to `cr_raw` when off). + pub cr: Vec, + /// The `x_eff` the chain fed into c1..c4. + pub x_eff: Vec, + /// `x_cunge`, the S19 Cunge X BEFORE the S19' stability cap, recomputed + /// from the same saved `top_width`/`celerity` the chain used. Equals + /// `x_eff` when the clamp is off. + pub x_cunge: Vec, + /// `k_musk / k_raw` — 1.0 wherever the S18' floor did not bite. + pub k_ratio: Vec, + pub c1: Vec, + pub c3: Vec, + /// Total network discharge `Σ_i q_t[i]` at each of the first + /// `trace_steps` timesteps (index 0 = the hot-start `Q_0` itself). + /// Used to measure how long an inflated cold start takes to wash out. + pub trace_total_q: Vec, +} + +fn to_vec(p: I::FloatTensorPrimitive) -> Vec { + Tensor::::from_primitive(TensorPrimitive::Float(p)) + .into_data() + .to_vec::() + .expect("f32 tensor") +} + +/// Route `n_steps` timesteps, harvesting diagnostics every `sample_every` +/// steps. `cfg` decides `enforce_positivity` — pass two configs to compare. +pub fn run_courant_probe( + cfg: &Config, + inp: &ProbeInputs, + n_steps: usize, + sample_every: usize, + trace_steps: usize, +) -> CourantReport +where + I::FloatTensorPrimitive: 'static, + I::Device: 'static, +{ + let dt = crate::routing::mmc::DT_SECONDS; + let n = inp.n_segments; + let discharge_lb = cfg.params.attribute_minimums.discharge; + + let t_avail = inp.q_prime.dims()[0]; + let steps = n_steps.min(t_avail.saturating_sub(1)); + // Mirrors `MuskingumCunge::forward`: one clamp on the whole forcing block, + // THEN the subdivision divisor (dividing first would floor each of the `m` + // pieces independently and create mass), and the initial state clamped once. + let q_prime = inp.q_prime.clone().clamp_min(discharge_lb); + let q_prime = match inp.pieces_per_row.as_ref() { + Some(d) => q_prime / d.clone().unsqueeze_dim::<2>(0), + None => q_prime, + }; + let mut q_t = inp.q0.clone().clamp_min(discharge_lb); + + let length: Vec = inp.length.clone().into_data().to_vec::().expect("f32"); + let slope: Vec = inp.slope.clone().into_data().to_vec::().expect("f32"); + + reset_negative_solve_stats(); + + let mut rep = CourantReport { + n_reaches: n, + n_steps: steps, + n_sampled_steps: 0, + neg_solves: 0, + total_solves: 0, + cr_raw: Vec::new(), + cr: Vec::new(), + x_eff: Vec::new(), + x_cunge: Vec::new(), + k_ratio: Vec::new(), + c1: Vec::new(), + c3: Vec::new(), + trace_total_q: Vec::new(), + }; + + let total_q = |q: &Tensor| -> f64 { + q.clone() + .into_data() + .to_vec::() + .expect("f32") + .iter() + .map(|&v| v as f64) + .sum() + }; + if trace_steps > 0 { + rep.trace_total_q.push(total_q(&q_t)); + } + + for t in 1..=steps { + let q_prime_t = q_prime + .clone() + .slice([(t - 1)..t, 0..n]) + .reshape([n]); + let mut x_eff_out: Option = None; + let mut leak_out = None; + let (q_next, saved) = forward_chain_inner::( + cfg, + &inp.pattern, + inp.n.clone(), + inp.q_spatial.clone(), + inp.p_spatial.clone(), + q_t.clone(), + q_prime_t, + inp.length.clone(), + inp.slope.clone(), + inp.x_storage.clone(), + None, + &mut leak_out, + &mut x_eff_out, + /* track_neg */ true, + ); + + if t % sample_every == 0 { + let k = to_vec::(saved[forward_saved_idx::K_MUSKINGUM].clone()); + let cel = to_vec::(saved[forward_saved_idx::CELERITY].clone()); + let c1 = to_vec::(saved[forward_saved_idx::C1].clone()); + let c3 = to_vec::(saved[forward_saved_idx::C3].clone()); + let x = to_vec::(x_eff_out.clone().expect("x_eff always written")); + let tw = to_vec::(saved[forward_saved_idx::TOP_WIDTH].clone()); + let qt = q_t.clone().into_data().to_vec::().expect("f32"); + for i in 0..n { + let k_raw = length[i] / cel[i]; + rep.cr_raw.push(dt / k_raw); + rep.cr.push(dt / k[i]); + rep.k_ratio.push(k[i] / k_raw); + rep.x_eff.push(x[i]); + // Mirrors S19 exactly (`forward_chain_inner`). + let w = qt[i] / (tw[i] * slope[i] * cel[i] * length[i] + 1e-12); + rep.x_cunge.push((0.5 * (1.0 - w)).clamp(0.0, 0.5)); + rep.c1.push(c1[i]); + rep.c3.push(c3[i]); + } + rep.n_sampled_steps += 1; + } + + q_t = Tensor::from_primitive(TensorPrimitive::Float(q_next)); + if t < trace_steps { + rep.trace_total_q.push(total_q(&q_t)); + } + } + + let (neg, total) = negative_solve_stats(); + rep.neg_solves = neg; + rep.total_solves = total; + rep +} diff --git a/src/routing/mmc.rs b/src/routing/mmc.rs index dbda5cb..1b8f84d 100644 --- a/src/routing/mmc.rs +++ b/src/routing/mmc.rs @@ -32,6 +32,34 @@ use crate::sparse::{triangular_csr_solve, AValuesAssembler, CsrPattern, SparseAd /// Hardcoded routing timestep in seconds. Matches `self.t` in `mmc.py:192`. pub const DT_SECONDS: f32 = 3600.0; +/// Per-row lateral-inflow divisor from a `parent_offset` map: every row owned +/// by parent `p` gets `m_p = parent_offset[p + 1] - parent_offset[p]`, the +/// number of sub-reach pieces the parent was split into. +/// +/// Returns `None` when no parent owns more than one row — the network is not +/// subdivided, every divisor would be `1.0`, and the caller can skip the +/// division entirely instead of emitting a no-op tensor op. +fn pieces_per_row_divisor(parent_offset: &[i32], n: usize) -> Option> { + let mut divisor = Vec::with_capacity(n); + let mut subdivided = false; + for w in parent_offset.windows(2) { + let m = w[1] - w[0]; + assert!( + m >= 1, + "parent_offset must be strictly increasing; parent owns {m} rows" + ); + subdivided |= m > 1; + divisor.extend(std::iter::repeat_n(m as f32, m as usize)); + } + assert_eq!( + divisor.len(), + n, + "parent_offset covers {} sub-reach rows but the network has {n}", + divisor.len() + ); + subdivided.then_some(divisor) +} + /// Static channel attributes and topology for a network. /// /// Adjacency, channel length, and slope come bundled inside `SparseAdjacency` @@ -93,12 +121,40 @@ pub struct MuskingumCunge { assembler: Option>, q_prime: Option, 2>>, + /// Per-row lateral-inflow divisor `[n]`: the piece count of each row's + /// parent reach, from `SparseAdjacency::parent_offset`. Built once at + /// `setup_inputs`; `None` when the network is not subdivided (every + /// divisor would be `1.0`), so the un-subdivided path skips the op and + /// stays byte-identical. + pieces_per_row: Option, 1>>, + /// Whether the cold-start solve `(I − N)·Q_0 = q'_0` sees the SAME divided + /// forcing `forward` routes. **Default `true`.** On an un-subdivided network + /// there is no divisor, so this is an exact no-op there; under subdivision + /// the undivided `q'_0` makes parent `p`'s outlet start at `m_p ×` its true + /// steady state. + /// + /// Measured (2026-08-05, 1,841 CONUS gauges / 184,676 sub-reach rows, + /// `max_pieces: 8`): the undivided cold start put 2.94× the correct total + /// discharge into the network, and the A/B difference took **221 hourly + /// steps to fall below 10 %** and 282 to fall below 5 % — against a + /// configured `warmup` of 5 days = 120 steps, at which point it was still + /// 41.7 %. The inflated state therefore leaks into the scored window, so it + /// is divided by default. `false` reproduces the un-divided behaviour for + /// `probe_courant --divide-hotstart` A/B runs. + pub divide_hotstart_by_pieces: bool, discharge_t: Option, 1>>, /// Eval-time zeta accumulation (leakance diagnostics). Off by default; /// `enable_zeta_accumulation` turns it on. Sums live on the inner backend /// (no autograd tape) and grow by one elementwise add per timestep. collect_zeta: bool, + + /// Per-timestep negative-discharge tracking. Off by default; + /// `enable_negative_discharge_tracking` turns it on. When off, the host + /// sync in `forward_chain_inner` is skipped entirely — zero added cost. + /// Enable in the training forward path; leave off in eval (the diagnostic + /// is only meaningful during training where we want to observe the rate). + track_negative_discharge: bool, zeta_abs_sum: Option>, zeta_net_sum: Option>, depth_sum: Option>, @@ -154,8 +210,11 @@ impl MuskingumCunge { pattern: None, assembler: None, q_prime: None, + pieces_per_row: None, + divide_hotstart_by_pieces: true, discharge_t: None, collect_zeta: false, + track_negative_discharge: false, zeta_abs_sum: None, zeta_net_sum: None, depth_sum: None, @@ -204,6 +263,16 @@ impl MuskingumCunge { ) .clamp_min(slope_min); + // Per-row lateral-inflow divisor, built once here rather than per + // timestep. A reach split into `m` pieces of length `L/m` gives each + // piece `q'/m` (see `forward`). + self.pieces_per_row = inputs + .adjacency + .parent_offset + .as_deref() + .and_then(|off| pieces_per_row_divisor(off, n)) + .map(|d| Tensor::, 1>::from_floats(d.as_slice(), &self.device)); + // Build CSR pattern + assembler constants directly from COO (O(nnz)). let pattern = Arc::new(CsrPattern::from_sparse(&inputs.adjacency)); self.assembler = Some(AValuesAssembler::::new(&pattern, &self.device)); @@ -262,6 +331,18 @@ impl MuskingumCunge { .clone() .slice([0..1, 0..n]) .reshape([n]); + // Give the cold start the same `q'/m` lateral inflow + // `forward` routes. Without it a subdivided network starts + // ~m× too wet at every parent outlet and drains that + // surplus for ~220 hourly steps — far past `warmup`. + // `None` divisor (un-subdivided) ⇒ exact no-op. + let q_prime_0 = match ( + self.divide_hotstart_by_pieces, + self.pieces_per_row.as_ref(), + ) { + (true, Some(d)) => q_prime_0 / d.clone(), + _ => q_prime_0, + }; // Hotstart: solve (I − N) · Q_0 = q'_0 via the same CSR solver // with c = 1 (all-ones vector), then clamp. let device = self.device.clone(); @@ -388,6 +469,7 @@ impl MuskingumCunge { k_d, d_gw, leakance_factor, self.impervious_mask.as_ref().cloned(), if self.collect_zeta { Some(&mut zeta_step) } else { None }, + self.track_negative_discharge, ); if let Some(diag) = zeta_step { fn add(slot: &mut Option>, v: Tensor) { @@ -428,6 +510,7 @@ impl MuskingumCunge { n, q_spatial, p_spatial, q_t, q_prime_clamp, length, slope, x_storage, + self.track_negative_discharge, ) } } @@ -441,6 +524,22 @@ impl MuskingumCunge { let discharge_lb = self.cfg.params.attribute_minimums.discharge; // Clamp once (single op + single tape node) instead of T times in-loop. let q_prime_clamped = q_prime.clamp_min(discharge_lb); + // Split lateral inflow evenly along a subdivided reach: a piece of + // length `L/m` receives `q'/m`. This is HEC-HMS's own treatment — its + // lateral term is `C4·(q_L·Δx)` with `q_L` an inflow per unit length — + // and it conserves each parent reach's total `q'` exactly, because the + // pieces chain in series so the outlet piece still carries the whole + // reach's runoff. + // + // MUST stay AFTER the clamp above. Clamping first applies the + // `discharge_lb` floor once, to the parent's inflow; dividing first + // would floor each of the `m` pieces independently, so a dry reach + // would inject `m · discharge_lb` instead of `discharge_lb` — mass + // created in proportion to the piece count. + let q_prime_clamped = match self.pieces_per_row.as_ref() { + Some(d) => q_prime_clamped / d.clone().unsqueeze_dim::<2>(0), + None => q_prime_clamped, + }; let initial = self .discharge_t .as_ref() @@ -479,6 +578,8 @@ impl MuskingumCunge { let mut columns: Vec, 2>> = Vec::with_capacity(num_timesteps); columns.push(initial.unsqueeze_dim::<2>(1)); + crate::routing::mmc_op::reset_negative_solve_stats(); + for t in 1..num_timesteps { let q_prime_t: Tensor, 1> = q_prime_clamped .clone() @@ -489,6 +590,45 @@ impl MuskingumCunge { self.discharge_t = Some(q_next); } + // Fix 1: distinguish three cases after the timestep loop. + // + // When `use_cuda_graphs` is true, `route_timestep` dispatches to + // `timestep_forward_via_graph`, which replays an on-device CUDA graph + // and never enters `forward_chain_inner`. Both counters stay at zero, + // which is *indistinguishable* from "measured, found zero negatives" — + // the silence would be a lie. Print an UNAVAILABLE notice so the + // output can never be misread as a zero-negative measurement. + // + // The fallback inside `timestep_forward_via_graph` (capture failed → + // direct launch) also carries this notice, because the user requested + // graphs and we cannot distinguish "all replays succeeded" from "some + // fell back silently". A config-based check is acceptable here because + // the leakance guard already rejects `use_cuda_graphs + leakance` at + // load time, so reaching this point with graphs on means the non- + // leakance graph path was requested. + let graphs_requested = self.cfg.params.use_cuda_graphs + && self.sparse_solver == SparseSolver::Cuda + && crate::sparse::dispatch::backend_is_cuda::(); + + if graphs_requested { + if self.track_negative_discharge { + eprintln!( + " negative solves: UNAVAILABLE — use_cuda_graphs is true; \ + the CUDA-graph path does not enter forward_chain_inner, so \ + no count was taken. Disable use_cuda_graphs to measure." + ); + } + } else { + let (neg, total) = crate::routing::mmc_op::negative_solve_stats(); + if total > 0 && neg > 0 { + eprintln!( + " negative solves before clamp: {neg}/{total} ({:.3}%) — Muskingum \ + coefficient sign violation (see .claude/PHYSICS-CORRECTIONS.md)", + 100.0 * neg as f64 / total as f64 + ); + } + } + Tensor::cat(columns, 1) } @@ -499,6 +639,18 @@ impl MuskingumCunge { self.collect_zeta = true; } + /// Turn on per-timestep negative-discharge tracking. When off (the + /// default), the host sync in `forward_chain_inner` is skipped entirely + /// — zero added cost on the forward. Enable in the training driver path + /// (per-micro-batch) so the negative-solve rate is visible. The CUDA-graph + /// path (`use_cuda_graphs: true`) never enters `forward_chain_inner`, so + /// enabling this flag there has no effect; `forward` will print an + /// UNAVAILABLE notice instead of a count. Training use only — the eval + /// path never enables this. + pub fn enable_negative_discharge_tracking(&mut self) { + self.track_negative_discharge = true; + } + /// Eval-time leakance diagnostic sums accumulated across `route_timestep` /// calls since construction. `None` until the first accumulated step. pub fn zeta_sums(&self) -> Option> { @@ -537,6 +689,27 @@ impl MuskingumCunge { self.pattern.as_ref() } + /// Inner-backend snapshot of every input `forward_chain_inner` consumes, + /// for the S18'/S19' Courant diagnostic in + /// [`crate::routing::courant_probe`]. Diagnostic only: strips the tape, + /// changes nothing. Call after `setup_inputs`. + pub fn probe_inputs(&self) -> crate::routing::courant_probe::ProbeInputs { + let n_seg = self.n_segments.expect("setup_inputs not called"); + crate::routing::courant_probe::ProbeInputs { + pattern: self.pattern.as_ref().expect("pattern").clone(), + n: self.n.as_ref().expect("n").clone().inner(), + q_spatial: self.q_spatial.as_ref().expect("q_spatial").clone().inner(), + p_spatial: self.p_spatial_broadcast(n_seg).inner(), + length: self.length.as_ref().expect("length").clone().inner(), + slope: self.slope.as_ref().expect("slope").clone().inner(), + x_storage: self.x_storage.as_ref().expect("x_storage").clone().inner(), + q_prime: self.q_prime.as_ref().expect("q_prime").clone().inner(), + pieces_per_row: self.pieces_per_row.as_ref().map(|d| d.clone().inner()), + q0: self.discharge_t.as_ref().expect("discharge_t").clone().inner(), + n_segments: n_seg, + } + } + fn p_spatial_broadcast(&self, n: usize) -> Tensor, 1> { let dims = self.p_spatial.dims(); if dims[0] == n { diff --git a/src/routing/mmc_op.rs b/src/routing/mmc_op.rs index 4ba9290..d35b3b1 100644 --- a/src/routing/mmc_op.rs +++ b/src/routing/mmc_op.rs @@ -8,6 +8,7 @@ //! Parents in fixed order: [n, q_spatial, p_spatial, q_t, q_prime_t]. use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use burn::backend::Autodiff; use burn::backend::autodiff::checkpoint::base::Checkpointer; @@ -19,6 +20,42 @@ use burn::tensor::{backend::Backend, Tensor, TensorPrimitive}; use crate::config::Config; use crate::sparse::{self, dispatch, primitive_to_vec, AValuesAssembler, CsrPattern}; +/// Count of solve outputs that came out NEGATIVE before the S28 +/// `clamp_min(discharge_lb)` rewrote them to `+1e-4`. +/// +/// Why this exists: Muskingum coefficients are only non-negative for +/// `2X <= Cr <= 2(1-X)` with `Cr = dt/K`. Measured on CONUS at mean flow with +/// `X = 0.3`, 69.8% of reaches sit outside that window (28.4% give `c1 < 0`, +/// 41.4% give `c3 < 0`), so negative discharge is expected — and the clamp +/// both CREATES MASS and removes the only symptom. Nothing measured this +/// before 2026-08-02. +/// +/// Diagnostic only: reads the solve to host, changes no numerics, and behaves +/// identically in both `ddr_match` modes. +static NEG_SOLVES: AtomicU64 = AtomicU64::new(0); +static TOTAL_SOLVES: AtomicU64 = AtomicU64::new(0); + +/// `(negative_count, total_count)` accumulated since the last reset. +pub fn negative_solve_stats() -> (u64, u64) { + (NEG_SOLVES.load(Ordering::Relaxed), TOTAL_SOLVES.load(Ordering::Relaxed)) +} + +/// Zero both counters. Call at the start of each `forward`. +pub fn reset_negative_solve_stats() { + NEG_SOLVES.store(0, Ordering::Relaxed); + TOTAL_SOLVES.store(0, Ordering::Relaxed); +} + +/// Safety margin pulling the S18'/S19' positivity clamp strictly INSIDE the +/// `2X <= Cr <= 2(1-X)` window (`params.enforce_positivity`). +/// +/// At `δ = 0` the clamp lands exactly on `c1 = 0` / `c3 = 0` and f32 roundoff +/// crosses it (measured min `c1` −3.3e−8, min `c3` −6.8e−8 over a 400k-draw +/// sweep). `δ = 1e-2` is ~400x f32 eps; it moves the measured minima to +/// +1.8e−4 / +5.0e−5 and costs only a 1% tightening of the X ceiling and a +/// 1% rise in the K floor. +pub const POSITIVITY_DELTA: f32 = 1e-2; + /// Inner-backend leakance inputs threaded into `forward_chain_inner`. #[derive(Clone)] pub(crate) struct LeakanceTensors { @@ -76,6 +113,16 @@ pub(crate) struct TimestepState { pub length: B::FloatTensorPrimitive, pub slope: B::FloatTensorPrimitive, pub x_storage: B::FloatTensorPrimitive, + /// The Muskingum X the forward ACTUALLY used at S19. Equals `x_storage` + /// when `ddr_match`; the Cunge-derived `clamp(0.5(1 − Q/(B·S·c·L)), 0, 0.5)` + /// otherwise. The backward reads it both to evaluate the c1..c4 chain and + /// to build the `[0, 0.5]` clamp mask (B19). + /// + /// Deliberately NOT part of the `forward_saved_idx` array: that array is + /// index-locked to `cuda_graph::PersistentScratch`'s `state_*` fields, and + /// the CUDA-graph path is `ddr_match: true`-only by config validation. It + /// travels through `forward_chain_inner`'s `x_eff_out` sink instead. + pub x_effective: B::FloatTensorPrimitive, // Forward intermediates (saved for backward). pub depth: B::FloatTensorPrimitive, pub top_width: B::FloatTensorPrimitive, @@ -107,6 +154,12 @@ pub(crate) struct TimestepState { pub discharge_lb: f32, pub dt: f32, pub use_cuda: bool, + pub ddr_match: bool, + /// Mirrors `forward_chain_inner`'s `enforce_pos` — i.e. the ALREADY-COMBINED + /// `!ddr_match && cfg.params.enforce_positivity`, not the raw config flag. + /// Tells the backward whether S18' (K floor) and S19' (three-way X min) ran, + /// so B18'/B19' apply their masks. `false` ⇒ the pre-clamp math, unchanged. + pub enforce_pos: bool, } #[derive(Debug)] @@ -121,6 +174,27 @@ pub(crate) struct ZetaGeomGrads { pub g_p_spatial: Tensor, } +/// The three extra chain-rule terms the Cunge-derived Muskingum `X` +/// (`ddr_match: false`, S19) contributes. Each is ADDED to the accumulator +/// that already carries that quantity's gradient — see B19 for the derivation +/// and for the ordering constraint each one imposes. +struct XGrads { + g_q_t: Tensor, + g_top_width: Tensor, + g_celerity: Tensor, +} + +/// The four extra chain-rule terms the trapezoidal celerity +/// (`ddr_match: false`, S17) contributes. Each is ADDED to the accumulator +/// that already carries that quantity's gradient through the `hyd_radius` +/// chain — see B17 for the derivation. +struct BetaGrads { + g_area: Tensor, + g_top_width: Tensor, + g_wp: Tensor, + g_side_slope: Tensor, +} + /// The five accumulated parent gradients produced by [`timestep_backward_core`], /// in parent order `[n, q_spatial, p_spatial, q_t, q_prime_t]`. pub(crate) struct FiveGrads { @@ -174,7 +248,10 @@ where let q_prime_t = wrap(state.q_prime_t.clone()); let length = wrap(state.length.clone()); let slope = wrap(state.slope.clone()); - let x_storage = wrap(state.x_storage.clone()); + // The X the forward ACTUALLY used at S19 — equal to `state.x_storage` + // under `ddr_match`, Cunge-derived otherwise. `state.x_storage` is + // deliberately never read here: it is not what c1..c4 were built from. + let x_eff = wrap(state.x_effective.clone()); let depth = wrap(state.depth.clone()); let top_width = wrap(state.top_width.clone()); @@ -200,6 +277,8 @@ where let bw_raw = wrap(state.bw_raw.clone()); let dt = state.dt; + // Already the combined `!ddr_match && params.enforce_positivity`. + let enforce_pos = state.enforce_pos; let bottom_width_lb = state.bottom_width_lb; let depth_lb = state.depth_lb; let velocity_lb = state.velocity_lb; @@ -298,9 +377,9 @@ where // and ∂L/∂num_i_from_ci = gci / denom. // =========================================================== let denom_sq = denom.clone() * denom.clone(); - let one_minus_x = -x_storage.clone() + 1.0; + let one_minus_x = -x_eff.clone() + 1.0; let two_k = k_muskingum.clone() * 2.0; - let two_kx = two_k.clone() * x_storage.clone(); + let two_kx = two_k.clone() * x_eff.clone(); let two_k_1mx = two_k.clone() * one_minus_x.clone(); // num_c1 = -2kx + dt @@ -331,25 +410,215 @@ where let g_2kx_total = g_2kx_from_c1 + g_2kx_from_c2; let g_2k1mx_total = g_2k1mx_from_c3 + g_2k1mx_from_denom; - // 2kx = 2k · x_storage → ∂(2kx)/∂(2k) = x_storage - let g_2k_from_2kx = g_2kx_total.clone() * x_storage.clone(); - // 2k(1-x) = 2k · (1 - x_storage) → ∂(2k(1-x))/∂(2k) = (1 - x_storage) + // =========================================================== + // B19. Muskingum X (`ddr_match: false` only). + // + // X feeds ONLY `two_kx = 2k·X` and `two_k_1mx = 2k·(1−X)`, so + // gX = 2k · (g_2kx_total − g_2k1mx_total) + // + // With W ≡ Q/(B·S·c·L) and X_raw = 0.5·(1 − W): + // ∂X/∂Q = −0.5·W/Q ∂X/∂B = +0.5·W/B ∂X/∂c = +0.5·W/c + // and 0 wherever the [0, 0.5] clamp saturated. S and L are constants, + // not parents, so they get no term. + // + // The three contributions are ADDED into accumulators that already + // carry gradient for those quantities: `gcelerity` (B18), `gtw_total` + // (before B7), and `gq_t_total` (final). Each MUST land before its + // accumulator is consumed — see the comment at each site. + // + // NOTE: this is a SECOND gradient path for `q_t`, which already + // reaches the loss through the S25 RHS and the S2 depth chain. + // + // --------------------------------------------------------------- + // B19'. Positivity clamp (`enforce_positivity`, S19'): + // + // x_eff = min(x_cunge, hi_a, hi_b) + // hi_a = cr·0.5·(1−δ) hi_b = (1 − 0.5·cr)·(1−δ) + // cr = Δt / k_musk + // + // A `min` routes the incoming gradient to exactly ONE branch per + // element, so the three branch masks must PARTITION the reaches — no + // element counted twice, none dropped. Tie-break priority (matching + // `branch_report` in `tests/positivity_clamp.rs`): + // + // Cunge > hi_a > hi_b + // mask_cunge = (x_cunge <= hi_a) && (x_cunge <= hi_b) + // mask_a = !mask_cunge && (hi_a <= hi_b) + // mask_b = !mask_cunge && !mask_a + // + // Ties are a measure-zero subgradient choice; the ordering just has to + // be deterministic and total, which the `<=` cascade guarantees. + // + // Consequences: + // * the existing `∂X/∂Q`, `∂X/∂B`, `∂X/∂c` terms are the Cunge + // branch's, so they are additionally masked by `mask_cunge`; + // * `hi_a`/`hi_b` open a NEW path `x_eff → cr → k_musk → celerity` + // that did not exist before the clamp. Its contribution is + // accumulated into `gk_from_x_cap` and folded into `gk_muskingum` + // BELOW (before the S18' floor mask and before B18 consumes it). + // * the `[0, 0.5]` clamp mask belongs to `x_cunge`, NOT to `x_eff`: + // when `hi_a`/`hi_b` win, `x_eff` can sit anywhere in `[0, 0.5]` + // while `x_cunge` is saturated, and vice versa. + // =========================================================== + // + // `hi_a`/`hi_b` → `k_muskingum`. `None` unless `enforce_pos`. + let mut gk_from_x_cap: Option> = None; + let x_grads = if state.ddr_match { + None + } else { + let gx = two_k.clone() * (g_2kx_total.clone() - g_2k1mx_total.clone()); + // Same expression as the forward's S19 (including the +1e-12). + let bscl = + top_width.clone() * slope.clone() * celerity.clone() * length.clone() + 1e-12; + let w = q_t.clone() / bscl.clone(); + let half_w = w.clone() * 0.5; + // `gx` is ∂L/∂x_eff. Split it across the S19' min branches and pick + // the tensor whose `[0, 0.5]` clamp mask applies to the Cunge term. + // Without the clamp, x_eff IS x_cunge and both reduce to identity. + let (gx, x_for_clamp) = if enforce_pos { + // Recomputed bit-identically to the forward (same ops, same + // saved operands), so the comparisons below reproduce the + // `min_pair` chain's choice exactly. + let x_cunge = ((-w + 1.0) * 0.5).clamp(0.0, 0.5); + let cr = k_muskingum.clone().recip() * dt; + let hi_a = cr.clone() * (0.5 * (1.0 - POSITIVITY_DELTA)); + let hi_b = (-cr * 0.5 + 1.0) * (1.0 - POSITIVITY_DELTA); + + let mask_cunge = x_cunge + .clone() + .lower_equal(hi_a.clone()) + .bool_and(x_cunge.clone().lower_equal(hi_b.clone())); + let mask_a = mask_cunge + .clone() + .bool_not() + .bool_and(hi_a.lower_equal(hi_b)); + let mask_b = mask_cunge + .clone() + .bool_not() + .bool_and(mask_a.clone().bool_not()); + + let g_hi_a = gx.clone().mask_fill(mask_a.bool_not(), 0.0); + let g_hi_b = gx.clone().mask_fill(mask_b.bool_not(), 0.0); + let gx_cunge = gx.mask_fill(mask_cunge.bool_not(), 0.0); + + // ∂hi_a/∂cr = +0.5(1−δ), ∂hi_b/∂cr = −0.5(1−δ) + let half_1md = 0.5 * (1.0 - POSITIVITY_DELTA); + let gcr = (g_hi_a - g_hi_b) * half_1md; + // cr = Δt/k_musk → ∂cr/∂k_musk = −Δt/k_musk² + gk_from_x_cap = Some(-gcr * dt / (k_muskingum.clone() * k_muskingum.clone())); + (gx_cunge, x_cunge) + } else { + (gx, x_eff.clone()) + }; + // Zero where the [0, 0.5] clamp saturated (X is then constant). + let unsat = x_for_clamp + .clone() + .greater_elem(0.0) + .bool_and(x_for_clamp.lower_elem(0.5)); + let gx = gx.mask_fill(unsat.bool_not(), 0.0); + Some(XGrads { + // ∂X/∂Q = −0.5·W/Q, but W = Q/(B·S·c·L) so Q cancels exactly. + // Written in the cancelled form on purpose: `q_t` is only + // floored at `discharge_lb` AFTER the solve, so the raw + // per-step q_t can legitimately be 0 and `W/Q` would be 0/0. + g_q_t: -gx.clone() * 0.5 / bscl, + g_top_width: gx.clone() * half_w.clone() / top_width.clone(), + g_celerity: gx * half_w / celerity.clone(), + }) + }; + + // 2kx = 2k · x_eff → ∂(2kx)/∂(2k) = x_eff + let g_2k_from_2kx = g_2kx_total.clone() * x_eff.clone(); + // 2k(1-x) = 2k · (1 - x_eff) → ∂(2k(1-x))/∂(2k) = (1 - x_eff) let g_2k_from_2k1mx = g_2k1mx_total.clone() * one_minus_x.clone(); let g_2k_total = g_2k_from_2kx + g_2k_from_2k1mx; // 2k = 2 · k_muskingum let gk_muskingum = g_2k_total * 2.0; + // ORDERING: B19's `x_eff → cr → k_musk` term (enforce_positivity only) + // MUST join here, i.e. AFTER the c1..c4 contribution and BEFORE the + // S18' floor mask below — both paths reach `celerity` through the SAME + // `clamp_min`, so a term added after the mask would leak gradient on + // reaches where the floor binds. + let gk_muskingum = match gk_from_x_cap { + Some(g) => gk_muskingum + g, + None => gk_muskingum, + }; + + // =========================================================== + // Geometry recomputed from the saved primitives. Hoisted above B18 + // because the `ddr_match: false` celerity (B17) needs A, T, P and + // sqrt(1+z²) too; B14/B13 below reuse the exact same tensors. + // wp = bw + 2·d·sqrt(1+ss²) (S13) + // area = R · wp (S14 inverted; area itself unsaved) + // =========================================================== + let one_plus_ss_sq = side_slope.clone() * side_slope.clone() + 1.0; + let sqrt_1_plus_ss_sq = one_plus_ss_sq.clone().sqrt(); + let wp = _bottom_width.clone() + depth.clone() * sqrt_1_plus_ss_sq.clone() * 2.0; + let area = hyd_radius.clone() * wp.clone(); + + // =========================================================== + // B18'. K floor (`enforce_positivity` only, S18'): + // k_musk = max(k_raw, k_floor) → ∂k_musk/∂k_raw = 1 where + // k_raw > k_floor, else 0 (the reach is sub-grid and K is pinned to + // the constant floor, so celerity gets NO gradient there). + // `k_raw = length/celerity` is recomputed from the saved `length` and + // `celerity` with the same op the forward used, so the mask reproduces + // the forward's `clamp_min` decision exactly. + // =========================================================== + let gk_raw = if enforce_pos { + let k_raw = length.clone() / celerity.clone(); + let k_floor = dt * (1.0 + POSITIVITY_DELTA) / 2.0; + let unfloored = k_raw.greater_elem(k_floor); + gk_muskingum.mask_fill(unfloored.bool_not(), 0.0) + } else { + gk_muskingum + }; // =========================================================== // B18. k_muskingum = length / celerity // ∂k/∂celerity = -length / celerity² // =========================================================== + // ORDERING: B19's ∂X/∂c term must join here, BEFORE B17 consumes + // `gcelerity` — celerity reaches the loss both through K (this line) + // and, under `ddr_match: false`, through X. let celerity_sq = celerity.clone() * celerity.clone(); - let gcelerity = -gk_muskingum * length.clone() / celerity_sq; + let gcelerity = -gk_raw * length.clone() / celerity_sq; + let gcelerity = match x_grads.as_ref() { + Some(xg) => gcelerity + xg.g_celerity.clone(), + None => gcelerity, + }; // =========================================================== - // B17. celerity = velocity_cl · 5/3 + // B17. celerity = velocity_cl · beta + // + // ddr_match=true : beta ≡ 5/3 (constant) → gvelocity_cl = gc · 5/3. + // + // ddr_match=false: beta = 5/3 − G, G ≡ (4/3)·A·u/(T·P), u = √(1+z²). + // ∂celerity/∂velocity_cl = beta (NOT the constant 5/3) + // gbeta = gcelerity · velocity_cl + // ∂beta/∂A = −G/A ∂beta/∂T = +G/T + // ∂beta/∂P = +G/P ∂beta/∂z = −G·z/(1+z²) + // The four terms are ADDITIONAL contributions folded into the + // existing gA (B12), gT (gtw_total), gP (B13) and gz (B9) + // accumulators — A, T, P and z all already carry gradient through + // the hyd_radius chain. // =========================================================== - let gvelocity_cl = gcelerity * (5.0 / 3.0); + let (gvelocity_cl, beta_grads) = if state.ddr_match { + (gcelerity * (5.0 / 3.0), None) + } else { + let g_term = area.clone() * sqrt_1_plus_ss_sq.clone() + / (top_width.clone() * wp.clone()) + * (4.0 / 3.0); + let beta = -g_term.clone() + (5.0 / 3.0); + let gbeta = gcelerity.clone() * _velocity_cl.clone(); + let grads = BetaGrads { + g_area: -gbeta.clone() * g_term.clone() / area.clone(), + g_top_width: gbeta.clone() * g_term.clone() / top_width.clone(), + g_wp: gbeta.clone() * g_term.clone() / wp.clone(), + g_side_slope: -gbeta * g_term * side_slope.clone() / one_plus_ss_sq.clone(), + }; + (gcelerity * beta, Some(grads)) + }; // =========================================================== // B16. velocity_cl = clamp(velocity_un, velocity_lb, 15) @@ -375,17 +644,9 @@ where // B14. R = area / wp // ∂R/∂area = 1 / wp // ∂R/∂wp = -area / wp² = -R/wp - // Need wp and area. We saved wp implicitly via area = R·wp. Re-derive wp: + // `wp` and `area` are recomputed above B18 (area is not saved; it is + // recovered exactly as R·wp). // =========================================================== - // Recompute wp from saved bottom_width + 2·depth·sqrt(1+ss²) — equivalently - // wp = area / hyd_radius (cheap and exact). - // Use area from saved? We didn't save area, but area = R · wp. We need wp directly. - // Recompute: wp = bottom_width + 2·depth·sqrt(1 + side_slope²). - let one_plus_ss_sq = side_slope.clone() * side_slope.clone() + 1.0; - let sqrt_1_plus_ss_sq = one_plus_ss_sq.clone().sqrt(); - let wp = _bottom_width.clone() + depth.clone() * sqrt_1_plus_ss_sq.clone() * 2.0; - let area = hyd_radius.clone() * wp.clone(); - let gr = gr_from_s15; let garea_from_r = gr.clone() / wp.clone(); let gwp_from_r = -gr * area.clone() / (wp.clone() * wp.clone()); @@ -396,7 +657,12 @@ where // ∂wp/∂d = 2·sqrt(1+ss²) // ∂wp/∂ss = 2·d · ss / sqrt(1+ss²) // =========================================================== - let gwp = gwp_from_r; + // `beta` (ddr_match=false) reads P directly at S17, so its ∂beta/∂P + // term joins gwp BEFORE the S13 decomposition. + let gwp = match beta_grads.as_ref() { + Some(bg) => gwp_from_r + bg.g_wp.clone(), + None => gwp_from_r, + }; let gbw_from_s13 = gwp.clone(); let gd_from_s13 = gwp.clone() * sqrt_1_plus_ss_sq.clone() * 2.0; let gss_from_s13 = gwp * depth.clone() * 2.0 * side_slope.clone() / sqrt_1_plus_ss_sq; @@ -407,7 +673,11 @@ where // ∂area/∂bw = d/2 // ∂area/∂d = (tw + bw)/2 // =========================================================== - let garea = garea_from_r; + // Likewise ∂beta/∂A joins garea BEFORE the S12 decomposition. + let garea = match beta_grads.as_ref() { + Some(bg) => garea_from_r + bg.g_area.clone(), + None => garea_from_r, + }; let half_d = depth.clone() * 0.5; let gtw_from_s12 = garea.clone() * half_d.clone(); let gbw_from_s12 = garea.clone() * half_d.clone(); @@ -435,7 +705,11 @@ where // B9. side_slope = clamp(side_slope_raw, 0.5, 50) // gradient passes where 0.5 < ss_raw < 50. // =========================================================== - let gss_combined = gss_from_s13 + gss_from_s10; + let gss_combined = match beta_grads.as_ref() { + // ∂beta/∂z enters through u = √(1+z²) at S17. + Some(bg) => gss_from_s13 + gss_from_s10 + bg.g_side_slope.clone(), + None => gss_from_s13 + gss_from_s10, + }; let mask_ss_lo = side_slope_raw.clone().greater_elem(0.5); let mask_ss_hi = side_slope_raw.clone().lower_elem(50.0); let mask_ss = mask_ss_lo.bool_and(mask_ss_hi); @@ -456,7 +730,16 @@ where / (two_d.clone() * depth.clone()); // Accumulate gtw before S7 (since S7 produces depth → uses tw in its derivative wrt q_eps). - let gtw_total = gtw_from_s12 + gtw_from_s10 + gtw_from_s8; + // ∂beta/∂T (S17) is the fourth contributor when ddr_match=false, and + // ∂X/∂B (S19, B19) is the fifth. Both MUST land before B7 consumes it. + let gtw_total = match beta_grads.as_ref() { + Some(bg) => gtw_from_s12 + gtw_from_s10 + gtw_from_s8 + bg.g_top_width.clone(), + None => gtw_from_s12 + gtw_from_s10 + gtw_from_s8, + }; + let gtw_total = match x_grads.as_ref() { + Some(xg) => gtw_total + xg.g_top_width.clone(), + None => gtw_total, + }; // =========================================================== // B7. top_width = p · depth^q_eps @@ -542,7 +825,14 @@ where if let Some(zg) = zeta_geom.as_ref() { gp_total = gp_total + zg.g_p_spatial.clone(); } + // q_t reaches the loss through the S25 RHS (c3·q_t), the S24 SpMV + // (N·q_t), the S2 depth chain — and, under `ddr_match: false`, a + // FOURTH path: the Cunge X at S19 (B19). let gq_t_total = gq_t_from_s25 + gq_t_from_s24 + gq_t_from_s2; + let gq_t_total = match x_grads.as_ref() { + Some(xg) => gq_t_total + xg.g_q_t.clone(), + None => gq_t_total, + }; // Touch unused intermediate bindings to silence dead-code warnings. let _ = (_q_spatial, _velocity_cl); @@ -781,6 +1071,10 @@ pub(crate) fn forward_chain_inner( xst_in: Tensor, leakance: Option>, leak_out: &mut Option>, + // Receives the S19 Muskingum X the chain actually used (see + // `TimestepState::x_effective`). Always written. + x_eff_out: &mut Option, + track_neg: bool, ) -> ( I::FloatTensorPrimitive, [I::FloatTensorPrimitive; NUM_SAVED_STATE], @@ -797,6 +1091,11 @@ where let velocity_lb = cfg.params.attribute_minimums.velocity; let discharge_lb = cfg.params.attribute_minimums.discharge; let use_cuda = cfg.params.sparse_solver == SparseSolver::Cuda; + let ddr_match = cfg.params.ddr_match; + // S18'/S19' positivity clamp. `ddr_match: true` must stay byte-identical to + // DDR (invariant 1), so the clamp is gated on BOTH flags even though + // `validate_enforce_positivity` already rejects the combination at load. + let enforce_pos = !ddr_match && cfg.params.enforce_positivity; let unwrap = |t: Tensor| -> I::FloatTensorPrimitive { match t.into_primitive() { @@ -845,14 +1144,90 @@ where * slope_in.clone().sqrt(); // S16 let velocity_cl = velocity_un.clone().clamp(velocity_lb, 15.0); - // S17 - let celerity = velocity_cl.clone() * (5.0_f32 / 3.0_f32); + // S17: celerity. + // ddr_match=true -> c = v·5/3, the wide-rectangular Kleitz-Seddon limit. + // Mirrors ddr/mmc.py:167. WRONG for the trapezoid built + // in S7-S13: kappa = b/y ~ 0.7-1.8 here, so the true + // ratio is ~1.30-1.36, not 1.667 (22-27% high). + // ddr_match=false -> exact trapezoidal c = dQ/dA = (dQ/dy)/T: + // beta = 5/3 - (4/3)·A·sqrt(1+z²)/(T·P) + // (-> 5/3 as b/y -> inf, -> 4/3 as b -> 0). + let celerity = if ddr_match { + velocity_cl.clone() * (5.0_f32 / 3.0_f32) + } else { + let root = (side_slope.clone().powf_scalar(2.0) + 1.0).sqrt(); + let beta = + -(_area.clone() * root) / (top_width.clone() * wp.clone()) * (4.0 / 3.0) + (5.0 / 3.0); + velocity_cl.clone() * beta + }; // S18..S23: Muskingum coefficients - let k_muskingum = length_in.clone() / celerity.clone(); - let one_minus_x = -xst_in.clone() + 1.0; + // + // S18': K floor (`enforce_positivity` only). A reach with K < dt/2 is + // SUB-GRID: the timestep cannot resolve its transit, and the unclamped + // scheme expressed that as oscillation which S28's `clamp_min(1e-4)` + // silently rewrote to +1e-4 (creating mass). Flooring K makes the + // coarse-graining explicit and puts Cr = dt/K inside (0, 2/(1+δ)], which + // is the feasible region for the S19' X cap below (it is what guarantees + // `hi_b > 0`, so no `clamp_min` is needed on the three-way min). + let k_raw = length_in.clone() / celerity.clone(); + let k_muskingum = if enforce_pos { + k_raw.clamp_min(dt * (1.0 + POSITIVITY_DELTA) / 2.0) + } else { + k_raw + }; + + // S19: Muskingum storage weight X. + // ddr_match=true -> the caller's constant (forward.rs sets 0.3). NOT + // Cunge-derived: severs the link between the scheme's + // numerical diffusion and the channel's physical + // hydraulic diffusivity, giving a median 28x + // over-diffusion on CONUS. It is exact only on the + // measure-zero locus Q/(B·S·c·L) == 0.4. + // ddr_match=false -> Cunge: X = clamp(0.5(1 - Q/(B·S·c·L)), 0, 0.5), + // which makes D_num = c·L·(0.5-X) equal D_phys = + // Q/(2·B·S). Measured X ~ 0.49 on CONUS, which also + // narrows the non-negative-coefficient window + // 2X <= Cr <= 2(1-X) from [0.6, 1.4] to ~[0.98, 1.02] + // (see Task 5, Courant sub-stepping). + // NOTE for the backward: X feeds ONLY `two_kx` and `two_k_1mx`, and + // depends on q_t, top_width and celerity. See B19 in + // `timestep_backward_core`. + let x_eff = if ddr_match { + xst_in.clone() + } else { + let w = qt_in.clone() + / (top_width.clone() * slope_in.clone() * celerity.clone() * length_in.clone() + + 1e-12); + let x_cunge = ((-w + 1.0) * 0.5).clamp(0.0, 0.5); + if enforce_pos { + // S19': stability cap. Muskingum's non-negative-coefficient window + // is 2X <= Cr <= 2(1-X); solving both sides for X gives + // X <= 0.5·Cr (<=> c1 = (dt - 2KX)/denom >= 0) + // X <= 1 - 0.5·Cr (<=> c3 = (2K(1-X) - dt)/denom >= 0) + // and c2, c4 > 0 unconditionally. With b >= 0 and the forward + // substitution x[i] = b[i] + c1[i]·Σ_up x[j] taken in topological + // order, c1, c3 >= 0 makes every x[i] >= 0 by induction — no + // negative solve can reach S28. + // + // The (1-δ) factor is MANDATORY, not cosmetic: at δ = 0 the cap + // lands exactly on c1 = 0 / c3 = 0 and f32 roundoff crosses it + // (measured min c1 = -3.3e-8, min c3 = -6.8e-8 over a 400k sweep). + // δ = 1e-2 (~400x f32 eps) puts the measured minima at +1.8e-4 and + // +5.0e-5 and costs a 1% tightening of the X ceiling. + let cr = k_muskingum.clone().recip() * dt; + let hi_a = cr.clone() * (0.5 * (1.0 - POSITIVITY_DELTA)); + let hi_b = (-cr * 0.5 + 1.0) * (1.0 - POSITIVITY_DELTA); + x_cunge.min_pair(hi_a).min_pair(hi_b) + } else { + x_cunge + } + }; + *x_eff_out = Some(unwrap(x_eff.clone())); + + let one_minus_x = -x_eff.clone() + 1.0; let two_k = k_muskingum.clone() * 2.0; - let two_kx = two_k.clone() * xst_in.clone(); + let two_kx = two_k.clone() * x_eff.clone(); let two_k_1mx = two_k.clone() * one_minus_x.clone(); let denom = two_k_1mx.clone() + dt; let c1 = (-two_kx.clone() + dt) / denom.clone(); @@ -915,6 +1290,17 @@ where ); let x_sol = wrap(x_sol_prim.clone()); + // Diagnostic: count negative solve outputs before the S28 clamp rewrites + // them to +1e-4. Gated on `track_neg` so the default (off) path pays zero + // cost — no host sync, no blocking device→host transfer. Mirror of the + // `collect_zeta` / `enable_zeta_accumulation` pattern on MuskingumCunge. + if track_neg { + let x_host: Vec = primitive_to_vec::(x_sol_prim.clone()); + let n_neg = x_host.iter().filter(|&&v| v < 0.0).count() as u64; + NEG_SOLVES.fetch_add(n_neg, Ordering::Relaxed); + TOTAL_SOLVES.fetch_add(x_host.len() as u64, Ordering::Relaxed); + } + // S28: q_next = max(x_sol, discharge_lb) let q_next = x_sol.clone().clamp_min(discharge_lb); let q_next_prim = unwrap(q_next); @@ -966,6 +1352,11 @@ where /// the named outputs. Keep the kernel order identical to /// [`forward_chain_inner`] so V9 bit-match still holds. If you change one, /// change both. +/// +/// This function is intentionally DDR-only (no `ddr_match` branch) until it +/// is revived; do not silently diverge from [`forward_chain_inner`] by adding +/// physics-correction branches here without also threading `ddr_match` and +/// updating the CUDA-graph capture path. #[allow(clippy::too_many_arguments, dead_code)] pub(crate) fn forward_chain_inner_pinned( cfg: &Config, @@ -1260,6 +1651,9 @@ where let x_sol = wrap(x_sol_prim.clone()); // S28: q_next = max(x_sol, discharge_lb) + // NOTE: negative-solve counter is NOT instrumented here — this function is + // #[allow(dead_code)] with no live callers. Only forward_chain_inner is + // instrumented. let q_next = x_sol.clone().clamp_min(discharge_lb); pin_clone(&q_next, pin); let q_next_prim = unwrap(q_next); @@ -1319,6 +1713,7 @@ pub fn timestep_forward( length_at: Tensor, 1>, slope_at: Tensor, 1>, x_storage_at: Tensor, 1>, + track_neg: bool, ) -> Tensor, 1> where I::FloatTensorPrimitive: 'static, @@ -1332,6 +1727,9 @@ where let velocity_lb = cfg.params.attribute_minimums.velocity; let discharge_lb = cfg.params.attribute_minimums.discharge; let use_cuda = cfg.params.sparse_solver == SparseSolver::Cuda; + let ddr_match = cfg.params.ddr_match; + // Must mirror `forward_chain_inner`'s gate EXACTLY (S18'/S19' vs B18'/B19'). + let enforce_pos = !ddr_match && cfg.params.enforce_positivity; // Extract AutodiffTensor (carries `primitive` + `node`). let unwrap_at = |t: Tensor, 1>| match t.into_primitive() { @@ -1361,6 +1759,7 @@ where Tensor::from_primitive(TensorPrimitive::Float(p)) }; + let mut x_eff_out: Option = None; let (q_next_prim, saved) = forward_chain_inner::( cfg, pattern, @@ -1374,7 +1773,10 @@ where wrap(xst_p.clone()), None, &mut None, + &mut x_eff_out, + track_neg, ); + let x_effective = x_eff_out.expect("forward_chain_inner always writes x_eff_out"); // Unpack saved-state array into named TimestepState fields. Indices MUST // match `forward_saved_idx`. @@ -1401,6 +1803,7 @@ where length: length_p, slope: slope_p, x_storage: xst_p, + x_effective, depth: depth_p, top_width: top_width_p, side_slope: side_slope_p, @@ -1430,6 +1833,8 @@ where discharge_lb, dt, use_cuda, + ddr_match, + enforce_pos, }; // Register the op on the autograd tape. @@ -1484,6 +1889,7 @@ pub fn timestep_forward_leakance( leakance_factor_at: Tensor, 1>, impervious_mask: Option>, zeta_out: Option<&mut Option>>, + track_neg: bool, ) -> Tensor, 1> where I::FloatTensorPrimitive: 'static, @@ -1497,6 +1903,9 @@ where let velocity_lb = cfg.params.attribute_minimums.velocity; let discharge_lb = cfg.params.attribute_minimums.discharge; let use_cuda = cfg.params.sparse_solver == SparseSolver::Cuda; + let ddr_match = cfg.params.ddr_match; + // Must mirror `forward_chain_inner`'s gate EXACTLY (S18'/S19' vs B18'/B19'). + let enforce_pos = !ddr_match && cfg.params.enforce_positivity; let unwrap_at = |t: Tensor, 1>| match t.into_primitive() { TensorPrimitive::Float(p) => p, @@ -1537,6 +1946,7 @@ where mask: impervious_mask, }; let mut leak_out: Option> = None; + let mut x_eff_out: Option = None; let (q_next_prim, saved) = forward_chain_inner::( cfg, @@ -1551,8 +1961,11 @@ where wrap(xst_p.clone()), Some(leakance), &mut leak_out, + &mut x_eff_out, + track_neg, ); let leak = leak_out.expect("forward_chain_inner must populate LeakanceSaved when leakance is Some"); + let x_effective = x_eff_out.expect("forward_chain_inner always writes x_eff_out"); use forward_saved_idx as fsi; let [ @@ -1597,6 +2010,7 @@ where length: length_p, slope: slope_p, x_storage: xst_p, + x_effective, depth: depth_p, top_width: top_width_p, side_slope: side_slope_p, @@ -1626,6 +2040,8 @@ where discharge_lb, dt, use_cuda, + ddr_match, + enforce_pos, }; let state = TimestepLeakanceState:: { base, leak }; @@ -1702,11 +2118,15 @@ where unsafe { crate::sparse::cusparse::ensure_cuda_cache::(pattern, &cache_device) }; // If no graph was installed (capture failed), fall through to direct launch. + // Pass track_neg=false: when graphs were requested, `forward` will print + // the UNAVAILABLE notice covering both the graph-replay and capture-failure + // cases, so there is no point incurring the host-sync cost here. if cache.graph_fwd.is_none() || cache.scratch.is_none() { return timestep_forward::( cfg, pattern, assembler, n_at, q_spatial_at, p_spatial_at, q_t_at, q_prime_t_at, length_at, slope_at, x_storage_at, + false, ); } @@ -1725,6 +2145,7 @@ where let velocity_lb = cfg.params.attribute_minimums.velocity; let discharge_lb = cfg.params.attribute_minimums.discharge; let use_cuda = cfg.params.sparse_solver == SparseSolver::Cuda; + let ddr_match = cfg.params.ddr_match; // Unwrap autograd primitives. let unwrap_at = |t: Tensor, 1>| match t.into_primitive() { @@ -1867,6 +2288,14 @@ where ] = state_arr; let _ = (fsi::DEPTH, fsi::BW_RAW); // index sanity touch + // The captured graph encodes the DDR chain only. Config load rejects + // `ddr_match: false` + `use_cuda_graphs: true` precisely so this holds. + debug_assert!( + ddr_match, + "CUDA-graph replay cannot serve ddr_match=false (captured graph has \ + neither the trapezoidal celerity nor the Cunge X branch)" + ); + // Build TimestepState. Backward needs the per-step `q_t`/`q_prime_t` // primitives (NOT scratch handles — scratch gets overwritten on next // replay, but the backward closure runs on `loss.backward()` after the @@ -1880,7 +2309,12 @@ where q_prime_t: qpt_p, length: length_p, slope: slope_p, - x_storage: xst_p, + // The captured graph is the `ddr_match: true` chain (S19 returns the + // caller's constant unchanged), and config load rejects + // `ddr_match: false` together with `use_cuda_graphs: true`, so + // `x_effective == x_storage` here by construction. + x_storage: xst_p.clone(), + x_effective: xst_p, depth: depth_p, top_width: top_width_p, side_slope: side_slope_p, @@ -1910,6 +2344,11 @@ where discharge_lb, dt, use_cuda, + ddr_match, + // The captured graph is the `ddr_match: true` chain (asserted above), + // and `enforce_positivity` requires `ddr_match: false` at config load, + // so S18'/S19' can never have run on this path. + enforce_pos: false, }; let result_prim = match TimestepOp @@ -1957,7 +2396,7 @@ where { let (_q_next, saved) = forward_chain_inner::( cfg, pattern, n_in, qsp_in, psp_in, qt_in, qpt_in, length_in, slope_in, xst_in, None, - &mut None, + &mut None, &mut None, false, ); // Indices K1 produces (skip 14..=17: A_VALUES, B_RHS, I_T, X_SOL). @@ -2020,7 +2459,7 @@ where { let (q_next_prim, saved) = forward_chain_inner::( cfg, pattern, n_in, qsp_in, psp_in, qt_in, qpt_in, length_in, slope_in, xst_in, None, - &mut None, + &mut None, &mut None, false, ); let to_vec = |prim: I::FloatTensorPrimitive| -> Vec { diff --git a/src/routing/mod.rs b/src/routing/mod.rs index ce58d71..7382c05 100644 --- a/src/routing/mod.rs +++ b/src/routing/mod.rs @@ -1,3 +1,4 @@ +pub mod courant_probe; pub mod leakance; pub mod mmc; pub mod utils; diff --git a/src/sparse/mod.rs b/src/sparse/mod.rs index 9537847..14fd37a 100644 --- a/src/sparse/mod.rs +++ b/src/sparse/mod.rs @@ -56,6 +56,16 @@ pub struct SparseAdjacency { /// Channel slope per reach (dimensionless), length `n`, aligned to /// topological order. Engine clamps to `attribute_minimums.slope`. pub slope: Vec, + /// Reach-subdivision parent map, length `n_parent + 1`: rows + /// `[parent_offset[p], parent_offset[p + 1])` are the sub-reach pieces of + /// parent reach `p`, contiguous and ordered upstream→downstream (so the + /// parent's outlet is `parent_offset[p + 1] - 1`). Mirrors the + /// `/parent_offset` array on `ConusAdjacencyStore`. + /// + /// `None` — and equally an identity map `0..=n` — means the network is not + /// subdivided; the engine then skips the lateral-inflow split entirely. + /// The last element must equal `n`. + pub parent_offset: Option>, } impl SparseAdjacency { @@ -86,7 +96,7 @@ impl SparseAdjacency { } } } - Self { n, rows, cols, values, length_m, slope } + Self { n, rows, cols, values, length_m, slope, parent_offset: None } } pub fn nnz(&self) -> usize { diff --git a/src/training/driver.rs b/src/training/driver.rs index a48f2c9..c93fffb 100644 --- a/src/training/driver.rs +++ b/src/training/driver.rs @@ -129,9 +129,11 @@ fn run_micro_batch( debug_assert_eq!(g, num_gauges); // Build obs tensor preserving NaN so the filter can detect them. - // Shape: obs_arr is (rho_days, G); trim first/last day → (t_days, G). + // Shape: obs_arr is (rho_days, G). Under the 2026-08-08 tau convention + // pooled day i is scored against obs day i (see tau_trim_and_downsample; + // the legacy pairing was obs day i+1). assert!( - t_days_full >= 2 + t_days, + t_days_full >= t_days, "obs/pred shape mismatch: obs rows={} pred t_days={}", t_days_full, t_days @@ -139,8 +141,7 @@ fn run_micro_batch( let mut obs_buf: Vec = Vec::with_capacity(g * t_days); for gi in 0..g { for ti in 0..t_days { - // obs row index after trim = ti + 1; column = gi. - obs_buf.push(obs_arr[(ti + 1, gi)]); + obs_buf.push(obs_arr[(ti, gi)]); } } let obs_t: Tensor, 2> = diff --git a/src/training/eval.rs b/src/training/eval.rs index 3b6c4a6..f63dd25 100644 --- a/src/training/eval.rs +++ b/src/training/eval.rs @@ -252,6 +252,31 @@ pub fn evaluate( chunk_idx += 1; } + // DIAGNOSTIC (opt-in, off by default): dump the PRE-TRIM hourly series so + // `params.tau` can be swept EXACTLY offline instead of re-running eval once + // per tau. Raw row-major f32 (n_gauges, n_hours) + a `.json` dims sidecar. + // An env var rather than a new parameter because `evaluate` has several + // call sites and this is a throwaway probe. + if let Ok(dump_path) = std::env::var("DDRS_HOURLY_DUMP") { + use std::io::Write; + let raw: Vec = predictions_full.iter().copied().collect(); + let bytes: &[u8] = unsafe { + std::slice::from_raw_parts(raw.as_ptr() as *const u8, std::mem::size_of_val(&raw[..])) + }; + std::fs::File::create(&dump_path) + .and_then(|mut f| f.write_all(bytes)) + .unwrap_or_else(|e| panic!("DDRS_HOURLY_DUMP write to {dump_path} failed: {e}")); + let meta = format!( + r#"{{"n_gauges":{n_all_gauges},"n_hours":{n_hours_full},"dtype":"f32","order":"C","tau_shipped":{}}}"#, + cfg.params.tau + ); + std::fs::write(format!("{dump_path}.json"), meta).ok(); + eprintln!( + " hourly dump -> {dump_path} ({n_all_gauges} x {n_hours_full} f32, {:.2} GB)", + (n_all_gauges * n_hours_full * 4) as f64 / 1e9 + ); + } + // End-of-pipeline tau-trim + daily downsample. Lift the f32 accumulator // into a BURN tensor for the existing tau_trim_and_downsample helper. let pred_full_vec: Vec = predictions_full.iter().copied().collect(); @@ -276,14 +301,17 @@ pub fn evaluate( .as_standard_layout() .to_owned(); - // Predictions after tau_trim_and_downsample: shape (G, n_days_full - 1). - // (Math: T_hours = n_days_full * 24; trim drops 24 hours total; /24 = n_days_full - 1.) - // To match observations_daily's (G, n_days_full - 2), drop the LAST day - // of predictions. (This SAFE CONSERVATIVE alignment is documented in the - // SP-5 plan Task 6 design note; Task 11 V4 will surface any drift.) + // Predictions after tau_trim_and_downsample: shape (G, n_days_full - 1), + // pooled day i ↔ store day i (2026-08-08 tau convention; T_hours = + // n_days_full * 24, trim drops 24 hours total). Observations above are + // sliced [1..-1] (store days 1..n_days_full-2), so drop the FIRST + // prediction day to pair store day 1..n_days_full-2 on both sides. + // (Under the legacy convention pooled day i ↔ store day i+1 and the + // LAST prediction day was dropped instead; the zarr time axis is + // unchanged by the convention switch.) let pd_dims = predictions_daily.dim(); let predictions_daily = predictions_daily - .slice(s![.., 0..pd_dims.1 - 1]) + .slice(s![.., 1..pd_dims.1]) .to_owned(); debug_assert_eq!( diff --git a/src/training/forward.rs b/src/training/forward.rs index 55d85e0..8dd0006 100644 --- a/src/training/forward.rs +++ b/src/training/forward.rs @@ -33,6 +33,75 @@ pub fn scatter_add_by_group( zeros.scatter(0, group_2d, gathered, IndexingUpdateOp::Add) } +// --------------------------------------------------------------------------- +// Reach subdivision: parent -> sub-reach gather +// --------------------------------------------------------------------------- + +/// Row → parent index for a subdivided network, or `None` when there is nothing +/// to gather. +/// +/// `parent_offset` is the reach-subdivision map in the batch's compressed +/// space: parent `p` owns rows `[off[p], off[p + 1])`. One row per parent means +/// the map is the identity, and returning `None` makes the caller skip the +/// gather **entirely** rather than run `select` with an identity index — that is +/// what keeps `params.subdivision.enabled: false` byte-identical. +fn parent_row_index(parent_offset: Option<&Vec>, n_rows: usize) -> Option> { + let off = parent_offset?; + let n_parent = off.len().checked_sub(1)?; + if n_parent == n_rows { + return None; + } + let mut idx: Vec = Vec::with_capacity(n_rows); + for (p, w) in off.windows(2).enumerate() { + for _ in w[0]..w[1] { + idx.push(p as i32); + } + } + debug_assert_eq!( + idx.len(), + n_rows, + "parent_offset covers {} rows but the network has {n_rows}", + idx.len() + ); + Some(idx) +} + +/// Expand every KAN output from parent resolution `[N_parent]` onto the +/// routing's sub-reach resolution `[n_rows]`. +/// +/// The head is run ONCE per MERIT reach. Sub-reaches inherit their parent's +/// hydraulics — `n`, `p_spatial`, `q_spatial`, slope — because MERIT carries no +/// within-reach variation, so duplicating the attribute rows would cost ~5x the +/// head compute at `max_pieces: 8` for numerically identical outputs. +/// +/// `select`'s backward is a **scatter-add**, so a parent receives the SUMMED +/// gradient from all of its pieces. That is the correct semantics for a shared +/// parameter (a broadcast that only forwarded one piece's gradient would not +/// be), and `tests/subdivision_integration.rs::gradient_sums_back_to_the_parent` +/// pins it. +/// +/// Returns `params` untouched — no op recorded on the tape — when the network +/// is not subdivided. +#[doc(hidden)] +pub fn gather_params_to_subreaches( + params: std::collections::HashMap>, + parent_offset: Option<&Vec>, + n_rows: usize, + device: &B::Device, +) -> std::collections::HashMap> { + let Some(idx) = parent_row_index(parent_offset, n_rows) else { + return params; + }; + let idx = Tensor::::from_data(TensorData::from(idx.as_slice()), device); + params + .into_iter() + .map(|(name, t)| { + let gathered = t.select(0, idx.clone()); + (name, gathered) + }) + .collect() +} + // --------------------------------------------------------------------------- // FrozenParams + forward_with_frozen_params // --------------------------------------------------------------------------- @@ -183,6 +252,10 @@ use crate::nn::kan_head::KanHead; /// with 47% of CONUS pinned at `n = 0.015`, and that was only visible after a /// 14 h run plus a `dump_parameters` pass. Logging it per micro-batch turns a /// post-mortem into a first-epoch observation. +/// +/// Deliberately NOT gathered onto sub-reaches: under `params.subdivision` the +/// statistic is one sample per MERIT reach, not one per piece, so long reaches +/// do not get extra votes on the median and at-floor fraction. pub fn manning_n_stats( cfg: &Config, tensors: &RoutingTensors>, @@ -232,13 +305,20 @@ pub fn forward( device: &I::Device, carry_state: bool, ) -> Tensor, 2> { - let params_map = head.forward(tensors.spatial_attributes.clone()); + 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. + let params_map = gather_params_to_subreaches( + head.forward(tensors.spatial_attributes.clone()), + tensors.adjacency.parent_offset.as_ref(), + n_active, + device, + ); let n_param = params_map.get("n").expect("MLP missing n").clone(); let q_param = params_map.get("q_spatial").expect("MLP missing q_spatial").clone(); let p_param = params_map.get("p_spatial").cloned(); - let n_active = tensors.adjacency.n; // Learnable Muskingum X: when the KAN emits `x_storage`, denormalize its // [0,1] output to the configured range so the routing learns its own // attenuation-vs-translation per reach (gradient already flows via the @@ -300,6 +380,9 @@ pub fn forward( carry_state, tensors.initial_state.clone(), ); + // Enable negative-discharge tracking so the count appears in the training + // log. When use_cuda_graphs is true, forward will print UNAVAILABLE instead. + engine.enable_negative_discharge_tracking(); let runoff = engine.forward(); // (N, T_hours) @@ -525,7 +608,15 @@ fn forward_eval_core( overrides: Option<&LeakanceOverride>, param_overrides: Option<&RoutingParamOverride>, ) -> Tensor { - let params_map = head.forward(tensors.spatial_attributes.clone()); + let n_active = tensors.adjacency.n; + // Parent → sub-reach gather, before the overrides below: `RoutingParamOverride` + // and `LeakanceOverride` are both sized to the routing network. + let params_map = gather_params_to_subreaches( + head.forward(tensors.spatial_attributes.clone()), + tensors.adjacency.parent_offset.as_ref(), + n_active, + device, + ); let n_param = params_map.get("n").expect("MLP missing n").clone(); let q_param = params_map.get("q_spatial").expect("MLP missing q_spatial").clone(); @@ -549,7 +640,6 @@ fn forward_eval_core( (n_param, q_param, p_param) }; - let n_active = tensors.adjacency.n; // Learnable Muskingum X (eval path mirrors `forward`): denormalize the // KAN's `x_storage` output when present, else the constant 0.3. let x_storage: Tensor = match params_map.get("x_storage") { diff --git a/src/training/loss.rs b/src/training/loss.rs index d24f61c..3938340 100644 --- a/src/training/loss.rs +++ b/src/training/loss.rs @@ -12,24 +12,33 @@ use crate::config::{LossConfig, LossKind}; /// Tau-trim then daily downsample via area-mode adaptive average pooling. /// -/// Mirrors DDR's `~/projects/ddr/src/ddr/io/functions.py:22`: +/// Pooling mirrors DDR's `~/projects/ddr/src/ddr/io/functions.py:22`: /// `F.interpolate(data.unsqueeze(1), size=(rho,), mode="area").squeeze(1)`. /// -/// Input shape `(G, T_hours)`. Slicing convention from DDR -/// `compute_daily_runoff`: `[13 + tau : -11 + tau]`. The trimmed length -/// does NOT need to be a multiple of 24 — fractional boundary hours are -/// handled by area-mode pooling. +/// Input shape `(G, T_hours)`. Slicing convention (since 2026-08-08): +/// `[tau : -(24 - tau)]` — pooled day `i` covers hours +/// `[tau + 24i, tau + 24(i+1))` and is scored against OBSERVATION DAY `i`. +/// `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)`. /// -/// Returns `(G, T_days)` where `T_days = T_hours_trimmed // 24` (matching -/// DDR's `num_days` computation at `scripts/train.py:78`). +/// Legacy mapping: the pre-2026-08-08 slice was `[13+tau : -11+tau]` with +/// pooled day `i` scored against obs day `i+1`; `tau_new = tau_old - 11` +/// (so old shipped 3 ≡ new −8, old optimum 20 ≡ new 9). DDR-Python's +/// `compute_daily_runoff` still uses the legacy form. The total trim is +/// 24 h under both conventions, so `T_days = T_hours/24 - 1` is unchanged. +/// +/// Returns `(G, T_days)` where `T_days = T_hours_trimmed // 24`. pub fn tau_trim_and_downsample( predictions_hourly: Tensor, tau: u32, ) -> Tensor { let dims = predictions_hourly.dims(); let (g, t_hours) = (dims[0], dims[1]); - let start = 13 + tau as usize; - let end = t_hours - 11 + tau as usize; + assert!(tau < 24, "tau must be in [0, 24) hours; got {tau}"); + let start = tau as usize; + let end = t_hours - (24 - tau as usize); assert!(start < end, "tau-trim window degenerate: [{start}, {end})"); let t_trimmed = end - start; let t_days = t_trimmed / 24; @@ -665,9 +674,9 @@ mod tests { #[test] fn tau_trim_matches_old_block_mean_on_divisible_input() { - // Verify the new area-pool body reduces to block-mean whenever - // the trimmed window IS a multiple of 24. tau=11, T=72 → - // trimmed window is hours [24..72] (length 48 = 2 days exactly). + // Verify the area-pool body reduces to block-mean whenever the + // trimmed window IS a multiple of 24. tau=0, T=72 → trimmed window + // is hours [0..48) (length 48 = 2 days exactly, day-aligned). let device = Default::default(); let v: Vec = (0..72).map(|x| x as f32).collect(); let input: Tensor = Tensor::::from_data( @@ -675,12 +684,47 @@ mod tests { &device, ) .reshape([1, 72]); - let out = tau_trim_and_downsample(input, 11); + let out = tau_trim_and_downsample(input, 0); let got: Vec = out.into_data().to_vec().unwrap(); - // Sliced = hours 24..72 (48 values: 24..=71). + // Sliced = hours 0..48. + // Day 0 = mean(0..=23) = 11.5 // Day 1 = mean(24..=47) = 35.5 - // Day 2 = mean(48..=71) = 59.5 - assert!((got[0] - 35.5).abs() < 1e-4, "got {}", got[0]); - assert!((got[1] - 59.5).abs() < 1e-4, "got {}", got[1]); + assert!((got[0] - 11.5).abs() < 1e-4, "got {}", got[0]); + assert!((got[1] - 35.5).abs() < 1e-4, "got {}", got[1]); + } + + #[test] + fn new_tau_equals_legacy_tau_plus_eleven_shifted_one_day() { + // Convention-change equivalence: new tau=t reproduces the legacy + // slice [13+(t+11) : -11+(t+11)] exactly, offset by one pooled day + // (legacy day i was scored against obs day i+1; new day i against + // obs day i). new[:, j+1] == legacy[:, j] for all j. + let device = Default::default(); + let t_hours = 24 * 10; + let v: Vec = (0..t_hours).map(|x| ((x * 37) % 101) as f32).collect(); + let input: Tensor = Tensor::::from_data( + burn::tensor::TensorData::new(v.clone(), [t_hours]), + &device, + ) + .reshape([1, t_hours]); + let tau_new: u32 = 9; // ≡ legacy tau 20 + let new_out: Vec = tau_trim_and_downsample(input, tau_new) + .into_data() + .to_vec() + .unwrap(); + // Legacy formula, computed by hand: start 13+20=33, end T-11+20=T+9… the + // legacy end offset (-11+tau) only stays in-bounds for tau<=11, so build + // the expected bins directly from the window definition instead: legacy + // day j covered hours [33 + 24j, 33 + 24(j+1)). + let n_days = new_out.len(); + for j in 0..n_days - 1 { + let s = 33 + 24 * j; + let expect: f32 = v[s..s + 24].iter().sum::() / 24.0; + assert!( + (new_out[j + 1] - expect).abs() < 1e-4, + "day {j}: legacy {expect} vs new[j+1] {}", + new_out[j + 1] + ); + } } } diff --git a/src/training/probe.rs b/src/training/probe.rs index 0fed3c2..020069c 100644 --- a/src/training/probe.rs +++ b/src/training/probe.rs @@ -21,7 +21,7 @@ use crate::data::dataset::RoutingTensors; use crate::nn::kan_head::KanHead; use crate::routing::utils::denormalize; use crate::routing::{MuskingumCunge, RoutingInputs, SpatialParameters}; -use crate::training::forward::scatter_add_by_group; +use crate::training::forward::{gather_params_to_subreaches, scatter_add_by_group}; /// Detach `t` from its autograd graph and re-lift it as a `require_grad` /// leaf. Values are bit-identical; only the tape topology changes. @@ -70,9 +70,16 @@ pub fn probe_forward( "probing leakance params requires params.use_leakance: true" ); - let params_map = head.forward(tensors.spatial_attributes.clone()); - let n_active = tensors.adjacency.n; + // The head runs at PARENT resolution; expand to routing rows so the lifted + // leaves below are per-reach-row, as every consumer here assumes. No-op + // unless the adjacency was built with `params.subdivision.enabled`. + let params_map = gather_params_to_subreaches( + head.forward(tensors.spatial_attributes.clone()), + tensors.adjacency.parent_offset.as_ref(), + n_active, + device, + ); // Build lifted leaves from the params_map for every name in `lift`. // Must come before x_storage so the lifted-or-fallback pattern below diff --git a/tests/celerity_beta.rs b/tests/celerity_beta.rs new file mode 100644 index 0000000..4a0379a --- /dev/null +++ b/tests/celerity_beta.rs @@ -0,0 +1,419 @@ +//! Trapezoidal kinematic-wave celerity (`ddr_match: false`). +//! +//! Part 1 — the physics. `beta = c / v = (dQ/dA) / (Q/A)` for a trapezoid under +//! Manning. DDR (and `ddr_match: true`) hardcodes `5/3`, which is the +//! *wide-rectangular* limit `b/y -> inf`. The channels this solver builds have +//! `kappa = b/y ~ 0.7-1.8`, so the true ratio is ~1.30-1.36 and `5/3` is +//! 22-27% too high. +//! +//! beta = 5/3 - (4/3)·A·sqrt(1+z²)/(T·P) +//! +//! Part 2 — the gradcheck. `celerity = velocity_cl · beta` adds four +//! hand-derived terms to the analytical backward in `src/routing/mmc_op.rs` +//! (into the area / top-width / wetted-perimeter / side-slope accumulators). +//! Those are validated against central finite differences below, with +//! `mock_cfg()` setting `ddr_match = false`. + +use std::sync::Arc; + +use burn::backend::{Autodiff, NdArray}; +use burn::tensor::Tensor; + +use ddrs::config::Config; +use ddrs::routing::mmc_op::timestep_forward; +use ddrs::sparse::{AValuesAssembler, CsrPattern, SparseAdjacency}; + +// =========================================================================== +// Part 1: the beta formula, in f64, independent of any ddrs code. +// =========================================================================== + +/// Trapezoid cross-sectional area: `A = (b + z·y)·y`. +fn area(b: f64, z: f64, y: f64) -> f64 { + (b + z * y) * y +} + +/// Trapezoid top width: `T = b + 2·z·y`. +fn top_width(b: f64, z: f64, y: f64) -> f64 { + b + 2.0 * z * y +} + +/// Trapezoid wetted perimeter: `P = b + 2·y·sqrt(1+z²)`. +fn wetted_perimeter(b: f64, z: f64, y: f64) -> f64 { + b + 2.0 * y * (1.0 + z * z).sqrt() +} + +/// `beta = 5/3 - (4/3)·A·sqrt(1+z²)/(T·P)`, the exact `c/v` ratio. +fn beta(b: f64, z: f64, y: f64) -> f64 { + let a = area(b, z, y); + let t = top_width(b, z, y); + let p = wetted_perimeter(b, z, y); + 5.0 / 3.0 - (4.0 / 3.0) * a * (1.0 + z * z).sqrt() / (t * p) +} + +/// Manning discharge with `n = 1`, `S = 1` (both cancel out of `c/v`). +fn manning_q(b: f64, z: f64, y: f64) -> f64 { + let a = area(b, z, y); + let p = wetted_perimeter(b, z, y); + a.powf(5.0 / 3.0) * p.powf(-2.0 / 3.0) +} + +/// `dQ/dA` by central differences in depth: `(dQ/dy) / (dA/dy)`. +fn fd_dq_da(b: f64, z: f64, y: f64) -> f64 { + let h = 1e-6 * y; + let dq = manning_q(b, z, y + h) - manning_q(b, z, y - h); + let da = area(b, z, y + h) - area(b, z, y - h); + dq / da +} + +/// `v = Q/A`. +fn velocity(b: f64, z: f64, y: f64) -> f64 { + manning_q(b, z, y) / area(b, z, y) +} + +#[test] +fn beta_wide_rectangular_limit_is_five_thirds() { + // b/y -> inf with z = 0 is the Kleitz-Seddon wide-rectangular case that + // DDR's hardcoded 5/3 actually describes. + let got = beta(1e6, 0.0, 2.0); + assert!( + (got - 5.0 / 3.0).abs() < 1e-5, + "wide rectangular beta should approach 5/3, got {got}" + ); +} + +#[test] +fn beta_triangular_limit_is_four_thirds() { + // b = 0: A = z·y², T = 2·z·y, P = 2·y·sqrt(1+z²) + // => (4/3)·A·sqrt(1+z²)/(T·P) = (4/3)·(1/4) = 1/3 exactly, for any z, y. + for &(z, y) in &[(2.0, 2.0), (0.5, 7.0), (3.3, 0.4)] { + let got = beta(0.0, z, y); + assert!( + (got - 4.0 / 3.0).abs() < 1e-12, + "triangular beta should be exactly 4/3, got {got} at z={z} y={y}" + ); + } +} + +#[test] +fn beta_matches_finite_difference_dq_da() { + // Realistic (b, z, y) triples from the trapezoids this solver builds. + for &(b, z, y) in &[ + (6.3_f64, 0.50_f64, 9.31_f64), + (13.8, 0.57, 8.09), + (8.8, 2.76, 6.38), + (21.4, 1.10, 3.02), + ] { + let analytic = beta(b, z, y) * velocity(b, z, y); + let numeric = fd_dq_da(b, z, y); + let rel = (analytic - numeric).abs() / numeric.abs(); + assert!( + rel < 1e-6, + "beta·v != dQ/dA at (b={b}, z={z}, y={y}): analytic={analytic} fd={numeric} rel={rel}" + ); + } +} + +#[test] +fn beta_is_bounded_by_one_and_five_thirds() { + // beta is NOT monotone in b/y and is not bounded below by the triangular + // 4/3: for narrow, near-rectangular sections (b/y -> 0, z -> 0) it decays + // toward 1 (A=by, T=b, P->2y => G->2/3). 5/3 remains the strict upper + // bound (wide rectangular). Sweep confirms inf(beta) = 1. + for &(b, z, y) in &[ + (6.3_f64, 0.50_f64, 9.31_f64), + (13.8, 0.57, 8.09), + (8.8, 2.76, 6.38), + (0.05, 0.0, 1.0), + ] { + let got = beta(b, z, y); + assert!( + got > 1.0 && got < 5.0 / 3.0, + "beta out of (1, 5/3) at (b={b}, z={z}, y={y}): {got}" + ); + } +} + +#[test] +fn hardcoded_five_thirds_is_twenty_to_thirty_percent_high() { + // Regression guard on the magnitude of the bug DDR carries. + let b = beta(13.8, 0.57, 8.09); + let overshoot = (5.0 / 3.0) / b - 1.0; + assert!( + (0.20..0.30).contains(&overshoot), + "expected 5/3 to be 20-30% high vs beta={b}, got {:.1}%", + overshoot * 100.0 + ); +} + +// =========================================================================== +// Part 2: gradcheck of the `ddr_match: false` backward. +// Scaffolding mirrors `tests/sp8_gradcheck.rs`; only `mock_cfg` differs. +// =========================================================================== + +type I = NdArray; +type AB = Autodiff; + +const N: usize = 4; +const EPS: f32 = 1e-3; +const REL_TOL: f32 = 5e-3; +const ABS_TOL: f32 = 1e-4; + +#[derive(Copy, Clone, Debug)] +enum Parent { + N, + QSpatial, + PSpatial, + QT, +} + +fn linear_chain_sparse() -> SparseAdjacency { + 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]) +} + +fn mock_cfg() -> Config { + let mut cfg = Config::default(); + // The whole point of this file: exercise the corrected-physics branch. + cfg.params.ddr_match = false; + cfg.params.parameter_ranges.n = [0.01, 0.1]; + cfg.params.parameter_ranges.q_spatial = [0.1, 0.9]; + cfg.params.parameter_ranges.p_spatial = [1.0, 200.0]; + cfg.params.attribute_minimums.velocity = 0.1; + cfg.params.attribute_minimums.depth = 0.01; + cfg.params.attribute_minimums.discharge = 0.001; + cfg.params.attribute_minimums.bottom_width = 0.1; + cfg.params.attribute_minimums.slope = 0.001; + cfg.params.defaults.insert("p_spatial".to_string(), 1.0); + cfg.params.log_space_parameters = vec![]; + cfg +} + +fn default_inputs() -> (Vec, Vec, Vec, Vec, Vec) { + let n_vec = vec![0.035f32; N]; + let qsp_vec = vec![0.4f32; N]; + let psp_vec = vec![20.0f32; N]; + let qt_vec = vec![100.0f32, 120.0, 140.0, 160.0]; + let qpt_vec = vec![10.0f32, 12.0, 14.0, 16.0]; + (n_vec, qsp_vec, psp_vec, qt_vec, qpt_vec) +} + +struct GradTensors { + n: Tensor, + qsp: Tensor, + psp: Tensor, + qt: Tensor, +} + +#[allow(clippy::too_many_arguments)] +fn run_forward_loss( + cfg: &Config, + pattern: &Arc, + assembler: &AValuesAssembler, + device: &::Device, + n_vec: &[f32], + qsp_vec: &[f32], + psp_vec: &[f32], + qt_vec: &[f32], + qpt_vec: &[f32], + length_vec: &[f32], + slope_vec: &[f32], + x_storage_vec: &[f32], + require_grad_parent: Option, +) -> (Tensor, GradTensors) { + let mk = |data: &[f32], req: bool| -> Tensor { + let t: Tensor = Tensor::from_floats(data, device); + if req { t.require_grad() } else { t } + }; + let n_t = mk(n_vec, matches!(require_grad_parent, Some(Parent::N))); + let qsp_t = mk(qsp_vec, matches!(require_grad_parent, Some(Parent::QSpatial))); + let psp_t = mk(psp_vec, matches!(require_grad_parent, Some(Parent::PSpatial))); + let qt_t = mk(qt_vec, matches!(require_grad_parent, Some(Parent::QT))); + let qpt_t = mk(qpt_vec, false); + let length_t = mk(length_vec, false); + let slope_t = mk(slope_vec, false); + let xst_t = mk(x_storage_vec, false); + + let q_next = timestep_forward::( + cfg, + pattern, + assembler, + n_t.clone(), + qsp_t.clone(), + psp_t.clone(), + qt_t.clone(), + qpt_t.clone(), + length_t, + slope_t, + xst_t, + false, + ); + + ( + q_next, + GradTensors { + n: n_t, + qsp: qsp_t, + psp: psp_t, + qt: qt_t, + }, + ) +} + +fn compute_analytical_grad(parent: Parent) -> Vec { + let cfg = mock_cfg(); + let adj = linear_chain_sparse(); + let device = ::Device::default(); + let pattern = Arc::new(CsrPattern::from_sparse(&adj)); + let assembler = AValuesAssembler::::new(&pattern, &device); + + let (n_vec, qsp_vec, psp_vec, qt_vec, qpt_vec) = default_inputs(); + let length_vec = adj.length_m.clone(); + let slope_vec = adj.slope.clone(); + let x_storage_vec = vec![0.3f32; N]; + + let (q_next, parents) = run_forward_loss( + &cfg, + &pattern, + &assembler, + &device, + &n_vec, + &qsp_vec, + &psp_vec, + &qt_vec, + &qpt_vec, + &length_vec, + &slope_vec, + &x_storage_vec, + Some(parent), + ); + + let loss = q_next.sum(); + let grads = loss.backward(); + + let g = match parent { + Parent::N => parents.n.grad(&grads).expect("grad on n"), + Parent::QSpatial => parents.qsp.grad(&grads).expect("grad on q_spatial"), + Parent::PSpatial => parents.psp.grad(&grads).expect("grad on p_spatial"), + Parent::QT => parents.qt.grad(&grads).expect("grad on q_t"), + }; + g.into_data().to_vec::().unwrap() +} + +fn compute_fd_grad(parent: Parent) -> Vec { + let cfg = mock_cfg(); + let adj = linear_chain_sparse(); + let device = ::Device::default(); + let pattern = Arc::new(CsrPattern::from_sparse(&adj)); + let assembler = AValuesAssembler::::new(&pattern, &device); + + let (n_vec, qsp_vec, psp_vec, qt_vec, qpt_vec) = default_inputs(); + let length_vec = adj.length_m.clone(); + let slope_vec = adj.slope.clone(); + let x_storage_vec = vec![0.3f32; N]; + + let eval_loss = |n: &[f32], qsp: &[f32], psp: &[f32], qt: &[f32], qpt: &[f32]| -> f32 { + let (q_next, _) = run_forward_loss( + &cfg, + &pattern, + &assembler, + &device, + n, + qsp, + psp, + qt, + qpt, + &length_vec, + &slope_vec, + &x_storage_vec, + None, + ); + let v: Vec = q_next.sum().into_data().to_vec::().unwrap(); + v[0] + }; + + let mut grad = vec![0.0f32; N]; + for i in 0..N { + let mut plus_n = n_vec.clone(); + let mut plus_qsp = qsp_vec.clone(); + let mut plus_psp = psp_vec.clone(); + let mut plus_qt = qt_vec.clone(); + let mut minus_n = n_vec.clone(); + let mut minus_qsp = qsp_vec.clone(); + let mut minus_psp = psp_vec.clone(); + let mut minus_qt = qt_vec.clone(); + + let (plus, minus, base) = match parent { + Parent::N => (&mut plus_n, &mut minus_n, &n_vec), + Parent::QSpatial => (&mut plus_qsp, &mut minus_qsp, &qsp_vec), + Parent::PSpatial => (&mut plus_psp, &mut minus_psp, &psp_vec), + Parent::QT => (&mut plus_qt, &mut minus_qt, &qt_vec), + }; + let eps = (EPS * base[i].abs()).max(EPS); + plus[i] = base[i] + eps; + minus[i] = base[i] - eps; + + let l_plus = eval_loss(&plus_n, &plus_qsp, &plus_psp, &plus_qt, &qpt_vec); + let l_minus = eval_loss(&minus_n, &minus_qsp, &minus_psp, &minus_qt, &qpt_vec); + grad[i] = (l_plus - l_minus) / (2.0 * eps); + } + grad +} + +fn compare_grads(name: &str, analytical: &[f32], fd: &[f32]) { + assert_eq!(analytical.len(), fd.len()); + println!("--- {name} ---"); + let mut worst_rel = 0.0f32; + let mut worst_abs = 0.0f32; + for i in 0..analytical.len() { + let a = analytical[i]; + let f = fd[i]; + let abs_diff = (a - f).abs(); + let denom = a.abs().max(f.abs()).max(1e-12); + let rel_diff = abs_diff / denom; + worst_abs = worst_abs.max(abs_diff); + worst_rel = worst_rel.max(rel_diff); + println!(" [{i}] analytical={a:.6e} fd={f:.6e} abs={abs_diff:.3e} rel={rel_diff:.3e}"); + } + println!(" worst abs={worst_abs:.3e} worst rel={worst_rel:.3e}"); + let pass = analytical.iter().zip(fd).all(|(&a, &f)| { + let abs_diff = (a - f).abs(); + let denom = a.abs().max(f.abs()).max(1e-12); + let rel_diff = abs_diff / denom; + rel_diff < REL_TOL || abs_diff < ABS_TOL + }); + assert!( + pass, + "{name}: gradcheck failed (worst rel={worst_rel:.3e}, abs={worst_abs:.3e})" + ); +} + +#[test] +fn gradcheck_beta_n() { + let a = compute_analytical_grad(Parent::N); + let fd = compute_fd_grad(Parent::N); + compare_grads("n", &a, &fd); +} + +#[test] +fn gradcheck_beta_q_spatial() { + let a = compute_analytical_grad(Parent::QSpatial); + let fd = compute_fd_grad(Parent::QSpatial); + compare_grads("q_spatial", &a, &fd); +} + +#[test] +fn gradcheck_beta_p_spatial() { + let a = compute_analytical_grad(Parent::PSpatial); + let fd = compute_fd_grad(Parent::PSpatial); + compare_grads("p_spatial", &a, &fd); +} + +#[test] +fn gradcheck_beta_q_t() { + let a = compute_analytical_grad(Parent::QT); + let fd = compute_fd_grad(Parent::QT); + compare_grads("q_t", &a, &fd); +} diff --git a/tests/cuda_backward_parity.rs b/tests/cuda_backward_parity.rs new file mode 100644 index 0000000..316a807 --- /dev/null +++ b/tests/cuda_backward_parity.rs @@ -0,0 +1,1248 @@ +//! CUDA-backend verification of the `ddr_match: false` physics backwards. +//! +//! Build with: `cargo test --features cuda --test cuda_backward_parity` +//! +//! # Why this file exists +//! +//! Three physics corrections landed with hand-written BURN-0.21 +//! `Backward` implementations rather than autograd-tape unrolling +//! (invariant 4, `docs/reference/burn-autograd.md`): +//! +//! 1. trapezoidal celerity `c = v·β`, `β = 5/3 − (4/3)·A·√(1+z²)/(T·P)` (S17) +//! 2. Cunge `X = clamp(0.5(1 − Q/(B·S·c·L)), 0, 0.5)` (S19) +//! 3. the positivity clamp — K floor + three-way X cap (S18'/S19') +//! +//! Every existing gradcheck for them (`tests/celerity_beta.rs`, +//! `tests/cunge_x.rs`, `tests/positivity_clamp.rs`, `tests/sparse_gradcheck.rs`) +//! declares `type I = NdArray` — **CPU only**. But `ddr_match: false` + +//! `use_cuda_graphs: false` + a CUDA backend is a legal, actively-used +//! configuration, so those backwards ship un-exercised on the hardware that +//! actually runs them. This file closes that gap. +//! +//! The CUDA-graph path is separately walled off and is NOT this file's problem: +//! `validate_ddr_match` rejects `ddr_match: false` + `use_cuda_graphs: true`, +//! and `validate_enforce_positivity` requires `enforce_positivity ⟹ !ddr_match`, +//! so transitively `enforce_positivity ⟹ !use_cuda_graphs`. That transitive +//! implication is asserted in Part D — it was previously only *implied* by two +//! independent validators and never tested. +//! +//! # Structure +//! +//! * **Part A** — native central-difference gradcheck with `Cuda` as +//! the inner backend. This is the load-bearing evidence: it shows CUDA +//! gradients are *correct*, not merely *consistent with CPU*. +//! * **Part B** — CUDA-vs-CPU analytic gradient parity on an identical fixture. +//! * **Part C** — non-vacuity guards, evaluated *on the CUDA backend*, so a +//! fixture that went degenerate only on GPU would be caught. +//! * **Part D** — the transitive `enforce_positivity ⟹ !use_cuda_graphs` guard. +//! +//! Every numeric helper below is generic in the inner backend `I` and is called +//! with BOTH `NdArray` and `Cuda`. That is deliberate: a parity +//! claim is only meaningful if both sides execute the same source. +//! +//! # Falsifiability +//! +//! This repo has a burned-in lesson: `tests/cunge_x.rs` was originally VACUOUS +//! (1000 m reaches saturated X's clamp on every reach, so all four tests passed +//! with the backward terms DELETED). The guards in Part C are the standing +//! defense, and they run as a PRECONDITION of every gradcheck rather than as +//! standalone tests that a future retune could leave behind. + +#![cfg(feature = "cuda")] + +use std::sync::Arc; + +use burn::backend::{Autodiff, NdArray}; +use burn::tensor::backend::Backend; +use burn::tensor::Tensor; + +use ddrs::config::Config; +use ddrs::routing::mmc_op::{timestep_forward, POSITIVITY_DELTA}; +use ddrs::sparse::{AValuesAssembler, CsrPattern, SparseAdjacency}; + +/// CPU reference backend. +type Cpu = NdArray; +/// GPU backend under test. `burn::backend::Cuda` needs the umbrella crate's +/// "cuda" feature; the direct crate is the convention here — see +/// `tests/kan_head_fixture_forward.rs` and `tests/cusparse_ptr_spike.rs`. +type Gpu = burn_cuda::Cuda; + +/// `crate::routing::mmc::DT_SECONDS`, restated so a change there fails loudly +/// here rather than silently retuning the fixture. +const DT: f32 = 3600.0; + +fn device() -> I::Device { + ::Device::default() +} + +// =========================================================================== +// Fixture — lifted from `tests/positivity_clamp.rs`, unchanged. +// +// The reach lengths span 400 m .. 150 km so `Cr = Δt·c/L` sweeps ~0.05 .. ~7.9: +// the short end sits deep in the `c3 < 0` regime, the long end deep in the +// `c1 < 0` regime, and the middle keeps at least one reach where Cunge X wins +// the three-way min strictly inside `[0, 0.5]`. +// =========================================================================== + +struct Fixture { + length: Vec, + slope: Vec, + n: Vec, + qsp: Vec, + psp: Vec, + qt: Vec, + qpt: Vec, +} + +/// The gradcheck fixture. `q_t` is GRADED along the chain, and that is +/// load-bearing, not decoration — see `tests/positivity_clamp.rs::grad_fixture`: +/// with a flat `q_t`, the partition identity `c1+c2+c3 = 1` on a chain where +/// `x_up ≈ i_t ≈ q_t` makes `q_next ≈ q_t` regardless of `(K, X)`, every +/// S17/S19/S18' effect cancels to f32 noise, and the reaches would be +/// gradient-DEAD. On GPU that would show up as "CUDA and CPU agree perfectly" +/// — both computing nothing. +/// +/// `q'` is small (0.01 m³/s) so the `c3 < 0` reaches are not masked by `c4·q'`. +fn grad_fixture() -> Fixture { + let length = vec![ + 400.0f32, 800.0, 1500.0, 2500.0, 3600.0, 5000.0, 8000.0, 20000.0, 60000.0, 150000.0, + ]; + let n = length.len(); + Fixture { + length, + slope: vec![0.001; n], + n: vec![0.035; n], + qsp: vec![0.4; n], + psp: vec![20.0; n], + qt: vec![6.0, 8.0, 15.0, 50.0, 120.0, 200.0, 300.0, 40.0, 400.0, 30.0], + qpt: vec![0.01; n], + } +} + +/// `enforce_positivity` is the only knob: `false` gives β celerity + Cunge X, +/// `true` adds S18'/S19'. There is deliberately no config where β is on and +/// Cunge X is off — both are gated on the single `ddr_match: false` flag — so +/// two configs cover all three patterns. +fn stress_cfg(enforce_positivity: bool) -> Config { + let mut cfg = Config::default(); + cfg.params.ddr_match = false; + cfg.params.enforce_positivity = enforce_positivity; + cfg.params.parameter_ranges.n = [0.01, 0.3]; + cfg.params.parameter_ranges.q_spatial = [0.1, 0.9]; + cfg.params.parameter_ranges.p_spatial = [1.0, 200.0]; + cfg.params.attribute_minimums.velocity = 0.01; + cfg.params.attribute_minimums.depth = 0.001; + cfg.params.attribute_minimums.discharge = 1e-4; + cfg.params.attribute_minimums.bottom_width = 0.01; + cfg.params.attribute_minimums.slope = 0.0001; + cfg.params.defaults.insert("p_spatial".to_string(), 1.0); + cfg.params.log_space_parameters = vec![]; + cfg +} + +fn chain(f: &Fixture) -> SparseAdjacency { + let n = f.length.len(); + 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, f.length.clone(), f.slope.clone()) +} + +/// The S1..S23 quantities the non-vacuity guards need, read out of the REAL +/// forward chain on backend `I` (not recomputed from a parallel model). +struct ChainOutputs { + side_slope: Vec, + velocity_clamped: Vec, + celerity: Vec, + k_muskingum: Vec, + top_width: Vec, + c1: Vec, + c2: Vec, + c3: Vec, +} + +fn run_chain(f: &Fixture, enforce_positivity: bool) -> ChainOutputs +where + I::FloatTensorPrimitive: 'static, + I::Device: 'static, +{ + let adj = chain(f); + let dev = device::(); + let pattern = Arc::new(CsrPattern::from_sparse(&adj)); + let mk = |d: &[f32]| -> Tensor { Tensor::from_floats(d, &dev) }; + let outs = ddrs::routing::mmc_op::__spike_forward_chain_k1_outputs::( + &stress_cfg(enforce_positivity), + &pattern, + mk(&f.n), + mk(&f.qsp), + mk(&f.psp), + mk(&f.qt), + mk(&f.qpt), + mk(&f.length), + mk(&f.slope), + mk(&vec![0.3f32; f.length.len()]), + ); + // k1 output order: [depth, top_width, side_slope, bottom_width, hyd_radius, + // velocity_un, velocity_cl, celerity, k_muskingum, denom, + // c1, c2, c3, c4, ...] + ChainOutputs { + top_width: outs[1].clone(), + side_slope: outs[2].clone(), + velocity_clamped: outs[6].clone(), + celerity: outs[7].clone(), + k_muskingum: outs[8].clone(), + c1: outs[10].clone(), + c2: outs[11].clone(), + c3: outs[12].clone(), + } +} + +fn run_solve(f: &Fixture, enforce_positivity: bool) -> Vec +where + I::FloatTensorPrimitive: 'static, + I::Device: 'static, +{ + let adj = chain(f); + let dev = device::(); + let pattern = Arc::new(CsrPattern::from_sparse(&adj)); + let mk = |d: &[f32]| -> Tensor { Tensor::from_floats(d, &dev) }; + let (_b, _i, x_sol, _q) = ddrs::routing::mmc_op::__spike_forward_chain_k23_outputs::( + &stress_cfg(enforce_positivity), + &pattern, + mk(&f.n), + mk(&f.qsp), + mk(&f.psp), + mk(&f.qt), + mk(&f.qpt), + mk(&f.length), + mk(&f.slope), + mk(&vec![0.3f32; f.length.len()]), + ); + x_sol +} + +// =========================================================================== +// Non-vacuity instrumentation (shared by Part A and Part C) +// =========================================================================== + +/// Which branch of `x_eff = min(x_cunge, hi_a, hi_b)` won on a reach. The +/// tie-break cascade `Cunge > hi_a > hi_b` mirrors the backward's mask cascade +/// in `mmc_op.rs` B19' exactly. +#[derive(Copy, Clone, PartialEq, Eq, Debug)] +enum Branch { + Cunge, + HiA, + HiB, +} + +struct Report { + branches: Vec, + floored: Vec, + cr_raw: Vec, + x_cunge: Vec, + /// `β = c / v_clamped`, recovered from the chain outputs. + beta: Vec, + side_slope: Vec, +} + +fn report(f: &Fixture, enforce_positivity: bool) -> Report +where + I::FloatTensorPrimitive: 'static, + I::Device: 'static, +{ + let out = run_chain::(f, enforce_positivity); + let k_floor = DT * (1.0 + POSITIVITY_DELTA) / 2.0; + let mut r = Report { + branches: Vec::new(), + floored: Vec::new(), + cr_raw: Vec::new(), + x_cunge: Vec::new(), + beta: Vec::new(), + side_slope: out.side_slope.clone(), + }; + for i in 0..f.length.len() { + let k_raw = f.length[i] / out.celerity[i]; + r.floored.push(k_raw < k_floor); + r.cr_raw.push(DT / k_raw); + r.beta.push(out.celerity[i] / out.velocity_clamped[i]); + let w = f.qt[i] / (out.top_width[i] * f.slope[i] * out.celerity[i] * f.length[i] + 1e-12); + let x_cunge = (0.5 * (1.0 - w)).clamp(0.0, 0.5); + r.x_cunge.push(x_cunge); + // `cr` uses the POST-floor K (S18'), matching how the forward composes + // S18' into S19'. + let cr = DT / out.k_muskingum[i]; + let hi_a = cr * 0.5 * (1.0 - POSITIVITY_DELTA); + let hi_b = (1.0 - 0.5 * cr) * (1.0 - POSITIVITY_DELTA); + r.branches.push(if x_cunge <= hi_a && x_cunge <= hi_b { + Branch::Cunge + } else if hi_a <= hi_b { + Branch::HiA + } else { + Branch::HiB + }); + } + r +} + +/// STANDING GUARD, run as a PRECONDITION of every gradcheck below. +/// +/// Five independent ways this file could go vacuous: +/// +/// 1. `β ≡ 5/3` (or constant) → the S17 backward terms `∂β/∂A,T,P,z` carry no +/// signal and could be deleted; +/// 2. `side_slope` saturated on its `[0.5, 50]` clamp everywhere → `∂β/∂z` is +/// masked to zero on every reach; +/// 3. every Cunge-branch win sitting on X's own `[0, 0.5]` clamp → exactly the +/// failure that made `tests/cunge_x.rs` vacuous at 1000 m; +/// 4. one branch of the three-way min always winning → the other two masks are +/// never exercised (positivity config only); +/// 5. every reach on the same side of the K floor → S18' is a no-op or a +/// constant (positivity config only). +fn assert_fixture_exercises_everything(label: &str, enforce_positivity: bool) +where + I::FloatTensorPrimitive: 'static, + I::Device: 'static, +{ + let f = grad_fixture(); + let r = report::(&f, enforce_positivity); + println!("--- non-vacuity [{label}] enforce_positivity={enforce_positivity} ---"); + for i in 0..r.branches.len() { + println!( + " [{i}] L={:>8.0} q_t={:>6.1} beta={:.4} z={:>8.4} Cr_raw={:>7.3} \ + floored={:<5} x_cunge={:.4} branch={:?}", + f.length[i], + f.qt[i], + r.beta[i], + r.side_slope[i], + r.cr_raw[i], + r.floored[i], + r.x_cunge[i], + r.branches[i] + ); + } + + // (1) β must be genuinely trapezoidal and must VARY. `5/3` is the + // wide-rectangular limit; a fixture sitting there tests nothing. + let beta_min = r.beta.iter().cloned().fold(f32::INFINITY, f32::min); + let beta_max = r.beta.iter().cloned().fold(f32::NEG_INFINITY, f32::max); + println!(" beta span: {beta_min:.4} .. {beta_max:.4} (5/3 = {:.4})", 5.0 / 3.0); + assert!( + beta_max < 5.0 / 3.0 - 0.05, + "[{label}] beta ~ 5/3 everywhere ({beta_max:.4}) — the S17 correction is untested" + ); + assert!( + beta_max - beta_min > 1e-3, + "[{label}] beta is constant ({beta_min:.5}..{beta_max:.5}) — \ + d(beta)/d(A,T,P,z) carry no signal" + ); + + // (2) `∂β/∂z` is masked wherever S9's clamp saturates. + let z_interior = r + .side_slope + .iter() + .filter(|&&z| z > 0.5 + 1e-3 && z < 50.0 - 1e-3) + .count(); + assert!( + z_interior > 0, + "[{label}] side_slope saturated on [0.5, 50] on every reach — d(beta)/dz is masked off" + ); + + // (3) Cunge X must be strictly interior somewhere, or `gX` is masked to + // zero and the S19 terms could be deleted (the cunge_x.rs trap). + let x_interior = r.x_cunge.iter().filter(|&&x| x > 0.02 && x < 0.48).count(); + assert!( + x_interior > 0, + "[{label}] every x_cunge sits on its own [0, 0.5] clamp — same vacuity trap as cunge_x.rs" + ); + + if !enforce_positivity { + return; + } + + // (4) three-way min branch mix. + let n = r.branches.len() as f32; + let frac = |b: Branch| r.branches.iter().filter(|&&x| x == b).count() as f32 / n; + let (fc, fa, fb) = (frac(Branch::Cunge), frac(Branch::HiA), frac(Branch::HiB)); + println!(" branch mix: cunge={fc:.2} hi_a={fa:.2} hi_b={fb:.2}"); + assert!( + fc > 0.05 && fa > 0.05 && fb > 0.05, + "[{label}] fixture is vacuous: branch mix cunge={fc:.2} hi_a={fa:.2} hi_b={fb:.2}" + ); + + // (5) the K floor must BIND on some reaches and not on others. + let f_floored = r.floored.iter().filter(|&&x| x).count() as f32 / n; + println!(" K floored fraction: {f_floored:.2}"); + assert!( + (0.05..0.95).contains(&f_floored), + "[{label}] fixture must STRADDLE the K floor, got {f_floored:.2}" + ); + + // The clamp must actually be doing work: with it ON no coefficient may be + // negative, and with it OFF the same fixture must violate both bounds. + let on = run_chain::(&f, true); + let off = run_chain::(&f, false); + let min_of = |v: &[f32]| v.iter().cloned().fold(f32::INFINITY, f32::min); + println!( + " min c1/c3: clamp ON ({:e}, {:e}) OFF ({:e}, {:e})", + min_of(&on.c1), + min_of(&on.c3), + min_of(&off.c1), + min_of(&off.c3) + ); + assert!( + min_of(&on.c1) >= 0.0 && min_of(&on.c3) >= 0.0, + "[{label}] clamp ON still produced a negative coefficient" + ); + assert!( + min_of(&off.c1) < 0.0 && min_of(&off.c3) < 0.0, + "[{label}] fixture must violate BOTH bounds when unclamped — the clamp is untested" + ); + // Mass conservation survives (the reason the clamp targets K/X, not c1/c3). + for i in 0..on.c1.len() { + let s = on.c1[i] + on.c2[i] + on.c3[i]; + assert!( + (s - 1.0).abs() < 1e-5, + "[{label}] reach {i}: c1+c2+c3 = {s} on the clamped path" + ); + } +} + +// =========================================================================== +// Gradient machinery — generic in the inner backend. +// =========================================================================== + +#[derive(Copy, Clone, Debug)] +enum Parent { + N, + QSpatial, + PSpatial, + QT, +} + +impl Parent { + fn name(self) -> &'static str { + match self { + Parent::N => "n", + Parent::QSpatial => "q_spatial", + Parent::PSpatial => "p_spatial", + Parent::QT => "q_t", + } + } +} + +const ALL_PARENTS: [Parent; 4] = [Parent::N, Parent::QSpatial, Parent::PSpatial, Parent::QT]; + +/// FD step as a pure FRACTION of the base value. +/// +/// An ABSOLUTE floor (as the sibling gradchecks use) would be a 2.9% +/// perturbation of `n = 0.035`, large enough to move a reach across the +/// `x_cunge` / `hi_b` min boundary — central differences would then average two +/// different slopes and disagree with the analytical gradient by ~30%, an FD +/// artifact rather than a backward bug. Matches `tests/positivity_clamp.rs`. +const REL_STEP: f32 = 3e-3; +/// FD-vs-analytic tolerance. Same value as the CPU gradchecks; NOT widened for +/// CUDA — see `NOISE_ULPS` for how GPU round-off is accounted for instead. +const REL_TOL: f32 = 5e-3; + +/// Ulps of the loss allowed for accumulated round-off in the weighted sum plus +/// the FD subtraction. +/// +/// `tests/positivity_clamp.rs` uses 16 on CPU. 64 here because the GPU's +/// transcendental ops are looser than glibc's: cubecl lowers `powf(x, y)` to +/// `exp2(y·log2(x))` on hardware SFU instructions with ~2 ulp each, versus +/// libm's ~0.5-1 ulp, and the S1..S17 chain contains 4 `powf`/`powf_scalar`, +/// 2 `sqrt` and 2 `recip`. 4x the CPU allowance is a conservative bound on +/// that (it would cover ~8 ulp of extra error per transcendental). +/// +/// This constant only *widens the band in which a disagreement is forgiven as +/// unmeasurable*; it does NOT relax `REL_TOL` for reaches whose gradient is +/// above the floor, and `compare_grads` fails outright if fewer than 4 reaches +/// remain above it. So it cannot be abused to make a broken backward pass — +/// it would instead trip the `resolved >= 4` power check. +const NOISE_ULPS: f32 = 64.0; + +struct GradTensors { + n: Tensor, 1>, + qsp: Tensor, 1>, + psp: Tensor, 1>, + qt: Tensor, 1>, +} + +/// Per-reach loss weights `w[i] = 1/max(q_next_base[i], 1)`, so every reach +/// contributes O(1) to `loss = Σ w[i]·q_next[i]` instead of the O(180) the +/// downstream reaches would otherwise contribute. Without this, `l₊ − l₋` at a +/// 3e-3 relative step lands at the f32 round-off floor of an O(600) loss. +/// +/// Computed ONCE on the CPU backend and reused for both backends: the weights +/// must be bit-identical across backends or Part B would be comparing two +/// different loss functions. +fn conditioning_weights(enforce_positivity: bool) -> Vec { + let base = run_solve::(&grad_fixture(), enforce_positivity); + base.iter().map(|&q| 1.0 / q.max(1.0)).collect() +} + +#[allow(clippy::too_many_arguments)] +fn run_forward_loss( + cfg: &Config, + pattern: &Arc, + assembler: &AValuesAssembler, + dev: &I::Device, + n_vec: &[f32], + qsp_vec: &[f32], + psp_vec: &[f32], + qt_vec: &[f32], + qpt_vec: &[f32], + length_vec: &[f32], + slope_vec: &[f32], + weights: &[f32], + require_grad_parent: Option, +) -> (Tensor, 1>, GradTensors) +where + I::FloatTensorPrimitive: 'static, + I::Device: 'static, +{ + let mk = |data: &[f32], req: bool| -> Tensor, 1> { + let t: Tensor, 1> = Tensor::from_floats(data, dev); + if req { + t.require_grad() + } else { + t + } + }; + let n_t = mk(n_vec, matches!(require_grad_parent, Some(Parent::N))); + let qsp_t = mk(qsp_vec, matches!(require_grad_parent, Some(Parent::QSpatial))); + let psp_t = mk(psp_vec, matches!(require_grad_parent, Some(Parent::PSpatial))); + let qt_t = mk(qt_vec, matches!(require_grad_parent, Some(Parent::QT))); + let qpt_t = mk(qpt_vec, false); + let length_t = mk(length_vec, false); + let slope_t = mk(slope_vec, false); + let xst_t = mk(&vec![0.3f32; n_vec.len()], false); + + let q_next = timestep_forward::( + cfg, + pattern, + assembler, + n_t.clone(), + qsp_t.clone(), + psp_t.clone(), + qt_t.clone(), + qpt_t.clone(), + length_t, + slope_t, + xst_t, + false, + ); + + let loss = q_next * mk(weights, false); + + ( + loss, + GradTensors { + n: n_t, + qsp: qsp_t, + psp: psp_t, + qt: qt_t, + }, + ) +} + +fn analytical_grad(parent: Parent, enforce_positivity: bool) -> Vec +where + I::FloatTensorPrimitive: 'static, + I::Device: 'static, +{ + analytical_grad_cfg::(&stress_cfg(enforce_positivity), parent, enforce_positivity) +} + +fn analytical_grad_cfg( + cfg: &Config, + parent: Parent, + enforce_positivity: bool, +) -> Vec +where + I::FloatTensorPrimitive: 'static, + I::Device: 'static, +{ + let f = grad_fixture(); + let adj = chain(&f); + let dev = device::(); + let pattern = Arc::new(CsrPattern::from_sparse(&adj)); + let assembler = AValuesAssembler::::new(&pattern, &dev); + let weights = conditioning_weights(enforce_positivity); + + let (loss, parents) = run_forward_loss::( + cfg, &pattern, &assembler, &dev, &f.n, &f.qsp, &f.psp, &f.qt, &f.qpt, &f.length, &f.slope, + &weights, Some(parent), + ); + + let grads = loss.sum().backward(); + let g = match parent { + Parent::N => parents.n.grad(&grads).expect("grad on n"), + Parent::QSpatial => parents.qsp.grad(&grads).expect("grad on q_spatial"), + Parent::PSpatial => parents.psp.grad(&grads).expect("grad on p_spatial"), + Parent::QT => parents.qt.grad(&grads).expect("grad on q_t"), + }; + g.into_data().convert::().into_vec::().unwrap() +} + +/// `(fd_grad, fd_noise_floor)`. The second vector is the smallest gradient +/// difference f32 central differences can resolve at each reach's step size: +/// `NOISE_ULPS · ulp(loss) / (2·eps_i)`. Anything below it is unmeasurable, not +/// evidence about the backward. +fn fd_grad( + parent: Parent, + enforce_positivity: bool, +) -> (Vec, Vec) +where + I::FloatTensorPrimitive: 'static, + I::Device: 'static, +{ + fd_grad_cfg::(&stress_cfg(enforce_positivity), parent, enforce_positivity) +} + +fn fd_grad_cfg( + cfg: &Config, + parent: Parent, + enforce_positivity: bool, +) -> (Vec, Vec) +where + I::FloatTensorPrimitive: 'static, + I::Device: 'static, +{ + let f = grad_fixture(); + let adj = chain(&f); + let dev = device::(); + let pattern = Arc::new(CsrPattern::from_sparse(&adj)); + let assembler = AValuesAssembler::::new(&pattern, &dev); + let weights = conditioning_weights(enforce_positivity); + + let eval_loss = |n: &[f32], qsp: &[f32], psp: &[f32], qt: &[f32]| -> f32 { + let (loss, _) = run_forward_loss::( + cfg, &pattern, &assembler, &dev, n, qsp, psp, qt, &f.qpt, &f.length, &f.slope, + &weights, None, + ); + loss.sum().into_data().convert::().into_vec::().unwrap()[0] + }; + let loss_mag = eval_loss(&f.n, &f.qsp, &f.psp, &f.qt).abs().max(1.0); + + let n_reach = f.length.len(); + let mut grad = vec![0.0f32; n_reach]; + let mut noise = vec![0.0f32; n_reach]; + for i in 0..n_reach { + let (mut pn, mut pq, mut pp, mut pt) = + (f.n.clone(), f.qsp.clone(), f.psp.clone(), f.qt.clone()); + let (mut mn, mut mq, mut mp, mut mt) = + (f.n.clone(), f.qsp.clone(), f.psp.clone(), f.qt.clone()); + let (plus, minus, base) = match parent { + Parent::N => (&mut pn, &mut mn, &f.n), + Parent::QSpatial => (&mut pq, &mut mq, &f.qsp), + Parent::PSpatial => (&mut pp, &mut mp, &f.psp), + Parent::QT => (&mut pt, &mut mt, &f.qt), + }; + let eps = REL_STEP * base[i].abs(); + assert!(eps > 0.0, "reach {i}: zero FD step (base = {})", base[i]); + plus[i] = base[i] + eps; + minus[i] = base[i] - eps; + grad[i] = (eval_loss(&pn, &pq, &pp, &pt) - eval_loss(&mn, &mq, &mp, &mt)) / (2.0 * eps); + noise[i] = NOISE_ULPS * loss_mag * f32::EPSILON / (2.0 * eps); + } + (grad, noise) +} + +/// `(worst_rel, worst_abs, resolved, pass)`. A reach passes when it either +/// agrees to `REL_TOL` or disagrees by less than what f32 central differences +/// can resolve at that step size — the second clause is a statement about +/// measurability, not a relaxed tolerance. +/// +/// Split out from `compare_grads` so the negative control can read the verdict +/// without the panic. +fn gradcheck_verdict( + name: &str, + analytical: &[f32], + fd: &[f32], + noise: &[f32], + verbose: bool, +) -> (f32, f32, usize, bool) { + assert_eq!(analytical.len(), fd.len()); + if verbose { + println!("--- {name} ---"); + } + let (mut worst_rel, mut worst_abs) = (0.0f32, 0.0f32); + let mut resolved = 0usize; + let mut pass = true; + for i in 0..analytical.len() { + let (a, d) = (analytical[i], fd[i]); + let abs_diff = (a - d).abs(); + let rel_diff = abs_diff / a.abs().max(d.abs()).max(1e-12); + let informative = a.abs().max(d.abs()) > noise[i]; + if informative { + resolved += 1; + worst_abs = worst_abs.max(abs_diff); + worst_rel = worst_rel.max(rel_diff); + } + if !(rel_diff < REL_TOL || abs_diff < noise[i]) { + pass = false; + } + if verbose { + println!( + " [{i}] analytical={a:.6e} fd={d:.6e} abs={abs_diff:.3e} rel={rel_diff:.3e} \ + fd_noise={:.3e}{}", + noise[i], + if informative { "" } else { " (below FD resolution)" } + ); + } + } + if verbose { + println!(" resolved reaches: {resolved}/{}", analytical.len()); + println!(" worst abs={worst_abs:.3e} worst rel={worst_rel:.3e}"); + } + (worst_rel, worst_abs, resolved, pass) +} + +fn compare_grads(name: &str, analytical: &[f32], fd: &[f32], noise: &[f32]) { + let (worst_rel, worst_abs, resolved, pass) = + gradcheck_verdict(name, analytical, fd, noise, true); + assert!( + resolved >= 4, + "{name}: only {resolved} reaches are above the FD noise floor — \ + the gradcheck has no power left" + ); + assert!( + pass, + "{name}: gradcheck failed (worst rel={worst_rel:.3e}, abs={worst_abs:.3e})" + ); +} + +// =========================================================================== +// Part A — native finite-difference gradcheck ON CUDA. +// +// The strongest evidence in this file: the analytical CUDA gradient is checked +// against central differences of the CUDA forward. This tests CORRECTNESS, not +// merely agreement with the CPU backend (two backends can be identically +// wrong, e.g. if a mask were inverted in shared source). +// +// Two configs cover all three physics patterns: +// * enforce_positivity=false -> beta celerity (S17) + Cunge X (S19) +// * enforce_positivity=true -> the above + the positivity clamp (S18'/S19') +// =========================================================================== + +fn cuda_gradcheck(parent: Parent, enforce_positivity: bool) { + let label = if enforce_positivity { + "beta + cunge-X + positivity clamp" + } else { + "beta + cunge-X" + }; + assert_fixture_exercises_everything::("cuda", enforce_positivity); + let (fd, noise) = fd_grad::(parent, enforce_positivity); + let analytic = analytical_grad::(parent, enforce_positivity); + compare_grads( + &format!("CUDA FD gradcheck: {} [{label}]", parent.name()), + &analytic, + &fd, + &noise, + ); +} + +#[test] +fn cuda_gradcheck_beta_and_cunge_x_n() { + cuda_gradcheck(Parent::N, false); +} + +#[test] +fn cuda_gradcheck_beta_and_cunge_x_q_spatial() { + cuda_gradcheck(Parent::QSpatial, false); +} + +#[test] +fn cuda_gradcheck_beta_and_cunge_x_p_spatial() { + cuda_gradcheck(Parent::PSpatial, false); +} + +/// `q_t` carries the Cunge `∂X/∂Q` term on top of its pre-existing S25 RHS, +/// S24 SpMV and S2 depth paths — the most heavily multiplexed parent. +#[test] +fn cuda_gradcheck_beta_and_cunge_x_q_t() { + cuda_gradcheck(Parent::QT, false); +} + +#[test] +fn cuda_gradcheck_positivity_clamp_n() { + cuda_gradcheck(Parent::N, true); +} + +#[test] +fn cuda_gradcheck_positivity_clamp_q_spatial() { + cuda_gradcheck(Parent::QSpatial, true); +} + +#[test] +fn cuda_gradcheck_positivity_clamp_p_spatial() { + cuda_gradcheck(Parent::PSpatial, true); +} + +/// Under the clamp, `q_t`'s Cunge path is switched OFF on every reach where +/// `hi_a`/`hi_b` win the three-way min — the mask cascade B19' is at its most +/// load-bearing here. +#[test] +fn cuda_gradcheck_positivity_clamp_q_t() { + cuda_gradcheck(Parent::QT, true); +} + +// =========================================================================== +// Part B — CUDA-vs-CPU analytic gradient parity. +// =========================================================================== + +/// Cross-backend relative tolerance for the ANALYTIC gradient. +/// +/// Derived, not tuned: +/// +/// * The two backends run the SAME source, so any difference is pure op-level +/// round-off. The gradient of one reach is a product/sum of roughly 25 f32 +/// chain-rule factors. Under linear error propagation with ≤1 ulp +/// (`f32::EPSILON ≈ 1.2e-7`) per elementary op and no catastrophic +/// cancellation — verified: every component compared here is O(1e-3)..O(25), +/// none is a near-cancelling difference — the worst-case accumulated relative +/// error is bounded by `25 · 1.2e-7 ≈ 3e-6`. +/// * `1e-5` is ~3x that analytic bound, which absorbs the extra slack in the +/// two transcendentals the BACKWARD adds on top of the forward (`ratio.log()` +/// in B6, `powf` in B15): cubecl lowers these onto the SFU (~2 ulp) whereas +/// `NdArray` calls libm (~0.5-1 ulp). +/// * MEASURED on this fixture: worst = 2.784e-7 over all four parents and both +/// configs (and exactly 0.0 for two of the four). So the tolerance sits ~36x +/// above observation — enough headroom for a driver or architecture change, +/// and 36x is small enough that it is still an assertion rather than a +/// formality. `negative_control_*` below proves it: a 1e-4 relative +/// perturbation — 360x smaller than any structural error could be — trips it. +/// * Sharpness: a wrong branch mask, a dropped `∂β/∂z` term or an inverted +/// K-floor mask changes the gradient by O(10%)-O(100%) on the affected +/// reaches, i.e. 1e4-1e5 times this tolerance. +const XBACKEND_REL_TOL: f32 = 1e-5; + +/// Reaches whose |gradient| is below this (relative to the largest component) +/// are excluded from the relative comparison — a relative metric on a component +/// that is ~0 on both backends is meaningless, and an ABSOLUTE check catches +/// the real failure there (one backend producing 0 where the other does not). +const XBACKEND_REL_FLOOR: f32 = 1e-6; + +/// Worst relative disagreement between two gradient vectors, restricted to the +/// components that carry signal. Also asserts the ZERO PATTERN matches: a +/// component that is exactly 0 on one backend and not the other is a masking +/// disagreement, never round-off. +/// +/// Split out from `compare_backends` so the negative controls can call the +/// metric without the tolerance assertion. +fn worst_backend_rel(name: &str, cpu: &[f32], gpu: &[f32], verbose: bool) -> f32 { + assert_eq!(cpu.len(), gpu.len()); + let scale = cpu + .iter() + .chain(gpu.iter()) + .fold(0.0f32, |m, v| m.max(v.abs())); + if verbose { + println!("--- {name} (scale = {scale:.6e}) ---"); + } + let mut worst_rel = 0.0f32; + let mut compared = 0usize; + for i in 0..cpu.len() { + let abs_diff = (cpu[i] - gpu[i]).abs(); + let denom = cpu[i].abs().max(gpu[i].abs()); + let rel = abs_diff / denom.max(1e-30); + let significant = denom > XBACKEND_REL_FLOOR * scale; + // A component that is EXACTLY zero on one backend and not the other is + // a masking disagreement, not round-off — always a failure. + assert!( + (cpu[i] == 0.0) == (gpu[i] == 0.0), + "{name} reach {i}: gradient is zero on exactly one backend \ + (cpu={:e}, gpu={:e}) — a mask disagrees across backends", + cpu[i], + gpu[i] + ); + if significant { + compared += 1; + worst_rel = worst_rel.max(rel); + } + if verbose { + println!( + " [{i}] cpu={:.7e} gpu={:.7e} abs={abs_diff:.3e} rel={rel:.3e}{}", + cpu[i], + gpu[i], + if significant { "" } else { " (below relative floor)" } + ); + } + } + if verbose { + println!(" compared {compared}/{} components, worst rel = {worst_rel:.3e}", cpu.len()); + } + assert!( + compared >= 4, + "{name}: only {compared} components are above the relative floor — no power left" + ); + worst_rel +} + +fn compare_backends(name: &str, cpu: &[f32], gpu: &[f32]) -> f32 { + let worst_rel = worst_backend_rel(name, cpu, gpu, true); + assert!( + worst_rel < XBACKEND_REL_TOL, + "{name}: CUDA/CPU analytic gradient parity failed, worst rel = {worst_rel:.3e} \ + (tol {XBACKEND_REL_TOL:.1e})" + ); + worst_rel +} + +fn parity_for(enforce_positivity: bool) -> f32 { + assert_fixture_exercises_everything::("cpu", enforce_positivity); + assert_fixture_exercises_everything::("cuda", enforce_positivity); + let mut worst = 0.0f32; + for p in ALL_PARENTS { + let cpu = analytical_grad::(p, enforce_positivity); + let gpu = analytical_grad::(p, enforce_positivity); + worst = worst.max(compare_backends( + &format!( + "d(loss)/d({}) [enforce_positivity={enforce_positivity}]", + p.name() + ), + &cpu, + &gpu, + )); + } + println!("WORST cross-backend rel (enforce_positivity={enforce_positivity}) = {worst:.3e}"); + worst +} + +#[test] +fn analytic_gradients_match_across_backends_beta_and_cunge_x() { + parity_for(false); +} + +#[test] +fn analytic_gradients_match_across_backends_positivity_clamp() { + parity_for(true); +} + +/// The production CUDA configuration is `sparse_solver: cuda`, which swaps the +/// host-side forward substitution for the cuSPARSE triangular solve — a +/// DIFFERENT `Backward` impl in `src/sparse/` sitting directly downstream of +/// the new physics terms. Everything above runs with the default +/// `sparse_solver: cpu` (`Config::default()`), so without this test the new +/// backwards would still be unverified in the combination that actually ships. +/// +/// Checked two ways: analytic-vs-FD **entirely inside the cuSPARSE path** +/// (so a solver-side gradient error cannot be cancelled by a matching forward +/// error), and analytic-vs-analytic against the host-solve path. +/// +/// The cuSPARSE path is provably taken rather than silently falling back: +/// `dispatch::effective_use_cuda::` degrades to CPU on exactly one +/// condition — `B != Cuda` — which is a compile-time-known `TypeId` +/// comparison, true here; and `cusparse_forward` itself has no fallback (it +/// `.expect()`s every cuSPARSE call). Note the observed cuSPARSE-vs-host +/// gradients are BIT-identical on this fixture: the chain has one off-diagonal +/// per row, so forward substitution reduces to the same `b[i] − a·x[i−1]` in +/// both implementations. The FD half of this test is therefore the one +/// carrying the discriminating power. +#[test] +fn cusparse_solver_path_carries_the_same_gradients() { + let mut cfg = stress_cfg(true); + cfg.params.sparse_solver = ddrs::config::SparseSolver::Cuda; + + for p in ALL_PARENTS { + let (fd, noise) = fd_grad_cfg::(&cfg, p, true); + let analytic = analytical_grad_cfg::(&cfg, p, true); + compare_grads( + &format!("cuSPARSE FD gradcheck: {} [positivity clamp]", p.name()), + &analytic, + &fd, + &noise, + ); + // Same physics, different triangular solve: the two must agree to the + // cross-backend tolerance, since both run on the same device and only + // the solve implementation differs. + let host_solve = analytical_grad::(p, true); + compare_backends( + &format!("cuSPARSE vs host-solve: d(loss)/d({})", p.name()), + &host_solve, + &analytic, + ); + } +} + +/// NEGATIVE CONTROL for Part B. A parity test that passes because the +/// tolerance is loose is worthless, so this pins the tolerance's SHARPNESS: +/// inject a 1e-4 relative error into a single significant component of the +/// CUDA gradient and confirm the harness rejects it. +/// +/// 1e-4 is deliberately far smaller than any structural defect could produce +/// (a dropped or mis-masked backward term moves a reach by 10%-100%), so +/// passing this control means the parity test would catch a real divergence +/// with ~3 orders of magnitude to spare. +#[test] +fn negative_control_backend_parity_rejects_a_tiny_perturbation() { + let cpu = analytical_grad::(Parent::QT, true); + let idx = (0..cpu.len()) + .max_by(|&a, &b| cpu[a].abs().partial_cmp(&cpu[b].abs()).unwrap()) + .unwrap(); + let mut perturbed = cpu.clone(); + perturbed[idx] *= 1.0 + 1e-4; + let worst = worst_backend_rel("negative control (perturbed)", &cpu, &perturbed, false); + println!( + "negative control: perturbing reach {idx} by 1e-4 gives worst rel = {worst:.3e} \ + (tol {XBACKEND_REL_TOL:.1e})" + ); + assert!( + worst >= XBACKEND_REL_TOL, + "the cross-backend tolerance {XBACKEND_REL_TOL:.1e} is too loose to detect a \ + 1e-4 relative perturbation (measured {worst:.3e})" + ); + // And the zero-pattern assertion inside `worst_backend_rel` must fire when + // a mask is dropped on one backend only — the failure mode that matters + // most, since a mask error zeroes a whole reach rather than nudging it. + let mut zeroed = cpu.clone(); + zeroed[idx] = 0.0; + let caught = std::panic::catch_unwind(|| { + worst_backend_rel("negative control (masked)", &cpu, &zeroed, false) + }) + .is_err(); + assert!(caught, "a gradient zeroed on exactly one backend must be rejected"); +} + +/// NEGATIVE CONTROL for Part A. Same idea one level up: perturb the ANALYTIC +/// gradient and confirm the FD comparison rejects it. +/// +/// The perturbation is 2%, chosen to sit just above the `REL_TOL` band while +/// still being an order of magnitude below what a deleted backward term does. +/// It is applied to the reach with the largest gradient, which is also the one +/// furthest above the FD noise floor — i.e. the harness cannot dodge it via +/// the `abs_diff < noise` escape clause. +#[test] +fn negative_control_fd_gradcheck_rejects_a_two_percent_error() { + let (fd, noise) = fd_grad::(Parent::QT, true); + let analytic = analytical_grad::(Parent::QT, true); + let (_, _, _, clean_pass) = + gradcheck_verdict("negative control (clean)", &analytic, &fd, &noise, false); + assert!(clean_pass, "the unperturbed CUDA gradcheck must pass first"); + + let idx = (0..analytic.len()) + .max_by(|&a, &b| analytic[a].abs().partial_cmp(&analytic[b].abs()).unwrap()) + .unwrap(); + let mut broken = analytic.clone(); + broken[idx] *= 1.02; + let (worst, _, _, pass) = + gradcheck_verdict("negative control (broken)", &broken, &fd, &noise, false); + println!( + "negative control: 2% error on reach {idx} gives worst rel = {worst:.3e} — \ + gradcheck pass = {pass}" + ); + assert!( + !pass, + "a 2% analytic-gradient error slipped through the FD gradcheck \ + (worst rel = {worst:.3e}, REL_TOL = {REL_TOL:.1e})" + ); +} + +/// Forward parity anchor: if the two backends' FORWARDS disagreed materially, +/// the gradient parity above would be comparing derivatives at two different +/// operating points and its tolerance would be uninterpretable. +#[test] +fn forward_matches_across_backends() { + for enforce in [false, true] { + let f = grad_fixture(); + let cpu = run_solve::(&f, enforce); + let gpu = run_solve::(&f, enforce); + let scale = cpu.iter().fold(0.0f32, |m, v| m.max(v.abs())); + let worst = cpu + .iter() + .zip(&gpu) + .fold(0.0f32, |m, (a, b)| m.max((a - b).abs())); + println!( + "x_sol forward [enforce_positivity={enforce}]: max abs diff = {worst:.3e} \ + (scale {scale:.3e}, rel {:.3e})", + worst / scale + ); + assert!( + worst / scale < 1e-5, + "forward diverges across backends: {worst:.3e} on scale {scale:.3e}" + ); + } +} + +// =========================================================================== +// Part C — CUDA-side sanity of the individual tensor ops the new code uses. +// +// The new forward/backward paths lean on: `min_pair`, `clamp`, `clamp_min`, +// `greater_elem`, `lower_elem`, `lower_equal`, `bool_and`, `bool_not`, +// `mask_fill`, `powf`, `powf_scalar`, `recip`, `sqrt`. Part A/B exercise them +// in composition; this test pins each one INDIVIDUALLY on CUDA so that a +// failure there localises immediately instead of surfacing as a mysterious +// gradient discrepancy. +// +// The values are chosen to hit the exact edges the physics hits: `min_pair` +// with ties, `clamp` saturating on both ends, `greater_elem` / `lower_elem` +// exactly ON the threshold (strict inequality — a `>=` would flip the mask), +// and `powf` with a non-integer exponent. +// =========================================================================== + +#[test] +fn cuda_tensor_ops_used_by_the_new_physics_are_correct() { + let dev = device::(); + let mk = |d: &[f32]| -> Tensor { Tensor::from_floats(d, &dev) }; + let get = |t: Tensor| -> Vec { t.into_data().convert::().into_vec().unwrap() }; + let getb = |t: Tensor| -> Vec { + t.into_data().convert::().into_vec().unwrap() + }; + + // min_pair — including an exact tie (index 2) and negatives (index 3). + let a = mk(&[1.0, 5.0, 3.0, -2.0]); + let b = mk(&[4.0, 2.0, 3.0, -7.0]); + assert_eq!(get(a.clone().min_pair(b.clone())), vec![1.0, 2.0, 3.0, -7.0]); + + // clamp / clamp_min, saturating on both ends. + assert_eq!( + get(mk(&[-1.0, 0.25, 0.9]).clamp(0.0, 0.5)), + vec![0.0, 0.25, 0.5] + ); + assert_eq!(get(mk(&[1.0, 5.0]).clamp_min(2.0)), vec![2.0, 5.0]); + + // greater_elem / lower_elem: STRICT. The value exactly on the threshold + // must be false on both — the K-floor and clamp masks rely on that to zero + // the gradient at a saturated point. + assert_eq!(getb(mk(&[1.0, 2.0, 3.0]).greater_elem(2.0)), vec![false, false, true]); + assert_eq!(getb(mk(&[1.0, 2.0, 3.0]).lower_elem(2.0)), vec![true, false, false]); + + // lower_equal (used by the three-way branch cascade): NON-strict. + assert_eq!( + getb(mk(&[1.0, 2.0, 3.0]).lower_equal(mk(&[2.0, 2.0, 2.0]))), + vec![true, true, false] + ); + + // bool_and / bool_not, the mask algebra of B19'. + let m1 = mk(&[1.0, 1.0, 0.0, 0.0]).greater_elem(0.5); + let m2 = mk(&[1.0, 0.0, 1.0, 0.0]).greater_elem(0.5); + assert_eq!(getb(m1.clone().bool_and(m2.clone())), vec![true, false, false, false]); + assert_eq!(getb(m1.clone().bool_not()), vec![false, false, true, true]); + + // mask_fill — the gradient-zeroing primitive every mask above feeds. + assert_eq!( + get(mk(&[7.0, 8.0, 9.0, 10.0]).mask_fill(m1.bool_not(), 0.0)), + vec![7.0, 8.0, 0.0, 0.0] + ); + + // powf (elementwise exponent), powf_scalar, recip, sqrt — checked against + // f64 references at tolerances tight enough to catch a wrong op, loose + // enough for SFU-precision transcendentals. + let p = get(mk(&[2.0, 9.0, 100.0]).powf(mk(&[0.5, 0.5, 0.6]))); + for (got, want) in p.iter().zip([2f64.sqrt(), 3.0, 100f64.powf(0.6)]) { + assert!( + (*got as f64 - want).abs() / want < 1e-5, + "powf: got {got}, want {want}" + ); + } + let ps = get(mk(&[8.0, 27.0]).powf_scalar(2.0 / 3.0)); + for (got, want) in ps.iter().zip([4.0f64, 9.0]) { + assert!( + (*got as f64 - want).abs() / want < 1e-5, + "powf_scalar: got {got}, want {want}" + ); + } + let rc = get(mk(&[2.0, 4.0, 1e-3]).recip()); + for (got, want) in rc.iter().zip([0.5f64, 0.25, 1000.0]) { + assert!( + (*got as f64 - want).abs() / want < 1e-6, + "recip: got {got}, want {want}" + ); + } + let sq = get(mk(&[1e-6, 2.0, 1e6]).sqrt()); + for (got, want) in sq.iter().zip([1e-3f64, 2f64.sqrt(), 1e3]) { + assert!( + (*got as f64 - want).abs() / want < 1e-6, + "sqrt: got {got}, want {want}" + ); + } +} + +// =========================================================================== +// Part D — the TRANSITIVE config guard. +// +// `enforce_positivity: true` + `use_cuda_graphs: true` must be impossible. +// Nothing asserts this today: it falls out of two independent validators, and +// which one fires depends on `ddr_match`. +// +// ddr_match omitted (default true) -> validate_enforce_positivity fires +// (enforce_positivity requires !ddr_match) +// ddr_match: false -> validate_ddr_match fires +// (!ddr_match requires !use_cuda_graphs) +// +// If EITHER validator were relaxed the combination would become reachable, and +// the CUDA-graph kernel — which implements DDR's formulation only, by design — +// would silently run the uncorrected forward against the corrected backward. +// This test does not depend on CUDA hardware, but it belongs here: it is the +// wall that makes "the graph kernel need not support the new physics" true. +// =========================================================================== + +fn load_params_yaml(name: &str, params: &str) -> ddrs::data::error::Result { + let yaml = format!( + "mode: training\ngeodataset: merit\nseed: 1\nnp_seed: 1\nparams:\n{params}" + ); + let path = std::env::temp_dir().join(name); + std::fs::write(&path, yaml).unwrap(); + Config::from_yaml_file(&path) +} + +#[test] +fn enforce_positivity_can_never_reach_the_cuda_graph_path() { + // Case 1: ddr_match left at its default (true). The enforce_positivity + // validator rejects before cuda_graphs is even relevant. + let err = load_params_yaml( + "ddrs_ep_graphs_default_ddr_match.yaml", + " enforce_positivity: true\n use_cuda_graphs: true\n", + ) + .expect_err("enforce_positivity + cuda_graphs (default ddr_match) must be rejected"); + let msg = format!("{err}"); + println!("case 1 rejection: {msg}"); + assert!( + msg.contains("enforce_positivity") && msg.contains("ddr_match"), + "expected the enforce_positivity/ddr_match conflict, got: {msg}" + ); + + // Case 2: ddr_match: false, which is the ONLY way enforce_positivity can + // load at all. Now the ddr_match validator rejects the graphs. + let err = load_params_yaml( + "ddrs_ep_graphs_ddr_match_false.yaml", + " ddr_match: false\n enforce_positivity: true\n use_cuda_graphs: true\n", + ) + .expect_err("ddr_match:false + enforce_positivity + cuda_graphs must be rejected"); + let msg = format!("{err}"); + println!("case 2 rejection: {msg}"); + assert!( + msg.contains("ddr_match") && msg.contains("use_cuda_graphs"), + "expected the ddr_match/cuda_graphs conflict, got: {msg}" + ); + + // Case 3: ddr_match: true + enforce_positivity: true, graphs OFF. Still + // rejected — the clamp changes forward output and would break the DDR + // sandbox ABSOLUTE MATCH (invariant 1). + let err = load_params_yaml( + "ddrs_ep_ddr_match_true.yaml", + " ddr_match: true\n enforce_positivity: true\n use_cuda_graphs: false\n", + ) + .expect_err("ddr_match:true + enforce_positivity must be rejected"); + println!("case 3 rejection: {err}"); + + // NON-VACUITY: the guard must not be rejecting everything. The one legal + // corner — corrected physics, clamp on, graphs off — must LOAD, and must + // land with the flags actually set (a validator that silently cleared + // `enforce_positivity` would pass every assertion above). + let cfg = load_params_yaml( + "ddrs_ep_legal.yaml", + " ddr_match: false\n enforce_positivity: true\n use_cuda_graphs: false\n", + ) + .expect("ddr_match:false + enforce_positivity:true + cuda_graphs:false must load"); + assert!(!cfg.params.ddr_match); + assert!(cfg.params.enforce_positivity); + assert!(!cfg.params.use_cuda_graphs); + + // And `enforce_positivity` must default OFF, so the guard is not merely + // describing a flag nobody can set. + let cfg = load_params_yaml("ddrs_ep_default.yaml", " ddr_match: false\n") + .expect("plain ddr_match:false must load"); + assert!(!cfg.params.enforce_positivity, "enforce_positivity must default to false"); +} diff --git a/tests/cunge_x.rs b/tests/cunge_x.rs new file mode 100644 index 0000000..7070c76 --- /dev/null +++ b/tests/cunge_x.rs @@ -0,0 +1,498 @@ +//! Cunge-derived Muskingum storage weight `X` (`ddr_match: false`). +//! +//! Part 1 — the physics. Cunge picks `X` so the Muskingum scheme's *numerical* +//! diffusion equals the channel's *physical* hydraulic diffusivity: +//! +//! D_num = c·Δx·(0.5 − X) D_phys = Q/(2·B·S₀) +//! => X = clamp( 0.5·(1 − Q/(B·S₀·c·Δx)), 0, 0.5 ) +//! +//! A constant `X` (DDR / `ddr_match: true`; `forward.rs` supplies 0.3) severs +//! that link — it is Cunge-optimal only on the measure-zero locus where +//! `Q/(B·S₀·c·Δx) == 0.4`, and over-diffuses by ~30x at the CONUS-median +//! operating point (documented median `D_num/D_phys` = 28x). +//! +//! Part 2 — the gradcheck. `X` now depends on `Q`, `B` (top width) and `c` +//! (celerity), so the analytical backward in `src/routing/mmc_op.rs` gains +//! three hand-derived terms. `q_t` in particular acquires a SECOND gradient +//! path (it already entered the RHS at S25 and the depth chain at S2). +//! Validated against central finite differences below with `mock_cfg()` +//! setting `ddr_match = false`. + +use std::sync::Arc; + +use burn::backend::{Autodiff, NdArray}; +use burn::tensor::Tensor; + +use ddrs::config::Config; +use ddrs::routing::mmc_op::timestep_forward; +use ddrs::sparse::{AValuesAssembler, CsrPattern, SparseAdjacency}; + +// =========================================================================== +// Part 1: the X formula, in f64, independent of any ddrs code. +// =========================================================================== + +/// `X = clamp( 0.5·(1 − Q/(B·S·c·Δx)), 0, 0.5 )`. +fn cunge_x(q: f64, b: f64, s: f64, c: f64, dx: f64) -> f64 { + (0.5 * (1.0 - q / (b * s * c * dx))).clamp(0.0, 0.5) +} + +/// Muskingum numerical diffusion `D_num = c·Δx·(0.5 − X)`. +fn d_num(c: f64, dx: f64, x: f64) -> f64 { + c * dx * (0.5 - x) +} + +/// Physical hydraulic diffusivity `D_phys = Q/(2·B·S₀)`. +fn d_phys(q: f64, b: f64, s: f64) -> f64 { + q / (2.0 * b * s) +} + +/// A CONUS-median-ish reach: `W = Q/(B·S·c·Δx) ≈ 0.0135` → Cunge `X ≈ 0.493`, +/// reproducing the documented "`X ≈ 0.49` almost everywhere" measurement. +/// +/// Self-consistency check on Q: `c = 1.4` under `β = 5/3` implies `v ≈ 0.84` +/// m/s, so `A = Q/v ≈ 11.9 m²`, i.e. ~0.3 m mean depth over a 40 m top width — +/// a plausible Manning `n ≈ 0.024` at `S = 2e-3`. +const MEDIAN_REACH: (f64, f64, f64, f64, f64) = (10.0, 40.0, 2e-3, 1.4, 6598.0); + +#[test] +fn cunge_x_matches_numerical_to_physical_diffusion() { + // Unclamped regime: the identity D_num == D_phys is exact by construction. + for &(q, b, s, c, dx) in &[ + MEDIAN_REACH, + (300.0_f64, 40.0_f64, 2e-3_f64, 1.4_f64, 6598.0_f64), + (5.0, 12.0, 5e-4, 0.9, 2100.0), + ] { + let x = cunge_x(q, b, s, c, dx); + assert!(x > 0.0 && x < 0.5, "operating point must be unclamped, got X={x}"); + let (dn, dp) = (d_num(c, dx, x), d_phys(q, b, s)); + assert!( + (dn / dp - 1.0).abs() < 1e-9, + "D_num={dn} != D_phys={dp} at (q={q}, b={b}, s={s}, c={c}, dx={dx})" + ); + } +} + +#[test] +fn constant_x_030_over_diffuses_at_the_median_reach() { + // Regression guard on the magnitude of the defect this task fixes. + let (q, b, s, c, dx) = MEDIAN_REACH; + let ratio = d_num(c, dx, 0.3) / d_phys(q, b, s); + assert!( + ratio > 2.0, + "expected heavy over-diffusion from constant X=0.3, got {ratio:.2}x" + ); + // The documented CONUS median is 28x; this operating point sits at ~30x. + assert!( + (20.0..40.0).contains(&ratio), + "expected the documented order-30x over-diffusion, got {ratio:.2}x" + ); +} + +#[test] +fn constant_x_030_is_correct_only_on_a_measure_zero_locus() { + // D_num(0.3)/D_phys = 0.4/W, so the constant is Cunge-optimal iff W = 0.4. + // That is why (300, 40, 2e-3, 1.4, 6598) — W = 0.406 — makes X=0.3 look + // fine: it is one point on that locus, NOT a representative reach. + let (b, s, c, dx) = (40.0_f64, 2e-3_f64, 1.4_f64, 6598.0_f64); + let q_tuned = 0.4 * b * s * c * dx; + let ratio = d_num(c, dx, 0.3) / d_phys(q_tuned, b, s); + assert!( + (ratio - 1.0).abs() < 1e-9, + "X=0.3 must be exact at W=0.4, got {ratio}" + ); + // Move Q by one order of magnitude in either direction and it breaks by + // the same order (D_num(0.3)/D_phys = 0.4/W is exactly inverse in Q). + for &f in &[0.1_f64, 10.0] { + let r = d_num(c, dx, 0.3) / d_phys(q_tuned * f, b, s); + assert!( + r.max(1.0 / r) > 9.0, + "X=0.3 should be ~{f}x wrong at {f}x the tuned Q, got {r}" + ); + } +} + +#[test] +fn cunge_x_clamps_into_zero_half() { + let (_, b, s, c, dx) = MEDIAN_REACH; + // Huge Q -> raw X goes negative -> clamped to 0 (pure advection). + assert_eq!(cunge_x(1e9, b, s, c, dx), 0.0); + // Tiny Q -> raw X approaches 0.5 from below, never exceeds it. + let x_tiny = cunge_x(1e-9, b, s, c, dx); + assert!(x_tiny <= 0.5, "X must not exceed 0.5, got {x_tiny}"); + assert!(x_tiny > 0.499, "tiny Q should push X to the 0.5 limit, got {x_tiny}"); + // Sweep: X stays inside [0, 0.5] over 12 decades of Q. + for e in -6..6 { + let x = cunge_x(10f64.powi(e), b, s, c, dx); + assert!((0.0..=0.5).contains(&x), "X={x} out of [0,0.5] at Q=1e{e}"); + } +} + +#[test] +fn muskingum_coefficients_sum_to_one_for_any_x() { + // Mass conservation is independent of X: the three numerators sum to the + // shared denominator identically. + for &x in &[0.0_f64, 0.3, 0.49, 0.5] { + let (k, dt) = (3295.0_f64, 3600.0_f64); + let denom = 2.0 * k * (1.0 - x) + dt; + let c1 = (dt - 2.0 * k * x) / denom; + let c2 = (dt + 2.0 * k * x) / denom; + let c3 = (2.0 * k * (1.0 - x) - dt) / denom; + assert!((c1 + c2 + c3 - 1.0).abs() < 1e-12, "x={x}"); + } +} + +#[test] +fn cunge_x_narrows_the_non_negative_coefficient_window() { + // Non-negative Muskingum coefficients require 2X <= Cr <= 2(1-X). + // X=0.3 -> [0.6, 1.4]; Cunge X~0.49 -> ~[0.98, 1.02]. This is the + // documented reason Task 5 (Courant sub-stepping) exists. + let (q, b, s, c, dx) = MEDIAN_REACH; + let x = cunge_x(q, b, s, c, dx); + let (lo, hi) = (2.0 * x, 2.0 * (1.0 - x)); + assert!(hi - lo < 0.1, "Cunge window [{lo}, {hi}] should be narrow"); + assert!( + hi - lo < (2.0 * (1.0 - 0.3) - 2.0 * 0.3) / 10.0, + "Cunge window must be >10x narrower than the X=0.3 window" + ); +} + +// =========================================================================== +// Part 2: gradcheck of the Cunge-X backward (`ddr_match: false`). +// Scaffolding mirrors `tests/celerity_beta.rs` Part 2 (itself from +// `tests/sp8_gradcheck.rs`). `mock_cfg()` sets `ddr_match = false`, so BOTH +// the trapezoidal celerity (S17, Task 3) and the Cunge X (S19, this task) are +// live — the two corrections are coupled through `celerity`, and X's +// `∂X/∂c` term feeds the very `gcelerity` accumulator B17 consumes. +// =========================================================================== + +type I = NdArray; +type AB = Autodiff; + +const N: usize = 4; +const EPS: f32 = 1e-3; +const REL_TOL: f32 = 5e-3; +const ABS_TOL: f32 = 1e-4; + +#[derive(Copy, Clone, Debug)] +enum Parent { + N, + QSpatial, + PSpatial, + QT, +} + +/// NOTE the 5000 m reach length — it is load-bearing, not cosmetic. +/// +/// `tests/celerity_beta.rs` uses 1000 m. At that length the fixture's +/// `W = Q/(B·S·c·L) ≈ 1.6`, so Cunge `X` saturates at the lower clamp on every +/// reach, `gX` is masked to zero, and the whole Part-2 gradcheck would pass +/// with the S19 backward terms DELETED (verified: it does). 5000 m puts +/// `W ≈ 0.32..0.51`, i.e. `X ≈ 0.25..0.34`, strictly interior — so the three +/// new terms are actually exercised. `cunge_x_operating_point_is_unclamped` +/// below is the standing guard on this. +fn linear_chain_sparse() -> SparseAdjacency { + 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![5000.0; N], vec![0.001; N]) +} + +fn mock_cfg() -> Config { + let mut cfg = Config::default(); + // The whole point of this file: exercise the corrected-physics branch. + cfg.params.ddr_match = false; + cfg.params.parameter_ranges.n = [0.01, 0.1]; + cfg.params.parameter_ranges.q_spatial = [0.1, 0.9]; + cfg.params.parameter_ranges.p_spatial = [1.0, 200.0]; + cfg.params.attribute_minimums.velocity = 0.1; + cfg.params.attribute_minimums.depth = 0.01; + cfg.params.attribute_minimums.discharge = 0.001; + cfg.params.attribute_minimums.bottom_width = 0.1; + cfg.params.attribute_minimums.slope = 0.001; + cfg.params.defaults.insert("p_spatial".to_string(), 1.0); + cfg.params.log_space_parameters = vec![]; + cfg +} + +fn default_inputs() -> (Vec, Vec, Vec, Vec, Vec) { + let n_vec = vec![0.035f32; N]; + let qsp_vec = vec![0.4f32; N]; + let psp_vec = vec![20.0f32; N]; + let qt_vec = vec![100.0f32, 120.0, 140.0, 160.0]; + let qpt_vec = vec![10.0f32, 12.0, 14.0, 16.0]; + (n_vec, qsp_vec, psp_vec, qt_vec, qpt_vec) +} + +struct GradTensors { + n: Tensor, + qsp: Tensor, + psp: Tensor, + qt: Tensor, +} + +#[allow(clippy::too_many_arguments)] +fn run_forward_loss( + cfg: &Config, + pattern: &Arc, + assembler: &AValuesAssembler, + device: &::Device, + n_vec: &[f32], + qsp_vec: &[f32], + psp_vec: &[f32], + qt_vec: &[f32], + qpt_vec: &[f32], + length_vec: &[f32], + slope_vec: &[f32], + x_storage_vec: &[f32], + require_grad_parent: Option, +) -> (Tensor, GradTensors) { + let mk = |data: &[f32], req: bool| -> Tensor { + let t: Tensor = Tensor::from_floats(data, device); + if req { t.require_grad() } else { t } + }; + let n_t = mk(n_vec, matches!(require_grad_parent, Some(Parent::N))); + let qsp_t = mk(qsp_vec, matches!(require_grad_parent, Some(Parent::QSpatial))); + let psp_t = mk(psp_vec, matches!(require_grad_parent, Some(Parent::PSpatial))); + let qt_t = mk(qt_vec, matches!(require_grad_parent, Some(Parent::QT))); + let qpt_t = mk(qpt_vec, false); + let length_t = mk(length_vec, false); + let slope_t = mk(slope_vec, false); + let xst_t = mk(x_storage_vec, false); + + let q_next = timestep_forward::( + cfg, + pattern, + assembler, + n_t.clone(), + qsp_t.clone(), + psp_t.clone(), + qt_t.clone(), + qpt_t.clone(), + length_t, + slope_t, + xst_t, + false, + ); + + ( + q_next, + GradTensors { + n: n_t, + qsp: qsp_t, + psp: psp_t, + qt: qt_t, + }, + ) +} + +fn compute_analytical_grad(parent: Parent) -> Vec { + let cfg = mock_cfg(); + let adj = linear_chain_sparse(); + let device = ::Device::default(); + let pattern = Arc::new(CsrPattern::from_sparse(&adj)); + let assembler = AValuesAssembler::::new(&pattern, &device); + + let (n_vec, qsp_vec, psp_vec, qt_vec, qpt_vec) = default_inputs(); + let length_vec = adj.length_m.clone(); + let slope_vec = adj.slope.clone(); + let x_storage_vec = vec![0.3f32; N]; + + let (q_next, parents) = run_forward_loss( + &cfg, + &pattern, + &assembler, + &device, + &n_vec, + &qsp_vec, + &psp_vec, + &qt_vec, + &qpt_vec, + &length_vec, + &slope_vec, + &x_storage_vec, + Some(parent), + ); + + let loss = q_next.sum(); + let grads = loss.backward(); + + let g = match parent { + Parent::N => parents.n.grad(&grads).expect("grad on n"), + Parent::QSpatial => parents.qsp.grad(&grads).expect("grad on q_spatial"), + Parent::PSpatial => parents.psp.grad(&grads).expect("grad on p_spatial"), + Parent::QT => parents.qt.grad(&grads).expect("grad on q_t"), + }; + g.into_data().to_vec::().unwrap() +} + +fn compute_fd_grad(parent: Parent) -> Vec { + let cfg = mock_cfg(); + let adj = linear_chain_sparse(); + let device = ::Device::default(); + let pattern = Arc::new(CsrPattern::from_sparse(&adj)); + let assembler = AValuesAssembler::::new(&pattern, &device); + + let (n_vec, qsp_vec, psp_vec, qt_vec, qpt_vec) = default_inputs(); + let length_vec = adj.length_m.clone(); + let slope_vec = adj.slope.clone(); + let x_storage_vec = vec![0.3f32; N]; + + let eval_loss = |n: &[f32], qsp: &[f32], psp: &[f32], qt: &[f32], qpt: &[f32]| -> f32 { + let (q_next, _) = run_forward_loss( + &cfg, + &pattern, + &assembler, + &device, + n, + qsp, + psp, + qt, + qpt, + &length_vec, + &slope_vec, + &x_storage_vec, + None, + ); + let v: Vec = q_next.sum().into_data().to_vec::().unwrap(); + v[0] + }; + + let mut grad = vec![0.0f32; N]; + for i in 0..N { + let mut plus_n = n_vec.clone(); + let mut plus_qsp = qsp_vec.clone(); + let mut plus_psp = psp_vec.clone(); + let mut plus_qt = qt_vec.clone(); + let mut minus_n = n_vec.clone(); + let mut minus_qsp = qsp_vec.clone(); + let mut minus_psp = psp_vec.clone(); + let mut minus_qt = qt_vec.clone(); + + let (plus, minus, base) = match parent { + Parent::N => (&mut plus_n, &mut minus_n, &n_vec), + Parent::QSpatial => (&mut plus_qsp, &mut minus_qsp, &qsp_vec), + Parent::PSpatial => (&mut plus_psp, &mut minus_psp, &psp_vec), + Parent::QT => (&mut plus_qt, &mut minus_qt, &qt_vec), + }; + let eps = (EPS * base[i].abs()).max(EPS); + plus[i] = base[i] + eps; + minus[i] = base[i] - eps; + + let l_plus = eval_loss(&plus_n, &plus_qsp, &plus_psp, &plus_qt, &qpt_vec); + let l_minus = eval_loss(&minus_n, &minus_qsp, &minus_psp, &minus_qt, &qpt_vec); + grad[i] = (l_plus - l_minus) / (2.0 * eps); + } + grad +} + +fn compare_grads(name: &str, analytical: &[f32], fd: &[f32]) { + assert_eq!(analytical.len(), fd.len()); + println!("--- {name} ---"); + let mut worst_rel = 0.0f32; + let mut worst_abs = 0.0f32; + for i in 0..analytical.len() { + let a = analytical[i]; + let f = fd[i]; + let abs_diff = (a - f).abs(); + let denom = a.abs().max(f.abs()).max(1e-12); + let rel_diff = abs_diff / denom; + worst_abs = worst_abs.max(abs_diff); + worst_rel = worst_rel.max(rel_diff); + println!(" [{i}] analytical={a:.6e} fd={f:.6e} abs={abs_diff:.3e} rel={rel_diff:.3e}"); + } + println!(" worst abs={worst_abs:.3e} worst rel={worst_rel:.3e}"); + let pass = analytical.iter().zip(fd).all(|(&a, &f)| { + let abs_diff = (a - f).abs(); + let denom = a.abs().max(f.abs()).max(1e-12); + let rel_diff = abs_diff / denom; + rel_diff < REL_TOL || abs_diff < ABS_TOL + }); + assert!( + pass, + "{name}: gradcheck failed (worst rel={worst_rel:.3e}, abs={worst_abs:.3e})" + ); +} + +#[test] +fn gradcheck_cunge_x_n() { + let a = compute_analytical_grad(Parent::N); + let fd = compute_fd_grad(Parent::N); + compare_grads("n (cunge X)", &a, &fd); +} + +#[test] +fn gradcheck_cunge_x_q_spatial() { + let a = compute_analytical_grad(Parent::QSpatial); + let fd = compute_fd_grad(Parent::QSpatial); + compare_grads("q_spatial (cunge X)", &a, &fd); +} + +#[test] +fn gradcheck_cunge_x_p_spatial() { + let a = compute_analytical_grad(Parent::PSpatial); + let fd = compute_fd_grad(Parent::PSpatial); + compare_grads("p_spatial (cunge X)", &a, &fd); +} + +/// The genuinely new gradient path this task adds: `q_t` now enters `X` at +/// S19 on top of the S25 RHS, the S24 SpMV and the S2 depth chain. +#[test] +fn gradcheck_cunge_x_q_t() { + let a = compute_analytical_grad(Parent::QT); + let fd = compute_fd_grad(Parent::QT); + compare_grads("q_t (cunge X)", &a, &fd); +} + +/// Standing guard against a VACUOUS gradcheck. +/// +/// Every S19 gradient term is multiplied by the `[0, 0.5]` clamp mask, so if +/// the fixture's operating point saturates `X`, all four gradcheck tests above +/// pass trivially whether or not the backward terms exist. This test pins the +/// fixture into the interior of the clamp using the top width and celerity the +/// forward chain actually produced. +#[test] +fn cunge_x_operating_point_is_unclamped() { + let cfg = mock_cfg(); + let adj = linear_chain_sparse(); + let device = ::Device::default(); + let pattern = Arc::new(CsrPattern::from_sparse(&adj)); + + let (n_vec, qsp_vec, psp_vec, qt_vec, qpt_vec) = default_inputs(); + let mk = |d: &[f32]| -> Tensor { Tensor::from_floats(d, &device) }; + + let outs = ddrs::routing::mmc_op::__spike_forward_chain_k1_outputs::( + &cfg, + &pattern, + mk(&n_vec), + mk(&qsp_vec), + mk(&psp_vec), + mk(&qt_vec), + mk(&qpt_vec), + mk(&adj.length_m), + mk(&adj.slope), + mk(&vec![0.3f32; N]), + ); + // k1 output order: [depth, top_width, side_slope, bottom_width, hyd_radius, + // velocity_un, velocity_cl, celerity, ...] + let top_width = &outs[1]; + let celerity = &outs[7]; + + for i in 0..N { + let (q, b, s, c, dx) = ( + qt_vec[i] as f64, + top_width[i] as f64, + adj.slope[i] as f64, + celerity[i] as f64, + adj.length_m[i] as f64, + ); + let w = q / (b * s * c * dx); + let x = cunge_x(q, b, s, c, dx); + println!(" [{i}] B={b:.3} c={c:.4} W={w:.4} X={x:.4}"); + assert!( + x > 0.02 && x < 0.48, + "reach {i}: Cunge X={x:.4} (W={w:.4}) is at/near a clamp bound — \ + the gradcheck would be vacuous. Retune the fixture." + ); + } +} diff --git a/tests/ddr_match_flag.rs b/tests/ddr_match_flag.rs new file mode 100644 index 0000000..bebc4f4 --- /dev/null +++ b/tests/ddr_match_flag.rs @@ -0,0 +1,114 @@ +//! `ddr_match` defaults to true so every existing config and the DDR sandbox +//! parity example keep their current behaviour (invariant 1). +use ddrs::config::Config; + +#[test] +fn ddr_match_defaults_to_true() { + let yaml = r#" +mode: training +geodataset: merit +seed: 42 +np_seed: 42 +params: + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] +"#; + let path = std::env::temp_dir().join("ddrs_ddr_match_default_test.yaml"); + std::fs::write(&path, yaml).unwrap(); + let cfg = Config::from_yaml_file(&path).expect("parse"); + assert!(cfg.params.ddr_match, "ddr_match must default to true"); +} + +#[test] +fn ddr_match_can_be_disabled() { + let yaml = r#" +mode: training +geodataset: merit +seed: 42 +np_seed: 42 +params: + ddr_match: false + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] +"#; + let path = std::env::temp_dir().join("ddrs_ddr_match_disabled_test.yaml"); + std::fs::write(&path, yaml).unwrap(); + let cfg = Config::from_yaml_file(&path).expect("parse"); + assert!(!cfg.params.ddr_match); +} + +#[test] +fn enforce_positivity_defaults_to_false() { + let yaml = r#" +mode: training +geodataset: merit +seed: 42 +np_seed: 42 +params: + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] +"#; + let path = std::env::temp_dir().join("ddrs_enforce_positivity_default_test.yaml"); + std::fs::write(&path, yaml).unwrap(); + let cfg = Config::from_yaml_file(&path).expect("parse"); + assert!( + !cfg.params.enforce_positivity, + "enforce_positivity must default to false so existing runs are unchanged" + ); +} + +#[test] +fn enforce_positivity_requires_corrected_physics() { + // ddr_match: true + enforce_positivity: true must be rejected at load: the + // clamp changes K and X, which would break compare_ddr_sandbox. + let yaml = r#" +mode: training +geodataset: merit +seed: 42 +np_seed: 42 +params: + ddr_match: true + enforce_positivity: true + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] +"#; + let path = std::env::temp_dir().join("ddrs_enforce_positivity_reject_test.yaml"); + std::fs::write(&path, yaml).unwrap(); + let err = Config::from_yaml_file(&path).expect_err("must reject"); + let msg = err.to_string(); + assert!( + msg.contains("enforce_positivity"), + "error must name the offending key, got: {msg}" + ); +} + +#[test] +fn enforce_positivity_loads_with_corrected_physics() { + let yaml = r#" +mode: training +geodataset: merit +seed: 42 +np_seed: 42 +params: + ddr_match: false + use_cuda_graphs: false + enforce_positivity: true + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] +"#; + let path = std::env::temp_dir().join("ddrs_enforce_positivity_accept_test.yaml"); + std::fs::write(&path, yaml).unwrap(); + let cfg = Config::from_yaml_file(&path).expect("ddr_match:false + enforce_positivity:true must load"); + assert!(cfg.params.enforce_positivity); + assert!(!cfg.params.ddr_match); +} diff --git a/tests/disagg_enabled.rs b/tests/disagg_enabled.rs new file mode 100644 index 0000000..aa66490 --- /dev/null +++ b/tests/disagg_enabled.rs @@ -0,0 +1,77 @@ +//! `kan_head.disaggregation.enabled` defaults to true so every existing config +//! keeps its behaviour: a bare block enables the head. `enabled: false` strips +//! the block at load, giving the flat repeat-24 (nearest) daily→hourly +//! fallback with the block left inert in YAML — the one-line ablation switch. +//! The section is `deny_unknown_fields`: phantom keys (e.g. the removed +//! `use_precip`) fail loudly instead of silently building a different head. +use ddrs::config::Config; + +const BASE: &str = r#" +mode: training +geodataset: merit +seed: 42 +np_seed: 42 +experiment: + batch_size: 4 + start_time: 1981/10/01 + end_time: 1982/09/30 + epochs: 1 + warmup: 5 +kan_head: + hidden_size: 21 + num_hidden_layers: 2 + input_var_names: [aridity] + learnable_parameters: [n] + disaggregation: + hidden_size: 16 +"#; + +const PARAMS: &str = r#" +params: + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] +"#; + +fn load(name: &str, extra_disagg: &str) -> ddrs::data::error::Result { + let yaml = format!("{BASE}{extra_disagg}{PARAMS}"); + let path = std::env::temp_dir().join(name); + std::fs::write(&path, yaml).unwrap(); + Config::from_yaml_file(&path) +} + +#[test] +fn bare_block_defaults_to_enabled() { + let cfg = load("ddrs_disagg_enabled_default_test.yaml", "").expect("parse"); + assert!( + cfg.kan_head.unwrap().disaggregation.is_some(), + "a bare disaggregation block must enable the head" + ); +} + +#[test] +fn enabled_false_strips_block() { + let cfg = load("ddrs_disagg_enabled_off_test.yaml", " enabled: false\n").expect("parse"); + assert!( + cfg.kan_head.unwrap().disaggregation.is_none(), + "enabled: false must strip the disaggregation block" + ); +} + +#[test] +fn enabled_true_keeps_block() { + let cfg = load("ddrs_disagg_enabled_on_test.yaml", " enabled: true\n").expect("parse"); + assert!(cfg.kan_head.unwrap().disaggregation.is_some()); +} + +#[test] +fn unknown_key_is_rejected() { + let err = load("ddrs_disagg_unknown_key_test.yaml", " use_precip: true\n") + .expect_err("phantom keys in the disaggregation block must fail load"); + let msg = err.to_string(); + assert!( + msg.contains("unknown field") && msg.contains("use_precip"), + "error should name the unknown field, got: {msg}" + ); +} diff --git a/tests/gauge_mass_conservation.rs b/tests/gauge_mass_conservation.rs new file mode 100644 index 0000000..214f2b1 --- /dev/null +++ b/tests/gauge_mass_conservation.rs @@ -0,0 +1,309 @@ +//! End-to-end mass-conservation check on the per-gauge prediction path. +//! +//! A USGS gauge measures ALL drainage area above it, and we do not know where +//! along its reach the gauge physically sits. The Muskingum-Cunge solve at the +//! gauge's own reach already accumulates every upstream contribution PLUS that +//! reach's own lateral inflow (mass conservation), so the gauge's extracted +//! prediction must equal the discharge at the gauge reach — which at steady +//! state equals the sum of `q_prime` over every reach in the subgraph. +//! +//! This is the test that would have caught the `outflow_idx` defect: +//! `collate::compress` returned the gauge's UPSTREAM neighbours instead of the +//! gauge reach itself, silently dropping the gauge reach's own local runoff +//! from every prediction. Real-world symptom: gauge `01457000` (366.8 km², +//! whose own reach is 250.1 km² = 68% of the basin) predicted 1.58 m³/s +//! against an observed 7.60 and a summed-Q' baseline of 7.38 — a constant +//! 0.215× suppression across 15 straight eval years. +//! +//! The defect is preserved (DDR parity) under `params.ddr_match: true` and +//! corrected under `false`, so both directions are asserted here. +//! +//! Network (mirrors 01457000's topology exactly): two headwaters draining into +//! the gauge reach. +//! +//! 0 ──┐ +//! ├──> 2 (gauge) +//! 1 ──┘ +//! +//! `ddr_match: true` sums reaches {0, 1} — 2/3 of the mass. `ddr_match: false` +//! reads reach {2} alone, which carries 3/3. + +use burn::backend::{Autodiff, NdArray}; +use burn::tensor::{Int, Tensor, TensorData}; + +use ddrs::config::Config; +use ddrs::data::collate::{compress, UnionedCoo}; +use ddrs::data::ids::{Comid, Staid}; +use ddrs::routing::{MuskingumCunge, RoutingInputs, SpatialParameters}; +use ddrs::sparse::SparseAdjacency; +use ddrs::training::forward::scatter_add_by_group; + +type I = NdArray; +type AB = Autodiff; + +/// Parent (MERIT) reaches in the fixture network. The *row* count grows when +/// the gauge reach is subdivided, but the mass balance is set by this. +const N: usize = 3; +/// Long enough for a 1 km / 0.1%-slope network at the configured dt to settle. +const T: usize = 64; +/// Constant lateral inflow per reach, m³/s. +const Q_PRIME: f32 = 10.0; + +/// Same knobs as `tests/sp8_gradcheck.rs::mock_cfg` — puts velocity/depth in +/// well-conditioned, non-saturated regimes for this network. +fn mock_cfg(ddr_match: bool) -> Config { + let mut cfg = Config::default(); + cfg.params.parameter_ranges.n = [0.01, 0.1]; + cfg.params.parameter_ranges.q_spatial = [0.1, 0.9]; + cfg.params.parameter_ranges.p_spatial = [1.0, 200.0]; + cfg.params.attribute_minimums.velocity = 0.1; + cfg.params.attribute_minimums.depth = 0.01; + cfg.params.attribute_minimums.discharge = 0.001; + cfg.params.attribute_minimums.bottom_width = 0.1; + cfg.params.attribute_minimums.slope = 0.001; + cfg.params.defaults.insert("p_spatial".to_string(), 1.0); + cfg.params.log_space_parameters = vec![]; + cfg.params.ddr_match = ddr_match; + cfg +} + +/// Two headwaters (0, 1) → gauge reach (2), with the gauge reach split into +/// `pieces` sub-reaches chained in series (reach subdivision, Task 3): +/// +/// ```text +/// 0 ─┐ +/// ├─> 2 → 3 → … → (1 + pieces) (gauge reach; outlet is the last) +/// 1 ─┘ +/// ``` +/// +/// Total gauge-reach length is held at 1000 m regardless of `pieces`, so the +/// steady state is unchanged — the physics is identical, only the discretization +/// differs. Returns the adjacency plus its CONUS-space `parent_offset`. +fn confluence_sparse(pieces: usize) -> (SparseAdjacency, Vec) { + assert!(pieces >= 1); + let n = 2 + pieces; + let mut rows: Vec = vec![2, 2]; + let mut cols: Vec = vec![0, 1]; + for k in 1..pieces { + cols.push((2 + k - 1) as i32); + rows.push((2 + k) as i32); + } + let piece_len = 1000.0 / pieces as f32; + let mut length_m = vec![1000.0_f32; 2]; + length_m.extend(std::iter::repeat_n(piece_len, pieces)); + let parent_offset: Vec = vec![0, 1, 2, n as i32]; + let adj = SparseAdjacency { + n, + values: vec![1.0; rows.len()], + rows, + cols, + length_m, + slope: vec![0.001; n], + parent_offset: Some(parent_offset.clone()), + }; + (adj, parent_offset) +} + +/// Run `collate::compress` over the same topology to get `outflow_idx` from +/// the REAL production code path (not a hand-rolled copy). +fn outflow_idx_from_collate(ddr_match: bool, pieces: usize) -> Vec> { + let gauge_comid = Comid(73005764); + let mut conus_order = vec![Comid(73006562), Comid(73006585)]; + conus_order.extend(std::iter::repeat_n(gauge_comid, pieces)); + let (adj, parent_offset) = confluence_sparse(pieces); + let edges: Vec<(usize, usize)> = adj + .rows + .iter() + .zip(adj.cols.iter()) + .map(|(&r, &c)| (r as usize, c as usize)) + .collect(); + // The subgraph builder resolves a gauge COMID to its parent's LAST row, so + // `gage_idx` is the outlet piece (`cache.rs::resolve_or_build`). + let unioned = UnionedCoo { + edges, + gauges: vec![(Staid::new("01457000"), adj.n - 1, "73005764".to_string())], + }; + compress(&unioned, &conus_order, ddr_match, Some(&parent_offset)) + .expect("compress") + .outflow_idx +} + +/// Route the network to steady state and extract the gauge series exactly as +/// training does. Returns `(gauge_series, all_reach_discharge_row_major_NxT)`. +fn route_and_extract(ddr_match: bool, pieces: usize) -> (Vec, Vec) { + let cfg = mock_cfg(ddr_match); + let device = ::Device::default(); + let (adjacency, _) = confluence_sparse(pieces); + let n = adjacency.n; + + // Constant q_prime on every ROW — exactly what `StreamflowStore::read_window` + // hands back, since every piece of a parent carries the parent's COMID. The + // engine is what divides by the piece count (`mmc.rs`, Task 5). + let q_prime: Tensor = + Tensor::from_data(TensorData::new(vec![Q_PRIME; T * n], [T, n]), &device); + + // Mid-range normalized params → n ≈ 0.055, q_spatial ≈ 0.5, p_spatial ≈ 100. + let mk = |v: f32| -> Tensor { Tensor::from_floats(vec![v; n].as_slice(), &device) }; + let mut engine = MuskingumCunge::::new(cfg, device.clone()); + engine.setup_inputs( + RoutingInputs { + adjacency, + x_storage: mk(0.3), + }, + q_prime, + SpatialParameters { + n: mk(0.5), + q_spatial: mk(0.5), + p_spatial: Some(mk(0.5)), + k_d: None, + d_gw: None, + leakance_factor: None, + impervious_mask: None, + }, + false, + None, + ); + + // (N, T) routed discharge. + let runoff: Tensor = engine.forward().inner(); + + let outflow_idx = outflow_idx_from_collate(ddr_match, pieces); + let flat: Vec = outflow_idx[0].iter().map(|&c| c as i32).collect(); + let groups: Vec = vec![0; flat.len()]; + let flat_t: Tensor = + Tensor::from_data(TensorData::from(flat.as_slice()), &device); + let group_t: Tensor = + Tensor::from_data(TensorData::from(groups.as_slice()), &device); + let gauge_q: Tensor = scatter_add_by_group(runoff.clone(), flat_t, group_t, 1); + + let series: Vec = gauge_q + .into_data() + .to_vec::() + .expect("gauge series to host"); + let all: Vec = runoff.into_data().to_vec::().expect("runoff to host"); + (series, all) +} + +/// Total lateral inflow entering the whole subgraph, m³/s. At steady state +/// every drop of it must appear at the gauge. +const EXPECTED: f32 = Q_PRIME * N as f32; // 30.0 + +#[test] +fn gauge_prediction_conserves_mass_when_not_ddr_match() { + let (series, all) = route_and_extract(false, 1); + let final_q = series[T - 1]; + println!( + "ddr_match=false: gauge={final_q:.4} m3/s expected={EXPECTED:.4} \ + ratio={:.4}", + final_q / EXPECTED + ); + + assert!( + (final_q - EXPECTED).abs() / EXPECTED < 1e-2, + "gauge prediction at steady state = {final_q} m³/s, expected {EXPECTED} \ + m³/s (sum of q_prime over all {N} reaches). Ratio {:.3}. A ratio near \ + {:.3} means outflow_idx is pointing at the gauge's upstream neighbours \ + instead of the gauge reach itself.", + final_q / EXPECTED, + (N as f32 - 1.0) / N as f32, + ); + + // And the extracted prediction must literally BE the gauge reach's own + // routed discharge, at every timestep — not a sum over other reaches. + for t in 0..T { + let reach2 = all[2 * T + t]; + assert!( + (series[t] - reach2).abs() <= 1e-4 * reach2.abs().max(1.0), + "t={t}: gauge prediction {} != gauge reach discharge {reach2}", + series[t], + ); + } +} + +#[test] +fn ddr_match_gauge_prediction_omits_the_gauge_reach() { + // Pins the defect that `ddr_match: true` faithfully reproduces, and proves + // the mass check above actually discriminates between the two conventions. + // Summing the two headwaters yields 2/3 of the network's lateral inflow; + // the gauge reach's own 10 m³/s is silently dropped. + let (series, all) = route_and_extract(true, 1); + let final_q = series[T - 1]; + let ddr_expected = Q_PRIME * (N as f32 - 1.0); // 20.0 + println!( + "ddr_match=true : gauge={final_q:.4} m3/s mass-conserving={EXPECTED:.4} \ + ratio={:.4}", + final_q / EXPECTED + ); + + assert!( + (final_q - ddr_expected).abs() / ddr_expected < 1e-2, + "ddr_match=true must sum the two headwaters only: got {final_q}, \ + expected {ddr_expected}" + ); + assert!( + (final_q - EXPECTED).abs() / EXPECTED > 0.2, + "ddr_match=true must NOT conserve mass — if it does, the test network \ + no longer discriminates the two outflow_idx conventions" + ); + // The gauge reach itself carries the full mass; it is simply not read. + let reach2_final = all[2 * T + (T - 1)]; + assert!( + (reach2_final - EXPECTED).abs() / EXPECTED < 1e-2, + "gauge reach discharge {reach2_final} should still conserve mass \ + ({EXPECTED}); the defect is in extraction, not routing" + ); +} + +/// Steady-state gauge discharge with the gauge's own reach split `pieces` ways. +fn gauge_steady_state(pieces: usize) -> f32 { + route_and_extract(false, pieces).0[T - 1] +} + +#[test] +fn gauge_conserves_mass_when_its_reach_is_subdivided() { + // Same topology as `gauge_prediction_conserves_mass_when_not_ddr_match`, + // but the gauge's own reach is split 4 ways. The answer must not change: + // the MC solve is mass-conserving down the internal chain, so the whole + // reach's runoff — upstream network plus its own `q'`, split `q'/4` across + // the pieces — arrives at the LAST piece. + let un_split = gauge_steady_state(1); + let split = gauge_steady_state(4); + println!( + "subdivision: 1 piece = {un_split:.6} m3/s 4 pieces = {split:.6} m3/s \ + expected = {EXPECTED:.6}" + ); + assert!((un_split - EXPECTED).abs() < 1e-3, "control changed: {un_split}"); + assert!( + (split - EXPECTED).abs() < 1e-3, + "subdivided gauge lost mass: got {split}, expected {EXPECTED}. Reading \ + any piece other than the outlet drops the downstream fraction of the \ + gauge reach's own lateral inflow — the inlet piece carries only \ + {:.1} m3/s here.", + Q_PRIME * (N as f32 - 1.0) + Q_PRIME / 4.0, + ); +} + +#[test] +fn subdivided_interior_pieces_carry_less_than_the_outlet() { + // Proves the test above actually discriminates: the four pieces of the + // gauge reach form a strictly increasing ramp (22.5, 25.0, 27.5, 30.0), so + // pointing `outflow_idx` at anything but the last piece is detectable. + let (_, all) = route_and_extract(false, 4); + let n = 2 + 4; + let piece_q: Vec = (2..n).map(|r| all[r * T + (T - 1)]).collect(); + println!("gauge-reach pieces at steady state: {piece_q:?}"); + for w in piece_q.windows(2) { + assert!( + w[1] > w[0] + 1e-3, + "pieces must accumulate downstream, got {piece_q:?}" + ); + } + let expected_inlet = Q_PRIME * (N as f32 - 1.0) + Q_PRIME / 4.0; // 22.5 + assert!( + (piece_q[0] - expected_inlet).abs() < 1e-2, + "inlet piece {} != {expected_inlet} (two headwaters + one quarter of \ + the gauge reach's own lateral inflow)", + piece_q[0] + ); + assert!((piece_q[3] - EXPECTED).abs() < 1e-2, "outlet {} != {EXPECTED}", piece_q[3]); +} diff --git a/tests/leakance_gradcheck.rs b/tests/leakance_gradcheck.rs index 8bf2ee0..f98fb32 100644 --- a/tests/leakance_gradcheck.rs +++ b/tests/leakance_gradcheck.rs @@ -166,6 +166,7 @@ fn run_forward( fac_t.clone(), mask, None, + false, ); ( diff --git a/tests/negative_discharge_counter.rs b/tests/negative_discharge_counter.rs new file mode 100644 index 0000000..9d88b95 --- /dev/null +++ b/tests/negative_discharge_counter.rs @@ -0,0 +1,166 @@ +//! The S28 clamp silently turns negative solves into +1e-4. This counter is +//! the only way to see how often Muskingum's non-negative-coefficient +//! condition (2X <= Cr <= 2(1-X)) is violated in a real run. +//! +//! Both test cases run inside ONE `#[test]` function to avoid races on the +//! process-global `NEG_SOLVES`/`TOTAL_SOLVES` atomics. Running this file +//! with `--test-threads=1` would also be safe but is not required here. +use std::sync::Arc; + +use burn::backend::NdArray; +use burn::tensor::Tensor; + +use ddrs::config::Config; +use ddrs::routing::mmc_op::{negative_solve_stats, reset_negative_solve_stats, timestep_forward}; +use ddrs::sparse::{AValuesAssembler, CsrPattern, SparseAdjacency}; + +type I = NdArray; + +/// 4-reach linear chain (reach k feeds reach k+1). +fn linear_chain(n: usize, length_m: Vec, slope: Vec) -> SparseAdjacency { + 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, length_m, slope) +} + +fn base_cfg() -> Config { + let mut cfg = Config::default(); + cfg.params.parameter_ranges.n = [0.01, 0.3]; + cfg.params.parameter_ranges.q_spatial = [0.1, 0.9]; + cfg.params.parameter_ranges.p_spatial = [1.0, 200.0]; + cfg.params.attribute_minimums.velocity = 0.01; + cfg.params.attribute_minimums.depth = 0.001; + cfg.params.attribute_minimums.discharge = 1e-4; + cfg.params.attribute_minimums.bottom_width = 0.01; + cfg.params.attribute_minimums.slope = 0.0001; + cfg.params.defaults.insert("p_spatial".to_string(), 1.0); + cfg.params.log_space_parameters = vec![]; + cfg +} + +#[test] +fn counter_starts_at_zero_and_positive_control() { + let device = ::Device::default(); + + // ----------------------------------------------------------------------- + // Case A: zero counter after reset (tautology guard) + // ----------------------------------------------------------------------- + reset_negative_solve_stats(); + let (neg, total) = negative_solve_stats(); + assert_eq!(neg, 0, "neg count must be zero after reset"); + assert_eq!(total, 0, "total count must be zero after reset"); + + // ----------------------------------------------------------------------- + // Case B: positive control — inputs that drive x_sol[headwater] < 0 + // + // Design target: c3 = (2K(1-X) - dt) / denom < 0 at every reach. + // K = L / celerity, celerity = v_clamped * 5/3 + // With L=100 m, slope=0.05, n=0.01, q_t=1e5 m³/s (large) the Manning + // formula gives a large velocity → large celerity → small K. + // At these numbers: depth ≈ large, v_Manning well above clamp → + // celerity ≈ 15 * 5/3 = 25 m/s → K ≈ 100/25 = 4 s. + // dt = 3600 s, X = 0.3 → 2K(1-X) = 5.6 s << 3600 → c3 < 0. + // The headwater reach has no upstream (i_t[0]=0), so + // b_rhs[0] = c2*0 + c3*q_t[0] + c4*q_prime_t[0] + // With q_t[0]=1e5 and q_prime_t[0]=1e-3, the large negative c3*q_t + // dominates → b_rhs[0] < 0 → x_sol[0] < 0 before the S28 clamp. + // ----------------------------------------------------------------------- + reset_negative_solve_stats(); + + const N: usize = 4; + let adj = linear_chain( + N, + vec![100.0_f32; N], // very short reaches (100 m) + vec![0.05_f32; N], // steep slope + ); + let pattern = Arc::new(CsrPattern::from_sparse(&adj)); + let assembler = AValuesAssembler::::new(&pattern, &device); + let cfg = base_cfg(); + + let mk = |v: f32| Tensor::, 1>::from_floats([v; N].as_ref(), &device); + // n=0.01 → low Manning → high velocity; q_spatial=0.5; p_spatial=20; x_storage=0.3 + let n_t = mk(0.01); + let qsp_t = mk(0.5); + let psp_t = mk(20.0); + let x_t = mk(0.3); + // Very large current discharge to maximise |c3 * q_t| + let qt_t: burn::tensor::Tensor, 1> = + Tensor::from_floats([1e5_f32; N].as_ref(), &device); + // Tiny inflow so c4*q_prime_t is negligible + let qpt_t: burn::tensor::Tensor, 1> = + Tensor::from_floats([1e-3_f32; N].as_ref(), &device); + + let _q_next = timestep_forward::( + &cfg, + &pattern, + &assembler, + n_t, + qsp_t, + psp_t, + qt_t, + qpt_t, + Tensor::from_floats(adj.length_m.as_slice(), &device), + Tensor::from_floats(adj.slope.as_slice(), &device), + x_t, + true, // track_neg = ON + ); + + let (neg_b, total_b) = negative_solve_stats(); + assert!( + neg_b > 0, + "positive control: expected at least one negative solve before the S28 clamp, \ + got 0 out of {total_b}. Check that c3 < 0 with the chosen L/slope/n/q_t." + ); + + // ----------------------------------------------------------------------- + // Case C: negative control — inputs that keep c3 >= 0 + // + // Design target: c3 >= 0 ⟺ 2K(1-X) >= dt = 3600 s + // With L=10000 m, low celerity (large n, low q): K = L/c ≈ 10000/0.1 = 1e5 s + // → 2K(1-X) = 1.4e5 >> 3600 → c3 > 0. + // With smooth inflow q_t ≈ q_prime_t all coefficients positive → + // x_sol > 0 at every reach, no negatives. + // ----------------------------------------------------------------------- + reset_negative_solve_stats(); + + let adj2 = linear_chain( + N, + vec![10_000.0_f32; N], // long reaches (10 km) + vec![0.001_f32; N], // gentle slope + ); + let pattern2 = Arc::new(CsrPattern::from_sparse(&adj2)); + let assembler2 = AValuesAssembler::::new(&pattern2, &device); + let cfg2 = base_cfg(); + + let mk2 = |v: f32| Tensor::, 1>::from_floats([v; N].as_ref(), &device); + let n_t2 = mk2(0.04); // typical Manning's n + let qsp_t2 = mk2(0.4); + let psp_t2 = mk2(20.0); + let x_t2 = mk2(0.3); + let qt_t2 = mk2(100.0); // moderate discharge + let qpt_t2 = mk2(10.0); // smooth inflow + + let _q_next2 = timestep_forward::( + &cfg2, + &pattern2, + &assembler2, + n_t2, + qsp_t2, + psp_t2, + qt_t2, + qpt_t2, + Tensor::from_floats(adj2.length_m.as_slice(), &device), + Tensor::from_floats(adj2.slope.as_slice(), &device), + x_t2, + true, // track_neg = ON + ); + + let (neg_c, total_c) = negative_solve_stats(); + assert_eq!( + neg_c, 0, + "negative control: expected zero negative solves with long reaches / gentle slope / \ + typical flow, got {neg_c} out of {total_c}" + ); +} diff --git a/tests/positivity_clamp.rs b/tests/positivity_clamp.rs new file mode 100644 index 0000000..552b037 --- /dev/null +++ b/tests/positivity_clamp.rs @@ -0,0 +1,817 @@ +//! Positivity clamp (`params.enforce_positivity`) — S18' K floor and S19' X cap. +//! +//! # The theorem +//! +//! With `Cr = Δt/K` and `denom = 2K(1−X) + Δt > 0`: +//! +//! ```text +//! c2 = (2KX + Δt)/denom > 0 always (K>0, X>=0, Δt>0) +//! c4 = 2Δt/denom > 0 always +//! c1 = (Δt − 2KX)/denom >= 0 <=> Cr >= 2X +//! c3 = (2K(1−X) − Δt)/denom >= 0 <=> Cr <= 2(1−X) +//! ``` +//! +//! S27 is forward substitution in topological order, +//! `x[i] = b[i] + c1[i]·Σ_{j∈up(i)} x[j]`, with +//! `b[i] = c2·(N q_t)[i] + c3·q_t[i] + c4·q'[i]`. Since `q_t > 0` (S28's +//! `clamp_min(1e-4)` and the hotstart) and `q' > 0`, `c1, c3 >= 0` gives +//! `b >= 0`, and induction over the topological order gives `x >= 0` at every +//! reach: headwaters have no upstream so `x = b >= 0`, and each subsequent +//! `x[i]` is a non-negative combination of `b[i]` and already-non-negative +//! upstream values. +//! +//! The clamp therefore targets the INPUTS (`K`, `X`), never the coefficients: +//! `c1+c2+c3 = 1` holds identically for any `(K, X)`, so mass is preserved +//! exactly — clamping `c3` would break that. +//! +//! # Non-vacuity +//! +//! `fixture_is_not_vacuous` is a standing guard, not decoration. The Cunge-X +//! gradcheck in `tests/cunge_x.rs` was initially VACUOUS because a 1000 m +//! fixture saturated `X`'s clamp on every reach, so all four tests passed with +//! the backward terms deleted. Here the analogue would be a fixture where one +//! branch of the three-way `min` always wins, or where every reach sits on the +//! same side of the K floor. The fixture below spans `Cr` from 0.051 to 7.9, +//! straddles the K floor 5/10, and lets all three `min` branches win. +//! +//! Part 2's gradcheck uses `grad_fixture()` — the same reaches with a graded +//! `q_t` — and re-asserts the same branch-mix / K-floor-straddle preconditions +//! before every comparison via `assert_fixture_exercises_everything`. + +use std::sync::Arc; + +use burn::backend::{Autodiff, NdArray}; +use burn::tensor::Tensor; + +use ddrs::config::Config; +use ddrs::routing::mmc_op::{timestep_forward, POSITIVITY_DELTA}; +use ddrs::sparse::{AValuesAssembler, CsrPattern, SparseAdjacency}; + +type I = NdArray; +type AB = Autodiff; + +/// `crate::routing::mmc::DT_SECONDS`, restated so a change there fails loudly +/// here rather than silently retuning the fixture. +const DT: f32 = 3600.0; + +/// `x_sol` from the PRE-CHANGE code path, captured bit-for-bit by running this +/// exact fixture against `src/routing/mmc_op.rs` at commit fa5bcb4 (the config +/// flag, before the forward gained S18'/S19'). This is what makes +/// `off_parity_bit_identical_when_disabled` a real regression test rather than +/// a self-consistency tautology. +const GOLDEN_X_SOL_BITS: [u32; 10] = [ + 0xc0645cca, // -3.5681634 + 0xbf9f4594, // -1.2443109 + 0x3f511646, // 0.8167461 + 0xbf7f58e2, // -0.99745 + 0x4143d4d4, // 12.23946 + 0x42857d42, // 66.744644 + 0x433334a2, // 179.2056 + 0x428d80f3, // 70.751854 + 0x42f2ee24, // 121.46512 + 0x42a20a08, // 81.01959 +]; + +// =========================================================================== +// Fixture +// =========================================================================== + +/// Which branch of `x_eff = min(x_cunge, hi_a, hi_b)` won on a reach. +#[derive(Copy, Clone, PartialEq, Eq, Debug)] +enum Branch { + /// The Cunge-derived X was already inside the stability window. + Cunge, + /// `hi_a = 0.5·Cr·(1−δ)` — the `c1 >= 0` constraint bound. + HiA, + /// `hi_b = (1 − 0.5·Cr)·(1−δ)` — the `c3 >= 0` constraint bound. + HiB, +} + +struct Fixture { + length: Vec, + slope: Vec, + n: Vec, + qsp: Vec, + psp: Vec, + qt: Vec, + qpt: Vec, +} + +/// A 10-reach linear chain deliberately spanning the whole Courant range. +/// +/// `Cr = Δt·c/L`, so length is the primary lever: 400 m gives `Cr ≈ 7.9` +/// (deep in the `c3 < 0` regime that produces negative solves) and 150 km +/// gives `Cr ≈ 0.051` (deep in the `c1 < 0` regime). `q_t` is graded so the +/// short reaches do not saturate Cunge `X` at its lower clamp, which keeps at +/// least one Cunge-branch win in the interior of `[0, 0.5]`. +/// +/// `q'` is small (0.01 m³/s) so that on the `c3 < 0` reaches the negative +/// `c3·q_t` term is not masked by `c4·q'` — that is what makes the positive +/// control produce actual negatives. +fn fixture() -> Fixture { + let length = vec![ + 400.0f32, 800.0, 1500.0, 2500.0, 3600.0, 5000.0, 8000.0, 20000.0, 60000.0, 150000.0, + ]; + let n = length.len(); + Fixture { + length, + slope: vec![0.001; n], + n: vec![0.035; n], + qsp: vec![0.4; n], + psp: vec![20.0; n], + qt: vec![6.0, 8.0, 15.0, 50.0, 120.0, 200.0, 100.0, 100.0, 100.0, 100.0], + qpt: vec![0.01; n], + } +} + +fn stress_cfg(enforce_positivity: bool) -> Config { + let mut cfg = Config::default(); + // The clamp only exists on the corrected-physics path (and `Config`'s + // loader rejects `ddr_match: true` + `enforce_positivity: true` outright). + cfg.params.ddr_match = false; + cfg.params.enforce_positivity = enforce_positivity; + cfg.params.parameter_ranges.n = [0.01, 0.3]; + cfg.params.parameter_ranges.q_spatial = [0.1, 0.9]; + cfg.params.parameter_ranges.p_spatial = [1.0, 200.0]; + cfg.params.attribute_minimums.velocity = 0.01; + cfg.params.attribute_minimums.depth = 0.001; + cfg.params.attribute_minimums.discharge = 1e-4; + cfg.params.attribute_minimums.bottom_width = 0.01; + cfg.params.attribute_minimums.slope = 0.0001; + cfg.params.defaults.insert("p_spatial".to_string(), 1.0); + cfg.params.log_space_parameters = vec![]; + cfg +} + +fn chain(f: &Fixture) -> SparseAdjacency { + let n = f.length.len(); + 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, f.length.clone(), f.slope.clone()) +} + +/// The S1..S23 quantities the fixture assertions need, read out of the real +/// forward chain (not recomputed from a parallel model). +struct ChainOutputs { + celerity: Vec, + k_muskingum: Vec, + top_width: Vec, + denom: Vec, + c1: Vec, + c2: Vec, + c3: Vec, +} + +fn run_chain(f: &Fixture, enforce_positivity: bool) -> ChainOutputs { + let adj = chain(f); + let device = ::Device::default(); + let pattern = Arc::new(CsrPattern::from_sparse(&adj)); + let mk = |d: &[f32]| -> Tensor { Tensor::from_floats(d, &device) }; + let outs = ddrs::routing::mmc_op::__spike_forward_chain_k1_outputs::( + &stress_cfg(enforce_positivity), + &pattern, + mk(&f.n), + mk(&f.qsp), + mk(&f.psp), + mk(&f.qt), + mk(&f.qpt), + mk(&f.length), + mk(&f.slope), + mk(&vec![0.3f32; f.length.len()]), + ); + // k1 output order: [depth, top_width, side_slope, bottom_width, hyd_radius, + // velocity_un, velocity_cl, celerity, k_muskingum, denom, + // c1, c2, c3, c4, ...] + ChainOutputs { + top_width: outs[1].clone(), + celerity: outs[7].clone(), + k_muskingum: outs[8].clone(), + denom: outs[9].clone(), + c1: outs[10].clone(), + c2: outs[11].clone(), + c3: outs[12].clone(), + } +} + +/// `x_sol` — the raw S27 solve, BEFORE S28's `clamp_min` rewrites negatives to +/// `+1e-4`. This is the quantity the whole task is about. +fn run_solve(f: &Fixture, enforce_positivity: bool) -> Vec { + let adj = chain(f); + let device = ::Device::default(); + let pattern = Arc::new(CsrPattern::from_sparse(&adj)); + let mk = |d: &[f32]| -> Tensor { Tensor::from_floats(d, &device) }; + let (_b_rhs, _i_t, x_sol, _q_next) = + ddrs::routing::mmc_op::__spike_forward_chain_k23_outputs::( + &stress_cfg(enforce_positivity), + &pattern, + mk(&f.n), + mk(&f.qsp), + mk(&f.psp), + mk(&f.qt), + mk(&f.qpt), + mk(&f.length), + mk(&f.slope), + mk(&vec![0.3f32; f.length.len()]), + ); + x_sol +} + +/// `(negatives, total)` in the raw solve. +fn run_stress_network(enforce_positivity: bool) -> (usize, usize) { + let x_sol = run_solve(&fixture(), enforce_positivity); + (x_sol.iter().filter(|&&v| v < 0.0).count(), x_sol.len()) +} + +/// Which `min` branch won on each reach, and whether the K floor bound, both +/// derived from the *actual* chain outputs. `x_cunge` uses the pre-floor +/// `celerity` (S17) while `cr` uses the post-floor `k_muskingum` (S18'), which +/// is exactly how the forward composes them. +fn branch_report(f: &Fixture) -> (Vec, Vec, Vec, Vec) { + let out = run_chain(f, true); + let k_floor = DT * (1.0 + POSITIVITY_DELTA) / 2.0; + let mut branches = Vec::new(); + let mut floored = Vec::new(); + let mut cr_raw = Vec::new(); + let mut x_cunge_all = Vec::new(); + for i in 0..f.length.len() { + let k_raw = f.length[i] / out.celerity[i]; + floored.push(k_raw < k_floor); + cr_raw.push(DT / k_raw); + let cr = DT / out.k_muskingum[i]; + let w = f.qt[i] / (out.top_width[i] * f.slope[i] * out.celerity[i] * f.length[i] + 1e-12); + let x_cunge = (0.5 * (1.0 - w)).clamp(0.0, 0.5); + x_cunge_all.push(x_cunge); + let hi_a = cr * 0.5 * (1.0 - POSITIVITY_DELTA); + let hi_b = (1.0 - 0.5 * cr) * (1.0 - POSITIVITY_DELTA); + branches.push(if x_cunge <= hi_a && x_cunge <= hi_b { + Branch::Cunge + } else if hi_a <= hi_b { + Branch::HiA + } else { + Branch::HiB + }); + } + (branches, floored, cr_raw, x_cunge_all) +} + +// =========================================================================== +// Tests +// =========================================================================== + +/// STANDING GUARD against a fixture that would let every test above pass with +/// the clamp deleted. Four independent ways this file could go vacuous: +/// +/// 1. one `min` branch always wins → the other two are never exercised; +/// 2. every reach on the same side of the K floor → S18' is a no-op or a +/// constant; +/// 3. `Cr` never leaves `[0.98, 1.02]` → nothing to clamp; +/// 4. every Cunge-branch win sits on `x_cunge`'s own `[0, 0.5]` clamp → the +/// branch is "active" but carries no signal (this is precisely the failure +/// mode that made `tests/cunge_x.rs` vacuous at 1000 m reaches). +#[test] +fn fixture_is_not_vacuous() { + let f = fixture(); + let (branches, floored, cr_raw, x_cunge) = branch_report(&f); + let n = branches.len() as f32; + let frac = |b: Branch| branches.iter().filter(|&&x| x == b).count() as f32 / n; + let (fc, fa, fb) = (frac(Branch::Cunge), frac(Branch::HiA), frac(Branch::HiB)); + let f_floored = floored.iter().filter(|&&x| x).count() as f32 / n; + let cr_min = cr_raw.iter().cloned().fold(f32::INFINITY, f32::min); + let cr_max = cr_raw.iter().cloned().fold(f32::NEG_INFINITY, f32::max); + println!("branch mix: cunge={fc:.2} hi_a={fa:.2} hi_b={fb:.2}"); + println!("K floored: {f_floored:.2} Cr_raw span: {cr_min:.3} .. {cr_max:.3}"); + for i in 0..branches.len() { + println!( + " [{i}] L={:>8.0} Cr_raw={:>7.3} floored={:<5} x_cunge={:.4} branch={:?}", + f.length[i], cr_raw[i], floored[i], x_cunge[i], branches[i] + ); + } + + assert!(fc > 0.05, "Cunge branch never wins ({fc:.2}) — S19' cap is unfalsifiable"); + assert!(fa > 0.05, "hi_a branch never wins ({fa:.2}) — the c1>=0 bound is untested"); + assert!(fb > 0.05, "hi_b branch never wins ({fb:.2}) — the c3>=0 bound is untested"); + assert!( + (0.05..0.95).contains(&f_floored), + "fixture must STRADDLE the K floor, got {f_floored:.2} floored" + ); + assert!(cr_min < 0.5, "fixture must reach well below Cr=1, got min {cr_min:.3}"); + assert!(cr_max > 2.0, "fixture must reach well above Cr=2, got max {cr_max:.3}"); + // At least one Cunge win must be strictly interior to X's own [0, 0.5] + // clamp, or the "Cunge branch wins" reaches carry no gradient signal. + let interior_cunge = branches + .iter() + .zip(&x_cunge) + .filter(|(&b, &x)| b == Branch::Cunge && x > 0.02 && x < 0.48) + .count(); + assert!( + interior_cunge > 0, + "every Cunge-branch win sits on X's own clamp — same vacuity trap as cunge_x.rs" + ); +} + +/// POSITIVE CONTROL. If this stops producing negatives the fixture is broken +/// and every other test in this file becomes meaningless. +#[test] +fn positive_control_negatives_exist_without_the_clamp() { + let (neg, tot) = run_stress_network(false); + println!("clamp OFF: {neg}/{tot} negative solves"); + assert!(neg > 0, "fixture must actually produce negatives, got {neg}/{tot}"); +} + +#[test] +fn clamp_drives_negatives_to_exactly_zero() { + let (neg, tot) = run_stress_network(true); + println!("clamp ON: {neg}/{tot} negative solves"); + assert_eq!(neg, 0, "expected exactly zero negative solves, got {neg}/{tot}"); + assert!(tot > 0, "fixture must have solved something"); +} + +/// Mass conservation. The whole reason the clamp targets `K` and `X` rather +/// than `c1`/`c3` is that this identity is preserved for free. +#[test] +fn partition_identity_survives_the_clamp() { + let out = run_chain(&fixture(), true); + for i in 0..out.c1.len() { + let s = out.c1[i] + out.c2[i] + out.c3[i]; + assert!( + (s - 1.0).abs() < 1e-5, + "reach {i}: c1+c2+c3 = {s} (c1={}, c2={}, c3={})", + out.c1[i], + out.c2[i], + out.c3[i] + ); + } +} + +/// The δ margin is what makes this hold in f32: at δ = 0 the cap lands exactly +/// on `c1 = 0` / `c3 = 0` and roundoff crosses it. +#[test] +fn coefficients_are_non_negative_with_margin() { + let out = run_chain(&fixture(), true); + let min_c1 = out.c1.iter().cloned().fold(f32::INFINITY, f32::min); + let min_c3 = out.c3.iter().cloned().fold(f32::INFINITY, f32::min); + println!("min c1 = {min_c1:e}, min c3 = {min_c3:e}"); + assert!(min_c1 >= 0.0, "min c1 = {min_c1:e} < 0"); + assert!(min_c3 >= 0.0, "min c3 = {min_c3:e} < 0"); + // The same coefficients WITHOUT the clamp must go negative, or this test + // would pass on an unclamped build. + let off = run_chain(&fixture(), false); + let off_c1 = off.c1.iter().cloned().fold(f32::INFINITY, f32::min); + let off_c3 = off.c3.iter().cloned().fold(f32::INFINITY, f32::min); + assert!( + off_c1 < 0.0 && off_c3 < 0.0, + "fixture must violate BOTH bounds when unclamped (min c1={off_c1:e}, min c3={off_c3:e})" + ); +} + +/// `enforce_positivity: false` must be a BIT-IDENTICAL no-op. `GOLDEN_X_SOL_BITS` +/// was captured from the code at commit fa5bcb4, i.e. before S18'/S19' existed. +#[test] +fn off_parity_bit_identical_when_disabled() { + let x_sol = run_solve(&fixture(), false); + assert_eq!(x_sol.len(), GOLDEN_X_SOL_BITS.len()); + for (i, (&v, &want)) in x_sol.iter().zip(&GOLDEN_X_SOL_BITS).enumerate() { + assert_eq!( + v.to_bits(), + want, + "reach {i}: disabled clamp changed x_sol — got {v:e} ({:#010x}), \ + pre-change {:e} ({want:#010x})", + v.to_bits(), + f32::from_bits(want) + ); + } + // Non-vacuity: the clamp must actually move this fixture, otherwise + // "byte-identical when disabled" is trivially true. + let on = run_solve(&fixture(), true); + assert!( + on.iter().zip(&x_sol).any(|(a, b)| a != b), + "clamp ON produced identical output — the off-parity test proves nothing" + ); +} + +/// S18' floors K at `Δt(1+δ)/2`, which is exactly the condition +/// `Cr <= 2/(1+δ) < 2`, which is what makes `hi_b = (1 − 0.5·Cr)(1−δ) > 0` and +/// therefore makes the three-way `min` positive without a `clamp_min`. +#[test] +fn k_floor_keeps_courant_below_two() { + let f = fixture(); + let out = run_chain(&f, true); + let off = run_chain(&f, false); + let k_floor = DT * (1.0 + POSITIVITY_DELTA) / 2.0; + let cr_cap = 2.0 / (1.0 + POSITIVITY_DELTA); + let mut any_floored = false; + for i in 0..f.length.len() { + assert!( + out.k_muskingum[i] >= k_floor * (1.0 - 1e-6), + "reach {i}: K = {} below floor {k_floor}", + out.k_muskingum[i] + ); + let cr = DT / out.k_muskingum[i]; + assert!(cr <= cr_cap * (1.0 + 1e-6), "reach {i}: Cr = {cr} exceeds {cr_cap}"); + if off.k_muskingum[i] < k_floor { + any_floored = true; + // Where the floor binds it must bind to exactly the floor value. + assert!( + (out.k_muskingum[i] - k_floor).abs() < 1e-3, + "reach {i}: floored K = {} != {k_floor}", + out.k_muskingum[i] + ); + } else { + // Where it does not bind, K must be untouched bit-for-bit. + assert_eq!( + out.k_muskingum[i].to_bits(), + off.k_muskingum[i].to_bits(), + "reach {i}: unfloored K was perturbed" + ); + } + } + assert!(any_floored, "no reach hit the K floor — the S18' branch is untested"); +} + +// =========================================================================== +// Part 2: gradcheck of B18'/B19' (the positivity-clamp backward). +// +// Scaffolding mirrors `tests/cunge_x.rs` Part 2. Central finite differences on +// the three learnable parents (`n`, `q_spatial`, `p_spatial`) plus `q_t`, which +// is the parent the clamp most affects: it reaches the loss through the S25 +// RHS, the S24 SpMV, the S2 depth chain AND the Cunge X — and under the clamp +// that last path is switched off on every reach where `hi_a`/`hi_b` win the min. +// +// The two facts under test: +// B18' gk_raw = gk_musk · mask(k_raw > k_floor) +// B19' gx_eff splits three ways; the `hi_a`/`hi_b` branches feed a NEW +// path `x_eff → cr → k_musk → celerity`. +// +// `assert_fixture_exercises_everything` runs before every comparison: with a +// single-branch or single-side-of-the-floor fixture these tests would pass with +// the new backward terms deleted, which is exactly how `tests/cunge_x.rs` went +// vacuous at 1000 m reaches. +// +// # Two f32 hazards this fixture forces us to handle explicitly +// +// 1. **Kink crossing.** `sum()` of a raw MC forward is O(600) here, so the +// sibling gradchecks' ABSOLUTE step `max(1e-3·x, 1e-3)` is a 2.9% +// perturbation of `n = 0.035`. That is large enough to move reach 0 across +// the `x_cunge` / `hi_b` boundary (it sits at 0.0044 vs 0.0099), so central +// differences average two different slopes and disagree with the analytical +// gradient by ~30% — an FD artifact, not a backward bug. `REL_STEP` is +// therefore a PURE RELATIVE step, small enough to stay on one branch. +// 2. **Cancellation.** With a small step, `l_plus − l_minus` on an O(600) loss +// lands at the f32 round-off floor. `conditioning_weights` rescales each +// reach's contribution to O(1) (loss ≈ 10) so the difference survives, and +// `fd_noise` states the surviving quantum explicitly instead of hiding it in +// a magic `ABS_TOL`: a disagreement is only forgiven when it is provably +// below what f32 central differences can resolve at that step size. +// =========================================================================== + +/// The gradcheck fixture: `fixture()` with a q_t profile that VARIES sharply +/// along the chain. +/// +/// `fixture()`'s flat `q_t = 100` on reaches 6..9 makes them gradient-DEAD, and +/// not because of any masking: `q_next = c1·Σx_up + c2·i_t + c3·q_t + c4·q'`, +/// and on a chain with `x_up ≈ i_t ≈ q_t` the partition identity `c1+c2+c3 = 1` +/// makes `q_next ≈ q_t` no matter what `(K, X)` do. Every S18'/S19' effect then +/// cancels to f32 noise, and the `hi_a` branch — which only ever wins on the +/// long, low-Courant reaches — would go untested. Grading `q_t` breaks the +/// cancellation and lifts those reaches above the FD floor. +/// +/// `fixture()` itself must NOT be retuned: `GOLDEN_X_SOL_BITS` is keyed to it. +fn grad_fixture() -> Fixture { + let mut f = fixture(); + f.qt = vec![6.0, 8.0, 15.0, 50.0, 120.0, 200.0, 300.0, 40.0, 400.0, 30.0]; + f +} + +/// FD step as a pure FRACTION of the base value (see hazard 1 above). +const REL_STEP: f32 = 3e-3; +const REL_TOL: f32 = 5e-3; +/// How many ulps of the loss to allow for accumulated round-off in the weighted +/// sum plus the subtraction. 10 reaches ⇒ a handful of ulps; 16 is slack. +const NOISE_ULPS: f32 = 16.0; + +/// Per-reach loss weights `w[i] = 1/max(q_next_base[i], 1)`, so every reach +/// contributes O(1) to `loss = Σ w[i]·q_next[i]` instead of the O(180) that the +/// downstream reaches would otherwise contribute. Constant (no grad) — this is +/// just a better-conditioned scalar loss, and the analytical backward sees it +/// as `grad_out = w` rather than `1`. +fn conditioning_weights() -> Vec { + let f = grad_fixture(); + let base = run_solve(&f, true); + base.iter().map(|&q| 1.0 / q.max(1.0)).collect() +} + +#[derive(Copy, Clone, Debug)] +enum Parent { + N, + QSpatial, + PSpatial, + QT, +} + +struct GradTensors { + n: Tensor, + qsp: Tensor, + psp: Tensor, + qt: Tensor, +} + +#[allow(clippy::too_many_arguments)] +fn run_forward_loss( + cfg: &Config, + pattern: &Arc, + assembler: &AValuesAssembler, + device: &::Device, + n_vec: &[f32], + qsp_vec: &[f32], + psp_vec: &[f32], + qt_vec: &[f32], + qpt_vec: &[f32], + length_vec: &[f32], + slope_vec: &[f32], + weights: &[f32], + require_grad_parent: Option, +) -> (Tensor, GradTensors) { + let mk = |data: &[f32], req: bool| -> Tensor { + let t: Tensor = Tensor::from_floats(data, device); + if req { t.require_grad() } else { t } + }; + let n_t = mk(n_vec, matches!(require_grad_parent, Some(Parent::N))); + let qsp_t = mk(qsp_vec, matches!(require_grad_parent, Some(Parent::QSpatial))); + let psp_t = mk(psp_vec, matches!(require_grad_parent, Some(Parent::PSpatial))); + let qt_t = mk(qt_vec, matches!(require_grad_parent, Some(Parent::QT))); + let qpt_t = mk(qpt_vec, false); + let length_t = mk(length_vec, false); + let slope_t = mk(slope_vec, false); + let xst_t = mk(&vec![0.3f32; n_vec.len()], false); + + let q_next = timestep_forward::( + cfg, + pattern, + assembler, + n_t.clone(), + qsp_t.clone(), + psp_t.clone(), + qt_t.clone(), + qpt_t.clone(), + length_t, + slope_t, + xst_t, + false, + ); + + // Conditioned scalar loss (see hazard 2 in the section header). + let loss = q_next * mk(weights, false); + + ( + loss, + GradTensors { + n: n_t, + qsp: qsp_t, + psp: psp_t, + qt: qt_t, + }, + ) +} + +fn compute_analytical_grad(parent: Parent) -> Vec { + let f = grad_fixture(); + let cfg = stress_cfg(true); + let adj = chain(&f); + let device = ::Device::default(); + let pattern = Arc::new(CsrPattern::from_sparse(&adj)); + let assembler = AValuesAssembler::::new(&pattern, &device); + let weights = conditioning_weights(); + + let (loss, parents) = run_forward_loss( + &cfg, &pattern, &assembler, &device, &f.n, &f.qsp, &f.psp, &f.qt, &f.qpt, &f.length, + &f.slope, &weights, Some(parent), + ); + + let grads = loss.sum().backward(); + let g = match parent { + Parent::N => parents.n.grad(&grads).expect("grad on n"), + Parent::QSpatial => parents.qsp.grad(&grads).expect("grad on q_spatial"), + Parent::PSpatial => parents.psp.grad(&grads).expect("grad on p_spatial"), + Parent::QT => parents.qt.grad(&grads).expect("grad on q_t"), + }; + g.into_data().to_vec::().unwrap() +} + +/// `(fd_grad, fd_noise_floor)`. The second vector is the smallest gradient +/// difference f32 central differences can resolve at each reach's step size: +/// `NOISE_ULPS · ulp(loss) / (2·eps_i)`. Anything below it is unmeasurable, not +/// evidence about the backward. +fn compute_fd_grad(parent: Parent) -> (Vec, Vec) { + let f = grad_fixture(); + let cfg = stress_cfg(true); + let adj = chain(&f); + let device = ::Device::default(); + let pattern = Arc::new(CsrPattern::from_sparse(&adj)); + let assembler = AValuesAssembler::::new(&pattern, &device); + let weights = conditioning_weights(); + + let eval_loss = |n: &[f32], qsp: &[f32], psp: &[f32], qt: &[f32]| -> f32 { + let (loss, _) = run_forward_loss( + &cfg, &pattern, &assembler, &device, n, qsp, psp, qt, &f.qpt, &f.length, &f.slope, + &weights, None, + ); + loss.sum().into_data().to_vec::().unwrap()[0] + }; + let loss_mag = eval_loss(&f.n, &f.qsp, &f.psp, &f.qt).abs().max(1.0); + + let n_reach = f.length.len(); + let mut grad = vec![0.0f32; n_reach]; + let mut noise = vec![0.0f32; n_reach]; + for i in 0..n_reach { + let (mut pn, mut pq, mut pp, mut pt) = + (f.n.clone(), f.qsp.clone(), f.psp.clone(), f.qt.clone()); + let (mut mn, mut mq, mut mp, mut mt) = + (f.n.clone(), f.qsp.clone(), f.psp.clone(), f.qt.clone()); + let (plus, minus, base) = match parent { + Parent::N => (&mut pn, &mut mn, &f.n), + Parent::QSpatial => (&mut pq, &mut mq, &f.qsp), + Parent::PSpatial => (&mut pp, &mut mp, &f.psp), + Parent::QT => (&mut pt, &mut mt, &f.qt), + }; + // Pure relative step — an absolute floor here would be a 2.9% + // perturbation of `n` and would cross reach 0's min-branch boundary. + let eps = REL_STEP * base[i].abs(); + assert!(eps > 0.0, "reach {i}: zero FD step (base = {})", base[i]); + plus[i] = base[i] + eps; + minus[i] = base[i] - eps; + grad[i] = (eval_loss(&pn, &pq, &pp, &pt) - eval_loss(&mn, &mq, &mp, &mt)) / (2.0 * eps); + noise[i] = NOISE_ULPS * loss_mag * f32::EPSILON / (2.0 * eps); + } + (grad, noise) +} + +fn compare_grads(name: &str, analytical: &[f32], fd: &[f32], noise: &[f32]) { + assert_eq!(analytical.len(), fd.len()); + println!("--- {name} ---"); + let (mut worst_rel, mut worst_abs) = (0.0f32, 0.0f32); + let mut resolved = 0usize; + for i in 0..analytical.len() { + let (a, d) = (analytical[i], fd[i]); + let abs_diff = (a - d).abs(); + let rel_diff = abs_diff / a.abs().max(d.abs()).max(1e-12); + // Only reaches whose gradient is above the FD noise floor carry + // information; the rest are reported but excluded from the worst-case. + let informative = a.abs().max(d.abs()) > noise[i]; + if informative { + resolved += 1; + worst_abs = worst_abs.max(abs_diff); + worst_rel = worst_rel.max(rel_diff); + } + println!( + " [{i}] analytical={a:.6e} fd={d:.6e} abs={abs_diff:.3e} rel={rel_diff:.3e} \ + fd_noise={:.3e}{}", + noise[i], + if informative { "" } else { " (below FD resolution)" } + ); + } + println!(" resolved reaches: {resolved}/{}", analytical.len()); + println!(" worst abs={worst_abs:.3e} worst rel={worst_rel:.3e}"); + // Guard against the noise floor swallowing the whole test. + assert!( + resolved >= 4, + "{name}: only {resolved} reaches are above the FD noise floor — \ + the gradcheck has no power left" + ); + let pass = (0..analytical.len()).all(|i| { + let (a, d) = (analytical[i], fd[i]); + let abs_diff = (a - d).abs(); + let rel_diff = abs_diff / a.abs().max(d.abs()).max(1e-12); + rel_diff < REL_TOL || abs_diff < noise[i] + }); + assert!( + pass, + "{name}: gradcheck failed (worst rel={worst_rel:.3e}, abs={worst_abs:.3e})" + ); +} + +/// Re-runs the non-vacuity guard as a PRECONDITION of every gradcheck, so a +/// future retune of `grad_fixture()` cannot quietly turn these into tautologies. +/// (`fixture_is_not_vacuous` guards the Part-1 fixture; this guards Part 2's.) +fn assert_fixture_exercises_everything() { + let f = grad_fixture(); + let (branches, floored, cr_raw, x_cunge) = branch_report(&f); + for i in 0..branches.len() { + println!( + " [{i}] L={:>8.0} q_t={:>6.1} Cr_raw={:>7.3} floored={:<5} x_cunge={:.4} branch={:?}", + f.length[i], f.qt[i], cr_raw[i], floored[i], x_cunge[i], branches[i] + ); + } + let n = branches.len() as f32; + let frac = |b: Branch| branches.iter().filter(|&&x| x == b).count() as f32 / n; + let (fc, fa, fb) = (frac(Branch::Cunge), frac(Branch::HiA), frac(Branch::HiB)); + let f_floored = floored.iter().filter(|&&x| x).count() as f32 / n; + assert!( + fc > 0.05 && fa > 0.05 && fb > 0.05, + "fixture is vacuous: branch mix cunge={fc:.2} hi_a={fa:.2} hi_b={fb:.2}" + ); + assert!( + (0.05..0.95).contains(&f_floored), + "fixture must straddle the K floor, got {f_floored:.2}" + ); +} + +#[test] +fn gradcheck_positivity_clamp_n() { + assert_fixture_exercises_everything(); + let (fd, noise) = compute_fd_grad(Parent::N); + compare_grads( + "n (positivity clamp)", + &compute_analytical_grad(Parent::N), + &fd, + &noise, + ); +} + +#[test] +fn gradcheck_positivity_clamp_q_spatial() { + assert_fixture_exercises_everything(); + let (fd, noise) = compute_fd_grad(Parent::QSpatial); + compare_grads( + "q_spatial (positivity clamp)", + &compute_analytical_grad(Parent::QSpatial), + &fd, + &noise, + ); +} + +#[test] +fn gradcheck_positivity_clamp_p_spatial() { + assert_fixture_exercises_everything(); + let (fd, noise) = compute_fd_grad(Parent::PSpatial); + compare_grads( + "p_spatial (positivity clamp)", + &compute_analytical_grad(Parent::PSpatial), + &fd, + &noise, + ); +} + +/// `q_t` is the parent the clamp bites hardest: it is the only one whose Cunge +/// path (`∂X/∂Q`) is switched OFF wherever `hi_a`/`hi_b` win the min. +#[test] +fn gradcheck_positivity_clamp_q_t() { + assert_fixture_exercises_everything(); + let (fd, noise) = compute_fd_grad(Parent::QT); + compare_grads( + "q_t (positivity clamp)", + &compute_analytical_grad(Parent::QT), + &fd, + &noise, + ); +} + +/// The assumption the whole B19' mask cascade rests on: the backward +/// RECOMPUTES `x_cunge`, `hi_a` and `hi_b` (they are not saved) and decides the +/// winning branch by comparing them. If that recomputation did not reproduce +/// the forward's `min_pair` chain, gradient would be routed to the wrong branch +/// on exactly the reaches where it matters most. +/// +/// `x_eff` itself is not exported, so it is recovered from the saved +/// `denom = 2K(1−X) + Δt` — an inversion, hence the 1e-4 tolerance rather than +/// a bit comparison. +/// +/// The test also pins the PARTITION property: the `Cunge > hi_a > hi_b` +/// tie-break cascade in `branch_report` mirrors the backward's, and assigning +/// exactly one `Branch` per reach is what makes "no element counted twice, none +/// dropped" true by construction. +#[test] +fn recomputed_min_branches_reproduce_the_forward_x() { + for f in [fixture(), grad_fixture()] { + let out = run_chain(&f, true); + let (branches, _floored, _cr_raw, x_cunge) = branch_report(&f); + for i in 0..f.length.len() { + let k = out.k_muskingum[i]; + // denom = 2K(1-X) + dt => X = 1 - (denom - dt)/(2K) + let x_eff = 1.0 - (out.denom[i] - DT) / (2.0 * k); + let cr = DT / k; + let hi_a = cr * 0.5 * (1.0 - POSITIVITY_DELTA); + let hi_b = (1.0 - 0.5 * cr) * (1.0 - POSITIVITY_DELTA); + let recomputed = x_cunge[i].min(hi_a).min(hi_b); + assert!( + (recomputed - x_eff).abs() < 1e-4, + "reach {i}: backward's min(x_cunge={:.6}, hi_a={hi_a:.6}, hi_b={hi_b:.6}) \ + = {recomputed:.6} but the forward used X = {x_eff:.6}", + x_cunge[i] + ); + // The branch the cascade picked must be the one that attains the min. + let winner = match branches[i] { + Branch::Cunge => x_cunge[i], + Branch::HiA => hi_a, + Branch::HiB => hi_b, + }; + assert_eq!( + winner.to_bits(), + recomputed.to_bits(), + "reach {i}: tie-break picked {:?} = {winner} but the min is {recomputed}", + branches[i] + ); + } + } +} diff --git a/tests/sp8_gradcheck.rs b/tests/sp8_gradcheck.rs index da91594..3edcb8f 100644 --- a/tests/sp8_gradcheck.rs +++ b/tests/sp8_gradcheck.rs @@ -116,6 +116,7 @@ fn run_forward_loss( length_t, slope_t, xst_t, + false, ); ( diff --git a/tests/subdivide.rs b/tests/subdivide.rs new file mode 100644 index 0000000..4e35b9b --- /dev/null +++ b/tests/subdivide.rs @@ -0,0 +1,746 @@ +//! `params.subdivision` — static reach subdivision (variable Δx). +//! +//! Defaults to disabled so every existing config keeps its current behaviour. +use ddrs::adjacency::build::ConusAdjacency; +use ddrs::adjacency::subdivide::{plan_reaches, reference_celerity, subdivide, ReachPlan}; +use ddrs::adjacency::zarr_write::{write_conus_store, write_conus_store_subdivided}; +use ddrs::config::{Config, Subdivision}; +use ddrs::data::{Comid, ConusAdjacencyStore}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use tempfile::TempDir; + +/// A minimal valid training config, with `extra` spliced into `params:`. +/// `extra` must already be indented two spaces (it sits at `params:` depth). +fn yaml_with_params(extra: &str) -> String { + format!( + r#" +mode: training +geodataset: merit +seed: 42 +np_seed: 42 +params: + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] +{extra}"# + ) +} + +/// Inline YAML → tempfile → `Config::from_yaml_file`, the style used by +/// `tests/ddr_match_flag.rs`. The filename is unique per call because cargo +/// runs the tests in this file concurrently. +fn try_load_cfg(yaml: &str) -> Result { + static SEQ: AtomicUsize = AtomicUsize::new(0); + let path = std::env::temp_dir().join(format!( + "ddrs_subdivision_{}_{}.yaml", + std::process::id(), + SEQ.fetch_add(1, Ordering::Relaxed) + )); + std::fs::write(&path, yaml).unwrap(); + Config::from_yaml_file(&path).map_err(|e| e.to_string()) +} + +fn load_cfg(yaml: &str) -> Config { + try_load_cfg(yaml).expect("parse") +} + +#[test] +fn subdivision_defaults_to_disabled() { + let cfg = load_cfg(&yaml_with_params("")); + assert!(!cfg.params.subdivision.enabled); + assert_eq!(cfg.params.subdivision.max_pieces, 8); +} + +#[test] +fn subdivision_rejects_max_pieces_below_one() { + let err = try_load_cfg(&yaml_with_params( + " subdivision:\n enabled: true\n max_pieces: 0\n", + )) + .expect_err("must reject"); + assert!(err.to_string().contains("max_pieces"), "got: {err}"); +} + +#[test] +fn subdivision_enabled_loads() { + let cfg = load_cfg(&yaml_with_params( + " subdivision:\n enabled: true\n max_pieces: 4\n", + )); + assert!(cfg.params.subdivision.enabled); + assert_eq!(cfg.params.subdivision.max_pieces, 4); +} + +// --------------------------------------------------------------------------- +// Task 2: reference celerity + the two-sided reach plan +// --------------------------------------------------------------------------- + +fn cfg(max_pieces: usize) -> Subdivision { + Subdivision { + enabled: true, + max_pieces, + ..Default::default() + } +} + +#[test] +fn celerity_rises_with_slope_and_area() { + let c = cfg(8); + let lo = reference_celerity(100.0, 1e-3, &c); + assert!( + reference_celerity(100.0, 1e-2, &c) > lo, + "steeper must be faster" + ); + assert!( + reference_celerity(10_000.0, 1e-3, &c) > lo, + "bigger must be faster" + ); + assert!( + lo > 0.0 && lo < 15.0, + "celerity {lo} outside physical range" + ); +} + +#[test] +fn long_reaches_split_and_are_capped() { + let c = cfg(4); + let p = plan_reaches(&[200_000.0], &[1e-4], &[30.0], 3600.0, &c); + assert_eq!(p.pieces[0], 4, "must clamp to max_pieces"); + assert_eq!( + p.length_m[0], 200_000.0, + "long reaches keep their true length" + ); +} + +#[test] +fn short_reaches_are_length_clamped_not_split() { + let c = cfg(8); + let dx = reference_celerity(10_000.0, 1e-2, &c) * 3600.0; + // The fixture must sit INSIDE `max_clamp_factor`, or that bound binds first + // and we would be testing the bound instead of the clamp. Half of dx_target + // is a 2x stretch, comfortably under the default 4x ceiling. + let raw = dx / 2.0; + let p = plan_reaches(&[raw], &[1e-2], &[10_000.0], 3600.0, &c); + assert_eq!(p.pieces[0], 1, "short reach must not split"); + assert!( + (p.length_m[0] - dx).abs() < 1e-3, + "expected clamp to dx_target {dx}, got {}", + p.length_m[0] + ); + assert!(p.length_m[0] > raw, "clamp must lengthen, not shorten"); +} + +#[test] +fn min_length_fraction_zero_disables_the_clamp() { + let mut c = cfg(8); + c.min_length_fraction = 0.0; + let p = plan_reaches(&[50.0], &[1e-2], &[10_000.0], 3600.0, &c); + assert_eq!(p.length_m[0], 50.0, "clamp must be off"); +} + +#[test] +fn disabled_is_an_exact_no_op() { + let mut c = cfg(8); + c.enabled = false; + let p = plan_reaches( + &[200_000.0, 50.0], + &[1e-4, 1e-2], + &[30.0, 10_000.0], + 3600.0, + &c, + ); + assert_eq!(p.pieces, vec![1, 1]); + assert_eq!( + p.length_m, + vec![200_000.0, 50.0], + "lengths must be untouched" + ); +} + +#[test] +fn degenerate_input_never_yields_zero_pieces_or_zero_length() { + let c = cfg(8); + let p = plan_reaches(&[5000.0, 0.0], &[0.0, 0.0], &[0.0, 0.0], 3600.0, &c); + assert!( + p.pieces.iter().all(|&v| v >= 1), + "pieces must be >= 1, got {:?}", + p.pieces + ); + assert!( + p.length_m.iter().all(|&v| v > 0.0), + "length must be > 0 (a 0 m reach gives K = 0 and c1 = 1), got {:?}", + p.length_m + ); +} + +// ── Bounds added after Task 2 review ──────────────────────────────────────── +// Two holes the original six tests did not cover. + +#[test] +fn zero_length_reach_survives_min_length_fraction_zero() { + // The combination `min_length_fraction: 0.0` + a 0 m reach previously gave + // l_eff = 0, hence K = L/c = 0 and c1 = 1, breaking the solve. MERIT + // contains sub-10 m reaches, so this is reachable, not hypothetical. + let mut c = cfg(8); + c.min_length_fraction = 0.0; + let p = plan_reaches(&[0.0], &[1e-3], &[100.0], 3600.0, &c); + assert!( + p.length_m[0] >= 1.0, + "absolute floor must apply even with the clamp disabled, got {}", + p.length_m[0] + ); + assert_eq!(p.pieces[0], 1); +} + +#[test] +fn clamp_cannot_stretch_a_reach_beyond_max_clamp_factor() { + // A steep headwater: dx_target is large, but a 100 m reach must not be + // rewritten into a multi-km channel. Unbounded, measured clamp factors + // reached p99 = 36x and max = 48,597x. + let mut c = cfg(8); + c.max_clamp_factor = 4.0; + let p = plan_reaches(&[100.0], &[1e-2], &[100.0], 3600.0, &c); + assert!( + p.length_m[0] <= 100.0 * 4.0 + 1e-3, + "stretched {}x, exceeding max_clamp_factor", + p.length_m[0] / 100.0 + ); + assert!(p.length_m[0] >= 100.0, "clamp must never shorten a reach"); +} + +#[test] +fn reference_celerity_stays_in_a_physical_flood_wave_band() { + let c = cfg(8); + // Steep + large: the case that previously produced 8.9 m/s and a 32 km dx. + let fast = reference_celerity(10_000.0, 1e-2, &c); + assert!( + (0.05..=5.0).contains(&fast), + "celerity {fast} m/s outside the physical flood-wave band" + ); + assert!( + fast * 3600.0 <= 18_000.0 + 1.0, + "dx_target {} m is larger than any plausible MERIT reach", + fast * 3600.0 + ); +} + +// --------------------------------------------------------------------------- +// Task 3: graph expansion +// --------------------------------------------------------------------------- + +/// Chain of 3 reaches: 0 -> 1 -> 2 (rows = downstream, cols = upstream, +/// so rows[k] >= cols[k]). +fn chain3() -> ConusAdjacency { + ConusAdjacency { + order: vec![100, 200, 300], + rows: vec![1, 2], + cols: vec![0, 1], + length_m: vec![3000.0, 6000.0, 900.0], + slope: vec![1e-3, 2e-3, 3e-3], + dropped_comids: vec![], + } +} + +/// Explicit plan so these tests exercise expansion alone, independent of the +/// celerity heuristic in the two-sided rule above. +fn plan(pieces: Vec, length_m: Vec) -> ReachPlan { + ReachPlan { pieces, length_m } +} + +/// Row → parent position, derived from `parent_offset`. Used to collapse the +/// expanded COO back to parent space. +fn row_parents(offset: &[i32]) -> Vec { + let mut out = Vec::new(); + for p in 0..offset.len() - 1 { + for _ in offset[p]..offset[p + 1] { + out.push(p); + } + } + out +} + +#[test] +fn expansion_preserves_total_length_and_slope() { + let s = subdivide(&chain3(), &plan(vec![3, 2, 1], chain3().length_m.clone())); + assert_eq!(s.length_m.len(), 6); + for (p, &m) in [3u32, 2, 1].iter().enumerate() { + let lo = s.parent_offset[p] as usize; + let hi = s.parent_offset[p + 1] as usize; + assert_eq!(hi - lo, m as usize, "parent {p} piece count"); + let total: f32 = s.length_m[lo..hi].iter().sum(); + assert!( + (total - chain3().length_m[p]).abs() < 1e-3, + "parent {p} length not conserved: {total}" + ); + assert!( + s.slope[lo..hi].iter().all(|&v| v == chain3().slope[p]), + "slope must be inherited unchanged" + ); + } +} + +#[test] +fn expansion_stays_lower_triangular_and_topological() { + let s = subdivide(&chain3(), &plan(vec![3, 2, 1], chain3().length_m.clone())); + for (&r, &c) in s.rows.iter().zip(s.cols.iter()) { + assert!(r > c, "edge {c}->{r} violates strict lower-triangular ordering"); + } +} + +#[test] +fn expansion_edge_count_is_original_plus_internal_links() { + let s = subdivide(&chain3(), &plan(vec![3, 2, 1], chain3().length_m.clone())); + // 2 original edges + (3-1) + (2-1) + (1-1) internal = 5 + assert_eq!(s.rows.len(), 5, "rows: {:?} cols: {:?}", s.rows, s.cols); +} + +#[test] +fn external_edges_land_on_parent_outlet_and_inlet() { + let s = subdivide(&chain3(), &plan(vec![3, 2, 1], chain3().length_m.clone())); + // parent0 rows 0..3, parent1 rows 3..5, parent2 row 5. + // edge 0->1 becomes outlet(0)=2 -> inlet(1)=3 + assert!( + s.rows.iter().zip(&s.cols).any(|(&r, &c)| c == 2 && r == 3), + "missing 2->3; rows {:?} cols {:?}", + s.rows, + s.cols + ); + // edge 1->2 becomes outlet(1)=4 -> inlet(2)=5 + assert!( + s.rows.iter().zip(&s.cols).any(|(&r, &c)| c == 4 && r == 5), + "missing 4->5; rows {:?} cols {:?}", + s.rows, + s.cols + ); +} + +#[test] +fn all_ones_is_an_exact_identity() { + let a = chain3(); + let s = subdivide(&a, &plan(vec![1, 1, 1], a.length_m.clone())); + assert_eq!(s.order, a.order); + assert_eq!(s.rows, a.rows); + assert_eq!(s.cols, a.cols); + assert_eq!(s.length_m, a.length_m); + assert_eq!(s.parent_offset, vec![0, 1, 2, 3]); +} + +#[test] +fn expansion_uses_the_clamped_length_not_the_raw_one() { + let a = chain3(); // reach 2 is only 900 m + let clamped = vec![3000.0, 6000.0, 4000.0]; // reach 2 stretched to 4 km + let s = subdivide(&a, &plan(vec![1, 1, 1], clamped)); + assert_eq!( + s.length_m[2], 4000.0, + "must use ReachPlan.length_m, not ConusAdjacency.length_m" + ); +} + +// ── Bounds added after Task 3 review ──────────────────────────────────────── +// A 3-reach chain is too easy a fixture to trust for the invariant that the +// forward-substitution solver silently depends on: it has no junction, and its +// piece counts happen to be monotonically decreasing. These three tests use a +// branching network instead. + +/// Two headwaters into a confluence, then an outlet — the shape a pure chain +/// cannot exercise. Topologically ordered, strictly lower triangular. +/// +/// ```text +/// COMID 10 (p0) ─┐ +/// ├─> COMID 30 (p2) ─> COMID 40 (p3) +/// COMID 20 (p1) ─┘ +/// ``` +fn junction4() -> ConusAdjacency { + ConusAdjacency { + order: vec![10, 20, 30, 40], + // rows = downstream, cols = upstream. + rows: vec![2, 2, 3], + cols: vec![0, 1, 2], + length_m: vec![1000.0, 2000.0, 4000.0, 8000.0], + slope: vec![5e-3, 4e-3, 3e-3, 1e-3], + dropped_comids: vec![], + } +} + +#[test] +fn junction_expansion_stays_strictly_lower_triangular() { + // Piece counts deliberately non-monotonic (2, 3, 1, 4) so the invariant + // cannot pass by an accident of ordering, and include an unsplit reach. + let a = junction4(); + let s = subdivide(&a, &plan(vec![2, 3, 1, 4], a.length_m.clone())); + assert_eq!(s.parent_offset, vec![0, 2, 5, 6, 10]); + assert_eq!(s.order.len(), 10); + for (&r, &c) in s.rows.iter().zip(s.cols.iter()) { + assert!( + r > c, + "edge {c}->{r} violates strict lower-triangular ordering; \ + rows {:?} cols {:?}", + s.rows, + s.cols + ); + } + // 3 external + (2-1) + (3-1) + (1-1) + (4-1) = 3 + 6 = 9 edges. + assert_eq!(s.rows.len(), 9); + // The confluence must survive: the inlet of parent 2 (row 5) still has two + // distinct upstream rows — the outlets of parents 0 and 1. + let ups: Vec = s + .rows + .iter() + .zip(&s.cols) + .filter(|(&r, _)| r == 5) + .map(|(_, &c)| c) + .collect(); + assert_eq!(ups, vec![s.outlet(0) as i32, s.outlet(1) as i32]); +} + +#[test] +fn collapsing_subedges_to_parents_reproduces_the_parent_edge_set() { + // The direction check that lower-triangularity CANNOT make: mapping + // `inlet(u) -> outlet(p)` instead of `outlet(u) -> inlet(p)` is still + // lower triangular, and reversing rows/cols on a general DAG can be too. + // Collapsing every sub-edge back to (parent_of_row, parent_of_col) must + // reproduce the original (rows, cols) pairs exactly, in order. + let a = junction4(); + let s = subdivide(&a, &plan(vec![2, 3, 1, 4], a.length_m.clone())); + let owner = row_parents(&s.parent_offset); + + let mut external: Vec<(usize, usize)> = Vec::new(); + for (&r, &c) in s.rows.iter().zip(s.cols.iter()) { + let (pr, pc) = (owner[r as usize], owner[c as usize]); + if pr == pc { + // Internal chain link: must join consecutive pieces of one parent. + assert_eq!(r, c + 1, "internal link {c}->{r} is not consecutive"); + } else { + external.push((pr, pc)); + // And it must be anchored at the true outlet/inlet, not any + // interior piece — otherwise part of the reach is bypassed. + assert_eq!(c as usize, s.outlet(pc), "upstream end is not the outlet"); + assert_eq!(r as usize, s.inlet(pr), "downstream end is not the inlet"); + } + } + let expected: Vec<(usize, usize)> = a + .rows + .iter() + .zip(&a.cols) + .map(|(&r, &c)| (r as usize, c as usize)) + .collect(); + assert_eq!( + external, expected, + "collapsed edge set does not match the parent graph — the flow \ + direction was inverted somewhere" + ); +} + +#[test] +fn every_subreach_row_carries_its_parents_comid() { + let a = junction4(); + let s = subdivide(&a, &plan(vec![2, 3, 1, 4], a.length_m.clone())); + assert_eq!(s.parent_order, a.order, "parent space must be untouched"); + for (row, &p) in row_parents(&s.parent_offset).iter().enumerate() { + assert_eq!( + s.order[row], a.order[p], + "row {row} should carry COMID of parent {p}" + ); + assert_eq!(s.slope[row], a.slope[p], "row {row} slope"); + } + for p in 0..a.order.len() { + assert_eq!(s.inlet(p), s.parent_offset[p] as usize); + assert_eq!(s.outlet(p), s.parent_offset[p + 1] as usize - 1); + assert_eq!(s.pieces(p), s.outlet(p) - s.inlet(p) + 1); + } +} + +// ─────────────────────────── Task 4: persist + cache key ────────────────────── + +/// Build → subdivide → write zarr into a `TempDir` → reload through the real +/// reader. Uses `write_conus_store_subdivided`, the same writer the managed +/// cache calls, so the test exercises the production path rather than a mock. +fn round_trip_store(adj: &ConusAdjacency, pieces: &[u32]) -> (ConusAdjacencyStore, TempDir) { + let s = subdivide(adj, &plan(pieces.to_vec(), adj.length_m.clone())); + let dir = tempfile::tempdir().expect("tempdir"); + write_conus_store_subdivided(&s, dir.path()).expect("write"); + let store = ConusAdjacencyStore::open(dir.path()).expect("open"); + // The TempDir must outlive the store's use; returning it keeps it alive. + (store, dir) +} + +#[test] +fn cache_key_changes_with_every_subdivision_field() { + use ddrs::adjacency::cache::content_key_for_test as key; + + let on = Subdivision { + enabled: true, + ..Default::default() + }; + let base = key("fab", "gag", None, &Subdivision::default()); + assert_ne!(base, key("fab", "gag", None, &on), "enabling must invalidate"); + + // Every field that feeds `plan_reaches` must invalidate, or a config edit + // silently reuses a graph built with different geometry — a failure that + // reads as a physics result rather than as a bug. + for (name, modified) in [ + ( + "max_pieces", + Subdivision { + max_pieces: 4, + ..on.clone() + }, + ), + ( + "reference_n", + Subdivision { + reference_n: 0.03, + ..on.clone() + }, + ), + ( + "q_coeff", + Subdivision { + reference_discharge_coefficient: 0.02, + ..on.clone() + }, + ), + ( + "q_exp", + Subdivision { + reference_discharge_exponent: 0.8, + ..on.clone() + }, + ), + ( + "min_len_fr", + Subdivision { + min_length_fraction: 0.5, + ..on.clone() + }, + ), + ( + "max_clamp_factor", + Subdivision { + max_clamp_factor: 3.0, + ..on.clone() + }, + ), + ] { + assert_ne!( + key("fab", "gag", None, &on), + key("fab", "gag", None, &modified), + "changing {name} must invalidate the cache" + ); + } +} + +#[test] +fn cache_key_is_stable_for_identical_subdivision() { + use ddrs::adjacency::cache::content_key_for_test as key; + let s = Subdivision { + enabled: true, + max_pieces: 4, + ..Default::default() + }; + assert_eq!(key("fab", "gag", None, &s), key("fab", "gag", None, &s)); +} + +#[test] +fn store_index_maps_comid_to_parent_not_subreach() { + // Built with pieces [3, 2, 1]; COMID 200 is parent 1. + let (store, _dir) = round_trip_store(&chain3(), &[3, 2, 1]); + assert_eq!(store.parent_offset, vec![0, 3, 5, 6]); + assert_eq!(store.n, 6, "sub-reach count"); + assert_eq!(store.n_parent(), 3, "parent count"); + assert_eq!( + store.parent_order, + vec![Comid(100), Comid(200), Comid(300)], + "parent space must hold each COMID exactly once" + ); + // `order` carries duplicates — this is why the index cannot come from it. + assert_eq!(store.order.len(), 6); + assert_eq!(store.order[0], store.order[1]); + + let p = store.index.position(&Comid(200)).expect("COMID 200 must resolve"); + assert_eq!( + p, 1, + "index must return the PARENT position, not a sub-reach row" + ); + // A gauge on COMID 200 reads row 4, the parent's outlet. + assert_eq!(store.outlet_row(p), 4); +} + +#[test] +fn round_tripped_store_matches_the_expansion_element_for_element() { + let a = chain3(); + let s = subdivide(&a, &plan(vec![3, 2, 1], a.length_m.clone())); + let (store, _dir) = round_trip_store(&a, &[3, 2, 1]); + assert_eq!(store.indices_0, s.rows); + assert_eq!(store.indices_1, s.cols); + assert_eq!(store.length_m.to_vec(), s.length_m); + assert_eq!(store.slope.to_vec(), s.slope); + assert_eq!(store.parent_offset, s.parent_offset); + assert_eq!(store.nnz, s.rows.len()); +} + +/// Subdivision off must be indistinguishable from the pre-subdivision store: +/// the parent map degenerates to the identity and `index` still resolves every +/// COMID to its own row. +#[test] +fn disabled_subdivision_round_trips_to_the_identity_parent_map() { + let a = chain3(); + let (store, _dir) = round_trip_store(&a, &[1, 1, 1]); + assert_eq!(store.n, 3); + assert_eq!(store.parent_offset, vec![0, 1, 2, 3]); + assert_eq!(store.parent_order, store.order); + for (i, c) in [Comid(100), Comid(200), Comid(300)].iter().enumerate() { + assert_eq!(store.index.position(c), Some(i)); + assert_eq!(store.outlet_row(i), i); + } +} + +/// A store written WITHOUT the parent map — i.e. every pre-existing cache and +/// every engine export — must keep loading, with the identity synthesized. +#[test] +fn store_without_parent_arrays_synthesizes_the_identity() { + let a = chain3(); + let dir = tempfile::tempdir().expect("tempdir"); + // `write_conus_store` is the pre-subdivision writer: no parent arrays. + write_conus_store(&a, dir.path()).expect("write legacy"); + assert!( + !dir.path().join("parent_order").exists(), + "fixture must genuinely lack the parent map" + ); + + let store = ConusAdjacencyStore::open(dir.path()).expect("open legacy store"); + assert_eq!(store.n, 3); + assert_eq!(store.n_parent(), 3); + assert_eq!(store.parent_order, store.order); + assert_eq!(store.parent_offset, vec![0, 1, 2, 3]); + assert_eq!(store.index.position(&Comid(300)), Some(2)); + assert_eq!(store.outlet_row(2), 2); +} + +/// The real pre-subdivision CONUS store on disk. Read-only; skipped when the +/// path is absent so a clean checkout still passes. +#[test] +fn real_pre_subdivision_conus_store_still_loads() { + const REAL: &str = "/home/tbindas/projects/ddr/data/merit_conus_adjacency.zarr"; + if !std::path::Path::new(REAL).exists() { + eprintln!("skipping: {REAL} not present"); + return; + } + let store = ConusAdjacencyStore::open(REAL).expect("real store must still open"); + assert_eq!(store.n, 346_321); + assert_eq!(store.nnz, 338_814); + assert_eq!(store.n_parent(), store.n, "no parent map on disk → identity"); + assert_eq!(store.parent_order, store.order); + assert_eq!(store.parent_offset.len(), store.n + 1); + assert_eq!(store.parent_offset[0], 0); + assert_eq!(*store.parent_offset.last().unwrap(), store.n as i32); + // The identity map must leave COMID lookups exactly where they were. + let probe = store.order[12_345]; + assert_eq!(store.index.position(&probe), Some(12_345)); + assert_eq!(store.outlet_row(12_345), 12_345); +} + +// ── the silent-inertness guard ─────────────────────────────────────────────── +// +// `cli::plan::resolve_adjacency` only reaches the managed builder — and hence +// `subdivide` — when `data_sources` carries NO explicit adjacency paths. With +// them set, `params.subdivision.enabled: true` would be a silent no-op: the run +// would route the un-split network while its manifest claimed otherwise. Config +// load rejects that combination, EXCEPT when the store really is subdivided. + +/// A config with a full `data_sources` block. `adjacency` is spliced in as-is +/// (two-space indented), so a test can supply explicit paths or a fabric. +fn yaml_with_sources(adjacency: &str, subdivision: &str) -> String { + format!( + r#" +mode: training +geodataset: merit +seed: 42 +np_seed: 42 +data_sources: + attributes: /tmp/attrs.nc + streamflow: /tmp/streamflow.ic + observations: /tmp/obs + gages: /tmp/gages.csv +{adjacency} +params: + parameter_ranges: + n: [0.015, 0.25] + q_spatial: [0.0, 1.0] + p_spatial: [1.0, 200.0] +{subdivision}"# + ) +} + +const SUBDIV_ON: &str = " subdivision:\n enabled: true\n"; + +#[test] +fn subdivision_rejects_explicit_adjacency_paths() { + let a = chain3(); + let dir = tempfile::tempdir().expect("tempdir"); + // Written by the pre-subdivision writer: no parent map at all. + write_conus_store(&a, dir.path()).expect("write legacy store"); + let adjacency = format!( + " conus_adjacency: {}\n gages_adjacency: {}\n", + dir.path().display(), + dir.path().display() + ); + + let err = try_load_cfg(&yaml_with_sources(&adjacency, SUBDIV_ON)) + .expect_err("explicit adjacency + subdivision must be rejected"); + assert!( + err.to_string().contains("SILENTLY INERT"), + "error must name the conflict, got: {err}" + ); + assert!( + err.to_string().contains("geospatial_fabric"), + "error must point at the remedy, got: {err}" + ); +} + +#[test] +fn subdivision_rejects_explicit_paths_to_an_identity_parent_map() { + // A managed store built with subdivision OFF still carries `parent_order`, + // but at n_parent == n. Asking to subdivide against it is just as inert. + let (_store, dir) = round_trip_store(&chain3(), &[1, 1, 1]); + let adjacency = format!( + " conus_adjacency: {}\n gages_adjacency: {}\n", + dir.path().display(), + dir.path().display() + ); + let err = try_load_cfg(&yaml_with_sources(&adjacency, SUBDIV_ON)) + .expect_err("identity parent map is not a subdivided store"); + assert!(err.to_string().contains("n_parent == n"), "got: {err}"); +} + +#[test] +fn subdivision_accepts_an_explicit_path_to_an_already_subdivided_store() { + let (_store, dir) = round_trip_store(&chain3(), &[3, 2, 1]); + let adjacency = format!( + " conus_adjacency: {}\n gages_adjacency: {}\n", + dir.path().display(), + dir.path().display() + ); + let cfg = try_load_cfg(&yaml_with_sources(&adjacency, SUBDIV_ON)) + .expect("a genuinely subdivided store is a legitimate explicit path"); + assert!(cfg.params.subdivision.enabled); +} + +#[test] +fn subdivision_accepts_a_fabric_only_config() { + let cfg = try_load_cfg(&yaml_with_sources( + " geospatial_fabric: /tmp/riv.dbf\n", + SUBDIV_ON, + )) + .expect("fabric-only reaches the managed builder"); + assert!(cfg.params.subdivision.enabled); +} + +#[test] +fn explicit_adjacency_paths_are_fine_when_subdivision_is_off() { + let cfg = try_load_cfg(&yaml_with_sources( + " conus_adjacency: /tmp/does_not_exist.zarr\n gages_adjacency: /tmp/nope.zarr\n", + "", + )) + .expect("the guard must only fire when subdivision is enabled"); + assert!(!cfg.params.subdivision.enabled); +} diff --git a/tests/subdivision_integration.rs b/tests/subdivision_integration.rs new file mode 100644 index 0000000..edc21c6 --- /dev/null +++ b/tests/subdivision_integration.rs @@ -0,0 +1,531 @@ +//! End-to-end checks for reach subdivision inside the routing core. +//! +//! A reach split into `m` pieces must have its lateral inflow split `m` ways: +//! the pieces are `L/m` long and chain in series, so each carries `q'/m` and +//! the parent's outlet piece still discharges the whole reach's runoff. This +//! mirrors HEC-HMS, whose lateral term is `C4·(q_L·Δx)` with `q_L` an inflow +//! per unit length. +//! +//! ```text +//! 1 piece 4 pieces +//! ┌──────────────┐ ┌────┬────┬────┬────┐ +//! │ q' │──> Q │q'/4│q'/4│q'/4│q'/4│──> Q +//! └──────────────┘ └────┴────┴────┴────┘ +//! L = 4 km each 1 km, Σq' = q' +//! ``` +//! +//! Both configurations must reach the SAME steady-state outflow. If the split +//! is missing, the 4-piece network manufactures 4× the mass. + +use std::collections::HashMap; + +use burn::backend::{Autodiff, NdArray}; +use burn::tensor::{Tensor, TensorData}; + +use ddrs::adjacency::build::ConusAdjacency; +use ddrs::adjacency::subdivide::{subdivide, ReachPlan}; +use ddrs::adjacency::zarr_write::write_conus_store_subdivided; +use ddrs::config::Config; +use ddrs::data::collate::UnionedCoo; +use ddrs::data::dataset::slice_reach_geometry; +use ddrs::data::{compress, ConusAdjacencyStore, Staid}; +use ddrs::routing::{MuskingumCunge, RoutingInputs, SpatialParameters}; +use ddrs::sparse::SparseAdjacency; +use ddrs::training::forward::gather_params_to_subreaches; + +type I = NdArray; +type AB = Autodiff; + +/// Total length of the single parent reach, metres. Split evenly across pieces. +const PARENT_LENGTH_M: f32 = 4000.0; +/// Long enough for a 4 km / 0.1%-slope chain at dt = 3600 s to settle. +const T: usize = 256; +/// Constant lateral inflow for the parent reach, m³/s. +const Q_PRIME: f32 = 10.0; + +/// Same knobs as `tests/gauge_mass_conservation.rs::mock_cfg` — puts +/// velocity/depth in well-conditioned, non-saturated regimes for this network. +fn mock_cfg() -> Config { + let mut cfg = Config::default(); + cfg.params.parameter_ranges.n = [0.01, 0.1]; + cfg.params.parameter_ranges.q_spatial = [0.1, 0.9]; + cfg.params.parameter_ranges.p_spatial = [1.0, 200.0]; + cfg.params.attribute_minimums.velocity = 0.1; + cfg.params.attribute_minimums.depth = 0.01; + cfg.params.attribute_minimums.discharge = 0.001; + cfg.params.attribute_minimums.bottom_width = 0.1; + cfg.params.attribute_minimums.slope = 0.001; + cfg.params.defaults.insert("p_spatial".to_string(), 1.0); + cfg.params.log_space_parameters = vec![]; + cfg +} + +/// One parent reach expanded into `pieces` sub-reaches chained +/// upstream→downstream: `0 → 1 → ... → pieces-1`. Lower-triangular and +/// topologically ordered, exactly as `adjacency::subdivide` emits. +/// +/// `with_parent_map` controls whether the engine is told about the expansion; +/// `false` is the "forgot to split q'" control. +fn subdivided_reach(pieces: usize, with_parent_map: bool) -> SparseAdjacency { + let n = pieces; + let mut dense = vec![0.0_f32; n * n]; + for i in 1..n { + dense[i * n + (i - 1)] = 1.0; // adj[i, i-1]: piece i-1 flows into piece i + } + let piece_len = PARENT_LENGTH_M / pieces as f32; + let mut adj = SparseAdjacency::from_dense(n, &dense, vec![piece_len; n], vec![0.001; n]); + if with_parent_map { + // A single parent owning all `pieces` rows. + adj.parent_offset = Some(vec![0, pieces as i32]); + } + adj +} + +/// Route to steady state and return the outlet piece's discharge (m³/s). +fn steady_state_outflow(pieces: usize, with_parent_map: bool) -> f32 { + let row = outlet_series(pieces, with_parent_map, None); + row[T - 1] +} + +/// The outlet piece's full discharge series (length `T`). Column 0 is the +/// cold-start `Q_0` that `setup_inputs` solved, so `[0]` measures the initial +/// condition and `[T-1]` the steady state. +/// +/// `divide_hotstart` overrides `MuskingumCunge::divide_hotstart_by_pieces`; +/// `None` leaves the shipped default in place. +fn outlet_series(pieces: usize, with_parent_map: bool, divide_hotstart: Option) -> Vec { + let n = pieces; + let device = ::Device::default(); + + // Constant q_prime on every sub-reach row, shape (T, N). The engine is what + // divides it by the piece count — the caller always supplies the parent's + // full lateral inflow, exactly as the q'-store read does. + let q_prime: Tensor = + Tensor::from_data(TensorData::new(vec![Q_PRIME; T * n], [T, n]), &device); + + // Mid-range normalized params → n ≈ 0.055, q_spatial ≈ 0.5, p_spatial ≈ 100. + let mk = |v: f32| -> Tensor { Tensor::from_floats(vec![v; n].as_slice(), &device) }; + let mut engine = MuskingumCunge::::new(mock_cfg(), device.clone()); + if let Some(d) = divide_hotstart { + engine.divide_hotstart_by_pieces = d; + } + engine.setup_inputs( + RoutingInputs { + adjacency: subdivided_reach(pieces, with_parent_map), + x_storage: mk(0.3), + }, + q_prime, + SpatialParameters { + n: mk(0.5), + q_spatial: mk(0.5), + p_spatial: Some(mk(0.5)), + k_d: None, + d_gw: None, + leakance_factor: None, + impervious_mask: None, + }, + false, + None, + ); + + // (N, T) routed discharge; the parent's outlet is the LAST piece. + let runoff: Tensor = engine.forward().inner(); + let all: Vec = runoff.into_data().to_vec::().expect("runoff to host"); + all[(n - 1) * T..n * T].to_vec() +} + +/// A 1-reach network with constant q' must reach the SAME steady-state outflow +/// whether or not it is subdivided — the pieces split the inflow m ways but +/// chain in series, so the outlet still carries the whole reach's runoff. +#[test] +fn subdivision_conserves_mass_at_steady_state() { + let un_split = steady_state_outflow(1, true); + let split = steady_state_outflow(4, true); + println!("1 piece = {un_split:.6} m3/s 4 pieces = {split:.6} m3/s q' = {Q_PRIME}"); + + assert!( + (un_split - Q_PRIME).abs() / Q_PRIME < 1e-3, + "control drifted: 1-piece steady state {un_split} != q' {Q_PRIME}" + ); + assert!( + (split - un_split).abs() / un_split < 1e-3, + "mass not conserved: 1 piece = {un_split}, 4 pieces = {split} \ + (ratio {:.4}; a ratio near 4 means q' was not divided by the piece count)", + split / un_split, + ); +} + +/// Proves the test above discriminates: without the parent map the engine +/// cannot know the reach was split, so every piece receives the parent's full +/// q' and the outlet manufactures `pieces ×` the mass. +#[test] +fn without_the_parent_map_a_split_reach_manufactures_mass() { + let split_unaware = steady_state_outflow(4, false); + println!("4 pieces, no parent map = {split_unaware:.6} m3/s"); + assert!( + (split_unaware - 4.0 * Q_PRIME).abs() / (4.0 * Q_PRIME) < 1e-3, + "expected 4x mass ({}), got {split_unaware} — if this now conserves \ + mass the divisor is being applied without a parent map", + 4.0 * Q_PRIME, + ); +} + +/// The disabled path must be a true no-op: an identity parent map (one row per +/// parent, `0..=n`) is byte-identical to no parent map at all, because +/// `pieces_per_row_divisor` returns `None` rather than a tensor of ones. +#[test] +fn identity_parent_map_is_bit_identical_to_none() { + let n = 4; + let device = ::Device::default(); + let run = |parent_offset: Option>| -> Vec { + let q_prime: Tensor = + Tensor::from_data(TensorData::new(vec![Q_PRIME; T * n], [T, n]), &device); + let mk = |v: f32| -> Tensor { Tensor::from_floats(vec![v; n].as_slice(), &device) }; + let mut adj = subdivided_reach(n, false); + adj.parent_offset = parent_offset; + let mut engine = MuskingumCunge::::new(mock_cfg(), device.clone()); + engine.setup_inputs( + RoutingInputs { adjacency: adj, x_storage: mk(0.3) }, + q_prime, + SpatialParameters { + n: mk(0.5), + q_spatial: mk(0.5), + p_spatial: Some(mk(0.5)), + k_d: None, + d_gw: None, + leakance_factor: None, + impervious_mask: None, + }, + false, + None, + ); + engine + .forward() + .inner() + .into_data() + .to_vec::() + .expect("runoff to host") + }; + + // Here the four rows are four independent PARENTS that happen to chain. + let identity = run(Some((0..=n as i32).collect())); + let none = run(None); + assert_eq!( + identity.iter().map(|v| v.to_bits()).collect::>(), + none.iter().map(|v| v.to_bits()).collect::>(), + "identity parent map must not perturb a single bit" + ); +} + +// --------------------------------------------------------------------------- +// Per-row channel geometry under subdivision +// --------------------------------------------------------------------------- + +/// Chain of 3 reaches `0 → 1 → 2` with deliberately DISTINCT lengths and +/// slopes, so a row that reads the wrong index is visible rather than lucky. +fn chain3() -> ConusAdjacency { + ConusAdjacency { + order: vec![100, 200, 300], + rows: vec![1, 2], + cols: vec![0, 1], + length_m: vec![3000.0, 6000.0, 900.0], + slope: vec![1e-3, 2e-3, 3e-3], + dropped_comids: vec![], + } +} + +/// `chain3` subdivided into `pieces`, written to a zarr store and reloaded — +/// the same round trip `ddrs plan` performs, so `parent_order` / `parent_offset` +/// come back through the real reader. +fn subdivided_store(pieces: &[u32]) -> (ConusAdjacencyStore, tempfile::TempDir) { + let a = chain3(); + let s = subdivide( + &a, + &ReachPlan { + pieces: pieces.to_vec(), + length_m: a.length_m.clone(), + }, + ); + let dir = tempfile::tempdir().expect("tempdir"); + write_conus_store_subdivided(&s, dir.path()).expect("write"); + let store = ConusAdjacencyStore::open(dir.path()).expect("open"); + (store, dir) +} + +/// Every sub-reach row must carry its OWN geometry: length `L/m` (of its +/// parent's possibly-clamped total) and its parent's slope. +/// +/// This is the regression for the parent/sub-reach index mix-up: the geometry +/// arrays are in SUB-REACH space, but `ConusAdjacencyStore::index` resolves a +/// COMID to a PARENT position. Slicing with the latter gives row `i` the +/// geometry of parent-position `i`, which is only correct when the two spaces +/// coincide — i.e. everywhere except under subdivision, which is what made the +/// bug silent. +#[test] +fn subdivided_rows_get_their_own_length_and_slope() { + let pieces = [3u32, 2, 1]; + let (store, _dir) = subdivided_store(&pieces); + assert_eq!(store.parent_offset, vec![0, 3, 5, 6]); + + // Activate the whole network: every subdivided edge plus a gauge on the + // outlet piece of the last parent. + let edges: Vec<(usize, usize)> = store + .indices_0 + .iter() + .zip(store.indices_1.iter()) + .map(|(&r, &c)| (r as usize, c as usize)) + .collect(); + let unioned = UnionedCoo { + edges, + gauges: vec![(Staid::new("gauge"), store.outlet_row(2), "300".to_string())], + }; + let compressed = compress(&unioned, &store.order, false, Some(&store.parent_offset)) + .expect("compress"); + assert_eq!( + compressed.divide_comids.len(), + 6, + "all six sub-reaches must be active for this fixture to be conclusive" + ); + + let (length_m, slope) = slice_reach_geometry(&store, &compressed); + println!("length_m = {length_m:?}\nslope = {slope:?}"); + + let parent = chain3(); + for (p, &m) in pieces.iter().enumerate() { + let lo = store.parent_offset[p] as usize; + let hi = store.parent_offset[p + 1] as usize; + let expect_len = parent.length_m[p] / m as f32; + for row in lo..hi { + assert!( + (length_m[row] - expect_len).abs() < 1e-3, + "row {row} (parent {p}, {m} pieces): length {} != L/m {expect_len}", + length_m[row], + ); + assert!( + (slope[row] - parent.slope[p]).abs() < 1e-9, + "row {row} (parent {p}): slope {} != parent slope {}", + slope[row], + parent.slope[p], + ); + } + } + + // Total length is conserved per parent: splitting must not create channel. + for (p, _) in pieces.iter().enumerate() { + let lo = store.parent_offset[p] as usize; + let hi = store.parent_offset[p + 1] as usize; + let total: f32 = length_m[lo..hi].iter().sum(); + assert!( + (total - parent.length_m[p]).abs() < 1e-2, + "parent {p}: pieces sum to {total}, parent is {}", + parent.length_m[p], + ); + } +} + +// --------------------------------------------------------------------------- +// Parent -> sub-reach KAN parameter gather +// --------------------------------------------------------------------------- + +/// Gather `parents` onto the rows described by `parent_offset`, forward only. +fn gather_for_test(parents: &[f32], parent_offset: &[i32]) -> Vec { + let device = ::Device::default(); + let n_rows = *parent_offset.last().expect("non-empty offset") as usize; + let mut params: HashMap> = HashMap::new(); + params.insert("n".to_string(), Tensor::from_floats(parents, &device)); + let out = gather_params_to_subreaches( + params, + Some(&parent_offset.to_vec()), + n_rows, + &device, + ); + out["n"].clone().into_data().into_vec().expect("to host") +} + +/// `d(Σ gathered)/d(parent p)`. Proves `select`'s backward is a scatter-add. +fn gather_grad_for_test(parents: &[f32], parent_offset: &[i32]) -> Vec { + let device = ::Device::default(); + let n_rows = *parent_offset.last().expect("non-empty offset") as usize; + let leaf: Tensor = Tensor::from_floats(parents, &device).require_grad(); + let mut params: HashMap> = HashMap::new(); + params.insert("n".to_string(), leaf.clone()); + let out = gather_params_to_subreaches( + params, + Some(&parent_offset.to_vec()), + n_rows, + &device, + ); + let grads = out["n"].clone().sum().backward(); + leaf.grad(&grads) + .expect("parent leaf must receive a gradient") + .into_data() + .into_vec() + .expect("to host") +} + +/// Sub-reaches inherit their parent's hydraulics verbatim — MERIT carries no +/// within-reach variation, so there is nothing better to give them. +#[test] +fn every_piece_inherits_its_parents_parameters() { + let gathered = gather_for_test(&[0.02, 0.05, 0.10], &[0, 3, 5, 6]); + println!("gathered = {gathered:?}"); + assert_eq!(gathered, vec![0.02, 0.02, 0.02, 0.05, 0.05, 0.10]); +} + +/// The gradient half of the shared-parameter contract: a parent that feeds `m` +/// pieces must receive the SUM of their gradients, not one piece's. Anything +/// else (a broadcast, or a gather that forgets the tape) makes long reaches +/// learn at a different rate than the objective actually implies. +#[test] +fn gradient_sums_back_to_the_parent() { + let g = gather_grad_for_test(&[0.02, 0.05, 0.10], &[0, 3, 5, 6]); + println!("d(sum)/d(parent) = {g:?}"); + assert_eq!( + g, + vec![3.0, 2.0, 1.0], + "scatter-add must sum piece gradients (expected the piece counts)" + ); +} + +/// The subdivision-off path must not perturb a value. +/// +/// Honest about its reach: this is a VALUE-level check. It would still pass if +/// the `n_parent == n_rows` short-circuit were deleted and the gather ran with +/// an identity index (that property is structural — `gather_params_to_subreaches` +/// returns the map itself, recording nothing on the tape). What it does catch is +/// a mis-built row→parent index, which in the identity case would reorder or +/// truncate the parameters. +#[test] +fn identity_parent_offset_leaves_values_untouched() { + let device = ::Device::default(); + let vals = [0.02_f32, 0.05, 0.10]; + let mut params: HashMap> = HashMap::new(); + params.insert("n".to_string(), Tensor::from_floats(vals.as_slice(), &device)); + + for offset in [None, Some(vec![0i32, 1, 2, 3])] { + let out = gather_params_to_subreaches(params.clone(), offset.as_ref(), 3, &device); + let got: Vec = out["n"].clone().into_data().into_vec().expect("to host"); + assert_eq!( + got.iter().map(|v| v.to_bits()).collect::>(), + vals.iter().map(|v| v.to_bits()).collect::>(), + "identity map must not perturb a single bit (offset {offset:?})" + ); + } +} + +// ── cold start under subdivision ───────────────────────────────────────────── +// +// `setup_inputs` cold-starts with `(I − N)·Q_0 = q'_0`. The `q'/m` split lives +// in `forward`, so without `divide_hotstart_by_pieces` the cold start feeds the +// UNDIVIDED `q'_0` into a chain of `m` pieces and parent `p`'s outlet begins at +// `m_p ×` its true steady state. +// +// Measured on the real network (2026-08-05, `probe_courant --max-pieces 8`, +// 1,841 CONUS gauges / 184,676 rows): the undivided start put 2.94× the correct +// total discharge into the network and needed **221 hourly steps to decay +// below a 10 % difference** (282 below 5 %, still >1 % at 500), against a +// configured `warmup` of 5 days = 120 steps — at which point it was still +// 41.7 % off. Hence the division is on by default. + +/// The cold start must be mass-consistent: subdividing a reach must not change +/// the discharge the solver starts that reach's outlet at. +#[test] +fn divided_hotstart_gives_a_subdivided_reach_the_same_initial_condition() { + let un_split = outlet_series(1, true, None)[0]; + let split = outlet_series(4, true, None)[0]; + println!("Q_0 outlet: 1 piece = {un_split:.6}, 4 pieces = {split:.6}"); + assert!( + (split - un_split).abs() / un_split < 1e-3, + "cold start not mass-consistent: 1 piece = {un_split}, 4 pieces = {split} \ + (ratio {:.4}; a ratio near 4 means the hot-start q'_0 was not divided)", + split / un_split + ); +} + +/// The discriminating control: with the division off, the cold start really is +/// inflated `m×`. This is the state the real-network wash-out was measured on. +#[test] +fn undivided_hotstart_inflates_the_initial_condition_by_the_piece_count() { + let un_split = outlet_series(1, true, Some(false))[0]; + let split = outlet_series(4, true, Some(false))[0]; + println!("Q_0 outlet (undivided hot start): 1 piece = {un_split:.6}, 4 pieces = {split:.6}"); + assert!( + (split / un_split - 4.0).abs() < 1e-2, + "expected a 4x inflated cold start, got ratio {:.4}", + split / un_split + ); +} + +/// Steady state is reached either way — the cold start only sets how long the +/// spin-up takes — so the fix must not move the converged answer. +#[test] +fn hotstart_division_does_not_move_the_steady_state() { + let divided = outlet_series(4, true, Some(true))[T - 1]; + let undivided = outlet_series(4, true, Some(false))[T - 1]; + println!("steady state: divided = {divided:.6}, undivided = {undivided:.6}"); + assert!( + (divided - undivided).abs() / divided < 1e-3, + "the cold start changed the steady state: {divided} vs {undivided}" + ); +} + +/// An un-subdivided network has no divisor, so the default must be an EXACT +/// no-op there — this is what keeps `params.subdivision.enabled: false` +/// byte-identical (and `compare_ddr_sandbox` an ABSOLUTE MATCH). +#[test] +fn hotstart_division_is_a_no_op_without_subdivision() { + assert_eq!( + outlet_series(1, true, Some(true)), + outlet_series(1, true, Some(false)), + "the hot-start divisor must not touch an un-subdivided network" + ); +} + +// ── the non-negativity window ──────────────────────────────────────────────── + +/// **Why "Cr ≈ 1 ⇒ non-negative coefficients" does not survive contact with +/// MERIT.** `c1 ≥ 0 ⟺ Cr ≥ 2X` and `c3 ≥ 0 ⟺ Cr ≤ 2(1−X)`, so BOTH hold only +/// inside a window of width `2(1−2X)`. That width collapses as `X → 0.5`, and +/// on the real network the Cunge `X` sits at a median of 0.492–0.497 — a window +/// **1.3–3.1 % wide** (measured 2026-08-05, `probe_courant`, 1,841 gauges: +/// `2(1−2X)` p50 = 0.0134 un-split, 0.0310 at `max_pieces: 8`). +/// +/// A build-time piece count sets `Δx` from a *reference* flow, while `Cr` +/// tracks the *routed* celerity, which varies several-fold within a single +/// storm. Landing inside a 1–3 % window is therefore not achievable by +/// subdivision, whatever the cap. Measured `frac c1 ≥ 0 AND c3 ≥ 0` on CONUS: +/// 3.1 % un-split → 0.9 % at cap 8. +#[test] +fn both_coefficients_are_non_negative_only_inside_a_window_that_collapses_at_x_half() { + // Muskingum coefficients, mirroring `mmc_op.rs:1077-1080` with Cr = dt/K. + fn coeffs(cr: f32, x: f32) -> (f32, f32) { + let k = 1.0 / cr; // dt = 1 in these units + let denom = 2.0 * k * (1.0 - x) + 1.0; + ((1.0 - 2.0 * k * x) / denom, (2.0 * k * (1.0 - x) - 1.0) / denom) + } + + for &x in &[0.30_f32, 0.45, 0.4966] { + let lo = 2.0 * x; + let hi = 2.0 * (1.0 - x); + let width = hi - lo; + assert!( + (width - 2.0 * (1.0 - 2.0 * x)).abs() < 1e-6, + "window width must be 2(1-2X)" + ); + // Inside the window both are non-negative... + let (c1, c3) = coeffs(0.5 * (lo + hi), x); + assert!(c1 >= 0.0 && c3 >= 0.0, "X={x}: mid-window gave c1={c1}, c3={c3}"); + // ...and stepping outside it, either side, breaks one of them. + let (c1_lo, _) = coeffs(lo * 0.98, x); + assert!(c1_lo < 0.0, "X={x}: Cr below 2X must give c1 < 0, got {c1_lo}"); + let (_, c3_hi) = coeffs(hi * 1.02, x); + assert!(c3_hi < 0.0, "X={x}: Cr above 2(1-X) must give c3 < 0, got {c3_hi}"); + } + + // The measured CONUS median X leaves a window barely 1 % wide. + let x = 0.4966_f32; + assert!( + 2.0 * (1.0 - 2.0 * x) < 0.02, + "the CONUS-median non-negativity window must be under 2% wide" + ); +} diff --git a/tests/training_step_layer_b.rs b/tests/training_step_layer_b.rs index df9677a..0eae33a 100644 --- a/tests/training_step_layer_b.rs +++ b/tests/training_step_layer_b.rs @@ -132,6 +132,7 @@ fn adjacency_from_fixture( values, length_m, slope, + parent_offset: None, } } @@ -166,7 +167,9 @@ fn layer_b_step1_subgraph_adjacency_matches_ddr() { let staid = Staid::new(STAID_STR); let g = gages.get(&staid).expect("gauge not in gages store"); let unioned = ddrs::data::collate::union_subgraphs(&[staid.clone()], &gages); - let compressed = ddrs::data::collate::compress(&unioned, &conus.order) + // `ddr_match: true` — this test compares against a DDR-generated fixture, + // so it must use DDR's `outflow_idx` convention. + let compressed = ddrs::data::collate::compress(&unioned, &conus.order, true, None) .expect("compress failed"); // Convert divide_comids to i64 for comparison. @@ -525,7 +528,10 @@ fn layer_b_step3_mc_forward_matches_ddr() { /// Note on tau-slicing divergence (spec C7): /// DDR uses `[13:-11+tau]` → for tau=3: [13:2128], 2115 hours → truncated /// to 88 days (3 hours dropped from the end by downsample). -/// DDRS uses `[13+tau:-11+tau]` → [16:2128], 2112 hours → exactly 88 days. +/// DDRS (LEGACY convention, inlined below for fixture parity) used +/// `[13+tau:-11+tau]` → [16:2128], 2112 hours → exactly 88 days. Since +/// 2026-08-08 `tau_trim_and_downsample` uses `[tau : -(24-tau)]` with +/// day i ↔ obs day i; the identical window is reproduced at tau=16. /// /// Both produce 88 daily samples, but DDRS's day 1 = DDR's hours 16-39 /// while DDR's day 1 = hours 13-36. The 3-hour offset gives a small but diff --git a/tests/training_step_layer_c.rs b/tests/training_step_layer_c.rs index 6d00a1a..63702d7 100644 --- a/tests/training_step_layer_c.rs +++ b/tests/training_step_layer_c.rs @@ -193,6 +193,7 @@ fn adjacency_from_fixture(conus: &ConusAdjacencyStore) -> SparseAdjacency { values, length_m, slope, + parent_offset: None, } } @@ -337,7 +338,10 @@ fn full_forward( /// Note: avoids calling `tau_trim_and_downsample` directly because it calls /// `squeeze::<2>()` which collapses BOTH size-1 dims when n_gauges=1, /// producing a 1D tensor instead of 2D. Inlined here to preserve the 2D -/// shape. See Layer B sub-test 4 for the same workaround. +/// shape. See Layer B sub-test 4 for the same workaround. This inlines the +/// LEGACY slice `[13+tau : -11+tau]` (DDR-fixture convention; legacy tau=3 +/// cuts the same hourly window as the 2026-08-08 convention's tau=16) — +/// do not "modernize" it, the DDR fixtures pin it. fn ddrs_pred_post_warmup(hourly_q: Tensor, tau: u32, warmup: usize) -> Tensor { let dims = hourly_q.dims(); let (g, t_hours) = (dims[0], dims[1]); diff --git a/tests/training_step_layer_d.rs b/tests/training_step_layer_d.rs index 65ba973..74e5eb0 100644 --- a/tests/training_step_layer_d.rs +++ b/tests/training_step_layer_d.rs @@ -190,6 +190,7 @@ fn adjacency_from_fixture(conus: &ConusAdjacencyStore) -> SparseAdjacency { values, length_m, slope, + parent_offset: None, } } @@ -308,8 +309,10 @@ fn full_forward( (head, gauge_q) } -/// Apply DDRS's tau-trim + daily downsample + warmup slice to hourly gauge Q. -/// Mirrors the same helper from `training_step_layer_c.rs`. +/// Apply the LEGACY tau-trim + daily downsample + warmup slice to hourly +/// gauge Q. Mirrors the same helper from `training_step_layer_c.rs`; kept on +/// the legacy `[13+tau : -11+tau]` slice because the DDR fixtures pin it +/// (see that file's note on the 2026-08-08 convention change). fn ddrs_pred_post_warmup(hourly_q: Tensor, tau: u32, warmup: usize) -> Tensor { let dims = hourly_q.dims(); let (g, t_hours) = (dims[0], dims[1]); diff --git a/tests/training_verification.rs b/tests/training_verification.rs index efea6d3..511eff3 100644 --- a/tests/training_verification.rs +++ b/tests/training_verification.rs @@ -153,7 +153,13 @@ fn v1_loss_matches_ddr_for_frozen_constant_params_small_batch() { forward_with_frozen_params::>(&cfg, &tensors, &frozen, &device, false); // 10. Tau-trim + daily downsample → (num_gauges, T_days). - let pred_daily = tau_trim_and_downsample(pred_hourly, cfg.params.tau); + // The DDR fixture was generated under the LEGACY tau convention at + // tau=3 (window [16 : T-8], pooled day i ↔ obs day i+1). Under the + // 2026-08-08 convention the identical window is tau=16, and this + // test's obs slice [1..-1] keeps the legacy day pairing — so the + // fixture comparison is unchanged bit-for-bit. Do NOT switch this + // to cfg.params.tau: the config default is now on the new scale. + let pred_daily = tau_trim_and_downsample(pred_hourly, 16); let [_g, t_days] = pred_daily.dims(); // 11. Convert BURN tensor → ndarray::Array2 (row-major, shape (G, T_days)). @@ -288,7 +294,13 @@ fn v2_loss_matches_ddr_for_frozen_constant_params_all_gauges() { forward_with_frozen_params::>(&cfg, &tensors, &frozen, &device, false); // 10. Tau-trim + daily downsample → (num_gauges, T_days). - let pred_daily = tau_trim_and_downsample(pred_hourly, cfg.params.tau); + // The DDR fixture was generated under the LEGACY tau convention at + // tau=3 (window [16 : T-8], pooled day i ↔ obs day i+1). Under the + // 2026-08-08 convention the identical window is tau=16, and this + // test's obs slice [1..-1] keeps the legacy day pairing — so the + // fixture comparison is unchanged bit-for-bit. Do NOT switch this + // to cfg.params.tau: the config default is now on the new scale. + let pred_daily = tau_trim_and_downsample(pred_hourly, 16); let [_g, t_days] = pred_daily.dims(); // 11. Convert BURN tensor → ndarray::Array2 (row-major, shape (G, T_days)).