repro_bundle/code/plot.py

"""Build figures + CSVs from the job artifacts (eval.json, efficiency.json,
per-run *_hist.json). PNGs for the poster, CSVs as figure-cell raw data."""
from __future__ import annotations

import argparse
import csv
import glob
import json
import os

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

C = {"L-SR1": "#4C72B0", "L-BFGS": "#DD8452", "Adam": "#55A868",
     "AdaHessian": "#C44E52", "proj": "#4C72B0", "noproj": "#8172B3"}


def savecsv(path, header, rows):
    with open(path, "w", newline="") as f:
        w = csv.writer(f); w.writerow(header); w.writerows(rows)


def fig_newton(ev, outdir):
    na = ev["newton_alignment"]
    it = list(range(1, len(na["lsr1_proj"]) + 1))
    plt.figure(figsize=(5, 3.4))
    plt.plot(it, na["lsr1_proj"], "-o", ms=3, color=C["proj"],
             label="L-SR1 (projection)")
    plt.plot(it, na["lsr1_noproj"], "-s", ms=3, color=C["noproj"],
             label="L-SR1 (no projection)")
    plt.plot(it, na["lbfgs"], "-^", ms=3, color=C["L-BFGS"], label="L-BFGS")
    plt.xlabel("optimization step"); plt.ylabel("cosine to Newton direction")
    plt.title(f"Newton-direction alignment (quadratics, N={na['N']})")
    plt.legend(fontsize=8); plt.grid(alpha=.3); plt.tight_layout()
    plt.savefig(f"{outdir}/newton_alignment.png", dpi=150); plt.close()
    savecsv(f"{outdir}/newton_alignment.csv",
            ["step", "lsr1_proj", "lsr1_noproj", "lbfgs"],
            list(zip(it, na["lsr1_proj"], na["lsr1_noproj"], na["lbfgs"])))


def fig_profiles(ev, outdir):
    pf = ev["performance"]
    taus, profs, aucs = pf["taus"], pf["profiles"], pf["aucs"]
    plt.figure(figsize=(5, 3.4))
    for s in ["L-SR1", "L-BFGS", "Adam", "AdaHessian"]:
        plt.plot(taus, profs[s], color=C[s],
                 label=f"{s} (AUC={aucs[s]:.2f})")
    plt.xscale("log"); plt.xlabel(r"performance ratio $\tau$")
    plt.ylabel(r"fraction of problems ($\rho$)")
    plt.title("Performance profiles (30-problem suite)")
    plt.legend(fontsize=8); plt.grid(alpha=.3); plt.ylim(0, 1.02)
    plt.tight_layout(); plt.savefig(f"{outdir}/performance_profiles.png", dpi=150)
    plt.close()
    rows = list(zip(taus, *[profs[s] for s in ["L-SR1", "L-BFGS", "Adam",
                                               "AdaHessian"]]))
    savecsv(f"{outdir}/performance_profiles.csv",
            ["tau", "L-SR1", "L-BFGS", "Adam", "AdaHessian"], rows)


def fig_convergence(ev, outdir):
    cv = ev["convergence"]
    plt.figure(figsize=(5, 3.4))
    plt.semilogy(cv["iters"], cv["lsr1"], "-o", ms=3, color=C["L-SR1"],
                 label="L-SR1")
    plt.semilogy(cv["iters"], cv["lbfgs"], "-^", ms=3, color=C["L-BFGS"],
                 label="L-BFGS")
    plt.semilogy(cv["iters"], cv["adam"], "-s", ms=3, color=C["Adam"],
                 label="Adam")
    plt.xlabel("optimization step"); plt.ylabel("median f - f*  (log)")
    plt.title("Convergence on quadratics (N=10)")
    plt.legend(fontsize=8); plt.grid(alpha=.3, which="both"); plt.tight_layout()
    plt.savefig(f"{outdir}/convergence.png", dpi=150); plt.close()
    savecsv(f"{outdir}/convergence.csv", ["iter", "L-SR1", "L-BFGS", "Adam"],
            list(zip(cv["iters"], cv["lsr1"], cv["lbfgs"], cv["adam"])))


def fig_efficiency(effpath, outdir):
    if not os.path.exists(effpath):
        return
    e = json.load(open(effpath))
    ref = e["paper_reference"]
    fig, ax = plt.subplots(1, 2, figsize=(6.4, 3.2))
    labels = ["LGD-style", "L-SR1"]
    ax[0].bar(labels, [e["lgd_ms"], e["lsr1_ms"]], color=[C["L-BFGS"], C["L-SR1"]])
    ax[0].set_ylabel("ms / inner step"); ax[0].set_title(
        f"Runtime (−{e['speedup_pct']:.0f}%)  [paper −{ref['speedup_pct']}%]")
    mr = e.get("mem_reduction_pct")
    ax[1].bar(labels, [e["lgd_mem_gib"], e["lsr1_mem_gib"]],
              color=[C["L-BFGS"], C["L-SR1"]])
    ax[1].set_ylabel("peak memory (GiB)")
    ax[1].set_title((f"Memory (−{mr:.0f}%)" if mr else "Memory") +
                    f"  [paper −{ref['mem_reduction_pct']}%]")
    plt.tight_layout(); plt.savefig(f"{outdir}/efficiency.png", dpi=150); plt.close()


def fig_training(outdir, artdir):
    hists = sorted(glob.glob(f"{artdir}/*_hist.json"))
    if not hists:
        return
    plt.figure(figsize=(5, 3.4))
    for h in hists:
        d = json.load(open(h))
        name = os.path.basename(h).replace("_hist.json", "")
        it = [r["iter"] for r in d["history"]]
        fl = [r["f_final"] for r in d["history"]]
        plt.plot(it, fl, label=name, lw=1)
    plt.xlabel("meta-iteration"); plt.ylabel("f_final (rollout)")
    plt.title("Meta-training progress"); plt.legend(fontsize=7)
    plt.grid(alpha=.3); plt.tight_layout()
    plt.savefig(f"{outdir}/training_curves.png", dpi=150); plt.close()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--eval", default="outputs/eval.json")
    ap.add_argument("--efficiency", default="outputs/efficiency.json")
    ap.add_argument("--artifacts", default="outputs")
    ap.add_argument("--out", default="figures")
    args = ap.parse_args()
    os.makedirs(args.out, exist_ok=True)
    ev = json.load(open(args.eval))
    fig_newton(ev, args.out)
    fig_profiles(ev, args.out)
    fig_convergence(ev, args.out)
    # NOTE: no efficiency figure -- Claim 6 is HMR-specific and not reproduced;
    # a synthetic runtime bar would misrepresent it (see Claim 6 page).
    fig_training(args.out, args.artifacts)
    print("wrote figures to", args.out, sorted(os.listdir(args.out)))


if __name__ == "__main__":
    main()