repro-bundle-v0/experiments/make_plots.py

import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402

RESULTS_DIR = Path(__file__).resolve().parents[1] / "results"


def plot_exp1():
    d = json.load(open(RESULTS_DIR / "exp1_bilevel_rate.json"))
    det_rows = d["deterministic"]["rows"]
    sto_rows = d["stochastic"]["rows"]
    det_slopes = d["deterministic"]["fitted_slopes"]
    sto_slopes = d["stochastic"]["fitted_slopes"]

    fig, ax = plt.subplots(1, 2, figsize=(11, 4.2))
    Ks_d = [r["K"] for r in det_rows]; grad_d = [r["grad_norm_sq"] for r in det_rows]
    Ks_s = [r["K"] for r in sto_rows]; grad_s = [r["grad_norm_sq"] for r in sto_rows]
    ax[0].loglog(Ks_d, grad_d, "o--", label=f"deterministic (slope={det_slopes['grad_norm_sq']:.2f})", color="#999999")
    ax[0].loglog(Ks_s, grad_s, "o-", label=f"stochastic (slope={sto_slopes['grad_norm_sq']:.2f})", color="#2166ac")
    ref = grad_s[0] * (np.array(Ks_s) / Ks_s[0]) ** (-2 / 3)
    ax[0].loglog(Ks_s, ref, "--", label=r"$O(K^{-2/3})$ (Thm 3.4)", color="#b2182b")
    ax[0].set_xlabel("K (iterations)"); ax[0].set_ylabel(r"$\min_k \|\nabla \mathcal{F}^*(\bar x^k)\|^2$")
    ax[0].set_title("Claim 2: Theorem 3.4 hypergradient rate"); ax[0].legend(fontsize=8); ax[0].grid(True, which="both", alpha=0.3)

    consx_d = [r["cons_x"] for r in det_rows]; consx_s = [r["cons_x"] for r in sto_rows]
    ax[1].loglog(Ks_d, consx_d, "o--", label=f"deterministic (slope={det_slopes['cons_x']:.1f})", color="#999999")
    ax[1].loglog(Ks_s, consx_s, "o-", label=f"stochastic (slope={sto_slopes['cons_x']:.2f})", color="#2166ac")
    ref2 = consx_s[0] * (np.array(Ks_s) / Ks_s[0]) ** (-1)
    ax[1].loglog(Ks_s, ref2, "--", label=r"$O(K^{-1})$ (Prop 3.5)", color="#b2182b")
    ax[1].set_xlabel("K (iterations)"); ax[1].set_ylabel(r"$\min_k \frac{1}{n}\sum_i\|x_i^k-\bar x^k\|^2$")
    ax[1].set_title("Claim 4: Proposition 3.5 consensus rate"); ax[1].legend(fontsize=8); ax[1].grid(True, which="both", alpha=0.3)
    fig.tight_layout()
    fig.savefig(RESULTS_DIR / "exp1_bilevel_rate.png", dpi=140)
    plt.close(fig)


def plot_exp2():
    d = json.load(open(RESULTS_DIR / "exp2_pushpull_rate.json"))
    rows = d["rows"]
    Ks = [r["K"] for r in rows]
    grad = [r["grad_norm_sq"] for r in rows]
    cons = [r["cons"] for r in rows]
    sens = d["eta_sensitivity"]
    stoch = d["stochastic"]

    fig, ax = plt.subplots(1, 2, figsize=(11, 4.2))
    ax[0].loglog(Ks, grad, "o--", label=r"deterministic grad$^2$ ($\eta_x$=0.01)", color="#999999")
    ax[0].loglog(stoch["checkpoints"], stoch["grad_norm_sq"], "o-",
                 label=f"stochastic grad$^2$ (slope={stoch['fitted_slope']:.2f})", color="#2166ac")
    ref = stoch["grad_norm_sq"][0] * (np.array(stoch["checkpoints"]) / stoch["checkpoints"][0]) ** (-1)
    ax[0].loglog(stoch["checkpoints"], ref, "--", label=r"$O(K^{-1})$ (Thm 3.6)", color="#b2182b")
    ax[0].set_xlabel("K (iterations)"); ax[0].set_ylabel("metric")
    ax[0].set_title("Claim 3: Theorem 3.6 Push-Pull rate"); ax[0].legend(fontsize=8); ax[0].grid(True, which="both", alpha=0.3)

    etas = [r["eta_x"] for r in sens]
    c25600 = [max(r["cons_at_25600"], 1e-40) for r in sens]
    ax[1].semilogy(etas, c25600, "o-", color="#762a83")
    ax[1].set_xlabel(r"fixed step size $\eta_x$"); ax[1].set_ylabel(r"consensus error at $K$=25600")
    ax[1].set_title("Step-size threshold (Thm 3.6)", fontsize=11)
    ax[1].grid(True, which="both", alpha=0.3)
    fig.tight_layout()
    fig.savefig(RESULTS_DIR / "exp2_pushpull_rate.png", dpi=140)
    plt.close(fig)


def plot_exp3():
    d = json.load(open(RESULTS_DIR / "exp3_fashion_mnist.json"))
    fig, ax = plt.subplots(1, 2, figsize=(11, 4.2))
    colors = {"cr_0.3": "#2166ac", "cr_0.5": "#b2182b"}
    for cr, v in d.items():
        ks_f = [p[0] for p in v["fab_test_acc"]]
        acc_f = [p[1] for p in v["fab_test_acc"]]
        ks_b = [p[0] for p in v["baseline_test_acc"]]
        acc_b = [p[1] for p in v["baseline_test_acc"]]
        c = colors.get(cr, "black")
        ax[0].plot(ks_f, acc_f, "-", color=c, label=f"FAB ({cr})")
        ax[0].plot(ks_b, acc_b, "--", color=c, label=f"baseline ({cr})")
    ax[0].set_xlabel("iteration"); ax[0].set_ylabel("held-out test accuracy")
    ax[0].set_title("Claim 5: Fashion-MNIST hyper-cleaning"); ax[0].legend(fontsize=8); ax[0].grid(alpha=0.3)

    crs = list(d.keys())
    w_corrupt = [d[cr]["learned_weight_by_corruption"]["mean_weight_corrupted_examples"] for cr in crs]
    w_clean = [d[cr]["learned_weight_by_corruption"]["mean_weight_clean_examples"] for cr in crs]
    xpos = np.arange(len(crs))
    ax[1].bar(xpos - 0.18, w_corrupt, width=0.35, label="corrupted examples", color="#b2182b")
    ax[1].bar(xpos + 0.18, w_clean, width=0.35, label="clean examples", color="#2166ac")
    ax[1].set_xticks(xpos); ax[1].set_xticklabels(crs)
    ax[1].set_ylabel("mean learned weight sigmoid(x)")
    ax[1].set_title("FAB down-weights corrupted examples"); ax[1].legend(fontsize=8); ax[1].grid(alpha=0.3, axis="y")
    fig.tight_layout()
    fig.savefig(RESULTS_DIR / "exp3_fashion_mnist.png", dpi=140)
    plt.close(fig)


if __name__ == "__main__":
    plot_exp1()
    plot_exp2()
    plot_exp3()
    print("Saved plots to", RESULTS_DIR)