Paper experiments: consolidated code (Stage 1/2 landscape, E-07/09/10/11, diversity, blend)

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 \
+        <out-dir> [<out-dir> ...]   # 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()

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).
+#   <worktree>/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

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_<member>        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 <dir>
+"""
+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="<f4").reshape(len(gauges), N_DAYS).astype(np.float64)
+        rel = np.abs(S[lab] - ref) / np.maximum(np.abs(ref), 1e-3)
+        check[lab] = {"run": rid, "max_rel": float(rel.max()), "p99_rel": float(np.quantile(rel, 0.99))}
+        print(f"  {lab} raw sum vs {rid} baseline: max rel diff {rel.max():.2e}, p99 {np.quantile(rel, .99):.2e}")
+
+    obs = np.fromfile(RUNS / MEMBERS[0][3] / "baseline" / "observations.f32", dtype="<f4").reshape(len(gauges), N_DAYS)
+    np.savez_compressed(out / "blend_inflow.npz", gauges=np.array(gauges), start=str(START), n_days=N_DAYS,
+                        S_blend=S_blend.astype(np.float32), obs=obs,
+                        full_cover=full_cover, n_up=n_up,
+                        **{f"S_{k}": v.astype(np.float32) for k, v in S.items()})
+    allc_comids = order[allc]
+    np.save(out / "comids_in_all_three_stores.npy", allc_comids)
+    (out / "blend_inflow.json").write_text(json.dumps({
+        "route": "exact (icechunk stores read directly, blend.rs renormalisation)",
+        "window": [str(START), str(START + N_DAYS - 1)], "n_gauges": len(gauges),
+        "n_full_cover_gauges": int(full_cover.sum()), "n_conus_reaches_all_three": int(allc.sum()),
+        "n_conus_reaches_none": int(none.sum()), "nan_counts": nan_counts, "baseline_check": check,
+        "members": [{"label": l, "path": p, "weight": w_, "run": r} for l, p, w_, r in MEMBERS],
+    }, indent=2))
+    print(f"wrote {out}/blend_inflow.npz")
+
+
+if __name__ == "__main__":
+    main()