repro-bundle/v1/src/sequential_test.py

"""Independent re-implementation of the sequential Markov-chain test of
Sethi et al. (2026), "Asymptotically Optimal Sequential Testing with Markovian
Data" (ICML 2026, arXiv:2602.17587, OpenReview YEckWPoS09).

Implemented from:
  * Algorithm 1 (Section 4, paper text lines 400-428)
  * Definition 3.2 (stationary-weighted KL divergence)
  * Theorem 3.3 (non-asymptotic instance-dependent lower bound, eq (3))
  * Proposition 3.1 (Poisson-solution bound, constant C_P)
  * Definition 2.2 (pseudo-spectral gap, Paulin 2015)
  * Appendix G.3 (parametric family construction)

Pure NumPy / SciPy. No dependence on the authors' notebooks.
"""
import numpy as np
from numpy.linalg import solve, lstsq
from scipy.linalg import eig as scipy_eig
from scipy.optimize import minimize_scalar

EPS = 1e-15


# ---------------------------------------------------------------------------
# Markov-chain primitives
# ---------------------------------------------------------------------------
def stationary_dist(P):
    """Stationary distribution pi of row-stochastic P, solving pi^T P = pi^T."""
    m = P.shape[0]
    A = np.vstack([P.T - np.eye(m), np.ones(m)])
    b = np.zeros(m + 1)
    b[-1] = 1.0
    pi, *_ = lstsq(A, b, rcond=None)
    pi = np.maximum(pi, 0.0)
    return pi / pi.sum()


def time_reversal(P, pi):
    """P* (time reversal): P*(i,j) = P(j,i) * pi(j) / pi(i)."""
    return (P.T * pi[None, :]) / np.maximum(pi[:, None], EPS)


def pseudo_spectral_gap(P, pi=None, K=40):
    """gamma_ps(P) = max_{k>=1} (1/k) * (1 - lambda_2((P*)^k P^k))
    where lambda_2 is the second-largest eigenvalue (the largest equals 1).

    Matches Paulin (2015) Definition 3.4 as quoted in the paper's Definition 2.2.
    For the i.i.d. case (rows all equal to pi) this returns 1.
    """
    if pi is None:
        pi = stationary_dist(P)
    Pstar = time_reversal(P, pi)
    m = P.shape[0]
    best = 0.0
    Pk = np.eye(m)
    Pstark = np.eye(m)
    for k in range(1, K + 1):
        Pk = Pk @ P
        Pstark = Pstark @ Pstar
        M = Pstark @ Pk
        ev = np.sort(np.real(scipy_eig(M, right=False)))
        lam2 = ev[-2]  # second-largest (largest is ~1)
        gap_k = (1.0 / k) * max(0.0, 1.0 - lam2)
        if gap_k > best:
            best = gap_k
    return best


def C_constant(P, pi=None):
    """C_P from Proposition 3.1 (paper eq (9)).

    || omega_{P,f} ||_inf <= C_P ||f||_inf, with
        C_P = 1 / [ (1-gamma_ps)^{1/(2 gamma_ps)} * sqrt(pi*) * (1 - sqrt(1-gamma_ps)) ]
    for gamma_ps in (0,1), and C_P = 2 for gamma_ps = 1.
    """
    if pi is None:
        pi = stationary_dist(P)
    pi_star = pi.min()
    gamma = pseudo_spectral_gap(P, pi)
    if gamma >= 1.0 - 1e-9:
        return 2.0
    g = min(gamma, 1.0 - 1e-12)
    base = (1.0 - g) ** (1.0 / (2.0 * g))
    return 1.0 / (base * np.sqrt(pi_star) * (1.0 - np.sqrt(1.0 - g)))


# ---------------------------------------------------------------------------
# Information quantities
# ---------------------------------------------------------------------------
def kl_row(q, p, eps=EPS):
    """D_KL(q || p) for discrete distributions given as rows."""
    q = np.clip(np.asarray(q, float), eps, 1.0)
    p = np.clip(np.asarray(p, float), eps, 1.0)
    q = q / q.sum()
    p = p / p.sum()
    return float(np.sum(q * np.log(q / p)))


