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 try:127 return np.linalg.solve(A.sum(axis=0), rhs)128 except np.linalg.LinAlgError: # singular aggregated information (failed local fits): report NaN129 return np.full(theta_local.shape[1], np.nan)130131132def csl(model, db, theta_init, rounds=1, machine=0, block_sizes=None, fit_kw=None):133 """Jordan, Lee & Yang (2019) communication-efficient surrogate likelihood, maximized on `machine`.134135 Round t: workers send gradients at theta_t (p numbers each); machine 0 maximizes136 L_tilde(theta) = L_0(theta) - <grad L_0(theta_t) - grad L_N(theta_t), theta>137 by Newton's method on its own block (local Hessian only). Returns iterates.138 """139 fit_kw = fit_kw or {}140 th = np.asarray(theta_init, float).copy(); path = [th.copy()]141 d0 = {k: v[machine][None] for k, v in db.items()}142 for _ in range(rounds):143 gN, _, g_k, _ = aggregated_score_hessian(model, db, th, block_sizes)144 shift = g_k[machine] - gN # p-vector145 # maximize L_0(theta) - shift' theta with Newton on machine 0146 x = th.copy()147 for it in range(100):148 g, H = model.score_hessian(x[None], d0)149 g = g[0] - shift; H = H[0]150 try:151 step = np.linalg.solve(-H, g)152 except np.linalg.LinAlgError:153 x = np.full_like(x, np.nan); break154 # backtracking on the surrogate objective155 f = lambda z: model.loglik(z[None], d0)[0] - shift @ z156 f0 = f(x); t = 1.0157 while f(x + t * step) < f0 - 1e-12 and t > 1e-8:158 t *= 0.5159 x = x + t * step160 if np.max(np.abs(g)) < 1e-10:161 break162 th = x; path.append(th.copy())163 if not np.all(np.isfinite(th)):164 break165 while len(path) < rounds + 1:166 path.append(np.full_like(th, np.nan))167 return path168169170def analytic_bias_correction(model, db, theta_local):171 """Cox–Snell / Cordeiro–McCullagh O(1/m) bias correction of each local MLE (GLM canonical link, or Exponential):172 b_hat = (1/2) I^{-1} [n^{-1} sum_i nabla^2 psi_i (I^{-1})] per block; theta_corrected = theta - b_hat / m.173 (Derivation: theta_hat - theta0 = -H^{-1} Psi - H^{-1}(V delta) - (1/2) H^{-1} nabla^2 Psi[delta,delta] + ...; with H = -I,174 E[delta delta^T] = I^{-1}/n and the V-term vanishing conditionally on the design for canonical GLMs.)175 Equivalent to the general second-order bias of M-estimators with the first (score-covariance) term dropped,176 which is exact for canonical GLMs conditional on the design.177 """178 K, m = next(iter(db.values())).shape[:2]179 g, H = model.score_hessian(theta_local, db)180 I = -H181 try:182 Iinv = np.linalg.inv(I) # p small; explicit inverse acceptable for a diagnostic quantity183 except np.linalg.LinAlgError:184 Iinv = np.linalg.pinv(I) # failed local fits: pseudo-inverse, result flagged through the status codes185 T = model.third_contraction(theta_local, db, Iinv) # (K,p) = n^{-1} sum nabla^2 psi_i [I^{-1}]186 b = 0.5 * np.einsum("kpq,kq->kp", Iinv, T) # per-block O(1) bias coefficient b_1 = (1/2) I^{-1} T187 return theta_local - b / m, b188189190# ----------------------------------------------------------------------------- one call = all estimators on one sample191def estimate_all(model, db, splits_max=3, q_max=2, block_sizes=None, fit_kw=None, extras=True):192 """Compute, on the same partitioned sample, the full family of estimators used in the Monte Carlo.193194 Returns dict name -> theta (p,), plus a dict of diagnostics (validity masks, Newton diagnostics, local fit stats).195 Names: raw (r=1,q=0) ... 'r{r}_q{q}' for r=1..splits_max, q=0..q_max with splits (1,2,4,..);196 'central' centralized MLE; 'aee' Lin–Xi; 'csl1','csl2' CSL rounds from the raw average;197 'cs_q{q}' analytic Cox–Snell corrected average + q Newton; 'stale_r{r}_q{q}' stale-Hessian variants.198 """199 fit_kw = fit_kw or {}200 splits = DEFAULT_SPLITS[splits_max]201 fits = local_split_fits(model, db, splits, **fit_kw)202 K, m = next(iter(db.values())).shape[:2]203 out, diag = {}, {}204 diag["local_status_counts"] = {s: np.bincount(fits[s][1].ravel(), minlength=4).tolist() for s in splits}205 # centralized MLE on the pooled data206 full = {k: v.reshape((1, K * m) + v.shape[2:]) for k, v in db.items()}207 th_c, info_c = model.fit_batched(full, tol=1e-12, maxit=100)208 out["central"] = th_c[0]; diag["central_status"] = int(info_c["status"][0])209 theta_full_local, status_full, H_local = fits[1]210 H_bar = H_local[:, 0].mean(axis=0) # average local Hessian at local MLEs (sent with the pilot)211 for r in range(1, splits_max + 1):212 sp = splits[:r]213 J, ok = worker_corrected_estimates({s: fits[s] for s in sp}, sp)214 diag[f"valid_r{r}"] = ok215 th0 = pilot_estimator(J, block_sizes, ok=None) # keep all workers (failures tracked separately)216 path, nd = global_newton(model, db, th0, q=q_max, block_sizes=block_sizes)217 for q in range(q_max + 1):218 out[f"r{r}_q{q}"] = path[q]219 diag[f"newton_r{r}"] = nd220 if extras:221 path_s, _ = global_newton(model, db, th0, q=q_max, hessian_mode="stale", stale_hessian=H_bar, block_sizes=block_sizes)222 for q in range(1, q_max + 1):223 out[f"stale_r{r}_q{q}"] = path_s[q]224 # valid-worker-only variant: workers with any non-converged sub-block fit are dropped from the pilot (weights renormalized)225 if ok.sum() < K and ok.sum() > 0:226 th0v = pilot_estimator(J, block_sizes, ok=ok)227 pathv, _ = global_newton(model, db, th0v, q=q_max, block_sizes=block_sizes)228 else:229 pathv = path230 for q in range(q_max + 1):231 out[f"v{r}_q{q}"] = pathv[q]232 diag[f"dropped_r{r}"] = int(K - ok.sum())233 if extras:234 out["aee"] = aee_lin_xi(theta_full_local[:, 0], H_local[:, 0])235 raw = out["r1_q0"]236 cpath = csl(model, db, raw, rounds=2, block_sizes=block_sizes)237 out["csl1"], out["csl2"] = cpath[1], cpath[2]238 try:239 th_cs, b = analytic_bias_correction(model, db, theta_full_local[:, 0])240 th0 = pilot_estimator(th_cs, block_sizes)241 pth, _ = global_newton(model, db, th0, q=q_max, block_sizes=block_sizes)242 for q in range(q_max + 1):243 out[f"cs_q{q}"] = pth[q]244 except NotImplementedError:245 pass246 return out, diag247