diff --git a/config/experiments/sr_n0_gamma_dhbv2_dist_s43.yaml b/config/experiments/sr_n0_gamma_dhbv2_dist_s43.yaml new file mode 100644 index 00000000..fd253faa --- /dev/null +++ b/config/experiments/sr_n0_gamma_dhbv2_dist_s43.yaml @@ -0,0 +1,336 @@ +# sr_n0_gamma_dhbv2_dist_s43 +# +# Seed-43 replicate of the dHBV2 distributed fixed-channel model (E-07, seed +# noise floor). Seed-42 run: 2026-09-13T17-22-30Z-train-and-test, trained from +# sr_n0_gamma_dhbv2_dist.yaml at commit 70f6457. This replicate runs a release +# build of 70f6457, so the pair shares a binary. Only `seed` and `np_seed` +# differ (42 -> 43); everything below this header is the seed-42 config +# verbatim, including its own header text. +# +# .claude/worktrees/e07-build-70f6457/target/release/ddrs \ +# --config config/experiments/sr_n0_gamma_dhbv2_dist_s43.yaml \ +# --workspace /home/tbindas/projects/ddrs/.ddrs \ +# run --workflow train-and-test --plot --backend cpu + +# sr_n0_gamma_dhbv2_dist +# +# Inflow-source arm for the Beven-challenges study (AGU 2026): the n_0 + gamma +# head with the channel fixed (p = 21, q = 0.65), trained on the +# dHBV2 distributed, AORC forcing, daily +# unit-catchment product, everything else identical to sr_n0_gamma.yaml +# (2,365 gauges of gages_3000.csv, seed 42, 500 updates). Five such arms, all +# AORC-forced, none on the dHBV2 UH retrospective. Question: how much of the +# learned roughness is the routing absorbing inflow bias, and do the per-gauge +# loss-landscape optima move with the inflow source? +# +# sr_n0_gamma +# +# The KAN emits ONLY n_0 and gamma. p_spatial is fixed at 21 and q_spatial at +# 0.65 (params.defaults), so the channel shape is prescribed and the head +# carries the two roughness parameters of n(d) = n_0 * (d/d_ref)^(-gamma). +# +# Why (2026-09-12, findings §36.5-36.7): with p and q free, a learned gamma +# rode the width channel (rho(p, gamma) 0.89, rho(q, gamma) 0.76) and the +# head traded q against gamma, so the sweep could not tell degeneracy from +# unidentifiability. Fixing p and q removes that direction by construction, +# and a two-output head is what the loss-landscape and adjoint instruments +# can probe over (n_0, gamma) directly. +# +# Matched control: sr_n0_only.yaml (n_0 alone, same fixed p and q, gamma = 0). +# +# ddrs --config config/experiments/sr_n0_gamma.yaml \ +# --workspace /home/tbindas/projects/ddrs/.ddrs \ +# run --workflow train-and-test --plot --backend cpu + +# head_shared_linear +# +# One arm of the head-topology experiment. Derived from the 500-update +# nse-batch baseline 2026-09-10T21-21-48Z-conus-train-and-test +# (NSE 0.7376 / KGE 0.7600) by changing the kan_head section and nothing +# else. Screen that motivated these arms: +# .ddrs/experiments/head-arch/2026-09-11-screen +# experiments/head_arch/analyze.py +# +# MATCHED CONTROL. Current architecture: one shared trunk, Linear(H,P) +# read-out. Measured trunk effective rank 4.36 of 21; init affine R^2 between +# logit(n) and logit(q) at the chance level for that rank. +# +# Run (note --workspace: without it .ddrs lands beside this config): +# ddrs --config config/experiments/head_shared_linear.yaml \ +# --workspace /home/tbindas/projects/ddrs/.ddrs \ +# run --workflow train-and-test + +# 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: 43 +np_seed: 43 + +# 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: + # CONUS MERIT inputs as available on the workstation. + 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 + # Icechunk dHBV2 UH retrospective (divide-major, daily m³/s) + streamflow: /mnt/ssd1/data/icechunk/daily_dhbv2_distributed_aorc2f_merit_unit_catchments_rechunked.ic + # Icechunk USGS daily observations + observations: /mnt/ssd1/data/icechunk/usgs_daily_observations + gages: /home/tbindas/projects/ddr/references/gage_info/gages_3000.csv + +experiment: + batch_size: 64 + start_time: 1981/10/01 + end_time: 1995/09/30 + epochs: 50 # 2026-09-10: 30 -> 50, with grad_accum_steps 4 this is 500 optimizer updates (was 60) + 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 + # 2026-09-10: 20 -> 4. The 2026-09-08 run made 60 optimizer updates in total + # (37 micro-batches/epoch at 2,365 kept gauges, 2 steps/epoch, 30 epochs) and + # findings §21 shows it never reached a stationary point: on BOTH the training + # and testing windows the loss still falls if roughness rises, at ~79% of + # gauges (trimmed gradient alignment 0.96/0.94). At accum 4 that is 10 + # steps/epoch, so 500 updates over 50 epochs, effective batch 4*64 = 256. + # Wall clock is forward-pass bound at 37 forwards/epoch either way: the + # 30-epoch run trained in 64 min, so 50 epochs is ~107 min. The "27.5 s per + # micro-batch" figure above is stale; measured here it is ~3.5 s. + grad_accum_steps: 4 + # 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. + # 2026-09-10: flattened. n froze at epoch 10 in the previous run, which is + # exactly where the old schedule cut 0.005 -> 0.001, so "the gradient + # vanished" and "the steps became too small" were not separable from the + # trajectory (§10 correction). §21 shows the gradient had NOT vanished. + # Keys are the FIRST epoch at which the lr applies. At 10 steps/epoch: + # epochs 1-20 @ 0.005 (200 updates) + # epochs 21-40 @ 0.002 (200 updates) + # epochs 41-50 @ 0.001 (100 updates) + # Displacement budget ~1.3 against the previous ~0.13. + # WATCH `n_at_floor`: lr 0.01 pinned 47.3% of CONUS at the 0.015 floor and + # lr 2e-3 collapsed 0% -> 25-36% by epoch 5 in earlier small-batch runs. + # Holding 0.005 for 200 updates at an effective batch of 256 is the real risk + # in this plan. If n_at_floor passes a few percent, drop the first rate to + # 0.002 and keep the step count. + learning_rate: + 1: 0.005 + 21: 0.002 + 41: 0.001 + 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 + # All three learned together. The baseline fixed p_spatial at 21 + # (docs/2026-09-08-fixed-p-assessment.md); §30 argued that learning p and q + # as a pair relocates the free parameter onto something falsifiable against + # remotely sensed river width, and §31 showed q alone was only n amplified + # fourfold. Because this differs from the baseline in TWO ways (learnable set + # and head topology), the shared-linear arm is the matched control: read + # these arms against each other, never against 0.7376/0.7600. + learnable_parameters: + - n + - gamma + # 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: + # Stage-roughness exponent. Non-negative because channels get SMOOTHER as + # they fill; capped at 0.5 because the depth exponent is 3/(5+3q+3*gamma) + # and roughness grows without bound as depth goes to zero. 0.183 reproduces + # observed at-a-station hydraulic geometry when paired with q ~ 0.65, and + # the sigmoid centre of this box is 0.25. + gamma: [0.0, 0.5] + 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 + # Fixed width exponent: the at-a-station Leopold & Maddock value b/f = 0.65 + # (design doc §2.2), which with gamma ~ 0.18 reproduces the observed + # at-a-station exponents. Must lie inside parameter_ranges.q_spatial. + q_spatial: 0.65 + 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/sr_n0_gamma_dhbv2_lumped_s43.yaml b/config/experiments/sr_n0_gamma_dhbv2_lumped_s43.yaml new file mode 100644 index 00000000..f47cbd39 --- /dev/null +++ b/config/experiments/sr_n0_gamma_dhbv2_lumped_s43.yaml @@ -0,0 +1,337 @@ +# sr_n0_gamma_dhbv2_lumped_s43 +# +# Seed-43 replicate of the dHBV2 lumped fixed-channel model (E-07, seed noise +# floor). Seed-42 run: 2026-09-13T13-56-50Z-train-and-test, trained from +# sr_n0_gamma_dhbv2_lumped.yaml at commit 4ebce97, whose src/ and Cargo.lock +# are identical to 70f6457. This replicate runs a release build of 70f6457, so +# the pair shares a binary. Only `seed` and `np_seed` differ (42 -> 43); +# everything below this header is the seed-42 config verbatim, including its +# own header text. +# +# .claude/worktrees/e07-build-70f6457/target/release/ddrs \ +# --config config/experiments/sr_n0_gamma_dhbv2_lumped_s43.yaml \ +# --workspace /home/tbindas/projects/ddrs/.ddrs \ +# run --workflow train-and-test --plot --backend cpu + +# sr_n0_gamma_dhbv2_lumped +# +# Inflow-source arm for the Beven-challenges study (AGU 2026): the n_0 + gamma +# head with the channel fixed (p = 21, q = 0.65), trained on the +# dHBV2 lumped, AORC forcing, daily +# unit-catchment product, everything else identical to sr_n0_gamma.yaml +# (2,365 gauges of gages_3000.csv, seed 42, 500 updates). Five such arms, all +# AORC-forced, none on the dHBV2 UH retrospective. Question: how much of the +# learned roughness is the routing absorbing inflow bias, and do the per-gauge +# loss-landscape optima move with the inflow source? +# +# sr_n0_gamma +# +# The KAN emits ONLY n_0 and gamma. p_spatial is fixed at 21 and q_spatial at +# 0.65 (params.defaults), so the channel shape is prescribed and the head +# carries the two roughness parameters of n(d) = n_0 * (d/d_ref)^(-gamma). +# +# Why (2026-09-12, findings §36.5-36.7): with p and q free, a learned gamma +# rode the width channel (rho(p, gamma) 0.89, rho(q, gamma) 0.76) and the +# head traded q against gamma, so the sweep could not tell degeneracy from +# unidentifiability. Fixing p and q removes that direction by construction, +# and a two-output head is what the loss-landscape and adjoint instruments +# can probe over (n_0, gamma) directly. +# +# Matched control: sr_n0_only.yaml (n_0 alone, same fixed p and q, gamma = 0). +# +# ddrs --config config/experiments/sr_n0_gamma.yaml \ +# --workspace /home/tbindas/projects/ddrs/.ddrs \ +# run --workflow train-and-test --plot --backend cpu + +# head_shared_linear +# +# One arm of the head-topology experiment. Derived from the 500-update +# nse-batch baseline 2026-09-10T21-21-48Z-conus-train-and-test +# (NSE 0.7376 / KGE 0.7600) by changing the kan_head section and nothing +# else. Screen that motivated these arms: +# .ddrs/experiments/head-arch/2026-09-11-screen +# experiments/head_arch/analyze.py +# +# MATCHED CONTROL. Current architecture: one shared trunk, Linear(H,P) +# read-out. Measured trunk effective rank 4.36 of 21; init affine R^2 between +# logit(n) and logit(q) at the chance level for that rank. +# +# Run (note --workspace: without it .ddrs lands beside this config): +# ddrs --config config/experiments/head_shared_linear.yaml \ +# --workspace /home/tbindas/projects/ddrs/.ddrs \ +# run --workflow train-and-test + +# 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: 43 +np_seed: 43 + +# 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: + # CONUS MERIT inputs as available on the workstation. + 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 + # Icechunk dHBV2 UH retrospective (divide-major, daily m³/s) + streamflow: /mnt/ssd1/data/icechunk/daily_dhbv2_lumped_aorc2f_merit_unit_catchments_rechunked.ic + # Icechunk USGS daily observations + observations: /mnt/ssd1/data/icechunk/usgs_daily_observations + gages: /home/tbindas/projects/ddr/references/gage_info/gages_3000.csv + +experiment: + batch_size: 64 + start_time: 1981/10/01 + end_time: 1995/09/30 + epochs: 50 # 2026-09-10: 30 -> 50, with grad_accum_steps 4 this is 500 optimizer updates (was 60) + 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 + # 2026-09-10: 20 -> 4. The 2026-09-08 run made 60 optimizer updates in total + # (37 micro-batches/epoch at 2,365 kept gauges, 2 steps/epoch, 30 epochs) and + # findings §21 shows it never reached a stationary point: on BOTH the training + # and testing windows the loss still falls if roughness rises, at ~79% of + # gauges (trimmed gradient alignment 0.96/0.94). At accum 4 that is 10 + # steps/epoch, so 500 updates over 50 epochs, effective batch 4*64 = 256. + # Wall clock is forward-pass bound at 37 forwards/epoch either way: the + # 30-epoch run trained in 64 min, so 50 epochs is ~107 min. The "27.5 s per + # micro-batch" figure above is stale; measured here it is ~3.5 s. + grad_accum_steps: 4 + # 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. + # 2026-09-10: flattened. n froze at epoch 10 in the previous run, which is + # exactly where the old schedule cut 0.005 -> 0.001, so "the gradient + # vanished" and "the steps became too small" were not separable from the + # trajectory (§10 correction). §21 shows the gradient had NOT vanished. + # Keys are the FIRST epoch at which the lr applies. At 10 steps/epoch: + # epochs 1-20 @ 0.005 (200 updates) + # epochs 21-40 @ 0.002 (200 updates) + # epochs 41-50 @ 0.001 (100 updates) + # Displacement budget ~1.3 against the previous ~0.13. + # WATCH `n_at_floor`: lr 0.01 pinned 47.3% of CONUS at the 0.015 floor and + # lr 2e-3 collapsed 0% -> 25-36% by epoch 5 in earlier small-batch runs. + # Holding 0.005 for 200 updates at an effective batch of 256 is the real risk + # in this plan. If n_at_floor passes a few percent, drop the first rate to + # 0.002 and keep the step count. + learning_rate: + 1: 0.005 + 21: 0.002 + 41: 0.001 + 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 + # All three learned together. The baseline fixed p_spatial at 21 + # (docs/2026-09-08-fixed-p-assessment.md); §30 argued that learning p and q + # as a pair relocates the free parameter onto something falsifiable against + # remotely sensed river width, and §31 showed q alone was only n amplified + # fourfold. Because this differs from the baseline in TWO ways (learnable set + # and head topology), the shared-linear arm is the matched control: read + # these arms against each other, never against 0.7376/0.7600. + learnable_parameters: + - n + - gamma + # 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: + # Stage-roughness exponent. Non-negative because channels get SMOOTHER as + # they fill; capped at 0.5 because the depth exponent is 3/(5+3q+3*gamma) + # and roughness grows without bound as depth goes to zero. 0.183 reproduces + # observed at-a-station hydraulic geometry when paired with q ~ 0.65, and + # the sigmoid centre of this box is 0.25. + gamma: [0.0, 0.5] + 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 + # Fixed width exponent: the at-a-station Leopold & Maddock value b/f = 0.65 + # (design doc §2.2), which with gamma ~ 0.18 reproduces the observed + # at-a-station exponents. Must lie inside parameter_ranges.q_spatial. + q_spatial: 0.65 + 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/experiments/adjoint-six-models/analyze.py b/experiments/adjoint-six-models/analyze.py new file mode 100644 index 00000000..64d15e6b --- /dev/null +++ b/experiments/adjoint-six-models/analyze.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python +"""E-09 analysis: celerity, transfer coefficient, inherited share, FD gate across the six arms. + +Usage (ddr venv, for xarray/pandas): + /home/tbindas/projects/ddr/.venv/bin/python experiments/adjoint-six-models/analyze.py \ + [ ...] # several dirs are merged (arms are disjoint) + +Reuses experiments/adjoint/plots.py (load, per_gauge_celerity, inherited_decomposition). +Writes results.md and the csv tables next to the first out-dir under analysis/. +""" +import json, sys +from pathlib import Path +import numpy as np +import pandas as pd + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "adjoint")) +import plots # noqa: E402 + +OLD41 = Path("/home/tbindas/projects/ddrs/.ddrs/experiments/adjoint-conus/2026-09-07T18-18-39Z/gauges.csv") +ORDER = ["uh-fixed-channel", "dhbv2-lumped", "dhbv2-distributed", "hydrodl-lstm", "nh-daily-lstm", "nh-hourly-mts-lstm"] + + +def main(): + outs = [Path(p) for p in sys.argv[1:]] + data, arms, val = {}, [], [] + for o in outs: + man, a, gauges, d = plots.load(o) + arms += [x for x in a if any(k[0] == x for k in d)] + data.update(d) + val += man["validation"] + arms = [a for a in ORDER if a in arms] + ana = outs[0] / "analysis" + ana.mkdir(exist_ok=True) + lines = [] + P = lines.append + + # FD gate + P("## FD gate (+5 percent inflow at three reaches, tolerance 5 percent)\n") + P("| arm | gauge | passed | max rel err (gauge / median / farthest reach) |") + P("|---|---|---|---|") + seen = {} + for v in val: + e = [c["rel_err"] for c in v["checks"]] + if v["arm"] in seen: + continue + seen[v["arm"]] = 1 + P(f"| {v['arm']} | {v['staid']} | {v['passed']} | {max(e)*100:.3f} % ({', '.join(f'{x*100:.3f}' for x in e)}) |") + P("") + + cel = plots.per_gauge_celerity(arms, gauges, data) + dec = plots.inherited_decomposition(arms, gauges, data) + cel.to_csv(ana / "celerity.csv", index=False) + dec.to_csv(ana / "inherited.csv", index=False) + old = set(pd.read_csv(OLD41, dtype=str).staid) + sets = {"all": set(gauges.staid), "old41_overlap": set(gauges.staid) & old} + for name, s in sets.items(): + P(f"## Population `{name}`: {len(s)} gauges\n") + c = cel[cel.staid.isin(s)] + P("Effective celerity per gauge (origin fit of kernel mean lag vs distance), median over gauges, m/s:\n") + P("| arm | n gauges (low) | low-flow median | high-flow median | transfer-proxy: median wet volume sens |") + P("|---|---|---|---|---|") + for a in arms: + lo = c[(c.arm == a) & (c.kind == "low")] + hi = c[(c.arm == a) & (c.kind == "high")] + P(f"| {a} | {len(lo)} | {lo.celerity_m_s.median():.3f} | {hi.celerity_m_s.median():.3f} | {lo.volume_sens_median.median():.3f} |") + P("") + for kind in ["low", "high"]: + w = c[c.kind == kind].pivot(index="staid", columns="arm", values="celerity_m_s")[arms].dropna() + if w.empty: + continue + ratio = w.max(axis=1) / w.min(axis=1) + slow = w.idxmin(axis=1).value_counts().reindex(arms).fillna(0).astype(int) + fast = w.idxmax(axis=1).value_counts().reindex(arms).fillna(0).astype(int) + P(f"{kind}-flow cross-arm max/min celerity ratio over {len(w)} gauges with all {len(arms)} arms: median {ratio.median():.2f}, IQR {ratio.quantile(.25):.2f} to {ratio.quantile(.75):.2f}") + P(f" slowest-arm counts: {dict(slow)}") + P(f" fastest-arm counts: {dict(fast)}\n") + # ratio among the five AORC-or-UH vs excluding hourly + sub = [a for a in arms if a != "nh-hourly-mts-lstm"] + if len(sub) < len(arms): + r2 = (w[sub].max(axis=1) / w[sub].min(axis=1)) + P(f" without the hourly arm ({len(sub)} arms): median ratio {r2.median():.2f}\n") + d = dec[dec.staid.isin(s)] + if not d.empty: + share = (d.assign(a=d.inherited.abs(), t=d.total.abs()).groupby(["arm", "staid"], as_index=False) + .agg(a=("a", "sum"), t=("t", "sum"), tr=("transfer_mean", "first"))) + share["share"] = share.a / share.t.replace(0, np.nan) + P("Downstream gauges: inherited share of bias and volume transfer coefficient (upstream gauge reach seen at downstream gauge):\n") + P("| arm | n | inherited share median (IQR) | transfer coefficient median (IQR) |") + P("|---|---|---|---|") + for a in arms: + x = share[share.arm == a] + if x.empty: + continue + P(f"| {a} | {len(x)} | {x.share.median():.2f} ({x.share.quantile(.25):.2f} to {x.share.quantile(.75):.2f}) | {x.tr.median():.3f} ({x.tr.quantile(.25):.3f} to {x.tr.quantile(.75):.3f}) |") + P("") + (ana / "results.md").write_text("\n".join(lines)) + print("\n".join(lines)) + + +if __name__ == "__main__": + main() diff --git a/experiments/adjoint-six-models/experiment.yaml b/experiments/adjoint-six-models/experiment.yaml new file mode 100644 index 00000000..36b8d682 --- /dev/null +++ b/experiments/adjoint-six-models/experiment.yaml @@ -0,0 +1,28 @@ +# E-09: adjoint influence maps on the six fixed-channel (n0 + gamma) models, over +# the GAGES-II nested-reference population. Same functionals and settings as +# experiments/adjoint-conus; only the arms change. Canonical run ids from +# research/findings/2026-09-30-paper-model-lineage.md (paper/e01-lineage@e8e7d3b). +# Arms are read-only; latest directory checkpoint (epoch_50_mb_9). +# /target/release/ddrs --workspace /home/tbindas/projects/ddrs/.ddrs experiment adjoint-six-models --backend cpu +study: adjoint +arms: + - { name: uh-fixed-channel, run: 2026-09-12T23-39-03Z-train-and-test } + - { name: dhbv2-lumped, run: 2026-09-13T13-56-50Z-train-and-test } + - { name: dhbv2-distributed, run: 2026-09-13T17-22-30Z-train-and-test } + - { name: hydrodl-lstm, run: 2026-09-13T17-21-58Z-train-and-test } + - { name: nh-daily-lstm, run: 2026-09-13T13-55-03Z-train-and-test } + - { name: nh-hourly-mts-lstm, run: 2026-09-13T14-14-46Z-train-and-test } + +adjoint: + gauges: + source: nested-reference + gages_ii_dbf: /mnt/ssd1/data/gage_shp_files/gagesII_9322_sept30_2011.dbf + window_days: 365 + anchors: { high: 2, low: 2 } + anchor_min_gap_days: 60 + lag_days: 30 + tail_days: 7 + water_year: 2000 + functionals: [kernel, volume, residual] + fd_perturb_fraction: 0.05 + fd_rel_tol: 0.05 diff --git a/experiments/blend_prediction/blend_inflow.py b/experiments/blend_prediction/blend_inflow.py new file mode 100644 index 00000000..c117a86e --- /dev/null +++ b/experiments/blend_prediction/blend_inflow.py @@ -0,0 +1,169 @@ +"""Per-gauge summed inflow of the three-product blend, computed offline (no training). + +Exact route. Reads the three icechunk Q' stores named in +config/experiments/sr_n0_gamma_blend.yaml (multi-product-forcing worktree) over the +test window, blends them per divide and day with the rule of +src/data/store/blend.rs (read only): + + q[t, j] = sum_m a[m, j] q_m[t, j], a[m, j] = w_m c[m, j] / sum_k w_k c[k, j] + +where c[m, j] = 1 when store m contains divide j, and a divide no store contains +reads the 0.001 fill. By linearity the gauge sum is accumulated store by store, +so no blended per-divide array is ever held in memory. + +For every gauge of the members' summed-Q' baseline (2,698 gauges, the gauge +adjacency `merit_gages_conus_adjacency.zarr`, upstream set = CONUS `order` at the +union of the subgraph's COO indices, as `GageSubgraph::upstream_comids`) it writes: + + S_blend the blend's summed inflow + S_ each member's raw summed inflow (0.001 fill where the store + lacks a divide), which must reproduce that member run's + baseline/predictions.f32 (checked and printed) + full_cover[g] every upstream divide is in all three stores + + /home/tbindas/projects/ddr/.venv/bin/python experiments/blend_prediction/blend_inflow.py --out +""" +import argparse +import json +import time +from pathlib import Path + +import icechunk +import numpy as np +import scipy.sparse as sp +import zarr + +RUNS = Path("/home/tbindas/projects/ddrs/.ddrs/runs") +CONUS = "/home/tbindas/projects/ddr/data/merit_conus_adjacency.zarr" +GAGES_ADJ = "/home/tbindas/projects/ddr/data/merit_gages_conus_adjacency.zarr" +MEMBERS = [ + # label, store, blend weight, member run whose baseline the raw sum must reproduce + ("uh", "/mnt/ssd1/data/icechunk/merit_dhbv2_UH_retrospective.ic", 0.50, + "2026-09-27T07-29-47Z-train-and-test"), + ("dist", "/mnt/ssd1/data/icechunk/daily_dhbv2_distributed_aorc2f_merit_unit_catchments_rechunked.ic", 0.25, + "2026-09-13T17-22-30Z-train-and-test"), + ("lstm", "/mnt/ssd1/data/icechunk/daily_lstm_merit_unit_catchments.ic", 0.25, + "2026-09-13T13-55-03Z-train-and-test"), +] +START, N_DAYS = np.datetime64("1995-10-01"), 5479 # WY1996-2010, both ends inclusive +FILL = 0.001 +BLOCK = 3080 * 8 + + +def open_store(path): + repo = icechunk.Repository.open(icechunk.local_filesystem_storage(path)) + return zarr.open_group(repo.readonly_session("main").store, mode="r") + + +def time_index(g): + t = np.asarray(g["time"][:]) + units = g["time"].attrs["units"] + origin = np.datetime64(units.split("since")[1].strip()[:10]) + assert np.all(np.diff(t) == 1), "time axis not contiguous" + i0 = int((START - origin).astype(int) - t[0]) + assert 0 <= i0 and i0 + N_DAYS <= t.size, (units, t[0], t.size) + return i0 + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--out", required=True) + out = Path(ap.parse_args().out) + out.mkdir(parents=True, exist_ok=True) + + conus = zarr.open_group(CONUS, mode="r") + order = np.asarray(conus["order"][:], dtype=np.int64) + n = order.size + pos_of = {int(c): i for i, c in enumerate(order)} + + bm = json.loads((RUNS / MEMBERS[0][3] / "baseline" / "manifest.json").read_text()) + gauges = list(bm["gage_ids"]) + gz = zarr.open_group(GAGES_ADJ, mode="r") + rows, cols = [], [] + for gi, s in enumerate(gauges): + sg = gz[s] + p = np.union1d(np.asarray(sg["indices_0"][:]), np.asarray(sg["indices_1"][:])) + rows.append(np.full(p.size, gi)); cols.append(p) + rows, cols = np.concatenate(rows), np.concatenate(cols) + A = sp.csr_matrix((np.ones(rows.size), (rows, cols)), shape=(len(gauges), n)).tocsc() + n_up = np.asarray(A.sum(1)).ravel() + print(f"{len(gauges)} gauges, {n:,} CONUS reaches, upstream divides per gauge median {np.median(n_up):.0f}") + + stores = [open_store(p) for _, p, _, _ in MEMBERS] + cover = np.zeros((len(MEMBERS), n), dtype=bool) + store_pos = [] + for k, g in enumerate(stores): + d = np.asarray(g["divide_id"][:], dtype=np.int64) + assert np.unique(d).size == d.size, "duplicate divide ids" + sp_ = np.array([pos_of.get(int(c), -1) for c in d]) + cover[k, sp_[sp_ >= 0]] = True + store_pos.append(sp_) + print(f" {MEMBERS[k][0]:>5}: {d.size:,} divides, {int((sp_ >= 0).sum()):,} in the CONUS network") + w = np.array([m[2] for m in MEMBERS])[:, None] + denom = (w * cover).sum(0) + a = np.where(denom > 0, w * cover / np.where(denom > 0, denom, 1.0), 0.0) # (members, n) + none = ~cover.any(0) + allc = cover.all(0) + print(f" CONUS reaches covered by all three {allc.sum():,}, by none {none.sum():,}") + + full_cover = (A @ (~allc).astype(float)) == 0 + print(f" gauges with every upstream divide in all three stores: {full_cover.sum()} of {len(gauges)}") + + S = {m[0]: np.zeros((len(gauges), N_DAYS)) for m in MEMBERS} + S_blend = np.zeros((len(gauges), N_DAYS)) + nan_counts = {} + for k, (lab, path, wk, _) in enumerate(MEMBERS): + g = stores[k]; Qr = g["Qr"] + i0 = time_index(g) + nd = Qr.shape[0] + n_nan = 0 + t = time.time() + for r0 in range(0, nd, BLOCK): + r1 = min(nd, r0 + BLOCK) + p = store_pos[k][r0:r1] + keep = p >= 0 + if not keep.any(): + continue + q = np.asarray(Qr[r0:r1, i0:i0 + N_DAYS], dtype=np.float64)[keep] + p = p[keep] + Ab = A[:, p] + used = np.asarray(Ab.sum(0)).ravel() > 0 + n_nan += int(np.isnan(q[used]).sum()) + q = np.nan_to_num(q, nan=0.0) # baseline nansum semantics + S[lab] += Ab @ q + S_blend += Ab @ (a[k, p][:, None] * q) + S[lab] += FILL * (A @ (~cover[k]).astype(float))[:, None] + nan_counts[lab] = n_nan + print(f" read {lab} in {time.time() - t:.0f} s, NaN values on upstream divides in window: {n_nan}") + S_blend += FILL * (A @ none.astype(float))[:, None] + + # The raw per-member sum must reproduce each member run's baseline. + check = {} + for lab, _, _, rid in MEMBERS: + b = RUNS / rid / "baseline" + m = json.loads((b / "manifest.json").read_text()) + assert list(m["gage_ids"]) == gauges, f"{rid}: baseline gauge order differs" + ref = np.fromfile(b / "predictions.f32", dtype=" routed flow); + 3. secondary: peak-shape diagnostic ratio from bias_absorber_per_gauge.csv; + 4. skill: blend model vs UH fixed-channel model (seeds 42 and 43), gauge-bootstrap CIs (levers/common.py), and + engine blend vs the offline output blend of the three member runs. + + /home/tbindas/projects/ddr/.venv/bin/python experiments/blend_prediction/blend_readout.py \ + --pred-data --readout +""" +import argparse +import importlib.util +import sys +from pathlib import Path + +import netCDF4 +import numpy as np +import pandas as pd + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE.parent / "levers")) +sys.path.insert(0, "/home/tbindas/projects/ddrs/.claude/skills/ddrs-eval-plots/scripts") +import common as C # noqa: E402 +from load_ddrs_predictions import load_baseline_f32 # noqa: E402 + +spec = importlib.util.spec_from_file_location("routing_lag", HERE.parent / "stage_roughness" / "routing_lag.py") +rl = importlib.util.module_from_spec(spec); spec.loader.exec_module(rl) +spec = importlib.util.spec_from_file_location("inflow_props", HERE / "inflow_props.py") +ip = importlib.util.module_from_spec(spec); spec.loader.exec_module(ip) + +BLEND = "2026-10-05T13-32-00Z-train-and-test" +DIST = "2026-09-13T17-22-30Z-train-and-test" +LSTM = "2026-09-13T13-55-03Z-train-and-test" +GAGES = "/home/tbindas/projects/ddr/references/gage_info/gages_3000.csv" +N_FLOOR, N_CEIL = 0.015, 0.2495 + + +def say(*a): + print(*a, flush=True) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--pred-data", required=True) + ap.add_argument("--readout", required=True) + a = ap.parse_args() + pdat, ro = Path(a.pred_data), Path(a.readout) + + # ---- 1. validity gates + say("== 1. validity gates (prediction doc section 5)") + z = np.load(pdat / "blend_inflow.npz") + bl = load_baseline_f32(C.RUNS / BLEND / "baseline") + gauges = z["gauges"].astype(str) + t_all = np.datetime64(str(z["start"])) + np.arange(int(z["n_days"])) + same_g = (bl.gage_ids.values.astype(str) == gauges).all() + tb = bl.time.values.astype("datetime64[D]") + ti = pd.Index(t_all).get_indexer(tb) + say(f"baseline gauges {bl.gage_ids.size}, days {tb.size}; gauge order equals offline npz: {same_g}; " + f"all baseline days inside npz span: {(ti >= 0).all()}") + run_S = bl.predictions.values.astype(np.float64) + off_S = z["S_blend"][:, ti[ti >= 0]].astype(np.float64) + run_S = run_S[:, ti >= 0] + rel = np.abs(run_S - off_S) / np.maximum(np.abs(off_S), 1e-9) + say(f"run baseline summed inflow vs offline S_blend: relative diff p50 {np.percentile(rel, 50):.2e} " + f"p99 {np.percentile(rel, 99):.2e} p99.9 {np.percentile(rel, 99.9):.2e} max {rel.max():.2e}; " + f"threshold p99 <= 1e-4 -> {'PASS' if np.percentile(rel, 99) <= 1e-4 else 'FAIL'}") + + ds = netCDF4.Dataset(ro / "dump" / "blend_kan_parameters.nc") + comid = np.asarray(ds["COMID"][:], dtype=np.int64) + n = np.asarray(ds["n"][:], dtype=np.float64) + full = np.load(pdat / "fields" / "comids_full_upstream_cover.npy") + mb = np.isin(comid, full) + say(f"n_0 finite everywhere: {np.isfinite(n).all()}; set (b) reaches {mb.sum():,}; " + f"at floor (<= {N_FLOOR}): {100 * np.mean(n[mb] <= N_FLOOR + 1e-6):.3f} %, at ceiling (>= {N_CEIL}): " + f"{100 * np.mean(n[mb] >= N_CEIL):.3f} %; min {n[mb].min():.4f} max {n[mb].max():.4f}") + + # ---- eval predictions + ids, tt, Pb, O = C.load_preds(BLEND) + ids42, t42, P42, O42 = C.load_preds(C.S42) + assert (t42 == tt).all() + P42 = C.align(ids42, P42, ids); O42 = C.align(ids42, O42, ids) + assert np.allclose(np.nan_to_num(O, nan=-1), np.nan_to_num(O42, nan=-1)) + say(f"eval gauges {len(ids)}, days {len(tt)} ({tt[0].date()}..{tt[-1].date()}); NaN in blend predictions: " + f"{int(np.isnan(Pb).sum())}; non-finite: {int((~np.isfinite(Pb)).sum())}") + i43, t43, P43, _ = C.load_preds(C.S43); assert (t43 == tt).all(); P43 = C.align(i43, P43, ids) + iD, tD, PD, _ = C.load_preds(DIST); assert (tD == tt).all(); PD = C.align(iD, PD, ids) + iL, tL, PL, _ = C.load_preds(LSTM); assert (tL == tt).all(); PL = C.align(iL, PL, ids) + + # ---- 2. router-added lag + say("\n== 2. router-added lag, run's own summed inflow -> its routed flow (best_lag, window -3..10 d)") + area = pd.read_csv(GAGES, dtype={"STAID": str}).set_index("STAID")["DRAIN_SQKM"].reindex(ids).values + bidx = pd.Index(bl.gage_ids.values.astype(str)).get_indexer(ids) + bt = pd.Index(tb).get_indexer(tt.values.astype("datetime64[D]")) + assert (bidx >= 0).all() and (bt >= 0).all() + Sb = bl.predictions.values[bidx][:, bt].astype(np.float64) + lags, flags = [], [] + for i in range(len(ids)): + l, _ = rl.best_lag(Sb[i], Pb[i]); lags.append(l); flags.append(ip.frac_lag(Sb[i], Pb[i], l)) + lag = pd.DataFrame(dict(gauge=ids, area=area, lag=lags, flag=flags)) + lag["cls"] = pd.cut(lag.area, ip.EDGES, labels=ip.CLASSES, right=False) + t = lag.groupby("cls", observed=True).agg(n=("lag", "size"), median=("lag", "median"), mean=("lag", "mean"), + sub_day_median=("flag", "median")) + say(t.round(3).to_string()) + lag.to_csv(ro / "router_lag_per_gauge.csv", index=False) + + # ---- 3. peak-shape diagnostic + say("\n== 3. peak-shape diagnostic (median over gauges of the basin-median n_0, >= 3 reaches)") + ba = pd.read_csv(ro / "fields" / "cross_arm_fields" / "bias_absorber_per_gauge.csv", dtype={"gauge": str}) + ba = ba[ba.n_reach >= 3] + ar = pd.Series(area, index=ids) + out = {} + for lo, hi, name in [(300, 3000, "300-3000"), (3000, 30000, "3000-30000")]: + g = [x for x in ids if lo <= ar[x] < hi] + u = ba[(ba.arm == "uh_0927") & ba.gauge.isin(g)].n0_basin + b = ba[(ba.arm == "blend") & ba.gauge.isin(g)].n0_basin + out[name] = (b.median() / u.median(), u.median(), b.median(), len(u), len(b)) + say(f"{name:>10}: UH median {u.median():.4f} ({len(u)} gauges), blend median {b.median():.4f} " + f"({len(b)} gauges), ratio {b.median() / u.median():.4f}") + r1, r2 = out["300-3000"][0], out["3000-30000"][0] + say(f"ratio(300-3000) {r1:.4f} < ratio(3000-30000) - 0.05 = {r2 - 0.05:.4f}? -> {r1 < r2 - 0.05}") + + # ---- 4. skill + say("\n== 4. skill on the test window (WY1996-2010), 2,365 gauges; bootstrap = paired over gauges, 2,000 draws") + M = {k: C.metrics(P, O) for k, P in {"blend": Pb, "uh42": P42, "uh43": P43, "dist": PD, "lstm": PL, + "off_w": 0.5 * P42 + 0.25 * PD + 0.25 * PL, + "off_eq": (P42 + PD + PL) / 3}.items()} + for k, m in M.items(): + say(f" {k:7s} NSE median {m.nse.median():.4f} KGE median {m.kge.median():.4f} " + f"(finite NSE {int(m.nse.notna().sum())}, KGE {int(m.kge.notna().sum())})") + rows = [] + for new, base in [("blend", "uh42"), ("blend", "uh43"), ("uh43", "uh42"), ("off_w", "uh42"), ("off_eq", "uh42"), + ("blend", "off_w"), ("blend", "off_eq"), ("blend", "dist"), ("blend", "lstm")]: + for met in ("nse", "kge"): + s = C.summarize(f"{new} vs {base}", M[new][met].values, M[base][met].values, met.upper()) + rows.append(s); say(C.fmt(s)) + pd.DataFrame(rows).to_csv(ro / "skill_summary.csv", index=False) + pd.DataFrame({"gauge": ids, "area": area, **{f"nse_{k}": m.nse.values for k, m in M.items()}, + **{f"kge_{k}": m.kge.values for k, m in M.items()}}).to_csv(ro / "skill_per_gauge.csv", index=False) + # size-class split of the blend-vs-UH(42) median change + say("\n blend minus UH42 median NSE by drainage-area class:") + d = pd.DataFrame({"area": area, "b": M["blend"].nse.values, "u": M["uh42"].nse.values, "ow": M["off_w"].nse.values}) + d["cls"] = pd.cut(d.area, ip.EDGES, labels=ip.CLASSES, right=False) + for c, g in d.groupby("cls", observed=True): + say(f" {c:>7} n={len(g):4d} blend {g.b.median():.4f} UH {g.u.median():.4f} offline {g.ow.median():.4f}") + + +if __name__ == "__main__": + main() diff --git a/experiments/blend_prediction/inflow_props.py b/experiments/blend_prediction/inflow_props.py new file mode 100644 index 00000000..4092be04 --- /dev/null +++ b/experiments/blend_prediction/inflow_props.py @@ -0,0 +1,156 @@ +"""Inflow properties of the blend and its members, by drainage-area class, before any blend training. + +Reads the per-gauge summed inflow written by blend_inflow.py (exact route) and, for +the members only, their routed test predictions. Population: the 2,365 gauges the +member runs were scored on (eval/predictions.zarr), over the eval window +1995-10-02..2010-09-29 (WY1996-2010 less the two edge days the eval drops). + +Per gauge and series (blend, UH retrospective, dHBV2 distributed, NH daily LSTM): + flag_* the same lags refined below one day by a parabola through the correlation + at the best whole-day lag and its two neighbours + lag_obs cross-correlation lag of summed inflow against observed flow, in days, + positive = observed is later, `best_lag` imported from + experiments/stage_roughness/routing_lag.py (window -3..10 d, >= 365 pairs) + fhv_obsday 100 (sum q on the top 2 % observed days - sum obs there) / sum obs there, + the definition cross_arm_fields.py uses (linear in q, so the blend's value + is the weighted mean of its members' exactly) + fhv_fdc the same on the flow-duration curve: top 2 % of sorted q against top 2 % + of sorted obs (Yilmaz et al. 2008), which does see peak flattening + alpha KGE alpha, std(q) / std(obs) + vol sum q / sum obs + r, nse correlation and NSE of the summed inflow +Members also get lag_routed = lag(summed -> routed), what their trained router adds. + + /home/tbindas/projects/ddr/.venv/bin/python experiments/blend_prediction/inflow_props.py --data +""" +import argparse +import importlib.util +import sys +from pathlib import Path + +import numpy as np +import pandas as pd + +sys.path.insert(0, "/home/tbindas/projects/ddrs/.claude/skills/ddrs-eval-plots/scripts") +from load_ddrs_predictions import load_predictions_zarr # noqa: E402 + +HERE = Path(__file__).resolve().parent +spec = importlib.util.spec_from_file_location("routing_lag", HERE.parent / "stage_roughness" / "routing_lag.py") +rl = importlib.util.module_from_spec(spec); spec.loader.exec_module(rl) +best_lag = rl.best_lag + +RUNS = Path("/home/tbindas/projects/ddrs/.ddrs/runs") +GAGES = "/home/tbindas/projects/ddr/references/gage_info/gages_3000.csv" +SERIES = {"blend": "S_blend", "uh": "S_uh", "dist": "S_dist", "lstm": "S_lstm"} +ROUTED = {"uh": "2026-09-27T07-29-47Z-train-and-test", "uh_0912": "2026-09-12T23-39-03Z-train-and-test", + "dist": "2026-09-13T17-22-30Z-train-and-test", "lstm": "2026-09-13T13-55-03Z-train-and-test"} +EDGES = [0, 300, 3000, 30000, np.inf] +CLASSES = ["<300", "300-3k", "3k-30k", ">30k"] + + +def corr_at(a, b, lag): + x, y = (a[: a.size - lag], b[lag:]) if lag >= 0 else (a[-lag:], b[: b.size + lag]) + m = np.isfinite(x) & np.isfinite(y) + return np.corrcoef(x[m], y[m])[0, 1] if m.sum() >= 365 else np.nan + + +def frac_lag(a, b, lag): + """Sub-day lag: parabola through the cross-correlation at lag-1, lag, lag+1.""" + r0, r1, r2 = corr_at(a, b, lag - 1), corr_at(a, b, lag), corr_at(a, b, lag + 1) + den = r0 - 2 * r1 + r2 + if not np.isfinite(den) or den >= 0: + return float(lag) + return lag + float(np.clip(0.5 * (r0 - r2) / den, -0.5, 0.5)) + + +def metrics(q, o): + m = np.isfinite(o) & (o > 0) & np.isfinite(q) + if m.sum() < 365: + return dict(fhv_obsday=np.nan, fhv_fdc=np.nan, alpha=np.nan, vol=np.nan, r=np.nan, nse=np.nan) + q, o = q[m], o[m] + k = max(1, int(0.02 * o.size)) + top = np.argsort(o)[-k:] + qs, os_ = np.sort(q)[-k:], np.sort(o)[-k:] + return dict(fhv_obsday=100 * (q[top].sum() - o[top].sum()) / o[top].sum(), + fhv_fdc=100 * (qs.sum() - os_.sum()) / os_.sum(), + alpha=q.std() / o.std(), vol=q.sum() / o.sum(), r=np.corrcoef(q, o)[0, 1], + nse=1 - np.sum((q - o) ** 2) / np.sum((o - o.mean()) ** 2)) + + +def main(): + ap = argparse.ArgumentParser(); ap.add_argument("--data", required=True) + d = Path(ap.parse_args().data) + z = np.load(d / "blend_inflow.npz") + gauges = z["gauges"].astype(str) + t_all = np.datetime64(str(z["start"])) + np.arange(int(z["n_days"])) + ev = load_predictions_zarr(RUNS / ROUTED["uh"] / "eval" / "predictions.zarr") + eg = ev.gage_ids.values.astype(str) + te = ev.time.values.astype("datetime64[D]") + gi = pd.Index(gauges).get_indexer(eg); assert (gi >= 0).all() + ti = pd.Index(t_all).get_indexer(te); assert (ti >= 0).all() + obs_b = z["obs"][gi][:, ti].astype(np.float64) + obs_e = ev.observations.values.astype(np.float64) + assert np.allclose(np.nan_to_num(obs_b, nan=-1), np.nan_to_num(obs_e, nan=-1), rtol=1e-6), "obs mismatch" + obs = np.where(obs_e <= 0, np.nan, obs_e) + area = pd.read_csv(GAGES, dtype={"STAID": str}).set_index("STAID")["DRAIN_SQKM"].reindex(eg).values + full = z["full_cover"][gi] + print(f"{eg.size} eval gauges, {te.size} days {te[0]}..{te[-1]}; full three-store cover at {full.sum()}") + + routed = {} + for lab, rid in ROUTED.items(): + e = load_predictions_zarr(RUNS / rid / "eval" / "predictions.zarr") + assert (e.gage_ids.values.astype(str) == eg).all() and (e.time.values.astype("datetime64[D]") == te).all() + routed[lab] = e.predictions.values.astype(np.float64) + + rows = [] + for s, key in SERIES.items(): + S = z[key][gi][:, ti].astype(np.float64) + for i in range(eg.size): + r = dict(series=s, gauge=eg[i], area=area[i], full_cover=bool(full[i])) + r["lag_obs"], r["r_lag_obs"] = best_lag(S[i], obs[i]) + r["flag_obs"] = frac_lag(S[i], obs[i], r["lag_obs"]) + r.update(metrics(S[i], obs[i])) + if s in routed: + r["lag_routed"], _ = best_lag(S[i], routed[s][i]) + r["flag_routed"] = frac_lag(S[i], routed[s][i], r["lag_routed"]) + if s == "uh": + r["lag_routed_0912"], _ = best_lag(S[i], routed["uh_0912"][i]) + rows.append(r) + df = pd.DataFrame(rows) + df["cls"] = pd.cut(df.area, EDGES, labels=CLASSES, right=False) + df.to_csv(d / "inflow_props_per_gauge.csv", index=False) + + def block(g): + return pd.Series({ + "n": len(g), + "lag_obs_med": g.lag_obs.median(), "lag_obs_mean": g.lag_obs.mean(), + "lead_ge1_pct": 100 * (g.lag_obs >= 1).mean(), + "flag_obs_med": g.flag_obs.median(), "flag_obs_mean": g.flag_obs.mean(), + "fhv_obsday": g.fhv_obsday.median(), "fhv_fdc": g.fhv_fdc.median(), + "alpha": g.alpha.median(), "vol": g.vol.median(), "r": g.r.median(), "nse": g.nse.median(), + "lag_routed_med": g.lag_routed.median() if "lag_routed" in g and g.lag_routed.notna().any() else np.nan, + "lag_routed_mean": g.lag_routed.mean() if "lag_routed" in g and g.lag_routed.notna().any() else np.nan, + "flag_routed_med": g.flag_routed.median() if "flag_routed" in g and g.flag_routed.notna().any() else np.nan, + }) + pd.set_option("display.width", 250); pd.set_option("display.max_columns", 40) + tab = df.groupby(["series", "cls"], observed=True).apply(block, include_groups=False) + tot = df.groupby("series").apply(block, include_groups=False) + tot.index = pd.MultiIndex.from_tuples([(s, "all") for s in tot.index], names=["series", "cls"]) + tab = pd.concat([tab, tot]).sort_index() + tab.to_csv(d / "inflow_props_by_class.csv") + print(tab.round(3).to_string()) + u = df[df.series == "uh"] + print(f"\nUH router-added lag, run 0927 vs 0912, 3k-30k km2: " + f"median {u[u.cls == '3k-30k'].lag_routed.median()} / {u[u.cls == '3k-30k'].lag_routed_0912.median()}, " + f"mean {u[u.cls == '3k-30k'].lag_routed.mean():.3f} / {u[u.cls == '3k-30k'].lag_routed_0912.mean():.3f}") + # blend vs weighted member mean, per gauge: how much does averaging flatten? + w = {"uh": 0.5, "dist": 0.25, "lstm": 0.25} + p = df.pivot(index="gauge", columns="series", values=["alpha", "fhv_fdc", "fhv_obsday", "vol"]) + for m in ["alpha", "fhv_fdc", "fhv_obsday", "vol"]: + wm = sum(w[k] * p[(m, k)] for k in w) + print(f"{m:>10}: blend median {p[(m, 'blend')].median():.3f}, weighted member mean (per gauge) median {wm.median():.3f}, " + f"blend below UH at {100 * (p[(m, 'blend')] < p[(m, 'uh')]).mean():.1f}% of gauges") + + +if __name__ == "__main__": + main() diff --git a/experiments/blend_prediction/member_fields.py b/experiments/blend_prediction/member_fields.py new file mode 100644 index 00000000..11769413 --- /dev/null +++ b/experiments/blend_prediction/member_fields.py @@ -0,0 +1,166 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.11" +# dependencies = ["numpy", "pandas", "netCDF4", "xarray", "zarr>=3", "scipy", "matplotlib", "icechunk>=2"] +# /// +"""Learned (n_0, gamma) fields of the blend's three members, recomputed from their checkpoint dumps. + +Two parts. + +1. Runs `experiments/stage_roughness/cross_arm_fields.py` unchanged on the member + runs. That script reads `/plot/kan_parameters.nc`; the UH fixed-channel run + 2026-09-27T07-29-47Z only has `plot/kan_parameters_epoch50.nc` (same final + checkpoint, `epoch_50_mb_9`), so the runs are presented through a read-only + symlink farm under --out and the script's RUNS constant is pointed at it. No + run directory is written. + +2. The two reach sets the prediction is stated on: + (a) every CONUS reach with a finite log10 drainage area (cross_arm_fields' set); + (b) the reaches whose whole upstream divide set, the reach itself included, is + in all three blend stores, so the blend reads exactly 0.50/0.25/0.25 on every + divide that drains to them. + For each member: n_0 and gamma median (with p25, p75), fraction of reaches at the + n_0 ceiling, Spearman rho(n_0, gamma), and medians by drainage-area class. + + uv run experiments/blend_prediction/member_fields.py --out [--add blend=] + +The primary statistic of research/specs/2026-09-30-blend-arm-n0-prediction.md is the +`n0_med` cell of row (label, b_full_cover) in member_fields.csv. +""" +import argparse +import contextlib +import importlib.util +import io +import json +from pathlib import Path + +import icechunk +import numpy as np +import pandas as pd +import zarr +from netCDF4 import Dataset +from scipy.stats import spearmanr + +RUNS = Path("/home/tbindas/projects/ddrs/.ddrs/runs") +HERE = Path(__file__).resolve().parent +CROSS = HERE.parent / "stage_roughness" / "cross_arm_fields.py" +ATTRS = "/home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc" +CONUS = "/home/tbindas/projects/ddr/data/merit_conus_adjacency.zarr" +STORES = [ + "/mnt/ssd1/data/icechunk/merit_dhbv2_UH_retrospective.ic", + "/mnt/ssd1/data/icechunk/daily_dhbv2_distributed_aorc2f_merit_unit_catchments_rechunked.ic", + "/mnt/ssd1/data/icechunk/daily_lstm_merit_unit_catchments.ic", +] +# label, run id, parameter dump inside the run (all are the epoch_50_mb_9 head) +MEMBERS = [ + ("uh_0927", "2026-09-27T07-29-47Z-train-and-test", "plot/kan_parameters_epoch50.nc"), + ("uh_0912", "2026-09-12T23-39-03Z-train-and-test", "plot/kan_parameters.nc"), + ("dist", "2026-09-13T17-22-30Z-train-and-test", "plot/kan_parameters.nc"), + ("lstm", "2026-09-13T13-55-03Z-train-and-test", "plot/kan_parameters.nc"), +] +SIZE = [(-np.inf, 1), (1, 2), (2, 3), (3, 4), (4, np.inf)] # log10 km2 +SIZE_LAB = ["<10", "10-100", "100-1e3", "1e3-1e4", ">1e4"] +N_CEIL = 0.2495 + + +def store_divides(path): + repo = icechunk.Repository.open(icechunk.local_filesystem_storage(path)) + g = zarr.open_group(repo.readonly_session("main").store, mode="r") + return np.asarray(g["divide_id"][:], dtype=np.int64) + + +def upstream_full_cover(): + """Per CONUS reach: is every divide upstream of it (itself included) in all three stores?""" + c = zarr.open_group(CONUS, mode="r") + order = np.asarray(c["order"][:], dtype=np.int64) + rows = np.asarray(c["indices_0"][:], dtype=np.int64) # downstream position + cols = np.asarray(c["indices_1"][:], dtype=np.int64) # upstream position + assert np.all(rows >= cols), "adjacency not lower-triangular" + own = np.ones(order.size, dtype=bool) + for p in STORES: + own &= np.isin(order, store_divides(p)) + full = own.copy() + # positions are topologically ordered, so processing edges by increasing + # downstream position settles every upstream reach before its receiver + for k in np.argsort(rows, kind="stable"): + if not full[cols[k]]: + full[rows[k]] = False + return pd.Series(own, index=order), pd.Series(full, index=order) + + +def summarize(df, lab, set_name): + r = {"member": lab, "set": set_name, "n_reaches": len(df), + "n0_med": df.n0.median(), "n0_p25": df.n0.quantile(.25), "n0_p75": df.n0.quantile(.75), + "n0_ceil_pct": 100 * np.mean(df.n0 >= N_CEIL), + "gamma_med": df.gamma.median(), "gamma_p25": df.gamma.quantile(.25), "gamma_p75": df.gamma.quantile(.75), + "gamma_zero_pct": 100 * np.mean(df.gamma < 0.001), + "rho_n0_gamma": spearmanr(df.n0, df.gamma).statistic} + for (lo, hi), sl in zip(SIZE, SIZE_LAB): + m = (df.la >= lo) & (df.la < hi) + r[f"n0_{sl}"] = df.n0[m].median(); r[f"gamma_{sl}"] = df.gamma[m].median(); r[f"n_{sl}"] = int(m.sum()) + return r + + +def main(): + ap = argparse.ArgumentParser(); ap.add_argument("--out", required=True) + ap.add_argument("--add", action="append", default=[], + help="extra run scored the same way, label= (dump plot/kan_parameters.nc); " + "for the blend read-out, after the run has finished") + args = ap.parse_args() + out = Path(args.out); out.mkdir(parents=True, exist_ok=True) + for x in args.add: + lab, rid, *dump = x.split("=") # label=[=] + MEMBERS.append((lab, rid, dump[0] if dump else "plot/kan_parameters.nc")) + + # Part 1: cross_arm_fields.py, unchanged, through a symlink farm. + farm = out / "runs_view" + for _, rid, nc in MEMBERS: + d = farm / rid; (d / "plot").mkdir(parents=True, exist_ok=True) + for sub in ("baseline", "eval"): + if not (d / sub).exists(): + (d / sub).symlink_to(RUNS / rid / sub) + if not (d / "plot" / "kan_parameters.nc").exists(): + (d / "plot" / "kan_parameters.nc").symlink_to(RUNS / rid / nc) # absolute nc overrides RUNS + spec = importlib.util.spec_from_file_location("cross_arm_fields", CROSS) + caf = importlib.util.module_from_spec(spec); spec.loader.exec_module(caf) + caf.RUNS = farm + import sys + argv = sys.argv + sys.argv = [str(CROSS)] + [f"{lab}={rid}" for lab, rid, _ in MEMBERS] + ["--out", str(out / "cross_arm_fields")] + buf = io.StringIO() + with contextlib.redirect_stdout(buf): + caf.main() + sys.argv = argv + (out / "cross_arm_fields.txt").write_text(buf.getvalue()) + print(buf.getvalue()) + + # Part 2: the two reach sets. + own, full = upstream_full_cover() + print(f"CONUS reaches: {own.size:,}; own divide in all three stores {own.sum():,}; " + f"whole upstream set in all three stores {full.sum():,}") + at = Dataset(ATTRS) + la_by = pd.Series(np.asarray(at["log10_uparea"][:], dtype=np.float64), + index=np.asarray(at["COMID"][:], dtype=np.int64)) + rows = [] + for lab, rid, nc in MEMBERS: + ds = Dataset(RUNS / rid / nc) + assert ds.getncattr("checkpoint").endswith("epoch_50_mb_9/head"), ds.getncattr("checkpoint") + c = np.asarray(ds["COMID"][:], dtype=np.int64) + df = pd.DataFrame({"n0": np.asarray(ds["n"][:], dtype=np.float64), + "gamma": np.asarray(ds["gamma"][:], dtype=np.float64), + "la": la_by.reindex(c).values}, index=c) + df = df[np.isfinite(df.la)] + rows.append(summarize(df, lab, "a_all")) + rows.append(summarize(df[full.reindex(df.index, fill_value=False).values], lab, "b_full_cover")) + t = pd.DataFrame(rows) + t.to_csv(out / "member_fields.csv", index=False) + pd.set_option("display.width", 250); pd.set_option("display.max_columns", 50) + print(t.round(4).to_string(index=False)) + np.save(out / "comids_full_upstream_cover.npy", full.index.values[full.values]) + (out / "member_fields.json").write_text(json.dumps({ + "n_conus": int(own.size), "n_own_cover": int(own.sum()), "n_full_upstream_cover": int(full.sum()), + "members": [{"label": l, "run": r, "dump": n} for l, r, n in MEMBERS]}, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/experiments/blend_prediction/uh_repro_score.py b/experiments/blend_prediction/uh_repro_score.py new file mode 100644 index 00000000..5817546b --- /dev/null +++ b/experiments/blend_prediction/uh_repro_score.py @@ -0,0 +1,38 @@ +"""Score the UH fixed-channel reproduction retrain against the existing UH model (2026-09-27T07-29-47Z). +Usage: python uh_repro_score.py ; reads dump/uh_repro_kan_parameters.nc and fields_repro/.""" +import sys +from pathlib import Path +import netCDF4 +import numpy as np +from scipy.stats import spearmanr + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE.parent / "levers")) +sys.path.insert(0, "/home/tbindas/projects/ddrs/.claude/skills/ddrs-eval-plots/scripts") +import common as C # noqa: E402 + +ro = Path(sys.argv[1]) +REPRO = "2026-10-07T22-05-06Z-train-and-test" +a = netCDF4.Dataset(ro / "dump" / "uh_repro_kan_parameters.nc") +b = netCDF4.Dataset("/home/tbindas/projects/ddrs/.ddrs/runs/2026-09-27T07-29-47Z-train-and-test/plot/kan_parameters_epoch50.nc") +ca = np.asarray(a["COMID"][:], dtype=np.int64); cb = np.asarray(b["COMID"][:], dtype=np.int64) +print("COMID order identical:", (ca == cb).all(), len(ca)) +m = np.isin(ca, np.load(ro / "fields_repro" / "comids_full_upstream_cover.npy")) +print("set (b) reaches:", m.sum()) +for v in ("n", "gamma"): + x = np.asarray(a[v][:], float); y = np.asarray(b[v][:], float) + d = np.abs(x - y) + print(f"{v}: set(b) repro median {np.median(x[m]):.6f} p25 {np.percentile(x[m],25):.6f} p75 {np.percentile(x[m],75):.6f} | " + f"UH median {np.median(y[m]):.6f} p25 {np.percentile(y[m],25):.6f} p75 {np.percentile(y[m],75):.6f}") + print(f" per-reach |diff| on (b): max {d[m].max():.3e} median {np.median(d[m]):.3e} p99 {np.percentile(d[m],99):.3e}; " + f"all reaches max {d.max():.3e}; Spearman(b) {spearmanr(x[m], y[m])[0]:.6f}; median diff of medians {np.median(x[m])-np.median(y[m]):+.2e}") +ids, tt, P, O = C.load_preds(REPRO) +i0, t0, P0, O0 = C.load_preds(C.S42) +assert (t0 == tt).all() +P0 = C.align(i0, P0, ids) +print("eval gauges", len(ids), "NaN preds", int(np.isnan(P).sum())) +M = C.metrics(P, O); M0 = C.metrics(P0, O) +for k, mm in (("retrain", M), ("existing UH", M0)): + print(f"{k:12s} NSE median {mm.nse.median():.4f} KGE median {mm.kge.median():.4f}") +dn = (M.nse - M0.nse).abs(); print("per-gauge |dNSE| max", float(dn.max()), "median", float(dn.median())) +print("max |pred diff| m3/s:", float(np.nanmax(np.abs(P - P0)))) diff --git a/experiments/blend_prediction/uh_reproduction.yaml b/experiments/blend_prediction/uh_reproduction.yaml new file mode 100644 index 00000000..a85a1c3e --- /dev/null +++ b/experiments/blend_prediction/uh_reproduction.yaml @@ -0,0 +1,328 @@ +# dam_release_full_off +# +# Full-population no-dam arm for the learned dam release: every reach routed as a +# channel. Matched arm: dam_release_full_learned.yaml (same seed, same config, plus +# the learned release at all 1,024 NID dam COMIDs). The pair measures the per-gauge +# change the learned release makes over the whole eval population, at gauges below a +# NID dam >= 10 MCM and at gauges without one. +# +# The sr_n0_gamma recipe verbatim (run 2026-09-12T23-39-03Z, config snapshot copied +# unchanged below this header): full gages_3000 population (2,365 gauges after +# the filters), train 1981-10-01..1995-09-30, test 1995-10-01..2010-09-30, seed 42. +# Compare only with its matched arm, on the same gauge CSV (trap T19). +# +# Run from the repo root, CPU (trap T13), main workspace: +# /target/release/ddrs --config config/experiments/dam_release_full_off.yaml \ +# --workspace /home/tbindas/projects/ddrs/.ddrs run --workflow train-and-test --backend cpu + +# sr_n0_gamma +# +# The KAN emits ONLY n_0 and gamma. p_spatial is fixed at 21 and q_spatial at +# 0.65 (params.defaults), so the channel shape is prescribed and the head +# carries the two roughness parameters of n(d) = n_0 * (d/d_ref)^(-gamma). +# +# Why (2026-09-12, findings §36.5-36.7): with p and q free, a learned gamma +# rode the width channel (rho(p, gamma) 0.89, rho(q, gamma) 0.76) and the +# head traded q against gamma, so the sweep could not tell degeneracy from +# unidentifiability. Fixing p and q removes that direction by construction, +# and a two-output head is what the loss-landscape and adjoint instruments +# can probe over (n_0, gamma) directly. +# +# Matched control: sr_n0_only.yaml (n_0 alone, same fixed p and q, gamma = 0). +# +# ddrs --config config/experiments/sr_n0_gamma.yaml \ +# --workspace /home/tbindas/projects/ddrs/.ddrs \ +# run --workflow train-and-test --plot --backend cpu + +# head_shared_linear +# +# One arm of the head-topology experiment. Derived from the 500-update +# nse-batch baseline 2026-09-10T21-21-48Z-conus-train-and-test +# (NSE 0.7376 / KGE 0.7600) by changing the kan_head section and nothing +# else. Screen that motivated these arms: +# .ddrs/experiments/head-arch/2026-09-11-screen +# experiments/head_arch/analyze.py +# +# MATCHED CONTROL. Current architecture: one shared trunk, Linear(H,P) +# read-out. Measured trunk effective rank 4.36 of 21; init affine R^2 between +# logit(n) and logit(q) at the chance level for that rank. +# +# Run (note --workspace: without it .ddrs lands beside this config): +# ddrs --config config/experiments/head_shared_linear.yaml \ +# --workspace /home/tbindas/projects/ddrs/.ddrs \ +# run --workflow train-and-test + +# 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: + # CONUS MERIT inputs as available on the workstation. + 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 + # Icechunk dHBV2 UH retrospective (divide-major, daily m³/s) + streamflow: /mnt/ssd1/data/icechunk/merit_dhbv2_UH_retrospective.ic + # Icechunk USGS daily observations + observations: /mnt/ssd1/data/icechunk/usgs_daily_observations + gages: /home/tbindas/projects/ddr/references/gage_info/gages_3000.csv + +experiment: + batch_size: 64 + start_time: 1981/10/01 + end_time: 1995/09/30 + epochs: 50 # 2026-09-10: 30 -> 50, with grad_accum_steps 4 this is 500 optimizer updates (was 60) + 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 + # 2026-09-10: 20 -> 4. The 2026-09-08 run made 60 optimizer updates in total + # (37 micro-batches/epoch at 2,365 kept gauges, 2 steps/epoch, 30 epochs) and + # findings §21 shows it never reached a stationary point: on BOTH the training + # and testing windows the loss still falls if roughness rises, at ~79% of + # gauges (trimmed gradient alignment 0.96/0.94). At accum 4 that is 10 + # steps/epoch, so 500 updates over 50 epochs, effective batch 4*64 = 256. + # Wall clock is forward-pass bound at 37 forwards/epoch either way: the + # 30-epoch run trained in 64 min, so 50 epochs is ~107 min. The "27.5 s per + # micro-batch" figure above is stale; measured here it is ~3.5 s. + grad_accum_steps: 4 + # 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. + # 2026-09-10: flattened. n froze at epoch 10 in the previous run, which is + # exactly where the old schedule cut 0.005 -> 0.001, so "the gradient + # vanished" and "the steps became too small" were not separable from the + # trajectory (§10 correction). §21 shows the gradient had NOT vanished. + # Keys are the FIRST epoch at which the lr applies. At 10 steps/epoch: + # epochs 1-20 @ 0.005 (200 updates) + # epochs 21-40 @ 0.002 (200 updates) + # epochs 41-50 @ 0.001 (100 updates) + # Displacement budget ~1.3 against the previous ~0.13. + # WATCH `n_at_floor`: lr 0.01 pinned 47.3% of CONUS at the 0.015 floor and + # lr 2e-3 collapsed 0% -> 25-36% by epoch 5 in earlier small-batch runs. + # Holding 0.005 for 200 updates at an effective batch of 256 is the real risk + # in this plan. If n_at_floor passes a few percent, drop the first rate to + # 0.002 and keep the step count. + learning_rate: + 1: 0.005 + 21: 0.002 + 41: 0.001 + 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 + # All three learned together. The baseline fixed p_spatial at 21 + # (docs/2026-09-08-fixed-p-assessment.md); §30 argued that learning p and q + # as a pair relocates the free parameter onto something falsifiable against + # remotely sensed river width, and §31 showed q alone was only n amplified + # fourfold. Because this differs from the baseline in TWO ways (learnable set + # and head topology), the shared-linear arm is the matched control: read + # these arms against each other, never against 0.7376/0.7600. + learnable_parameters: + - n + - gamma + # 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: + # Stage-roughness exponent. Non-negative because channels get SMOOTHER as + # they fill; capped at 0.5 because the depth exponent is 3/(5+3q+3*gamma) + # and roughness grows without bound as depth goes to zero. 0.183 reproduces + # observed at-a-station hydraulic geometry when paired with q ~ 0.65, and + # the sigmoid centre of this box is 0.25. + gamma: [0.0, 0.5] + 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 + # Fixed width exponent: the at-a-station Leopold & Maddock value b/f = 0.65 + # (design doc §2.2), which with gamma ~ 0.18 reproduces the observed + # at-a-station exponents. Must lie inside parameter_ranges.q_spatial. + q_spatial: 0.65 + 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/experiments/cross_product/cp_common.py b/experiments/cross_product/cp_common.py new file mode 100644 index 00000000..f622a18b --- /dev/null +++ b/experiments/cross_product/cp_common.py @@ -0,0 +1,148 @@ +"""Shared table, loaders and statistics for the cross-product analysis (E-10, E-08). + +The six fixed-channel models: the same n_0 + gamma head, channel (p = 21, q = 0.65), 2,365 gauges, +network, seed and 500-update recipe; only the unit-catchment inflow product differs. Run ids are the +outline's section 6 (`~/papers/ddr_equifinality/outline_2026-09-25.md`), unless the E-01 lineage doc +re-points them. +""" +from __future__ import annotations + +import sys +from pathlib import Path + +import numpy as np +import pandas as pd + +sys.path.insert(0, "/home/tbindas/projects/ddrs/.claude/skills/ddrs-eval-plots/scripts") +from load_ddrs_predictions import load_baseline_f32, load_predictions_zarr # noqa: E402 + +RUNS = Path("/home/tbindas/projects/ddrs/.ddrs/runs") +OUT = Path("/home/tbindas/projects/ddrs/.ddrs/experiments/cross_product/2026-09-30") +HERE = Path(__file__).resolve().parent +FIG = HERE.parents[1] / "research" / "figures" / "cross-product" +GAGES = "/home/tbindas/projects/ddr/references/gage_info/gages_3000.csv" +GAGES_ADJ = "/home/tbindas/projects/ddr/data/merit_gages_conus_adjacency.zarr" +ATTRS = "/home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc" +IC = "/mnt/ssd1/data/icechunk/" + +# key, descriptive label, run id, unit-catchment Q' store (from each run's config.yaml data_sources.streamflow) +MODELS = [ + ("uh", "UH fixed-channel", "2026-09-12T23-39-03Z-train-and-test", IC + "merit_dhbv2_UH_retrospective.ic"), + ("dhbv2_lumped", "dHBV2 lumped", "2026-09-13T13-56-50Z-train-and-test", + IC + "daily_dhbv2_lumped_aorc2f_merit_unit_catchments_rechunked.ic"), + ("dhbv2_dist", "dHBV2 distributed", "2026-09-13T17-22-30Z-train-and-test", + IC + "daily_dhbv2_distributed_aorc2f_merit_unit_catchments_rechunked.ic"), + ("hydrodl_lstm", "hydroDL LSTM", "2026-09-13T17-21-58Z-train-and-test", + IC + "daily_hydrodl_lstm_aorc2f_merit_unit_catchments_rechunked.ic"), + ("nh_daily_lstm", "NH daily LSTM", "2026-09-13T13-55-03Z-train-and-test", IC + "daily_lstm_merit_unit_catchments.ic"), + ("nh_hourly_mts", "NH hourly MTS-LSTM", "2026-09-13T14-14-46Z-train-and-test", IC + "hourly_lstm_merit_unit_catchments.ic"), +] +KEYS = [m[0] for m in MODELS] +LABEL = {m[0]: m[1] for m in MODELS} +RUN = {m[0]: m[2] for m in MODELS} +STORE = {m[0]: m[3] for m in MODELS} + +TEST0, TEST1 = pd.Timestamp("1995-10-01"), pd.Timestamp("2010-09-30") # WY1996-2010 +AREA_CLASSES = [("< 300", 0, 300), ("300-3,000", 300, 3000), ("3,000-30,000", 3000, 30000), ("> 30,000", 30000, np.inf)] +N_BOX = (0.015, 0.25) +G_BOX = (0.0, 0.5) + + +def load_series(): + """Routed, unrouted (summed Q') and observed daily series on one gauge x day grid for all six models. + + Returns (gauges, days, R[k] [g, t], S[k] [g, t], O [g, t]). Gauges and days are the intersection of every + model's eval predictions and baseline. Observations come from the UH fixed-channel run and are checked + identical (NaN-aware) in every other run, both in the eval zarr and in the baseline. + """ + ev = {k: load_predictions_zarr(RUNS / RUN[k] / "eval" / "predictions.zarr") for k in KEYS} + bl = {k: load_baseline_f32(RUNS / RUN[k] / "baseline") for k in KEYS} + g = ev[KEYS[0]].gage_ids.values + for k in KEYS: + g = np.intersect1d(g, ev[k].gage_ids.values) + g = np.intersect1d(g, bl[k].gage_ids.values) + t = ev[KEYS[0]].time.values + for k in KEYS: + t = np.intersect1d(t, ev[k].time.values) + t = np.intersect1d(t, bl[k].time.values) + R, S = {}, {} + O = ev[KEYS[0]].observations.sel(gage_ids=g, time=t).values.astype(np.float64) + for k in KEYS: + R[k] = ev[k].predictions.sel(gage_ids=g, time=t).values.astype(np.float64) + S[k] = bl[k].predictions.sel(gage_ids=g, time=t).values.astype(np.float64) + for name, oo in (("eval", ev[k].observations), ("baseline", bl[k].observations)): + ok = oo.sel(gage_ids=g, time=t).values.astype(np.float64) + same = np.array_equal(np.isnan(ok), np.isnan(O)) and np.allclose(ok[np.isfinite(ok)], O[np.isfinite(O)], rtol=1e-6) + assert same, f"observations differ: {k} {name}" + return g, pd.DatetimeIndex(t), R, S, O + + +def gauge_area(gauges): + meta = pd.read_csv(GAGES, dtype={"STAID": str}).set_index("STAID") + return meta["DRAIN_SQKM"].reindex(gauges).to_numpy(dtype=float) + + +def basin_comids(gauges): + """Upstream COMIDs per gauge from the gauge subgraph store the six runs trained on.""" + import zarr + gz = zarr.open_group(GAGES_ADJ, mode="r") + out = {} + for s in gauges: + try: + out[s] = np.asarray(gz[str(s)]["order"][:], dtype=np.int64) + except KeyError: + pass + return out + + +def load_field(key): + """Per-reach learned n_0 and gamma (plot/kan_parameters.nc), indexed by COMID, as cross_arm_fields.py reads them.""" + from netCDF4 import Dataset + ds = Dataset(RUNS / RUN[key] / "plot" / "kan_parameters.nc") + c = np.asarray(ds["COMID"][:], dtype=np.int64) + return pd.DataFrame({"n0": np.asarray(ds["n"][:], dtype=np.float64), + "gamma": np.asarray(ds["gamma"][:], dtype=np.float64)}, index=c) + + +def water_year(t): + return np.asarray(t.year + (t.month >= 10)) + + +def bands(d, tt): + """Split d[g, t] into (water-year mean, 30-day MA of the rest, remainder). + + Copied from `bands()` in experiments/levers/blend_mechanism.py on branch multi-product-forcing + (worktree ~/projects/ddrs/.claude/worktrees/agent-a92e512a7c47c97b4, file untracked there as of + 2026-09-30). One change: d may hold NaN on days without an observation, so the water-year mean is a + nanmean and the moving average skips NaN (pandas rolling mean counts only finite values toward + min_periods), and all three bands are NaN on the days d is NaN. With no NaN the result is identical to the + original. + """ + wy = water_year(tt) + low = np.full_like(d, np.nan) + for y in np.unique(wy): + s = wy == y + with np.errstate(invalid="ignore"): + low[:, s] = np.nanmean(d[:, s], axis=1, keepdims=True) + rest = d - low + mid = pd.DataFrame(rest.T).rolling(30, center=True, min_periods=1).mean().to_numpy().T + ok = np.isfinite(d) + mid = np.where(ok, mid, np.nan) + return np.where(ok, low, np.nan), mid, rest - mid + + +def boot_median(x, B=2000, seed=0, stat=np.nanmedian): + """Statistic of x with a gauge-bootstrap 95 % percentile CI. Returns (est, lo, hi, n).""" + x = np.asarray(x, dtype=float) + x = x[np.isfinite(x)] + if x.size == 0: + return np.nan, np.nan, np.nan, 0 + rng = np.random.default_rng(seed) + idx = rng.integers(0, x.size, size=(B, x.size)) + bs = stat(x[idx], axis=1) + return float(stat(x)), float(np.percentile(bs, 2.5)), float(np.percentile(bs, 97.5)), int(x.size) + + +def fmt_ci(t, nd=3): + est, lo, hi, n = t + return f"{est:.{nd}f} [{lo:.{nd}f}, {hi:.{nd}f}]" diff --git a/experiments/cross_product/e08_fields.py b/experiments/cross_product/e08_fields.py new file mode 100755 index 00000000..72267e46 --- /dev/null +++ b/experiments/cross_product/e08_fields.py @@ -0,0 +1,169 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.11" +# dependencies = ["numpy", "pandas", "xarray", "zarr>=3", "netCDF4", "scipy", "matplotlib"] +# /// +"""E-08: n_0 level vs spatial pattern across the six fixed-channel models, per reach. + +Fields as cross_arm_fields.py loads them (plot/kan_parameters.nc: `n` is n_0, `gamma` is the stage exponent), +restricted to reaches with a finite log10_uparea in the attribute table, as there. Two reach sets: + all every such reach (the section 38.2 population, so its numbers can be reproduced); + filtered reaches covered by every product (COMID on each store's divide_id axis, so ddrs read real Q' and + not its 0.001 m3/s fill) with median daily Q' >= 1e-3 m3/s over WY1996-2010 in every product + (qprime_medians.py). +Per pair of products: Spearman rho of n_0 and of gamma, median per-reach ratio of n_0 (and the ratio of the +medians), median per-reach difference of gamma. Per reach: range-normalised cross-product spread (paper Eq. +normspread as cross_arm_fields.py implements it: (max - min) / box width, box n_0 [0.015, 0.25], gamma +[0, 0.5]). Plus a two-way decomposition of log n_0 over (reach, product): product main effect (level), reach +main effect (pattern shared by all products), and the interaction (product-specific pattern). + + experiments/cross_product/e08_fields.py +Writes experiments/cross_product/e08_fields.txt and research/figures/cross-product/n0_rank_matrix.png. +""" +import io +import itertools + +import numpy as np +import pandas as pd +from netCDF4 import Dataset +from scipy.stats import spearmanr + +import cp_common as C + +FLOOR = 1e-3 + + +def summarise(name, F, idx, say): + N = np.stack([F[k].n0.reindex(idx).to_numpy() for k in C.KEYS]) + G = np.stack([F[k].gamma.reindex(idx).to_numpy() for k in C.KEYS]) + say(f"\n== {name}: {idx.size:,} reaches") + say(f"{'product':>20} {'n0 median':>10} {'n0 p10-p90':>14} {'at n0 ceiling':>13} {'gamma median':>12} {'gamma = 0':>9}") + for i, k in enumerate(C.KEYS): + say(f"{C.LABEL[k]:>20} {np.median(N[i]):10.4f} {np.percentile(N[i], 10):6.3f}-{np.percentile(N[i], 90):.3f} " + f"{100 * np.mean(N[i] >= 0.2495):12.1f}% {np.median(G[i]):12.3f} {100 * np.mean(G[i] < 0.001):8.1f}%") + rows = [] + say(f"\n{'pair':>42} {'rho n0':>7} {'rho gamma':>9} {'median n0 ratio':>15} {'ratio of medians':>16} {'median gamma diff':>17}") + for i, j in itertools.combinations(range(len(C.KEYS)), 2): + a, b = C.KEYS[i], C.KEYS[j] + rn = spearmanr(N[i], N[j]).statistic + rg = spearmanr(G[i], G[j]).statistic + mr = np.median(N[i] / N[j]) + rm = np.median(N[i]) / np.median(N[j]) + gd = np.median(G[i] - G[j]) + rows.append(dict(a=a, b=b, rho_n0=rn, rho_gamma=rg, med_ratio=mr, ratio_med=rm, gamma_diff=gd)) + say(f"{C.LABEL[a] + ' / ' + C.LABEL[b]:>42} {rn:+7.3f} {rg:+9.3f} {mr:15.3f} {rm:16.3f} {gd:+17.3f}") + R = pd.DataFrame(rows) + sn = (N.max(0) - N.min(0)) / (C.N_BOX[1] - C.N_BOX[0]) + sg = (G.max(0) - G.min(0)) / (C.G_BOX[1] - C.G_BOX[0]) + say(f"\nrange-normalised cross-product spread per reach: n0 median {np.median(sn):.3f} (p10 {np.percentile(sn, 10):.3f}, " + f"p90 {np.percentile(sn, 90):.3f}); gamma median {np.median(sg):.3f} (p10 {np.percentile(sg, 10):.3f}, " + f"p90 {np.percentile(sg, 90):.3f})") + say(f"pairwise rho n0: median {R.rho_n0.median():.3f}, range {R.rho_n0.min():+.3f}..{R.rho_n0.max():+.3f}; " + f"rho gamma: median {R.rho_gamma.median():.3f}, range {R.rho_gamma.min():+.3f}..{R.rho_gamma.max():+.3f}") + # level vs pattern: two-way decomposition of log n0 + L = np.log(N) + grand = L.mean() + prod_eff = L.mean(1, keepdims=True) - grand + reach_eff = L.mean(0, keepdims=True) - grand + inter = L - grand - prod_eff - reach_eff + ss_tot = ((L - grand) ** 2).sum() + ss_p = (prod_eff ** 2).sum() * L.shape[1] + ss_r = (reach_eff ** 2).sum() * L.shape[0] + ss_i = (inter ** 2).sum() + say(f"log n0 two-way decomposition (shares of total SS): product level {ss_p / ss_tot:.3f}, shared reach pattern " + f"{ss_r / ss_tot:.3f}, product-specific pattern {ss_i / ss_tot:.3f} (sum {(ss_p + ss_r + ss_i) / ss_tot:.3f})") + # same with rank-transformed fields per product: how much of the pattern is shared once level is removed entirely + rk = np.stack([pd.Series(x).rank().to_numpy() for x in N]) + rk = (rk - rk.mean(1, keepdims=True)) + shared = rk.mean(0, keepdims=True) + say(f"rank-pattern share explained by the cross-product mean rank: " + f"{1 - ((rk - shared) ** 2).sum() / (rk ** 2).sum():.3f}") + return R, N, G, sn, sg + + +def main(): + buf = io.StringIO() + + def say(*a): + print(*a) + print(*a, file=buf) + + at = Dataset(C.ATTRS) + ac = np.asarray(at["COMID"][:], dtype=np.int64) + ua = np.asarray(at["log10_uparea"][:], dtype=np.float64) + la = pd.Series(ua, index=ac) + F = {} + for k in C.KEYS: + f = C.load_field(k) + ok = np.isfinite(la.reindex(f.index).to_numpy()) + F[k] = f[ok] + idx_all = F[C.KEYS[0]].index + for k in C.KEYS[1:]: + idx_all = idx_all.intersection(F[k].index) + + cov = idx_all + say("coverage and wetness filter, per product (on the all-reach set):") + for k in C.KEYS: + z = np.load(C.OUT / f"qprime_median_{k}.npz") + med = pd.Series(z["median"], index=z["divide_id"]) + on = idx_all.isin(med.index) + wet = on & (med.reindex(idx_all).to_numpy() >= FLOOR) + say(f"{C.LABEL[k]:>20}: store has {med.size:,} divides; {on.sum():,} of {idx_all.size:,} reaches covered; " + f"{wet.sum():,} covered with median Q' >= {FLOOR:g} m3/s") + cov = cov[(cov.isin(med.index)) & (med.reindex(cov).to_numpy() >= FLOOR)] + covered_all = idx_all + for k in C.KEYS: + z = np.load(C.OUT / f"qprime_median_{k}.npz") + covered_all = covered_all[covered_all.isin(z["divide_id"])] + say(f"covered by every product: {covered_all.size:,}; and wet in every product: {cov.size:,}") + + Ra, *_ = summarise("all reaches (section 38.2 population)", F, idx_all, say) + Rc, *_ = summarise("covered by every product, no flow floor", F, covered_all, say) + Rf, N, G, sn, sg = summarise(f"filtered: covered by every product, median Q' >= {FLOOR:g} m3/s in all six", F, cov, say) + + # what the filter changed + say("\n== What the filter changed (filtered minus all-reach)") + m = Ra.merge(Rf, on=["a", "b"], suffixes=("_all", "_filt")) + for _, r in m.iterrows(): + say(f"{C.LABEL[r.a] + ' / ' + C.LABEL[r.b]:>42} rho n0 {r.rho_n0_all:+.3f} -> {r.rho_n0_filt:+.3f} " + f"rho gamma {r.rho_gamma_all:+.3f} -> {r.rho_gamma_filt:+.3f} median n0 ratio {r.med_ratio_all:.3f} -> {r.med_ratio_filt:.3f}") + + # the reaches the filter removes: what are they? + drop = idx_all.difference(cov) + say(f"\nreaches removed by the filter: {drop.size:,}; median log10 upstream area removed {np.median(la.reindex(drop)):.2f} " + f"vs kept {np.median(la.reindex(cov)):.2f}") + (C.HERE / "e08_fields.txt").write_text(buf.getvalue()) + pd.concat([Ra.assign(set="all"), Rc.assign(set="covered"), Rf.assign(set="filtered")]).to_csv( + C.OUT / "e08_pairs.csv", index=False) + + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + K = len(C.KEYS) + M = np.full((K, K), np.nan) + for _, r in Rf.iterrows(): + i, j = C.KEYS.index(r.a), C.KEYS.index(r.b) + M[i, j] = r.rho_n0 + M[j, i] = r.rho_gamma + fig, ax = plt.subplots(figsize=(7.6, 6.4)) + im = ax.imshow(M, cmap="RdBu_r", vmin=-1, vmax=1) + for i in range(K): + for j in range(K): + if i != j: + ax.text(j, i, f"{M[i, j]:.2f}", ha="center", va="center", fontsize=9, + color="white" if abs(M[i, j]) > 0.6 else "black") + labs = [C.LABEL[k] for k in C.KEYS] + ax.set_xticks(range(K), labs, rotation=35, ha="right", fontsize=8.5) + ax.set_yticks(range(K), labs, fontsize=8.5) + ax.set_title(f"per-reach Spearman rho between products, {cov.size:,} covered wet reaches\n" + "upper triangle: n0; lower triangle: gamma", fontsize=10) + fig.colorbar(im, ax=ax, shrink=0.8, label="Spearman rho") + fig.tight_layout() + C.FIG.mkdir(parents=True, exist_ok=True) + fig.savefig(C.FIG / "n0_rank_matrix.png", dpi=150, facecolor="white") + print(f"wrote {C.FIG / 'n0_rank_matrix.png'}") + + +if __name__ == "__main__": + main() diff --git a/experiments/cross_product/e10_bias.py b/experiments/cross_product/e10_bias.py new file mode 100755 index 00000000..458a9dcf --- /dev/null +++ b/experiments/cross_product/e10_bias.py @@ -0,0 +1,496 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.11" +# dependencies = ["numpy", "pandas", "xarray", "zarr>=3", "netCDF4", "scipy", "matplotlib"] +# /// +"""E-10: bias decomposition per inflow product, the router's added lag, n_0 vs timing lead, and error by time scale. + +Six fixed-channel models (cp_common.MODELS), test water years 1996-2010, the 2,365 eval gauges. + +Per gauge and product, for the unrouted summed inflow (the run's summed-Q' baseline) and for the routed flow, +each against USGS observations: + volume ratio sum(sim) / sum(obs) on jointly valid days + FHV DDR/ddrs flow-duration peak bias: sort each series, top 2 %, 100 sum(s - o) / sum(o) + (src/training/metrics.rs::Metrics::compute). Timing-free. + event peak bias the section 38.2 variant: sim on the observed top-2 % days (cross_arm_fields.py). Carries timing. + NSE, KGE (2009) + lags best cross-correlation lag (days, -3..10, positive = second series later) as in + experiments/stage_roughness/routing_lag.py::best_lag: summed -> obs (timing lead of the + product), routed -> obs (residual), summed -> routed (what the router adds). +Medians with gauge-bootstrap 95 % CIs, overall and by drainage-area class. + +n_0 vs lead: across products (n = 6, exploratory) and within each product (Spearman over gauges, basin-median +n_0 over the gauge's upstream reaches). + +Error by time scale: each product's routed residual r_k = routed_k - obs split with bands() (water-year mean, +30-day centred MA of the rest, remainder) and into the part common to all six products c = mean_k r_k and the +product-specific part s_k = r_k - c. Per gauge, SS(r_k) = SS(c) + SS(s_k) + 2 exactly, and each part +splits into its three bands plus band cross terms, so the nine shares sum to one per gauge. + + experiments/cross_product/e10_bias.py +Writes /e10_per_gauge.csv, /e10_bands_per_gauge.csv, experiments/cross_product/e10_bias.txt and +research/figures/cross-product/{n0_vs_lead,band_shares}.png. +""" +import io +import sys + +import numpy as np +import pandas as pd +from scipy.stats import spearmanr + +import cp_common as C + +LAGS = np.arange(-3, 11) + + +def best_lag(a, b, min_days=365): + """Vectorised over gauges: lag maximising corr(a(t), b(t + lag)), NaN-safe. + + Same lag range, sign convention, 365-day minimum and tie rule (first lag reaching the maximum wins, + scanning -3..10) as routing_lag.py::best_lag, which loops one gauge at a time. + """ + G, T = a.shape + best = np.zeros(G, dtype=float) + best_r = np.full(G, -np.inf) + rall = np.full((G, LAGS.size), np.nan) + for j, lag in enumerate(LAGS): + if lag >= 0: + x, y = a[:, :T - lag], b[:, lag:] + else: + x, y = a[:, -lag:], b[:, :T + lag] + m = np.isfinite(x) & np.isfinite(y) + n = m.sum(1) + xs, ys = np.where(m, x, 0.0), np.where(m, y, 0.0) + mx, my = xs.sum(1) / np.maximum(n, 1), ys.sum(1) / np.maximum(n, 1) + dx, dy = np.where(m, x - mx[:, None], 0.0), np.where(m, y - my[:, None], 0.0) + with np.errstate(invalid="ignore", divide="ignore"): + r = (dx * dy).sum(1) / np.sqrt((dx ** 2).sum(1) * (dy ** 2).sum(1)) + r = np.where(n >= min_days, r, np.nan) + rall[:, j] = r + upd = np.isfinite(r) & (r > best_r) + best[upd], best_r[upd] = lag, r[upd] + best[~np.isfinite(best_r)] = np.nan + # sub-day refinement (not in routing_lag.py): vertex of the parabola through the peak and its two neighbours + j = np.searchsorted(LAGS, np.nan_to_num(best, nan=LAGS[0])).clip(1, LAGS.size - 2) + rm, r0, rp = rall[np.arange(G), j - 1], rall[np.arange(G), j], rall[np.arange(G), j + 1] + with np.errstate(invalid="ignore", divide="ignore"): + off = 0.5 * (rm - rp) / (rm - 2 * r0 + rp) + interior = (best > LAGS[0]) & (best < LAGS[-1]) + frac = np.where(interior & np.isfinite(off) & (np.abs(off) <= 0.5), LAGS[j] + off, best) + frac[~np.isfinite(best)] = np.nan + return best, best_r, frac + + +def skill(P, O, min_days=365): + """Per-gauge volume ratio, FHV (flow-duration), event peak bias, NSE, KGE on jointly finite days.""" + G = P.shape[0] + out = np.full((G, 5), np.nan) + for i in range(G): + m = np.isfinite(P[i]) & np.isfinite(O[i]) + if m.sum() < min_days: + continue + p, o = P[i, m], O[i, m] + so = np.sum((o - o.mean()) ** 2) + nse = 1 - np.sum((p - o) ** 2) / so if so > 0 else np.nan + r = np.corrcoef(p, o)[0, 1] + alpha, beta = p.std() / o.std(), p.mean() / o.mean() + kge = 1 - np.sqrt((r - 1) ** 2 + (alpha - 1) ** 2 + (beta - 1) ** 2) + ps, os_ = np.sort(p), np.sort(o) + h = int(round(0.98 * p.size)) + fhv = 100 * np.sum(ps[h:] - os_[h:]) / np.sum(os_[h:]) + k = max(1, int(0.02 * o.size)) + top = np.argsort(o)[-k:] + epb = 100 * (p[top].sum() - o[top].sum()) / o[top].sum() + out[i] = (p.sum() / o.sum(), fhv, epb, nse, kge) + return out + + +def ss(x): + return np.nansum(x ** 2, axis=1) + + +def dot(x, y): + return np.nansum(x * y, axis=1) + + +def band_split(r, tt): + """Per gauge: SS of the three bands and twice their summed cross products, for residual r[g, t].""" + lo, mi, hi = C.bands(r, tt) + return ss(lo), ss(mi), ss(hi), 2 * (dot(lo, mi) + dot(lo, hi) + dot(mi, hi)) + + +def main(): + C.OUT.mkdir(parents=True, exist_ok=True) + C.FIG.mkdir(parents=True, exist_ok=True) + buf = io.StringIO() + + def say(*a): + print(*a) + print(*a, file=buf) + + gauges, tt, R, S, O = C.load_series() + area = C.gauge_area(gauges) + say(f"{gauges.size} gauges x {tt.size} days, {tt[0].date()}..{tt[-1].date()}; " + f"observations identical across the six runs (eval zarr and baseline)") + say(f"gauges with drainage area: {np.isfinite(area).sum()}; classes: " + + ", ".join(f"{n} km2 {np.sum((area >= lo) & (area < hi))}" for n, lo, hi in C.AREA_CLASSES)) + Olag = np.where(O <= 0, np.nan, O) # routing_lag.py treats obs <= 0 as missing + + # basin-median n_0 and gamma per gauge and product + comids = C.basin_comids(gauges) + rows = [] + for k in C.KEYS: + f = C.load_field(k) + sk_s, sk_r = skill(S[k], O), skill(R[k], O) + l_so, r_so, f_so = best_lag(S[k], Olag) + l_ro, r_ro, f_ro = best_lag(R[k], Olag) + l_sr, _, f_sr = best_lag(S[k], R[k]) + for i, g in enumerate(gauges): + c = comids.get(g) + sub = f.reindex(c).dropna() if c is not None else f.iloc[:0] + rows.append(dict(model=k, gauge=g, area=area[i], + vol_sum=sk_s[i, 0], fhv_sum=sk_s[i, 1], epb_sum=sk_s[i, 2], nse_sum=sk_s[i, 3], kge_sum=sk_s[i, 4], + vol_routed=sk_r[i, 0], fhv_routed=sk_r[i, 1], epb_routed=sk_r[i, 2], nse_routed=sk_r[i, 3], + kge_routed=sk_r[i, 4], lead=l_so[i], r_sum_obs=r_so[i], lag_routed_obs=l_ro[i], + r_routed_obs=r_ro[i], lag_router=l_sr[i], lead_frac=f_so[i], lag_router_frac=f_sr[i], + lag_routed_obs_frac=f_ro[i], + n0_basin=sub.n0.median() if len(sub) >= 3 else np.nan, + gamma_basin=sub.gamma.median() if len(sub) >= 3 else np.nan, n_reach=len(sub))) + df = pd.DataFrame(rows) + df.to_csv(C.OUT / "e10_per_gauge.csv", index=False) + + # ---- tables: unrouted and routed skill / bias, by area class + metrics = [("vol", "volume ratio", 3), ("fhv", "FHV %", 1), ("epb", "event peak bias %", 1), + ("nse", "NSE", 3), ("kge", "KGE", 3)] + classes = [("all", 0, np.inf)] + C.AREA_CLASSES + say("\n== Unrouted summed inflow vs observed, and routed vs observed: median [gauge-bootstrap 95 % CI]") + for col, name, nd in metrics: + say(f"\n-- {name}") + say(f"{'product':>20} {'area km2':>13} {'n':>5} {'unrouted':>24} {'routed':>24} {'routed - unrouted (paired)':>26}") + for k in C.KEYS: + d = df[df.model == k] + for cname, lo, hi in classes: + dd = d[(d.area >= lo) & (d.area < hi)] if cname != "all" else d + a, b = dd[f"{col}_sum"].to_numpy(), dd[f"{col}_routed"].to_numpy() + say(f"{C.LABEL[k]:>20} {cname:>13} {np.isfinite(b).sum():>5} {C.fmt_ci(C.boot_median(a), nd):>24} " + f"{C.fmt_ci(C.boot_median(b), nd):>24} {C.fmt_ci(C.boot_median(b - a), nd):>26}") + + say("\n== Lags (days, positive = later): timing lead of the summed inflow (summed -> obs), lag the router adds " + "(summed -> routed), residual (routed -> obs). median [CI] / mean [CI]") + lagcols = [("lead", "summed->obs"), ("lag_router", "summed->routed"), ("lag_routed_obs", "routed->obs")] + say(f"{'product':>20} {'area km2':>13} {'n':>5} " + " ".join(f"{n:>34}" for _, n in lagcols) + " router = lead") + for k in C.KEYS: + d = df[df.model == k] + for cname, lo, hi in classes: + dd = d[(d.area >= lo) & (d.area < hi)] if cname != "all" else d + cells = [] + for c, _ in lagcols: + x = dd[c].to_numpy() + med, mean = C.boot_median(x), C.boot_median(x, stat=np.nanmean) + cells.append(f"{med[0]:.0f} [{med[1]:.0f},{med[2]:.0f}] / {C.fmt_ci(mean, 2)}") + ok = np.isfinite(dd.lead) & np.isfinite(dd.lag_router) + say(f"{C.LABEL[k]:>20} {cname:>13} {ok.sum():>5} " + " ".join(f"{c:>34}" for c in cells) + + f" {100 * np.mean(dd.lead[ok] == dd.lag_router[ok]):5.1f}%") + say("\n-- the same lags refined to sub-day by a parabola through the correlation peak (not in routing_lag.py): " + "median [CI]") + lagf = [("lead_frac", "summed->obs"), ("lag_router_frac", "summed->routed"), ("lag_routed_obs_frac", "routed->obs")] + say(f"{'product':>20} {'area km2':>13} {'n':>5} " + " ".join(f"{n:>22}" for _, n in lagf)) + for k in C.KEYS: + d = df[df.model == k] + for cname, lo, hi in classes: + dd = d[(d.area >= lo) & (d.area < hi)] if cname != "all" else d + say(f"{C.LABEL[k]:>20} {cname:>13} {np.isfinite(dd.lead_frac).sum():>5} " + + " ".join(f"{C.fmt_ci(C.boot_median(dd[c].to_numpy()), 2):>22}" for c, _ in lagf)) + + SFS = {} + # ---- volume bookkeeping: the routed run scales the gauge reach's own inflow by FLOW_SCALE, the baseline does not + say("\n== Volume: how much of routed/unrouted is the FLOW_SCALE bookkeeping (gauge_reach_inflow.py), and what is " + "left for the router. median [CI]") + say(f"{'product':>20} {'area km2':>13} {'n':>5} {'routed / unrouted':>24} {'unrouted(FLOW_SCALE) / unrouted':>32} " + f"{'routed / unrouted(FLOW_SCALE)':>30} {'vol ratio, unrouted(FS)':>24}") + for k in C.KEYS: + z = np.load(C.OUT / f"gauge_reach_inflow_{k}.npz", allow_pickle=True) + gi = pd.Index(z["gauge"]).get_indexer(gauges) + pos = pd.DatetimeIndex(z["days"]).get_indexer(tt) + assert (gi >= 0).all() and (pos >= 0).all() + Sfs = S[k] - (1 - z["flow_scale"][gi, None]) * z["q"][gi][:, pos] + SFS[k] = Sfs + m = np.isfinite(O) & np.isfinite(R[k]) & np.isfinite(S[k]) + vr = np.where(m, R[k], 0).sum(1) / np.where(m, S[k], 0).sum(1) + vf = np.where(m, Sfs, 0).sum(1) / np.where(m, S[k], 0).sum(1) + vrf = np.where(m, R[k], 0).sum(1) / np.where(m, Sfs, 0).sum(1) + vo = np.where(m, Sfs, 0).sum(1) / np.where(m, O, 0).sum(1) + enough = m.sum(1) >= 365 + df.loc[df.model == k, "vol_sum_fs"] = np.where(enough, vo, np.nan) + df.loc[df.model == k, "router_vol"] = np.where(enough, vrf, np.nan) + for cname, lo, hi in classes: + s = enough & (((area >= lo) & (area < hi)) if cname != "all" else True) + say(f"{C.LABEL[k]:>20} {cname:>13} {s.sum():>5} {C.fmt_ci(C.boot_median(vr[s])):>24} " + f"{C.fmt_ci(C.boot_median(vf[s])):>32} {C.fmt_ci(C.boot_median(vrf[s])):>30} {C.fmt_ci(C.boot_median(vo[s])):>24}") + df.to_csv(C.OUT / "e10_per_gauge.csv", index=False) + say("\nDoes the router absorb volume bias? If it did, its volume change (routed / like-for-like unrouted) would fall " + "where the inflow over-predicts. Spearman over gauges:") + say(f"{'product':>20} {'n':>5} {'rho(router vol, inflow vol ratio)':>34} {'rho(router vol, basin n0)':>26} " + f"{'rho(router vol, area)':>22} {'median router vol':>18} {'median inflow vol ratio (FS)':>29}") + for k in C.KEYS: + d = df[(df.model == k) & np.isfinite(df.router_vol) & np.isfinite(df.vol_sum_fs) & np.isfinite(df.n0_basin)] + say(f"{C.LABEL[k]:>20} {len(d):5d} {spearmanr(d.router_vol, d.vol_sum_fs).statistic:+34.3f} " + f"{spearmanr(d.router_vol, d.n0_basin).statistic:+26.3f} {spearmanr(d.router_vol, d.area).statistic:+22.3f} " + f"{d.router_vol.median():18.3f} {d.vol_sum_fs.median():29.3f}") + pv = df.groupby("model")[["router_vol", "vol_sum_fs", "n0_basin"]].median() + say(f" across products (n = 6): rho(median router vol, median inflow vol ratio) " + f"{spearmanr(pv.router_vol, pv.vol_sum_fs).statistic:+.3f}; rho(median router vol, median basin n0) " + f"{spearmanr(pv.router_vol, pv.n0_basin).statistic:+.3f}") + + # ---- population check on the summed-Q' baseline medians quoted in findings section 38.1 + say("\n== Baseline population: the manifest's precomputed summed-Q' medians over all baseline gauges vs the eval set") + import json + for k in C.KEYS: + man = json.loads((C.RUNS / C.RUN[k] / "baseline" / "manifest.json").read_text()) + ids = np.array(man["gage_ids"]) + nse = np.array([np.nan if v is None else v for v in man["metrics"]["nse"]], float) + kge = np.array([np.nan if v is None else v for v in man["metrics"]["kge"]], float) + sel = np.isin(ids, gauges) + say(f"{C.LABEL[k]:>20}: all {ids.size} gauges NSE {np.nanmedian(nse):.3f} KGE {np.nanmedian(kge):.3f} " + f"eval {sel.sum()} gauges NSE {np.nanmedian(nse[sel]):.3f} KGE {np.nanmedian(kge[sel]):.3f}") + + # ---- n_0 vs timing lead + say("\n== n_0 vs timing lead") + prod = [] + for k in C.KEYS: + d = df[df.model == k] + ok = np.isfinite(d.n0_basin) & np.isfinite(d.lead) + prod.append(dict(model=k, n0_med=d.n0_basin.median(), lead_med=d.lead.median(), lead_mean=d.lead.mean(), + router_mean=d.lag_router.mean(), vol_med=d.vol_sum.median(), fhv_med=d.fhv_sum.median(), epb_med=d.epb_sum.median(), + n0_net=C.load_field(k).n0.median(), n=int(ok.sum()), + rho_lead=spearmanr(d.n0_basin[ok], d.lead[ok]).statistic, + lead_frac_med=d.lead_frac.median(), lead_frac_mean=d.lead_frac.mean(), + rho_lead_frac=spearmanr(d.n0_basin[ok], d.lead_frac[ok], nan_policy="omit").statistic, + rho_router=spearmanr(d.n0_basin[ok & np.isfinite(d.lag_router)], d.lag_router[ok & np.isfinite(d.lag_router)]).statistic, + rho_vol=spearmanr(d.n0_basin[ok], d.vol_sum[ok], nan_policy="omit").statistic, + rho_fhv=spearmanr(d.n0_basin[ok], d.fhv_sum[ok], nan_policy="omit").statistic, + rho_epb=spearmanr(d.n0_basin[ok], d.epb_sum[ok], nan_policy="omit").statistic)) + P = pd.DataFrame(prod).set_index("model") + P["abs_logvol"] = np.abs(np.log(P.vol_med)) + say("Across products (n = 6, EXPLORATORY: six points, Spearman needs |rho| >= 0.89 for p < 0.05):") + say(f"{'product':>20} {'basin n0 med':>12} {'network n0 med':>14} {'lead median':>11} {'lead mean':>9} " + f"{'router mean':>11} {'vol ratio med':>13} {'FHV med %':>9} {'event pk %':>9}") + for k in C.KEYS: + p = P.loc[k] + say(f"{C.LABEL[k]:>20} {p.n0_med:12.3f} {p.n0_net:14.3f} {p.lead_med:11.0f} {p.lead_mean:9.2f} " + f"{p.router_mean:11.2f} {p.vol_med:13.3f} {p.fhv_med:9.1f} {p.epb_med:9.1f}") + say(f" sub-day lead (median / mean): " + ", ".join(f"{C.LABEL[k]} {P.loc[k].lead_frac_med:.2f} / {P.loc[k].lead_frac_mean:.2f}" for k in C.KEYS)) + for xc, nm in (("lead_frac_med", "median sub-day lead"), ("lead_frac_mean", "mean sub-day lead"), ("lead_mean", "mean lead"), ("router_mean", "mean router lag"), ("vol_med", "median volume ratio"), + ("fhv_med", "median FHV"), ("epb_med", "median event peak bias"), + ("abs_logvol", "|log median volume ratio|")): + say(f" rho(basin-median n0, {nm}) across products = {spearmanr(P.n0_med, P[xc]).statistic:+.3f} " + f"Pearson {np.corrcoef(P.n0_med, P[xc])[0, 1]:+.3f}") + say(f" with network-median n0 instead: rho(n0, mean lead) {spearmanr(P.n0_net, P.lead_mean).statistic:+.3f}, " + f"rho(n0, median sub-day lead) {spearmanr(P.n0_net, P.lead_frac_med).statistic:+.3f}, " + f"rho(n0, mean sub-day lead) {spearmanr(P.n0_net, P.lead_frac_mean).statistic:+.3f}, " + f"rho(n0, median volume ratio) {spearmanr(P.n0_net, P.vol_med).statistic:+.3f}") + say(f" the confound: rho(mean lead, median volume ratio) across products {spearmanr(P.lead_mean, P.vol_med).statistic:+.3f}, " + f"rho(median sub-day lead, median volume ratio) {spearmanr(P.lead_frac_med, P.vol_med).statistic:+.3f}, " + f"rho(mean sub-day lead, median volume ratio) {spearmanr(P.lead_frac_mean, P.vol_med).statistic:+.3f}") + say("Within each product (Spearman over gauges of basin-median n0 against the product's own per-gauge error):") + say(f"{'product':>20} {'n':>5} {'rho(n0, lead)':>13} {'rho(n0, router lag)':>19} {'rho(n0, vol)':>12} " + f"{'rho(n0, FHV)':>12} {'rho(n0, event peak)':>19} {'rho(n0, sub-day lead)':>21}") + for k in C.KEYS: + p = P.loc[k] + say(f"{C.LABEL[k]:>20} {int(p.n):5d} {p.rho_lead:+13.3f} {p.rho_router:+19.3f} {p.rho_vol:+12.3f} {p.rho_fhv:+12.3f} {p.rho_epb:+19.3f} {p.rho_lead_frac:+21.3f}") + ok = np.isfinite(df.n0_basin) & np.isfinite(df.lead) & np.isfinite(df.fhv_sum) + say(f" pooled over products and gauges (n = {ok.sum()}): rho(n0, lead) {spearmanr(df.n0_basin[ok], df.lead[ok]).statistic:+.3f} " + f"rho(n0, FHV) {spearmanr(df.n0_basin[ok], df.fhv_sum[ok]).statistic:+.3f} " + f"rho(n0, event peak) {spearmanr(df.n0_basin[ok], df.epb_sum[ok], nan_policy='omit').statistic:+.3f} " + f"rho(n0, vol) {spearmanr(df.n0_basin[ok], df.vol_sum[ok], nan_policy='omit').statistic:+.3f}") + # demeaned within product: does a gauge that leads more than its product's typical gauge get a rougher basin? + dm = df[ok].copy() + for c in ("n0_basin", "lead"): + dm[c + "_dm"] = dm[c] - dm.groupby("model")[c].transform("median") + say(f" pooled after removing each product's median (within-product only): rho(n0, lead) " + f"{spearmanr(dm.n0_basin_dm, dm.lead_dm).statistic:+.3f}") + say(f" pooled sub-day: rho(n0, sub-day lead) {spearmanr(df.n0_basin, df.lead_frac, nan_policy='omit').statistic:+.3f}") + + # ---- error by time scale + say("\n== Error by time scale: routed residual r = routed - obs, bands = water-year mean / 30-day MA / remainder") + Rr = {k: np.where(np.isfinite(O), R[k] - O, np.nan) for k in C.KEYS} + Sr = {k: np.where(np.isfinite(O), SFS[k] - O, np.nan) for k in C.KEYS} # like-for-like unrouted (FLOW_SCALE applied) + nvalid = np.isfinite(O).sum(1) + keep = nvalid >= 365 + common = np.mean([Rr[k] for k in C.KEYS], axis=0) + cb = band_split(common, tt) + ssc = ss(common) + frames = [] + for k in C.KEYS: + spec = Rr[k] - common + sb = band_split(spec, tt) + rb = band_split(Rr[k], tt) + ub = band_split(Sr[k], tt) + tot, tot_u, sss = ss(Rr[k]), ss(Sr[k]), ss(spec) + cross_cs = 2 * dot(common, spec) + frames.append(pd.DataFrame(dict( + model=k, gauge=gauges, area=area, ss_routed=tot, ss_unrouted=tot_u, + r_low=rb[0] / tot, r_mid=rb[1] / tot, r_high=rb[2] / tot, r_x=rb[3] / tot, + u_low=ub[0] / tot, u_mid=ub[1] / tot, u_high=ub[2] / tot, + c_low=cb[0] / tot, c_mid=cb[1] / tot, c_high=cb[2] / tot, c_x=cb[3] / tot, + s_low=sb[0] / tot, s_mid=sb[1] / tot, s_high=sb[2] / tot, s_x=sb[3] / tot, cs_x=cross_cs / tot, + cl_of_c=cb[0] / ssc, cm_of_c=cb[1] / ssc, ch_of_c=cb[2] / ssc, + sl_of_s=sb[0] / sss, sm_of_s=sb[1] / sss, sh_of_s=sb[2] / sss))[keep]) + B = pd.concat(frames, ignore_index=True) + ident = B[["c_low", "c_mid", "c_high", "c_x", "s_low", "s_mid", "s_high", "s_x", "cs_x"]].sum(1) + assert np.allclose(ident, 1.0, atol=1e-6), f"identity failed: {ident.min()} {ident.max()}" + say(f"gauges with >= 365 observed days: {keep.sum()}; per-gauge identity SS(r_k) = sum of nine parts holds to 1e-6") + B.to_csv(C.OUT / "e10_bands_per_gauge.csv", index=False) + + say("\n-- Band shares of each product's own routed MSE (median [CI] over gauges); band cross terms make the " + "medians not sum to 1, the band-sum column is median (low+mid+high)") + say(f"{'product':>20} {'n':>5} {'water-year mean':>22} {'30-day MA':>22} {'event (< ~30 d)':>22} {'band sum':>8}") + for k in C.KEYS: + d = B[B.model == k] + say(f"{C.LABEL[k]:>20} {len(d):5d} {C.fmt_ci(C.boot_median(d.r_low)):>22} {C.fmt_ci(C.boot_median(d.r_mid)):>22} " + f"{C.fmt_ci(C.boot_median(d.r_high)):>22} {np.median(d.r_low + d.r_mid + d.r_high):8.3f}") + say("\n-- What routing changed, per band: median over gauges of SS_band(unrouted residual) / SS_band(routed residual); " + "unrouted = summed inflow with the gauge reach scaled by FLOW_SCALE, so only the router differs") + say(f"{'product':>20} {'WY mean unrouted/routed':>24} {'30-d unrouted/routed':>21} {'event unrouted/routed':>22} " + f"{'MSE unrouted/routed':>20}") + for k in C.KEYS: + d = B[B.model == k] + say(f"{C.LABEL[k]:>20} {C.fmt_ci(C.boot_median(d.u_low / d.r_low)):>24} {C.fmt_ci(C.boot_median(d.u_mid / d.r_mid)):>21} " + f"{C.fmt_ci(C.boot_median(d.u_high / d.r_high)):>22} {C.fmt_ci(C.boot_median(d.ss_unrouted / d.ss_routed)):>20}") + say(" share of gauges where routing reduced the band's SS (WY mean / 30-d / event): " + "; ".join( + f"{C.LABEL[k]} {np.mean(B[B.model == k].u_low > B[B.model == k].r_low):.2f} / {np.mean(B[B.model == k].u_mid > B[B.model == k].r_mid):.2f}" + f" / {np.mean(B[B.model == k].u_high > B[B.model == k].r_high):.2f}" for k in C.KEYS)) + + # exact average over the six products: mean_k SS(r_k) = SS(c) + mean_k SS(s_k), no cross term + say("\n-- The average product (exact, no cross term): mean_k SS(r_k) = SS(common) + mean_k SS(specific_k). " + "Shares of mean_k SS(r_k), median [CI] over gauges") + tot_avg = np.mean([ss(Rr[k]) for k in C.KEYS], axis=0) + spec_b = [band_split(Rr[k] - common, tt) for k in C.KEYS] + spec_avg = [np.mean([b[j] for b in spec_b], axis=0) for j in range(4)] + A = pd.DataFrame(dict(gauge=gauges, area=area, common=ssc / tot_avg, specific=np.mean([ss(Rr[k] - common) for k in C.KEYS], axis=0) / tot_avg, + c_low=cb[0] / tot_avg, c_mid=cb[1] / tot_avg, c_high=cb[2] / tot_avg, c_x=cb[3] / tot_avg, + s_low=spec_avg[0] / tot_avg, s_mid=spec_avg[1] / tot_avg, s_high=spec_avg[2] / tot_avg, + s_x=spec_avg[3] / tot_avg))[keep] + assert np.allclose(A.common + A.specific, 1.0, atol=1e-6) + say(f" common {C.fmt_ci(C.boot_median(A.common))} specific {C.fmt_ci(C.boot_median(A.specific))} (n = {len(A)})") + say(f" common: WY mean {C.fmt_ci(C.boot_median(A.c_low))} 30-d {C.fmt_ci(C.boot_median(A.c_mid))} event {C.fmt_ci(C.boot_median(A.c_high))}") + say(f" specific: WY mean {C.fmt_ci(C.boot_median(A.s_low))} 30-d {C.fmt_ci(C.boot_median(A.s_mid))} event {C.fmt_ci(C.boot_median(A.s_high))}") + say(" means over gauges (additive): " + " ".join(f"{c} {A[c].mean():.3f}" for c in ("c_low", "c_mid", "c_high", "c_x", "s_low", "s_mid", "s_high", "s_x"))) + for cname, lo, hi in C.AREA_CLASSES: + aa = A[(A.area >= lo) & (A.area < hi)] + say(f" {cname:>13} km2 (n = {len(aa)}): common {np.median(aa.common):.3f}, specific {np.median(aa.specific):.3f}, " + f"specific event {np.median(aa.s_high):.3f}, common event {np.median(aa.c_high):.3f}") + A.to_csv(C.OUT / "e10_bands_average_product.csv", index=False) + + say("\n-- Common vs product-specific: share of each product's routed MSE (mean over gauges, additive; the nine " + "columns sum to 1), then medians") + cols = ["c_low", "c_mid", "c_high", "c_x", "s_low", "s_mid", "s_high", "s_x", "cs_x"] + say(f"{'product':>20} " + " ".join(f"{c:>7}" for c in cols) + f" {'common':>7} {'specific':>8} {'2':>7}") + for k in C.KEYS: + d = B[B.model == k] + mu = d[cols].mean() + say(f"{C.LABEL[k]:>20} " + " ".join(f"{mu[c]:7.3f}" for c in cols) + + f" {mu[['c_low', 'c_mid', 'c_high', 'c_x']].sum():7.3f} {mu[['s_low', 's_mid', 's_high', 's_x']].sum():8.3f} {mu.cs_x:7.3f}") + mu = B[cols].mean() + say(f"{'mean of six':>20} " + " ".join(f"{mu[c]:7.3f}" for c in cols)) + say("medians [CI] of the headline parts, share of each product's routed MSE:") + say(f"{'product':>20} {'common, all bands':>22} {'specific, all bands':>22} {'specific event':>22} {'common event':>22}") + for k in C.KEYS: + d = B[B.model == k] + cs = d[["c_low", "c_mid", "c_high", "c_x"]].sum(1) + sp = d[["s_low", "s_mid", "s_high", "s_x"]].sum(1) + say(f"{C.LABEL[k]:>20} {C.fmt_ci(C.boot_median(cs)):>22} {C.fmt_ci(C.boot_median(sp)):>22} " + f"{C.fmt_ci(C.boot_median(d.s_high)):>22} {C.fmt_ci(C.boot_median(d.c_high)):>22}") + say("\n-- Band composition within each part (median over gauges of SS_band / SS_part)") + d = B[B.model == C.KEYS[0]] + say(f"{'common part':>20} WY mean {C.fmt_ci(C.boot_median(d.cl_of_c))} 30-d {C.fmt_ci(C.boot_median(d.cm_of_c))} " + f"event {C.fmt_ci(C.boot_median(d.ch_of_c))} (n = {len(d)})") + for k in C.KEYS: + d = B[B.model == k] + say(f"{'specific, ' + C.LABEL[k]:>20} WY mean {C.fmt_ci(C.boot_median(d.sl_of_s.astype(float)))} " + f"30-d {C.fmt_ci(C.boot_median(d.sm_of_s.astype(float)))} event {C.fmt_ci(C.boot_median(d.sh_of_s.astype(float)))}") + say("\n-- by area class: median share of routed MSE that is product-specific event-scale / common (all bands)") + say(f"{'product':>20} " + " ".join(f"{c[0]:>24}" for c in C.AREA_CLASSES)) + for k in C.KEYS: + d = B[B.model == k] + cells = [] + for _, lo, hi in C.AREA_CLASSES: + dd = d[(d.area >= lo) & (d.area < hi)] + cells.append(f"{np.median(dd.s_high):.2f} / {np.median(dd[['c_low', 'c_mid', 'c_high', 'c_x']].sum(1)):.2f} (n={len(dd)})") + say(f"{C.LABEL[k]:>20} " + " ".join(f"{c:>24}" for c in cells)) + + (C.HERE / "e10_bias.txt").write_text(buf.getvalue()) + figures(df, P, B, A) + + +def figures(df, P, B, A): + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + colors = dict(zip(C.KEYS, ["#1f77b4", "#d62728", "#2ca02c", "#9467bd", "#ff7f0e", "#8c564b"])) + # Figure 1: n_0 vs timing lead (sub-day lead) + fig, ax = plt.subplots(1, 2, figsize=(13.5, 5.4), gridspec_kw=dict(width_ratios=[1.35, 1])) + edges = np.arange(-1.25, 4.26, 0.5) + centres = 0.5 * (edges[1:] + edges[:-1]) + for k in C.KEYS: + d = df[(df.model == k) & np.isfinite(df.lead_frac) & np.isfinite(df.n0_basin)] + ax[0].scatter(d.lead_frac, d.n0_basin, s=3, alpha=0.12, color=colors[k], rasterized=True) + b = np.digitize(d.lead_frac, edges) - 1 + inb = (b >= 0) & (b < centres.size) + g = d[inb].groupby(b[inb]).n0_basin + med, cnt = g.median(), g.size() + med = med[cnt >= 25].sort_index() + ax[0].plot(centres[med.index], med.values, "-o", ms=3, color=colors[k], lw=1.8, + label=f"{C.LABEL[k]} (within-product rho {P.loc[k].rho_lead_frac:+.2f})") + ax[0].set_xlabel("timing lead of the unrouted summed inflow over observed (days, sub-day cross-correlation lag)") + ax[0].set_ylabel("basin-median learned n0") + ax[0].set_xlim(-1.5, 4.5) + ax[0].set_title("within each product: gauges (points), median n0 per half-day bin with >= 25 gauges (lines)", fontsize=10) + ax[0].legend(fontsize=7.5, loc="upper right") + ax[0].grid(alpha=0.3) + offs = {"uh": (8, -12), "dhbv2_lumped": (-60, 8), "dhbv2_dist": (8, -6), "hydrodl_lstm": (-40, 10), + "nh_daily_lstm": (8, 6), "nh_hourly_mts": (-30, 12)} + for k in C.KEYS: + p = P.loc[k] + ax[1].scatter(p.lead_frac_mean, p.n0_med, s=90, color=colors[k], zorder=3, edgecolor="k") + ax[1].annotate(C.LABEL[k], (p.lead_frac_mean, p.n0_med), textcoords="offset points", xytext=offs[k], fontsize=8) + rho = spearmanr(P.n0_med, P.lead_frac_mean).statistic + ax[1].set_xlabel("mean sub-day timing lead over gauges (days)") + ax[1].set_ylabel("median basin n0 over gauges") + ax[1].set_title(f"across products, n = 6 (exploratory): Spearman {rho:+.2f}", fontsize=10) + ax[1].grid(alpha=0.3) + fig.tight_layout() + fig.savefig(C.FIG / "n0_vs_lead.png", dpi=150, facecolor="white") + plt.close(fig) + + # Figure 2: per product, the common and the product-specific part of the routed MSE, each split by time scale + bands_ = [("low", "water-year mean"), ("mid", "30-day"), ("high", "event (< ~30 d)"), ("x", "band cross terms")] + shade = {"c": ["#08306b", "#2171b5", "#6baed6", "#c6dbef"], "s": ["#7f2704", "#d94801", "#fd8d3c", "#fdd0a2"]} + groups = [(C.LABEL[k], B[B.model == k]) for k in C.KEYS] + [("average product\n(exact, no cross term)", A)] + fig, ax = plt.subplots(figsize=(13, 5.6)) + w = 0.38 + for gi_, (name, d) in enumerate(groups): + for part, dx in (("c", -w / 2), ("s", w / 2)): + bottom = 0.0 + for bi, (b, bname) in enumerate(bands_): + v = d[f"{part}_{b}"].mean() + lab = (f"{'common' if part == 'c' else 'product-specific'}: {bname}") if gi_ == 0 else None + ax.bar(gi_ + dx, v, w, bottom=bottom, color=shade[part][bi], edgecolor="white", lw=0.5, label=lab) + bottom += v + if "cs_x" in d: + ax.plot(gi_, d.cs_x.mean(), "kD", ms=6, label="2 (cross term)" if gi_ == 0 else None) + ax.axhline(0, color="k", lw=0.8) + ax.set_xticks(range(len(groups)), [g[0] for g in groups], fontsize=8.5) + ax.set_ylabel("share of the routed MSE (mean over gauges)") + ax.set_title("left bar: error common to all six products (mean residual); right bar: product-specific remainder", fontsize=10) + ax.legend(fontsize=7.5, bbox_to_anchor=(1.01, 1), loc="upper left") + ax.grid(axis="y", alpha=0.3) + fig.tight_layout() + fig.savefig(C.FIG / "band_shares.png", dpi=150, facecolor="white") + plt.close(fig) + print(f"wrote {C.FIG / 'n0_vs_lead.png'} and {C.FIG / 'band_shares.png'}") + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/experiments/cross_product/gauge_reach_inflow.py b/experiments/cross_product/gauge_reach_inflow.py new file mode 100755 index 00000000..1ed3d769 --- /dev/null +++ b/experiments/cross_product/gauge_reach_inflow.py @@ -0,0 +1,66 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.11" +# dependencies = ["numpy", "pandas", "xarray", "zarr>=3", "icechunk>=2,<3", "netCDF4", "scipy", "matplotlib"] +# /// +"""Daily lateral inflow of each gauge's own reach, per product, over WY1996-2010. + +Why: the routed prediction scales the gauge reach's lateral inflow by FLOW_SCALE, the fraction of that unit +catchment upstream of the gauge (src/data/collate.rs::build_flow_scale, mirroring DDR readers.py), but the +summed-Q' baseline (src/baseline/summed_q_prime.rs) sums the whole unit catchment. Part of the routed-minus- +unrouted volume difference is therefore this bookkeeping, not routing. With the gauge reach's own inflow q_g, +the like-for-like unrouted series is S - (1 - FLOW_SCALE) q_g. + + experiments/cross_product/gauge_reach_inflow.py +Writes /gauge_reach_inflow_.npz (gauge, comid, flow_scale, days, q[g, t]); a COMID missing from the +product store gets the reader's 0.001 m3/s fill, as in ddrs. +""" +import icechunk +import numpy as np +import pandas as pd +import zarr + +import cp_common as C +from qprime_medians import open_store, window + + +def main(): + meta = pd.read_csv(C.GAGES, dtype={"STAID": str}).set_index("STAID") + ev = C.load_predictions_zarr(C.RUNS / C.RUN["uh"] / "eval" / "predictions.zarr") + gauges = ev.gage_ids.values + comid = meta["COMID"].reindex(gauges).to_numpy(dtype=np.int64) + fs = meta["FLOW_SCALE"].reindex(gauges).to_numpy(dtype=float) + days = pd.date_range(C.TEST0, C.TEST1, freq="D") + for k in C.KEYS: + dst = C.OUT / f"gauge_reach_inflow_{k}.npz" + if dst.exists(): + print(f"{k}: exists") + continue + g = open_store(C.STORE[k]) + q = g["Qr"] + ids = np.asarray(g["divide_id"][:], dtype=np.int64) + pos = pd.Index(ids).get_indexer(comid) + per_day, s0, n_days = window(g) + out = np.full((gauges.size, n_days), 0.001) + cd = q.chunks[0] + have = np.where(pos >= 0)[0] + for a in range(0, ids.size, cd): + sel = have[(pos[have] >= a) & (pos[have] < a + cd)] + if sel.size == 0: + continue + b = min(a + cd, ids.size) + if per_day == 1: + x = q[a:b, s0:s0 + n_days].astype(np.float64) + else: + x = np.empty((b - a, n_days)) + step = q.chunks[1] // 24 + for d in range(0, n_days, step): + e = min(d + step, n_days) + x[:, d:e] = q[a:b, s0 + 24 * d:s0 + 24 * e].astype(np.float64).reshape(b - a, e - d, 24).mean(2) + out[sel] = x[pos[sel] - a] + np.savez(dst, gauge=gauges, comid=comid, flow_scale=fs, days=days.values, q=out) + print(f"{k}: {have.size}/{gauges.size} gauge reaches covered; median FLOW_SCALE {np.nanmedian(fs):.3f}") + + +if __name__ == "__main__": + main() diff --git a/experiments/cross_product/qprime_medians.py b/experiments/cross_product/qprime_medians.py new file mode 100755 index 00000000..8713421b --- /dev/null +++ b/experiments/cross_product/qprime_medians.py @@ -0,0 +1,79 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.11" +# dependencies = ["numpy", "pandas", "xarray", "zarr>=3", "icechunk>=2,<3", "netCDF4", "scipy", "matplotlib"] +# /// +"""Per-divide coverage and median daily lateral inflow Q' of each product over WY1996-2010. + +E-08 keeps reaches covered by every product (the COMID is on the product store's divide_id axis, so ddrs +reads real Q' rather than its 0.001 m3/s fill) whose median daily Q' over the test water years is at least +1e-3 m3/s in every product. Daily stores are read directly; the hourly MTS-LSTM store is averaged to +UTC calendar days first, so all six medians are medians of daily-mean flow. + + experiments/cross_product/qprime_medians.py +Writes /qprime_median_.npz (divide_id, median, nan_frac). +""" +import time + +import icechunk +import numpy as np +import pandas as pd +import zarr + +import cp_common as C + + +def open_store(path): + repo = icechunk.Repository.open(icechunk.local_filesystem_storage(path)) + return zarr.open_group(repo.readonly_session(branch="main").store, mode="r") + + +def window(g): + units = g["time"].attrs["units"] + step, origin = units.split(" since ") + t0 = pd.Timestamp(origin) + n_days = (C.TEST1 - C.TEST0).days + 1 + d0 = (C.TEST0 - t0).days + if step == "hours": + return 24, d0 * 24, n_days + assert step == "days", units + return 1, d0, n_days + + +def main(): + C.OUT.mkdir(parents=True, exist_ok=True) + for k in C.KEYS: + dst = C.OUT / f"qprime_median_{k}.npz" + if dst.exists(): + print(f"{k}: exists, skipping") + continue + tic = time.time() + g = open_store(C.STORE[k]) + q = g["Qr"] + dims = tuple(q.metadata.dimension_names or q.attrs.get("_ARRAY_DIMENSIONS")) + assert dims == ("divide_id", "time"), dims + ids = np.asarray(g["divide_id"][:], dtype=np.int64) + per_day, s0, n_days = window(g) + cd = q.chunks[0] + med = np.empty(ids.size, dtype=np.float64) + nanf = np.empty(ids.size, dtype=np.float64) + tchunk = q.chunks[1] if per_day == 24 else n_days + for a in range(0, ids.size, cd): + b = min(a + cd, ids.size) + if per_day == 1: + daily = q[a:b, s0:s0 + n_days].astype(np.float64) + else: + daily = np.empty((b - a, n_days)) + for d in range(0, n_days, tchunk // 24): + e = min(d + tchunk // 24, n_days) + x = q[a:b, s0 + 24 * d:s0 + 24 * e].astype(np.float64) + daily[:, d:e] = x.reshape(b - a, e - d, 24).mean(2) + nanf[a:b] = np.isnan(daily).mean(1) + med[a:b] = np.nanmedian(daily, axis=1) + np.savez(dst, divide_id=ids, median=med, nan_frac=nanf) + print(f"{k}: {ids.size:,} divides, median of medians {np.nanmedian(med):.4g} m3/s, " + f"{np.mean(med >= 1e-3):.3f} at >= 1e-3, any NaN {np.mean(nanf > 0):.4f}, {time.time() - tic:.0f} s") + + +if __name__ == "__main__": + main() diff --git a/experiments/landscape-blend-smoke/experiment.yaml b/experiments/landscape-blend-smoke/experiment.yaml new file mode 100644 index 00000000..cedb70b0 --- /dev/null +++ b/experiments/landscape-blend-smoke/experiment.yaml @@ -0,0 +1,24 @@ +# SMOKE TEST ONLY (two gauges, not a census). Checks that the landscape reads the blend run's weighted +# inflow list (UH retrospective 0.50, dHBV2 distributed 0.25, NH daily LSTM 0.25) and returns finite +# curvatures on the (n_0, gamma) axes of the three-product blend model. The census itself uses all 2,365 +# gauges (user decision 2026-10-01) and is a separate task. +# Modelled on experiments/landscape-p21c-strat400/experiment.yaml and landscape-n0-gamma/experiment.yaml. +study: landscape +arms: + - { name: blend, run: 2026-10-05T13-32-00Z-train-and-test, checkpoint: epoch_50_mb_9 } +landscape: + axes: [n, gamma, q_spatial] # slot 1 is gamma; q_spatial (fixed at 0.65) occupies slot 2 and is carried, never perturbed + gauges: + source: list + staids: ["01563500", "01567000"] # Juniata River at Mapleton Depot and Newport + period: testing + window_days: 1826 + water_year: 1996 + n_windows: 1 + alpha_max: 1.0986123 + fd_step: 0.025 + newton_iters: 0 + grid: 0 + tolerances: [0.05, 0.10] + max_clamped: 0.05 + series: false diff --git a/experiments/landscape-stage1-kge-all/experiment.yaml b/experiments/landscape-stage1-kge-all/experiment.yaml new file mode 100644 index 00000000..e6d99f5a --- /dev/null +++ b/experiments/landscape-stage1-kge-all/experiment.yaml @@ -0,0 +1,25 @@ +# E-02 Stage 1 census: UH free-channel model (n, p, q learned, gamma = 0, 500 updates, +# batch-NSE trained), run 2026-09-12T03-53-34Z-train-and-test (bitwise twin of +# 2026-09-11T23-24-04Z; lineage: paper/e01-lineage@e8e7d3b research/findings/2026-09-30-paper-model-lineage.md). +# Census objective is KGE (decision D-6, 2026-10-01) on every gauge the evaluator scores +# (2,365; no NSE floor, no flow floor). Axes n, p, q; h = 0.025; WY1996-2000; Newton + eigen. +# grid 0: Newton + Hessian eigendecomposition only; the 25x25 figure grids are a separate stratified bundle. +# Latest directory checkpoint (no `checkpoint:` key). Shard with scripts/landscape_shards.sh. +study: landscape +arms: + - { name: uh-free-channel, run: 2026-09-12T03-53-34Z-train-and-test } +landscape: + axes: [n, p_spatial, q_spatial] + gauges: + source: all + period: testing + window_days: 1826 # five water years + water_year: 1996 # window starts 1995-10-01 + n_windows: 1 + alpha_max: 2.3025851 + fd_step: 0.025 + newton_iters: 12 + tolerances: [0.05, 0.10] + max_clamped: 0.05 + grid: 0 + objective: kge diff --git a/experiments/landscape-stage1-nse-all/experiment.yaml b/experiments/landscape-stage1-nse-all/experiment.yaml new file mode 100644 index 00000000..87b6deb6 --- /dev/null +++ b/experiments/landscape-stage1-nse-all/experiment.yaml @@ -0,0 +1,25 @@ +# E-02 Stage 1 NSE-objective rerun (only the objective differs from landscape-stage1-kge-all): UH free-channel model (n, p, q learned, gamma = 0, 500 updates, +# batch-NSE trained), run 2026-09-12T03-53-34Z-train-and-test (bitwise twin of +# 2026-09-11T23-24-04Z; lineage: paper/e01-lineage@e8e7d3b research/findings/2026-09-30-paper-model-lineage.md). +# Objective nse-batch (the training objective), on every gauge the evaluator scores +# (2,365; no NSE floor, no flow floor). Axes n, p, q; h = 0.025; WY1996-2000; Newton + eigen. +# grid 0: Newton + Hessian eigendecomposition only; the 25x25 figure grids are a separate stratified bundle. +# Latest directory checkpoint (no `checkpoint:` key). Shard with scripts/landscape_shards.sh. +study: landscape +arms: + - { name: uh-free-channel, run: 2026-09-12T03-53-34Z-train-and-test } +landscape: + axes: [n, p_spatial, q_spatial] + gauges: + source: all + period: testing + window_days: 1826 # five water years + water_year: 1996 # window starts 1995-10-01 + n_windows: 1 + alpha_max: 2.3025851 + fd_step: 0.025 + newton_iters: 12 + tolerances: [0.05, 0.10] + max_clamped: 0.05 + grid: 0 + objective: nse-batch diff --git a/experiments/landscape/stage1_channel_geometry.py b/experiments/landscape/stage1_channel_geometry.py new file mode 100644 index 00000000..42ca485d --- /dev/null +++ b/experiments/landscape/stage1_channel_geometry.py @@ -0,0 +1,199 @@ +#!/usr/bin/env python +"""Stage 1 channel-geometry tables and figures from the merged (n, p, q) census. + +Reads /uh-free-channel/gauges/*.nc (one per gauge) and writes +/stage1_per_gauge.csv, /stage1_tables.json, and three figures. +Definitions are stated in research/findings/*-stage1-census-findings.md. + +Usage: + ~/projects/ddr/.venv/bin/python experiments/landscape/stage1_channel_geometry.py [--out DIR] +""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +import numpy as np # noqa: E402 +import pandas as pd # noqa: E402 +import xarray as xr # noqa: E402 + +LN2 = float(np.log(2.0)) +ALPHA_MAX = 2.3025851 +NAMES = ["n", "p", "q"] + + +def load(merged: Path) -> pd.DataFrame: + rows = [] + for f in sorted((merged / "uh-free-channel" / "gauges").glob("*.nc")): + d = xr.open_dataset(f, decode_timedelta=False) + r = dict(staid=d.attrs["staid"], n_reach=int(d.attrs["n_reach"]), + obs_mean=float(d.attrs["obs_mean_q_m3s"]), + loss0=float(d.loss0), loss_star=float(d.loss_star), + nse0=float(d.nse0), kge0=float(d.kge0), nse_star=float(d.nse_star), + kge_star=float(d.kge_star), clamped=float(d.attrs["clamped_frac_star"]), + hit_bound=int(d.attrs["hit_range_bound"]), + grad_fallback=int(d.attrs["used_gradient_fallback"])) + g0 = d.grad0.values.astype(float) + H0 = d.hess0.values.astype(float) + Hs = d.hess_star.values.astype(float) + a = d.alpha_star.values.astype(float) + ev = d.eigval_star.values.astype(float) + V = d.eigvec_star.values.astype(float) + for i, nm in enumerate(NAMES): + r[f"alpha_{nm}"] = a[i] + r[f"g0_{nm}"] = g0[i] + r[f"H0_{nm}"] = H0[i, i] + r[f"Hs_{nm}"] = Hs[i, i] + # loss change for a factor-two move, trained point (quadratic form + # |g| ln2 + 0.5 |H| ln2^2), relative to the loss there + r[f"dL2_{nm}"] = (abs(g0[i]) * LN2 + 0.5 * abs(H0[i, i]) * LN2**2) / r["loss0"] + r[f"edge_{nm}"] = int(abs(a[i]) >= ALPHA_MAX - 1e-3) + for k in range(3): + r[f"lam{k+1}"] = ev[k] + # sign-normalised stiff eigenvector (v1_n >= 0) and flat eigenvector + v1 = V[:, 0] * (1 if V[0, 0] >= 0 else -1) + v3 = V[:, 2] + for i, nm in enumerate(NAMES): + r[f"v1_{nm}"] = v1[i] + r[f"v3_{nm}"] = v3[i] + # curvature along the depth-invariant direction (1,1,0)/sqrt2 and its + # orthogonal partner (1,-1,0)/sqrt2, at the optimum and the trained point + u_same = np.array([1, 1, 0]) / np.sqrt(2) + u_opp = np.array([1, -1, 0]) / np.sqrt(2) + for tag, H in (("s", Hs), ("0", H0)): + r[f"Huu_np_same_{tag}"] = float(u_same @ H @ u_same) + r[f"Huu_np_opp_{tag}"] = float(u_opp @ H @ u_opp) + r["indef_star"] = int(ev.min() <= 0) + r["indef_trained"] = int(np.linalg.eigvalsh(H0).min() <= 0) + r["cos_v1_same"] = abs(float(v1 @ u_same)) + r["cos_v3_same"] = abs(float(v3 @ u_same)) + rows.append(r) + return pd.DataFrame(rows) + + +def bands(df: pd.DataFrame): + # KGE bands on the trained-point KGE over the five-year window + edges = [-np.inf, 0.0, 0.3, 0.6, np.inf] + labels = ["KGE < 0", "0 to 0.3", "0.3 to 0.6", ">= 0.6"] + return pd.cut(df["kge0"], edges, labels=labels, right=False) + + +def summarize(sub: pd.DataFrame) -> dict: + out = {"n_gauges": int(len(sub))} + if len(sub) == 0: + return out + for nm in NAMES: + out[f"flat_{nm}_1pct"] = float((sub[f"dL2_{nm}"] < 0.01).mean()) + out[f"H0_concave_{nm}"] = float((sub[f"H0_{nm}"] < 0).mean()) + out[f"Hs_concave_{nm}"] = float((sub[f"Hs_{nm}"] < 0).mean()) + out[f"edge_{nm}"] = float(sub[f"edge_{nm}"].mean()) + out["flat_q_0.5pct"] = float((sub.dL2_q < 0.005).mean()) + out["flat_q_2pct"] = float((sub.dL2_q < 0.02).mean()) + out["q_flat_and_n_curved"] = float(((sub.dL2_q < 0.01) & (sub.dL2_n >= 0.01)).mean()) + out["n_flat_and_q_curved"] = float(((sub.dL2_n < 0.01) & (sub.dL2_q >= 0.01)).mean()) + for tag in ("0", "s"): + for nm in ("p", "q"): + # ratio of absolute diagonal curvature to n's; medians over gauges + # with a nonzero n curvature + r = (sub[f"H{tag}_{nm}"].abs() / sub[f"H{tag}_n"].abs()).replace([np.inf], np.nan) + out[f"ratio_{nm}_to_n_{'trained' if tag=='0' else 'optimum'}_median"] = float(r.median()) + out[f"ratio_{nm}_to_n_{'trained' if tag=='0' else 'optimum'}_q25_q75"] = [float(r.quantile(.25)), float(r.quantile(.75))] + out["indef_star"] = float(sub.indef_star.mean()) + out["indef_trained"] = float(sub.indef_trained.mean()) + out["lam3_le0_star"] = float((sub.lam3 <= 0).mean()) + out["lam1_le0_star"] = float((sub.lam1 <= 0).mean()) + for nm in NAMES: + out[f"v1_{nm}_median"] = float(sub[f"v1_{nm}"].median()) + out[f"v1_{nm}_q25_q75"] = [float(sub[f"v1_{nm}"].quantile(.25)), float(sub[f"v1_{nm}"].quantile(.75))] + # exponents of the stiff combination n^a p^b q^c scaled so a = 1 and a = 0.8 norm + good = sub[sub.v1_n.abs() > 0.2] + out["stiff_b_over_a_median"] = float((good.v1_p / good.v1_n).median()) + out["stiff_c_over_a_median"] = float((good.v1_q / good.v1_n).median()) + out["stiff_n_dominant_share"] = float((sub[["v1_n", "v1_p", "v1_q"]].abs().idxmax(axis=1) == "v1_n").mean()) + out["cos_v1_with_n_equals_p_median"] = float(sub.cos_v1_same.median()) + out["cos_v3_with_n_equals_p_median"] = float(sub.cos_v3_same.median()) + out["Huu_same_over_opp_star_median"] = float((sub.Huu_np_same_s / sub.Huu_np_opp_s).median()) + out["Huu_same_over_Hnn_star_median"] = float((sub.Huu_np_same_s / sub.Hs_n).median()) + out["Huu_opp_over_Hnn_star_median"] = float((sub.Huu_np_opp_s / sub.Hs_n).median()) + out["hit_bound"] = float(sub.hit_bound.mean()) + out["kge0_median"] = float(sub.kge0.median()) + out["nse0_median"] = float(sub.nse0.median()) + return out + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("merged", type=Path) + ap.add_argument("--out", type=Path, default=None) + a = ap.parse_args() + out = a.out or a.merged / "figures" / "stage1" + out.mkdir(parents=True, exist_ok=True) + df = load(a.merged) + df["band"] = bands(df) + df.to_csv(out / "stage1_per_gauge.csv", index=False) + + tables = {"all": summarize(df), "nse0_gt_0.3": summarize(df[df.nse0 > 0.3])} + for b in df.band.cat.categories: + tables[f"kge0 {b}"] = summarize(df[df.band == b]) + tables["kge0 > 0.3 (well fit, KGE)"] = summarize(df[df.kge0 > 0.3]) + (out / "stage1_tables.json").write_text(json.dumps(tables, indent=1)) + + wf = df[df.nse0 > 0.3] + # --- Figure 1: curvature ratios of p and q to n at the optimum, by KGE band + fig, ax = plt.subplots(1, 3, figsize=(15, 4.4)) + cols = {"KGE < 0": "#b2182b", "0 to 0.3": "#ef8a62", "0.3 to 0.6": "#67a9cf", ">= 0.6": "#2166ac"} + for nm, axx in zip(("p", "q"), ax[:2]): + for b, c in cols.items(): + s = df[df.band == b] + r = (s[f"Hs_{nm}"].abs() / s.Hs_n.abs()).replace([np.inf], np.nan).dropna() + r = np.sort(r.values) + axx.plot(r, np.arange(1, len(r) + 1) / len(r), color=c, label=f"{b} (n={len(s)})") + axx.set_xscale("log"); axx.axvline(1, color="k", lw=.6, ls=":") + axx.set_xlabel(f"|H_{nm}{nm}| / |H_nn| at the optimum"); axx.set_ylabel("cumulative share of gauges") + axx.set_title(f"curvature of {nm} as a fraction of n's") + ax[0].legend(fontsize=8) + # dL2 for q + for b, c in cols.items(): + s = df[df.band == b] + r = np.sort(s.dL2_q.values) + ax[2].plot(r, np.arange(1, len(r) + 1) / len(r), color=c) + ax[2].set_xscale("log"); ax[2].axvline(0.01, color="k", lw=.8, ls="--") + ax[2].set_xlabel("relative loss change for a factor-two move in q (trained point)") + ax[2].set_ylabel("cumulative share of gauges"); ax[2].set_title("how flat is q (dashed: 1 %)") + fig.tight_layout(); fig.savefig(out / "fig1_curvature_ratios.png", dpi=130); plt.close(fig) + + # --- Figure 2: stiff eigenvector components + fig, ax = plt.subplots(1, 2, figsize=(11, 4.6)) + s = df[df.v1_n.abs() > 0.0] + sc = ax[0].scatter(s.v1_p, s.v1_n, c=np.log10(s.lam1.clip(lower=1e-6)), s=7, cmap="viridis", alpha=.6) + ax[0].set_xlabel("stiff eigenvector, p component"); ax[0].set_ylabel("stiff eigenvector, n component") + ax[0].axhline(0, color="k", lw=.4); ax[0].axvline(0, color="k", lw=.4) + fig.colorbar(sc, ax=ax[0], label="log10 stiff eigenvalue") + ax[0].set_title("stiff direction (sign: n component >= 0)") + ax[1].hist(np.degrees(np.arctan2(s.v1_p, s.v1_n)), bins=72, color="#2166ac") + ax[1].set_xlabel("angle of stiff direction in the (n, p) plane (deg; 0 = pure n)") + ax[1].set_ylabel("gauges"); ax[1].set_title("direction of the stiff axis") + fig.tight_layout(); fig.savefig(out / "fig2_stiff_eigenvector.png", dpi=130); plt.close(fig) + + # --- Figure 3: share of gauges whose parameter is curved enough, vs bar + bars = np.logspace(-3, 0, 60) + fig, ax = plt.subplots(figsize=(6.5, 4.4)) + for sub, ls, lab in ((wf, "-", "NSE at trained point > 0.3"), (df, "--", "all gauges")): + for nm, c in (("p", "#1b7837"), ("q", "#762a83")): + ratio = sub[f"Hs_{nm}"] / sub.Hs_n.abs() + share = [float(((sub[f"Hs_{nm}"] > 0) & (ratio >= b)).mean()) for b in bars] + ax.plot(bars, share, color=c, ls=ls, label=f"{nm}, {lab}") + ax.set_xscale("log"); ax.set_xlabel("bar: H_ii / |H_nn| at the optimum, with H_ii > 0") + ax.set_ylabel("share of gauges at or above the bar") + ax.set_title("share of gauges where p and q are curved enough"); ax.legend(fontsize=7) + fig.tight_layout(); fig.savefig(out / "fig3_share_vs_bar.png", dpi=130); plt.close(fig) + print(json.dumps({k: v["n_gauges"] for k, v in tables.items()})) + + +if __name__ == "__main__": + main() diff --git a/experiments/landscape/stage1_objective_compare.py b/experiments/landscape/stage1_objective_compare.py new file mode 100644 index 00000000..65e73e7c --- /dev/null +++ b/experiments/landscape/stage1_objective_compare.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python +"""Side-by-side of the KGE-objective and NSE-objective Stage 1 censuses. + +Usage: stage1_objective_compare.py --out DIR +Reads /figures/stage1/{stage1_per_gauge.csv,stage1_tables.json}. +""" +import argparse, json +from pathlib import Path +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np, pandas as pd + +ap = argparse.ArgumentParser(); ap.add_argument("kge", type=Path); ap.add_argument("nse", type=Path) +ap.add_argument("--out", type=Path, required=True); a = ap.parse_args(); a.out.mkdir(parents=True, exist_ok=True) +T = {k: json.loads((p / "figures/stage1/stage1_tables.json").read_text()) for k, p in (("KGE", a.kge), ("NSE", a.nse))} +D = {k: pd.read_csv(p / "figures/stage1/stage1_per_gauge.csv", dtype={"staid": str}) for k, p in (("KGE", a.kge), ("NSE", a.nse))} +pops = ["nse0_gt_0.3", "all", "kge0 KGE < 0", "kge0 0 to 0.3", "kge0 0.3 to 0.6", "kge0 >= 0.6"] +keys = ["n_gauges", "flat_q_0.5pct", "flat_q_1pct", "flat_q_2pct", "flat_p_1pct", "flat_n_1pct", "H0_concave_n", "H0_concave_p", "H0_concave_q", + "Hs_concave_n", "Hs_concave_p", "Hs_concave_q", "edge_p", "edge_q", "ratio_p_to_n_trained_median", "ratio_q_to_n_trained_median", + "ratio_p_to_n_optimum_median", "ratio_q_to_n_optimum_median", "indef_trained", "indef_star", "lam1_le0_star", + "v1_n_median", "v1_p_median", "v1_q_median", "stiff_b_over_a_median", "stiff_c_over_a_median", "stiff_n_dominant_share", "hit_bound"] +rows = [] +for k in keys: + r = {"quantity": k} + for pop in pops: + for o in ("KGE", "NSE"): + r[f"{pop} | {o}"] = T[o][pop].get(k) + rows.append(r) +pd.DataFrame(rows).to_csv(a.out / "stage1_objective_compare.csv", index=False) +# extra: stiff angle, same-gauge agreement +for o, d in D.items(): + ok = d[(d.lam1 > 0) & (d.nse0 > 0.3)] + print(o, "stiff angle median deg (well fit)", round(float(np.degrees(np.arctan2(ok.v1_p, ok.v1_n)).median()), 1)) +m = D["KGE"].merge(D["NSE"], on="staid", suffixes=("_k", "_n")) +wf = m[m.nse0_k > 0.3] +print("same-gauge: n", len(wf), "flat_q agree", float(((wf.dL2_q_k < .01) == (wf.dL2_q_n < .01)).mean()), + "spearman dL2_q", float(wf.dL2_q_k.corr(wf.dL2_q_n, method="spearman")), + "median loss0 KGE/NSE", float(wf.loss0_k.median()), float(wf.loss0_n.median()), + "alpha_q* same sign", float((np.sign(wf.alpha_q_k) == np.sign(wf.alpha_q_n)).mean()), + "spearman alpha_n*", float(wf.alpha_n_k.corr(wf.alpha_n_n, method="spearman")), + "spearman alpha_q*", float(wf.alpha_q_k.corr(wf.alpha_q_n, method="spearman"))) +# absolute loss change for factor two in q: dL2*loss0 +for o, d in D.items(): + wf2 = d[d.nse0 > 0.3]; ab = wf2.dL2_q * wf2.loss0 + print(o, "abs dL (factor-2 in q) median", float(ab.median()), "share abs<0.001", float((ab < 1e-3).mean()), "share abs<0.0025",float((ab<2.5e-3).mean()), "loss0 median", float(wf2.loss0.median())) +# figure: CDF of dL2_q by objective (well fit) and flat share by KGE band +fig, ax = plt.subplots(1, 3, figsize=(15, 4.4)) +for o, c in (("KGE", "#2166ac"), ("NSE", "#b2182b")): + d = D[o]; wf2 = d[d.nse0 > 0.3] + for nm, ls in (("q", "-"), ("p", "--"), ("n", ":")): + r = np.sort(wf2[f"dL2_{nm}"].values); ax[0].plot(r, np.arange(1, len(r) + 1) / len(r), color=c, ls=ls, label=f"{nm}, {o}") +ax[0].set_xscale("log"); ax[0].axvline(.01, color="k", lw=.7); ax[0].set_xlabel("relative loss change, factor-two move"); ax[0].set_ylabel("cumulative share (NSE>0.3 gauges)"); ax[0].legend(fontsize=7); ax[0].set_title("how flat, by objective") +bl = ["kge0 KGE < 0", "kge0 0 to 0.3", "kge0 0.3 to 0.6", "kge0 >= 0.6", "nse0_gt_0.3"] +x = np.arange(len(bl)) +for i, (o, c) in enumerate((("KGE", "#2166ac"), ("NSE", "#b2182b"))): + ax[1].bar(x + i * .38 - .19, [T[o][b]["flat_q_1pct"] for b in bl], .38, color=c, label=o) +ax[1].axhline(.85, color="k", ls="--", lw=.8); ax[1].text(4.45, .86, "old 85 %", ha="right", fontsize=8) +ax[1].set_xticks(x); ax[1].set_xticklabels(["KGE<0", "0-0.3", "0.3-0.6", ">=0.6", "NSE>0.3"], fontsize=8); ax[1].set_ylabel("share with q flat (1 %)"); ax[1].legend(); ax[1].set_title("q flat share by band") +sc = ax[2].scatter(wf.dL2_q_k, wf.dL2_q_n, s=5, alpha=.4); ax[2].set_xscale("log"); ax[2].set_yscale("log") +lo, hi = 1e-6, 10; ax[2].plot([lo, hi], [lo, hi], "k", lw=.6); ax[2].axvline(.01, color="gray", lw=.5); ax[2].axhline(.01, color="gray", lw=.5) +ax[2].set_xlabel("KGE objective"); ax[2].set_ylabel("NSE objective"); ax[2].set_title("per-gauge q flatness, same gauge") +fig.tight_layout(); fig.savefig(a.out / "fig_objective_compare.png", dpi=130) diff --git a/experiments/prediction_diversity/pair_diversity.py b/experiments/prediction_diversity/pair_diversity.py new file mode 100644 index 00000000..3544c700 --- /dev/null +++ b/experiments/prediction_diversity/pair_diversity.py @@ -0,0 +1,388 @@ +"""Seed pairs vs inflow-product pairs: parameter distance against prediction diversity. + +For every pair of trained fixed-channel models (n_0 + gamma head, p = 21, q = 0.65), on the test years +WY1996-2010 and the gauges common to both: + +prediction diversity + resid_corr median over gauges of corr(p_a - o, p_b - o), the per-gauge correlation of routed residuals + gain NSE of the equal-weight average minus the members' mean NSE, per gauge. For weights 1/2, + sum_i w_i (p_i - o)^2 = (pbar - o)^2 + sum_i w_i (p_i - pbar)^2 at every step, so + gain = sum_t sum_i w_i (p_i - pbar)^2 / sum_t (o - obar)^2 >= 0: the ambiguity (spread) term of + Krogh & Vedelsby (1995) divided by var(obs). The identity is asserted per gauge, as in + experiments/levers/blend_mechanism.py (multi-product-forcing worktree). Median over gauges with a + paired gauge-bootstrap 95 % CI. + band shares the gain split by time scale with blend_mechanism.py's bands(): water-year mean (year-to-year + volume), 30-day centred moving average of the rest (seasonal), remainder (event scale, < ~30 d). +parameter distance + n0_ratio larger / smaller network-median n_0, over the filtered reaches + rho_n0 per-reach Spearman rho of n_0 over the filtered reaches (covered by every product, median + local Q' >= 1e-3 m3/s in every product; reach_filter.py) + +Seed pairs whose seed-43 run has no eval output yet are written as PENDING rows. +Seed crossing: for the products with a seed-43 twin, every product pair is also scored at all four seed +combinations, to test whether a shared seed inflates product-pair pattern agreement (also the basin-median +n_0 rank correlation over gauges). +Writes pairs.csv, pairs.txt and prediction_diversity.png next to this script. + + /home/tbindas/projects/ddr/.venv/bin/python pair_diversity.py +""" +import itertools +import pathlib +import sys + +import numpy as np +import pandas as pd +from netCDF4 import Dataset +from scipy.stats import spearmanr + +HERE = pathlib.Path(__file__).resolve().parent +sys.path.insert(0, str(HERE.parent / "levers")) +import common as C # noqa: E402 (load_preds, align, water_year) + +RUNS = C.RUNS +ATTRS = "/home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc" +GAGES_ADJ = "/home/tbindas/projects/ddr/data/merit_gages_conus_adjacency.zarr" # every run's gages_adjacency +CACHE = pathlib.Path("/home/tbindas/projects/ddrs/.ddrs/experiments/prediction-diversity") +B, SEED = 2000, 1 + +# Descriptive name -> (run id, inflow product, seed). Product runs from the paper's lineage table. +MODELS = { + "UH fixed-channel model": ("2026-09-12T23-39-03Z-train-and-test", "uh", 42), + "dHBV2 lumped model": ("2026-09-13T13-56-50Z-train-and-test", "dhbv2_lumped", 42), + "dHBV2 distributed model": ("2026-09-13T17-22-30Z-train-and-test", "dhbv2_dist", 42), + "hydroDL LSTM model": ("2026-09-13T17-21-58Z-train-and-test", "hydrodl_lstm", 42), + "NH daily LSTM model": ("2026-09-13T13-55-03Z-train-and-test", "nh_daily_lstm", 42), + "NH hourly MTS-LSTM model": ("2026-09-13T14-14-46Z-train-and-test", "nh_hourly_lstm", 42), + # seed pairs: the seed-42 member shares a binary with its seed-43 twin + "UH fixed-channel model, seed 42 (09-27 rerun)": ("2026-09-27T07-29-47Z-train-and-test", "uh", 42), + "UH fixed-channel model, seed 43": ("2026-09-27T10-31-30Z-train-and-test", "uh", 43), +} +SHORT = {"UH fixed-channel model": "UH", "dHBV2 lumped model": "dHBV2 lumped", + "dHBV2 distributed model": "dHBV2 dist", "hydroDL LSTM model": "hydroDL LSTM", + "NH daily LSTM model": "NH daily LSTM", "NH hourly MTS-LSTM model": "NH hourly LSTM", + "UH fixed-channel model, seed 42 (09-27 rerun)": "UH s42", "UH fixed-channel model, seed 43": "UH s43"} +PRODUCT_MODELS = list(MODELS)[:6] +# E-07 (branch paper/e07-seed43) seed-43 replicates of the two dHBV2 models, release build of 70f6457 (the seed-42 +# commit), finished 2026-10-01. +# seed42 = None keeps the 09-13 product run as the seed-42 member; set it if E-07 retrained seed 42 on a new binary. +E07 = { + "dHBV2 distributed model": dict(product="dhbv2_dist", seed42=None, seed43="2026-10-01T02-33-37Z-train-and-test"), + "dHBV2 lumped model": dict(product="dhbv2_lumped", seed42=None, seed43="2026-10-01T02-33-01Z-train-and-test"), +} +SEED_PAIRS = [("UH fixed-channel model, seed 42 (09-27 rerun)", "UH fixed-channel model, seed 43", "UH seeds 42/43")] +for _name, _e in E07.items(): + _short = SHORT[_name] + _a = _name + if _e["seed42"]: + _a = f"{_name}, seed 42 (E-07 rerun)" + MODELS[_a] = (_e["seed42"], _e["product"], 42) + SHORT[_a] = f"{_short} s42" + _b = None + if _e["seed43"]: + _b = f"{_name}, seed 43" + MODELS[_b] = (_e["seed43"], _e["product"], 43) + SHORT[_b] = f"{_short} s43" + SEED_PAIRS.append((_a, _b, f"{_short} seeds 42/43")) +NOT_AORC = {"uh"} # every other product is AORC-forced + + +def bands(d, tt): + """Split d[g, t] into (water-year mean, 30-day MA of the rest, remainder). + + Copied verbatim from experiments/levers/blend_mechanism.py (multi-product-forcing worktree, + /home/tbindas/projects/ddrs/.claude/worktrees/agent-a92e512a7c47c97b4, uncommitted there on 2026-09-30). + """ + wy = C.water_year(tt) + low = np.empty_like(d) + for y in np.unique(wy): + s = wy == y + low[:, s] = d[:, s].mean(1, keepdims=True) + rest = d - low + mid = pd.DataFrame(rest.T).rolling(30, center=True, min_periods=1).mean().to_numpy().T + return low, mid, rest - mid + + +def boot_median(x, rng_seed=SEED): + x = np.asarray(x) + x = x[np.isfinite(x)] + idx = np.random.default_rng(rng_seed).integers(0, x.size, size=(B, x.size)) + m = np.median(x[idx], axis=1) + return float(np.percentile(m, 2.5)), float(np.percentile(m, 97.5)) + + +def load_fields(run): + ds = Dataset(RUNS / run / "plot" / "kan_parameters.nc") + return (np.asarray(ds["COMID"][:], dtype=np.int64), np.asarray(ds["n"][:], dtype=np.float64), + np.asarray(ds["gamma"][:], dtype=np.float64)) + + +def pair_row(a, b, kind, P, O, tt, F, keep): + Pa, Pb = P[a], P[b] + m = np.isfinite(O) & np.isfinite(Pa) & np.isfinite(Pb) + nv = m.sum(1) + om = np.where(m, O, 0.0) + mo = om.sum(1) / np.maximum(nv, 1) + sst = (np.where(m, O - mo[:, None], 0.0) ** 2).sum(1) + pbar = 0.5 * (Pa + Pb) + nse_a = 1 - (np.where(m, Pa - O, 0.0) ** 2).sum(1) / sst + nse_b = 1 - (np.where(m, Pb - O, 0.0) ** 2).sum(1) / sst + nse_avg = 1 - (np.where(m, pbar - O, 0.0) ** 2).sum(1) / sst + spread = 0.5 * ((np.where(m, Pa - pbar, 0.0) ** 2).sum(1) + (np.where(m, Pb - pbar, 0.0) ** 2).sum(1)) / sst + gain = nse_avg - 0.5 * (nse_a + nse_b) + ok = (nv >= 365) & np.isfinite(nse_avg) + assert np.allclose(gain[ok], spread[ok], atol=1e-8), "ambiguity identity failed" + # residual correlation per gauge + ea, eb = np.where(m, Pa - O, 0.0), np.where(m, Pb - O, 0.0) + ma, mb = ea.sum(1) / np.maximum(nv, 1), eb.sum(1) / np.maximum(nv, 1) + da, db = np.where(m, ea - ma[:, None], 0.0), np.where(m, eb - mb[:, None], 0.0) + rc = (da * db).sum(1) / np.sqrt((da ** 2).sum(1) * (db ** 2).sum(1)) + rc[~ok] = np.nan + # spread by time scale (deviation of a from the average; b's is its negative, same energy) + parts = np.stack([(np.where(m, x, 0.0) ** 2).sum(1) / sst for x in bands(Pa - pbar, tt)]) + share = parts / parts.sum(0) + g = gain[ok] + glo, ghi = boot_median(g) + rlo, rhi = boot_median(rc[ok]) + # parameters + (na, ga), (nb, gb) = F[a], F[b] + meda, medb = np.median(na[keep]), np.median(nb[keep]) + med_a, med_b, med_avg = np.median(nse_a[ok]), np.median(nse_b[ok]), np.median(nse_avg[ok]) + pa, pb = MODELS[a][1], MODELS[b][1] + return dict( + kind=kind, a=a, b=b, run_a=MODELS[a][0][:20], run_b=MODELS[b][0][:20], + forcing=("" if kind.startswith("seed") else + "both AORC" if pa not in NOT_AORC and pb not in NOT_AORC else "UH vs AORC"), + n_gauges=int(ok.sum()), + nse_a=med_a, nse_b=med_b, nse_avg=med_avg, avg_minus_better=med_avg - max(med_a, med_b), + resid_corr=float(np.median(rc[ok])), resid_corr_lo=rlo, resid_corr_hi=rhi, + gain=float(np.median(g)), gain_lo=glo, gain_hi=ghi, gain_mean=float(np.mean(g)), + share_year=float(np.median(share[0, ok])), share_season=float(np.median(share[1, ok])), + share_event=float(np.median(share[2, ok])), + n_reaches=int(keep.sum()), n0_med_a=meda, n0_med_b=medb, n0_ratio=max(meda, medb) / min(meda, medb), + rho_n0=float(spearmanr(na[keep], nb[keep]).statistic), + rho_gamma=float(spearmanr(ga[keep], gb[keep]).statistic), + rho_n0_all=float(spearmanr(na, nb).statistic), + n0_ratio_all=max(np.median(na), np.median(nb)) / min(np.median(na), np.median(nb)), + ) + + +def main(): + # predictions, aligned to the UH model's gauge order and time axis + ref_ids, tt, P0, O = C.load_preds(MODELS["UH fixed-channel model"][0]) + P, F = {}, {} + comid0 = None + for name, (run, _, _) in MODELS.items(): + ids, t, Pr, Or = C.load_preds(run) + assert (t == tt).all() and set(ids) == set(ref_ids) + assert np.array_equal(C.align(ids, Or, ref_ids), O, equal_nan=True) + P[name] = C.align(ids, Pr, ref_ids) + c, n0, g = load_fields(run) + comid0 = c if comid0 is None else comid0 + assert np.array_equal(c, comid0) + F[name] = (n0, g) + # the two seed-42 UH runs (different binaries, same src for routing) are the same model + dup_p = float(np.nanmax(np.abs(P["UH fixed-channel model"] - P["UH fixed-channel model, seed 42 (09-27 rerun)"]))) + dup_n = float(np.max(np.abs(F["UH fixed-channel model"][0] - F["UH fixed-channel model, seed 42 (09-27 rerun)"][0]))) + lines = [f"UH 2026-09-12T23-39-03Z vs 2026-09-27T07-29-47Z: max |dQ| {dup_p:.3g} m3/s, max |dn_0| {dup_n:.3g}"] + + fl = np.load(CACHE / "reach_flow_medians.npz") + assert np.array_equal(fl["COMID"], comid0) + prods = ["uh", "dhbv2_lumped", "dhbv2_dist", "hydrodl_lstm", "nh_daily_lstm", "nh_hourly_lstm"] + M = np.stack([fl[k] for k in prods]) + covered = np.isfinite(M).all(0) + keep = covered & (np.nan_to_num(M, nan=0.0) >= 1e-3).all(0) + lines.append(f"reaches: {comid0.size:,} in the network, {int(covered.sum()):,} covered by every product, " + f"{int(keep.sum()):,} also with median local Q' >= 1e-3 m3/s in every product") + + import zarr + gz = zarr.open_group(GAGES_ADJ, mode="r") + pos = {c: i for i, c in enumerate(comid0)} + basin = [np.array([pos[c] for c in np.asarray(gz[s]["order"][:], dtype=np.int64) if c in pos]) for s in ref_ids] + + def basin_rho(a, b): + """Spearman over gauges of the basin-median n_0 (median over each gauge's upstream reaches).""" + bma = np.array([np.median(F[a][0][ix]) for ix in basin]) + bmb = np.array([np.median(F[b][0][ix]) for ix in basin]) + return float(spearmanr(bma, bmb).statistic), bma, bmb + + rows = [] + for a, b, lab in SEED_PAIRS: + if b is None: + rows.append(dict(kind="seed (PENDING)", a=a, b=None, label=lab)) + continue + rows.append(dict(pair_row(a, b, "seed", P, O, tt, F, keep), label=lab)) + for a, b in itertools.combinations(PRODUCT_MODELS, 2): + rows.append(pair_row(a, b, "product", P, O, tt, F, keep)) + for b in PRODUCT_MODELS[1:]: # robustness: the UH member swapped for its seed-43 twin + rows.append(pair_row("UH fixed-channel model, seed 43", b, "product (UH seed 43)", P, O, tt, F, keep)) + # seed crossing: the three products with a seed-43 twin, every pair of products at every seed combination. + # If a shared seed inflates the product-pair pattern agreement, same-seed rows (42/42, 43/43) beat + # cross-seed rows (42/43, 43/42) for every product pair. + twins = {"UH": ("UH fixed-channel model", "UH fixed-channel model, seed 43")} + for _name in E07: + if E07[_name]["seed43"]: + twins[SHORT[_name]] = (_name if not E07[_name]["seed42"] else f"{_name}, seed 42 (E-07 rerun)", + f"{_name}, seed 43") + for pa, pb in itertools.combinations(twins, 2): + for sa, sb in [(0, 0), (1, 1), (0, 1), (1, 0)]: + a, b = twins[pa][sa], twins[pb][sb] + r = pair_row(a, b, "seed cross", P, O, tt, F, keep) + r.update(label=f"{pa} s{42 + sa} / {pb} s{42 + sb}", products=f"{pa} / {pb}", + same_seed=sa == sb, rho_basin=basin_rho(a, b)[0]) + rows.append(r) + df = pd.DataFrame(rows) + df.to_csv(HERE / "pairs.csv", index=False, float_format="%.6g") + + done = df[df.n_gauges.notna()] + hdr = (f"{'pair':44s} {'kind':9s} {'gauges':>6s} {'NSE a/b':>13s} {'avg-best':>8s} {'resid r':>7s} " + f"{'gain [95% CI]':>24s} {'yr/ssn/evt':>14s} {'reaches':>7s} {'n0 a/b':>13s} {'ratio':>5s} {'rho n0':>6s}") + lines += ["", hdr] + for _, r in df.iterrows(): + if not np.isfinite(r.get("n_gauges", np.nan)): + lines.append(f"{r.label:44s} PENDING seed-43 run from E-07 (branch paper/e07-seed43) not finished") + continue + if r.kind == "seed cross": + continue + nm = f"{SHORT[r.a]} / {SHORT[r.b]}" + lines.append( + f"{nm:44s} {r.kind[:9]:9s} {int(r.n_gauges):6d} {r.nse_a:.3f}/{r.nse_b:.3f} {r.avg_minus_better:+8.4f} " + f"{r.resid_corr:7.3f} {r.gain:+.4f} [{r.gain_lo:+.4f},{r.gain_hi:+.4f}] " + f"{r.share_year:.2f}/{r.share_season:.2f}/{r.share_event:.2f} {int(r.n_reaches):7d} " + f"{r.n0_med_a:.3f}/{r.n0_med_b:.3f} {r.n0_ratio:5.2f} {r.rho_n0:+6.3f}") + prod = done[done.kind == "product"] + seed = done[done.kind == "seed"] + lines += ["", + f"product pairs (n = {len(prod)}): resid r {prod.resid_corr.min():.3f}-{prod.resid_corr.max():.3f}, " + f"gain {prod.gain.min():+.4f} to {prod.gain.max():+.4f}, n0 ratio {prod.n0_ratio.min():.2f}-" + f"{prod.n0_ratio.max():.2f}, rho n0 {prod.rho_n0.min():+.3f} to {prod.rho_n0.max():+.3f}", + f"seed pairs done (n = {len(seed)}): resid r {', '.join(f'{x:.3f}' for x in seed.resid_corr)}, " + f"gain {', '.join(f'{x:+.4f}' for x in seed.gain)}, n0 ratio " + f"{', '.join(f'{x:.3f}' for x in seed.n0_ratio)}, rho n0 {', '.join(f'{x:+.3f}' for x in seed.rho_n0)}"] + for lab, d in prod.groupby("forcing"): + lines.append(f" {lab} (n = {len(d)}): median resid r {d.resid_corr.median():.3f}, median gain " + f"{d.gain.median():+.4f}") + rs = spearmanr(prod.n0_ratio, prod.gain).statistic + rr = spearmanr(prod.rho_n0, prod.resid_corr).statistic + rq = spearmanr(np.minimum(prod.nse_a, prod.nse_b), prod.gain).statistic + lines.append(f"across the 15 product pairs (exploratory, pairs share members): Spearman(n0 ratio, gain) {rs:+.2f}, " + f"Spearman(rho n0, resid r) {rr:+.2f}, Spearman(weaker member's NSE, gain) {rq:+.2f}") + s43 = done[done.kind == "product (UH seed 43)"].set_index("b") + s42 = prod[prod.a == "UH fixed-channel model"].set_index("b") + dd = (s43[["resid_corr", "gain", "rho_n0"]] - s42.loc[s43.index, ["resid_corr", "gain", "rho_n0"]]).abs() + lines.append(f"UH member swapped to seed 43, 5 product pairs: max |d resid r| {dd.resid_corr.max():.4f}, " + f"max |d gain| {dd.gain.max():.4f}, max |d rho n0| {dd.rho_n0.max():.4f}") + sx = done[done.kind == "seed cross"] + if len(sx): + lines += ["", "seed crossing (products with a seed-43 twin): same seed vs different seed, per product pair", + f" {'pair':34s} {'rho n0':>7s} {'basin rho':>9s} {'resid r':>7s} {'gain':>8s} {'ratio':>5s}"] + for _, r in sx.iterrows(): + lines.append(f" {r.label:34s} {r.rho_n0:+7.3f} {r.rho_basin:+9.3f} {r.resid_corr:7.3f} {r.gain:+8.4f} " + f"{r.n0_ratio:5.2f}") + g = sx.groupby(["products", "same_seed"])[["rho_n0", "rho_basin", "resid_corr", "gain"]].mean().unstack() + for pp in g.index: + lines.append(f" {pp:34s} mean same-seed minus mean cross-seed: rho n0 " + f"{g.loc[pp, ('rho_n0', True)] - g.loc[pp, ('rho_n0', False)]:+.3f}, basin rho " + f"{g.loc[pp, ('rho_basin', True)] - g.loc[pp, ('rho_basin', False)]:+.3f}, resid r " + f"{g.loc[pp, ('resid_corr', True)] - g.loc[pp, ('resid_corr', False)]:+.4f}, gain " + f"{g.loc[pp, ('gain', True)] - g.loc[pp, ('gain', False)]:+.4f}") + lines += ["", "per-reach n_0 agreement by drainage area (filtered reaches); log n_0 differences vs its spread"] + at = Dataset(ATTRS) + la = pd.Series(np.asarray(at["log10_uparea"][:], dtype=np.float64), + index=np.asarray(at["COMID"][:], dtype=np.int64)).reindex(comid0).values + detail = [(a, b) for a, b, _ in SEED_PAIRS if b is not None] + detail += [("UH fixed-channel model", "dHBV2 distributed model"), ("UH fixed-channel model, seed 43", "dHBV2 distributed model"), + ("UH fixed-channel model", "dHBV2 lumped model"), ("UH fixed-channel model, seed 43", "dHBV2 lumped model"), + ("NH daily LSTM model", "NH hourly MTS-LSTM model")] + for a, b in detail: + (na, ga), (nb, gb) = F[a], F[b] + _, bma, bmb = basin_rho(a, b) + lx, ly = np.log(na[keep]), np.log(nb[keep]) + iqr = np.subtract(*np.percentile(np.r_[lx, ly], [75, 25])) + by = [] + for lo, hi in [(0, 2), (2, 3), (3, 4), (4, 9)]: + s = keep & (la >= lo) & (la < hi) + by.append(f"10^{lo}-10^{hi} km2 ({int(s.sum()):,}) {spearmanr(na[s], nb[s]).statistic:+.2f}") + lines.append(f" {SHORT[a]} / {SHORT[b]}: rho n0 by area " + ", ".join(by) + + f"; median |d ln n0| {np.median(np.abs(lx - ly)):.3f} = {np.median(np.abs(lx - ly)) / iqr:.2f} " + f"of the pooled IQR of ln n0; rho gamma {spearmanr(ga[keep], gb[keep]).statistic:+.2f}; " + f"basin-median n0 over {len(ref_ids):,} gauges: rho {spearmanr(bma, bmb).statistic:+.2f}, " + f"median |d ln| {np.median(np.abs(np.log(bma / bmb))):.3f}") + txt = "\n".join(lines) + (HERE / "pairs.txt").write_text(txt + "\n") + print(txt) + plot(done) + + +def plot(df): + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + from matplotlib.ticker import FixedLocator, NullFormatter, ScalarFormatter + + # pairs named on the figure; the rest are identified in pairs.txt and the findings table + named = {("UH fixed-channel model", "dHBV2 distributed model"): "UH / dHBV2 distributed", + ("NH daily LSTM model", "NH hourly MTS-LSTM model"): "NH daily / NH hourly LSTM", + ("UH fixed-channel model", "dHBV2 lumped model"): "UH / dHBV2 lumped", + ("dHBV2 lumped model", "NH daily LSTM model"): "dHBV2 lumped / NH daily LSTM", + ("dHBV2 lumped model", "hydroDL LSTM model"): "dHBV2 lumped / hydroDL LSTM"} + OFFSETS = {"UH seeds 42/43": [(9, -2), (9, -2)], + "dHBV2 dist seeds 42/43": [(9, -4), (9, -4)], + "dHBV2 lumped seeds 42/43": [(9, -6), (9, -4)], + "UH / dHBV2 lumped": [(-20, -14), (6, 4)], + "dHBV2 lumped / NH daily LSTM": [None, (-125, 9)], + "dHBV2 lumped / hydroDL LSTM": [(6, -10), (6, -10)], + "NH daily / NH hourly LSTM": [(6, 4), (6, 6)]} + fig, ax = plt.subplots(1, 2, figsize=(13, 5.6)) + style = {"seed": dict(marker="*", s=300, c="black", label="seed pair: same inflow product, seed 42 vs 43", zorder=6), + "UH vs AORC": dict(marker="o", s=70, c="#d95f02", label="product pair: UH retrospective vs an AORC product"), + "both AORC": dict(marker="s", s=60, c="#1b9e77", label="product pair: two AORC-forced products")} + used = set() + sub = df[df.kind.isin(["seed", "product"])] + for _, r in sub.iterrows(): + key = "seed" if r.kind == "seed" else r.forcing + st = dict(style[key]) + lab = st.pop("label") + lab, _ = (None, used) if lab in used else (lab, used.add(lab)) + xy = [(r.n0_ratio, r.gain), (1 - r.rho_n0, 1 - r.resid_corr)] + ax[0].errorbar(r.n0_ratio, r.gain, yerr=[[r.gain - r.gain_lo], [r.gain_hi - r.gain]], fmt="none", + ecolor=st["c"], alpha=0.6, zorder=4) + ax[0].scatter(*xy[0], label=lab, edgecolor="white", linewidth=0.6, **st) + ax[1].scatter(*xy[1], edgecolor="white", linewidth=0.6, **st) + name = r.label if r.kind == "seed" else named.get((r.a, r.b)) + if name: + off = OFFSETS.get(name, [(6, 4), (6, 4)]) + for k in range(2): + if off[k] is not None: # None: too crowded on that panel, named in the table only + ax[k].annotate(name, xy[k], xytext=off[k], textcoords="offset points", fontsize=7.5, color="0.15") + seed = sub[sub.kind == "seed"] + for _, r in seed.iterrows(): # the seed floor on both parameter axes + ax[0].axvline(r.n0_ratio, color="0.45", ls="--", lw=0.7) + ax[1].axvline(1 - r.rho_n0, color="0.45", ls="--", lw=0.7) + ax[1].text(1 - seed.rho_n0.max(), 0.75, " dashed: the seed\n pairs' distances", fontsize=7.5, color="0.3", va="top") + ax[0].text(seed.n0_ratio.max(), 0.25, " dashed: the seed\n pairs' distances", fontsize=7.5, color="0.3", va="top") + ax[0].set_xscale("log") + ax[0].set_yscale("log") + ax[0].xaxis.set_major_locator(FixedLocator([1, 1.5, 2, 3, 4])) + ax[0].xaxis.set_major_formatter(ScalarFormatter()) + ax[0].xaxis.set_minor_formatter(NullFormatter()) + ax[0].set_xlabel("level distance: ratio of the two models' network-median $n_0$ (larger / smaller)") + ax[0].set_ylabel("averaging gain: NSE of the two-model average minus\nthe members' mean NSE (median over gauges, 95 % CI)") + ax[0].set_title("roughness level vs averaging gain") + ax[1].set_xlabel("pattern distance: 1 - per-reach Spearman $\\rho$ of $n_0$") + ax[1].set_ylabel("residual disagreement: 1 - median per-gauge\ncorrelation of routed residuals") + ax[1].set_yscale("log") + ax[1].set_title("roughness pattern vs residual disagreement") + ax[0].set_ylim(5e-5, 0.3) + ax[1].set_ylim(1e-4, 1) + for a in ax: + a.grid(alpha=0.3, which="both") + ax[0].legend(fontsize=7.5, loc="lower right") + fig.suptitle("Fixed-channel routing models, test WY1996-2010, 2,365 gauges, 107,225 reaches: " + "seed pairs vs inflow-product pairs", fontsize=11) + fig.tight_layout() + fig.savefig(HERE / "prediction_diversity.png", dpi=200, bbox_inches="tight", facecolor="white") + print(f"wrote {HERE / 'prediction_diversity.png'}") + + +if __name__ == "__main__": + main() diff --git a/experiments/prediction_diversity/reach_filter.py b/experiments/prediction_diversity/reach_filter.py new file mode 100644 index 00000000..780234c2 --- /dev/null +++ b/experiments/prediction_diversity/reach_filter.py @@ -0,0 +1,90 @@ +"""Per-reach coverage and median local inflow for every inflow product, test years WY1996-2010. + +The reach filter for the n_0 comparisons (the E-08 filter in the paper outline): a reach counts when its +divide has a Q' series in every product's store ("covered by every product") and its median local Q' over +the test window is at least 1e-3 m3/s in every product. Reaches without a Q' series are read by ddrs as the +1e-3 fill, so their learned n_0 never meets real water. + +The hourly MTS-LSTM store is reduced to daily means before the median (the routing is scored daily). + +Writes /reach_flow_medians.npz: COMID (the 346,321-reach parameter order of plot/kan_parameters.nc), +and per product a float32 median (NaN where the store has no series for that reach). + + /home/tbindas/projects/ddr/.venv/bin/python reach_filter.py +""" +import pathlib +import re +import time + +import icechunk +import numpy as np +import zarr +from netCDF4 import Dataset + +RUNS = pathlib.Path("/home/tbindas/projects/ddrs/.ddrs/runs") +CACHE = pathlib.Path("/home/tbindas/projects/ddrs/.ddrs/experiments/prediction-diversity") +T0, T1 = np.datetime64("1995-10-01"), np.datetime64("2010-10-01") +PRODUCTS = { # product label -> a seed-42 run that read it (store path comes from that run's config snapshot) + "uh": "2026-09-12T23-39-03Z-train-and-test", + "dhbv2_lumped": "2026-09-13T13-56-50Z-train-and-test", + "dhbv2_dist": "2026-09-13T17-22-30Z-train-and-test", + "hydrodl_lstm": "2026-09-13T17-21-58Z-train-and-test", + "nh_daily_lstm": "2026-09-13T13-55-03Z-train-and-test", + "nh_hourly_lstm": "2026-09-13T14-14-46Z-train-and-test", +} + + +def store_path(run): + m = re.search(r"^\s+streamflow:\s*(\S+)", (RUNS / run / "config.yaml").read_text(), re.M) + return m.group(1) + + +def median_local(path): + """(divide_id, median over the test window of daily local Q') for one store.""" + repo = icechunk.Repository.open(icechunk.local_filesystem_storage(path)) + g = zarr.open_group(repo.readonly_session("main").store, mode="r") + a = g["Qr"] + div = np.asarray(g["divide_id"][:], dtype=np.int64) + units = g["time"].attrs["units"] + step = "h" if units.startswith("hours") else "D" + origin = np.datetime64(units.split("since")[-1].strip().split()[0]) + tt = origin + np.asarray(g["time"][:]).astype(f"timedelta64[{step}]") + i0, i1 = int(np.searchsorted(tt, T0)), int(np.searchsorted(tt, T1)) + assert a.shape[0] == div.size, "expected divide-major Qr" + nday = (i1 - i0) // 24 if step == "h" else i1 - i0 + assert nday == 5479, nday + med = np.empty(div.size, dtype=np.float32) + cb = a.chunks[0] + for s in range(0, div.size, cb): + x = np.asarray(a[s:s + cb, i0:i1], dtype=np.float32) + if step == "h": + x = x.reshape(x.shape[0], nday, 24).mean(2) + med[s:s + cb] = np.nanmedian(x, axis=1) + return div, med + + +def main(): + CACHE.mkdir(parents=True, exist_ok=True) + comid = np.asarray(Dataset(RUNS / PRODUCTS["uh"] / "plot" / "kan_parameters.nc")["COMID"][:], dtype=np.int64) + pos = {c: i for i, c in enumerate(comid)} + out = {"COMID": comid} + for lab, run in PRODUCTS.items(): + t = time.time() + path = store_path(run) + div, med = median_local(path) + col = np.full(comid.size, np.nan, dtype=np.float32) + ix = np.array([pos.get(int(d), -1) for d in div]) + col[ix[ix >= 0]] = med[ix >= 0] + out[lab] = col + print(f"{lab:16s} {path}: {div.size:,} divides, {int((ix >= 0).sum()):,} in the network, " + f"median of medians {np.nanmedian(col):.4f} m3/s, {time.time() - t:.0f} s", flush=True) + np.savez(CACHE / "reach_flow_medians.npz", **out) + M = np.stack([out[k] for k in PRODUCTS]) + covered = np.isfinite(M).all(0) + keep = covered & (np.nan_to_num(M, nan=0.0) >= 1e-3).all(0) + print(f"covered by every product: {int(covered.sum()):,} of {comid.size:,}; " + f"and median local Q' >= 1e-3 m3/s in every product: {int(keep.sum()):,}") + + +if __name__ == "__main__": + main() diff --git a/experiments/seed_replicates/seed_pairs.py b/experiments/seed_replicates/seed_pairs.py new file mode 100644 index 00000000..2f84cb53 --- /dev/null +++ b/experiments/seed_replicates/seed_pairs.py @@ -0,0 +1,279 @@ +"""E-07 seed-pair read-out: how far skill, n_0 and gamma move when only the training seed changes. + +Three fixed-channel models (n_0 + gamma head, p = 21, q = 0.65, 2,365 gauges, 500 Adam updates), each +trained at seed 42 and seed 43 on the same binary, config and inflow store: + + UH fixed-channel model 2026-09-12T23-39-03Z (canonical seed 42 per E-01) / 2026-09-27T10-31-30Z + dHBV2 distributed model 2026-09-13T17-22-30Z / 2026-10-01T02-33-37Z + dHBV2 lumped model 2026-09-13T13-56-50Z / 2026-10-01T02-33-01Z + +Skill. Every number is recomputed from eval/predictions.zarr (loader from the ddrs-eval-plots skill) in the +ddrs convention E-01 uses (paper/e01-lineage@e8e7d3b, research/findings/2026-09-30-paper-model-lineage.md): +drop the first experiment.warmup = 5 days, then score with ddr.validation.Metrics (NSE; KGE 2009 with +r, sd ratio, mean ratio). The recomputed medians are checked against each manifest's median_nse_finite / +median_kge_finite. Per pair: + - difference of medians (seed 43 minus seed 42), with a paired gauge bootstrap (B = 2000, resampling + gauges jointly for both seeds) 95 % percentile CI; + - median of the per-gauge difference (same bootstrap), the share of gauges where seed 43 is better, and + the spread of |per-gauge difference|. +Residuals. Per gauge, corr(p42 - o, p43 - o) over days where o and both predictions are finite, gauges with +>= 365 such days (prompt 05's convention, experiments/prediction_diversity/pair_diversity.py on +paper/prediction-diversity@8d8000c, there on the full series). +Fields. n_0 (`n`) and gamma from plot/kan_parameters.nc. Network medians over all 346,321 reaches and over +E-08's filtered reaches; per-reach Spearman rho of n_0 and of gamma on the filtered reaches (and on all +reaches, the section 38.2 population, for comparison); rho of n_0 by drainage-area class; the share of +reaches at the n_0 ceiling (>= 0.2495) and at gamma = 0 (< 0.001), as e08_fields.py counts them. +Reach filter (E-08, paper/e10-e08-cross-product@283140e, experiments/cross_product/qprime_medians.py and +e08_fields.py): finite log10_uparea, COMID on every one of the six product stores' divide_id axes, and +median daily Q' over WY1996-2010 >= 1e-3 m3/s in all six products. Per-product medians read from that +commit's cache, /home/tbindas/projects/ddrs/.ddrs/experiments/cross_product/2026-09-30/. Expected 107,225. + + /home/tbindas/projects/ddr/.venv/bin/python experiments/seed_replicates/seed_pairs.py + +Writes seed_pairs.txt and seed_pairs.csv next to this script, and per-gauge tables under +/home/tbindas/projects/ddrs/.ddrs/experiments/seed_replicates//. +""" +from __future__ import annotations + +import io +import json +import sys +from datetime import date +from pathlib import Path + +import numpy as np +import pandas as pd +from netCDF4 import Dataset +from scipy.stats import spearmanr + +sys.path.insert(0, "/home/tbindas/projects/ddrs/.claude/skills/ddrs-eval-plots/scripts") +from load_ddrs_predictions import load_predictions_zarr # noqa: E402 +from ddr.validation import Metrics # noqa: E402 + +HERE = Path(__file__).resolve().parent +RUNS = Path("/home/tbindas/projects/ddrs/.ddrs/runs") +QMED = Path("/home/tbindas/projects/ddrs/.ddrs/experiments/cross_product/2026-09-30") +OUT = Path("/home/tbindas/projects/ddrs/.ddrs/experiments/seed_replicates") / date.today().isoformat() +ATTRS = "/home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc" +PRODUCTS = ["uh", "dhbv2_lumped", "dhbv2_dist", "hydrodl_lstm", "nh_daily_lstm", "nh_hourly_mts"] +WARMUP = 5 +B, BOOT_SEED = 2000, 0 +MIN_DAYS = 365 +N_CEIL, G_ZERO = 0.2495, 0.001 + +PAIRS = [ + ("UH fixed-channel model", "2026-09-12T23-39-03Z-train-and-test", "2026-09-27T10-31-30Z-train-and-test"), + ("dHBV2 distributed model", "2026-09-13T17-22-30Z-train-and-test", "2026-10-01T02-33-37Z-train-and-test"), + ("dHBV2 lumped model", "2026-09-13T13-56-50Z-train-and-test", "2026-10-01T02-33-01Z-train-and-test"), +] +# E-01: bitwise twin of the canonical UH seed-42 run; checked here, not used further. +UH_TWIN = "2026-09-27T07-29-47Z-train-and-test" +AREA_BINS = [(0, 2), (2, 3), (3, 4), (4, 9)] + + +def load_eval(run): + ds = load_predictions_zarr(RUNS / run / "eval" / "predictions.zarr") + return (np.asarray(ds.gage_ids.values), np.asarray(ds.time.values), + ds.predictions.values.astype(np.float64), ds.observations.values.astype(np.float64)) + + +def load_fields(run): + ds = Dataset(RUNS / run / "plot" / "kan_parameters.nc") + return (np.asarray(ds["COMID"][:], dtype=np.int64), np.asarray(ds["n"][:], dtype=np.float64), + np.asarray(ds["gamma"][:], dtype=np.float64)) + + +def reach_filter(comid): + """E-08's filtered reach set, as a boolean mask over `comid`.""" + at = Dataset(ATTRS) + la = pd.Series(np.asarray(at["log10_uparea"][:], dtype=np.float64), + index=np.asarray(at["COMID"][:], dtype=np.int64)).reindex(comid).to_numpy() + keep = np.isfinite(la) + covered = keep.copy() + for k in PRODUCTS: + z = np.load(QMED / f"qprime_median_{k}.npz") + med = pd.Series(z["median"], index=z["divide_id"]) + on = np.isin(comid, z["divide_id"]) + covered &= on + keep &= on & (med.reindex(comid).to_numpy() >= 1e-3) + return keep, covered, la + + +def paired_boot(a, b): + """Difference of medians and median of differences, each with a paired gauge-bootstrap 95 % CI.""" + ok = np.isfinite(a) & np.isfinite(b) + a, b = a[ok], b[ok] + idx = np.random.default_rng(BOOT_SEED).integers(0, a.size, size=(B, a.size)) + dm = np.median(b[idx], axis=1) - np.median(a[idx], axis=1) + md = np.median((b - a)[idx], axis=1) + return dict( + n=int(a.size), med_a=float(np.median(a)), med_b=float(np.median(b)), + diff_of_medians=float(np.median(b) - np.median(a)), + dom_lo=float(np.percentile(dm, 2.5)), dom_hi=float(np.percentile(dm, 97.5)), + median_of_diff=float(np.median(b - a)), + mod_lo=float(np.percentile(md, 2.5)), mod_hi=float(np.percentile(md, 97.5)), + share_b_better=float(np.mean(b > a)), + absdiff_p50=float(np.median(np.abs(b - a))), absdiff_p90=float(np.percentile(np.abs(b - a), 90)), + share_absdiff_gt_005=float(np.mean(np.abs(b - a) > 0.05)), + ) + + +def resid_corr(Pa, Pb, O): + m = np.isfinite(O) & np.isfinite(Pa) & np.isfinite(Pb) + nv = m.sum(1) + ea, eb = np.where(m, Pa - O, 0.0), np.where(m, Pb - O, 0.0) + ma, mb = ea.sum(1) / np.maximum(nv, 1), eb.sum(1) / np.maximum(nv, 1) + da, db = np.where(m, ea - ma[:, None], 0.0), np.where(m, eb - mb[:, None], 0.0) + with np.errstate(invalid="ignore", divide="ignore"): + rc = (da * db).sum(1) / np.sqrt((da ** 2).sum(1) * (db ** 2).sum(1)) + rc[nv < MIN_DAYS] = np.nan + return rc + + +def boot_median(x): + x = x[np.isfinite(x)] + idx = np.random.default_rng(BOOT_SEED).integers(0, x.size, size=(B, x.size)) + m = np.median(x[idx], axis=1) + return float(np.median(x)), float(np.percentile(m, 2.5)), float(np.percentile(m, 97.5)), int(x.size) + + +def main(): + buf = io.StringIO() + + def say(*a): + print(*a) + print(*a, file=buf) + + OUT.mkdir(parents=True, exist_ok=True) + + # UH twin check (E-01 says bitwise identical) + g0, t0, P0, O0 = load_eval(PAIRS[0][1]) + g1, t1, P1, O1 = load_eval(UH_TWIN) + assert np.array_equal(g0, g1) and np.array_equal(t0, t1) + c0, n0a, ga0 = load_fields(PAIRS[0][1]) + c1, n0b, gb1 = load_fields(UH_TWIN) + assert np.array_equal(c0, c1) + say(f"UH seed-42 twin check, {PAIRS[0][1][:20]} vs {UH_TWIN[:20]}: max |dQ| {np.nanmax(np.abs(P0 - P1)):.3g} m3/s, " + f"max |d n_0| {np.max(np.abs(n0a - n0b)):.3g}, max |d gamma| {np.max(np.abs(ga0 - gb1)):.3g}") + + comid = c0 + keep, covered, la = reach_filter(comid) + say(f"reaches: {comid.size:,} in the network, {int(covered.sum()):,} with finite log10_uparea and covered by every " + f"product store, {int(keep.sum()):,} also with median daily Q' >= 1e-3 m3/s in all six (E-08 filter)") + assert int(keep.sum()) == 107_225, "E-08 filter count changed" + + rows, gauge_tabs = [], [] + for name, r42, r43 in PAIRS: + say(f"\n== {name}: seed 42 {r42[:20]}, seed 43 {r43[:20]}") + ga, ta, Pa, Oa = load_eval(r42) + gb, tb, Pb, Ob = load_eval(r43) + assert np.array_equal(ta, tb), "time axes differ" + assert set(ga) == set(gb), "gauge sets differ" + ix = {s: i for i, s in enumerate(gb)} + order = [ix[s] for s in ga] + Pb, Ob = Pb[order], Ob[order] + assert np.array_equal(np.isnan(Oa), np.isnan(Ob)) and np.allclose(Oa[np.isfinite(Oa)], Ob[np.isfinite(Ob)]) + Pa, Pb, O = Pa[:, WARMUP:], Pb[:, WARMUP:], Oa[:, WARMUP:] + ma, mb = Metrics(pred=Pa, target=O), Metrics(pred=Pb, target=O) + for run, m in ((r42, ma), (r43, mb)): + man = json.loads((RUNS / run / "manifest.json").read_text())["metrics"] + say(f" {run[:20]}: recomputed median NSE {np.nanmedian(m.nse):.4f} KGE {np.nanmedian(m.kge):.4f}; " + f"manifest {man['median_nse_finite']:.4f} / {man['median_kge_finite']:.4f}; " + f"{int(np.isfinite(m.nse).sum())} gauges with finite NSE") + s_nse, s_kge = paired_boot(ma.nse, mb.nse), paired_boot(ma.kge, mb.kge) + for lab, s in (("NSE", s_nse), ("KGE", s_kge)): + say(f" {lab} ({s['n']} gauges): median {s['med_a']:.4f} -> {s['med_b']:.4f}, difference of medians " + f"{s['diff_of_medians']:+.4f} [{s['dom_lo']:+.4f}, {s['dom_hi']:+.4f}]; median per-gauge difference " + f"{s['median_of_diff']:+.4f} [{s['mod_lo']:+.4f}, {s['mod_hi']:+.4f}]; seed 43 better at " + f"{100 * s['share_b_better']:.1f} %; |difference| p50 {s['absdiff_p50']:.4f}, p90 {s['absdiff_p90']:.4f}, " + f"> 0.05 at {100 * s['share_absdiff_gt_005']:.1f} %") + rc = resid_corr(Pa, Pb, O) + rmed, rlo, rhi, rn = boot_median(rc) + say(f" routed residual correlation per gauge ({rn} gauges with >= {MIN_DAYS} days): median {rmed:.4f} " + f"[{rlo:.4f}, {rhi:.4f}], p10 {np.nanpercentile(rc, 10):.4f}, p90 {np.nanpercentile(rc, 90):.4f}, " + f"min {np.nanmin(rc):.4f}") + # full-series residual correlation, prompt 05's window, for the cross-check + rc_full = resid_corr(load_eval(r42)[2], load_eval(r43)[2][order], Oa) + say(f" same on the full series (no warmup drop): median {np.nanmedian(rc_full):.4f}") + + ca, na, gma = load_fields(r42) + cb, nb, gmb = load_fields(r43) + assert np.array_equal(ca, comid) and np.array_equal(cb, comid) + fld = dict( + n0_med_all_a=np.median(na), n0_med_all_b=np.median(nb), + n0_med_f_a=np.median(na[keep]), n0_med_f_b=np.median(nb[keep]), + g_med_all_a=np.median(gma), g_med_all_b=np.median(gmb), + g_med_f_a=np.median(gma[keep]), g_med_f_b=np.median(gmb[keep]), + rho_n0_f=spearmanr(na[keep], nb[keep]).statistic, rho_g_f=spearmanr(gma[keep], gmb[keep]).statistic, + rho_n0_all=spearmanr(na, nb).statistic, rho_g_all=spearmanr(gma, gmb).statistic, + ceil_f_a=np.mean(na[keep] >= N_CEIL), ceil_f_b=np.mean(nb[keep] >= N_CEIL), + g0_f_a=np.mean(gma[keep] < G_ZERO), g0_f_b=np.mean(gmb[keep] < G_ZERO), + dln_f=np.median(np.abs(np.log(na[keep]) - np.log(nb[keep]))), + g_p10_f_a=np.percentile(gma[keep], 10), g_p90_f_a=np.percentile(gma[keep], 90), + g_p10_f_b=np.percentile(gmb[keep], 10), g_p90_f_b=np.percentile(gmb[keep], 90), + n0_p10_f_a=np.percentile(na[keep], 10), n0_p90_f_a=np.percentile(na[keep], 90), + n0_p10_f_b=np.percentile(nb[keep], 10), n0_p90_f_b=np.percentile(nb[keep], 90), + ) + fld = {k: float(v) for k, v in fld.items()} + say(f" n_0 network median, all {comid.size:,} reaches: {fld['n0_med_all_a']:.4f} -> {fld['n0_med_all_b']:.4f}; " + f"filtered {int(keep.sum()):,}: {fld['n0_med_f_a']:.4f} -> {fld['n0_med_f_b']:.4f} " + f"(ratio {max(fld['n0_med_f_a'], fld['n0_med_f_b']) / min(fld['n0_med_f_a'], fld['n0_med_f_b']):.3f}); " + f"filtered p10-p90 {fld['n0_p10_f_a']:.3f}-{fld['n0_p90_f_a']:.3f} -> {fld['n0_p10_f_b']:.3f}-{fld['n0_p90_f_b']:.3f}") + say(f" gamma network median, all: {fld['g_med_all_a']:.4f} -> {fld['g_med_all_b']:.4f}; filtered: " + f"{fld['g_med_f_a']:.4f} -> {fld['g_med_f_b']:.4f}; filtered p10-p90 {fld['g_p10_f_a']:.3f}-{fld['g_p90_f_a']:.3f} -> " + f"{fld['g_p10_f_b']:.3f}-{fld['g_p90_f_b']:.3f}") + say(f" filtered reaches at the n_0 ceiling: {100 * fld['ceil_f_a']:.1f} % -> {100 * fld['ceil_f_b']:.1f} %; " + f"at gamma = 0: {100 * fld['g0_f_a']:.1f} % -> {100 * fld['g0_f_b']:.1f} %") + say(f" per-reach Spearman rho, filtered: n_0 {fld['rho_n0_f']:+.3f}, gamma {fld['rho_g_f']:+.3f}; all reaches: " + f"n_0 {fld['rho_n0_all']:+.3f}, gamma {fld['rho_g_all']:+.3f}; median |d ln n_0| (filtered) {fld['dln_f']:.3f}") + by = [] + for lo, hi in AREA_BINS: + s = keep & (la >= lo) & (la < hi) + by.append(f"10^{lo}-10^{hi} km2 ({int(s.sum()):,}) {spearmanr(na[s], nb[s]).statistic:+.2f}") + say(" rho n_0 by drainage area (filtered): " + ", ".join(by)) + + rows.append(dict(model=name, run_s42=r42, run_s43=r43, n_reaches_filtered=int(keep.sum()), + **{f"nse_{k}": v for k, v in s_nse.items()}, **{f"kge_{k}": v for k, v in s_kge.items()}, + resid_r_median=rmed, resid_r_lo=rlo, resid_r_hi=rhi, resid_r_n=rn, **fld)) + gauge_tabs.append(pd.DataFrame(dict(model=name, gauge=ga, nse_s42=ma.nse, nse_s43=mb.nse, kge_s42=ma.kge, + kge_s43=mb.kge, resid_r=rc))) + + # Product pairs among the three models, at matched seeds (42/42, 43/43) and crossed seeds (42/43, 43/42). + # All six product models in section 38.2 / E-08 share seed 42, i.e. the same head initialisation and gauge + # batch order; the crossed pairs show how much of their per-reach agreement that shared seed carries. + say("\n== product pairs among these three models, matched vs crossed seeds (filtered reaches)") + say(f"{'pair':48s} {'seeds':>6s} {'rho n0':>7s} {'rho gamma':>9s} {'n0 ratio':>8s} {'resid r':>8s}") + runs = {name: {42: r42, 43: r43} for name, r42, r43 in PAIRS} + F = {(name, s): load_fields(r)[1:] for name, d in runs.items() for s, r in d.items()} + E = {} + for name, d in runs.items(): + for s, r in d.items(): + g, _, P, Oo = load_eval(r) + ix = {x: i for i, x in enumerate(g)} + E[(name, s)] = P[[ix[x] for x in g0]][:, WARMUP:] + Oref = O0[:, WARMUP:] + xrows = [] + names = [p[0] for p in PAIRS] + for i in range(3): + for j in range(i + 1, 3): + a, b = names[i], names[j] + for sa, sb in ((42, 42), (43, 43), (42, 43), (43, 42)): + (na, ga_), (nb, gb_) = F[(a, sa)], F[(b, sb)] + ma_, mb_ = np.median(na[keep]), np.median(nb[keep]) + rcx = np.nanmedian(resid_corr(E[(a, sa)], E[(b, sb)], Oref)) + r = dict(a=a, b=b, seed_a=sa, seed_b=sb, rho_n0=spearmanr(na[keep], nb[keep]).statistic, + rho_gamma=spearmanr(ga_[keep], gb_[keep]).statistic, + n0_ratio=max(ma_, mb_) / min(ma_, mb_), resid_r=rcx) + xrows.append(r) + say(f"{a + ' / ' + b:48s} {sa}/{sb:<3d} {r['rho_n0']:+7.3f} {r['rho_gamma']:+9.3f} {r['n0_ratio']:8.3f} " + f"{r['resid_r']:8.4f}") + pd.DataFrame(xrows).to_csv(HERE / "product_pairs_by_seed.csv", index=False, float_format="%.6g") + + pd.DataFrame(rows).to_csv(HERE / "seed_pairs.csv", index=False, float_format="%.6g") + pd.concat(gauge_tabs).to_csv(OUT / "seed_pairs_per_gauge.csv", index=False, float_format="%.6g") + (HERE / "seed_pairs.txt").write_text(buf.getvalue()) + print(f"wrote {HERE / 'seed_pairs.txt'}, {HERE / 'seed_pairs.csv'}, {OUT / 'seed_pairs_per_gauge.csv'}") + + +if __name__ == "__main__": + main() diff --git a/research/scripts/e11_full_store_equal.py b/research/scripts/e11_full_store_equal.py new file mode 100644 index 00000000..ec7d7473 --- /dev/null +++ b/research/scripts/e11_full_store_equal.py @@ -0,0 +1,13 @@ +"""E-11: full-store equality of the 08-14 original lumped store and its 09-13 rechunk, plus whole-store stats.""" +import icechunk, zarr, numpy as np, warnings; warnings.filterwarnings("ignore") +B="/mnt/ssd1/data/icechunk/" +def op(n): return zarr.open_group(icechunk.Repository.open(icechunk.local_filesystem_storage(B+n)).readonly_session("main").store, mode="r")["Qr"] +a=op("daily_dhbv2_lumped_aorc2f_merit_unit_catchments.ic"); b=op("daily_dhbv2_lumped_aorc2f_merit_unit_catchments_rechunked.ic") +n=a.shape[0]; ne=0; tot=0; s=0.0; nan=0; fl=0; allfl=[]; mn=np.inf; mx=0.0; neg=0 +for i in range(0,n,3080): + x=a[i:i+3080,:]; y=b[i:i+3080,:] + ne+=int((x!=y).sum()); tot+=x.size; s+=float(x.sum(dtype="f8")); nan+=int(np.isnan(x).sum()); fl+=int((x==np.float32(1e-6)).sum()) + allfl.append((x.max(1)<=np.float32(1e-6))); mn=min(mn,float(x.min())); mx=max(mx,float(x.max())); neg+=int((x<0).sum()) +allfl=np.concatenate(allfl) +print(f"rows {n} values {tot} unequal {ne} mean {s/tot:.4f} m3/s NaN {nan} floor frac {fl/tot:.4f} min {mn:.2e} max {mx:.1f} neg {neg} divides entirely at floor {int(allfl.sum())}") +np.save("/tmp/e11_allfloor_rows.npy", np.where(allfl)[0]) diff --git a/research/scripts/e11_store_audit.py b/research/scripts/e11_store_audit.py new file mode 100644 index 00000000..eed0f915 --- /dev/null +++ b/research/scripts/e11_store_audit.py @@ -0,0 +1,86 @@ +"""E-11 store audit. Read-only. Run with /home/tbindas/projects/ddr/.venv/bin/python.""" +import icechunk, zarr, numpy as np, xarray as xr, json, warnings +warnings.filterwarnings("ignore") +B = "/mnt/ssd1/data/icechunk/" +S = { + "lumped_orig_0814": "daily_dhbv2_lumped_aorc2f_merit_unit_catchments.ic", + "lumped_rechunk": "daily_dhbv2_lumped_aorc2f_merit_unit_catchments_rechunked.ic", + "distributed": "daily_dhbv2_distributed_aorc2f_merit_unit_catchments_rechunked.ic", + "UH_retro": "merit_dhbv2_UH_retrospective.ic", + "hydroDL_lstm": "daily_hydrodl_lstm_aorc2f_merit_unit_catchments_rechunked.ic", +} +G = {k: zarr.open_group(icechunk.Repository.open(icechunk.local_filesystem_storage(B+v)).readonly_session("main").store, mode="r") for k, v in S.items()} +ids = {k: g["divide_id"][:] for k, g in G.items()} +for k in ids: print(k, "divide_id equal to lumped_rechunk:", bool((ids[k] == ids["lumped_rechunk"]).all()), "n", len(ids[k])) +t = G["lumped_rechunk"]["time"][:] +print("time", t[0], t[-1], "step==1:", bool((np.diff(t) == 1).all()), "n", len(t), "-> 1980-01-01 +", t[-1], "d =", np.datetime64("1980-01-01") + int(t[-1])) +for k in G: print(k, "time equal:", bool((G[k]["time"][:] == t).all())) +A = xr.open_dataset("/home/tbindas/projects/ddr/data/merit_global_attributes_v2.nc") +cs = A["catchsize"].sel(COMID=ids["lumped_rechunk"]).values.astype("f8") # km2 +print("catchsize km2 quantiles", np.nanquantile(cs, [0, .05, .5, .95, 1]), "NaN", int(np.isnan(cs).sum())) +UHPOS = np.full(int(max(ids["lumped_rechunk"].max(), ids["UH_retro"].max()))+1, -1, dtype=np.int64); print("UH id sets equal:", set(ids["UH_retro"].tolist()) == set(ids["lumped_rechunk"].tolist())) +UHPOS[ids["UH_retro"]] = np.arange(len(ids["UH_retro"])) +blocks = [3, 13, 23, 33, 43, 53] +tr = slice(5479*0, None) +rows = {k: [] for k in G} +eq_ok = True +res = {} +LAGMAX = 10 +FLOOR = np.float32(1e-6) +agg = {k: dict(n=0, nan=0, floor=0, allfloor=0, sum=0.0) for k in G} +cmp = {k: dict(corr=[], lag=[], ratio=[]) for k in ["distributed", "UH_retro", "hydroDL_lstm"]} +depth_bins = np.array([0, 3, 10, 30, 100, 1e9]) +dep = {k: [[] for _ in range(5)] for k in G} +shape_stats = {k: dict(cv=[], p98r=[]) for k in G} +orig_eq = 0; orig_n = 0 +for b in blocks: + s = slice(b*3080, (b+1)*3080) + X = {k: g["Qr"][s, :] for k, g in G.items() if k != "UH_retro"} + # UH_retro stores divides in a different row order: align by divide_id + pos = UHPOS[ids["lumped_rechunk"][s]] + assert (pos >= 0).all() + lo, hi = pos.min(), pos.max()+1 + X["UH_retro"] = G["UH_retro"]["Qr"][lo:hi, :][pos-lo] + orig_eq += int((X["lumped_orig_0814"] == X["lumped_rechunk"]).sum()); orig_n += X["lumped_rechunk"].size + area = cs[s] + for k, x in X.items(): + a = agg[k]; a["n"] += x.size; a["nan"] += int(np.isnan(x).sum()); a["floor"] += int((x == FLOOR).sum()) + a["allfloor"] += int((x.max(1) <= FLOOR).sum()); a["sum"] += float(np.nansum(x, dtype="f8")) + m = np.nanmean(x, 1, dtype="f8"); depth = m*86.4/area + bi = np.digitize(area, depth_bins[1:-1]) + for i in range(5): + sel = bi == i + if sel.any(): dep[k][i].append(depth[sel & np.isfinite(depth)]) + L = X["lumped_rechunk"].astype("f8"); mL = L.mean(1) + ok = mL > 0.01 + for k in cmp: + Y = X[k].astype("f8"); Ys = Y[ok]; Ls = L[ok] + cmp[k]["ratio"].append(Ls.mean(1)/np.maximum(Ys.mean(1), 1e-12)) + Lz = Ls - Ls.mean(1, keepdims=True); Yz = Ys - Ys.mean(1, keepdims=True) + den = np.sqrt((Lz**2).sum(1)*(Yz**2).sum(1)) + 1e-12 + cc = {} + for lag in range(-LAGMAX, LAGMAX+1): # lag>0: lumped leads (lumped[t] matches product[t+lag]) + if lag >= 0: num = (Lz[:, :Lz.shape[1]-lag]*Yz[:, lag:]).sum(1) + else: num = (Lz[:, -lag:]*Yz[:, :Yz.shape[1]+lag]).sum(1) + cc[lag] = num/den + M = np.stack([cc[l] for l in range(-LAGMAX, LAGMAX+1)], 1) + cmp[k]["corr"].append(cc[0]); cmp[k]["lag"].append(np.argmax(M, 1)-LAGMAX) + for k, x in X.items(): + xs = x[ok].astype("f8") + shape_stats[k]["cv"].append(xs.std(1)/xs.mean(1)) + shape_stats[k]["p98r"].append(np.quantile(xs, .98, axis=1)/xs.mean(1)) +print("\noriginal(08-14 snapshot) vs rechunked identical values: %d / %d" % (orig_eq, orig_n)) +print("\nper-store sample stats (6 row blocks = %d divides)" % (len(blocks)*3080)) +for k, a in agg.items(): + print(f"{k:18s} mean {a['sum']/a['n']:.4f} m3/s NaN {a['nan']} ==1e-6 floor frac {a['floor']/a['n']:.4f} divides entirely floor {a['allfloor']}") +print("\nlumped vs product, divides with lumped mean>0.01 m3/s") +for k, c in cmp.items(): + r = np.concatenate(c["ratio"]); cr = np.concatenate(c["corr"]); lg = np.concatenate(c["lag"]) + print(f"vs {k:13s} n={len(r)} mean-ratio median {np.median(r):.3f} [IQR {np.quantile(r,.25):.3f},{np.quantile(r,.75):.3f}] daily corr median {np.median(cr):.3f} best-lag distribution (days, + = lumped leads) " + str({int(l): round(float((lg == l).mean()), 3) for l in range(-LAGMAX, LAGMAX+1) if (lg == l).mean() > 0.005})) +print("\nrunoff depth (mm/day, mean over 1980-2020) by divide area class") +names = ["<3 km2", "3-10", "10-30", "30-100", ">100"] +for i, n in enumerate(names): + print(f"{n:8s} " + " ".join(f"{k}: {np.median(np.concatenate(dep[k][i])):.3f} (n={sum(len(z) for z in dep[k][i])})" for k in ["lumped_rechunk", "distributed", "UH_retro", "hydroDL_lstm"] if dep[k][i])) +print("\nper-divide flashiness (divides with lumped mean>0.01)") +for k, s_ in shape_stats.items(): + print(f"{k:18s} CV median {np.median(np.concatenate(s_['cv'])):.2f} p98/mean median {np.median(np.concatenate(s_['p98r'])):.2f}") diff --git a/research/scripts/e11_unrouted_counterfactual.py b/research/scripts/e11_unrouted_counterfactual.py new file mode 100644 index 00000000..378d60b2 --- /dev/null +++ b/research/scripts/e11_unrouted_counterfactual.py @@ -0,0 +1,49 @@ +"""E-11: what explains the low unrouted NSE of the lumped product? Counterfactual edits of the summed-Q' series. +Read-only on runs. Run with ddrs-py/.venv python (numpy, xarray, zarr) or the ddr venv.""" +import json, numpy as np, zarr, sys +R = "/home/tbindas/projects/ddrs/.ddrs/runs/" +RUNS = {"dHBV2 lumped": "2026-09-13T13-56-50Z-train-and-test", "dHBV2 distributed": "2026-09-13T17-22-30Z-train-and-test", + "UH": "2026-09-12T23-39-03Z-train-and-test", "hydroDL LSTM": "2026-09-13T17-21-58Z-train-and-test"} +def eval_ids(run): + z = zarr.open_group(R+run+"/eval/predictions.zarr", mode="r") + a = np.asarray(z["gage_ids"][:]) + if a.dtype == np.uint8: return ["".join(map(chr, r)).strip("\x00 ") for r in a] + return [str(x) for x in a] +def nse(p, o): + m = np.isfinite(o) & np.isfinite(p) + num = np.where(m, (p-o)**2, 0).sum(1); den = np.where(m, (o-np.nanmean(np.where(m, o, np.nan), 1, keepdims=True))**2, 0).sum(1) + return 1-num/den +def delay(x, d): + return x if d == 0 else np.concatenate([np.repeat(x[:, :1], d, 1), x[:, :-d]], 1) +def trail(x, k): + c = np.cumsum(np.concatenate([np.repeat(x[:, :1], k-1, 1), x], 1), 1); c = np.concatenate([np.zeros((x.shape[0], 1)), c], 1) + return (c[:, k:]-c[:, :-k])/k +out = {} +for name, run in RUNS.items(): + m = json.load(open(R+run+"/baseline/manifest.json")); ids = m["gage_ids"]; n = m["n_days"] + P = np.fromfile(R+run+"/baseline/predictions.f32", "f4").reshape(len(ids), n).astype("f8") + O = np.fromfile(R+run+"/baseline/observations.f32", "f4").reshape(len(ids), n).astype("f8") + ev = eval_ids(run); ix = {g: i for i, g in enumerate(ids)}; sel = [ix[g] for g in ev] + P = P[sel, 1:-1]; O = O[sel, 1:-1] # 5477 d, as E-10 + W = 10 # common scoring window start so shifted variants are comparable + res = {"gauges": len(sel), "raw_full": float(np.nanmedian(nse(P, O)))} + P0, O0 = P[:, W:], O[:, W:] + res["raw"] = float(np.nanmedian(nse(P0, O0))) + for d in (1, 2, 3): + res[f"delay{d}"] = float(np.nanmedian(nse(delay(P, d)[:, W:], O0))) + best = np.max(np.stack([nse(delay(P, d)[:, W:], O0) for d in range(0, 6)], 1), 1) + res["oracle_delay0to5"] = float(np.nanmedian(best)) + for k in (2, 3, 5): + res[f"trail{k}"] = float(np.nanmedian(nse(trail(P, k)[:, W:], O0))) + res["delay1+trail3"] = float(np.nanmedian(nse(trail(delay(P, 1), 3)[:, W:], O0))) + res["delay2+trail3"] = float(np.nanmedian(nse(trail(delay(P, 2), 3)[:, W:], O0))) + mo = np.isfinite(O0) + vr = np.where(mo, P0, 0).sum(1)/np.where(mo, O0, 0).sum(1) + res["volume_rescale_only"] = float(np.nanmedian(nse(P0/vr[:, None], O0))) + res["delay1+trail3+volume"] = float(np.nanmedian(nse(trail(delay(P, 1), 3)[:, W:]/vr[:, None], O0))) + res["share_nse_below0"] = float((nse(P0, O0) < 0).mean()) + out[name] = res +keys = list(next(iter(out.values())).keys()) +print("%-24s" % "median NSE" + "".join("%20s" % k for k in out)) +for k in keys: + print("%-24s" % k + "".join("%20.4f" % out[n][k] if k != "gauges" else "%20d" % out[n][k] for n in out)) diff --git a/scripts/landscape_shards.sh b/scripts/landscape_shards.sh index cb803647..555f0d3c 100755 --- a/scripts/landscape_shards.sh +++ b/scripts/landscape_shards.sh @@ -5,7 +5,7 @@ # Usage: scripts/landscape_shards.sh [backend] # # Starts K processes: -# target/release/ddrs --workspace .ddrs experiment --backend --shard i/K +# target/release/ddrs --workspace $LANDSCAPE_WORKSPACE (default .ddrs) experiment --backend --shard i/K # for i in 0..K-1, each as its own `systemd-run --user --unit=ddrs--shard-i # --collect --quiet` unit. Each shard's stdout+stderr is logged to # /home/tbindas/.claude/jobs/d68b0903/tmp/.shard-i.log; once a shard's @@ -20,7 +20,8 @@ backend="${3:-cpu}" repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" ddrs_bin="${repo_root}/target/release/ddrs" -log_dir="/home/tbindas/.claude/jobs/d68b0903/tmp" +log_dir="${LANDSCAPE_LOG_DIR:-/home/tbindas/.claude/jobs/d68b0903/tmp}" +nice_level="${LANDSCAPE_NICE:-0}" mkdir -p "$log_dir" if [[ ! -x "$ddrs_bin" ]]; then @@ -37,9 +38,9 @@ for ((i = 0; i < K; i++)); do unit="ddrs-${bundle}-shard-${i}" log="${log_dir}/${bundle}.shard-${i}.log" : >"$log" - systemd-run --user --unit="$unit" --collect --quiet \ + systemd-run --user --unit="$unit" --collect --quiet --nice="${nice_level}" \ --working-directory="$repo_root" \ - -- bash -c "exec '${ddrs_bin}' --workspace .ddrs experiment '${bundle}' --backend '${backend}' --shard '${i}/${K}' >'${log}' 2>&1" + -- bash -c "exec '${ddrs_bin}' --workspace '${LANDSCAPE_WORKSPACE:-.ddrs}' experiment '${bundle}' --backend '${backend}' --shard '${i}/${K}' >'${log}' 2>&1" echo " shard ${i}/${K}: unit ${unit}, log ${log}" >&2 done diff --git a/src/experiment/landscape/mod.rs b/src/experiment/landscape/mod.rs index fc04f02d..9aead914 100644 --- a/src/experiment/landscape/mod.rs +++ b/src/experiment/landscape/mod.rs @@ -96,6 +96,14 @@ pub struct LandscapeSpec { /// of `axes`. Default empty (every axis a log multiplier). #[serde(default)] pub additive_axes: std::collections::HashMap, + /// Axis parameters to hold at their trained field even though the arm's + /// head learns them: a listed parameter keeps its slot in `axes` but is + /// inactive (never perturbed, no Newton step, dropped from the Hessian + /// eigendecomposition), exactly as a parameter the run did not learn. + /// Must be a subset of `axes`. Default empty: every learnable axis is + /// active, which is the behaviour of every earlier study. + #[serde(default)] + pub frozen: Vec, /// Also compute per-reach `g_i = dL/d ln x_i` (x in n, p, q) at α = 0 and /// α*, plus `dist_to_gauge_m`, for the "perturbations in a watershed" /// study: where in the network the gauge still constrains parameters, @@ -266,6 +274,11 @@ impl LandscapeSpec { return Err(format!("landscape additive_axes half-width for `{name}` must be > 0, got {hw}").into()); } } + for f in &self.frozen { + if !self.axes.contains(f) { + return Err(format!("landscape `frozen` entry `{f}` is not one of the axes {:?}", self.axes).into()); + } + } Ok(()) } @@ -550,7 +563,7 @@ where // Physical units per unit alpha for each additively swept slot; `None` // keeps the log-multiplier convention. See `LandscapeSpec::additive_axes`. let additive: [Option; 3] = std::array::from_fn(|k| spec.additive_axes.get(&spec.axes[k]).map(|hw| hw / spec.alpha_max)); - let obj = Objective::::build(ctx, &staid, starts, window_days, &spec.objective, spec.deriv_weight, &spec.axes, additive)?; + let obj = Objective::::build(ctx, &staid, starts, window_days, &spec.objective, spec.deriv_weight, &spec.axes, additive, &spec.frozen)?; let h = spec.fd_step; let active = obj.active(); let n_active_params = active.iter().filter(|&&a| a).count(); @@ -1084,6 +1097,17 @@ mod tests { assert!(b.validate_axes().is_err()); } + #[test] + fn frozen_defaults_empty_and_must_be_a_subset_of_axes() { + let d: LandscapeSpec = serde_yaml::from_str("gauges: {}\n").unwrap(); + assert!(d.frozen.is_empty()); + assert!(d.validate_axes().is_ok()); + let ok: LandscapeSpec = serde_yaml::from_str("gauges: {}\nfrozen: [p_spatial]\n").unwrap(); + assert!(ok.validate_axes().is_ok()); + let bad: LandscapeSpec = serde_yaml::from_str("gauges: {}\nfrozen: [gamma]\n").unwrap(); + assert!(bad.validate_axes().is_err()); + } + #[test] fn plane_skip_reason_names_the_fixed_parameter() { // p_spatial (index 1) fixed. diff --git a/src/experiment/landscape/objective.rs b/src/experiment/landscape/objective.rs index d95e44ac..c67ee6ff 100644 --- a/src/experiment/landscape/objective.rs +++ b/src/experiment/landscape/objective.rs @@ -181,6 +181,7 @@ where deriv_weight: f32, axes: &[String; 3], additive: [Option; 3], + frozen: &[String], ) -> Result { let device = &ctx.device; for a in axes { @@ -225,7 +226,7 @@ where `params.stage_roughness.gamma` is a constant, not a field to perturb." .into()); } - let active: [bool; 3] = std::array::from_fn(|k| learns(&axes[k])); + let active: [bool; 3] = std::array::from_fn(|k| axis_active(learns(&axes[k]), &axes[k], frozen)); // The routed daily series `forward_loss` compares `obs` against has // length `window_days - 2` regardless of `tau`: `n_hourly = // (window_days - 1) * 24` (`RhoWindow::n_hourly`, @@ -737,6 +738,13 @@ fn has_valid_observation(obs: &[f32], warmup: usize, d: usize) -> bool { (warmup..d).any(|i| i < obs.len() && obs[i].is_finite() && obs[i] >= 0.0) } +/// Whether an axis slot is swept: the head learns it and it is not listed in +/// `LandscapeSpec::frozen`. A frozen learnable parameter keeps its trained +/// per-reach field (`trained_field` still reads the head output). +fn axis_active(learnable: bool, name: &str, frozen: &[String]) -> bool { + learnable && !frozen.iter().any(|f| f == name) +} + /// Trained physical field for one channel parameter, length `n_active`. /// `Some(head_out)` denormalizes the KAN head's `[0,1]` output to the /// physical range (the parameter is in `kan_head.learnable_parameters`); @@ -996,6 +1004,16 @@ pub fn solve3(a: [[f32; 3]; 3], b: [f32; 3]) -> Option<[f32; 3]> { #[cfg(test)] mod tests { + #[test] + fn axis_active_requires_learnable_and_not_frozen() { + let none: Vec = vec![]; + let frozen = vec!["p_spatial".to_string()]; + assert!(axis_active(true, "p_spatial", &none)); + assert!(!axis_active(false, "p_spatial", &none)); + assert!(!axis_active(true, "p_spatial", &frozen)); + assert!(axis_active(true, "n", &frozen)); + } + use super::*; use burn::backend::NdArray; use std::collections::HashMap;