scripts/linreg.py

"""Linear regression under log loss and under the beta-divergence loss
(beta = 1.5), plus the four step-size tuning rules that Algorithm 1 dispatches on.

beta-divergence for a Gaussian working model f(y;theta,z)=N(y; z'theta, sigma^2):
    l(theta;x) = -1/(beta-1) f^{beta-1} + (1/beta) int f(y';theta,z)^beta dy'
The integral is (2 pi sigma^2)^{-(beta-1)/2} beta^{-1/2}, which does NOT depend on
theta, so it drops out of every derivative (the authors' released code estimates
it by Monte Carlo; the closed form used here is exact).  With r = y - z'theta,

    grad l = -f^{beta-1} r z / sigma^2
    hess l =  f^{beta-1} [ 1 - (beta-1) r^2 / sigma^2 ] z z' / sigma^2
"""

import numpy as np
from scipy.optimize import minimize


class LinReg:
    def __init__(self, X, y, sigma=1.0, loss="log", beta=1.5, z_samples=None):
        self.X = np.asarray(X, float)
        self.y = np.asarray(y, float)
        self.N, self.D = self.X.shape
        self.sigma = sigma
        self.loss = loss
        self.beta = beta
        self.Gamma = np.zeros((self.D, self.D))  # no prior, as in Table 3
        # 'beta_mc' replicates the authors' released implementation, which replaces
        # the (theta-independent) integral (1/beta) int f(z;theta)^beta dz by the
        # plain average (1/beta) mean_j f(z_j;theta)^beta over FIXED draws
        # z_j ~ N(0,(5 sigma)^2).  That average is not an unbiased estimator of the
        # integral and, unlike the integral, does depend on theta, so it changes the
        # loss.  Both variants are reported.
        self.z = None if z_samples is None else np.asarray(z_samples, float)

    def _f_pow(self, r):
        s2 = self.sigma**2
        c = (2 * np.pi * s2) ** (-0.5)
        return (c ** (self.beta - 1)) * np.exp(-(self.beta - 1) * r**2 / (2 * s2))

    def _mc_terms(self, mu):
        """returns (dOmega/dmu, d2Omega/dmu^2) for the authors' fixed-z surrogate"""
        s2 = self.sigma ** 2
        c = (2 * np.pi * s2) ** (-0.5)
        d = self.z[None, :] - mu[:, None]                      # (N, n_mc)
        fb = (c ** self.beta) * np.exp(-self.beta * d ** 2 / (2 * s2))
        g = (fb * d / s2).mean(1)
        h = (fb / s2 * (self.beta * d ** 2 / s2 - 1.0)).mean(1)
        return g, h

    def value(self, th):
        r = self.y - self.X @ th
        if self.loss == "log":
            return float((r**2).mean() / (2 * self.sigma**2))
        v = float((-self._f_pow(r) / (self.beta - 1)).mean())
        if self.loss == "beta_mc":
            s2 = self.sigma ** 2
            c = (2 * np.pi * s2) ** (-0.5)
            mu = self.X @ th
            d = self.z[None, :] - mu[:, None]
            fb = (c ** self.beta) * np.exp(-self.beta * d ** 2 / (2 * s2))
            v += float((fb.mean(1) / self.beta).mean())
        return v

    def grad_n(self, th, X=None, y=None):
        X = self.X if X is None else X
        y = self.y if y is None else y
        r = y - X @ th
        if self.loss == "log":
            w = -r / self.sigma**2
        else:
            w = -self._f_pow(r) * r / self.sigma**2
            if self.loss == "beta_mc":
                mu = X @ th
                s2 = self.sigma ** 2
                c = (2 * np.pi * s2) ** (-0.5)
                d = self.z[None, :] - mu[:, None]
                fb = (c ** self.beta) * np.exp(-self.beta * d ** 2 / (2 * s2))
                w = w + (fb * d / s2).mean(1)
        return w[:, None] * X

    def hess_n(self, th):
        r = self.y - self.X @ th
        if self.loss == "log":
            w = np.ones(self.N) / self.sigma**2
        else:
            w = (
                self._f_pow(r)
                * (1 - (self.beta - 1) * r**2 / self.sigma**2)
                / self.sigma**2
            )
            if self.loss == "beta_mc":
                w = w + self._mc_terms(self.X @ th)[1]
        return w[:, None, None] * np.einsum("nd,ne->nde", self.X, self.X)

    def fit(self, th0=None):
        if self.loss == "log":
            return np.linalg.lstsq(self.X, self.y, rcond=None)[0]
        th0 = np.linalg.lstsq(self.X, self.y, rcond=None)[0] if th0 is None else th0
        r = minimize(
            self.value,
            th0,
            jac=lambda t: self.grad_n(t).mean(0),
            method="L-BFGS-B",
            options=dict(maxiter=5000, ftol=1e-16, gtol=1e-12),
        )
        return r.x

    # ------------------------------------------------------------------
    def sgd_paths(self, th0, Lam, B, n_iters, R, rng, burn_frac=0.5):
        """R independent preconditioned-SGD chains; returns per-chain tail covariance."""
        X, y = self.X, self.y
        TH = np.repeat(np.asarray(th0, float)[None, :], R, 0)
        LamT = np.asarray(Lam, float).T
        start = int(burn_frac * n_iters)
        acc = np.zeros((R, self.D, self.D))
        s = np.zeros((R, self.D))
        cnt = 0
        diverged = np.zeros(R, bool)
        s2 = self.sigma**2
        for t in range(n_iters):
            IDX = rng.integers(0, self.N, size=(R, B))
            Xb, yb = X[IDX], y[IDX]
            r = yb - np.einsum("rbd,rd->rb", Xb, TH)
            if self.loss == "log":
                w = -r / s2
            else:
                w = -self._f_pow(r) * r / s2
                if self.loss == "beta_mc":
                    mu = np.einsum("rbd,rd->rb", Xb, TH)
                    c = (2 * np.pi * s2) ** (-0.5)
                    dd = self.z[None, None, :] - mu[:, :, None]
                    fb = (c ** self.beta) * np.exp(-self.beta * dd**2 / (2 * s2))
                    w = w + (fb * dd / s2).mean(2)
            g = np.einsum("rb,rbd->rd", w, Xb) / B
            TH = TH - g @ LamT
            bad = ~np.isfinite(TH).all(1) | (np.abs(TH).max(1) > 1e8)
            diverged |= bad
            TH[bad] = 0.0
            if t >= start:
                acc += np.einsum("rd,re->rde", TH, TH)
                s += TH
                cnt += 1
        covs = acc / cnt - np.einsum("rd,re->rde", s / cnt, s / cnt)
        return covs, diverged