def f_P_vector(Q, P):
    """f_P(i) = D_KL(Q(i, .) || P(i, .)) for each state i (Theorem 3.3)."""
    return np.array([kl_row(Q[i], P[i]) for i in range(Q.shape[0])])


def D_M(Q, P, pi_Q=None):
    """D_M(Q,P) = sum_i pi_i D_KL(Q(i,.) || P(i,.))  (Definition 3.2)."""
    if pi_Q is None:
        pi_Q = stationary_dist(Q)
    f = f_P_vector(Q, P)
    return float(pi_Q @ f), f


def poisson_solution(Q, f, pi_Q=None):
    """Solution omega to (I - Q) omega = f - (pi.f) 1, normalized by pi.omega = 0
    (paper eq (2))."""
    if pi_Q is None:
        pi_Q = stationary_dist(Q)
    m = Q.shape[0]
    rhs = f - (pi_Q @ f) * np.ones(m)
    A = np.eye(m) - Q
    A[-1, :] = pi_Q
    rhs[-1] = 0.0
    w = solve(A, rhs)
    return w


def nonasymptotic_lower_bound(Q, P_set_info, alpha, pi_Q=None):
    """Theorem 3.3, eq (3):  E[tau] >= log(1/alpha)/D_M^inf - 2 C_Q / pi*.

    Returns dict with the bound and all components.  `P_set_info` must provide
        D_M_inf(Q, P)  : float           (inf over the null set)
        P_star(Q)      : ndarray (m,m)   (the P attaining the inf)
    """
    if pi_Q is None:
        pi_Q = stationary_dist(Q)
    pi_star = pi_Q.min()
    C_Q = C_constant(Q, pi_Q)
    D_inf, P_star = P_set_info["D_M_inf"](Q), P_set_info["P_star"](Q)
    lb = np.log(1.0 / alpha) / D_inf - 2.0 * C_Q / pi_star
    return {
        "alpha": alpha,
        "D_M_inf": D_inf,
        "C_Q": C_Q,
        "pi_star": pi_star,
        "gamma_ps": pseudo_spectral_gap(Q, pi_Q),
        "asymptotic_part": np.log(1.0 / alpha) / D_inf,
        "correction": 2.0 * C_Q / pi_star,
        "bound": max(0.0, lb),   # paper takes max with 0 (line 1467-1473)
        "bound_raw": lb,
        "P_star": P_star,
    }


# ---------------------------------------------------------------------------
# Parametric family P_theta (Appendix G.3)
# ---------------------------------------------------------------------------
def build_P_theta(theta, P0, f):
    """P_theta from P0 and feature f via Perron-Frobenius normalization
    (paper Appendix G.3).

    tilde_theta(i,j) = P0(i,j) exp(theta * f_j);  P_theta = tilde * v / (rho * v)
    where (rho, v) is the Perron-Frobenius eigenpair of tilde.
    """
    tilde = P0 * np.exp(theta * f[None, :])
    evals, evecs = scipy_eig(tilde.T)
    idx = int(np.argmax(np.real(evals)))
    rho = float(np.real(evals[idx]))
    v = np.real(evecs[:, idx])
    v = np.abs(v)
    Ptheta = (tilde * v[None, :]) / (rho * v[:, None])
    Ptheta = np.maximum(Ptheta, 0.0)
    Ptheta = Ptheta / Ptheta.sum(axis=1, keepdims=True)
    return Ptheta, rho, v


def make_parametric_null(theta_bounds, P0, f):
    """Returns a dict with the operations the test / lower bound need for the
    parametric null P = {P_theta : theta in theta_bounds}."""

    def D_M_inf(Q, pi_Q=None):
        if pi_Q is None:
            pi_Q = stationary_dist(Q)

        def obj(theta):
            Ptheta, _, _ = build_P_theta(theta, P0, f)
            return float(pi_Q @ f_P_vector(Q, Ptheta))

        res = minimize_scalar(obj, bounds=theta_bounds, method="bounded",
                              options={"xatol": 1e-10})
        Pstar, _, _ = build_P_theta(res.x, P0, f)
        return float(res.fun), res.x, Pstar

    return {"D_M_inf": D_M_inf, "theta_bounds": theta_bounds}


