scripts/make_figs.py

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

import json

import matplotlib

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

OUT = "/home/ubuntu/samuel/ulmc-kl-repro/outputs"
FIG = "/home/ubuntu/samuel/ulmc-kl-repro/figs"
ACC = "#1f6feb"
ACC2 = "#d2691e"
GREY = "#666666"
plt.rcParams.update(
    {
        "font.size": 9,
        "axes.grid": True,
        "grid.alpha": 0.25,
        "figure.dpi": 160,
        "savefig.bbox": "tight",
    }
)


def L(name):
    return json.load(open(f"{OUT}/{name}.json"))


sc = L("scaling2")
pr = L("prescription")

# ---------------------------------------------------------------- fig 1
fig, ax = plt.subplots(1, 3, figsize=(10.5, 3.1))
r = sc["ulmc"]["d_fixed_trH"]["rows"]
d = np.array([x["d"] for x in r], float)
N = np.array([x["N_eps"] for x in r], float)
ax[0].loglog(d, N / N[0], "o-", color=ACC, label="ULMC (measured)")
r2 = sc["rmd"]["d_fixed_trH"]["rows"]
d2 = np.array([x["d"] for x in r2], float)
N2 = np.array([x["N_eps"] for x in r2], float)
ax[0].loglog(d2, N2 / N2[0], "s-", color=ACC2, label="RMD (measured)")
ax[0].loglog(
    d, np.ones_like(d), "--", color=GREY, label=r"dimension-free: $\propto tr(H)^{1/2}$"
)
ax[0].loglog(d, np.sqrt(d / d[0]), ":", color="crimson", label=r"$d^{1/2}$ bound")
ax[0].set_xlabel("dimension $d$   (tr(H)=10, $\\alpha$=0.01, $\\beta$=1 fixed)")
ax[0].set_ylabel("$N_\\epsilon$ / $N_\\epsilon(d{=}11)$")
ax[0].set_title("Claim 1: iterations vs $d$ at fixed tr(H)")
ax[0].legend(fontsize=6.5)

rc = sc["ulmc"]["d_control_H_eq_betaI"]["rows"]
dc = np.array([x["d"] for x in rc], float)
Nc = np.array([x["N_eps"] for x in rc], float)
ax[1].loglog(dc, Nc / Nc[0], "o-", color=ACC, label="ULMC, $H=\\beta I$")
ax[1].loglog(
    dc,
    np.sqrt(dc / dc[0]),
    "--",
    color=GREY,
    label=r"$\propto tr(H)^{1/2}=(\beta d)^{1/2}$",
)
ax[1].set_xlabel("dimension $d$   (tr(H) = $\\beta d$ grows)")
ax[1].set_ylabel("$N_\\epsilon$ / $N_\\epsilon(d{=}10)$")
ax[1].set_title("Control: real $d$-dependence is detected")
ax[1].legend(fontsize=7)

bl = L("bias_law")["dimension_free_bias_bound"]
rat = [e["ratio"] for e in bl["examples"]]
ax[2].axhline(1.0, color="crimson", ls="--", lw=1.2, label="bound $h^2\\,tr(H)/256$")
ax[2].plot(range(len(rat)), rat, "o", color=ACC, ms=4)
ax[2].set_ylim(0, 1.15)
ax[2].set_xlabel("random spectra (200 sampled, $d \\leq 400$)")
ax[2].set_ylabel("bias / ($h^2$ tr(H)/256)")
ax[2].set_title(
    f"exact ULMC bias law\nmax ratio {bl['max_ratio_bias_over_h2trH_div256']:.3f} $\\leq$ 1"
)
ax[2].legend(fontsize=7)
fig.tight_layout()
fig.savefig(f"{FIG}/fig1_dimension_free.png")

# ---------------------------------------------------------------- fig 2
fig, ax = plt.subplots(1, 3, figsize=(10.5, 3.1))
for i, (key, xl, pred_u, pred_r) in enumerate(
    [
        ("eps", "$\\epsilon$", -1.0, -2 / 3),
        ("trH", "tr(H)", 0.5, 1 / 3),
        ("kappa", "$\\kappa$", 1.5, 1.0),
    ]
):
    for sch, col, mk in (("ulmc", ACC, "o"), ("rmd", ACC2, "s")):
        rr = sc[sch][key]["rows"]
        x = np.array([q[key if key != "eps" else "eps"] for q in rr], float)
        y = np.array([q["N_eps"] for q in rr], float)
        Lc = np.array([q["L"] for q in rr], float)
        e = sc[sch][key]["fit_exponent_log_corrected"]
        ax[i].loglog(
            x,
            y / Lc,
            mk + "-",
            color=col,
            ms=4,
            label=f"{sch.upper()} fit {e:+.3f} (thm {sc[sch][key]['predicted']:+.3f})",
        )
    ax[i].set_xlabel(xl)
    ax[i].set_ylabel("$N_\\epsilon\\,/\\,\\log$ factor")
    ax[i].legend(fontsize=6.5)
ax[0].set_title("Claim 2/3: $\\epsilon$ exponent")
ax[1].set_title("tr(H) exponent")
ax[2].set_title("$\\kappa$ exponent  (ULMC: 1.0 vs claimed 1.5)")
fig.tight_layout()
fig.savefig(f"{FIG}/fig2_exponents.png")

# ---------------------------------------------------------------- fig 3
fig, ax = plt.subplots(1, 2, figsize=(7.6, 3.1))
for sch, col, mk in (("ulmc", ACC, "o"), ("rmd", ACC2, "s")):
    rr = pr[sch]["rows"]
    k = np.array([q["kappa"] for q in rr], float)
    sf = np.array([q["shortfall_factor"] for q in rr], float)
    ax[0].loglog(
        k,
        sf,
        mk + "-",
        color=col,
        ms=4,
        label=f"{sch.upper()} (Thm {'4.3' if sch == 'ulmc' else '5.2'})",
    )
ax[0].loglog(k, np.sqrt(k / k[0]), ":", color="crimson", label=r"$\kappa^{1/2}$")
ax[0].axhline(1, color=GREY, ls="--", lw=1)
ax[0].set_xlabel("$\\kappa$")
ax[0].set_ylabel("steps actually needed / steps prescribed")
ax[0].set_title("Theorem's own $(h,N)$ pairing")
ax[0].legend(fontsize=7)
for sch, col, mk, th in (("ulmc", ACC, "o", 0.5), ("rmd", ACC2, "s", 1.0)):
    rr = pr[sch + "_simulated_time"]["rows"]
    k = np.array([q["kappa"] for q in rr], float)
    nh = np.array([q["Nh"] for q in rr], float)
    e = pr[sch + "_simulated_time"]["Nh_exponent_in_kappa"]
    ax[1].loglog(
        k,
        nh,
        mk + "-",
        color=col,
        ms=4,
        label=f"{sch.upper()} measured {e:.2f} (thm {th})",
    )
ax[1].loglog(
    k,
    nh[0] * (k / k[0]) ** 0.5,
    ":",
    color="crimson",
    label="Thm 4.3 implies $\\kappa^{1/2}$",
)
ax[1].set_xlabel("$\\kappa$")
ax[1].set_ylabel("minimum simulated time $N h$")
ax[1].set_title("Required burn-in time")
ax[1].legend(fontsize=6.5)
fig.tight_layout()
fig.savefig(f"{FIG}/fig3_prescription.png")

# ---------------------------------------------------------------- fig 4
lm = L("lemma61")
fig, ax = plt.subplots(1, 2, figsize=(7.6, 3.1))
s = np.geomspace(1.02, 1e3, 400)
ax[0].semilogx(
    s,
    2 * s / (1 + s - np.log(s)),
    color=ACC,
    lw=1.8,
    label="$E_\\mu[p^\\top H p]\\,/\\,(tr H+\\beta KL)$",
)
ax[0].axhline(1.0, color=GREY, ls="--", label="Lemma 6.1 as printed")
ax[0].axhline(
    lm["C_star_closed_form"],
    color="crimson",
    ls="-.",
    label="$C^*=2e^2/(e^2-1)=2.3130$",
)
ax[0].axvline(np.e**2, color="crimson", ls=":", lw=1)
ax[0].set_xlabel("momentum covariance $s$ of $\\mu$  ($H=\\beta I$)")
ax[0].set_ylabel("LHS / RHS")
ax[0].set_title("Claim 5: Lemma 6.1 violated for $s>1$")
ax[0].legend(fontsize=6.5)
rg = lm["random_gaussian_search"]
nq = lm["nonquadratic_nonconvex_V_quadrature"]
lbl = [
    "Gaussian $\\mu$\ngradient",
    "Gaussian $\\mu$\nmomentum",
    "non-quadratic\nnon-convex $V$",
]
val = [rg["worst_ratio_gradient"], rg["worst_ratio_momentum"], nq["worst_ratio"]]
ax[1].bar(lbl, val, color=[ACC, ACC, ACC2], width=0.55)
ax[1].axhline(1.0, color=GREY, ls="--", label="printed constant 1")
ax[1].axhline(lm["C_star_closed_form"], color="crimson", ls="-.", label="$C^*$")
ax[1].set_ylabel("worst observed LHS/RHS")
ax[1].set_title("worst ratio found; none exceeds $C^*$")
ax[1].legend(fontsize=7)
fig.tight_layout()
fig.savefig(f"{FIG}/fig4_lemma61.png")

# ---------------------------------------------------------------- fig 5
try:
    nqd = L("nonquadratic")
    fig, ax = plt.subplots(1, 2, figsize=(7.6, 3.1))
    b = nqd["bias_law_nonquadratic"]
    for sch, col, mk in (
        ("ulmc", ACC, "o"),
        ("lmc", "crimson", "^"),
        ("composite", ACC2, "s"),
    ):
        ax[0].loglog(
            b[sch]["h"],
            b[sch]["floor"],
            mk + "-",
            color=col,
            ms=4,
            label=f"{sch.upper()}  slope {b[sch]['h_exponent']:.2f}",
        )
    ax[0].set_xlabel("step size $h$")
    ax[0].set_ylabel("stationary KL bias (1 stiff coord.)")
    ax[0].set_title("Non-quadratic $V$: ULMC $O(h^2)$ vs LMC $O(h)$")
    ax[0].legend(fontsize=7)
    hh = nqd["head_to_head_nonquadratic"]
    rr = [r for r in hh["rows"] if r["ulmc"]["N_eps"]]
    e = np.array([r["eps"] for r in rr])
    for sch, col, mk in (
        ("ulmc", ACC, "o"),
        ("lmc", "crimson", "^"),
        ("composite", ACC2, "s"),
    ):
        ax[1].loglog(
            e,
            [r[sch]["N_eps"] for r in rr],
            mk + "-",
            color=col,
            ms=4,
            label=f"{sch.upper()} slope {hh[sch + '_eps_exponent']:.2f}",
        )
    ax[1].set_xlabel("$\\epsilon$")
    ax[1].set_ylabel("$N_\\epsilon$")
    ax[1].set_title(
        f"Claim 6: ridge-separable, tr(H)={hh['trH']:.0f} $\\ll$ "
        f"$\\beta d$={hh['beta_times_d']:.0f}"
    )
    ax[1].legend(fontsize=7)
    fig.tight_layout()
    fig.savefig(f"{FIG}/fig5_composite.png")
except FileNotFoundError:
    print("nonquadratic.json not ready")

# ---------------------------------------------------------------- fig 6
gc = L("genconvex")
fig, ax = plt.subplots(1, 2, figsize=(7.6, 3.1))
for sch, col, mk in (("ulmc", ACC, "o"), ("rmd", ACC2, "s")):
    rr = gc[sch]["d_fixed_trH"]["rows"]
    d = np.array([q["d"] for q in rr], float)
    N = np.array([q["N_eps"] for q in rr], float)
    ax[0].loglog(
        d,
        N / N[0],
        mk + "-",
        color=col,
        ms=4,
        label=f"{sch.upper()} exponent {gc[sch]['d_fixed_trH']['fit_exponent']:+.3f}",
    )
ax[0].loglog(d, np.sqrt(d / d[0]), ":", color="crimson", label="$d^{1/2}$")
ax[0].axhline(1, color=GREY, ls="--")
ax[0].set_xlabel("dimension $d$  (tr(H)=6, $W$=3 fixed, $\\alpha\\to0$)")
ax[0].set_ylabel("$N_\\epsilon$ / $N_\\epsilon(d{=}20)$")
ax[0].set_title("Claim 4: dimension-free at $\\alpha\\to0$")
ax[0].legend(fontsize=7)
for sch, col, mk, p in (("ulmc", ACC, "o", -4.0), ("rmd", ACC2, "s", -3.0)):
    rr = gc[sch]["eps"]["rows"]
    e = np.array([q["eps"] for q in rr], float)
    N = np.array([q["N_eps"] for q in rr], float)
    ax[1].loglog(
        e,
        N,
        mk + "-",
        color=col,
        ms=4,
        label=f"{sch.upper()} measured {gc[sch]['eps']['fit_exponent']:.2f}",
    )
    ax[1].loglog(
        e,
        N[-1] * (e / e[-1]) ** p,
        ":",
        color=col,
        lw=1,
        label=f"Thm bound $\\epsilon^{{{p:.0f}}}$",
    )
ax[1].set_xlabel("$\\epsilon$")
ax[1].set_ylabel("$N_\\epsilon$")
ax[1].set_title("$\\alpha=0$ rate: bound holds, far from tight")
ax[1].legend(fontsize=6)
fig.tight_layout()
fig.savefig(f"{FIG}/fig6_genconvex.png")
print("figures written")