scripts/make_figs.py

"""Poster figures from the claim JSON outputs. CPU, matplotlib Agg."""

import json
import sys
import numpy as np
import matplotlib

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

OUT = "/home/ubuntu/samuel/sgmcmc-uq-repro/outputs"
FIG = "/home/ubuntu/samuel/sgmcmc-uq-repro/figs"
ACC, ACC2, EMPH, GREY = "#0E6E5C", "#7BC4B4", "#C9782A", "#7A7A7A"
plt.rcParams.update(
    {
        "font.size": 13,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.linewidth": 1.1,
        "figure.dpi": 170,
        "savefig.bbox": "tight",
    }
)


def J(n):
    return json.load(open(f"{OUT}/{n}"))


# ---------------------------------------------------------------- fig 1 ----
c1 = J("claim1_thm41.json")
lam = np.array([r["lam"] for r in c1["sweep"]])
err = np.array([r["rel_cov_err"] for r in c1["sweep"]])
se = np.array([r["rel_cov_err_se"] for r in c1["sweep"]])
sd = np.array([r["rel_sd_err"] for r in c1["sweep"]])
fig, ax = plt.subplots(figsize=(6.4, 4.6))
ax.errorbar(
    lam,
    err,
    yerr=se,
    marker="o",
    color=ACC,
    lw=2.2,
    ms=7,
    capsize=3,
    label=r"$\|\Sigma_\theta-\Sigma_\psi\|_F/\|\Sigma_\theta\|_F$",
)
ax.plot(
    lam,
    sd,
    marker="s",
    color=ACC2,
    lw=2.0,
    ms=6,
    ls="--",
    label=r"$\max_d\,|\sigma_{\theta,d}-\sigma_{\psi,d}|/\sigma_{\theta,d}$",
)
k = err[0] / np.sqrt(lam[0])
ax.plot(
    lam,
    k * np.sqrt(lam),
    color=EMPH,
    lw=2.4,
    ls=":",
    label=r"theorem rate $C_v\lambda^{1/2}$ (anchored at $\lambda_{\min}$)",
)
ax.set_xscale("log")
ax.set_yscale("log")
ax.set_xlabel(r"step size $\lambda$")
ax.set_ylabel("relative error")
ax.set_title(
    f"Thm 4.1: fitted exponent {c1['fitted_exponent_cov']:.2f} " r"$\geq 1/2$",
    fontsize=13,
)
ax.legend(fontsize=10, frameon=False, loc="upper left")
ax.grid(alpha=0.25, which="both")
fig.savefig(f"{FIG}/claim1_thm41.png")
plt.close(fig)

# ---------------------------------------------------------------- fig 2 ----
c3 = J("claim3_w2.json")
fig, axs = plt.subplots(1, 2, figsize=(10.4, 4.3))
L = c3["lambda_sweep"]
x = np.array([r["lam"] for r in L])
u = np.array([r["w2_coupling_upper"] for r in L])
lo = np.array([r["w2_marginal_lower"] for r in L])
axs[0].fill_between(x, lo, u, color=ACC2, alpha=0.45, label="$W_2$ bracket")
axs[0].plot(x, u, "o-", color=ACC, lw=2.2, ms=6, label="synchronous coupling (upper)")
axs[0].plot(x, u[0] * (x / x[0]), ls=":", color=EMPH, lw=2.4, label=r"slope $1$")
axs[0].set_xscale("log")
axs[0].set_yscale("log")
axs[0].set_xlabel(r"$\lambda$   ($B=16$)")
axs[0].set_ylabel(r"$W_2(\pi_\theta,\pi_\psi)$")
axs[0].set_title(
    f"fitted {c3['fitted_exponent_lambda']:.3f} (predicted 1)", fontsize=12
)
Bs = c3["batch_sweep"]
xb = np.array([r["B"] for r in Bs])
ub = np.array([r["w2_coupling_upper"] for r in Bs])
lob = np.array([r["w2_marginal_lower"] for r in Bs])
axs[1].fill_between(xb, lob, ub, color=ACC2, alpha=0.45)
axs[1].plot(xb, ub, "o-", color=ACC, lw=2.2, ms=6)
axs[1].plot(
    xb, ub[0] * (xb / xb[0]) ** -1.0, ls=":", color=EMPH, lw=2.4, label=r"slope $-1$"
)
axs[1].set_xscale("log")
axs[1].set_yscale("log")
axs[1].set_xlabel(r"batch size $B$   ($\lambda=0.05$)")
axs[1].set_title(f"fitted {c3['fitted_exponent_B']:.3f} (predicted $-1$)", fontsize=12)
for a in axs:
    a.grid(alpha=0.25, which="both")
    a.legend(fontsize=10, frameon=False)
