scripts/make_figures.py

"""Figures for the reproduction poster / logbook."""

import sys, os, json

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import numpy as np
import matplotlib

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

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

ACC = "#0F766E"  # teal accent
ACC2 = "#B45309"  # amber
ACC3 = "#334155"  # slate
plt.rcParams.update(
    {
        "font.size": 13,
        "axes.grid": True,
        "grid.alpha": 0.25,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "figure.dpi": 150,
        "savefig.bbox": "tight",
    }
)


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


c2 = load("claim2_thm56.json")
c3 = load("claim3_thm512.json")
c4 = load("claim4_thm61_62.json")
c4b = load("claim4b_blocks.json")
c5 = load("claim5_covering.json")
c6 = load("claim6_prop51.json")
blk = load("block_sweep.json")

# ============================================================ FIG 1 : minimax rates
fig, ax = plt.subplots(1, 2, figsize=(11.5, 4.3))
h = c3["hard_family_minimax_per_coordinate"]
N = np.array([r["N"] for r in h])
ax[0].loglog(
    N,
    [r["erm_worstcase_per_coord_EXACT"] for r in h],
    "o-",
    color=ACC,
    label="ERM worst case (exact)  slope %.2f" % c3["hard_N_exponent_ERM"][0],
)
ax[0].loglog(
    N,
    [r["sga_worstcase_per_coord_MC"] for r in h],
    "s-",
    color=ACC2,
    label="SGA + averaging  slope %.2f" % c3["hard_N_exponent_SGA"][0],
)
ax[0].loglog(
    N,
    [r["bayes_lower_bound_per_coord"] for r in h],
    "^--",
    color=ACC3,
    label="exact Bayes minimax LB  slope %.2f" % c3["hard_N_exponent_LB"][0],
)
ax[0].set_xlabel("training sample size $N$")
ax[0].set_ylabel("excess risk per coupling constraint")
ax[0].set_title("Minimax regime: rate in $N$ (predicted $-1/2$)", fontsize=13)
ax[0].legend(fontsize=9.5, frameon=False)

sr = c3["hard_family_s_scaling"]
S = np.array([r["s"] for r in sr])
ax[1].loglog(
    S,
    [r["erm_risk"] for r in sr],
    "o-",
    color=ACC,
    label="ERM  slope %.2f" % c3["hard_s_exponent_ERM"][0],
)
ax[1].loglog(
    S,
    [r["sga_risk"] for r in sr],
    "s-",
    color=ACC2,
    label="SGA + averaging  slope %.2f" % c3["hard_s_exponent_SGA"][0],
)
ax[1].loglog(
    S,
    [r["minimax_lb"] for r in sr],
    "^--",
    color=ACC3,
    label="exact minimax LB  slope 1.00",
)
ref = np.array([r["minimax_lb"] for r in sr])[0] * (S / S[0]) ** 1.5
ax[1].loglog(S, ref, ":", color="#94A3B8", label="Thm 5.5 shape $s^{1.5}$")
ax[1].set_xlabel("number of coupling constraints $s$")
ax[1].set_ylabel("excess risk,  $N=256$")
ax[1].set_title("Minimax regime: rate in $s$ (LB predicts $+1$)", fontsize=13)
ax[1].legend(fontsize=9.5, frameon=False)
fig.tight_layout()
fig.savefig(os.path.join(FIG, "fig1_minimax_rates.png"))
plt.close(fig)

# ============================================================ FIG 2 : block MILP family
fig, ax = plt.subplots(1, 2, figsize=(11.5, 4.3))
cells = blk["cells"]
S_LIST = blk["S_LIST"]
cmap = plt.cm.viridis(np.linspace(0.1, 0.85, len(S_LIST)))
for k, s in enumerate(S_LIST):
    sub = [c for c in cells if c["s"] == s]
    ax[0].loglog(
        [c["N"] for c in sub],
        [c["erm_excess_mean"] for c in sub],
        "o-",
        color=cmap[k],
        label="$s=%d$  slope %.2f"
        % (s, blk["fits"]["erm_N_exponent_per_s"][str(s)][0]),
    )
ax[0].set_xlabel("training sample size $N$")
ax[0].set_ylabel("ERM excess risk  $\\mathcal{E}(\\hat\\pi)$")
ax[0].set_title("Block-decomposable MILP family: exact ERM risk", fontsize=13)
ax[0].legend(fontsize=9, frameon=False)

Nsel = blk["N_LIST"][-1]
sub = [c for c in cells if c["N"] == Nsel]
S = np.array([c["s"] for c in sub])
erm = np.array([c["erm_excess_mean"] for c in sub])
ax[1].loglog(
    S,
    erm,
    "o-",
    color=ACC,
    label="ERM measured, slope %.2f"
    % blk["fits"]["erm_s_exponent_per_N"][str(Nsel)][0],
)
ax[1].loglog(
    S,
    [c["thm55_constant_bound"] for c in sub],
    "--",
    color=ACC3,
    label="Thm 5.5 constant $12\\sqrt{\\pi}B\\pi_{\\max}s^{1.5}/\\sqrt{N}$",
)
ax[1].loglog(
    S,
    [c["thm512_constant_bound"] for c in sub],
    ":",
    color=ACC2,
    label="Thm 5.12 constant $2B\\pi_{\\max}s/\\sqrt{N}$",
)
ax[1].set_xlabel("number of coupling constraints $s$")
ax[1].set_ylabel("excess risk,  $N=%d$" % Nsel)
ax[1].set_title("Measured risk vs the theorems' explicit constants", fontsize=13)
ax[1].legend(fontsize=9, frameon=False)
fig.tight_layout()
fig.savefig(os.path.join(FIG, "fig2_block_family.png"))
plt.close(fig)

