src/make_figures.py

#!/usr/bin/env python3
"""Generate figures from stored experiment outputs (deterministic, data-driven)."""
import json, os
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

os.makedirs("figures", exist_ok=True)

# --- Fig 1: Claim 1 synthetic vs real acceptance length ---
try:
    d = json.load(open("outputs/claim1_throughput_bias.json"))
    fig, ax = plt.subplots(1, 2, figsize=(11, 4.2))
    r, s = d["real"], d["synthetic"]
    ax[0].bar(["real\n(SPEED-Bench)", "synthetic\n(random tokens)"],
              [r["mean_AL"], s["mean_AL"]],
              yerr=[r["std_AL"], s["std_AL"]], capsize=6,
              color=["#2a7fbf", "#bf6a2a"])
    ax[0].set_ylabel("mean acceptance length")
    ax[0].set_title("Claim 1: synthetic vs real AL\n(overestimation = %.1f%%, within noise)" % d["overestimation_pct"])
    cats = sorted(d["per_category_real_mean_AL"].items(), key=lambda x: -x[1])
    ax[1].barh([c for c, _ in cats][::-1], [v for _, v in cats][::-1], color="#2a7fbf")
    ax[1].axvline(s["mean_AL"], color="#bf6a2a", ls="--", label="synthetic mean")
    ax[1].set_xlabel("mean acceptance length"); ax[1].set_title("Real AL by category"); ax[1].legend()
    plt.tight_layout(); plt.savefig("figures/claim1_throughput_bias.png", dpi=130); plt.close()
    print("figures/claim1_throughput_bias.png")
except Exception as e:
    print("fig1 skip:", e)

# --- Fig 2: Claim 2a diversity ---
try:
    d = json.load(open("outputs/claim2_diversity.json"))
    pc = d["per_category"]
    cats = list(pc.keys())
    g = [pc[c]["greedy_sim"] for c in cats]
    r = [pc[c]["random_sim"] for c in cats]
    third_key = "lowdiv_sim" if "lowdiv_sim" in pc[cats[0]] else "full_sim"
    third_lbl = "low-diversity (max-similarity)" if third_key == "lowdiv_sim" else "full pool (all 80)"
    lo = [pc[c][third_key] for c in cats]
    x = np.arange(len(cats)); w = 0.27
    fig, ax = plt.subplots(figsize=(11, 4.6))
    ax.bar(x - w, g, w, label="greedy diverse (SPEED-Bench Alg.1)", color="#2a9d3f")
    ax.bar(x, r, w, label="random selection", color="#888")
    ax.bar(x + w, lo, w, label=third_lbl, color="#bf2a2a")
    ax.set_xticks(x); ax.set_xticklabels(cats, rotation=40, ha="right")
    ax.set_ylabel("avg pairwise cosine similarity")
    ax.set_title("Claim 2a: diversity selection lowers intra-set similarity "
                 "(greedy < random in %d/%d categories)" % (d["overall"]["n_better"], d["overall"]["n_total"]))
    ax.legend()
    plt.tight_layout(); plt.savefig("figures/claim2a_diversity.png", dpi=130); plt.close()
    print("figures/claim2a_diversity.png")
except Exception as e:
    print("fig2 skip:", e)

# --- Fig 3: Claim 2b draft length vs batch (roofline mechanism + CPU-measured null) ---
try:
    rf = json.load(open("outputs/claim2b_roofline.json"))
    meas = json.load(open("outputs/claim2b_draft_length.json"))
    pb = rf["per_batch"]
    batches = sorted((int(b) for b in pb), key=int)
    fig, ax = plt.subplots(1, 2, figsize=(11, 4.4))
    for b in batches:
        if b not in (1, 8, 32, 64, 256):  # declutter
            continue
        tp = pb[str(b)]["throughput_by_k"]
        ks = sorted((int(k) for k in tp), key=int)
        y = np.array([tp[str(k)] for k in ks]); y = y / y.max()
        ax[0].plot(ks, y, marker="o", label="batch %d" % b)
    ax[0].set_xlabel("draft length k"); ax[0].set_ylabel("normalized throughput")
    ax[0].set_title("Roofline throughput vs draft length (alpha=%.2f)" % rf["alpha"]); ax[0].legend(fontsize=8)
    ks_star = [rf["kstar_by_batch"][str(b)] for b in batches]
    ax[1].plot(batches, ks_star, marker="s", color="#2a9d3f", label="roofline model (k* decreases)")
    mb = sorted((int(b) for b in meas["per_batch"]), key=int)
    mstar = [meas["kstar_by_batch"][str(b)] for b in mb]
    ax[1].plot(mb, mstar, marker="x", ls="--", color="#888", label="CPU measured (compute-bound: k*=1)")
    ax[1].set_xscale("log", base=2); ax[1].set_xlabel("batch size"); ax[1].set_ylabel("optimal draft length k*")
    ax[1].set_ylim(0, 4); ax[1].legend(fontsize=8)
    ax[1].set_title("Claim 2b: optimal draft length vs batch")
    plt.tight_layout(); plt.savefig("figures/claim2b_draft_length.png", dpi=130); plt.close()
    print("figures/claim2b_draft_length.png")
except Exception as e:
    print("fig3 skip:", e)