repro-bundle/v1/experiments/make_plots.py

"""Generate all plots for the reproduction logbook."""
import json
import sys
from pathlib import Path

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

OUT = Path(__file__).resolve().parent.parent / "outputs"


def load(name):
    return json.loads((OUT / name).read_text())


def fig_m5_asymptotic():
    """Claim 2: E[tau] vs log(1/alpha) for m=5, with the 1/D asymptote."""
    d = load("alpha_sweep_m5.json")
    th = d["theory"]
    L = np.array([r["log_inv_alpha"] for r in d["alpha_results"]])
    E = np.array([r["mean_tau"] for r in d["alpha_results"]])
    Estd = np.array([r["std_tau"] for r in d["alpha_results"]])
    rate = th["asymptotic_rate"]
    fig, ax = plt.subplots(figsize=(7, 5))
    ax.errorbar(L, E, yerr=Estd, fmt="o-", color="C0", capsize=3,
                label=r"measured $\mathbb{E}[\tau_\alpha]$")
    ax.plot(L, L / th["D_M_inf"], "k--", lw=1.5,
            label=rf"asymptotic LB $\log(1/\alpha)/D_M^\inf$ (slope $1/D={rate:.3f}$)")
    ax.plot(L, L / th["naive_BPI_denom"], "r:", lw=1.5,
            label=rf"naive BPI $\log(1/\alpha)/\|f\|_\infty$ (slope ${1/th['naive_BPI_denom']:.3f}$)")
    ax.set_xscale("log"); ax.set_yscale("log")
    ax.set_xlabel(r"$\log(1/\alpha)$"); ax.set_ylabel(r"expected stopping time $\mathbb{E}[\tau_\alpha]$")
    ax.set_title(f"Claim 2 (m=5): $\\mathbb{{E}}[\\tau_\\alpha]$ tracks $\\log(1/\\alpha)/D_M^\\inf$ as $\\alpha\\to 0$\n"
                 f"$D_M^\\inf={th['D_M_inf']:.3f}$, $1/D={rate:.3f}$, naive $1/\\|f\\|_\\infty={1/th['naive_BPI_denom']:.3f}$")
    ax.grid(True, which="both", alpha=0.3); ax.legend(loc="upper left", fontsize=9)
    fig.tight_layout(); fig.savefig(OUT / "fig_claim2_m5.png", dpi=130)
    plt.close(fig)
    print("saved fig_claim2_m5.png")


def fig_ratio_convergence():
    """Claim 1+2: ratio E[tau]/log(1/alpha) -> 1/D as alpha->0 (m=5),
    and stays far above 1/D for m=64 (overhead-dominated, non-asymptotic)."""
    d5 = load("alpha_sweep_m5.json")
    d64 = load("alpha_sweep_m64.json")
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5))
    for d, ax, name in [(d5, ax1, "m=5"), (d64, ax2, "m=64")]:
        th = d["theory"]
        L = np.array([r["log_inv_alpha"] for r in d["alpha_results"]])
        R = np.array([r["mean_tau"] / r["log_inv_alpha"] for r in d["alpha_results"]])
        ax.semilogx(L, R, "o-", color="C0", label=r"measured $\mathbb{E}[\tau]/\log(1/\alpha)$")
        ax.axhline(th["asymptotic_rate"], color="k", ls="--",
                   label=rf"asymptotic $1/D_M^\inf={th['asymptotic_rate']:.3f}$")
        ax.axhline(1.0 / th["naive_BPI_denom"], color="r", ls=":",
                   label=rf"naive BPI $1/\|f\|_\infty={1/th['naive_BPI_denom']:.3f}$")
        ax.set_xlabel(r"$\log(1/\alpha)$")
        ax.set_ylabel(r"$\mathbb{E}[\tau_\alpha]/\log(1/\alpha)$")
        ax.set_title(f"{name}: $D_M^\\inf={th['D_M_inf']:.3f}$, "
                     f"$C_Q={th['C_Q']:.1f}$, $\\pi^*={th['pi_min']:.2e}$")
        ax.grid(True, which="both", alpha=0.3); ax.legend(fontsize=9)
    fig.suptitle("Claim 1+2: ratio $\\mathbb{E}[\\tau]/\\log(1/\\alpha)$ "
                 "$\\to 1/D_M^\\inf$ as $\\alpha\\to 0$ (left); overhead-dominated "
                 "for large m (right)", y=1.02)
    fig.tight_layout(); fig.savefig(OUT / "fig_ratio_convergence.png", dpi=130, bbox_inches="tight")
    plt.close(fig)
    print("saved fig_ratio_convergence.png")


