Skip to content

Rework disagg head to KAN + basin-normalized precip; H5/H6 loss-landscape equifinality infra#26

Merged
taddyb merged 9 commits into
masterfrom
loss-landscape
Jul 14, 2026
Merged

Rework disagg head to KAN + basin-normalized precip; H5/H6 loss-landscape equifinality infra#26
taddyb merged 9 commits into
masterfrom
loss-landscape

Conversation

@taddyb

@taddyb taddyb commented Jul 11, 2026

Copy link
Copy Markdown
Owner

Summary

  • Reworks the daily→hourly disaggregation head from an MLP to a KAN (Linear→KanLayer×N→Linear), and fixes precip normalization to basin-normalize by climatological meanP before log1p (the prior per-window z-score silently canceled the basin-normalization signal — verified empirically). Drops the now-unused use_attributes/use_precip/use_temp toggles since the head always consumes (daily Q', that day's precip).
  • Adds --mode eval-loss and --mode landscape to probe_zeta_gradient to support the H5 (parameter-transfer) and H6 (loss-landscape) selective-equifinality hypotheses, plus a param_dump store module for reading back dump_parameters NetCDF output as swap donors.
  • Runs and documents the H5/H6 campaign: docs/2026-07-09-h5-h6-equifinality-findings.md (v1) and docs/2026-07-09-h5-h6-equifinality-v2-findings.md (audit-corrected re-analysis of the same raw CSVs — verdicts unchanged at INCONCLUSIVE for both, but the evidentiary basis and surface-geometry reading change substantially).
  • Adds sensitivity/mass-balance verification examples for the new disagg head (kan_disagg_mass_balance_real, kan_disagg_trained_sensitivity, kan_sensitivity_sweep, disagg_boundary_verification).

Test plan

  • cargo test --release --lib --bins --tests — all suites pass (0 failures; two pre-existing doctest failures in src/cuda_graph/geometry_kernel.rs/src/sparse/cusparse.rs are unrelated — those files are untouched by this branch)
  • cargo run --release --example compare_ddr_sandbox — ABSOLUTE MATCH (max abs diff 1.5e-5 m³/s, well under the 1e-3 bar)

🤖 Generated with Claude Code

taddyb added 7 commits July 8, 2026 13:16
…ffolding

Adds param_dump::load_comid_field to read back dump_parameters NetCDF output
as a donor source for parameter swaps, re-exports the pieces forward.rs's
override plumbing needs, and pins forward_eval's 'own' composition against
the training-loop L1 loss (tests/eval_loss_own_parity.rs) ahead of the H5/H6
probe work.
…H5/H6 infra)

Stage 6 (--mode eval-loss): evaluates the training L1 loss over a fixed
deterministic rho-window/gauge plan under own/n-swap/geo-swap/full-swap
parameter compositions, injecting donor n/q_spatial/p_spatial from a
--donor-params-nc dump in place of the checkpoint's own KAN-head output.

Stage 7 (--mode landscape): (alpha, beta)-scaled loss-surface grid scan plus
a log-space linear parameter barrier between two arms' dumps, for the H6
loss-landscape hypothesis.

Registers the H5/H6 selective-equifinality spec
(docs/superpowers/specs/2026-07-08-landscape-hypotheses-h5-h6-draft.md) and
the corresponding ddrs-eval-plots skill references.
…ed v2 re-analysis

docs/2026-07-09-h5-h6-equifinality-findings.md: v1 findings from the
registered H5 (parameter-transfer) and H6 (loss-landscape) runs — both
INCONCLUSIVE per the pre-registered bars.

