scripts/common.py

"""
Independent re-implementation of the SG(L)D discrete-time proxy theory of
"Accurate Large-sample Uncertainty Quantification using Stochastic Gradient MCMC"
(arXiv 2606.00293, OpenReview Zkj9ctQdMM).

Everything here is written from the paper's equations only:
  Eq. (3)   CT tuning              Lambda = (Sigma H + H Sigma) Chat^{-1}
  Eq. (5)   discrete-quadratic     L H S + S H L = L (C + H S H) L
  Eq. (6)   Ziyin large-sample/well-specified noise
  Eq. (11)  Prop 4.2 stationary covariance (with temperature)
  Eq. (12)  Thm 4.3 exact minibatch noise covariance
  Eq. (B.3) Prop B.1 momentum stationary covariance

No code from the authors' repository is used (their repo was downloaded only to
read the experimental protocol for Table 3: which columns of Boston are used,
standardisation, epochs, batch sizes).
"""

import numpy as np

SQRT = np.sqrt


# --------------------------------------------------------------------------
# generic helpers
# --------------------------------------------------------------------------
def sym(M):
    return 0.5 * (M + M.T)


def relf(A, B):
    """relative Frobenius error ||A-B||_F / ||B||_F"""
    return float(np.linalg.norm(A - B, "fro") / np.linalg.norm(B, "fro"))


def kron_sum_JJ(J):
    """K = (1/N) sum_n (J_n kron J_n) so that K @ vec(S) = vec((1/N) sum_n J_n S J_n).

    Uses vec(A S B) = (B^T kron A) vec(S) with column-major (Fortran) vec.
    We use row-major vec (numpy .ravel()), for which vec(A S B) = (A kron B^T) vec(S).
    J_n is symmetric so B^T = J_n.
    """
    N, D, _ = J.shape
    K = np.zeros((D * D, D * D))
    for n in range(N):
        K += np.kron(J[n], J[n])
    return K / N


# --------------------------------------------------------------------------
# Theorem 4.3 -- exact minibatch noise covariance, Eq. (12)
# --------------------------------------------------------------------------
class NoiseModel:
    """Exact expected minibatch-gradient noise covariance C_psi of Eq. (12).

        C_psi = (1/B) [ Ihat - Gamma th th^T Gamma^T / N^2
                        + (1/N) sum_n J_n Sigma J_n - Jbar Sigma Jbar ]

    (times (N-B)/(N-1) for sampling without replacement).
    """

    def __init__(self, grads, J, theta_hat, Gamma=None, N=None):
        self.g = np.asarray(grads, float)  # (N,D) per-sample grads at theta_hat
        self.J = np.asarray(J, float)  # (N,D,D) per-sample Hessians at theta_hat
        self.N = int(self.g.shape[0]) if N is None else N
        self.D = self.g.shape[1]
        self.theta_hat = np.asarray(theta_hat, float)
        self.Gamma = (
            np.zeros((self.D, self.D)) if Gamma is None else np.asarray(Gamma, float)
        )
        self.Ihat = self.g.T @ self.g / self.N
        self.Jbar = self.J.mean(0)
        self.H = self.Jbar + self.Gamma / self.N  # Hessian of L = (1/N) sum l_n + R/N
        gt = self.Gamma @ self.theta_hat
        self.const = self.Ihat - np.outer(gt, gt) / self.N**2
        self.K = kron_sum_JJ(self.J)  # vec-operator for (1/N) sum J_n S J_n

    def C(self, Sigma, B, replace=True):
        quad = (self.K @ Sigma.ravel()).reshape(
            self.D, self.D
        ) - self.Jbar @ Sigma @ self.Jbar
        out = (self.const + quad) / B
        if not replace:
            out = out * (self.N - B) / (self.N - 1.0)
        return sym(out)

    def C_bruteforce(self, Sigma, B, replace=True, n_mc=200000, rng=None):
        """Monte-Carlo of E_{psi~pi}[ Cov_batch( G(psi) ) ] with psi ~ N(theta_hat, Sigma).

        Completely independent of Eq. (12): the per-sample proxy gradients are
        formed explicitly and their (exact, given psi) batch covariance averaged.
        """
        rng = np.random.default_rng(0) if rng is None else rng
        Lc = np.linalg.cholesky(Sigma + 1e-14 * np.eye(self.D))
        acc = np.zeros((self.D, self.D))
        chunk = max(1, min(n_mc, 2000))
        done = 0
        while done < n_mc:
            k = min(chunk, n_mc - done)
            U = rng.standard_normal((k, self.D)) @ Lc.T  # psi - theta_hat
            # v_n(u) = grad_n + J_n u   for every n  -> (k,N,D)
            V = self.g[None, :, :] + np.einsum("nde,ke->knd", self.J, U)
            vbar = V.mean(1)  # (k,D)
            S2 = np.einsum("knd,kne->kde", V, V) / self.N
            acc += (S2 - np.einsum("kd,ke->kde", vbar, vbar)).sum(0)
            done += k
        out = acc / n_mc / B
        if not replace:
            out = out * (self.N - B) / (self.N - 1.0)
        return sym(out)


