| """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 |
|
|
|
|
| |
| |
| |
| @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) |
|
|
|
|
| |
| |
| |
| 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, |
| ) |
|
|
|
|
| |
| |
| |
| 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) |
|
|
| |
| 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) |
| |
| 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) |
|
|
|
|
| |
| |
| |
| 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 |
| 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 |
|
|
|
|
| |
| |
| |
| 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)} |
|
|
|
|
| |
| |
| |
| 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 "" |
|
|
