scripts/make_figures.py

#!/usr/bin/env python3
"""
make_figures.py  --  Generate all ACML 2026 paper figures.

Output: paper/figures/fig_*.png  (300 dpi, black-and-white, JMLR column widths)

Run from the paper/ directory:
    python make_figures.py
"""

import csv, json, math, sys
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
from pathlib import Path

# ---------------------------------------------------------------------------
# Global ACML / JMLR style  (black-and-white, 10pt serif)
# ---------------------------------------------------------------------------
plt.rcParams.update({
    "font.family":        "serif",
    "font.size":          9,
    "axes.titlesize":     9,
    "axes.labelsize":     9,
    "xtick.labelsize":    8,
    "ytick.labelsize":    8,
    "legend.fontsize":    7.5,
    "legend.framealpha":  0.85,
    "figure.dpi":         300,
    "savefig.dpi":        300,
    "savefig.bbox":       "tight",
    "axes.linewidth":     0.8,
    "grid.linewidth":     0.4,
    "lines.linewidth":    1.4,
    "lines.markersize":   4,
    "text.usetex":        False,
    "axes.spines.top":    False,
    "axes.spines.right":  False,
})

COL1 = 3.25   # single column (inches)
COL2 = 6.75   # full width

DATA = Path(__file__).parent / "data"
OUT  = Path(__file__).parent / "figures"
OUT.mkdir(exist_ok=True)

LS = ["-", "--", "-.", ":"]
MK = ["o", "s", "^", "D", "v", "P"]
GR = ["#000000", "#444444", "#888888", "#AAAAAA"]


def save(name):
    p = OUT / f"fig_{name}.png"
    plt.savefig(p)
    plt.close()
    print(f"  saved {p.name}")


# ---------------------------------------------------------------------------
# Fig 1 -- VRAM vs sequence length
# ---------------------------------------------------------------------------
def fig_vram():
    rows = list(csv.DictReader(open(DATA / "results/roofline/vram_comparison.csv")))

    fno_x, fno_y = [], []
    sdpa_real_x, sdpa_real_y = [], []
    sdpa_th_x, sdpa_th_y = [], []

    for r in rows:
        if not r["vram_mb"]:
            continue
        T = int(r["seq_len"])
        v = int(r["vram_mb"]) / 1024
        if "FNO+GLA" in r["model"]:
            fno_x.append(T); fno_y.append(v)
        elif "theoretical" in r["model"]:
            sdpa_th_x.append(T); sdpa_th_y.append(v)
        elif "SDPA" in r["model"]:
            sdpa_real_x.append(T); sdpa_real_y.append(v)

    fig, ax = plt.subplots(figsize=(COL1, 2.6))
    ax.plot(fno_x, fno_y, LS[0]+MK[0], color="k", label="FELA 1.13B")
    if sdpa_real_x:
        ax.plot(sdpa_real_x, sdpa_real_y, MK[1], color="k",
                markerfacecolor="white", ms=6, zorder=5, label="GPT-2 XL (OOM at 2K)")
    if sdpa_th_x:
        ax.plot(sdpa_th_x, sdpa_th_y, LS[1], color="k", alpha=0.7,
                label="GPT-2 XL (extrap.)")
    ax.axhline(24, color="k", lw=0.7, ls=":", alpha=0.5)
    ax.text(512, 25.5, "A10G 24 GB", fontsize=7, color="gray")
    ax.set_xscale("log", base=2)
    ax.set_yscale("log")
    ax.set_xlabel("Sequence length (tokens)")
    ax.set_ylabel("Peak VRAM (GB)")
    ax.set_title("VRAM vs. Sequence Length")
    ax.xaxis.set_major_formatter(mticker.FuncFormatter(lambda x, _: f"{int(x):,}"))
    ax.legend(loc="upper left")
    ax.grid(True, which="both", alpha=0.2)
    save("vram_measured")


# ---------------------------------------------------------------------------
# Fig 2 -- Long-context BPB
# ---------------------------------------------------------------------------
def fig_longctx():
    rows = list(csv.DictReader(
        open(DATA / "results/paper_eval_22b/longctx_bpb.tsv"), delimiter="\t"))
    xs = [int(r["ctx_len"]) for r in rows]
    ys = [float(r["bpb"]) / 4.0 for r in rows]   # bits/token -> bits/byte

    fig, ax = plt.subplots(figsize=(COL1, 2.5))
    ax.plot(xs, ys, LS[0]+MK[0], color="k")
    dy = ys[0] - ys[-1]
    ax.annotate("", xy=(xs[-1], ys[-1]), xytext=(xs[-1], ys[0]),
                arrowprops=dict(arrowstyle="<->", color="k", lw=0.8))
    ax.text(xs[-1]*1.08, (ys[0]+ys[-1])/2,
            f"-{dy:.2f}\n({dy/ys[0]*100:.0f}%)", fontsize=7, va="center")
    ax.set_xscale("log", base=2)
    ax.set_xlabel("Context length (tokens)")
    ax.set_ylabel("WikiText-103 BPB (bits/byte)")
    ax.set_title("Long-Context Utilisation")
    ax.xaxis.set_major_formatter(mticker.FuncFormatter(lambda x, _: f"{int(x):,}"))
    ax.grid(True, alpha=0.2)
    save("longctx_bpb")