# --------------------------------------------------------------------------
# Proposition 4.2 / Eq. (11) -- exact stationary covariance of the proxy
# --------------------------------------------------------------------------
def stationary_cov_proxy(nm, Lam, B, beta=np.inf, replace=True):
    """Solve  S = (I-Lam H) S (I-Lam H)^T + Lam C(S) Lam + 2 Lam / beta  exactly.

    This is the one-step second-moment identity of the proxy recursion; it is
    algebraically equivalent to Eq. (11).  Affine in S -> solved as a D^2 linear
    system (no fixed-point iteration, no small-lambda approximation).
    """
    D = nm.D
    M = np.eye(D) - Lam @ nm.H
    fac = 1.0 if replace else (nm.N - B) / (nm.N - 1.0)
    # linear part
    A = np.zeros((D * D, D * D))
    E = np.zeros((D, D))
    for i in range(D):
        for j in range(D):
            E[:] = 0.0
            E[i, j] = 1.0
            quad = (nm.K @ E.ravel()).reshape(D, D) - nm.Jbar @ E @ nm.Jbar
            out = M @ E @ M.T + Lam @ (fac * quad / B) @ Lam
            A[:, i * D + j] = out.ravel()
    b = (Lam @ (fac * nm.const / B) @ Lam).ravel()
    if np.isfinite(beta):
        b = b + (2.0 / beta * Lam).ravel()
    S = np.linalg.solve(np.eye(D * D) - A, b).reshape(D, D)
    return sym(S)


def eq11_residual(nm, Lam, S, B, beta=np.inf, replace=True):
    """||LHS-RHS||_F / ||LHS||_F for Eq. (11) with C_psi from Eq. (12)."""
    C = nm.C(S, B, replace)
    lhs = Lam @ nm.H @ S + S @ nm.H @ Lam
    rhs = Lam @ (C + nm.H @ S @ nm.H) @ Lam
    if np.isfinite(beta):
        rhs = rhs + 2.0 / beta * Lam
    return float(np.linalg.norm(lhs - rhs, "fro") / np.linalg.norm(lhs, "fro"))


