repro-bundle/v1/experiments/claim3_structure.py

"""Claim 3: the lower bound incorporates BOTH the stationary distribution AND the
transition structure of the unknown Markov chain.

The bound's denominator is  D_M^inf(Q, P) = inf_{P in P} sum_i pi_i KL(Q(i,.), P(i,.))
                              ^^^ stationary weights          ^^^ per-row transition KL

We verify the "both incorporated" claim four ways:

  (A) Decomposition: for the m=5 instance, tabulate pi_i and f(i)=KL(Q_i,P*_i)
      across states; both vary, and D_M = sum_i pi_i f(i) exactly.
  (B) Same-pi / different-rows: lazy chain Q_eps = (1-eps) I + eps (1 pi^T)
      keeps pi fixed while varying the transition rows; D_M(Q_eps, P) varies
      with eps (transitions matter, pi held constant).
  (C) Different-pi: across the parametric family (theta_Q in {-0.4..-0.8}),
      both pi_Q and the row-KL vector f change; D_M^inf changes; the empirical
      E[tau] tracks 1/D_M^inf.
  (D) Weight choice matters: paper's pi-weighted sum vs uniform weights vs
      sup_w (naive BPI ||f||_inf) give different rates; only the pi-weighted
      rate matches the measured E[tau]/log(1/alpha).
"""
import json
import sys
import time
from pathlib import Path

import numpy as np

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


def kl_rows(Q, P):
    return np.array([st.kl_row(Q[i], P[i]) for i in range(Q.shape[0])])


def part_A_decomposition(m, theta_Q, seed):
    """Show D_M(Q, P*) = sum_i pi_i f(i) with both pi_i and f(i) varying."""
    P0, f, Q = build_instance(m, theta_Q, (0.4, 0.8), seed)
    pi_Q = st.stationary_dist(Q)
    info = st.make_parametric_null((0.4, 0.8), P0, f)
    D_inf, theta_star, P_star = info["D_M_inf"](Q, pi_Q)
    f_vec = kl_rows(Q, P_star)
    contribs = pi_Q * f_vec
    rows = []
    for i in range(m):
        rows.append({
            "state": i,
            "pi_i": float(pi_Q[i]),
            "f_i_KL": float(f_vec[i]),
            "contribution_pi_i_x_f_i": float(contribs[i]),
        })
    return {
        "D_M_inf": float(D_inf),
        "sum_pi_i_f_i": float(contribs.sum()),
        "pi_min": float(pi_Q.min()),
        "pi_max": float(pi_Q.max()),
        "f_min": float(f_vec.min()),
        "f_max": float(f_vec.max()),
        "pi_varies": bool(pi_Q.max() - pi_Q.min() > 1e-6),
        "f_varies": bool(f_vec.max() - f_vec.min() > 1e-3),
        "per_state": rows,
    }


def part_B_same_pi_different_rows(pi, m, epsilons):
    """Lazy chain Q_eps = (1-eps) I + eps (1 pi^T) has stationary pi for all eps
    but different transition rows.  D_M(Q_eps, P) against P = (1 pi^T) (the
    i.i.d. chain with stationary pi) varies with eps => transitions matter.
    """
    pi = np.asarray(pi, float)
    P_iid = np.tile(pi, (m, 1))   # null: i.i.d. with stationary pi
    rows = []
    for eps in epsilons:
        Q = (1.0 - eps) * np.eye(m) + eps * np.tile(pi, (m, 1))
        # verify stationary
        pi_check = st.stationary_dist(Q)
        pi_err = float(np.max(np.abs(pi_check - pi)))
        f_vec = kl_rows(Q, P_iid)
        D = float(pi @ f_vec)
        rows.append({
            "eps": float(eps),
            "D_M_Q_eps_P": D,
            "1_over_D": 1.0 / D if D > 0 else float("inf"),
            "max_pi_deviation": pi_err,
            "f_min": float(f_vec.min()),
            "f_max": float(f_vec.max()),
            "f_mean": float(f_vec.mean()),
        })
    return {"pi": pi.tolist(), "P_null": "i.i.d. rows = pi",
            "rows": rows,
            "D_varies_with_eps": bool(max(r["D_M_Q_eps_P"] for r in rows)
                                     - min(r["D_M_Q_eps_P"] for r in rows) > 1e-6),
            "pi_held_constant": bool(max(r["max_pi_deviation"] for r in rows) < 1e-6)}


