1"""Likelihood models with batched (per-block) Newton fitting.23Data layout4-----------5A dataset is a dict of arrays whose leading axis indexes observations: GLMs use {'X': (n, p), 'y': (n,)};6the exponential model uses {'x': (n,)}. A *batched* dataset has one extra leading axis for blocks:7{'X': (B, n, p), 'y': (B, n)}. All model methods accept batched data and batched parameters (B, p).89Status codes returned by `fit_batched`: 0 converged, 1 max iterations, 2 singular/indefinite Hessian,103 non-finite or diverged (e.g. logistic separation).11"""12from __future__ import annotations13import numpy as np1415STATUS_OK, STATUS_MAXIT, STATUS_SINGULAR, STATUS_DIVERGED = 0, 1, 2, 3161718def gather_blocks(data: dict, blocks: np.ndarray) -> dict:19 """Index a dataset by an integer array `blocks` of shape (K, m) -> batched dataset (K, m, ...)."""20 return {k: v[blocks] for k, v in data.items()}212223def subset(data: dict, idx) -> dict:24 return {k: v[idx] for k, v in data.items()}252627def _batched_solve(A, b):28 """Solve A x = b for batched A (B,p,p), b (B,p); returns x and a boolean mask of failures."""29 B, p = b.shape30 x = np.zeros_like(b); fail = np.zeros(B, dtype=bool)31 try:32 x = np.linalg.solve(A, b[..., None])[..., 0]33 bad = ~np.all(np.isfinite(x), axis=1)34 except np.linalg.LinAlgError:35 bad = np.ones(B, dtype=bool)36 if np.any(bad):37 for i in np.flatnonzero(bad):38 try:39 x[i] = np.linalg.solve(A[i], b[i]); fail[i] = not np.all(np.isfinite(x[i]))40 except np.linalg.LinAlgError:41 fail[i] = True; x[i] = 0.042 return x, fail434445class Model:46 """Abstract model. Subclasses implement the vectorised primitives below."""47 name = "abstract"48 p = None4950 # ---- primitives (batched: theta (B,p), data batched) ----51 def loglik(self, theta, data): # -> (B,)52 raise NotImplementedError53 def score_hessian(self, theta, data): # -> (B,p), (B,p,p) averages per block54 raise NotImplementedError55 def third_contraction(self, theta, data, Sigma):56 """n^{-1} sum_i nabla^2 psi_i(theta)[Sigma] -> (B,p); psi_i = per-observation score."""57 raise NotImplementedError58 def outer_scores(self, theta, data):59 """n^{-1} sum_i psi_i psi_i' -> (B,p,p) (for sandwich covariances)."""60 raise NotImplementedError61 def simulate(self, rng, n, theta0, **kw) -> dict:62 raise NotImplementedError63 def init(self, data): # starting value (B,p)64 raise NotImplementedError65 def n_obs(self, data):66 raise NotImplementedError67 def degenerate(self, theta, data):68 """Boolean mask (B,): a stationary point that should NOT be accepted as an MLE (e.g. logistic separation:69 fitted probabilities numerically 0/1 so the score vanishes while the Hessian collapses)."""70 return np.zeros(theta.shape[0], dtype=bool)7172 # ---- generic Newton with step halving ----73 def fit_batched(self, data, init=None, tol=1e-9, maxit=60, max_norm=1e4):74 theta = np.array(self.init(data) if init is None else init, dtype=float)75 if theta.ndim == 1:76 theta = theta[None, :]77 B = theta.shape[0]78 status = np.full(B, STATUS_MAXIT, dtype=int)79 iters = np.zeros(B, dtype=int)80 active = np.ones(B, dtype=bool)81 ll = self.loglik(theta, data)82 for it in range(maxit):83 g, H = self.score_hessian(theta, data)84 gmax = np.max(np.abs(g), axis=1)85 conv = active & (gmax < tol)86 if np.any(conv):87 deg = conv & self.degenerate(theta, data)88 status[deg] = STATUS_DIVERGED; conv &= ~deg89 status[conv] = STATUS_OK; active &= ~(conv | (status == STATUS_DIVERGED))90 if not np.any(active):91 break92 step, fail = _batched_solve(-H, g)93 if np.any(fail & active):94 status[fail & active] = STATUS_SINGULAR; active &= ~fail95 # ascent check: if the direction is not an ascent direction, fall back to gradient96 asc = np.einsum('bp,bp->b', step, g)97 bad = active & (asc <= 0)98 step[bad] = g[bad]99 t = np.ones(B)100 new = theta + t[:, None] * step101 ll_new = self.loglik(new, data)102 for _ in range(30):103 worse = active & (~np.isfinite(ll_new) | (ll_new < ll - 1e-12))104 if not np.any(worse):105 break106 t[worse] *= 0.5107 new = theta + t[:, None] * step108 ll_new = self.loglik(new, data)109 theta = np.where(active[:, None], new, theta)110 ll = np.where(active, ll_new, ll)111 iters[active] += 1112 div = active & (~np.all(np.isfinite(theta), axis=1) | (np.linalg.norm(theta, axis=1) > max_norm))113 status[div] = STATUS_DIVERGED; active &= ~div114 g, H = self.score_hessian(theta, data)115 info = {"status": status, "iters": iters, "grad_max": np.max(np.abs(g), axis=1),116 "hessian": H, "cond": np.linalg.cond(-H)}117 return theta, info118119 def fit(self, data, init=None, **kw):120 db = {k: v[None] for k, v in data.items()}121 th, info = self.fit_batched(db, init=None if init is None else np.asarray(init)[None], **kw)122 return th[0], {k: (v[0] if isinstance(v, np.ndarray) else v) for k, v in info.items()}123124125class Exponential(Model):126 """X ~ Exp(rate theta); average log-likelihood log(theta) - theta * xbar; MLE 1/xbar."""127 name = "exponential"; p = 1128129 def n_obs(self, data):130 return data['x'].shape[-1]131 def simulate(self, rng, n, theta0, **kw):132 return {'x': rng.exponential(scale=1.0 / float(np.ravel(theta0)[0]), size=n)}133 def init(self, data):134 return 1.0 / np.mean(data['x'], axis=-1)[:, None]135 def loglik(self, theta, data):136 th = theta[:, 0]137 return np.log(th) - th * np.mean(data['x'], axis=-1)138 def score_hessian(self, theta, data):139 th = theta[:, 0]140 g = (1.0 / th - np.mean(data['x'], axis=-1))[:, None]141 H = (-1.0 / th ** 2)[:, None, None]142 return g, H143 def third_contraction(self, theta, data, Sigma):144 th = theta[:, 0]145 return (2.0 / th ** 3 * Sigma[:, 0, 0])[:, None]146 def outer_scores(self, theta, data):147 th = theta[:, 0][:, None]148 psi = 1.0 / th - data['x']149 return np.mean(psi ** 2, axis=-1)[:, None, None]150 def mle(self, data):151 """Closed form, batched."""152 return (1.0 / np.mean(data['x'], axis=-1))[:, None]153154155class _GLM(Model):156 """Canonical-link GLM with per-observation score x_i (y_i - mu_i), observed Hessian -sum x_i x_i' mu'(eta_i)/n."""157 def n_obs(self, data):158 return data['y'].shape[-1]159 def _eta(self, theta, data):160 return np.einsum('bnp,bp->bn', data['X'], theta)161 def score_hessian(self, theta, data):162 X = data['X']; y = data['y']; n = y.shape[-1]163 eta = self._eta(theta, data)164 mu, w = self._mu_w(eta)165 g = np.einsum('bnp,bn->bp', X, y - mu) / n166 H = -np.einsum('bnp,bn,bnq->bpq', X, w, X) / n167 return g, H168 def third_contraction(self, theta, data, Sigma):169 X = data['X']; n = X.shape[1]170 eta = self._eta(theta, data)171 mu2 = self._mu2(eta) # mu''(eta)172 quad = np.einsum('bnp,bpq,bnq->bn', X, Sigma, X) # x_i' Sigma x_i173 return -np.einsum('bnp,bn,bn->bp', X, mu2, quad) / n174 def outer_scores(self, theta, data):175 X = data['X']; y = data['y']; n = y.shape[-1]176 mu, _ = self._mu_w(self._eta(theta, data))177 r = y - mu178 return np.einsum('bnp,bn,bnq->bpq', X, r * r, X) / n179 def fisher(self, theta, data):180 """Expected information = -observed Hessian for canonical links (B,p,p)."""181 _, H = self.score_hessian(theta, data); return -H182183184class Logistic(_GLM):185 name = "logistic"186 def __init__(self, p):187 self.p = p188 def init(self, data):189 B = data['y'].shape[0]190 return np.zeros((B, self.p))191 def _mu_w(self, eta):192 mu = 1.0 / (1.0 + np.exp(-eta)); return mu, mu * (1.0 - mu)193 def _mu2(self, eta):194 mu = 1.0 / (1.0 + np.exp(-eta)); return mu * (1.0 - mu) * (1.0 - 2.0 * mu)195 def degenerate(self, theta, data, eta_max=30.0):196 """Separation safeguard: some fitted |eta| > eta_max (probability within 1e-13 of 0 or 1) or a numerically197 singular information matrix. Legitimate designs in this paper never reach |eta| = 30."""198 eta = self._eta(theta, data)199 _, H = self.score_hessian(theta, data)200 return (np.max(np.abs(eta), axis=-1) > eta_max) | (np.linalg.cond(-H) > 1e12)201 def loglik(self, theta, data):202 eta = self._eta(theta, data); y = data['y']203 return np.mean(y * eta - np.logaddexp(0.0, eta), axis=-1)204 def simulate(self, rng, n, theta0, cov=None, intercept=True, **kw):205 theta0 = np.asarray(theta0, float); p = theta0.size206 q = p - 1 if intercept else p207 Z = rng.standard_normal((n, q)) if cov is None else rng.multivariate_normal(np.zeros(q), cov, size=n)208 X = np.hstack([np.ones((n, 1)), Z]) if intercept else Z209 eta = X @ theta0210 y = (rng.random(n) < 1.0 / (1.0 + np.exp(-eta))).astype(float)211 return {'X': X, 'y': y}212213214class Poisson(_GLM):215 name = "poisson"216 def __init__(self, p):217 self.p = p218 def init(self, data):219 B = data['y'].shape[0]220 th = np.zeros((B, self.p)); th[:, 0] = np.log(np.maximum(np.mean(data['y'], axis=-1), 1e-3))221 return th222 def _mu_w(self, eta):223 mu = np.exp(np.minimum(eta, 50.0)); return mu, mu224 def _mu2(self, eta):225 return np.exp(np.minimum(eta, 50.0))226 def loglik(self, theta, data):227 eta = self._eta(theta, data); y = data['y']228 return np.mean(y * eta - np.exp(np.minimum(eta, 50.0)), axis=-1)229 def simulate(self, rng, n, theta0, cov=None, intercept=True, **kw):230 theta0 = np.asarray(theta0, float); p = theta0.size231 q = p - 1 if intercept else p232 Z = rng.standard_normal((n, q)) if cov is None else rng.multivariate_normal(np.zeros(q), cov, size=n)233 X = np.hstack([np.ones((n, 1)), Z]) if intercept else Z234 y = rng.poisson(np.exp(X @ theta0)).astype(float)235 return {'X': X, 'y': y}236237238def make_model(name: str, p: int = 1) -> Model:239 name = name.lower()240 if name == "exponential":241 return Exponential()242 if name == "logistic":243 return Logistic(p)244 if name == "poisson":245 return Poisson(p)246 raise ValueError(name)247