# --------------------------------------------------------------------------
# Proposition B.1 -- momentum
# --------------------------------------------------------------------------
def stationary_cov_momentum(nm, Lam, B, kappa, beta=np.inf, replace=True):
    """Exact stationary covariance of the momentum proxy Eq. (B.2).

    Joint state z=(u,m), u=psi-theta_hat:
        m_t = kappa m_{t-1} + H u_{t-1} + eta_t
        u_t = (I-Lam H) u_{t-1} - kappa Lam m_{t-1} - Lam eta_t + zeta_t
    E[z z^T] satisfies a linear fixed point because Cov(eta|u) is linear in u u^T.
    Returns (Sigma_psi, P_full, cross) where cross = E[(u_t)(u_{t-1})^T].
    """
    D = nm.D
    fac = 1.0 if replace else (nm.N - B) / (nm.N - 1.0)
    A = np.zeros((2 * D, 2 * D))
    A[:D, :D] = np.eye(D) - Lam @ nm.H
    A[:D, D:] = -kappa * Lam
    A[D:, :D] = nm.H
    A[D:, D:] = kappa * np.eye(D)

    def noise_cov(Cu, temp=False):
        Q = np.zeros((2 * D, 2 * D))
        Q[:D, :D] = Lam @ Cu @ Lam
        Q[:D, D:] = -Lam @ Cu
        Q[D:, :D] = -Cu @ Lam
        Q[D:, D:] = Cu
        if temp and np.isfinite(beta):
            Q[:D, :D] = Q[:D, :D] + 2.0 / beta * Lam
        return Q

    n = 2 * D
    Amat = np.zeros((n * n, n * n))
    E = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            E[:] = 0.0
            E[i, j] = 1.0
            Su = E[:D, :D].copy()
            quad = (nm.K @ Su.ravel()).reshape(D, D) - nm.Jbar @ Su @ nm.Jbar
            Cu = fac * quad / B
            Amat[:, i * n + j] = (A @ E @ A.T + noise_cov(Cu)).ravel()
    Q0 = noise_cov(fac * nm.const / B, temp=True)
    P = np.linalg.solve(np.eye(n * n) - Amat, Q0.ravel()).reshape(n, n)
    P = sym(P)
    Sig = P[:D, :D]
    # E[u_t u_{t-1}^T] = A_uu Sigma + A_um E[m_{t-1} u_{t-1}^T]
    cross = A[:D, :D] @ Sig + A[:D, D:] @ P[D:, :D]
    return Sig, P, cross


def eqB3_residual(nm, Lam, S, B, kappa, beta=np.inf, replace=True, temp="paper"):
    """Relative residual of Eq. (B.3) of Proposition B.1, exactly as printed:

      (1-k)(L H S + S H L) + k/(1-k^2)(L H L H S + S H L H L)
          = L C L + (1+k^2)/(1-k^2) L H S H L + (1+k^2) 2L/beta

    temp='corrected' substitutes the temperature term derived in this
    reproduction: (1-k)^2 (2L/beta) + k/(1+k) [ L H (2L/beta) + (2L/beta) H L ].
    """
    H, C = nm.H, nm.C(S, B, replace)
    k = kappa
    lhs = (1 - k) * (Lam @ H @ S + S @ H @ Lam) + k / (1 - k**2) * (
        Lam @ H @ Lam @ H @ S + S @ H @ Lam @ H @ Lam
    )
    rhs = Lam @ C @ Lam + (1 + k**2) / (1 - k**2) * (Lam @ H @ S @ H @ Lam)
    if np.isfinite(beta):
        q = 2.0 * Lam / beta
        if temp == "paper":
            rhs = rhs + (1 + k**2) * q
        else:
            rhs = rhs + (1 - k) ** 2 * q + k / (1 + k) * (Lam @ H @ q + q @ H @ Lam)
    scale = max(np.linalg.norm(lhs, "fro"), np.linalg.norm(rhs, "fro"))
    return float(np.linalg.norm(lhs - rhs, "fro") / scale)


def solve_B3_for_Sigma(nm, Lam, B, kappa, beta=np.inf, replace=True):
    """Solve Eq. (B.3) for Sigma (it is linear in Sigma once C(Sigma) is substituted)."""
    D = nm.D
    H = nm.H
    fac = 1.0 if replace else (nm.N - B) / (nm.N - 1.0)
    A = np.zeros((D * D, D * D))
    E = np.zeros((D, D))
    for i in range(D):
        for j in range(D):
            E[:] = 0.0
            E[i, j] = 1.0
            quad = (nm.K @ E.ravel()).reshape(D, D) - nm.Jbar @ E @ nm.Jbar
            Cq = fac * quad / B
            lhs = (1 - kappa) * (Lam @ H @ E + E @ H @ Lam) + kappa / (
                1 - kappa**2
            ) * (Lam @ H @ Lam @ H @ E + E @ H @ Lam @ H @ Lam)
            rhs = Lam @ Cq @ Lam + (1 + kappa**2) / (1 - kappa**2) * (
                Lam @ H @ E @ H @ Lam
            )
            A[:, i * D + j] = (lhs - rhs).ravel()
    b = (Lam @ (fac * nm.const / B) @ Lam).ravel()
    if np.isfinite(beta):
        b = b + ((1 + kappa**2) * 2.0 * Lam / beta).ravel()
    return sym(np.linalg.solve(A, b).reshape(D, D))


