Rework disagg head to KAN + basin-normalized precip; H5/H6 loss-landscape equifinality infra#26
Merged
Conversation
…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.
…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).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Linear→KanLayer×N→Linear), and fixes precip normalization to basin-normalize by climatologicalmeanPbeforelog1p(the prior per-window z-score silently canceled the basin-normalization signal — verified empirically). Drops the now-unuseduse_attributes/use_precip/use_temptoggles since the head always consumes(daily Q', that day's precip).--mode eval-lossand--mode landscapetoprobe_zeta_gradientto support the H5 (parameter-transfer) and H6 (loss-landscape) selective-equifinality hypotheses, plus aparam_dumpstore module for reading backdump_parametersNetCDF output as swap donors.docs/2026-07-09-h5-h6-equifinality-findings.md(v1) anddocs/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).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 insrc/cuda_graph/geometry_kernel.rs/src/sparse/cusparse.rsare 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