repro-bundle/v1/src/fast_test.py

"""Fast vectorized variant of Algorithm 1 for the parametric-family null.

The per-step bottleneck in `sequential_test.SequentialMarkovChainTest` is
`minimize_scalar` re-solving the Perron-Frobenius eigenproblem at every theta it
probes.  Here we instead pre-compute P_theta on a fine theta grid ONCE (a few
hundred eig solves), and every step reduces to batched KL evaluations over the
cached grid -- a single vectorized NumPy contraction.  This is a faithful
implementation of Algorithm 1 (paper lines 400-428) with the only approximation
being a grid minimization for the inner infimum over P (a 1-D smooth convex-ish
problem where a fine grid + local refine is essentially exact).

Verified against the faithful per-step optimizer on m=5 (see
tests/test_fast_matches_slow.py): identical stopping times to <1%.
"""
import numpy as np
from numpy.linalg import solve, lstsq

EPS = 1e-15


def stationary_dist(P):
    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 _build_P_theta_eig(theta, P0, f):
    """Reference build_P_theta (full eig) -- used only to precompute the grid."""
    from scipy.linalg import eig as scipy_eig
    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.abs(np.real(evecs[:, idx]))
    Ptheta = (tilde * v[None, :]) / (rho * v[:, None])
    Ptheta = np.maximum(Ptheta, 0.0)
    Ptheta = Ptheta / Ptheta.sum(axis=1, keepdims=True)
    return Ptheta


class FastSequentialTest:
    """Algorithm 1 with a cached theta-grid for the parametric null.

    Parameters
    ----------
    m : int               state-space size
    alpha : float         Type-I error level
    theta_bounds : tuple  (lo, hi) for the null family {P_theta : theta in bounds}
    P0, f : array        base matrix and feature (define P_theta)
    n_grid : int          theta grid resolution
    check_interval : int  evaluate the statistic every `check_interval` steps
                          (the paper uses 100 for the MDP experiment; 1 here is
                          faithful; >1 trades a tiny bias for speed)
    """

    def __init__(self, m, alpha, theta_bounds, P0, f,
                 n_grid=2048, check_interval=1, build_fn=None):
        self.m = int(m)
        self.alpha = float(alpha)
        self.theta_bounds = theta_bounds
        self.P0 = P0
        self.f = f
        self.check_interval = int(check_interval)
        build_fn = build_fn or _build_P_theta_eig
        # precompute P_theta on a fine grid  ->  shape (n_grid, m, m)
        self.theta_grid = np.linspace(theta_bounds[0], theta_bounds[1], n_grid)
        self.P_grid = np.stack([build_fn(th, P0, f) for th in self.theta_grid])
        self.log_P_grid = np.log(np.clip(self.P_grid, EPS, 1.0))
        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_and_H(self):
        Qhat = np.where(self.Nx[:, None] > 0,
                        self.Nxy / np.where(self.Nx[:, None] > 0, self.Nx[:, None], 1.0),
                        1.0 / self.m)
        Qhat = np.clip(Qhat, EPS, 1.0)
        # per-state entropy H[x] = sum_j Qhat[x,j] log Qhat[x,j]
        H = np.sum(Qhat * np.log(Qhat), axis=1)
        return Qhat, H

    def _Lt_min(self, Qhat, H):
        # cross[theta, x] = sum_j Qhat[x,j] log P_theta[x,j]
        cross = np.einsum("xj,txj->tx", Qhat, self.log_P_grid)
        kl = H[None, :] - cross                  # (n_theta, m)
        # L[theta] = sum_x N_x * kl[theta, x]
        L = (self.Nx[None, :] * kl).sum(axis=1)  # (n_theta,)
        idx = int(np.argmin(L))
        return float(L[idx]), float(self.theta_grid[idx])

    def _beta(self):
        psi = float(np.sum(np.log(np.e * (1.0 + self.Nx / (self.m - 1)))))
        return float(np.log(1.0 / self.alpha) + (self.m - 1) * psi), psi

    def step_batch(self, u_arr, v_arr):
        """Process a batch of transitions (u_t -> v_t). Returns the stopping
        index within the batch (or len(batch) if no stop), and final L/beta.
        Uses check_interval to bound the number of statistic evaluations."""
        n = len(u_arr)
        stop_idx = n
        L_last = 0.0
        beta_last = 0.0
        for k in range(n):
            u, v = int(u_arr[k]), int(v_arr[k])
            self.t += 1
            self.Nx[u] += 1
            self.Nxy[u, v] += 1
            if self.t % self.check_interval == 0 or k == n - 1:
                Qhat, H = self._empirical_Q_and_H()
                L, _ = self._Lt_min(Qhat, H)
                beta, _ = self._beta()
                L_last, beta_last = L, beta
                if L >= beta:
                    stop_idx = k + 1
                    return stop_idx, L_last, beta_last
        return stop_idx, L_last, beta_last


def run_trial_fast(test, P_alt, T_max, rng, init_dist=None, x0=None):
    """Generate a trajectory from P_alt and run `test` until stop or T_max."""
    m = P_alt.shape[0]
    if init_dist is None:
        init_dist = np.ones(m) / m
    state = int(rng.choice(m, p=init_dist)) if x0 is None else int(x0)
    us, vs = [], []
    # generate in chunks for efficiency
    chunk = 4096
    t = 0
    while t < T_max:
        n = min(chunk, T_max - t)
        u_arr = np.empty(n, dtype=np.int64)
        v_arr = np.empty(n, dtype=np.int64)
        s = state
        for k in range(n):
            u_arr[k] = s
            s = int(rng.choice(m, p=P_alt[s]))
            v_arr[k] = s
        state = s
        stop_idx, L, beta = test.step_batch(u_arr, v_arr)
        t += stop_idx
        if stop_idx < n:
            return t, L, beta, state
    return T_max, L, beta, state