# --------------------------------------------------------------------------
# tuning rules of Algorithm 1 (solve for Lambda given a target covariance)
# --------------------------------------------------------------------------
from scipy.optimize import root  # noqa: E402


def solve_Lambda(H, Sigma, Cnoise_fn, beta=np.inf, tol=1e-12):
    """Solve  Lam H S + S H Lam = Lam (C(S) + H S H) Lam + 2 Lam/beta  for Lam.

    Powell hybrid (scipy.optimize.root, method='hybr') on vec(Lam), warm-started
    from the closed form Lam0 = (S H + H S)(C + H S H)^{-1}, exactly as described
    in Section 6 of the paper.
    """
    D = H.shape[0]
    C = Cnoise_fn(Sigma)
    R = C + H @ Sigma @ H
    Lam0 = (Sigma @ H + H @ Sigma) @ np.linalg.inv(R)

    def F(v):
        L = v.reshape(D, D)
        r = L @ H @ Sigma + Sigma @ H @ L - L @ R @ L
        if np.isfinite(beta):
            r = r - 2.0 / beta * L
        return r.ravel()

    sol = root(F, Lam0.ravel(), method="hybr", tol=tol)
    L = sol.x.reshape(D, D)
    return L, float(np.linalg.norm(F(sol.x))), bool(sol.success), Lam0


# --------------------------------------------------------------------------
# models
# --------------------------------------------------------------------------
class Logistic:
    """Logistic regression with Gaussian prior R(theta)=0.5 theta^T Gamma theta.

    L(theta) = (1/N) sum_n l_n(theta) + R(theta)/N.
    All Assumption (A)-(C) constants are available in closed form, which lets us
    evaluate the explicit constant of Eq. (F.12).
    """

    def __init__(self, X, y, Gamma):
        self.X, self.y = np.asarray(X, float), np.asarray(y, float)
        self.N, self.D = self.X.shape
        self.Gamma = np.asarray(Gamma, float)

    @staticmethod
    def _s(z):
        return 1.0 / (1.0 + np.exp(-z))

    def grad_n(self, th, idx=None):
        X = self.X if idx is None else self.X[idx]
        y = self.y if idx is None else self.y[idx]
        return (self._s(X @ th) - y)[:, None] * X

    def grad_batch(self, TH, IDX):
        """TH (R,D), IDX (R,B) -> mean minibatch gradient of L, shape (R,D)."""
        Xb = self.X[IDX]  # (R,B,D)
        yb = self.y[IDX]  # (R,B)
        z = np.einsum("rbd,rd->rb", Xb, TH)
        w = self._s(z) - yb
        return (
            np.einsum("rb,rbd->rd", w, Xb) / IDX.shape[1] + TH @ self.Gamma.T / self.N
        )

    def hess_n(self, th):
        p = self._s(self.X @ th)
        w = p * (1 - p)
        return w[:, None, None] * np.einsum("nd,ne->nde", self.X, self.X)

    def gradL(self, th):
        return self.grad_n(th).mean(0) + self.Gamma @ th / self.N

    def hessL(self, th):
        return self.hess_n(th).mean(0) + self.Gamma / self.N

    def map_estimate(self, tol=1e-13, iters=100):
        th = np.zeros(self.D)
        for _ in range(iters):
            g, Hm = self.gradL(th), self.hessL(th)
            step = np.linalg.lstsq(Hm, g, rcond=None)[0]
            th = th - step
            if np.linalg.norm(step) < tol:
                break
        return th

    def constants(self, th):
        """Assumption (A)-(C) / Lemma F.2 constants."""
        nrm = np.linalg.norm(self.X, axis=1)
        Ln = nrm**2 / 4.0
        Mn = nrm**3 / (6.0 * np.sqrt(3.0))  # sup_z |sigma''(z)| = 1/(6 sqrt 3)
        Hh = self.hessL(th)
        g = self.grad_n(th)
        return dict(
            L=float(
                np.linalg.eigvalsh(
                    np.einsum("nd,ne->de", self.X, self.X) / (4 * self.N)
                    + self.Gamma / self.N
                )[-1]
            ),
            Ln_max=float(Ln.max()),
            Mbar=float(Mn.mean()),
            Mbar2=float((Mn**2).mean()),
            mu_global=float(np.linalg.eigvalsh(self.Gamma / self.N)[0]),
            mu_hat=float(np.linalg.eigvalsh(Hh)[0]),
            L_hat=float(np.linalg.eigvalsh(Hh)[-1]),
            tau4=float(((np.linalg.norm(g, axis=1) ** 4).mean()) ** 0.25),
        )


