bundle_v1/scripts/claim5_analyze.py

"""Claim 5 — analyze GPU-grid results vs paper Table 1.

Reads per-run CSVs (cifar_{opt}_lam{lam}_seed{s}.csv) from the grid, computes
AUC (mean over epochs) per metric, aggregates across seeds, renders the
comparison table + training curves (paper Fig. 8 analogue), and checks the
paper's directional claims per optimizer:
  acc AUC up | sigreg AUC down | rank AUC up | dormancy AUC down.
"""
import argparse
import csv
import glob
import json
import os
import re
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(__file__))
from plot_style import SERIES, apply_style

import matplotlib.pyplot as plt

PAPER = {  # Table 1 (arXiv:2602.19373v3, App. C.1)
    ("adam", 0.0): dict(acc=33.0, sig=49.1, rank=149.0, dorm=18.5),
    ("adam", 1.0): dict(acc=50.0, sig=7.6, rank=262.9, dorm=0.4),
    ("radam", 0.0): dict(acc=29.8, sig=42.3, rank=148.8, dorm=12.6),
    ("radam", 1.0): dict(acc=42.8, sig=8.4, rank=251.8, dorm=0.6),
    ("kron", 0.0): dict(acc=62.6, sig=28.5, rank=167.4, dorm=7.9),
    ("kron", 1.0): dict(acc=72.7, sig=7.2, rank=252.2, dorm=0.9),
}


def load_runs(out_dir):
    runs = {}
    for path in sorted(glob.glob(os.path.join(out_dir, "cifar_*.csv"))):
        m = re.match(r"cifar_(\w+)_lam([\d.]+)_seed(\d+)\.csv",
                     os.path.basename(path))
        if not m:
            continue
        opt, lam, seed = m.group(1), float(m.group(2)), int(m.group(3))
        with open(path) as f:
            rows = list(csv.DictReader(f))
        runs[(opt, lam, seed)] = rows
    return runs


def auc(rows, key, scale=1.0):
    return float(np.mean([float(r[key]) for r in rows])) * scale


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out_dir", default="repro_isogaussian_drl/outputs/claim5_gpu")
    ap.add_argument("--plot_dir", default="repro_isogaussian_drl/outputs/claim5")
    ap.add_argument("--dorm_key", default="dormant_mlp",
                    choices=["dormant_mlp", "dormant_all", "dormant_cnn"])
    args = ap.parse_args()
    os.makedirs(args.plot_dir, exist_ok=True)
    apply_style()

    runs = load_runs(args.out_dir)
    if not runs:
        print("no runs found"); sys.exit(1)
    epochs = max(len(rows) for rows in runs.values())
    print(f"loaded {len(runs)} runs, {epochs} epochs")

    # per-(opt, lam): aggregate AUCs across seeds
    agg = {}
    for (opt, lam, seed), rows in runs.items():
        # Paper Table-1 units: every column is 100 x the time-averaged raw
        # metric (verified against Figs. 3/8 axis ranges) — including feature
        # rank, i.e. "149.0" = raw RankMe 1.49. Scale rank by 100 to match.
        a = dict(acc=auc(rows, "train_acc", 100), sig=auc(rows, "train_sigreg", 100),
                 rank=auc(rows, "rankme", 100), dorm=auc(rows, args.dorm_key, 100))
        agg.setdefault((opt, lam), []).append((seed, a))

    table_rows = []
    checks = []
    for opt in ("adam", "radam", "kron"):
        for lam in (0.0, 1.0):
            if (opt, lam) not in agg:
                continue
            seeds = agg[(opt, lam)]
            mean = {k: float(np.mean([a[k] for _, a in seeds])) for k in seeds[0][1]}
            spread = {k: float(np.ptp([a[k] for _, a in seeds])) for k in seeds[0][1]}
            p = PAPER[(opt, lam)]
            table_rows.append(dict(optimizer=opt, lam=lam, n_seeds=len(seeds),
                                   **{f"{k}_mean": v for k, v in mean.items()},
                                   **{f"{k}_range": v for k, v in spread.items()},
                                   **{f"{k}_paper": p[k] for k in p}))
        if (opt, 0.0) in agg and (opt, 1.0) in agg:
            m0 = {k: float(np.mean([a[k] for _, a in agg[(opt, 0.0)]])) for k in "acc sig rank dorm".split()}
            m1 = {k: float(np.mean([a[k] for _, a in agg[(opt, 1.0)]])) for k in "acc sig rank dorm".split()}
            ok = (m1["acc"] > m0["acc"] and m1["sig"] < m0["sig"]
                  and m1["rank"] > m0["rank"] and m1["dorm"] < m0["dorm"])
            checks.append((opt, ok, m0, m1))
            print(f"[{opt}] SIGReg effect: acc {m0['acc']:.1f}->{m1['acc']:.1f} | "
                  f"sig {m0['sig']:.1f}->{m1['sig']:.1f} | rank {m0['rank']:.1f}->"
                  f"{m1['rank']:.1f} | dorm {m0['dorm']:.1f}->{m1['dorm']:.1f} "
                  f"-> {'PASS (all 4 directions match Table 1)' if ok else 'MISMATCH'}")

    with open(os.path.join(args.plot_dir, "table1_comparison.csv"), "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(table_rows[0]))
        w.writeheader(); w.writerows(table_rows)

    # Fig. 8 analogue: curves per optimizer (seed-averaged)
    metrics = [("train_acc", "Train Accuracy", 1.0), ("train_sigreg", "SIGReg Loss", 1.0),
               ("rankme", "Feature Rank (RankMe)", 1.0), (args.dorm_key, "Dormant Neurons [%]", 100.0)]
    fig, axes = plt.subplots(1, 4, figsize=(16, 3.2))
    colors = {("adam", 0.0): SERIES[0], ("adam", 1.0): SERIES[0],
              ("radam", 0.0): SERIES[3], ("radam", 1.0): SERIES[3],
              ("kron", 0.0): SERIES[5], ("kron", 1.0): SERIES[5]}
    for ax, (key, label, scale) in zip(axes, metrics):
        for (opt, lam), seeds in sorted(agg.items()):
            series = []
            for seed, _ in seeds:
                rows = runs[(opt, lam, seed)]
                series.append([float(r[key]) * scale for r in rows])
            n = min(map(len, series))
            y = np.mean([s[:n] for s in series], axis=0)
            ax.plot(np.arange(n), y, color=colors[(opt, lam)],
                    linestyle="-" if lam > 0 else ":",
                    linewidth=1.8,
                    label=f"{opt}{'+SIGReg' if lam > 0 else ''}")
        for e in (20, 40, 60, 80):
            ax.axvline(e, color="#bbb", linewidth=0.7, linestyle="--")
        ax.set_xlabel("Epoch"); ax.set_title(label)
    axes[0].legend(fontsize=7, ncols=2)
    fig.suptitle("Non-stationary CIFAR-10 (labels reshuffled every 20 epochs) — reproduction", y=1.04)
    fig.tight_layout()
    fig.savefig(os.path.join(args.plot_dir, "claim5_curves.png"), bbox_inches="tight")

    ok = all(c[1] for c in checks) and len(checks) == 3
    with open(os.path.join(args.plot_dir, "claim5_verdict.json"), "w") as f:
        json.dump({"per_optimizer_pass": {c[0]: c[1] for c in checks}}, f, indent=2)
    print("CLAIM 5 DIRECTIONAL CHECK:", "PASS" if ok else "PARTIAL/FAIL")
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()