1#!/usr/bin/env python2"""Regenerate every figure of the manuscript from structured results. Figures whose inputs are absent are skipped3and reported. Usage: python scripts/reproduce_figures.py --results results/raw/ [--runs experiments/runs]4"""5from __future__ import annotations67import argparse, glob, json, sys8from collections import defaultdict9from pathlib import Path1011import numpy as np1213sys.path.insert(0, str(Path(__file__).resolve().parents[1]))14from scripts.plotstyle import CAT, SEQ_BLUE, setup, save # noqa: E4021516import matplotlib.pyplot as plt # noqa: E4021718MADE: list[str] = []19SKIPPED: list[str] = []202122def load_metrics(run_dir):23 p = Path(run_dir) / "metrics.jsonl"24 return [json.loads(l) for l in open(p)] if p.exists() else []252627def runs_matching(runs_root, prefix):28 return sorted(d for d in glob.glob(f"{runs_root}/{prefix}_seed*") if Path(d, "manifest.json").exists())293031# ---------------------------------------------------------------- 1. learnability: val loss by budget vs step32def fig_learnability(runs_root, prefixes, name, title):33 panels = [(p, runs_matching(runs_root, p)) for p in prefixes]34 panels = [(p, r) for p, r in panels if r]35 if not panels:36 SKIPPED.append(name); return37 fig, axes = plt.subplots(1, len(panels), figsize=(2.6 * len(panels), 2.4), sharey=True, squeeze=False)38 for ax, (p, runs) in zip(axes[0], panels):39 for run in runs[:1]:40 rows = [r for r in load_metrics(run) if any(k.startswith("val_loss_b") for k in r)]41 budgets = sorted({int(k[10:]) for r in rows for k in r if k.startswith("val_loss_b")})42 for i, b in enumerate(budgets):43 xs = [r["step"] + 1 for r in rows if f"val_loss_b{b}" in r]; ys = [r[f"val_loss_b{b}"] for r in rows if f"val_loss_b{b}" in r]44 ax.plot(xs, ys, marker="o", color=SEQ_BLUE[min(i + 1, 6)] if b else "#52514e", label=f"budget {b}")45 ax.set_title(p.replace("_", " "), loc="left"); ax.set_xlabel("training step")46 axes[0][0].set_ylabel("held-out answer-span loss (nats)"); axes[0][-1].legend(loc="upper right")47 fig.suptitle(title, x=0.01, ha="left", fontsize=8.5)48 MADE.append(save(fig, name))495051# ---------------------------------------------------------------- 2. labels: gain distribution + variance decomposition52def fig_labels(analyzed_glob):53 files = sorted(glob.glob(analyzed_glob))54 if not files:55 SKIPPED.append("labels"); return56 d = json.loads(Path(files[-1]).read_text())57 part = d.get("calib") or d.get("train")58 aud = part.get("audit", {})59 fig, axes = plt.subplots(1, 3, figsize=(9.0, 2.5), gridspec_kw={"wspace": 0.45})60 ax = axes[0]61 depths = sorted(int(k) for k in part["by_prefix_depth"])62 ax.bar([str(k) for k in depths], [part["by_prefix_depth"][str(k)]["frac_pos"] for k in depths], color=CAT[0], width=0.6)63 ax.axhline(0.5, color="#52514e", lw=0.8, ls="--"); ax.set_ylim(0, 0.7)64 ax.set_xlabel("prefix depth r (edits already applied)"); ax.set_ylabel("fraction of gains > 0"); ax.set_title("Sign of logged gains", loc="left")65 ax = axes[1]66 labels = ["between states", "within state (action)"]67 vals = [aud.get("var_between_states", 0), aud.get("var_within_state_actions", 0)]68 tot = sum(vals) or 169 ax.bar(labels, [v / tot for v in vals], color=[CAT[0], CAT[1]], width=0.6)70 for i, v in enumerate(vals):71 ax.text(i, v / tot + 0.02, f"{100*v/tot:.0f}%", ha="center", fontsize=7)72 ax.set_ylim(0, 1.05); ax.set_ylabel("share of gain variance"); ax.set_title("Variance of audited gains", loc="left")73 ax = axes[2]74 names = ["random\naction", "perfect stop,\nrandom action", "one-step\noracle"]75 vals = [aud.get("random_action_mean", 0), aud.get("oracle_state_level_mean", 0), aud.get("oracle_one_step_mean", 0)]76 ax.bar(names, vals, color=[CAT[2], CAT[3], CAT[0]], width=0.6)77 for i, v in enumerate(vals):78 ax.text(i, v + 0.0005, f"{v:.4f}", ha="center", fontsize=7)79 ax.set_ylabel("mean gain per state (nats)"); ax.set_title("Oracle decomposition", loc="left")80 MADE.append(save(fig, "labels_analysis"))818283# ---------------------------------------------------------------- 3. routing Pareto: loss vs extra ops84def fig_routing(results_dir):85 files = sorted(glob.glob(f"{results_dir}/routing_*.json"))86 if not files:87 SKIPPED.append("routing_pareto"); return88 for f in files:89 d = json.loads(Path(f).read_text()); recs = d["records"]90 fig, ax = plt.subplots(figsize=(4.2, 2.8))91 groups = defaultdict(list)92 for r in recs:93 if r.get("policy") in ("oracle_two_step",):94 continue95 groups[r["policy"]].append(r)96 order = ["budget0", "random_forced", "uniform", "fixed_pair", "entropy", "norm", "controller", "controller_locked", "oracle1_greedy"]97 for i, pol in enumerate([p for p in order if p in groups]):98 rs = sorted(groups[pol], key=lambda r: r["extra_ops_ratio"])99 xs = [r["extra_ops_ratio"] for r in rs]; ys = [r["loss"] for r in rs]100 style = dict(marker="o", color=CAT[i % 8], ls="-" if len(rs) > 1 else "none")101 if pol == "oracle1_greedy":102 style.update(ls="--", color="#52514e")103 if pol == "controller_locked":104 style.update(marker="*", markersize=9, ls="none")105 ax.plot(xs, ys, label=pol.replace("_", " ") + (" (diagnostic)" if "oracle" in pol else ""), **style)106 ax.set_xlabel("extra operations per token (× backbone)"); ax.set_ylabel("held-out loss (nats)")107 ax.set_title(f"Quality vs analytic cost — {d['split']}", loc="left"); ax.legend(ncol=2)108 MADE.append(save(fig, "routing_" + Path(f).stem.replace("routing_", "")))109110111# ---------------------------------------------------------------- 4. latency112def fig_latency(results_dir):113 files = sorted(glob.glob(f"{results_dir}/latency_*.json"))114 if not files:115 SKIPPED.append("latency"); return116 for f in files:117 d = json.loads(Path(f).read_text()); recs = [r for r in d["records"] if r["batch"] == 1]118 fig, ax = plt.subplots(figsize=(4.2, 2.6))119 groups = defaultdict(list)120 for r in recs:121 groups[r["policy"]].append(r)122 for i, (pol, rs) in enumerate(groups.items()):123 rs = sorted(rs, key=lambda r: r["budget"])124 ax.errorbar([r["budget"] for r in rs], [r["median_ms"] for r in rs], yerr=[[0] * len(rs), [r["p95_ms"] - r["median_ms"] for r in rs]],125 marker="o", color=CAT[i % 8], label=pol.replace("_", " "), capsize=2, lw=1.2)126 ax.set_xlabel("executed edits per token (budget)"); ax.set_ylabel("ms per token (median, +p95)")127 ax.set_title(f"Batch-1 decode latency — {d['device']}", loc="left"); ax.legend()128 MADE.append(save(fig, Path(f).stem))129130131# ---------------------------------------------------------------- 5. exact match by regime/policy132def fig_exact_match(results_dir):133 files = sorted(glob.glob(f"{results_dir}/eval_*_dev_*.json")) + sorted(glob.glob(f"{results_dir}/eval_*_test_*.json"))134 files = [f for f in files if "exact_match" in Path(f).read_text()[:400] or True]135 data = defaultdict(dict)136 for f in files:137 d = json.loads(Path(f).read_text())138 if "exact_match" not in d:139 continue140 for pol, em in d["exact_match"].items():141 data[(d["system"], pol)][d["split"]] = em["accuracy"]142 if not data:143 SKIPPED.append("exact_match"); return144 splits = sorted({s for v in data.values() for s in v})145 fig, ax = plt.subplots(figsize=(6.4, 2.6))146 keys = sorted(data)147 width = 0.8 / max(len(keys), 1)148 for i, k in enumerate(keys):149 xs = np.arange(len(splits)) + i * width150 ax.bar(xs, [data[k].get(s, np.nan) for s in splits], width=width * 0.95, color=CAT[i % 8], label=f"{k[0]} / {k[1]}")151 ax.set_xticks(np.arange(len(splits)) + 0.4 - width / 2); ax.set_xticklabels([s.replace("_", "\n") for s in splits])152 ax.set_ylabel("exact match"); ax.set_title("Free-generation exact match by regime", loc="left"); ax.legend(ncol=2, fontsize=6)153 MADE.append(save(fig, "exact_match_regimes"))154155156# ---------------------------------------------------------------- 6. pretraining curves157def fig_pretrain(runs_root):158 runs = runs_matching(runs_root, "pretrain103m") + runs_matching(runs_root, "pretrain_enlarged102m")159 if not runs:160 SKIPPED.append("pretrain"); return161 fig, ax = plt.subplots(figsize=(4.2, 2.6))162 for i, run in enumerate(runs):163 rows = load_metrics(run)164 tr = [(r["tokens"], r["loss"]) for r in rows if "loss" in r]; va = [(r["tokens"], r["val_loss"]) for r in rows if "val_loss" in r]165 name = Path(run).name.split("_seed")[0]166 if tr:167 ax.plot([t / 1e6 for t, _ in tr], [l for _, l in tr], color=CAT[i], alpha=0.35, lw=0.8)168 if va:169 ax.plot([t / 1e6 for t, _ in va], [l for _, l in va], color=CAT[i], marker="o", label=f"{name} (val)")170 ax.set_xlabel("training tokens (M)"); ax.set_ylabel("cross-entropy (nats)"); ax.set_title("Stage A pretraining (FineWeb-Edu)", loc="left"); ax.legend()171 MADE.append(save(fig, "pretrain_curves"))172173174def main(argv=None):175 ap = argparse.ArgumentParser(description=__doc__)176 ap.add_argument("--results", default="results/raw/"); ap.add_argument("--runs", default="experiments/runs"); ap.add_argument("--analyzed", default="results/analyzed")177 args = ap.parse_args(argv); setup()178 fig_learnability(args.runs, ["p4_editor_v1_scaled", "p4_editor_v1_Ustd02", "p4_editor_v1_Ustd05"], "learnability_bootstrap", "Stage B bootstrap on the frozen tiny backbone (dev-iid, random forced actions)")179 fig_learnability(args.runs, ["p4v2_prisme", "p4v2_gr", "p4v2_dense"], "learnability_v2", "Stage B on the v2 tiny backbone: sparse editor vs global recurrence vs dense writes")180 fig_labels(f"{args.analyzed}/labels_*.json")181 fig_routing(args.results); fig_latency(args.results); fig_exact_match(args.results); fig_pretrain(args.runs)182 print("made:", MADE); print("skipped (inputs absent):", SKIPPED)183184185if __name__ == "__main__":186 main()187