def sgd_paths(
    model,
    theta_hat,
    Lam,
    B,
    T,
    R,
    rng,
    beta=np.inf,
    proxy_nm=None,
    burn=0,
    record_every=1,
    coupled=True,
):
    """Run R parallel chains of true SG(L)D and (optionally) the proxy on the SAME
    minibatches and the SAME injected noise (synchronous coupling).

    Returns dict with pooled tail iterates of theta and psi and the coupling
    distance E||theta-psi||^2 at stationarity.
    """
    D = model.D
    TH = np.repeat(theta_hat[None, :], R, axis=0)
    PS = TH.copy()
    H = proxy_nm.H if proxy_nm is not None else None
    gh = model.grad_n(theta_hat)  # (N,D)
    Jn = proxy_nm.J if proxy_nm is not None else None
    accT = []
    accP = []
    d2 = []
    noise_scale = 0.0 if not np.isfinite(beta) else 1.0
    Lch = None
    if noise_scale:
        Lch = np.linalg.cholesky(2.0 / beta * Lam)
    for t in range(T):
        IDX = rng.integers(0, model.N, size=(R, B))
        gT = model.grad_batch(TH, IDX)
        if proxy_nm is not None and coupled:
            # proxy minibatch gradient: mean_{n in S} [ grad_n(that) + J_n (psi-that) ] + Gamma psi /N
            U = PS - theta_hat
            gb = gh[IDX].mean(1)
            Ju = np.einsum("rbde,re->rbd", Jn[IDX], U).mean(1)
            gP = (
                gb
                + Ju
                + U @ (model.Gamma.T / model.N)
                + theta_hat @ (model.Gamma.T / model.N)
            )
        xi = rng.standard_normal((R, D)) if noise_scale else None
        TH = TH - gT @ Lam.T
        if noise_scale:
            TH = TH + xi @ Lch.T
        if proxy_nm is not None and coupled:
            PS = PS - gP @ Lam.T
            if noise_scale:
                PS = PS + xi @ Lch.T
        if t >= burn and (t - burn) % record_every == 0:
            accT.append(TH.copy())
            if proxy_nm is not None and coupled:
                accP.append(PS.copy())
                d2.append(((TH - PS) ** 2).sum(1).mean())
    out = dict(theta=np.concatenate(accT, 0))
    if accP:
        out["psi"] = np.concatenate(accP, 0)
        out["d2"] = np.array(d2)
    return out


