chore: capture initial SDIL project state

1 files changed, 191 insertions, 0 deletions

diff --git a/experiments/analyze.py b/experiments/analyze.py
new file mode 100644
index 0000000..67b6f04
--- /dev/null
+++ b/experiments/analyze.py
@@ -0,0 +1,191 @@
+"""Aggregate results/*.json into summary tables (+ figures if matplotlib is
+present). Run on any box with access to the shared NFS home.
+Usage: python experiments/analyze.py [results_dir]
+"""
+import glob
+import json
+import os
+import sys
+from collections import defaultdict
+
+RES = sys.argv[1] if len(sys.argv) > 1 else "results"
+
+try:
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+    HAVE_MPL = True
+except Exception:
+    HAVE_MPL = False
+
+
+def load_all(res):
+    runs = []
+    for p in sorted(glob.glob(os.path.join(res, "*.json"))):
+        try:
+            with open(p) as f:
+                d = json.load(f)
+            d["_path"] = p
+            d["_tag"] = os.path.basename(p)[:-5]
+            runs.append(d)
+        except Exception as e:
+            print("skip", p, e)
+    return runs
+
+
+def mean_final_cos(run):
+    f = run.get("final", {})
+    c = f.get("cos_r_negg")
+    return sum(c) / len(c) if c else None
+
+
+def group(runs, keyfn):
+    g = defaultdict(list)
+    for r in runs:
+        k = keyfn(r)
+        if k is not None:
+            g[k].append(r)
+    return g
+
+
+def agg(vals):
+    vals = [v for v in vals if v is not None]
+    if not vals:
+        return (None, None, 0)
+    m = sum(vals) / len(vals)
+    sd = (sum((v - m) ** 2 for v in vals) / len(vals)) ** 0.5
+    return (m, sd, len(vals))
+
+
+def fmt(m, sd, n):
+    if m is None:
+        return "   n/a  "
+    return f"{m:.3f}±{sd:.3f}(n{n})"
+
+
+def table_main(runs):
+    print("\n================ MAIN COMPARISON (mnist depth3 width256) ================")
+    sel = [r for r in runs if r["args"].get("depth") == 3 and r["args"].get("width") == 256
+           and r["args"].get("dataset") == "mnist" and r["args"].get("nuis_rho", 0) == 0
+           and r["_tag"].startswith(("main_", "A_main_"))]
+    by_mode = group(sel, lambda r: r["args"]["mode"])
+    print(f"{'mode':6s} {'test_acc':22s} {'mean_cos(r,-g)':22s}")
+    for mode in ["bp", "dfa", "sdil"]:
+        rs = by_mode.get(mode, [])
+        acc = agg([r["final"].get("test_acc") for r in rs])
+        cos = agg([mean_final_cos(r) for r in rs])
+        print(f"{mode:6s} {fmt(*acc):22s} {fmt(*cos):22s}")
+
+
+def table_nuisance(runs):
+    print("\n================ RESIDUALIZATION UNDER NUISANCE ================")
+    print("(sdil; does subtracting the soma-predictable baseline preserve alignment as rho grows?)")
+    sel = [r for r in runs if r["_tag"].startswith(("nuis_", "A_nuis_"))]
+    by = group(sel, lambda r: (r["args"].get("nuis_rho"), r["args"].get("use_residual")))
+    rhos = sorted({k[0] for k in by})
+    print(f"{'rho':>5s} | {'residual=1 acc':22s} {'cos':16s} | {'residual=0 acc':22s} {'cos':16s}")
+    for rho in rhos:
+        r1 = by.get((rho, 1), [])
+        r0 = by.get((rho, 0), [])
+        a1 = agg([r["final"].get("test_acc") for r in r1]); c1 = agg([mean_final_cos(r) for r in r1])
+        a0 = agg([r["final"].get("test_acc") for r in r0]); c0 = agg([mean_final_cos(r) for r in r0])
+        print(f"{rho:5.1f} | {fmt(*a1):22s} {fmt(*c1):16s} | {fmt(*a0):22s} {fmt(*c0):16s}")
+
+
+def table_trainedA(runs):
+    print("\n================ TRAINED vs FIXED APICAL PATHWAY ================")
+    fixed = [r for r in runs if "fixedA" in r["_tag"]]
+    af = agg([r["final"].get("test_acc") for r in fixed]); cf = agg([mean_final_cos(r) for r in fixed])
+    print(f"{'fixed-A (DFA-like)':22s} acc {fmt(*af):22s} cos {fmt(*cf):16s}")
+    tr = [r for r in runs if "trainedA_nd" in r["_tag"]]
+    by = group(tr, lambda r: r["args"].get("pert_ndirs"))
+    for nd in sorted(by):
+        rs = by[nd]
+        a = agg([r["final"].get("test_acc") for r in rs]); c = agg([mean_final_cos(r) for r in rs])
+        print(f"trained-A ndirs={nd:<3d}     acc {fmt(*a):22s} cos {fmt(*c):16s}")
+
+
+def table_depth(runs):
+    print("\n================ DEPTH SCALING ================")
+    sel = [r for r in runs if r["_tag"].startswith("depth_")]
+    if not sel:
+        return
+    by = group(sel, lambda r: (r["args"]["mode"], r["args"]["depth"]))
+    depths = sorted({k[1] for k in by})
+    print(f"{'depth':>6s} | " + " | ".join(f"{m:>26s}" for m in ["bp", "dfa", "sdil"]))
+    for d in depths:
+        cells = []
+        for m in ["bp", "dfa", "sdil"]:
+            rs = by.get((m, d), [])
+            a = agg([r["final"].get("test_acc") for r in rs])
+            c = agg([mean_final_cos(r) for r in rs])
+            cells.append(f"a{fmt(*a)[:5]} c{('%.2f'%c[0]) if c[0] is not None else ' n/a'}")
+        print(f"{d:6d} | " + " | ".join(f"{c:>26s}" for c in cells))
+
+
+def plots(runs):
+    if not HAVE_MPL:
+        print("\n[matplotlib not available -> skipping figures]")
+        return
+    figdir = os.path.join(RES, "figs")
+    os.makedirs(figdir, exist_ok=True)
+
+    # Fig 1: alignment trajectory sdil vs dfa (seed0 main)
+    plt.figure(figsize=(6, 4))
+    for tag, lab in [("A_main_sdil_s0", "SDIL (trained A)"), ("A_main_dfa_s0", "DFA (fixed A)"),
+                     ("main_sdil_mnist_s0", "SDIL"), ("main_dfa_mnist_s0", "DFA")]:
+        r = next((x for x in runs if x["_tag"] == tag), None)
+        if not r:
+            continue
+        xs = [s["step"] for s in r["steps"] if "cos_r_negg" in s]
+        ys = [sum(s["cos_r_negg"]) / len(s["cos_r_negg"]) for s in r["steps"] if "cos_r_negg" in s]
+        if xs:
+            plt.plot(xs, ys, label=lab)
+    plt.xlabel("step"); plt.ylabel("mean cos(r, -grad)"); plt.title("Gradient alignment over training")
+    plt.legend(); plt.grid(alpha=.3); plt.tight_layout()
+    plt.savefig(os.path.join(figdir, "alignment_traj.png"), dpi=120); plt.close()
+
+    # Fig 2: acc & alignment vs rho (residual vs not)
+    sel = [r for r in runs if r["_tag"].startswith(("nuis_", "A_nuis_"))]
+    if sel:
+        by = group(sel, lambda r: (r["args"].get("nuis_rho"), r["args"].get("use_residual")))
+        rhos = sorted({k[0] for k in by})
+        fig, ax = plt.subplots(1, 2, figsize=(10, 4))
+        for ur, lab in [(1, "residual (SDIL)"), (0, "raw apical")]:
+            accs = [agg([r["final"].get("test_acc") for r in by.get((rho, ur), [])])[0] for rho in rhos]
+            coss = [agg([mean_final_cos(r) for r in by.get((rho, ur), [])])[0] for rho in rhos]
+            ax[0].plot(rhos, accs, marker="o", label=lab)
+            ax[1].plot(rhos, coss, marker="o", label=lab)
+        ax[0].set_xlabel("nuisance rho"); ax[0].set_ylabel("test acc"); ax[0].set_title("Accuracy vs nuisance")
+        ax[1].set_xlabel("nuisance rho"); ax[1].set_ylabel("mean cos(r,-grad)"); ax[1].set_title("Alignment vs nuisance")
+        for a in ax:
+            a.legend(); a.grid(alpha=.3)
+        plt.tight_layout(); plt.savefig(os.path.join(figdir, "nuisance.png"), dpi=120); plt.close()
+
+    # Fig 3: alignment vs depth
+    sel = [r for r in runs if r["_tag"].startswith("depth_")]
+    if sel:
+        by = group(sel, lambda r: (r["args"]["mode"], r["args"]["depth"]))
+        depths = sorted({k[1] for k in by})
+        plt.figure(figsize=(6, 4))
+        for m in ["dfa", "sdil"]:
+            ys = [agg([mean_final_cos(r) for r in by.get((m, d), [])])[0] for d in depths]
+            plt.plot(depths, ys, marker="o", label=m)
+        plt.xlabel("hidden depth"); plt.ylabel("mean cos(r,-grad)"); plt.title("Alignment vs depth")
+        plt.legend(); plt.grid(alpha=.3); plt.tight_layout()
+        plt.savefig(os.path.join(figdir, "depth.png"), dpi=120); plt.close()
+    print(f"\n[figures written to {figdir}]")
+
+
+def main():
+    runs = load_all(RES)
+    print(f"loaded {len(runs)} runs from {RES}")
+    table_main(runs)
+    table_nuisance(runs)
+    table_trainedA(runs)
+    table_depth(runs)
+    plots(runs)
+
+
+if __name__ == "__main__":
+    main()