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 _cat_cols(formula: str) -> list[str]:68 import re69 return re.findall(r"C\(([A-Za-z_][A-Za-z0-9_]*)\)", formula)707172def map_unseen_levels(train: pd.DataFrame, test: pd.DataFrame, formula: str) -> pd.DataFrame:73 """Test rows whose categorical fixed-effect level is absent from the training fold get the training mode (A-004);74 the alternative (dropping them) would change the test sample across models."""75 test = test.copy()76 for col in _cat_cols(formula):77 if col in test.columns:78 seen = set(train[col].dropna().unique()); mode = train[col].mode().iloc[0]79 test.loc[~test[col].isin(seen), col] = mode80 return test818283def cv_compare(data: pd.DataFrame, formulas: dict[str, str], y: str, family, fold_col: str | None = None,84 n_repeats: int = 10, k: int = 5, groups: str | None = None, fractional: bool = False, block_col: str | None = None) -> pd.DataFrame:85 """Out-of-sample comparison. If fold_col is given → leave-one-group-out on that column; else repeated k-fold.86 Returns one row per (model, fold) with OOS log-likelihood and, for binary y, AUC/Brier."""87 rows = []88 allvars = sorted(set(sum([_vars(f, data) for f in formulas.values()], [])) | {y} | ({groups} if groups else set()))89 d = data.dropna(subset=allvars).reset_index(drop=True)90 if fold_col:91 folds = [(f"loo_{g}", d[fold_col] != g, d[fold_col] == g) for g in sorted(d[fold_col].unique())]92 else:93 rng = np.random.default_rng(SEED); folds = []94 units = d[block_col].values if block_col else np.arange(len(d)) # polity-blocked folds when block_col is given95 uniq = np.unique(units)96 for r in range(n_repeats):97 perm = rng.permutation(len(uniq)); parts = np.array_split(perm, k)98 for i, te in enumerate(parts):99 mask = np.isin(units, uniq[te])100 folds.append((f"rep{r}_fold{i}", ~mask, mask))101 for name, formula in formulas.items():102 for fid, tr, te in folds:103 dtr, dte = d[tr], d[te]104 if dte[y].notna().sum() == 0 or len(dtr) < 20:105 continue106 try:107 mod = smf.glm(formula, dtr, family=family).fit()108 p = mod.predict(map_unseen_levels(dtr, dte, formula))109 ok = np.isfinite(p)110 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())}111 m.update({"model": name, "fold": fid, "n_train": int(len(dtr))}); rows.append(m)112 except Exception as e:113 rows.append({"model": name, "fold": fid, "error": str(e)[:80]})114 return pd.DataFrame(rows)115116117def fold_win_share(cv: pd.DataFrame, model: str, baseline: str, metric: str = "oos_loglik") -> dict:118 a = cv[cv.model == model].set_index("fold")[metric]; b = cv[cv.model == baseline].set_index("fold")[metric]119 j = a.to_frame("m").join(b.to_frame("b"), how="inner").dropna()120 if len(j) == 0:121 return {"model": model, "baseline": baseline, "folds": 0, "win_share": np.nan, "mean_diff": np.nan}122 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())}123124125def holm(pvals: dict[str, float]) -> dict[str, float]:126 items = sorted(pvals.items(), key=lambda kv: kv[1]); m = len(items); out = {}; running = 0127 for i, (k, p) in enumerate(items):128 adj = min(1.0, (m - i) * p); running = max(running, adj); out[k] = running129 return out130131132def conley_ols(formula: str, data: pd.DataFrame, lat: str, lon: str, cutoff_km: float):133 """OLS with Conley (1999) spatial HAC standard errors, Bartlett kernel with distance cutoff."""134 d = data.dropna(subset=_vars(formula, data) + [lat, lon]).reset_index(drop=True)135 mod = smf.ols(formula, d); res = mod.fit()136 X = mod.exog; u = res.resid.values; n, kx = X.shape137 la, lo = np.radians(d[lat].values), np.radians(d[lon].values)138 dlat = la[:, None] - la[None, :]; dlon = lo[:, None] - lo[None, :]139 h = np.sin(dlat / 2) ** 2 + np.cos(la)[:, None] * np.cos(la)[None, :] * np.sin(dlon / 2) ** 2140 dist = 2 * 6371.0 * np.arcsin(np.sqrt(np.clip(h, 0, 1)))141 w = np.clip(1 - dist / cutoff_km, 0, None)142 Xu = X * u[:, None]143 meat = Xu.T @ w @ Xu144 bread = np.linalg.inv(X.T @ X)145 V = bread @ meat @ bread146 se = np.sqrt(np.diag(V))147 return res, pd.Series(se, index=res.params.index), len(d)148149150def hinge_fit(data: pd.DataFrame, y: str, base_formula: str, pvar: str, family, grid=None):151 """Profile-likelihood hinge: y ~ base + P + max(P-c,0), c chosen on the given data (call on the TRAINING fold)."""152 d = data.dropna(subset=_vars(base_formula, data) + [y, pvar]).copy()153 grid = np.quantile(d[pvar], np.linspace(0.15, 0.85, 15)) if grid is None else grid154 best = None155 for c in grid:156 d["_hinge"] = np.clip(d[pvar] - c, 0, None)157 try:158 m = smf.glm(f"{base_formula} + _hinge", d, family=family).fit()159 if best is None or m.llf > best[0]:160 best = (m.llf, c, m)161 except Exception:162 continue163 return best164