repro-bundle/v1/tests/test_core.py

"""Unit tests for src/sequential_test.py.

Validates primitives against closed-form cases (i.i.d. chain, identical P=Q),
cross-checks the parametric-family information rate against the value printed
by the authors' own G3 notebook (I = 2.0251949227282857 for m=64, seed 123),
and runs Algorithm 1 end-to-end on a tiny instance.
"""
import json
import sys
import numpy as np
from pathlib import Path

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


def _near(a, b, tol=1e-9):
    return abs(a - b) <= tol * max(1.0, abs(b))


def test_stationary_dist():
    rng = np.random.default_rng(0)
    P = rng.random((6, 6))
    P = P / P.sum(axis=1, keepdims=True)
    pi = st.stationary_dist(P)
    assert _near(pi.sum(), 1.0)
    assert np.allclose(pi @ P, pi, atol=1e-10), pi @ P - pi
    assert np.all(pi > 0)
    print("  stationary_dist OK: pi =", np.round(pi, 4))


def test_pseudo_spectral_gap_iid():
    # i.i.d. chain: rows all equal to pi => gamma_ps == 1, C_P == 2.
    pi = np.array([0.1, 0.2, 0.3, 0.4])
    P = np.tile(pi, (4, 1))
    g = st.pseudo_spectral_gap(P)
    assert _near(g, 1.0, 1e-9), g
    C = st.C_constant(P)
    assert _near(C, 2.0, 1e-9), C
    print(f"  iid chain: gamma_ps={g:.6f}  C_P={C:.6f} (expected 1.0 / 2.0)")


def test_pseudo_spectral_gap_nontrivial():
    rng = np.random.default_rng(7)
    P = rng.random((5, 5))
    P = P / P.sum(axis=1, keepdims=True)
    g = st.pseudo_spectral_gap(P)
    assert 0.0 < g < 1.0, g
    C = st.C_constant(P)
    assert C > 2.0, (g, C)   # for gamma_ps < 1, C_P > 2 (mixing correction)
    print(f"  nontrivial chain: gamma_ps={g:.4f}  C_P={C:.4f} (> 2 as expected)")


def test_D_M_zero_for_identical():
    rng = np.random.default_rng(1)
    Q = rng.random((5, 5)); Q = Q / Q.sum(axis=1, keepdims=True)
    dm, f = st.D_M(Q, Q)
    assert _near(dm, 0.0, 1e-9), dm
    assert np.allclose(f, 0.0, atol=1e-9)
    print("  D_M(Q,Q)=0 OK")


def test_poisson_solution_iid_constant():
    # i.i.d. Q (rows = pi) with f = constant c => PE rhs = c - c = 0 => omega = 0.
    pi = np.array([0.2, 0.3, 0.5])
    Q = np.tile(pi, (3, 1))
    f = np.array([1.7, 1.7, 1.7])
    w = st.poisson_solution(Q, f)
    assert np.allclose(w, 0.0, atol=1e-9), w
    print("  Poisson solution for iid+constant f = 0 OK")


def test_build_P_theta_stochastic():
    rng = np.random.default_rng(123)
    m = 5
    P0 = rng.random((m, m)); P0 = P0 / P0.sum(axis=1, keepdims=True)
    f = np.array([1.0, 1.0, 0.0, -1.0, -1.0])
    for theta in [-0.6, -0.4, 0.4, 0.6, 0.8]:
        Pt, rho, v = st.build_P_theta(theta, P0, f)
        assert np.allclose(Pt.sum(axis=1), 1.0, atol=1e-9), Pt.sum(axis=1)
        assert np.all(Pt >= 0)
    # cross-check matrices match the authors' build_P_theta to machine precision
    import official_reference as off
    for theta in [-0.6, -0.4, 0.4, 0.6, 0.8]:
        Pt_mine, _, _ = st.build_P_theta(theta, P0, f)
        Pt_off, _, _ = off.build_P_theta(theta, P0, f)
        assert np.allclose(Pt_mine, Pt_off, atol=1e-10), (theta, np.abs(Pt_mine - Pt_off).max())
    print("  build_P_theta row-stochastic + matches official function OK")


