repro-bundle/v1/experiments/exact_bound.py

"""Compute the EXACT non-asymptotic lower bound (Theorem 3.3, first inequality)
for the m=5 parametric instance, using the actual Poisson solution and the
empirical distributions of X_0 and X_tau.

The exact bound is:
  E[tau] >= log(1/alpha)/D_M(Q, P*) - E_Q[omega(X_0) - omega(X_tau)] / D_M(Q, P*)

This is much tighter than the worst-case bound (eq 3) which upper-bounds the
correction by 2*C_Q/pi* (very loose when pi* is small).
"""
import json
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
import sequential_test as st
from run_alpha_sweep import build_instance, theory_quantities


def exact_bound_for_alpha(Q, P_star, pi_Q, alpha, dist_X0, dist_Xtau):
    """Theorem 3.3 first inequality: exact non-asymptotic bound."""
    f_P = st.f_P_vector(Q, P_star)
    omega = st.poisson_solution(Q, f_P, pi_Q)
    D_M_val = float(pi_Q @ f_P)
    E_w0 = float(dist_X0 @ omega)
    E_wtau = float(dist_Xtau @ omega)
    correction = E_w0 - E_wtau
    bound = np.log(1.0 / alpha) / D_M_val - correction / D_M_val
    return {
        "alpha": alpha,
        "log_inv_alpha": float(np.log(1.0 / alpha)),
        "D_M": D_M_val,
        "E_omega_X0": E_w0,
        "E_omega_Xtau": E_wtau,
        "correction_E_w0_minus_E_wtau": correction,
        "correction_over_D": correction / D_M_val,
        "bound_raw": float(bound),
        "bound_clipped": float(max(0.0, bound)),
    }


def main():
    out_dir = Path(__file__).resolve().parent.parent / "outputs"
    m = 5
    theta_Q = -0.6
    seed = 123
    P0, f, Q = build_instance(m, theta_Q, (0.4, 0.8), seed)
    pi_Q = st.stationary_dist(Q)
    tq = theory_quantities(Q, (0.4, 0.8), P0, f)
    P_star = st.build_P_theta(tq["theta_star"], P0, f)[0]

    # dist_X0 = uniform (as in the paper's experiments)
    dist_X0 = np.ones(m) / m

    # Load the m5 sweep to get empirical dist of X_tau per alpha
    sweep = json.loads((out_dir / "alpha_sweep_m5.json").read_text())
    # we didn't save X_tau distribution; approximate dist_Xtau ~ pi_Q (stationary)
    # since tau >> mixing time for small alpha. This is the paper's own approximation.
    dist_Xtau_approx = pi_Q  # conservative: omega(pi) = 0 by PE normalization

    alphas = [r["alpha"] for r in sweep["alpha_results"]]
    mean_taus = {r["alpha"]: r["mean_tau"] for r in sweep["alpha_results"]}

    print("=== Exact non-asymptotic bound (Theorem 3.3, first inequality) ===")
    print(f"m={m}, theta_Q={theta_Q}, D_M(Q,P*)={tq['D_M_inf']:.6f}")
    print(f"Poisson solution omega = {st.poisson_solution(Q, st.f_P_vector(Q, P_star), pi_Q)}")
    print(f"E[omega(X_0)] (uniform init) = {dist_X0 @ st.poisson_solution(Q, st.f_P_vector(Q, P_star), pi_Q):.6f}")
    print(f"E[omega(X_tau)] ~ pi_Q.omega = 0 (PE normalization)")
    print()
    results = []
    for a in alphas:
        b = exact_bound_for_alpha(Q, P_star, pi_Q, a, dist_X0, dist_Xtau_approx)
        b["mean_tau"] = mean_taus[a]
        b["satisfies"] = bool(mean_taus[a] >= b["bound_clipped"])
        b["gap"] = float(mean_taus[a] - b["bound_clipped"])
        results.append(b)
        print(f"  alpha={a:<10g}  bound={b['bound_raw']:.4f}  E[tau]={mean_taus[a]:.2f}  "
              f"gap={b['gap']:.2f}  satisfies={b['satisfies']}")

    # Compare exact vs worst-case bound
    print("\n=== Comparison: exact vs worst-case (eq 3) bound ===")
    worst_correction = 2.0 * tq["C_Q"] / tq["pi_min"]
    exact_correction = abs(results[0]["correction_E_w0_minus_E_wtau"])
    print(f"  exact correction |E[w(X0)]-E[w(Xtau)]| = {exact_correction:.6f}")
    print(f"  worst-case 2*C_Q/pi* = {worst_correction:.4f}")
    print(f"  ratio worst/exact = {worst_correction / max(exact_correction, 1e-12):.1f}x")

    out = {"m": m, "theta_Q": theta_Q, "D_M_inf": tq["D_M_inf"],
           "C_Q": tq["C_Q"], "pi_min": tq["pi_min"],
           "exact_correction": exact_correction,
           "worst_case_correction": worst_correction,
           "results": results}
    (out_dir / "exact_bound_m5.json").write_text(json.dumps(out, indent=2))
    print(f"\nsaved -> {out_dir / 'exact_bound_m5.json'}")


if __name__ == "__main__":
    main()