# ---------------------------------------------------------------------------
# Fig 3 -- Throughput
# ---------------------------------------------------------------------------
def fig_throughput():
    rows = list(csv.DictReader(
        open(DATA / "results/paper_eval_22b/throughput.tsv"), delimiter="\t"))
    xs = [int(r["seq_len"]) for r in rows]
    ys = [int(r["tok_per_sec"]) / 1000 for r in rows]

    base = ys[0]
    sdpa_xs = [512, 1024, 2048, 4096, 8192, 16384, 32768]
    sdpa_ys = [base * (512 / T) for T in sdpa_xs]

    fig, ax = plt.subplots(figsize=(COL1, 2.5))
    ax.plot(xs, ys, LS[0]+MK[0], color="k", label="FELA 1.13B")
    ax.plot(sdpa_xs, sdpa_ys, LS[1]+MK[1], color="k",
            markerfacecolor="white", label="SDPA (theoretical)")
    ax.set_xscale("log", base=2)
    ax.set_xlabel("Sequence length (tokens)")
    ax.set_ylabel("Throughput (K tok/s)")
    ax.set_title("Prefill Throughput")
    ax.xaxis.set_major_formatter(mticker.FuncFormatter(lambda x, _: f"{int(x):,}"))
    ax.legend()
    ax.grid(True, alpha=0.2)
    save("throughput")


# ---------------------------------------------------------------------------
# Fig 4 -- Needle-in-haystack heatmap
# ---------------------------------------------------------------------------
def fig_needle():
    data = json.load(open(DATA / "results/paper_eval_22b/needle_heatmap.json"))
    rows = data if isinstance(data, list) else data.get("results", [])

    ctx_lens = sorted(set(int(r["ctx_len"]) for r in rows))
    depths   = sorted(set(float(r["depth"]) for r in rows))

    mat = np.full((len(depths), len(ctx_lens)), np.nan)
    for r in rows:
        ci = ctx_lens.index(int(r["ctx_len"]))
        di = depths.index(float(r["depth"]))
        acc = r.get("accuracy")
        if acc is not None and not (isinstance(acc, float) and math.isnan(acc)):
            mat[di, ci] = float(acc)

    fig, ax = plt.subplots(figsize=(COL1, 2.8))
    im = ax.imshow(mat, aspect="auto", cmap="Greys", vmin=0, vmax=1,
                   origin="lower", interpolation="nearest")
    plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label="Accuracy")
    ax.set_xticks(range(len(ctx_lens)))
    ax.set_xticklabels(
        [f"{c//1024}K" if c >= 1024 else str(c) for c in ctx_lens],
        rotation=45, ha="right")
    ax.set_yticks(range(len(depths)))
    ax.set_yticklabels([f"{int(d*100)}%" for d in depths])
    ax.set_xlabel("Context length")
    ax.set_ylabel("Needle depth")
    ax.set_title("Passkey Retrieval (FELA 1.13B)")
    save("needle_heatmap")