fig.suptitle(r"Cor 4.6: $W_2(\pi_\theta,\pi_\psi)\leq A\lambda/B$", fontsize=14)
fig.savefig(f"{FIG}/claim3_w2.png")
plt.close(fig)

# ---------------------------------------------------------------- fig 3 ----
c5 = J("claim5_boston.json")
rows = [r for r in c5["results"] if r["loss"] == "log" and "cov_err_median" in r]
meth = ["CT", "LR+WS", "DQ+const", "DQ+exact", "DQ+exact-cf"]
paper = {
    ("CT", 16): 0.247,
    ("CT", 50): 0.589,
    ("LR+WS", 16): 9.23e8,
    ("LR+WS", 50): 1.40e7,
    ("DQ+exact", 16): 0.337,
    ("DQ+exact", 50): 0.352,
    ("DQ+exact-cf", 16): 0.337,
    ("DQ+exact-cf", 50): 0.352,
}
fig, ax = plt.subplots(figsize=(8.2, 4.4))
w = 0.36
xs = np.arange(len(meth))
for j, B in enumerate([16, 50]):
    vals = []
    for mth in meth:
        r = [q for q in rows if q["method"] == mth and q["B"] == B][0]
        v = r["cov_err_median"]
        vals.append(v if np.isfinite(v) else 1e12)
    ax.bar(
        xs + (j - 0.5) * w,
        vals,
        w,
        color=[ACC, ACC2][j],
        label=f"reproduced, $B={B}$",
        edgecolor="white",
    )
for j, B in enumerate([16, 50]):
    for i, mth in enumerate(meth):
        if (mth, B) in paper:
            ax.plot(
                [xs[i] + (j - 0.5) * w - w / 2, xs[i] + (j - 0.5) * w + w / 2],
                [paper[(mth, B)]] * 2,
                color=EMPH,
                lw=2.6,
                label="paper Table 3" if (i == 0 and j == 0) else None,
            )
ax.set_yscale("log")
ax.set_xticks(xs)
ax.set_xticklabels(
    ["CT", "LR+WS", "DQ+const", "DQ+exact\n(root solve)", "DQ+exact\n(closed form)"],
    fontsize=10,
)
ax.axhline(1.0, color=GREY, lw=1, ls="--")
ax.set_ylabel(r"$\|\hat S-S_\star\|_F/\|S_\star\|_F$   (median of 30)")
ax.set_title(
    "Table 3, Boston housing, log loss  (bars: this reproduction; "
    "orange rules: paper)",
    fontsize=12,
)
ax.legend(fontsize=10, frameon=False, ncol=2)
ax.grid(alpha=0.25, axis="y", which="both")
fig.savefig(f"{FIG}/claim5_boston.png")
plt.close(fig)

# ---------------------------------------------------------------- fig 4 ----
c6 = J("claim6_propB1.json")
g = c6["grid"]
ks = sorted(set(r["kappa"] for r in g))


def agg(key):
    return [np.median([r[key] for r in g if r["kappa"] == kk]) for kk in ks]


