1"""Shared estimation utilities for Phase 4 (statsmodels-based, deterministic)."""2from __future__ import annotations3import numpy as np4import pandas as pd5import statsmodels.api as sm6import statsmodels.formula.api as smf7from sklearn.metrics import roc_auc_score89SEED = 20261101101112def fit_formula(formula: str, data: pd.DataFrame, family, groups: str | None = None):13 """GLM (or OLS if family is None) with cluster-robust SE by `groups`. Returns the fitted result."""14 d = data.dropna(subset=_vars(formula, data) + ([groups] if groups else []))15 if family is None:16 mod = smf.ols(formula, d)17 else:18 mod = smf.glm(formula, d, family=family)19 if groups:20 g = pd.factorize(d[groups])[0]21 res = mod.fit(cov_type="cluster", cov_kwds={"groups": g})22 else:23 res = mod.fit(cov_type="HC1")24 res._n_used = len(d); res._data_used = d25 return res262728def _vars(formula: str, data: pd.DataFrame) -> list[str]:29 import re30 toks = set(re.findall(r"[A-Za-z_][A-Za-z0-9_]*", formula))31 return [t for t in toks if t in data.columns]323334def coef_row(res, name: str, label: str, extra: dict | None = None) -> dict:35 b = res.params.get(name, np.nan); se = res.bse.get(name, np.nan); p = res.pvalues.get(name, np.nan)36 row = {"model": label, "term": name, "coef": b, "se": se, "ci_lo": b - 1.96 * se, "ci_hi": b + 1.96 * se, "p": p,37 "n": getattr(res, "_n_used", np.nan), "llf": getattr(res, "llf", np.nan)}38 if extra:39 row.update(extra)40 return row414243def oos_loglik_binomial(y, p, trials=None):44 """Out-of-sample log-likelihood for binary or fractional outcomes (per observation mean)."""45 p = np.clip(np.asarray(p, float), 1e-6, 1 - 1e-6); y = np.asarray(y, float)46 ll = y * np.log(p) + (1 - y) * np.log(1 - p)47 return float(np.mean(ll))484950def binary_metrics(y, p) -> dict:51 y = np.asarray(y, float); p = np.clip(np.asarray(p, float), 1e-6, 1 - 1e-6)52 out = {"oos_loglik": oos_loglik_binomial(y, p), "brier": float(np.mean((y - p) ** 2)), "n": int(len(y)), "events": int(y.sum())}53 if 0 < y.sum() < len(y):54 out["auc"] = float(roc_auc_score(y, p))55 # calibration slope: logit(y) on logit(p)56 lp = np.log(p / (1 - p))57 try:58 cal = sm.Logit(y, sm.add_constant(lp)).fit(disp=0)59 out["calibration_slope"] = float(cal.params[1]); out["calibration_intercept"] = float(cal.params[0])60 except Exception:61 out["calibration_slope"] = np.nan62 else:63 out["auc"] = np.nan; out["calibration_slope"] = np.nan64 return out656667def cv_compare(data: pd.DataFrame, formulas: dict[str, str], y: str, family, fold_col: str | None = None,68 n_repeats: int = 10, k: int = 5, groups: str | None = None, fractional: bool = False) -> pd.DataFrame:69 """Out-of-sample comparison. If fold_col is given → leave-one-group-out on that column; else repeated k-fold.70 Returns one row per (model, fold) with OOS log-likelihood and, for binary y, AUC/Brier."""71 rows = []72 allvars = sorted(set(sum([_vars(f, data) for f in formulas.values()], [])) | {y} | ({groups} if groups else set()))73 d = data.dropna(subset=allvars).reset_index(drop=True)74 if fold_col:75 folds = [(f"loo_{g}", d[fold_col] != g, d[fold_col] == g) for g in sorted(d[fold_col].unique())]76 else:77 rng = np.random.default_rng(SEED); folds = []78 for r in range(n_repeats):79 perm = rng.permutation(len(d)); parts = np.array_split(perm, k)80 for i, te in enumerate(parts):81 mask = np.zeros(len(d), bool); mask[te] = True82 folds.append((f"rep{r}_fold{i}", ~mask, mask))83 for name, formula in formulas.items():84 for fid, tr, te in folds:85 dtr, dte = d[tr], d[te]86 if dte[y].notna().sum() == 0 or len(dtr) < 20:87 continue88 try:89 mod = smf.glm(formula, dtr, family=family).fit()90 # drop test rows with factor levels unseen in training (predict would fail)91 p = mod.predict(dte)92 ok = np.isfinite(p)93 m = binary_metrics(dte[y][ok], p[ok]) if not fractional else {"oos_loglik": oos_loglik_binomial(dte[y][ok], p[ok]), "n": int(ok.sum())}94 m.update({"model": name, "fold": fid, "n_train": int(len(dtr))}); rows.append(m)95 except Exception as e:96 rows.append({"model": name, "fold": fid, "error": str(e)[:80]})97 return pd.DataFrame(rows)9899100def fold_win_share(cv: pd.DataFrame, model: str, baseline: str, metric: str = "oos_loglik") -> dict:101 a = cv[cv.model == model].set_index("fold")[metric]; b = cv[cv.model == baseline].set_index("fold")[metric]102 j = a.to_frame("m").join(b.to_frame("b"), how="inner").dropna()103 if len(j) == 0:104 return {"model": model, "baseline": baseline, "folds": 0, "win_share": np.nan, "mean_diff": np.nan}105 return {"model": model, "baseline": baseline, "folds": int(len(j)), "win_share": float((j.m > j.b).mean()), "mean_diff": float((j.m - j.b).mean())}106107108def holm(pvals: dict[str, float]) -> dict[str, float]:109 items = sorted(pvals.items(), key=lambda kv: kv[1]); m = len(items); out = {}; running = 0110 for i, (k, p) in enumerate(items):111 adj = min(1.0, (m - i) * p); running = max(running, adj); out[k] = running112 return out113114115def conley_ols(formula: str, data: pd.DataFrame, lat: str, lon: str, cutoff_km: float):116 """OLS with Conley (1999) spatial HAC standard errors, Bartlett kernel with distance cutoff."""117 d = data.dropna(subset=_vars(formula, data) + [lat, lon]).reset_index(drop=True)118 mod = smf.ols(formula, d); res = mod.fit()119 X = mod.exog; u = res.resid.values; n, kx = X.shape120 la, lo = np.radians(d[lat].values), np.radians(d[lon].values)121 dlat = la[:, None] - la[None, :]; dlon = lo[:, None] - lo[None, :]122 h = np.sin(dlat / 2) ** 2 + np.cos(la)[:, None] * np.cos(la)[None, :] * np.sin(dlon / 2) ** 2123 dist = 2 * 6371.0 * np.arcsin(np.sqrt(np.clip(h, 0, 1)))124 w = np.clip(1 - dist / cutoff_km, 0, None)125 Xu = X * u[:, None]126 meat = Xu.T @ w @ Xu127 bread = np.linalg.inv(X.T @ X)128 V = bread @ meat @ bread129 se = np.sqrt(np.diag(V))130 return res, pd.Series(se, index=res.params.index), len(d)131132133def hinge_fit(data: pd.DataFrame, y: str, base_formula: str, pvar: str, family, grid=None):134 """Profile-likelihood hinge: y ~ base + P + max(P-c,0), c chosen on the given data (call on the TRAINING fold)."""135 d = data.dropna(subset=_vars(base_formula, data) + [y, pvar]).copy()136 grid = np.quantile(d[pvar], np.linspace(0.15, 0.85, 15)) if grid is None else grid137 best = None138 for c in grid:139 d["_hinge"] = np.clip(d[pvar] - c, 0, None)140 try:141 m = smf.glm(f"{base_formula} + _hinge", d, family=family).fit()142 if best is None or m.llf > best[0]:143 best = (m.llf, c, m)144 except Exception:145 continue146 return best147