# ---------------------------------------------------------------------------
# Fig 5 -- Chunk-size ablation: 3-panel
# ---------------------------------------------------------------------------
def fig_chunk():
    chunks = [64, 256, 512, 1024]

    # BPB from per-variant wikitext.json
    bpb_vals = []
    for c in chunks:
        wt = json.load(open(DATA / f"results/paper_eval_22b/chunk_ablation/variant_{c}/wikitext.json"))
        # bpb field is bits/token; divide by 4 (avg bytes/BPE token) -> bits/byte
        raw = wt.get("bpb") or wt.get("val_bpb") or wt.get("wikitext_bpb")
        bpb_vals.append(float(raw) / 4.0 if raw is not None else float("nan"))

    # Retrieval from chunk_ablation_FIXED needle_fixed.json  (4K ctx, depth 90%)
    retr_vals = []
    for c in chunks:
        needle = json.load(open(
            DATA / f"results/chunk_ablation_fixed/variant_{c}/needle_fixed.json"))
        hit = [float(r["accuracy"]) for r in needle
               if int(r["ctx_len"]) == 4096
               and abs(float(r["depth"]) - 0.9) < 0.01
               and r.get("accuracy") is not None
               and not math.isnan(float(r["accuracy"]))]
        retr_vals.append(hit[0] if hit else float("nan"))

    # Throughput: max tok/s from per-variant throughput.json
    tput_vals = []
    for c in chunks:
        tp = json.load(open(DATA / f"results/paper_eval_22b/chunk_ablation/variant_{c}/throughput.json"))
        vals = []
        def _collect(obj):
            if isinstance(obj, dict):
                for v in obj.values(): _collect(v)
            elif isinstance(obj, list):
                for v in obj: _collect(v)
            elif isinstance(obj, (int, float)) and 1000 < float(obj) < 1e7:
                vals.append(float(obj))
        _collect(tp)
        tput_vals.append(max(vals) / 1000 if vals else float("nan"))

    print(f"  chunk BPB:   {bpb_vals}")
    print(f"  chunk retr:  {retr_vals}")
    print(f"  chunk tput:  {tput_vals}")

    fig, axes = plt.subplots(1, 3, figsize=(COL2, 2.4))
    xs = list(range(len(chunks)))
    xlabels = [str(c) for c in chunks]
    hatches = ["", "//", "xx", ".."]

    specs = [
        (bpb_vals,  "BPB (lower is better)",      "WikiText-103 BPB"),
        (retr_vals, "Retrieval accuracy",           "Retrieval (4K, 90% depth)"),
        (tput_vals, "Throughput (K tok/s)",         "Max Throughput"),
    ]
    for ax, (vals, ylabel, title) in zip(axes, specs):
        # Only plot bars where value is not nan
        for xi, (v, h) in enumerate(zip(vals, hatches)):
            if not math.isnan(v):
                bar = ax.bar(xi, v, color="white", edgecolor="black",
                             linewidth=0.8, hatch=h)
                ax.text(xi, v * 1.015,
                        f"{v:.3f}" if v < 10 else f"{v:.0f}",
                        ha="center", va="bottom", fontsize=7)
        ax.set_xticks(xs)
        ax.set_xticklabels(xlabels)
        ax.set_xlabel("GLA chunk size")
        ax.set_ylabel(ylabel)
        ax.set_title(title)
        ax.grid(True, axis="y", alpha=0.25)
        if not all(math.isnan(v) for v in vals):
            valid = [v for v in vals if not math.isnan(v)]
            pad = (max(valid) - min(valid)) * 0.15 or max(valid) * 0.1
            ax.set_ylim(min(valid) * 0.92, max(valid) + pad * 3)

    fig.suptitle("GLA Chunk Size Ablation (109M, 1B tokens)", y=1.02)
    plt.tight_layout()
    save("chunk_ablation")


# ---------------------------------------------------------------------------
# Fig 6 -- FNO filter spectrum
# ---------------------------------------------------------------------------
def fig_fno_spectrum():
    data = json.load(open(DATA / "results/interpretability/fno_filter_spectrum.json"))

    groups = [
        ("Early (L0-2)",  ["layer_0", "layer_1", "layer_2"]),
        ("Mid (L4-6)",    ["layer_4", "layer_5", "layer_6"]),
        ("Late (L8-10)",  ["layer_8", "layer_9", "layer_10"]),
    ]

    fig, axes = plt.subplots(1, 3, figsize=(COL2, 2.5), sharey=False)
    for ax, (title, lnames) in zip(axes, groups):
        for i, lname in enumerate(lnames):
            if lname not in data:
                continue
            mags = np.array(data[lname]["magnitudes"])
            freqs = np.arange(len(mags)) / len(mags)
            lnum = lname.split("_")[1]
            ax.plot(freqs, mags, ls=LS[i], color=GR[i],
                    label=f"L{lnum}", alpha=0.9)
        ax.set_yscale("log")
        ax.set_xlabel("Normalised frequency")
        if ax is axes[0]:
            ax.set_ylabel("|H(f)| (mean over channels)")
        ax.set_title(title)
        ax.legend(fontsize=7)
        ax.grid(True, alpha=0.2)

    fig.suptitle("FNO Filter Spectra: Coarse-to-Fine Hierarchy", y=1.02)
    plt.tight_layout()
    save("fno_spectrum")


