1"""Distributed estimators (in-memory, statistically faithful to the coordinator/worker protocol).23All routines take a batched dataset `db` = {'X': (K, m, p), 'y': (K, m)} (or {'x': (K, m)}) so that block k is4`db[...][k]`. Worker-side quantities are computed block by block; coordinator-side quantities are block5averages. Nothing here measures communication; see `communication.py` (accounting) and `protocol.py` (real runs).6"""7from __future__ import annotations8import numpy as np9from .jackknife import vandermonde_weights, DEFAULT_SPLITS10from .likelihoods import STATUS_OK111213# ----------------------------------------------------------------------------- local fits14def local_split_fits(model, db, splits=(1,), **fit_kw):15 """Fit the MLE on every sub-block of every block for each split multiplier s in `splits`.1617 Returns dict s -> (theta (K, s, p), status (K, s), hessian (K, s, p, p)).18 """19 K = next(iter(db.values())).shape[0]20 m = model.n_obs(db)21 out = {}22 for s in splits:23 if m % s:24 raise ValueError(f"m={m} not divisible by s={s}")25 sub = {k: v.reshape((K * s, m // s) + v.shape[2:]) for k, v in db.items()}26 th, info = model.fit_batched(sub, **fit_kw)27 out[s] = (th.reshape(K, s, -1), info["status"].reshape(K, s), info["hessian"].reshape(K, s, th.shape[1], th.shape[1]))28 return out293031def worker_corrected_estimates(fits, splits):32 """J_{r,k} = sum_j w_j * mean_l theta_hat(sub-block l of level j) -> (K, p); also validity mask (K,)."""33 w = vandermonde_weights(splits)34 J = 0.0; ok = None35 for wj, s in zip(w, splits):36 th, st, _ = fits[s]37 J = J + wj * th.mean(axis=1)38 good = np.all(st == STATUS_OK, axis=1)39 ok = good if ok is None else (ok & good)40 return J, ok414243def pilot_estimator(J, block_sizes=None, ok=None):44 """Coordinator average of the worker estimates (weights m_k/N if sizes given). Invalid workers can be dropped."""45 J = np.asarray(J)46 if block_sizes is None:47 wts = np.ones(J.shape[0])48 else:49 wts = np.asarray(block_sizes, float)50 if ok is not None:51 wts = wts * ok52 wts = wts / wts.sum()53 return np.einsum('k,kp->p', wts, J)545556# ----------------------------------------------------------------------------- aggregated derivatives57def aggregated_score_hessian(model, db, theta, block_sizes=None):58 """Workers evaluate (score, Hessian) of their block at the common theta; coordinator averages (weighted by m_k)."""59 K = next(iter(db.values())).shape[0]60 th = np.broadcast_to(np.asarray(theta, float), (K, len(theta)))61 g, H = model.score_hessian(th, db)62 if block_sizes is None:63 return g.mean(axis=0), H.mean(axis=0), g, H64 w = np.asarray(block_sizes, float); w = w / w.sum()65 return np.einsum('k,kp->p', w, g), np.einsum('k,kpq->pq', w, H), g, H666768def newton_step(theta, g, H, damping=None):69 """theta - H^{-1} g via a linear solve (never an explicit inverse). Returns (new theta, diagnostics)."""70 try:71 step = np.linalg.solve(-H, g)72 ok = np.all(np.isfinite(step))73 except np.linalg.LinAlgError:74 step = np.zeros_like(g); ok = False75 diag = {"solve_ok": ok, "cond": float(np.linalg.cond(-H)) if ok else np.inf, "step_norm": float(np.linalg.norm(step))}76 if damping is not None and ok:77 step = step * damping78 return theta + step, diag798081def global_newton(model, db, theta0, q=1, hessian_mode="exact", block_sizes=None, stale_hessian=None,82 quantize_bits=None, single_machine=0, rng=None):83 """q globally aggregated Newton corrections from theta0.8485 hessian_mode:86 'exact' workers send score and Hessian at the broadcast iterate (p + p(p+1)/2 numbers each round);87 'stale' Hessian fixed = `stale_hessian` (e.g. average of local Hessians at local MLEs, sent once with the pilot);88 workers send only scores each round (p numbers);89 'single' Hessian of worker `single_machine` evaluated at the iterate (no Hessian uplink; CSL-type curvature);90 'quantized' exact aggregated Hessian rounded to `quantize_bits` significant bits per entry (uniform relative91 quantization), scores at full precision.92 Returns list of iterates [theta0, theta1, ..., thetaq] and per-step diagnostics.93 """94 th = np.asarray(theta0, float).copy(); path = [th.copy()]; diags = []95 for _ in range(q):96 g, H, g_k, H_k = aggregated_score_hessian(model, db, th, block_sizes)97 if hessian_mode == "stale":98 H = np.asarray(stale_hessian)99 elif hessian_mode == "single":100 H = H_k[single_machine]101 elif hessian_mode == "quantized":102 H = quantize_relative(H, quantize_bits)103 elif hessian_mode != "exact":104 raise ValueError(hessian_mode)105 th, d = newton_step(th, g, H)106 d["grad_norm"] = float(np.linalg.norm(g))107 path.append(th.copy()); diags.append(d)108 return path, diags109110111def quantize_relative(A, bits):112 """Round each entry to `bits` significant binary digits (mantissa truncation), keeping the sign and exponent."""113 if bits is None:114 return A115 A = np.asarray(A, float)116 mant, exp = np.frexp(A)117 scale = 2.0 ** bits118 return np.ldexp(np.round(mant * scale) / scale, exp)119120121# ----------------------------------------------------------------------------- baselines122def aee_lin_xi(theta_local, hess_local):123 """Lin & Xi (2011) aggregated estimating equation: (sum_k A_k)^{-1} sum_k A_k theta_k with A_k = -H_k(theta_k)."""124 A = -np.asarray(hess_local) # (K,p,p)125 rhs = np.einsum('kpq,kq->p', A, theta_local)126 return np.linalg.solve(A.sum(axis=0), rhs)127128129def csl(model, db, theta_init, rounds=1, machine=0, block_sizes=None, fit_kw=None):130 """Jordan, Lee & Yang (2019) communication-efficient surrogate likelihood, maximized on `machine`.131132 Round t: workers send gradients at theta_t (p numbers each); machine 0 maximizes133 L_tilde(theta) = L_0(theta) - <grad L_0(theta_t) - grad L_N(theta_t), theta>134 by Newton's method on its own block (local Hessian only). Returns iterates.135 """136 fit_kw = fit_kw or {}137 th = np.asarray(theta_init, float).copy(); path = [th.copy()]138 d0 = {k: v[machine][None] for k, v in db.items()}139 for _ in range(rounds):140 gN, _, g_k, _ = aggregated_score_hessian(model, db, th, block_sizes)141 shift = g_k[machine] - gN # p-vector142 # maximize L_0(theta) - shift' theta with Newton on machine 0143 x = th.copy()144 for it in range(100):145 g, H = model.score_hessian(x[None], d0)146 g = g[0] - shift; H = H[0]147 step = np.linalg.solve(-H, g)148 # backtracking on the surrogate objective149 f = lambda z: model.loglik(z[None], d0)[0] - shift @ z150 f0 = f(x); t = 1.0151 while f(x + t * step) < f0 - 1e-12 and t > 1e-8:152 t *= 0.5153 x = x + t * step154 if np.max(np.abs(g)) < 1e-10:155 break156 th = x; path.append(th.copy())157 return path158159160def analytic_bias_correction(model, db, theta_local):161 """Cox–Snell / Cordeiro–McCullagh O(1/m) bias correction of each local MLE (GLM canonical link, or Exponential):162 b_hat = (1/2) I^{-1} [n^{-1} sum_i nabla^2 psi_i (I^{-1})] per block; theta_corrected = theta - b_hat / m.163 (Derivation: theta_hat - theta0 = -H^{-1} Psi - H^{-1}(V delta) - (1/2) H^{-1} nabla^2 Psi[delta,delta] + ...; with H = -I,164 E[delta delta^T] = I^{-1}/n and the V-term vanishing conditionally on the design for canonical GLMs.)165 Equivalent to the general second-order bias of M-estimators with the first (score-covariance) term dropped,166 which is exact for canonical GLMs conditional on the design.167 """168 K, m = next(iter(db.values())).shape[:2]169 g, H = model.score_hessian(theta_local, db)170 I = -H171 Iinv = np.linalg.inv(I) # p small; explicit inverse acceptable for a diagnostic quantity172 T = model.third_contraction(theta_local, db, Iinv) # (K,p) = n^{-1} sum nabla^2 psi_i [I^{-1}]173 b = 0.5 * np.einsum("kpq,kq->kp", Iinv, T) # per-block O(1) bias coefficient b_1 = (1/2) I^{-1} T174 return theta_local - b / m, b175176177# ----------------------------------------------------------------------------- one call = all estimators on one sample178def estimate_all(model, db, splits_max=3, q_max=2, block_sizes=None, fit_kw=None, extras=True):179 """Compute, on the same partitioned sample, the full family of estimators used in the Monte Carlo.180181 Returns dict name -> theta (p,), plus a dict of diagnostics (validity masks, Newton diagnostics, local fit stats).182 Names: raw (r=1,q=0) ... 'r{r}_q{q}' for r=1..splits_max, q=0..q_max with splits (1,2,4,..);183 'central' centralized MLE; 'aee' Lin–Xi; 'csl1','csl2' CSL rounds from the raw average;184 'cs_q{q}' analytic Cox–Snell corrected average + q Newton; 'stale_r{r}_q{q}' stale-Hessian variants.185 """186 fit_kw = fit_kw or {}187 splits = DEFAULT_SPLITS[splits_max]188 fits = local_split_fits(model, db, splits, **fit_kw)189 K, m = next(iter(db.values())).shape[:2]190 out, diag = {}, {}191 diag["local_status_counts"] = {s: np.bincount(fits[s][1].ravel(), minlength=4).tolist() for s in splits}192 # centralized MLE on the pooled data193 full = {k: v.reshape((1, K * m) + v.shape[2:]) for k, v in db.items()}194 th_c, info_c = model.fit_batched(full, tol=1e-12, maxit=100)195 out["central"] = th_c[0]; diag["central_status"] = int(info_c["status"][0])196 theta_full_local, status_full, H_local = fits[1]197 H_bar = H_local[:, 0].mean(axis=0) # average local Hessian at local MLEs (sent with the pilot)198 for r in range(1, splits_max + 1):199 sp = splits[:r]200 J, ok = worker_corrected_estimates({s: fits[s] for s in sp}, sp)201 diag[f"valid_r{r}"] = ok202 th0 = pilot_estimator(J, block_sizes, ok=None) # keep all workers (failures tracked separately)203 path, nd = global_newton(model, db, th0, q=q_max, block_sizes=block_sizes)204 for q in range(q_max + 1):205 out[f"r{r}_q{q}"] = path[q]206 diag[f"newton_r{r}"] = nd207 if extras:208 path_s, _ = global_newton(model, db, th0, q=q_max, hessian_mode="stale", stale_hessian=H_bar, block_sizes=block_sizes)209 for q in range(1, q_max + 1):210 out[f"stale_r{r}_q{q}"] = path_s[q]211 if extras:212 out["aee"] = aee_lin_xi(theta_full_local[:, 0], H_local[:, 0])213 raw = out["r1_q0"]214 cpath = csl(model, db, raw, rounds=2, block_sizes=block_sizes)215 out["csl1"], out["csl2"] = cpath[1], cpath[2]216 try:217 th_cs, b = analytic_bias_correction(model, db, theta_full_local[:, 0])218 th0 = pilot_estimator(th_cs, block_sizes)219 pth, _ = global_newton(model, db, th0, q=q_max, block_sizes=block_sizes)220 for q in range(q_max + 1):221 out[f"cs_q{q}"] = pth[q]222 except NotImplementedError:223 pass224 return out, diag225