# ---------------------------------------------------------------- tunings --
def sandwich(J, I, N):
    Ji = np.linalg.inv(J)
    return Ji @ I @ Ji / N


def tune_Lambda(
    method,
    H,
    S,
    Ihat,
    Jhat,
    B,
    nm=None,
    A_cov=None,
    sig2=None,
    beta_temp=np.inf,
    sigma=1.0,
):
    """Return (Lambda, residual, closed_form_start) for one row of Algorithm 1 line 6.

    LR+WS uses Ziyin et al. (2022) Eq. (6),
        C = B^{-1} ( A Sigma A + Tr[A Sigma] A + sigma_r^2 A ),
    which the paper states for the sigma = 1 parameterisation of the Gaussian
    working model.  For a general working scale sigma the per-sample Hessian is
    z z' / sigma^2 and the score is -r z / sigma^2, so every term carries an extra
    sigma^{-4}; that scale-covariant form is what is implemented here (it reduces
    to Eq. (6) exactly when sigma = 1).
    """
    from common import solve_Lambda

    s4 = sigma**4
    if method == "CT":
        # Eq. (3):  Lambda = (Sigma H + H Sigma) Chat^{-1},  Chat = Cov(G_1(that)) = Ihat/B
        Lam = (S @ H + H @ S) @ np.linalg.inv(Ihat / B)
        return Lam, 0.0, Lam
    if method == "DQ+const":
        Cf = lambda Sig: Jhat / B  # noqa: E731
    elif method == "LR+WS":
        Cf = lambda Sig: (
            A_cov @ Sig @ A_cov
            + np.trace(A_cov @ Sig) * A_cov  # noqa: E731
            + sig2 * A_cov
        ) / (B * s4)
    elif method == "DQ+exact":
        Cf = lambda Sig: nm.C(Sig, B)  # noqa: E731
    else:
        raise ValueError(method)
    Lam, res, ok, Lam0 = solve_Lambda(H, S, Cf, beta=beta_temp)
    return Lam, res, Lam0