# ---------------------------------------------------------------------------
# Algorithm 1: Sequential Markov Chain Test
# ---------------------------------------------------------------------------
class SequentialMarkovChainTest:
    """Faithful re-implementation of Algorithm 1 (paper lines 400-428).

    The null is a parametric family {P_theta : theta in theta_bounds} and the
    statistic L_t is the infimum over the null of the visitation-weighted KL.
    """

    def __init__(self, m, alpha, theta_bounds, P0, f, build_P_theta_fn=None):
        self.m = int(m)
        self.alpha = float(alpha)
        self.theta_bounds = theta_bounds
        self.P0 = P0
        self.f = f
        self.build_P_theta = build_P_theta_fn or build_P_theta
        self.reset()

    def reset(self):
        self.t = 0
        self.Nx = np.zeros(self.m)
        self.Nxy = np.zeros((self.m, self.m))

    def empirical_Q(self):
        Qhat = np.zeros((self.m, self.m))
        for x in range(self.m):
            if self.Nx[x] > 0:
                Qhat[x] = self.Nxy[x] / self.Nx[x]
            else:
                Qhat[x] = np.ones(self.m) / self.m
        return Qhat

    def compute_psi(self):
        # psi_t = sum_x log( e * (1 + N_x / (m-1)) )   (paper line 415)
        return float(np.sum(np.log(np.e * (1.0 + self.Nx / (self.m - 1)))))

    def compute_beta(self, psi):
        return float(np.log(1.0 / self.alpha) + (self.m - 1) * psi)

    def _Lt_objective(self, theta, Qhat):
        Ptheta, _, _ = self.build_P_theta(theta, self.P0, self.f)
        kl = np.sum(Qhat * np.log((Qhat + EPS) / (Ptheta + EPS)), axis=1)
        mask = self.Nx > 0
        return float(np.sum(self.Nx[mask] * kl[mask]))

    def compute_Lt(self, Qhat):
        res = minimize_scalar(self._Lt_objective, bounds=self.theta_bounds,
                              args=(Qhat,), method="bounded",
                              options={"xatol": 1e-10})
        return float(res.fun), float(res.x)

    def step(self, u, v):
        """Process transition (u -> v). Returns (stop, L_t, beta_t, theta_hat)."""
        self.t += 1
        self.Nx[u] += 1
        self.Nxy[u, v] += 1
        Qhat = self.empirical_Q()
        psi = self.compute_psi()
        beta = self.compute_beta(psi)
        L, theta_hat = self.compute_Lt(Qhat)
        return (L >= beta), L, beta, theta_hat


# ---------------------------------------------------------------------------
# Markov chain simulator
# ---------------------------------------------------------------------------
class MarkovGenerator:
    def __init__(self, P, init_dist=None, rng=None):
        self.P = P
        self.m = P.shape[0]
        self.rng = rng if rng is not None else np.random.default_rng()
        if init_dist is None:
            init_dist = np.ones(self.m) / self.m
        self.state = self.rng.choice(self.m, p=init_dist)

    def current_state(self):
        return self.state

    def step(self):
        self.state = self.rng.choice(self.m, p=self.P[self.state])
        return self.state


def run_trial(test, gen, T_max, rng=None):
    """Run one trial of Algorithm 1 against `gen`. Returns stopping time tau and
    the final statistic/threshold at stop."""
    if rng is not None:
        gen.rng = rng
    t = 0
    x_prev = gen.current_state()
    L_last = beta_last = 0.0
    while t < T_max:
        x = gen.step()
        stop, L, beta, theta = test.step(x_prev, x)
        L_last, beta_last = L, beta
        if stop:
            return t + 1, L, beta, x
        x_prev = x
        t += 1
    return T_max, L_last, beta_last, x_prev