docs/2026-07-09-h5-h6-equifinality-v2-findings.md: audit-corrected
re-analysis of the same raw CSVs via scripts/h5_h6_audit_analysis.py — same
verdicts, but the noise-floor comparison, control-pair result, and H6
surface-geometry reading all change substantially (see doc §1 "what this
pass corrects").
… attribute space

H6-style treatment with the learned parameter (n, q_spatial, p_spatial) as
the height field instead of loss: median value per (log10 drainage area x
log10 mean slope) bin, one panel per arm plus a cross-arm range panel.
Replaces the daily->hourly disaggregation head's MLP with a KAN
(Linear -> KanLayer(H,H) x num_hidden_layers -> Linear), matching the
routing head's architecture. The head now always consumes
(daily Q', that day's precip) — drops the disagg_use_attributes/
disagg_use_precip/disagg_use_temp toggles since attribute- and
temp-conditioning are gone and precip is no longer optional.

normalize_precip now basin-normalizes by each reach's climatological mean
precip (meanP) before log1p, instead of a per-window z-score on top of a
plain log1p — the z-score was silently canceling the basin-normalization
signal (scale-invariant), so the network never saw the basin-relative
magnitude. Requires 'meanP' in kan_head.input_var_names whenever
kan_head.disaggregation is set.

Updates all tracked experiment configs to drop the removed disagg toggles
and adds config/experiments/kan_disagg_conus_no_leakance.yaml plus two new
verification examples (disagg_boundary_verification,
kan_sensitivity_sweep).
… disagg head

kan_disagg_mass_balance_real: checks the trained head's hourly output sums
back to the daily Q' it disaggregated, on real checkpoint + real data.

kan_disagg_trained_sensitivity: probes whether the trained head's output
shape actually responds to precip timing (synthetic flat-background +
single-hour-spike sweep).
… rate

meanP is mm/year but precip is mm/hr; dividing directly crushed real storms
to ~1e-3, ~1000x smaller than the disagg head's other input feature
(log(daily_q), typically O(1)) — numerically inert at the input layer and
undetectable to the trained checkpoint. Convert meanP to an hourly rate
(/ hours-per-year) before dividing so both operands share a timescale.

examples/kan_disagg_real_storm_shift.rs: circularly shifts one real day's
actual AORC precip vector by k hours (holding that day's real daily_q
fixed) and re-runs the trained disagg head per shift, to check whether the
output peak tracks precip timing under realistic (bursty, in-distribution)
input rather than the synthetic flat-background+spike sweep in
kan_disagg_trained_sensitivity.rs.
taddyb added 2 commits July 12, 2026 20:57
…aign

Adds a CamelsHourlyStore reader for real USGS/CAMELS-hourly discharge +
precip, and standalone binaries (pretrain_disagg, _blend, _window72,
_capacity) to train the production DisaggHead directly against real hourly
ground truth instead of only end-to-end inside the full CONUS routing model,
where the loss signal for sub-daily shape is indirect and noisy. Includes
boundary-blend and 72h-lookback storm-compare diagnostics, a gauge
reconciliation check, and a precip-normalization discriminator used to
validate the pretraining feature contract before wiring it into production.
…tart pretrained head

Adds `chunk_days` to DisaggHead/KanHead/config: relaxes daily-exact mass
conservation to a multi-day joint chunk so storms spanning a day boundary
disaggregate as one continuous event instead of clipping at hour 0/23
(chunk_days=1 is byte-identical to prior behavior). Ablation showed the
deeper, capacity-boosted head pretrained standalone on real hourly data
clearly wins on peak-timing and peak-magnitude error, so `kan_head.disaggregation`
gains `pretrained_checkpoint` + `freeze` to warm-start (and optionally freeze,
via Module::no_grad()) the disagg submodule from that standalone checkpoint
instead of training it end-to-end from a random init — isolating whether a
previous production regression came from a bad head or a bad joint gradient.

Also drops the meanP basin-normalization divisor from production precip
normalization (dataset.rs) in favor of the plain per-window z-score, matching
the pretraining feature contract exactly so a frozen head sees the same
input distribution it was trained on.

New experiment configs: kan_disagg_conus_chunk3 (72h joint window, unfrozen),
kan_disagg_conus_frozen_chunk1 / _finetune_chunk1 (capacity-boosted head,
frozen vs fine-tuned from the same pretrained init).
@taddyb
taddyb merged commit ec84987 into master Jul 14, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant