repro-bundle-v0/repro/make_figures.py

"""Figures for the ConstantStepsizeSA logbook."""

import json
import os

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

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
RES = os.path.join(ROOT, "results")
FIGS = os.path.join(HERE, "figures")
os.makedirs(FIGS, exist_ok=True)

R = json.load(open(os.path.join(RES, "claims.json")))


def ratio(alphas, w1):
    a = np.array(alphas)
    return np.array(w1) / (np.sqrt(a) * np.log(1.0 / a))


def fig_rates():
    fig, ax = plt.subplots(1, 2, figsize=(11, 4.4))
    series = [
        ("SGD (iid, Prop 3.1)", R["claim2_sgd"]["alphas"],
         R["claim2_sgd"]["w1"], R["claim2_sgd"]["slope"], "#2980b9"),
        ("linear SA (Prop 3.2)", R["claim3"]["linear"]["alphas"],
         R["claim3"]["linear"]["w1"], R["claim3"]["linear"]["slope"],
         "#27ae60"),
        ("contractive SA (Prop 3.3)", R["claim3"]["contractive"]["alphas"],
         R["claim3"]["contractive"]["w1"],
         R["claim3"]["contractive"]["slope"], "#8e44ad"),
        ("SGD (Markov, Prop 4.1)", R["claim5_markov"]["alphas"],
         R["claim5_markov"]["w1_correct"], R["claim5_markov"]["slope"],
         "#c0392b"),
    ]
    for name, al, w1, m, c in series:
        ax[0].loglog(al, w1, "o-", color=c, label=f"{name}: slope {m:.2f}")
    a = np.array(R["claim2_sgd"]["alphas"])
    ax[0].loglog(a, 0.16 * np.sqrt(a) * np.log(1 / a), "k--", lw=1,
                 label=r"$U\sqrt{\alpha}\log(1/\alpha)$")
    ax[0].set_xlabel(r"stepsize $\alpha$")
    ax[0].set_ylabel(r"$W_1(Y^{(\alpha)}, \mathcal{N}(0,\Sigma_Y))$")
    ax[0].set_title("Wasserstein rate across all four settings")
    ax[0].legend(fontsize=8)
    for name, al, w1, m, c in series:
        ax[1].semilogx(al, ratio(al, w1), "o-", color=c, label=name)
    ax[1].set_xlabel(r"$\alpha$")
    ax[1].set_ylabel(r"$W_1 / (\sqrt{\alpha}\log(1/\alpha))$")
    ax[1].set_title("Bound ratio: bounded (in fact decreasing)")
    ax[1].legend(fontsize=8)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "w1_rates.png"), dpi=140)
    plt.close(fig)


def fig_tails():
    t = R["claim4_tails"]
    fig, ax = plt.subplots(1, 2, figsize=(11, 4.4))
    for row in t["rows"]:
        ax[0].semilogy(t["a_levels"], row["deltas"], "o-",
                       label=rf"$\alpha$={row['alpha']}")
    ax[0].set_xlabel("deviation level a")
    ax[0].set_ylabel(r"$|P(\langle Y,\zeta\rangle > a\sigma) - P(Z > a)|$")
    ax[0].set_title("Tail error decays in the deviation level")
    ax[0].legend(fontsize=8)
    ax[1].loglog(t["alphas"], t["sup_scaled"], "o-", color="#c0392b",
                 label=f"sup over a of a*|Delta|: slope {t['slope']:.2f}")
    a = np.array(t["alphas"])
    ax[1].loglog(a, 0.09 * a ** 0.25 * np.sqrt(np.log(1 / a)), "k--", lw=1,
                 label=r"$C\,\alpha^{1/4}\log^{1/2}(1/\alpha)$")
    ax[1].set_xlabel(r"$\alpha$")
    ax[1].set_ylabel(r"$\sup_a a\,|\Delta(a)|$")
    ax[1].set_title("Berry-Esseen scaling in the stepsize")
    ax[1].legend(fontsize=8)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "tails.png"), dpi=140)
    plt.close(fig)


def fig_markov():
    m = R["claim5_markov"]
    fig, ax = plt.subplots(figsize=(6, 4.4))
    ax.loglog(m["alphas"], m["w1_correct"], "o-", color="#27ae60",
              label=r"vs $\mathcal{N}(0,\Sigma_Y)$ (long-run cov)")
    ax.loglog(m["alphas"], m["w1_naive"], "s--", color="#c0392b",
              label=r"vs naive marginal-cov Gaussian (floor $\approx$ 0.09)")
    ax.set_xlabel(r"$\alpha$")
    ax.set_ylabel(r"$W_1$")
    ax.set_title("Markovian noise: correct vs naive Gaussian limit")
    ax.legend(fontsize=9)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "markov.png"), dpi=140)
    plt.close(fig)


def fig_gibbs():
    g = R["claim6_gibbs"]
    fig, ax = plt.subplots(1, 2, figsize=(11, 4.4))
    for h, c in [(2, "#2980b9"), (4, "#c0392b")]:
        d = g[f"h{h}"]
        ax[0].loglog(d["alphas"], d["mean_abs"], "o-", color=c,
                     label=(f"h={h}: slope {d['scale_slope']:.4f} "
                            f"(theory {d['theory_slope']:.2f})"))
    ax[0].set_xlabel(r"$\alpha$")
    ax[0].set_ylabel(r"$\mathbb{E}|X^{(\alpha)}|$")
    ax[0].set_title(r"Scaling $\alpha^{1/h}$ (Prop 5.1)")
    ax[0].legend(fontsize=9)
    d4 = g["h4"]
    ax[1].semilogx(d4["alphas"], d4["w1_gibbs"], "o-", color="#27ae60",
                   label="W1 to Gibbs law")
    ax[1].semilogx(d4["alphas"], d4["w1_best_gauss"], "s--", color="#c0392b",
                   label="W1 to best-fit Gaussian")
    ax[1].set_xlabel(r"$\alpha$")
    ax[1].set_ylabel(r"$W_1$ of $X/\alpha^{1/4}$")
    ax[1].set_title(f"h=4 limit is Gibbs, not Gaussian "
                    f"(kurtosis {g['h4_kurtosis']:.2f} vs 2.19)")
    ax[1].legend(fontsize=9)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "gibbs.png"), dpi=140)
    plt.close(fig)


if __name__ == "__main__":
    fig_rates()
    fig_tails()
    fig_markov()
    fig_gibbs()
    print("figures written to", FIGS)