def part_C_different_pi(m, theta_Q_list, alpha, n_trials, seed, T_max=20000):
    """Across theta_Q in the alternative set, both pi_Q and the row-KLs change;
    D_M^inf changes; E[tau] tracks 1/D_M^inf."""
    P0, f, _ = build_instance(m, -0.6, (0.4, 0.8), seed)
    rows = []
    for theta_Q in theta_Q_list:
        Q, _, _ = st.build_P_theta(theta_Q, P0, f)
        pi_Q = st.stationary_dist(Q)
        info = st.make_parametric_null((0.4, 0.8), P0, f)
        D_inf, theta_star, P_star = info["D_M_inf"](Q, pi_Q)
        f_vec = kl_rows(Q, P_star)
        # run Algorithm 1
        taus = []
        shared = ft.FastSequentialTest(m, alpha, (0.4, 0.8), P0, f,
                                       n_grid=8192, check_interval=1)
        for tr in range(n_trials):
            rng = np.random.default_rng(int(seed * 100003 + tr * 7919 + theta_Q * 1e6))
            shared.reset()
            tau, _, _, _ = ft.run_trial_fast(shared, Q, T_max, rng)
            taus.append(int(tau))
        taus = np.array(taus)
        rows.append({
            "theta_Q": float(theta_Q),
            "pi_min": float(pi_Q.min()),
            "pi_max": float(pi_Q.max()),
            "f_min": float(f_vec.min()),
            "f_max": float(f_vec.max()),
            "pi_varies_across_states": float(pi_Q.max() - pi_Q.min()),
            "f_varies_across_states": float(f_vec.max() - f_vec.min()),
            "D_M_inf": float(D_inf),
            "1_over_D_M_inf": 1.0 / float(D_inf),
            "mean_tau": float(np.mean(taus)),
            "std_tau": float(np.std(taus, ddof=1)),
            "predicted_tau_asymptotic": float(np.log(1.0 / alpha) / D_inf),
            "ratio_tau_over_log": float(np.mean(taus) / np.log(1.0 / alpha)),
        })
    return {"alpha": alpha, "n_trials": n_trials, "rows": rows}


def part_D_weight_choice(m, theta_Q, alpha_list, n_trials, seed, T_max=200000):
    """Compare the asymptotic rate 1/D_M (pi-weighted) against:
       - uniform-weight rate  1 / mean_i f(i)
       - naive BPI rate       1 / ||f||_inf  (sup_w weighting)
    Only the pi-weighted rate matches the measured E[tau]/log(1/alpha) as
    alpha -> 0.  Sweeps alpha to show the measured ratio converges to the
    paper's pi-weighted rate, not to the naive BPI rate.
    """
    P0, f, Q = build_instance(m, theta_Q, (0.4, 0.8), seed)
    pi_Q = st.stationary_dist(Q)
    info = st.make_parametric_null((0.4, 0.8), P0, f)
    D_inf, theta_star, P_star = info["D_M_inf"](Q, pi_Q)
    f_vec = kl_rows(Q, P_star)
    rates = {
        "paper_pi_weighted_1_over_D_M": 1.0 / float(D_inf),
        "uniform_weight_1_over_mean_f": 1.0 / float(f_vec.mean()),
        "naive_BPI_1_over_max_f": 1.0 / float(f_vec.max()),
    }
    shared = ft.FastSequentialTest(m, 1.0, (0.4, 0.8), P0, f, n_grid=8192, check_interval=1)
    sweep = []
    for alpha in alpha_list:
        taus = []
        shared.alpha = float(alpha)
        shared.reset()
        for tr in range(n_trials):
            rng = np.random.default_rng(seed * 100003 + tr * 7919 + int(-np.log10(alpha) * 13))
            shared.reset()
            tau, _, _, _ = ft.run_trial_fast(shared, Q, T_max, rng)
            taus.append(int(tau))
        measured = float(np.mean(taus) / np.log(1.0 / alpha))
        sweep.append({
            "alpha": float(alpha),
            "log_inv_alpha": float(np.log(1.0 / alpha)),
            "mean_tau": float(np.mean(taus)),
            "measured_ratio": measured,
            "dist_to_paper_rate": abs(measured - rates["paper_pi_weighted_1_over_D_M"]),
            "dist_to_naive_BPI_rate": abs(measured - rates["naive_BPI_1_over_max_f"]),
            "closer_to_paper": bool(
                abs(measured - rates["paper_pi_weighted_1_over_D_M"])
                < abs(measured - rates["naive_BPI_1_over_max_f"])),
        })
    return {
        "alpha_list": list(alpha_list), "n_trials": n_trials,
        "D_M_inf": float(D_inf), "f_vector": f_vec.tolist(),
        "pi_Q": pi_Q.tolist(),
        "rates_predicted": rates,
        "sweep": sweep,
        "all_closer_to_paper": bool(all(s["closer_to_paper"] for s in sweep)),
    }


