scripts/make_figures.py

"""Figures for the logbook / poster.  Reads outputs/*.json written by the claim
scripts.  CPU only, matplotlib Agg."""

import json
import math
import os
import sys

import matplotlib

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

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(ROOT, "outputs")
FIGS = os.path.join(ROOT, "figs")
os.makedirs(FIGS, exist_ok=True)

ACCENT = "#7a1f3d"
ACCENT2 = "#1f6f8b"
GREY = "#555555"
plt.rcParams.update(
    {
        "font.size": 13,
        "axes.grid": True,
        "grid.alpha": 0.25,
        "figure.dpi": 130,
        "savefig.bbox": "tight",
    }
)


def load(name):
    with open(os.path.join(OUT, name)) as f:
        return json.load(f)


# --------------------------------------------------------------------- fig 1
c5 = load("claim5_thm311.json")
tab = c5["worst_case_table"]
fig, ax = plt.subplots(figsize=(6.4, 4.4))
for d, col in zip((8, 32, 128), (ACCENT, ACCENT2, "#c9761c")):
    sub = [t for t in tab if t["d_gamma"] == d]
    ax.loglog(
        [t["n"] for t in sub],
        [t["worst_E_loss"] for t in sub],
        "o-",
        color=col,
        label=f"$d_\\gamma$ = {d}",
        lw=2,
        ms=5,
    )
ref_n = np.array([8, 128000])
ax.loglog(
    ref_n, 0.347 * 8 / ref_n, "--", color=GREY, lw=1.4, label="$0.347\\, d_\\gamma/n$"
)
ax.set_xlabel("sample size $n$")
ax.set_ylabel("worst-case $E[L^\\gamma_D(M)]$")
ax.set_title(
    "Thm 3.11: median-of-three achieves $O(d_\\gamma/n)$\n"
    "fitted exponent $-0.999$ ($R^2 = 1.000$)",
    fontsize=12,
)
ax.legend(fontsize=11)
fig.savefig(os.path.join(FIGS, "fig1_median3_rate.png"))
plt.close(fig)

# --------------------------------------------------------------------- fig 2
c6 = load("claim6_thm312.json")
rows = [r for r in c6["proper_sample_complexity"] if r["d_gamma"] == 32]
x = np.array([r["ln_1_over_eps"] for r in rows])
y = np.array([r["n_star_eps_over_d"] for r in rows])
fit = c6["linear_fit_d32"]
c5sc = [s for s in c5["sample_complexity"] if s["d_gamma"] == 32]
fig, ax = plt.subplots(figsize=(6.4, 4.4))
ax.plot(x, y, "o", color=ACCENT, ms=8, label="proper learner (measured)")
xs = np.linspace(x.min(), x.max(), 50)
ax.plot(
    xs,
    fit["slope_vs_ln_1_over_eps"] * xs + fit["intercept"],
    "-",
    color=ACCENT,
    lw=2,
    label=f"${fit['slope_vs_ln_1_over_eps']:.2f}\\,\\ln(1/\\epsilon) {fit['intercept']:.2f}$  ($R^2$={fit['r2']:.4f})",
)
ax.plot(
    [math.log(1 / s["eps"]) for s in c5sc],
    [s["n_eps_over_d_over_eps"] for s in c5sc],
    "s-",
    color=ACCENT2,
    ms=7,
    lw=2,
    label="median-of-three (measured), flat at 0.347",
)
ax.set_xlabel("$\\ln(1/\\epsilon)$")
ax.set_ylabel("$n^*(\\epsilon)\\,\\epsilon\\,/\\,d_\\gamma$")
ax.set_title(
    "Thm 3.12: proper learning pays an extra $\\ln(1/\\epsilon)$\n"
    "aggregation does not ($d_\\gamma = 32$)",
    fontsize=12,
)
ax.legend(fontsize=10, loc="upper left")
fig.savefig(os.path.join(FIGS, "fig2_proper_log_factor.png"))
plt.close(fig)

# --------------------------------------------------------------------- fig 3
c2 = load("claim2_thm35.json")
c3 = load("claim3_thm38.json")
c4 = load("claim4_thm310.json")
fig, axes = plt.subplots(1, 3, figsize=(13.5, 3.9))

