plot + summary generator

make_report.py

@@ -1,15 +1,10 @@
-"""Aggregate eval JSONs into plots, a summary CSV, and a drafted report.
-
-Reads results/*_eval.json (and out/*_history.json if present), writes:
-  plots/acc_vs_k.png        accuracy vs hop count k (1..64) per family
-  plots/acc_vs_iters.png    accuracy vs inference-time unroll count
-  plots/acc_vs_state.png    accuracy vs state-space size
-  plots/training.png        in-distribution accuracy vs training step
-  results/summary.csv       one row per run with budget + headline accuracies
-  results/REPORT_DRAFT.md   tables filled from the numbers (prose written by
-                            the experimenter afterwards)
-
-Usage: python make_report.py --results-dir results [--hist-dir out]
+"""Generate figures and the results tables from the evaluation JSONs.
+
+Reads results/*_eval.json and results/*_history.json (downloaded from the Hub
+repo), writes plots to results/plots/ and results/summary_table.md.
+All numbers reported are empirical; no claims beyond the data.
+
+Usage: python make_report.py [--results-dir results] [--out-dir results/plots]
 """
 
 import argparse
@@ -18,218 +13,206 @@ import json
 import os
 from collections import defaultdict
 
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
 import numpy as np
 
-FAMILIES = {
-    "fixed4x1": "fixed-4 (4P)",
-    "rec2x2": "rec-2x2 (2P)",
-    "rec1x4": "rec-1x4 (1P)",
-    "rec1xdynamic": "rec-1xdyn (1P)",
-    "rec4x4": "rec-4x4 (iso-param, 4x flops/step)",
-    "fixed4-long": "fixed-4-long (iso-flop, 4x steps)",
+FAMILY = {
+    "fixed4x1": ("fixed-4 (matched FLOPs)", "#d62728"),
+    "rec1x4": ("rec 1x4 (matched FLOPs)", "#1f77b4"),
+    "rec2x2": ("rec 2x2 (matched FLOPs)", "#2ca02c"),
+    "rec1xdynamic": ("rec 1xdynamic R~U{2..8}", "#17becf"),
+    "rec4x4": ("rec 4x4 (iso-param with fixed-4)", "#9467bd"),
+    "fixed4-long": ("fixed-4 long (iso-FLOP with rec4x4)", "#ff7f0e"),
 }
 
 
-def family_of(name):
-    for fam in FAMILIES:
-        if name.startswith(fam):
-            return fam
+def load(results_dir):
+    runs = {}
+    for p in sorted(glob.glob(os.path.join(results_dir, "*_eval.json"))):
+        with open(p) as fh:
+            ev = json.load(fh)
+        name = ev["run_name"]
+        base = name.rsplit("-seed", 1)[0]
+        seed = int(name.rsplit("-seed", 1)[1]) if "-seed" in name else -1
+        runs.setdefault(base, {})[seed] = ev
+    return runs
+
+
+def mean_over_seeds(runs, base, get):
+    seeds = runs.get(base, {})
+    if not seeds:
+        return None, None, 0
+    vals = np.array([[get(ev) for ev in seeds[s].values()] for s in sorted(seeds)])
+    # get applied per eval dict -> we actually flatten manually below
     return None
 
 
-def load_evals(results_dir):
-    evals = []
-    for path in sorted(glob.glob(os.path.join(results_dir, "*_eval.json"))):
-        with open(path) as fh:
-            evals.append(json.load(fh))
-    return evals
-
+def stack(runs, base, key):
+    """values across seeds for a dict-valued metric, e.g. acc_vs_k."""
+    seeds = sorted(runs.get(base, {}))
+    if not seeds:
+        return None, None
+    ks = list(runs[base][seeds[0]][key].keys())
+    arr = np.array([[runs[base][s][key][k] for k in ks] for s in seeds])
+    return ks, arr  # arr shape [n_seeds, n_ks]
 
-def acc_vs_k_plot(evals, out):
-    import matplotlib
-    matplotlib.use("Agg")
-    import matplotlib.pyplot as plt
 
+def plot_acc_vs_k(runs, out_dir):
     fig, ax = plt.subplots(figsize=(7, 4.5))
-    by_fam = defaultdict(lambda: defaultdict(list))
-    for ev in evals:
-        fam = family_of(ev["run_name"])
-        if fam is None:
+    rows = []
+    for base, (label, color) in FAMILY.items():
+        ks, arr = stack(runs, base, "acc_vs_k")
+        if arr is None:
             continue
-        for k, v in ev["acc_vs_k"].items():
-            by_fam[fam][int(k)].append(v)
-    for fam, per_k in by_fam.items():
-        ks = sorted(per_k)
-        mean = np.array([np.mean(per_k[k]) for k in ks])
-        std = np.array([np.std(per_k[k]) for k in ks])
-        ax.plot(ks, mean, label=FAMILIES[fam], lw=1.8)
-        ax.fill_between(ks, mean - std, mean + std, alpha=0.15)
-    ax.axhline(1 / 64, color="gray", ls=":", lw=1, label="chance (1/64)")
-    ax.axvline(8, color="gray", ls="--", lw=1, label="max train k")
-    ax.set_xlabel("hop count k (f applied k times)")
+        m, sd = arr.mean(0), arr.std(0)
+        x = [int(k) for k in ks]
+        ax.plot(x, m, label=label, color=color)
+        ax.fill_between(x, m - sd, m + sd, color=color, alpha=0.18)
+        rows.append((label, m))
+    ax.axhline(1 / 64, ls=":", c="gray", lw=1)
+    ax.text(45, 1 / 64 + 0.008, "chance (1/64)", fontsize=8, color="gray")
+    ax.set_xscale("log")
+    ax.set_xticks([1, 2, 4, 8, 16, 32, 64])
+    ax.set_xticklabels([1, 2, 4, 8, 16, 32, 64])
+    ax.set_xlabel("hop count k (chain length)")
     ax.set_ylabel("exact-match accuracy")
-    ax.set_title("OOD chain length (trained on k=1..8, fresh random f)")
     ax.set_ylim(0, 1.02)
-    ax.legend(fontsize=7)
+    ax.set_title("Accuracy vs chain length (k), trained on k<=8, S=64")
+    ax.legend(fontsize=7, loc="upper right")
+    ax.grid(alpha=0.3)
     fig.tight_layout()
-    fig.savefig(out, dpi=150)
-
+    fig.savefig(os.path.join(out_dir, "acc_vs_k.png"), dpi=150)
+    plt.close(fig)
+    return rows
 
-def acc_vs_iters_plot(evals, out):
-    import matplotlib
-    matplotlib.use("Agg")
-    import matplotlib.pyplot as plt
 
+def plot_acc_vs_iters(runs, out_dir):
     fig, ax = plt.subplots(figsize=(7, 4.5))
-    by_fam = defaultdict(lambda: defaultdict(list))
-    for ev in evals:
-        fam = family_of(ev["run_name"])
-        if fam is None:
+    colors = {8: "#1f77b4", 16: "#ff7f0e", 32: "#2ca02c", 64: "#d62728"}
+    for base in ("rec1x4", "rec1xdynamic", "rec4x4", "fixed4x1", "fixed4-long"):
+        ks, arr = stack(runs, base, "acc_vs_iters")
+        if arr is None:
             continue
-        for r, per_k in ev["acc_vs_iters"].items():
-            for k, v in per_k.items():
-                if int(k) in (8, 16, 32, 64):
-                    by_fam[(fam, int(k))][int(r)].append(v)
-    palette = {}
-    for (fam, k), per_r in sorted(by_fam.items()):
-        rs = sorted(per_r)
-        mean = np.array([np.mean(per_r[r]) for r in rs])
-        lbl = f"{FAMILIES[fam].split(' (')[0]}, k={k}"
-        color = palette.setdefault(fam, None)
-        line, = ax.plot(rs, mean, lw=1.4, label=lbl)
-        palette[fam] = line.get_color()
-    ax.set_xlabel("inference-time unrolls R (trained R: fixed-4=1, rec models=4 or dyn 2..8)")
+        # arr[seed, R, k] flattened dict: acc_vs_iters is {R: {k: acc}}
+        Rs = sorted(int(r) for r in runs[base][0]["acc_vs_iters"])
+        for ki, k in enumerate(("8", "16", "32", "64")):
+            if k not in runs[base][0]["acc_vs_iters"]["1"]:
+                continue
+            m = np.array([[runs[base][s]["acc_vs_iters"][str(r)][k]
+                           for r in Rs] for s in sorted(runs[base])]).mean(0)
+            label = f"{FAMILY.get(base, (base,))[0]} k={k}"
+            ax.plot(Rs, m, label=label,
+                    color=colors[k], ls="--" if base.startswith("fixed") else "-",
+                    alpha=0.9 if base in ("rec1x4", "rec4x4") else 0.5)
+    ax.set_xlabel("inference-time unroll count R")
     ax.set_ylabel("exact-match accuracy")
-    ax.set_title("OOD inference-time recurrence (more/fewer unrolls than trained)")
-    ax.legend(fontsize=6)
+    ax.set_xticks(range(1, 17))
+    ax.set_title("Accuracy vs inference-time recurrence count (OOD iterations)")
+    ax.legend(fontsize=6, ncol=2)
+    ax.grid(alpha=0.3)
     fig.tight_layout()
-    fig.savefig(out, dpi=150)
-
+    fig.savefig(os.path.join(out_dir, "acc_vs_iters.png"), dpi=150)
+    plt.close(fig)
 
-def acc_vs_state_plot(evals, out):
-    import matplotlib
-    matplotlib.use("Agg")
-    import matplotlib.pyplot as plt
 
-    fig, ax = plt.subplots(figsize=(6, 4.5))
-    markers = {8: "o", 32: "s"}
-    by_fam = defaultdict(list)
-    for ev in evals:
-        fam = family_of(ev["run_name"])
-        if fam is None:
+def plot_state_size(runs, out_dir):
+    fig, axes = plt.subplots(1, 3, figsize=(12, 3.8), sharey=True)
+    for ax, base in zip(axes, ("fixed4x1", "rec1x4", "rec1xdynamic")):
+        found = runs.get(base)
+        if not found:
             continue
-        row = {}
-        for Sstr, per_k in ev["acc_vs_state_size"].items():
-            row[int(Sstr)] = {int(k): v for k, v in per_k.items()}
-        by_fam[fam].append(row)
-    for fam, rows in by_fam.items():
-        for k, mk in markers.items():
-            xs, ys = [], []
-            for row in rows:
-                for S, per_k in sorted(row.items()):
-                    if k in per_k:
-                        xs.append(S)
-                        ys.append(per_k[k])
-            if xs:
-                ax.plot(xs, ys, mk + "-", label=f"{FAMILIES[fam].split(' (')[0]} k={k}",
-                        alpha=0.8, lw=1.3, markersize=4)
-    ax.set_xlabel("state-space size S' (trained S=64)")
-    ax.set_ylabel("exact-match accuracy")
-    ax.set_title("OOD state-space size")
-    ax.legend(fontsize=6)
+        ev = found[sorted(found)[0]]
+        for Sp, style in (("16", "--"), ("32", "-."), ("64", "-")):
+            if Sp not in ev["acc_vs_state_size"]:
+                continue
+            d = ev["acc_vs_state_size"][Sp]
+            ks = [int(k) for k in d]
+            ax.plot(ks, [d[k] for k in d], style, label=f"S'={Sp}")
+        ax.set_xscale("log")
+        ax.set_xticks([1, 2, 4, 8, 16, 32, 64])
+        ax.set_xticklabels([1, 2, 4, 8, 16, 32, 64])
+        ax.set_xlabel("k")
+        ax.set_title(FAMILY.get(base, (base,))[0], fontsize=9)
+        ax.grid(alpha=0.3)
+    axes[0].set_ylabel("exact-match accuracy")
+    axes[0].legend(fontsize=8)
+    fig.suptitle("Accuracy vs state-space size (seed 0 shown)", fontsize=10)
     fig.tight_layout()
-    fig.savefig(out, dpi=150)
+    fig.savefig(os.path.join(out_dir, "acc_vs_state_size.png"), dpi=150)
+    plt.close(fig)
 
 
-def training_plot(evals, hist_dir, out):
-    import matplotlib
-    matplotlib.use("Agg")
-    import matplotlib.pyplot as plt
-
+def plot_training(runs, hist_dir, out_dir):
     fig, ax = plt.subplots(figsize=(7, 4.5))
-    by_fam = defaultdict(lambda: defaultdict(list))
-    for path in sorted(glob.glob(os.path.join(hist_dir, "*_history.json"))):
-        name = os.path.basename(path).replace("_history.json", "")
-        fam = family_of(name)
-        if fam is None:
+    for p in sorted(glob.glob(os.path.join(hist_dir, "*_history.json"))):
+        with open(p) as fh:
+            h = json.load(fh)
+        name = os.path.basename(p).replace("_history.json", "")
+        base = name.rsplit("-seed", 1)[0]
+        if base not in FAMILY:
             continue
-        h = json.load(open(path))
-        for rec in h["id_acc"]:
-            by_fam[fam][rec["step"]].append(rec["mean"])
-    for fam, per_step in by_fam.items():
-        steps = sorted(per_step)
-        mean = np.array([np.mean(per_step[s]) for s in steps])
-        ax.plot(steps, mean, label=FAMILIES[fam], lw=1.8)
+        steps = [r["step"] for r in h["id_acc"]]
+        ax.plot(steps, [r["mean"] for r in h["id_acc"]],
+                label=name, color=FAMILY[base][1], alpha=0.7)
     ax.set_xlabel("training step")
-    ax.set_ylabel("mean id accuracy (k=1..8)")
-    ax.set_title("Training progress")
-    ax.legend(fontsize=7)
+    ax.set_ylabel("in-distribution acc (k=1..8 mean)")
+    ax.set_title("In-distribution accuracy during training")
+    ax.legend(fontsize=6)
+    ax.grid(alpha=0.3)
     fig.tight_layout()
-    fig.savefig(out, dpi=150)
+    fig.savefig(os.path.join(out_dir, "training_id_acc.png"), dpi=150)
+    plt.close(fig)
 
 
-def summary_csv(evals, out):
-    import csv
-    rows = []
-    for ev in evals:
-        fam = family_of(ev["run_name"])
-        if fam is None:
+def summary_table(runs, out_dir):
+    lines = ["| model | params | train FLOPs (est.) | seeds | in-dist acc "
+             "k=1..8 | acc k=16 | acc k=32 | acc k=64 | best R@k=64 |",
+             "|---|---|---|---|---|---|---|---|---|"]
+    for base, (label, _) in FAMILY.items():
+        seeds = sorted(runs.get(base, {}))
+        if not seeds:
             continue
-        ak = {int(k): v for k, v in ev["acc_vs_k"].items()}
-        idm = float(np.mean([ak[k] for k in range(1, 9)]))
-        best_r, best64 = max(
-            ((int(r), per_k["64"]) for r, per_k in ev["acc_vs_iters"].items()),
-            key=lambda t: t[1])
-        rows.append({
-            "run_name": ev["run_name"], "family": fam,
-            "params_total": ev["budget"]["params_total"],
-            "train_flops_total": ev["budget"]["train_flops_total"],
-            "id_acc_k1_8": round(idm, 4),
-            "acc_k16": round(ak.get(16, float("nan")), 4),
-            "acc_k32": round(ak.get(32, float("nan")), 4),
-            "acc_k64": round(ak.get(64, float("nan")), 4),
-            "best_acc_k64_over_R": round(best64, 4),
-            "best_R_for_k64": best_r,
-            "state16_acc_k32": round(
-                ev["acc_vs_state_size"].get("16", {}).get("32", float("nan")), 4),
-        })
-    with open(out, "w", newline="") as fh:
-        w = csv.DictWriter(fh, fieldnames=list(rows[0].keys()))
-        w.writeheader()
-        w.writerows(rows)
-    return rows
-
-
-def draft_report(evals, rows, out):
-    import csv as _csv
-    from io import StringIO
-    buf = StringIO()
-    w = _csv.DictWriter(buf, fieldnames=list(rows[0].keys()))
-    w.writeheader()
-    w.writerows(rows)
-    md = ["# Results tables (auto-generated draft)",
-          "",
-          "Numbers come from `results/*_eval.json`; interpretation and the",
-          "empirical-vs-theoretical distinction are written up in README.md.",
-          "",
-          "## Summary table",
-          "", "```csv", buf.getvalue().strip(), "```", ""]
-    with open(out, "w") as fh:
-        fh.write("\n".join(md))
+        ev0 = runs[base][seeds[0]]
+        idm = np.mean([[ev["acc_vs_k"][str(k)] for k in range(1, 9)]
+                       for s in seeds for ev in [runs[base][s]]])
+        ood = {k: np.mean([runs[base][s]["acc_vs_k"][str(k)] for s in seeds])
+               for k in (16, 32, 64)}
+        best_r = None
+        try:
+            per_r = {int(r): np.mean([runs[base][s]["acc_vs_iters"][r]["64"]
+                                      for s in seeds]) for r in
+                     ev0["acc_vs_iters"]}
+            best_r = max(per_r, key=per_r.get)
+        except Exception:
+            pass
+        lines.append(
+            f"| {label} | {ev0['budget']['params_total']:,} | "
+            f"{ev0['budget']['train_flops_total']:.2e} | {len(seeds)} | "
+            f"{idm:.3f} | {ood[16]:.3f} | {ood[32]:.3f} | {ood[64]:.3f} | "
+            f"R={best_r} |")
+    txt = "\n".join(lines)
+    with open(os.path.join(out_dir, "..", "summary_table.md"), "w") as fh:
+        fh.write(txt + "\n")
+    print(txt)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--results-dir", default="results")
+    ap.add_argument("--out-dir", default="results/plots")
+    a = ap.parse_args()
+    os.makedirs(a.out_dir, exist_ok=True)
+    runs = load(a.results_dir)
+    print("found families:", {b: list(s) for b, s in runs.items()})
+    plot_acc_vs_k(runs, a.out_dir)
+    plot_acc_vs_iters(runs, a.out_dir)
+    plot_state_size(runs, a.out_dir)
+    plot_training(runs, a.results_dir, a.out_dir)
+    summary_table(runs, a.out_dir)
+    print("plots written to", a.out_dir)
 
 
 if __name__ == "__main__":
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--results-dir", default="results")
-    ap.add_argument("--hist-dir", default="out")
-    args = ap.parse_args()
-    os.makedirs("plots", exist_ok=True)
-    evals = load_evals(args.results_dir)
-    print(f"loaded {len(evals)} eval files")
-    acc_vs_k_plot(evals, "plots/acc_vs_k.png")
-    acc_vs_iters_plot(evals, "plots/acc_vs_iters.png")
-    acc_vs_state_plot(evals, "plots/acc_vs_state.png")
-    training_plot(evals, args.hist_dir, "plots/training.png")
-    rows = summary_csv(evals, "results/summary.csv")
-    draft_report(evals, rows, "results/REPORT_DRAFT.md")
-    print("wrote plots/ and results/summary.csv, results/REPORT_DRAFT.md")
+    main()