fig, ax = plt.subplots(figsize=(6.6, 4.4))
kk = np.array(ks)
kk_p = np.where(kk == 0, 3e-8, kk)
ax.plot(
    kk_p,
    np.maximum(agg("B3_resid_SGD"), 1e-17),
    "o-",
    color=ACC,
    lw=2.2,
    ms=7,
    label=r"Eq. (B.3), SGD ($\beta=\infty$)",
)
ax.plot(
    kk_p,
    np.maximum(agg("B3_resid_SGLD_paper"), 1e-17),
    "s-",
    color=EMPH,
    lw=2.4,
    ms=7,
    label=r"Eq. (B.3) as printed, SGLD ($\beta=N$)",
)
ax.plot(
    kk_p,
    np.maximum(agg("B3_resid_SGLD_corrected"), 1e-17),
    "^--",
    color=ACC2,
    lw=2.2,
    ms=7,
    label="corrected temperature term, SGLD",
)
ax.set_xscale("log")
ax.set_yscale("log")
ax.set_xlabel(r"momentum $\kappa$   ($\kappa=0$ plotted at the left edge)")
ax.set_ylabel("relative residual of Eq. (B.3)")
ax.axhline(1e-14, color=GREY, lw=1, ls=":")
ax.text(4e-8, 2.5e-14, "machine precision", fontsize=9, color=GREY)
ax.set_title(
    "Prop B.1: exact for SGD and as $\\kappa\\to0$; the SGLD\n"
    "temperature term is wrong for $\\kappa>0$",
    fontsize=12,
)
ax.legend(fontsize=9.5, frameon=False, loc="center left")
ax.grid(alpha=0.25, which="both")
fig.savefig(f"{FIG}/claim6_propB1.png")
plt.close(fig)

# ---------------------------------------------------------------- fig 5 ----
c2 = J("claim2_thm43.json")
c4 = J("claim4_alg1.json")
fig, axs = plt.subplots(1, 2, figsize=(10.4, 4.2))
labels = [
    "empirical MC\n(1.8M grads)",
    "exact finite-domain\n(156 configs)",
    "exhaustive minibatch\nenumeration",
]
vals = [
    max(r["rel_err_formula_vs_MC"] for r in c2["T1_empirical_mc"]),
    c2["T2_exact_finite_domain"]["worst_rel_err"],
    max(r["rel_err"] for r in c2["T3_exhaustive"]),
]
axs[0].bar(range(3), vals, color=[ACC2, ACC, ACC], edgecolor="white")
axs[0].set_yscale("log")
axs[0].set_xticks(range(3))
axs[0].set_xticklabels(labels, fontsize=9.5)
axs[0].axhline(1e-14, color=GREY, ls=":", lw=1)
axs[0].set_ylabel("worst relative error of Eq. (12)")
axs[0].set_title("Thm 4.3 verification", fontsize=12)
axs[0].grid(alpha=0.25, axis="y", which="both")
s2 = [r for r in c4["stage2"] if r["B"] in (16, 200)]
mm = ["CT", "LR+WS", "DQ+const", "DQ+exact"]
for j, B in enumerate([16, 200]):
    v = [
        max(
            [q for q in s2 if q["method"] == m and q["B"] == B][0][
                "exact_proxy_cov_vs_target"
            ],
            1e-16,
        )
        for m in mm
    ]
    axs[1].bar(
        np.arange(4) + (j - 0.5) * 0.36,
        v,
        0.36,
        color=[ACC, ACC2][j],
        label=f"$B={B}$",
        edgecolor="white",
    )
axs[1].set_yscale("log")
axs[1].set_xticks(range(4))
axs[1].set_xticklabels(mm, fontsize=10)
axs[1].set_ylabel(r"$\|\Sigma_\psi(\Lambda)-\hat S\|_F/\|\hat S\|_F$")
axs[1].set_title(
    "Alg 1 stage 2: does the tuned $\\Lambda$ hit the target?", fontsize=12
)
axs[1].legend(fontsize=10, frameon=False)
axs[1].grid(alpha=0.25, axis="y", which="both")
fig.savefig(f"{FIG}/claim24_alg1.png")
plt.close(fig)
print("figures written to", FIG)