# ---------------------------------------------------------------------------
# Fig 7 -- GLA gate hierarchy
# ---------------------------------------------------------------------------
def fig_gla_gates():
    data = json.load(open(DATA / "results/interpretability/gla_gate_heatmap.json"))
    texts = list(data.keys())
    layer_names = list(data[texts[0]].keys())
    layer_nums  = [int(ln.split("_")[1]) for ln in layer_names]

    fig, axes = plt.subplots(1, len(texts), figsize=(COL2, 2.4), sharey=True)
    if len(texts) == 1:
        axes = [axes]

    for ax, text in zip(axes, texts):
        for li, (lname, lnum) in enumerate(zip(layer_names, layer_nums)):
            gates = np.array(data[text][lname])   # [T, H]
            decay = np.exp(gates)
            mean_decay = decay.mean(axis=0)
            ax.plot(np.arange(len(mean_decay)), mean_decay,
                    ls=LS[li], marker=MK[li], color=GR[li],
                    label=f"L{lnum}", markersize=3)
        ax.set_xlabel("Head")
        ax.set_ylim(0, 1.05)
        short = text[:24] + "..." if len(text) > 24 else text
        ax.set_title(short, fontsize=7)
        ax.grid(True, alpha=0.2)
        ax.axhline(1.0, color="k", lw=0.4, ls=":")

    axes[0].set_ylabel("Mean gate decay (1=remember)")
    axes[0].legend(fontsize=7, title="Layer")
    fig.suptitle("GLA Gate Hierarchy: Deeper Layers Retain More Context", y=1.02)
    plt.tight_layout()
    save("gla_gates")


# ---------------------------------------------------------------------------
# Fig 8 -- Logit lens
# ---------------------------------------------------------------------------
def fig_logit_lens():
    data = json.load(open(DATA / "results/explain/logit_lens.json"))
    texts = list(data.keys())

    fig, ax = plt.subplots(figsize=(COL1, 2.8))
    all_layer_ids = None
    for ti, text in enumerate(texts[:3]):
        layers_data = data[text]["layers"]
        # final token = top-1 at last layer
        final_token = layers_data[-1]["top_k"][0]["token"]
        layer_ids = [l["layer"] for l in layers_data]
        # 1 if top-1 token already matches final prediction, 0 otherwise
        matches = [1 if l["top_k"][0]["token"] == final_token else 0
                   for l in layers_data]
        short = text[:20] + "..."
        ax.plot(layer_ids, matches, LS[ti]+MK[ti], color=GR[ti],
                label=short, lw=1.5)
        all_layer_ids = layer_ids

    ax.set_xlabel("Layer")
    ax.set_ylabel("Top 1 token matches final prediction")
    ax.set_title("Logit Lens: When Does Prediction Commit?")
    ax.set_xticks(all_layer_ids)
    ax.set_yticks([0, 1])
    ax.set_yticklabels(["No", "Yes"])
    ax.set_ylim(-0.15, 1.25)
    ax.axvline(6, color="k", lw=0.6, ls=":", alpha=0.5)
    ax.text(6.15, 0.08, "all commit\nby layer 6", fontsize=7, color="gray")
    ax.legend(fontsize=7)
    ax.grid(True, alpha=0.2)
    save("logit_lens")


# ---------------------------------------------------------------------------
# Fig 9 -- Training curve
# ---------------------------------------------------------------------------
def fig_training():
    rows = list(csv.DictReader(open(DATA / "results/gpu_v1/training_data.csv")))
    steps = [int(r["step"]) for r in rows if r["step"]]
    loss  = [float(r["loss"]) for r in rows if r["loss"]]
    val_steps = [int(r["step"]) for r in rows if r.get("val_bpb")]
    val_bpb   = [float(r["val_bpb"]) for r in rows if r.get("val_bpb")]

    fig, ax1 = plt.subplots(figsize=(COL1, 2.5))
    ax1.plot(steps, loss, color="k", lw=0.6, alpha=0.4)
    if len(loss) > 50:
        smooth = np.convolve(loss, np.ones(50)/50, mode="valid")
        ax1.plot(steps[49:], smooth, color="k", lw=1.5, label="Train loss")
    ax1.set_xlabel("Step")
    ax1.set_ylabel("Loss (nats)")
    ax1.set_title("Training Curve -- FELA 1.13B")
    ax1.grid(True, alpha=0.2)

    if val_bpb:
        ax2 = ax1.twinx()
        ax2.plot(val_steps, val_bpb, LS[1]+MK[1], color="k",
                 markerfacecolor="white", ms=5, label="Val BPB")
        ax2.set_ylabel("Val BPB")
        lines1, labs1 = ax1.get_legend_handles_labels()
        lines2, labs2 = ax2.get_legend_handles_labels()
        ax1.legend(lines1+lines2, labs1+labs2, fontsize=7, loc="upper right")
    else:
        ax1.legend(fontsize=7)
    save("training_curve")


