"""Model fitting helpers used across initiative pages. All routines return small, consistent dataclasses so the page-side code can render results uniformly. """ from __future__ import annotations from dataclasses import dataclass, field from typing import Iterable, Optional import numpy as np import pandas as pd import statsmodels.api as sm import statsmodels.formula.api as smf # -------------------------------------------------------------------------- # Result containers # -------------------------------------------------------------------------- @dataclass class CoefRow: term: str coef: float se: float t: float p: float ci_low: float ci_high: float def as_dict(self) -> dict: return self.__dict__ @dataclass class FitResult: """Lightweight wrapper around a fitted statsmodels regression.""" coefs: pd.DataFrame n_obs: int n_clusters: Optional[int] r2: Optional[float] formula: str raw: object = field(repr=False, default=None) @classmethod def from_results(cls, res, formula: str, *, raw_terms: Iterable[str] | None = None, n_clusters: Optional[int] = None) -> "FitResult": ci = res.conf_int() rows = [] for term in res.params.index: if raw_terms is not None and term not in raw_terms: continue rows.append(CoefRow( term=term, coef=float(res.params[term]), se=float(res.bse[term]), t=float(res.tvalues[term]), p=float(res.pvalues[term]), ci_low=float(ci.loc[term, 0]), ci_high=float(ci.loc[term, 1]), ).as_dict()) coefs = pd.DataFrame(rows) try: r2 = float(res.rsquared) except AttributeError: r2 = None return cls(coefs=coefs, n_obs=int(res.nobs), n_clusters=n_clusters, r2=r2, formula=formula, raw=res) # -------------------------------------------------------------------------- # Two-way fixed effects DiD / TWFE # -------------------------------------------------------------------------- def fit_twfe( df: pd.DataFrame, *, outcome: str, unit: str = "fips", period: str = "year", treatment_terms: list[str], controls: Optional[list[str]] = None, cluster_col: Optional[str] = None, ) -> FitResult: """Two-way fixed effects regression. Specification: y_{it} = alpha_i + lambda_t + sum_k beta_k * Treat_k_{it} + gamma X_{it} + e Implemented via OLS with explicit dummies (small panels, fast enough). """ controls = controls or [] needed = [outcome, unit, period] + treatment_terms + controls needed = [c for c in needed if c is not None] work = df.dropna(subset=needed).copy() formula_parts = [outcome, "~"] formula_parts.append(" + ".join(treatment_terms + controls)) formula_parts.append(f" + C({unit}) + C({period})") formula = " ".join(formula_parts) model = smf.ols(formula, data=work) if cluster_col is not None: groups = work[cluster_col] res = model.fit(cov_type="cluster", cov_kwds={"groups": groups}) n_clusters = int(work[cluster_col].nunique()) else: res = model.fit() n_clusters = None return FitResult.from_results( res, formula, raw_terms=treatment_terms + controls, n_clusters=n_clusters, ) # -------------------------------------------------------------------------- # Event study around a treatment time # -------------------------------------------------------------------------- def build_event_time( df: pd.DataFrame, *, unit: str, period: str, treat_unit_col: str, event_period_col: str, leads: int = 4, lags: int = 6, reference_lead: int = -1, ) -> pd.DataFrame: """Construct event-time indicators for a staggered-adoption event study. Returns the original df augmented with columns ev_m4, ..., ev_p6 plus a binned ev_minus / ev_plus for periods outside the window. Reference period (reference_lead) is omitted so coefficients are interpreted relative to that period. """ out = df.copy() et = out[period] - out[event_period_col] et = et.where(out[treat_unit_col] == 1) # NaN for never-treated # Bin tails et_binned = et.copy() et_binned = et_binned.where(et_binned >= -leads, other=-(leads + 99)) et_binned = et_binned.where(et_binned <= lags, other=(lags + 99)) cols = [] for k in range(-leads, lags + 1): if k == reference_lead: continue col = f"ev_{('m' if k < 0 else 'p')}{abs(k)}" out[col] = ((et_binned == k) & (out[treat_unit_col] == 1)).astype(int) cols.append(col) # binned tails out["ev_pre"] = ((et_binned == -(leads + 99)) & (out[treat_unit_col] == 1)).astype(int) out["ev_post"] = ((et_binned == (lags + 99)) & (out[treat_unit_col] == 1)).astype(int) return out, cols + ["ev_pre", "ev_post"] def event_study_table( fit: FitResult, *, leads: int = 4, lags: int = 6, reference_lead: int = -1, ) -> pd.DataFrame: """Reshape an event-study fit into a long table for plotting.""" rows = [] rows.append({ "event_time": reference_lead, "coef": 0.0, "se": 0.0, "ci_low": 0.0, "ci_high": 0.0, "p": 1.0, }) for k in range(-leads, lags + 1): if k == reference_lead: continue col = f"ev_{('m' if k < 0 else 'p')}{abs(k)}" m = fit.coefs[fit.coefs["term"] == col] if m.empty: continue m = m.iloc[0] rows.append({ "event_time": k, "coef": m["coef"], "se": m["se"], "ci_low": m["ci_low"], "ci_high": m["ci_high"], "p": m["p"], }) return pd.DataFrame(rows).sort_values("event_time").reset_index(drop=True) # -------------------------------------------------------------------------- # Logistic regression with marginal effects # -------------------------------------------------------------------------- def fit_logit_marginal( df: pd.DataFrame, *, outcome: str, predictors: list[str], ) -> tuple[FitResult, pd.DataFrame]: """Logit + average marginal effects (AME).""" work = df.dropna(subset=[outcome] + predictors).copy() formula = f"{outcome} ~ {' + '.join(predictors)}" res = smf.logit(formula, data=work).fit(disp=False) fit = FitResult.from_results(res, formula, raw_terms=predictors) me = res.get_margeff(at="overall", method="dydx") me_df = pd.DataFrame({ "term": me.results_table_data[0][1:], "marginal_effect": me.margeff, "se": me.margeff_se, "p": me.pvalues, }) if False else None # The API differs across versions; build manually: margeff = res.get_margeff() me_summary = margeff.summary_frame() me_summary = me_summary.reset_index().rename(columns={ "index": "term", "dy/dx": "marginal_effect", "Std. Err.": "se", "Pr(>|z|)": "p", }) return fit, me_summary # -------------------------------------------------------------------------- # Parallel-trends pre-period test # -------------------------------------------------------------------------- def parallel_trends_pvalue( df: pd.DataFrame, *, outcome: str, unit: str, period: str, treat_unit_col: str, pre_periods: list[int], ) -> dict: """Joint F-test on pre-period treated × period interactions.""" pre = df[df[period].isin(pre_periods)].copy() pre["t"] = pre[period] - min(pre_periods) pre["interact"] = pre[treat_unit_col] * pre["t"] formula = f"{outcome} ~ {treat_unit_col} + C({period}) + interact + C({unit})" res = smf.ols(formula, data=pre).fit() coef = float(res.params.get("interact", np.nan)) se = float(res.bse.get("interact", np.nan)) p = float(res.pvalues.get("interact", np.nan)) return {"coef": coef, "se": se, "p": p, "n": int(res.nobs)} # -------------------------------------------------------------------------- # Convenience: pretty p-value # -------------------------------------------------------------------------- def stars(p: float) -> str: if pd.isna(p): return "" if p < 0.001: return "***" if p < 0.01: return "**" if p < 0.05: return "*" if p < 0.1: return "·" return ""