def main():
    out_dir = Path(__file__).resolve().parent.parent / "outputs"
    out_dir.mkdir(parents=True, exist_ok=True)
    seed = 123
    print("=== Claim 3 ablation ===", flush=True)

    print("\n[A] Decomposition D_M = sum_i pi_i f(i), m=5...", flush=True)
    A = part_A_decomposition(5, -0.6, seed)
    print(f"  D_M^inf = {A['D_M_inf']:.6f} = sum pi_i f_i = {A['sum_pi_i_f_i']:.6f}", flush=True)
    print(f"  pi in [{A['pi_min']:.4f}, {A['pi_max']:.4f}]  f in [{A['f_min']:.4f}, {A['f_max']:.4f}]",
          flush=True)
    print(f"  pi varies: {A['pi_varies']}   f varies: {A['f_varies']}", flush=True)

    print("\n[B] Same pi, different rows (lazy chain Q_eps)...", flush=True)
    B = part_B_same_pi_different_rows([0.3, 0.2, 0.2, 0.15, 0.15], 5,
                                      [0.05, 0.1, 0.2, 0.4, 0.6, 0.8, 0.95, 0.99])
    print(f"  pi held constant (max dev {max(r['max_pi_deviation'] for r in B['rows']):.2e})",
          flush=True)
    for r in B["rows"]:
        print(f"  eps={r['eps']:<5}  D_M={r['D_M_Q_eps_P']:.4f}  "
              f"1/D={r['1_over_D']:.4f}  f in [{r['f_min']:.3f},{r['f_max']:.3f}]",
              flush=True)
    print(f"  D varies with eps (pi fixed): {B['D_varies_with_eps']}", flush=True)

    print("\n[C] Different theta_Q -> different pi AND rows; E[tau] tracks 1/D...",
          flush=True)
    C = part_C_different_pi(5, [-0.4, -0.5, -0.6, -0.7, -0.8], 1e-6, 30, seed)
    print(f"  alpha=1e-6, 30 trials each:", flush=True)
    for r in C["rows"]:
        print(f"  theta_Q={r['theta_Q']:<5}  D_M^inf={r['D_M_inf']:.4f}  "
              f"1/D={r['1_over_D_M_inf']:.4f}  E[tau]={r['mean_tau']:.1f}+-{r['std_tau']:.1f}  "
              f"ratio={r['ratio_tau_over_log']:.4f}  (pi range "
              f"{r['pi_min']:.3f}-{r['pi_max']:.3f}, f range "
              f"{r['f_min']:.3f}-{r['f_max']:.3f})", flush=True)

    print("\n[D] Weight choice: pi-weighted vs uniform vs naive BPI...", flush=True)
    D = part_D_weight_choice(5, -0.6, [1e-6, 1e-12, 1e-40, 1e-100, 1e-250], 30, seed)
    print(f"  D_M^inf = {D['D_M_inf']:.4f}; predicted rates:", flush=True)
    for k, v in D["rates_predicted"].items():
        print(f"    {k:<40} = {v:.4f}", flush=True)
    print(f"  measured E[tau]/log(1/a) across alpha:", flush=True)
    for s in D["sweep"]:
        tag = "paper" if s["closer_to_paper"] else "naive"
        print(f"    alpha={s['alpha']:<9g}  measured={s['measured_ratio']:.4f}  "
              f"(d_paper={s['dist_to_paper_rate']:.3f}, d_naive={s['dist_to_naive_BPI_rate']:.3f}) "
              f"-> closer to {tag}", flush=True)
    print(f"  all points closer to paper's pi-weighted rate: {D['all_closer_to_paper']}",
          flush=True)

    out = out_dir / "claim3_ablation.json"
    out.write_text(json.dumps({"A": A, "B": B, "C": C, "D": D}, indent=2))
    print(f"\nsaved -> {out}", flush=True)


if __name__ == "__main__":
    main()