# ---------------------------------------------------------------------------
# Fig 10 -- Wall-clock throughput: FELA vs SDPA (measured head-to-head)
# ---------------------------------------------------------------------------
def fig_wallclock():
    rows = list(csv.DictReader(open(DATA / "results/gpu_v1/gpu_crossover_v1.csv")))

    fela_x, fela_y = [], []
    sdpa_x, sdpa_y = [], []

    for r in rows:
        T = int(r["seq_len"])
        if r["fno_oom"] == "True" or not r["fno_gla_tok_s"]:
            pass
        else:
            fela_x.append(T); fela_y.append(float(r["fno_gla_tok_s"]) / 1000)
        if r["sdpa_oom"] == "True" or not r["sdpa_tok_s"]:
            pass
        else:
            sdpa_x.append(T); sdpa_y.append(float(r["sdpa_tok_s"]) / 1000)

    fig, ax = plt.subplots(figsize=(COL1, 2.6))
    ax.plot(fela_x, fela_y, LS[0]+MK[0], color="k", label="FELA 1.13B")
    ax.plot(sdpa_x, sdpa_y, LS[1]+MK[1], color="k",
            markerfacecolor="white", label="GPT-2 XL SDPA")

    # Mark SDPA OOM point
    if sdpa_x:
        ax.axvline(sdpa_x[-1], color="k", lw=0.7, ls=":", alpha=0.5)
        ax.text(sdpa_x[-1]*1.05, min(sdpa_y)*1.1, "SDPA\nOOM", fontsize=7, color="gray")

    ax.set_xscale("log", base=2)
    ax.set_xlabel("Sequence length (tokens)")
    ax.set_ylabel("Throughput (K tok/s)")
    ax.set_title("Measured Throughput: FELA vs. SDPA")
    ax.xaxis.set_major_formatter(mticker.FuncFormatter(lambda x, _: f"{int(x):,}"))
    ax.legend()
    ax.grid(True, alpha=0.2)
    save("wallclock")


# ---------------------------------------------------------------------------
# Fig 11 -- Inference FLOPs
# ---------------------------------------------------------------------------
def fig_flops():
    rows = list(csv.DictReader(
        open(DATA / "results/paper_eval_22b/gmacs.tsv"), delimiter="\t"))
    fela = [r for r in rows if "FELA" in r.get("model", "")]
    xs_f = [int(r["seq_len"]) for r in fela]
    ys_f = [float(r["gmacs"]) / 1000 for r in fela]   # GMAC -> TFLOP approx

    sdpa_xs = sorted(set(xs_f + [512, 2048, 8192, 32768, 65536, 131072]))
    # GPT-2 XL: 48 layers, 25 heads, 64 head-dim; 2*N*L*H*D per-token inference
    sdpa_ys = [2 * T * 48 * 25 * 64 / 1e9 for T in sdpa_xs]

    fig, ax = plt.subplots(figsize=(COL1, 2.5))
    ax.plot(xs_f, ys_f, LS[0]+MK[0], color="k", label="FELA (prefill)")
    ax.plot(sdpa_xs, sdpa_ys, LS[1]+MK[1], color="k",
            markerfacecolor="white", label="SDPA per-token")
    ax.set_xscale("log", base=2)
    ax.set_yscale("log")
    ax.set_xlabel("Sequence length (tokens)")
    ax.set_ylabel("GFLOPs")
    ax.set_title("Compute vs. Sequence Length")
    ax.xaxis.set_major_formatter(mticker.FuncFormatter(lambda x, _: f"{int(x):,}"))
    ax.legend()
    ax.grid(True, which="both", alpha=0.2)
    save("flops")


# ---------------------------------------------------------------------------
if __name__ == "__main__":
    target = sys.argv[1] if len(sys.argv) > 1 else None
    FIGS = [
        ("vram_measured",  fig_vram),
        ("longctx_bpb",    fig_longctx),
        ("throughput",     fig_throughput),
        ("wallclock",      fig_wallclock),
        ("needle_heatmap", fig_needle),
        ("chunk_ablation", fig_chunk),
        ("fno_spectrum",   fig_fno_spectrum),
        ("gla_gates",      fig_gla_gates),
        ("logit_lens",     fig_logit_lens),
        ("training_curve", fig_training),
        ("flops",          fig_flops),
    ]
    for name, fn in FIGS:
        if target and target != name:
            continue
        print(f"[{name}]")
        try:
            fn()
        except Exception as e:
            import traceback
            print(f"  ERROR: {e}")
            traceback.print_exc()
    print("done")