scripts/make_poster_figs.py

"""Single-panel poster figures (larger fonts, no titles -- captions live on the poster)."""

import json
import os

import matplotlib

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

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(ROOT, "outputs")
P = os.path.join(ROOT, "poster")
os.makedirs(P, exist_ok=True)
plt.rcParams.update(
    {
        "font.size": 15,
        "figure.dpi": 150,
        "axes.grid": True,
        "grid.alpha": 0.3,
        "axes.labelsize": 16,
    }
)

# ---------------- eOT ----------------
d = json.load(open(os.path.join(OUT, "summary_456.json")))["claim4"]
fig, ax = plt.subplots(figsize=(7.2, 4.6))
for c in d["configs"]:
    if c["mode"] == "dual" and c["C"] == 1e-3:
        lbl = rf"baseline dual, $\sigma^2$={c['sigma_sq']:g}, C=$10^{{-3}}$"
        col = "#c0392b"
    elif c["mode"] == "semidual" and c["alpha_init"] == "paper_code":
        lbl = rf"proposed semi-dual, $\rho$={c['rho']:g}, C={c['C']:g}"
        col = "#2c7f4b"
    else:
        continue
    it = np.array(c["iters"])
    m = np.array(c["mean_gap"])
    s = np.array(c["std_gap"])
    ls = "--" if c["mode"] == "dual" else "-"
    ax.plot(
        it, m, ls, color=col, lw=2.2, alpha=0.55 if c["rho"] == 0.03 else 1.0, label=lbl
    )
    ax.fill_between(it, np.maximum(m - s, 1e-4), m + s, color=col, alpha=0.12)
ax.set_xscale("log")
ax.set_yscale("log")
ax.set_xlabel("kernel-SGD iteration")
ax.set_ylabel("optimality gap")
ax.legend(fontsize=11, loc="lower left")
fig.tight_layout()
fig.savefig(os.path.join(P, "fig_eot.png"), bbox_inches="tight")
plt.close(fig)

# ---------------- stepsize heatmap ----------------
d5 = json.load(open(os.path.join(OUT, "summary_456.json")))["claim5"]
rows = d5["rows"]
batches = [10, 100, 1000]
labels = ["baseline"] + [
    rf"proposed $\rho$={r:g}" for r in [1e-4, 1e-3, 1e-2, 1e-1, 0.5, 0.9]
]
M = np.full((len(labels), 3), np.nan)
for r in rows:
    lb = "baseline" if r["method"] == "baseline" else rf"proposed $\rho$={r['rho']:g}"
    M[labels.index(lb), batches.index(r["batch"])] = (
        np.log10(r["max_stable_eta"]) if r["max_stable_eta"] else np.nan
    )
fig, ax = plt.subplots(figsize=(7.0, 4.8))
im = ax.imshow(M, cmap="viridis", aspect="auto", vmin=-8, vmax=-4)
ax.set_xticks(range(3))
ax.set_xticklabels([f"|D|={b}" for b in batches], fontsize=15)
ax.set_yticks(range(len(labels)))
ax.set_yticklabels(labels, fontsize=13)
for i in range(M.shape[0]):
    for j in range(M.shape[1]):
        ax.text(
            j,
            i,
            "diverges" if np.isnan(M[i, j]) else rf"$10^{{{int(M[i,j])}}}$",
            ha="center",
            va="center",
            color="k" if np.isnan(M[i, j]) else "w",
            fontsize=14,
            fontweight="bold",
        )
ax.grid(False)
cb = fig.colorbar(im, ax=ax)
cb.set_label(r"$\log_{10}$ largest stable $\eta$", fontsize=13)
fig.tight_layout()
fig.savefig(os.path.join(P, "fig_stepsize.png"), bbox_inches="tight")
plt.close(fig)

# ---------------- MNIST ----------------
d6 = json.load(open(os.path.join(OUT, "claim6_uot_dro_big.json")))
recs = d6["runs"]
fig, ax = plt.subplots(figsize=(7.2, 4.6))
colors = {"baseline": "#c0392b", "proposed": "#2c7f4b", "erm": "#7f8c8d"}
styles = {1e-3: "-", 1e-4: "--", 1e-5: ":", 1e-6: "-."}
for meth in ["baseline", "proposed", "erm"]:
    for lr in sorted({r["lr"] for r in recs if r["method"] == meth}, reverse=True):
        rs = [r for r in recs if r["method"] == meth and r["lr"] == lr]
        if meth == "baseline" and lr in (1e-6,):
            continue
        L = min(len(r["hist"]) for r in rs)
        steps = [h["step"] for h in rs[0]["hist"][:L]]
        F = np.array(
            [
                [h["F"] if h["F"] is not None else np.nan for h in r["hist"][:L]]
                for r in rs
            ]
        )
        div = any(r["diverged"] for r in rs)
        ax.plot(
            steps,
            np.nanmean(F, 0),
            styles.get(lr, "-"),
            color=colors[meth],
            lw=2.2,
            label=rf"{meth} $\eta$={lr:g}" + (" - overflow" if div else ""),
        )
        if div:
            k = (
                np.nanargmin(np.where(np.isnan(np.nanmean(F, 0)), np.inf, 0))
                if False
                else len(steps) - 1
            )
            ax.plot(
                [steps[k]], [np.nanmean(F, 0)[k]], "x", color=colors[meth], ms=14, mew=3
            )
ax.set_xlabel("SGD step (batch size 1)")
ax.set_ylabel(r"train objective $F(\theta)$, Eq. (20)")
ax.legend(fontsize=11)
fig.tight_layout()
fig.savefig(os.path.join(P, "fig_mnist.png"), bbox_inches="tight")
plt.close(fig)
print("poster figures written")