repro-bundle/v1/src/official_reference.py

"""Verbatim copies of the authors' G3 notebook functions, used only as an
independent reference to cross-check src/sequential_test.py on identical inputs.

Source: official_code/G3_ParametricFamily.ipynb cells 4, 10, 12.
Pasted exactly as written (no modifications) so any numerical discrepancy is
attributable to my re-implementation, not to a transcription difference.
"""
import numpy as np
from scipy.linalg import eig
from scipy.optimize import minimize_scalar


# --- Official cell 4: build_P_theta -----------------------------------------
def build_P_theta(theta, P0, f):
    tilde = P0 * np.exp(theta * f[np.newaxis, :])
    eigvals, eigvecs = eig(tilde.T)
    idx = np.argmax(np.real(eigvals))
    rho = np.real(eigvals[idx])
    v = np.real(eigvecs[:, idx])
    v = np.abs(v)
    Ptheta = (tilde * v[np.newaxis, :]) / (rho * v[:, np.newaxis])
    Ptheta = np.maximum(Ptheta, 0)
    Ptheta = Ptheta / Ptheta.sum(axis=1, keepdims=True)
    return Ptheta, rho, v


# --- Official cell 10: stationary_dist --------------------------------------
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, *_ = np.linalg.lstsq(A, b, rcond=None)
    pi = np.maximum(pi, 0)
    return pi / pi.sum()


# --- Official cell 11: kl_row, f_P_from_QP, D_M -----------------------------
def kl_row(q, p, eps=1e-12):
    q = np.clip(q, eps, 1)
    p = np.clip(p, eps, 1)
    return np.sum(q * np.log(q / p))


def f_P_from_QP(Q, P):
    m = Q.shape[0]
    f = np.zeros(m)
    for i in range(m):
        f[i] = kl_row(Q[i], P[i])
    return f


def D_M_official(Q, P):
    pi = stationary_dist(Q)
    f = f_P_from_QP(Q, P)
    return np.dot(pi, f)


# --- Official cell 12: compute_information_rate -----------------------------
def compute_information_rate(Q_TPM, Theta_P, P0, f, build_P_theta):
    pi_Q = stationary_dist(Q_TPM)

    def objective(theta):
        P_theta, _, _ = build_P_theta(theta, P0, f)
        eps = 1e-12
        ratio = Q_TPM / np.maximum(P_theta, eps)
        with np.errstate(divide="ignore", invalid="ignore"):
            kl_matrix = Q_TPM * np.log(ratio)
        kl_matrix[~np.isfinite(kl_matrix)] = 0.0
        return np.sum(pi_Q[:, None] * kl_matrix)

    res = minimize_scalar(objective, bounds=Theta_P, method="bounded")
    return res.fun