def test_D_M_inf_matches_official_functions():
    """Cross-check my build_P_theta / D_M_inf against the authors' own G3
    functions (src/official_reference.py, pasted verbatim) on IDENTICAL inputs.

    The notebook's printed I=2.0251949227282857 used an unseeded legacy
    `np.random.randn(m)` for the feature f, so it is not bit-reproducible.
    Instead we verify functional agreement: my functions == official functions
    to machine precision on a shared (P0, f, Q).
    """
    import official_reference as off
    rng = np.random.default_rng(123)
    m = 64
    P0 = rng.random((m, m)); P0 = P0 / P0.sum(axis=1, keepdims=True)
    f = rng.standard_normal(m)
    # build_P_theta: compare matrices
    for theta in [-0.6, -0.4, 0.4, 0.6, 0.8]:
        P_mine, _, _ = st.build_P_theta(theta, P0, f)
        P_off, _, _ = off.build_P_theta(theta, P0, f)
        assert np.allclose(P_mine, P_off, atol=1e-10), (theta, np.abs(P_mine-P_off).max())
    Q, _, _ = st.build_P_theta(-0.6, P0, f)
    # D_M: compare per-P values
    for theta in [0.4, 0.6, 0.8]:
        Pth, _, _ = st.build_P_theta(theta, P0, f)
        d_mine, _ = st.D_M(Q, Pth)
        d_off = off.D_M_official(Q, Pth)
        assert _near(d_mine, d_off, 1e-9), (theta, d_mine, d_off)
    # D_M^inf: compare optimized value
    D_mine, _, _ = st.make_parametric_null((0.4, 0.8), P0, f)["D_M_inf"](Q)
    D_off = off.compute_information_rate(Q, (0.4, 0.8), P0, f, off.build_P_theta)
    rel = abs(D_mine - D_off) / abs(D_off)
    print(f"  D_M^inf  mine={D_mine:.10f}  official_fn={D_off:.10f}  rel_err={rel:.2e}")
    # eps (1e-15 vs 1e-12 in KL clipping) + optimizer xatol => agree to ~1e-5
    assert rel < 1e-3, rel
    # brute-force grid sanity (coarse) — independent of any optimizer
    grid = np.linspace(0.4, 0.8, 2001)
    pi_Q = st.stationary_dist(Q)
    vals = [float(pi_Q @ st.f_P_vector(Q, st.build_P_theta(th, P0, f)[0])) for th in grid]
    D_grid = min(vals)
    assert abs(D_grid - D_mine) < 5e-3, (D_grid, D_mine)
    print(f"  brute-force grid min = {D_grid:.6f}  (consistent with optimizer)")


def test_algorithm1_null_rarely_stops():
    """Under the null (data from P_theta with theta in Theta_P) the false-reject
    rate should be <= alpha (up to MC noise)."""
    rng = np.random.default_rng(202)
    m = 5
    P0 = rng.random((m, m)); P0 = P0 / P0.sum(axis=1, keepdims=True)
    f = np.array([1.0, 1.0, 0.0, -1.0, -1.0])
    theta_null = 0.6   # in Theta_P = (0.4, 0.8)
    P_null, _, _ = st.build_P_theta(theta_null, P0, f)
    alpha = 0.05
    n_trials = 40
    T_max = 2000
    stops = 0
    for _ in range(n_trials):
        gen = st.MarkovGenerator(P_null, rng=np.random.default_rng())
        test = st.SequentialMarkovChainTest(m, alpha, (0.4, 0.8), P0, f)
        tau, _, _, _ = st.run_trial(test, gen, T_max)
        if tau < T_max:
            stops += 1
    rate = stops / n_trials
    print(f"  null false-reject rate = {rate:.3f} (alpha={alpha}, {n_trials} trials) "
          f"-> {'OK' if rate <= 3*alpha else 'CHECK'}")
    # Type-I control with a 3-sigma style slack (sequential test, conservative boundary)
    assert rate <= 3 * alpha, rate


def test_algorithm1_stops_under_alternative():
    """Under the alternative (theta_Q = -0.6) the test stops well before T_max."""
    rng = np.random.default_rng(321)
    m = 5
    P0 = rng.random((m, m)); P0 = P0 / P0.sum(axis=1, keepdims=True)
    f = np.array([1.0, 1.0, 0.0, -1.0, -1.0])
    Q, _, _ = st.build_P_theta(-0.6, P0, f)
    alpha = 0.01
    T_max = 20000
    taus = []
    for _ in range(10):
        gen = st.MarkovGenerator(Q, rng=np.random.default_rng())
        test = st.SequentialMarkovChainTest(m, alpha, (0.4, 0.8), P0, f)
        tau, _, _, _ = st.run_trial(test, gen, T_max)
        taus.append(tau)
    print(f"  alt stopping times (m=5,alpha=0.01): mean={np.mean(taus):.1f} "
          f"min={min(taus)} max={max(taus)} (T_max={T_max})")
    assert np.mean(taus) < T_max, taus


if __name__ == "__main__":
    tests = [
        ("stationary_dist", test_stationary_dist),
        ("pseudo_spectral_gap_iid", test_pseudo_spectral_gap_iid),
        ("pseudo_spectral_gap_nontrivial", test_pseudo_spectral_gap_nontrivial),
        ("D_M_zero_for_identical", test_D_M_zero_for_identical),
        ("poisson_solution_iid_constant", test_poisson_solution_iid_constant),
        ("build_P_theta_stochastic", test_build_P_theta_stochastic),
        ("D_M_inf_matches_official_functions", test_D_M_inf_matches_official_functions),
        ("algorithm1_null_rarely_stops", test_algorithm1_null_rarely_stops),
        ("algorithm1_stops_under_alternative", test_algorithm1_stops_under_alternative),
    ]
    print(f"Running {len(tests)} unit tests\n")
    failed = 0
    for name, fn in tests:
        print(f"[{name}]")
        try:
            fn()
        except AssertionError as e:
            print(f"  FAIL: {e}")
            failed += 1
        except Exception as e:
            print(f"  ERROR: {type(e).__name__}: {e}")
            failed += 1
    print(f"\n{len(tests)-failed}/{len(tests)} passed")
    sys.exit(1 if failed else 0)