"""Likelihood models with batched (per-block) Newton fitting. Data layout ----------- A dataset is a dict of arrays whose leading axis indexes observations: GLMs use {'X': (n, p), 'y': (n,)}; the exponential model uses {'x': (n,)}. A *batched* dataset has one extra leading axis for blocks: {'X': (B, n, p), 'y': (B, n)}. All model methods accept batched data and batched parameters (B, p). Status codes returned by `fit_batched`: 0 converged, 1 max iterations, 2 singular/indefinite Hessian, 3 non-finite or diverged (e.g. logistic separation). """ from __future__ import annotations import numpy as np STATUS_OK, STATUS_MAXIT, STATUS_SINGULAR, STATUS_DIVERGED = 0, 1, 2, 3 def gather_blocks(data: dict, blocks: np.ndarray) -> dict: """Index a dataset by an integer array `blocks` of shape (K, m) -> batched dataset (K, m, ...).""" return {k: v[blocks] for k, v in data.items()} def subset(data: dict, idx) -> dict: return {k: v[idx] for k, v in data.items()} def _batched_solve(A, b): """Solve A x = b for batched A (B,p,p), b (B,p); returns x and a boolean mask of failures.""" B, p = b.shape x = np.zeros_like(b); fail = np.zeros(B, dtype=bool) try: x = np.linalg.solve(A, b[..., None])[..., 0] bad = ~np.all(np.isfinite(x), axis=1) except np.linalg.LinAlgError: bad = np.ones(B, dtype=bool) if np.any(bad): for i in np.flatnonzero(bad): try: x[i] = np.linalg.solve(A[i], b[i]); fail[i] = not np.all(np.isfinite(x[i])) except np.linalg.LinAlgError: fail[i] = True; x[i] = 0.0 return x, fail class Model: """Abstract model. Subclasses implement the vectorised primitives below.""" name = "abstract" p = None # ---- primitives (batched: theta (B,p), data batched) ---- def loglik(self, theta, data): # -> (B,) raise NotImplementedError def score_hessian(self, theta, data): # -> (B,p), (B,p,p) averages per block raise NotImplementedError def third_contraction(self, theta, data, Sigma): """n^{-1} sum_i nabla^2 psi_i(theta)[Sigma] -> (B,p); psi_i = per-observation score.""" raise NotImplementedError def outer_scores(self, theta, data): """n^{-1} sum_i psi_i psi_i' -> (B,p,p) (for sandwich covariances).""" raise NotImplementedError def simulate(self, rng, n, theta0, **kw) -> dict: raise NotImplementedError def init(self, data): # starting value (B,p) raise NotImplementedError def n_obs(self, data): raise NotImplementedError def degenerate(self, theta, data): """Boolean mask (B,): a stationary point that should NOT be accepted as an MLE (e.g. logistic separation: fitted probabilities numerically 0/1 so the score vanishes while the Hessian collapses).""" return np.zeros(theta.shape[0], dtype=bool) # ---- generic Newton with step halving ---- def fit_batched(self, data, init=None, tol=1e-9, maxit=60, max_norm=1e4): theta = np.array(self.init(data) if init is None else init, dtype=float) if theta.ndim == 1: theta = theta[None, :] B = theta.shape[0] status = np.full(B, STATUS_MAXIT, dtype=int) iters = np.zeros(B, dtype=int) active = np.ones(B, dtype=bool) ll = self.loglik(theta, data) for it in range(maxit): g, H = self.score_hessian(theta, data) gmax = np.max(np.abs(g), axis=1) conv = active & (gmax < tol) if np.any(conv): deg = conv & self.degenerate(theta, data) status[deg] = STATUS_DIVERGED; conv &= ~deg status[conv] = STATUS_OK; active &= ~(conv | (status == STATUS_DIVERGED)) if not np.any(active): break step, fail = _batched_solve(-H, g) if np.any(fail & active): status[fail & active] = STATUS_SINGULAR; active &= ~fail # ascent check: if the direction is not an ascent direction, fall back to gradient asc = np.einsum('bp,bp->b', step, g) bad = active & (asc <= 0) step[bad] = g[bad] t = np.ones(B) new = theta + t[:, None] * step ll_new = self.loglik(new, data) for _ in range(30): worse = active & (~np.isfinite(ll_new) | (ll_new < ll - 1e-12)) if not np.any(worse): break t[worse] *= 0.5 new = theta + t[:, None] * step ll_new = self.loglik(new, data) theta = np.where(active[:, None], new, theta) ll = np.where(active, ll_new, ll) iters[active] += 1 div = active & (~np.all(np.isfinite(theta), axis=1) | (np.linalg.norm(theta, axis=1) > max_norm)) status[div] = STATUS_DIVERGED; active &= ~div g, H = self.score_hessian(theta, data) info = {"status": status, "iters": iters, "grad_max": np.max(np.abs(g), axis=1), "hessian": H, "cond": np.linalg.cond(-H)} return theta, info def fit(self, data, init=None, **kw): db = {k: v[None] for k, v in data.items()} th, info = self.fit_batched(db, init=None if init is None else np.asarray(init)[None], **kw) return th[0], {k: (v[0] if isinstance(v, np.ndarray) else v) for k, v in info.items()} class Exponential(Model): """X ~ Exp(rate theta); average log-likelihood log(theta) - theta * xbar; MLE 1/xbar.""" name = "exponential"; p = 1 def n_obs(self, data): return data['x'].shape[-1] def simulate(self, rng, n, theta0, **kw): return {'x': rng.exponential(scale=1.0 / float(np.ravel(theta0)[0]), size=n)} def init(self, data): return 1.0 / np.mean(data['x'], axis=-1)[:, None] def loglik(self, theta, data): th = theta[:, 0] return np.log(th) - th * np.mean(data['x'], axis=-1) def score_hessian(self, theta, data): th = theta[:, 0] g = (1.0 / th - np.mean(data['x'], axis=-1))[:, None] H = (-1.0 / th ** 2)[:, None, None] return g, H def third_contraction(self, theta, data, Sigma): th = theta[:, 0] return (2.0 / th ** 3 * Sigma[:, 0, 0])[:, None] def outer_scores(self, theta, data): th = theta[:, 0][:, None] psi = 1.0 / th - data['x'] return np.mean(psi ** 2, axis=-1)[:, None, None] def mle(self, data): """Closed form, batched.""" return (1.0 / np.mean(data['x'], axis=-1))[:, None] class _GLM(Model): """Canonical-link GLM with per-observation score x_i (y_i - mu_i), observed Hessian -sum x_i x_i' mu'(eta_i)/n.""" def n_obs(self, data): return data['y'].shape[-1] def _eta(self, theta, data): return np.einsum('bnp,bp->bn', data['X'], theta) def score_hessian(self, theta, data): X = data['X']; y = data['y']; n = y.shape[-1] eta = self._eta(theta, data) mu, w = self._mu_w(eta) g = np.einsum('bnp,bn->bp', X, y - mu) / n H = -np.einsum('bnp,bn,bnq->bpq', X, w, X) / n return g, H def third_contraction(self, theta, data, Sigma): X = data['X']; n = X.shape[1] eta = self._eta(theta, data) mu2 = self._mu2(eta) # mu''(eta) quad = np.einsum('bnp,bpq,bnq->bn', X, Sigma, X) # x_i' Sigma x_i return -np.einsum('bnp,bn,bn->bp', X, mu2, quad) / n def outer_scores(self, theta, data): X = data['X']; y = data['y']; n = y.shape[-1] mu, _ = self._mu_w(self._eta(theta, data)) r = y - mu return np.einsum('bnp,bn,bnq->bpq', X, r * r, X) / n def fisher(self, theta, data): """Expected information = -observed Hessian for canonical links (B,p,p).""" _, H = self.score_hessian(theta, data); return -H class Logistic(_GLM): name = "logistic" def __init__(self, p): self.p = p def init(self, data): B = data['y'].shape[0] return np.zeros((B, self.p)) def _mu_w(self, eta): mu = 1.0 / (1.0 + np.exp(-eta)); return mu, mu * (1.0 - mu) def _mu2(self, eta): mu = 1.0 / (1.0 + np.exp(-eta)); return mu * (1.0 - mu) * (1.0 - 2.0 * mu) def degenerate(self, theta, data, eta_max=30.0): """Separation safeguard: some fitted |eta| > eta_max (probability within 1e-13 of 0 or 1) or a numerically singular information matrix. Legitimate designs in this paper never reach |eta| = 30.""" eta = self._eta(theta, data) _, H = self.score_hessian(theta, data) return (np.max(np.abs(eta), axis=-1) > eta_max) | (np.linalg.cond(-H) > 1e12) def loglik(self, theta, data): eta = self._eta(theta, data); y = data['y'] return np.mean(y * eta - np.logaddexp(0.0, eta), axis=-1) def simulate(self, rng, n, theta0, cov=None, intercept=True, **kw): theta0 = np.asarray(theta0, float); p = theta0.size q = p - 1 if intercept else p Z = rng.standard_normal((n, q)) if cov is None else rng.multivariate_normal(np.zeros(q), cov, size=n) X = np.hstack([np.ones((n, 1)), Z]) if intercept else Z eta = X @ theta0 y = (rng.random(n) < 1.0 / (1.0 + np.exp(-eta))).astype(float) return {'X': X, 'y': y} class Poisson(_GLM): name = "poisson" def __init__(self, p): self.p = p def init(self, data): B = data['y'].shape[0] th = np.zeros((B, self.p)); th[:, 0] = np.log(np.maximum(np.mean(data['y'], axis=-1), 1e-3)) return th def _mu_w(self, eta): mu = np.exp(np.minimum(eta, 50.0)); return mu, mu def _mu2(self, eta): return np.exp(np.minimum(eta, 50.0)) def loglik(self, theta, data): eta = self._eta(theta, data); y = data['y'] return np.mean(y * eta - np.exp(np.minimum(eta, 50.0)), axis=-1) def simulate(self, rng, n, theta0, cov=None, intercept=True, **kw): theta0 = np.asarray(theta0, float); p = theta0.size q = p - 1 if intercept else p Z = rng.standard_normal((n, q)) if cov is None else rng.multivariate_normal(np.zeros(q), cov, size=n) X = np.hstack([np.ones((n, 1)), Z]) if intercept else Z y = rng.poisson(np.exp(X @ theta0)).astype(float) return {'X': X, 'y': y} def make_model(name: str, p: int = 1) -> Model: name = name.lower() if name == "exponential": return Exponential() if name == "logistic": return Logistic(p) if name == "poisson": return Poisson(p) raise ValueError(name)