def fig_m64_table3():
    """Claim 1 (non-asymptotic): m=64 E[tau] nearly flat across alpha (overhead
    dominates), reproducing paper Table 3's pattern. Compare to asymptotic LB."""
    d = load("alpha_sweep_m64.json")
    th = d["theory"]
    L = np.array([r["log_inv_alpha"] for r in d["alpha_results"]])
    E = np.array([r["mean_tau"] for r in d["alpha_results"]])
    Estd = np.array([r["std_tau"] for r in d["alpha_results"]])
    fig, ax = plt.subplots(figsize=(7, 5))
    ax.errorbar(L, E, yerr=Estd, fmt="o-", color="C0", capsize=3,
                label=r"measured $\mathbb{E}[\tau_\alpha]$ (m=64)")
    ax.plot(L, L / th["D_M_inf"], "k--", lw=1.5,
            label=rf"asymptotic LB $\log(1/\alpha)/D_M^\inf$")
    ax.set_xscale("log")
    ax.set_xlabel(r"$\log(1/\alpha)$"); ax.set_ylabel(r"$\mathbb{E}[\tau_\alpha]$")
    ax.set_title("Claim 1 (m=64): $\\mathbb{E}[\\tau_\\alpha]$ is nearly flat in $\\alpha$\n"
                 "(overhead $(m-1)\\psi_t$ dominates $\\log(1/\\alpha)$; "
                 "matches paper Table 3 pattern)")
    ax.grid(True, which="both", alpha=0.3); ax.legend(fontsize=9)
    fig.tight_layout(); fig.savefig(OUT / "fig_claim1_m64.png", dpi=130)
    plt.close(fig)
    print("saved fig_claim1_m64.png")


def fig_claim3_tracking():
    """Claim 3 (part C): E[tau] tracks 1/D_M^inf across theta_Q."""
    d = load("claim3_ablation.json")
    rows = d["C"]["rows"]
    D = np.array([r["D_M_inf"] for r in rows])
    E = np.array([r["mean_tau"] for r in rows])
    Estd = np.array([r["std_tau"] for r in rows])
    invD = 1.0 / D
    alpha = d["C"]["alpha"]
    fig, ax = plt.subplots(figsize=(7, 5))
    ax.errorbar(invD, E, yerr=Estd, fmt="o-", color="C0", capsize=3,
                label=r"measured $\mathbb{E}[\tau_\alpha]$")
    ax.plot(invD, invD * np.log(1.0 / alpha), "k--", lw=1.5,
            label=rf"asymptotic LB $\log(1/\alpha)/D_M^\inf$ ($\alpha={alpha}$)")
    ax.set_xlabel(r"$1/D_M^\inf(Q,\mathcal{P})$"); ax.set_ylabel(r"$\mathbb{E}[\tau_\alpha]$")
    ax.set_title("Claim 3: $\\mathbb{E}[\\tau]$ tracks $1/D_M^\\inf$ across $\\theta_Q$\n"
                 "(both $\\pi_Q$ and the transition rows vary; bound stays predictive)")
    ax.grid(True, alpha=0.3); ax.legend(fontsize=9)
    fig.tight_layout(); fig.savefig(OUT / "fig_claim3_tracking.png", dpi=130)
    plt.close(fig)
    print("saved fig_claim3_tracking.png")


def fig_mcmc_trajectory():
    """MCMC (Figure 1): L_t trajectory under Q_bad (slope D_M^inf) vs Q_good (flat)."""
    d = load("mcmc_figure1.json")
    tb = d["traj_bad_run0"]; tg = d["traj_good_run0"]
    D_bad = d["D_M_inf_Q_bad"]
    fig, ax = plt.subplots(figsize=(8, 5))
    t = np.array(tb["t"]); L = np.array(tb["L"]); beta = np.array(tb["beta"])
    ax.plot(t, L, "C3-", label=r"$L_t$ under $Q_{bad}$ (alternative)")
    ax.plot(t, beta, "k--", label=r"boundary $\beta_t$")
    ax.plot(t, D_bad * t, "C3:", lw=1.5,
            label=rf"theoretical slope $D_M^\inf={D_bad:.4f}$")
    tg_t = np.array(tg["t"]); tg_L = np.array(tg["L"]); tg_beta = np.array(tg["beta"])
    ax.plot(tg_t, tg_L, "C2-", label=r"$L_t$ under $Q_{good}$ (null)")
    ax.plot(tg_t, tg_beta, "k--", alpha=0.4)
    ax.set_xlabel("time $t$"); ax.set_ylabel("statistic $L_t$")
    ax.set_title("MCMC misspecification (paper Figure 1): $L_t$ grows with slope "
                 "$D_M^\\inf$ under alt,\nstays below boundary under null (0/20 false rejections)")
    ax.grid(True, alpha=0.3); ax.legend(fontsize=9)
    fig.tight_layout(); fig.savefig(OUT / "fig_mcmc_figure1.png", dpi=130)
    plt.close(fig)
    print("saved fig_mcmc_figure1.png")


def main():
    fig_m5_asymptotic()
    fig_ratio_convergence()
    fig_m64_table3()
    fig_claim3_tracking()
    fig_mcmc_trajectory()


if __name__ == "__main__":
    main()