# --------------------------------------------------------------------------
# fast synchronously-coupled simulator for logistic regression
# --------------------------------------------------------------------------
def run_coupled(
    model, theta_hat, Lam, B, T, R, rng, beta=np.inf, burn=None, thin=0, nblock=8
):
    """Run R chains of the TRUE SG(L)D and of the PROXY on identical minibatches
    and identical injected noise (synchronous coupling).

    Returns accumulated stationary moments plus a thinned sample for marginal-W2.
    Blocks of chains give jackknife error bars.
    """
    N, D = model.N, model.D
    X, y, Gam = model.X, model.y, model.Gamma / model.N
    J_w = None
    burn = T // 4 if burn is None else burn
    TH = np.repeat(theta_hat[None, :], R, axis=0)
    PS = TH.copy()
    gh = model.grad_n(theta_hat)  # (N,D)
    p_hat = 1.0 / (1.0 + np.exp(-X @ theta_hat))
    w_hat = p_hat * (1 - p_hat)  # J_n = w_hat[n] x_n x_n^T
    LamT = Lam.T
    noisy = np.isfinite(beta)
    Lch = np.linalg.cholesky(2.0 / beta * Lam) if noisy else None
    blk = np.arange(R) % nblock
    accT = np.zeros((nblock, D, D))
    accP = np.zeros((nblock, D, D))
    sT = np.zeros((nblock, D))
    sP = np.zeros((nblock, D))
    accD2 = np.zeros(nblock)
    cnt = np.zeros(nblock)
    keepT, keepP = [], []
    for t in range(T):
        IDX = rng.integers(0, N, size=(R, B))
        Xb = X[IDX]  # (R,B,D)
        zT = np.einsum("rbd,rd->rb", Xb, TH)
        gT = (
            np.einsum("rb,rbd->rd", 1.0 / (1.0 + np.exp(-zT)) - y[IDX], Xb) / B
            + TH @ Gam.T
        )
        U = PS - theta_hat
        sU = np.einsum("rbd,rd->rb", Xb, U)
        gP = (
            gh[IDX].mean(1)
            + np.einsum("rb,rbd->rd", w_hat[IDX] * sU, Xb) / B
            + U @ Gam.T
            + theta_hat @ Gam.T
        )
        TH = TH - gT @ LamT
        PS = PS - gP @ LamT
        if noisy:
            xi = rng.standard_normal((R, D)) @ Lch.T
            TH = TH + xi
            PS = PS + xi
        if t >= burn:
            dv = TH - PS
            for b in range(nblock):
                mk = blk == b
                accT[b] += TH[mk].T @ TH[mk]
                sT[b] += TH[mk].sum(0)
                accP[b] += PS[mk].T @ PS[mk]
                sP[b] += PS[mk].sum(0)
                accD2[b] += (dv[mk] ** 2).sum()
                cnt[b] += mk.sum()
            if thin and (t - burn) % thin == 0:
                keepT.append(TH.copy())
                keepP.append(PS.copy())
    return dict(
        accT=accT,
        accP=accP,
        sT=sT,
        sP=sP,
        accD2=accD2,
        cnt=cnt,
        thetaS=np.concatenate(keepT, 0) if keepT else None,
        psiS=np.concatenate(keepP, 0) if keepP else None,
    )


def cov_from_blocks(acc, s, cnt, blocks=None):
    """Pooled covariance over a subset of chain-blocks."""
    idx = slice(None) if blocks is None else blocks
    A = acc[idx].sum(0)
    m = s[idx].sum(0)
    n = cnt[idx].sum()
    return A / n - np.outer(m / n, m / n)


def jackknife(fn, nblock):
    """Leave-one-block-out jackknife standard error of scalar statistic fn(mask)."""
    full = fn(np.ones(nblock, bool))
    vals = []
    for b in range(nblock):
        mk = np.ones(nblock, bool)
        mk[b] = False
        vals.append(fn(mk))
    vals = np.array(vals)
    se = np.sqrt((nblock - 1) / nblock * ((vals - vals.mean()) ** 2).sum())
    return full, float(se)


def w2_marginal_lower(A, Bs):
    """max_d  W2 of the 1-D marginals -- a valid lower bound on the joint W2."""
    D = A.shape[1]
    out = 0.0
    for d in range(D):
        a = np.sort(A[:, d])
        b = np.sort(Bs[:, d])
        n = min(len(a), len(b))
        qa = np.quantile(a, np.linspace(0.0005, 0.9995, 4000))
        qb = np.quantile(b, np.linspace(0.0005, 0.9995, 4000))
        out = max(out, float(np.sqrt(((qa - qb) ** 2).mean())))
    return out