g2 = [g for g in c2["grid_exact"] if g["eps"] == 0.01]
axes[0].bar(
    [str(g["d_gamma"]) for g in g2],
    [g["ratio_E_loss_over_eps"] for g in g2],
    color=ACCENT,
)
axes[0].axhline(1.0, color="k", ls="--", lw=1.5)
axes[0].set_title(
    "Thm 3.5 (proper aggregation)\n$E[L]/\\epsilon$ at $n = d_\\gamma/(32\\epsilon)$",
    fontsize=11,
)
axes[0].set_xlabel("$d_\\gamma$")
axes[0].set_ylim(0, 2.4)

g3 = [g for g in c3["grid"] if g["eps"] == 0.01 and g["m_hypotheses"] == 100]
axes[1].bar(
    [str(g["d_gamma"]) for g in g3], [g["E_loss_over_eps"] for g in g3], color=ACCENT2
)
axes[1].axhline(1.0, color="k", ls="--", lw=1.5)
axes[1].set_title(
    "Thm 3.8 (finite interpolating aggregation)\n$E[L]/\\epsilon$ at $n = d_\\gamma/(128\\epsilon)$, $m=100$",
    fontsize=11,
)
axes[1].set_xlabel("$d_\\gamma$")
axes[1].set_ylim(0, 2.4)

g4 = [g for g in c4["grid"] if g["eps"] == 0.02]
labels = [f"{g['n_prime']}/{g['m_hypotheses']}" for g in g4]
axes[2].bar(labels, [g["E_loss"] for g in g4], color="#c9761c")
axes[2].axhline(0.98, color="k", ls="--", lw=1.5)
axes[2].axhline(c4["random_guessing_uniform_loss"], color=ACCENT, ls=":", lw=2)
axes[2].set_ylim(0.8, 1.02)
axes[2].set_title(
    "Thm 3.10 ($d_\\gamma=\\infty$, OIG $\\leq$ 3)\n$E[L] \\geq 1-\\epsilon$; dotted = random guessing",
    fontsize=11,
)
axes[2].set_xlabel("$n'$ / $m$")
axes[2].tick_params(axis="x", labelrotation=60, labelsize=8)
fig.tight_layout()
fig.savefig(os.path.join(FIGS, "fig3_lower_bounds.png"))
plt.close(fig)

print("wrote", os.listdir(FIGS))

# --------------------------------------------------------------------- fig 4
fig, axes = plt.subplots(1, 2, figsize=(9.2, 3.9))
sep = c3["separation"]
axes[0].bar([str(s["d_gamma"]) for s in sep], [s["gap_factor"] for s in sep], color=ACCENT)
axes[0].set_xlabel("$d_\\gamma$")
axes[0].set_ylabel("finite-agg loss / general-learner loss")
axes[0].set_title("Thm 3.8 separation grows with $d_\\gamma$\n(at $n = d_\\gamma/(128\\epsilon)$, $\\epsilon = 0.01$)", fontsize=11)

thr = c3["empirical_threshold"]
for d, col in zip((32, 64, 128), (ACCENT, ACCENT2, "#c9761c")):
    sub = [t for t in thr if t["d_gamma"] == d]
    axes[1].loglog([t["d_gamma"] / t["eps"] for t in sub], [t["n_star"] for t in sub],
                   "o-", color=col, lw=2, ms=6, label=f"$d_\\gamma$ = {d}")
xs = np.array([500, 15000])
axes[1].loglog(xs, 0.339 * xs, "--", color=GREY, lw=1.4, label="$0.339\\, d_\\gamma/\\epsilon$")
axes[1].set_xlabel("$d_\\gamma/\\epsilon$")
axes[1].set_ylabel("$n^*(\\epsilon)$ for finite aggregation")
axes[1].set_title("Threshold collapses onto $\\Theta(d_\\gamma/\\epsilon)$\nCV of $n^*\\epsilon/d_\\gamma$ = 1.2%", fontsize=11)
axes[1].legend(fontsize=10)
fig.tight_layout()
fig.savefig(os.path.join(FIGS, "fig4_separation.png"))
plt.close(fig)
print("fig4 done")