Files
jspace/scripts/fig_kcurves.py
T
2026-07-15 18:34:18 +02:00

112 lines
4.5 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""Small-multiples: pass@1 vs loop depth k, easy (pastel) + hard (dark),
one panel per regime arm — the full design-space sweep. All loop arms on
the same 250-item MBPP eval, e400 checkpoints. Final panel: no-recurrence
references (FF control, plan-distill; distill on the 500-item set).
Pending arms render as empty boxes and fill in on re-render.
"""
import json
import re
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
OUT = Path(__file__).resolve().parent.parent / "results-loop"
LOGD = Path("/home/nils/gpuq/spark-gpu0")
def cells(fname):
try:
d = json.load(open(OUT / fname))
except FileNotFoundError:
return None
return {int(k): (v["by_label"].get("easy"), v["by_label"].get("hard"))
for k, v in d["ks"].items()}
def cells_from_log(logfile):
out = {}
pat = re.compile(r"k=(\d+): pass@1=[\d.]+ by_label=\{'easy': ([\d.]+), "
r"'hard': ([\d.]+)")
try:
for line in open(logfile):
m = pat.search(line.replace('"', "'"))
if m:
out[int(m.group(1))] = (float(m.group(2)), float(m.group(3)))
except FileNotFoundError:
pass
return out or None
PANELS = [
("untrained α-merge\ntraining-free · ρ=0.3",
cells("eval_code_untrained.json"), "#8a8f98"),
("trained merge ★\nfixed A,B · curriculum",
cells("eval_code_trained.json"), "#2b6cb0"),
("tied-alpha\nlearned a, B=(1a) tied · curriculum",
cells("eval_code_ta_e400.json"), "#0987a0"),
("random-depth merge\nfixed A,B · NO curriculum",
cells("eval_code_merge_rk16_e400.json")
or cells_from_log(LOGD / "zzz_a_randk.log"), "#3182ce"),
("unconstrained rec\nlearned A,B · noise h₀ · ρ→4.5",
cells("eval_code_rec16_partial.json"), "#c53030"),
("Parcae rec\nρ<1 enforced · learned B · ρ→0.29",
cells("eval_code_parcae16_e400.json"), "#2f855a"),
("per-depth merges\ntime-varying · anchored",
cells("eval_code_pd4_e400.json")
or cells_from_log(LOGD / "perdepth_arm.log"), "#805ad5"),
("noise-s₀ merge\nfixed A,B · curriculum · noise s₀",
cells("eval_code_merge_ns_e400.json")
or cells_from_log(LOGD / "zzz_b_noises0.log"), "#b83280"),
("4× capacity merge\nMLP h=2048 · curriculum",
cells("eval_code_merge_h2048_e400.json")
or cells_from_log(LOGD / "zzz_c_h2048.log"), "#5f6b7a"),
("Parcae rec, seed 1",
cells("eval_code_parcae16_s1_e400.json")
or cells_from_log(LOGD / "zzz_d_parcae_s1.log"), "#276749"),
]
fig, axes = plt.subplots(2, 5, figsize=(16, 6.6), sharey=True)
axf = axes.flatten()
for ax, (title, kc, dark) in zip(axf, PANELS):
ax.grid(True, color="#ececec", lw=0.7)
ax.set_axisbelow(True)
for s in ("top", "right"):
ax.spines[s].set_visible(False)
ax.set_title(title, fontsize=8.6)
ax.set_ylim(0, 1.02)
if not kc:
ax.text(0.5, 0.5, "pending", transform=ax.transAxes, ha="center",
color="#999", fontsize=10)
continue
ks = sorted(kc)
xs = range(len(ks))
ax.plot(xs, [kc[k][0] for k in ks], "-o", color=dark, alpha=0.32,
lw=2.2, ms=4.5)
ax.plot(xs, [kc[k][1] for k in ks], "-s", color=dark, lw=2.2, ms=4.5)
# references: base levels (dotted) and distill hard (dashed amber)
ax.axhline(kc[ks[0]][0], color="#999", lw=0.8, ls=":")
ax.axhline(kc[ks[0]][1], color="#999", lw=0.8, ls=":")
ax.axhline(0.457, color="#b7791f", lw=1.0, ls="--", alpha=0.7)
ax.set_xticks(list(xs))
ax.set_xticklabels([str(k) for k in ks], fontsize=7.5)
ax.set_xlabel("k", fontsize=8)
for i in (0, 5):
axf[i].set_ylabel("pass@1", fontsize=9)
from matplotlib.lines import Line2D
axf[0].legend(handles=[
Line2D([], [], color="#555", alpha=0.32, marker="o", lw=2.2, label="easy"),
Line2D([], [], color="#555", marker="s", lw=2.2, label="hard"),
Line2D([], [], color="#b7791f", ls="--", lw=1.0,
label="distill hard (45.7%, no loop)"),
Line2D([], [], color="#999", ls=":", lw=0.8, label="k=0 base levels"),
], fontsize=7, frameon=False, loc="center right")
fig.suptitle("The design space, measured — depth curves for every regime "
"(same frozen band, data, 250-item eval; ★ = recommended recipe; "
"amber dashes = what distillation reaches with no loop at all)",
fontsize=11, y=1.0)
fig.tight_layout()
fig.savefig(OUT / "fig_kcurves.png", dpi=140, facecolor="white",
bbox_inches="tight")
print("wrote", OUT / "fig_kcurves.png")