# ============================================================ FIG 3 : warm start
fig, ax = plt.subplots(1, 2, figsize=(11.5, 4.3))
rows = c4b["rows"]
S_LIST = sorted({r["s"] for r in rows})
cmap = plt.cm.plasma(np.linspace(0.1, 0.8, len(S_LIST)))
for k, s in enumerate(S_LIST):
    sub = [r for r in rows if r["s"] == s]
    ax[0].loglog(
        [r["N"] for r in sub],
        [r["exact_excess_risk"] for r in sub],
        "o-",
        color=cmap[k],
        label="$s=%d$  slope %.2f" % (s, c4b["N_exponents"][str(s)][0]),
    )
ax[0].set_xlabel("training sample size $N$")
ax[0].set_ylabel("warm-start excess risk")
ax[0].set_title(
    "Thm 6.1: $\\mathcal{E}(\\hat\\phi)=O(s/N)$ on block MILPs", fontsize=13
)
ax[0].legend(fontsize=9, frameon=False)

lb = c4["C_bayes_per_coordinate"]
Ns = sorted(int(k) for k in lb.keys())
per = [lb[str(n)]["bayes_per_coordinate"] for n in Ns]
s0 = 16
ax[1].loglog(
    Ns,
    [s0 * p for p in per],
    "^--",
    color=ACC3,
    label="exact minimax LB (Thm 6.2)  slope %.2f" % c4["C_N_exponent"][0],
)
ax[1].loglog(
    Ns,
    [s0 * 1.0 / (4 * n) for n in Ns],
    "--",
    color=ACC2,
    label="Popoviciu UB $s\\pi_{\\max}^2/4N$ (Thm 6.1)",
)
sub16 = [r for r in rows if r["s"] == 16]
ax[1].loglog(
    [r["N"] for r in sub16],
    [r["exact_excess_risk"] for r in sub16],
    "o-",
    color=ACC,
    label="measured ERM risk, block MILPs",
)
ax[1].set_xlabel("training sample size $N$")
ax[1].set_ylabel("excess risk,  $s=16$")
ax[1].set_title("$\\Theta(s/N)$ sandwich: upper and lower bounds", fontsize=13)
ax[1].legend(fontsize=9, frameon=False)
fig.tight_layout()
fig.savefig(os.path.join(FIG, "fig3_warmstart.png"))
plt.close(fig)

# ============================================================ FIG 4 : Prop 5.1 + covering
fig, ax = plt.subplots(1, 2, figsize=(11.5, 4.3))
aud = c6["f_boundary_audit"]
rho = [a["claimed_B_over_true_B"] for a in aud]
ax[0].plot(
    rho,
    [a["max_ratio"] for a in aud],
    "o-",
    color=ACC,
    label="max $\\|g\\|_2/(2B'\\sqrt{s})$",
)
ax[0].plot(
    rho,
    [a["violation_frac"] for a in aud],
    "s-",
    color=ACC2,
    label="fraction of violations",
)
ax[0].axhline(1.0, color=ACC3, ls="--", lw=1)
ax[0].invert_xaxis()
ax[0].set_xlabel("claimed $B'$ / true Assumption-4.1 $B$")
ax[0].set_ylabel("ratio to the bound")
ax[0].set_title(
    "Prop 5.1(ii) boundary audit: the bound fails\nexactly when Assumption 4.1 is broken",
    fontsize=12,
)
ax[0].legend(fontsize=9.5, frameon=False)
ax[0].grid(alpha=0.25)

nets = c5["explicit_nets"]
packs = c5["packing_lower_estimates"]
xs = np.arange(len(nets))
ax[1].bar(
    xs - 0.22,
    [n["lemma53_bound"] for n in nets],
    width=0.22,
    color=ACC3,
    label="Lemma 5.3 bound",
)
ax[1].bar(
    xs,
    [n["log_net_size"] for n in nets],
    width=0.22,
    color=ACC,
    label="explicit $\\delta$-net built (verified)",
)
pk = {(p["s"], p["delta"]): p["log_packing_lower_est"] for p in packs}
ax[1].bar(
    xs + 0.22,
    [pk[(n["s"], n["delta"])] for n in nets],
    width=0.22,
    color=ACC2,
    label="greedy $2\\delta$-packing (lower est.)",
)
ax[1].set_xticks(xs)
ax[1].set_xticklabels(
    ["$s$=%d\n$\\delta$=%g" % (n["s"], n["delta"]) for n in nets], fontsize=8
)
ax[1].set_ylabel("$\\log N(\\delta,\\mathcal{U},\\|\\cdot\\|_{2,N})$")
ax[1].set_title("Lemma 5.3 holds, with 8-14 nats of slack", fontsize=12)
ax[1].legend(fontsize=9, frameon=False)
ax[1].grid(alpha=0.25, axis="y")
fig.tight_layout()
fig.savefig(os.path.join(FIG, "fig4_prop51_covering.png"))
plt.close(fig)

print("figures written to", FIG, os.listdir(FIG))