Add Sun-Abraham event study, Bacon decomposition, and updated DID module

- Sun & Abraham (2021) event study estimator
- Goodman-Bacon (2021) decomposition for TWFE DID
- Updated __init__.py exports

313 tests passing, 81% coverage.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: brycewang-stanford

src/statspai/__init__.py

@@ -30,7 +30,7 @@
 from .regression.ols import regress
 from .regression.iv import ivreg, IVRegression
 from .causal.causal_forest import CausalForest, causal_forest
-from .did import did, did_2x2, callaway_santanna
+from .did import did, did_2x2, callaway_santanna, sun_abraham, bacon_decomposition
 from .rd import rdrobust, rdplot
 from .synth import synth, SyntheticControl
 from .matching import match, MatchEstimator
@@ -59,6 +59,8 @@
     "did",
     "did_2x2",
     "callaway_santanna",
+    "sun_abraham",
+    "bacon_decomposition",
     # RD
     "rdrobust",
     "rdplot",

src/statspai/did/__init__.py

@@ -2,13 +2,10 @@
 Difference-in-Differences (DID) module for StatsPAI.
 
 Provides estimators for:
-- Classic 2x2 DID (two groups, two periods)
-- Staggered DID with heterogeneous treatment effects (Callaway & Sant'Anna, 2021)
-
-Planned:
-- Sun & Abraham (2021) interaction-weighted event study
-- Goodman-Bacon (2021) decomposition
-- Doubly Robust DID (Sant'Anna & Zhao, 2020) — standalone
+- Classic 2×2 DID (two groups, two periods)
+- Callaway & Sant'Anna (2021) — staggered DID with DR/IPW/REG
+- Sun & Abraham (2021) — interaction-weighted event study
+- Goodman-Bacon (2021) — TWFE decomposition diagnostic
 """
 
 from typing import Optional, List
@@ -18,6 +15,8 @@
 from ..core.results import CausalResult
 from .did_2x2 import did_2x2
 from .callaway_santanna import callaway_santanna
+from .sun_abraham import sun_abraham
+from .bacon import bacon_decomposition
 
 
 def did(
@@ -128,14 +127,28 @@ def did(
             base_period=base_period, alpha=alpha,
         )
 
+    if method in ('sun_abraham', 'sa', 'sunab'):
+        if id is None:
+            raise ValueError(
+                "'id' (unit identifier) is required for Sun-Abraham."
+            )
+        return sun_abraham(
+            data, y=y, g=treat, t=time, i=id,
+            covariates=covariates, cluster=cluster,
+            alpha=alpha,
+        )
+
     raise ValueError(
         f"Unknown DID method: '{method}'. "
-        "Available: '2x2', 'callaway_santanna' (or 'cs')."
+        "Available: '2x2', 'callaway_santanna' (or 'cs'), "
+        "'sun_abraham' (or 'sa')."
     )
 
 
 __all__ = [
     'did',
     'did_2x2',
     'callaway_santanna',
+    'sun_abraham',
+    'bacon_decomposition',
 ]

src/statspai/did/bacon.py

@@ -0,0 +1,274 @@
+"""
+Goodman-Bacon (2021) decomposition for two-way fixed effects DID.
+
+Decomposes the standard TWFE DID estimator into a weighted average of
+all possible 2×2 DID comparisons, revealing which comparisons drive
+the estimate and whether "forbidden" comparisons (using already-treated
+units as controls) contribute negative weights.
+
+This is a **diagnostic** tool — it does not provide a bias-corrected
+estimator. Use Callaway-Sant'Anna or Sun-Abraham for estimation.
+
+References
+----------
+Goodman-Bacon, A. (2021).
+"Difference-in-Differences with Variation in Treatment Timing."
+*Journal of Econometrics*, 225(2), 254-277.
+
+Goodman-Bacon, A., Goldring, T. and Nichols, A. (2019).
+"BACONDECOMP: Stata module to perform the Bacon decomposition
+of difference-in-differences estimation."
+"""
+
+from typing import Optional, List, Dict, Any, Tuple
+
+import numpy as np
+import pandas as pd
+
+from ..core.results import CausalResult
+
+
+def bacon_decomposition(
+    data: pd.DataFrame,
+    y: str,
+    treat: str,
+    time: str,
+    id: str,
+    alpha: float = 0.05,
+) -> Dict[str, Any]:
+    """
+    Goodman-Bacon (2021) decomposition of the TWFE DID estimator.
+
+    Decomposes the overall TWFE coefficient into a weighted sum of
+    2×2 DID comparisons between different treatment timing groups.
+
+    Parameters
+    ----------
+    data : pd.DataFrame
+        Balanced panel data.
+    y : str
+        Outcome variable.
+    treat : str
+        Binary treatment indicator (0 before treatment, 1 after).
+    time : str
+        Time period variable.
+    id : str
+        Unit identifier.
+    alpha : float, default 0.05
+        Significance level.
+
+    Returns
+    -------
+    dict
+        Keys:
+        - ``beta_twfe``: overall TWFE estimate
+        - ``decomposition``: pd.DataFrame with columns
+          [type, treated, control, estimate, weight]
+        - ``weighted_sum``: Σ(weight × estimate) — should equal beta_twfe
+        - ``n_comparisons``: number of 2×2 sub-comparisons
+        - ``negative_weight_share``: fraction of weight on comparisons
+          where already-treated units serve as controls (the "forbidden"
+          comparisons that can bias TWFE)
+
+    Examples
+    --------
+    >>> result = bacon_decomposition(df, y='outcome', treat='treated',
+    ...                              time='year', id='unit')
+    >>> print(result['decomposition'])
+    >>> print(f"TWFE = {result['beta_twfe']:.4f}")
+    >>> print(f"Negative weight share = {result['negative_weight_share']:.1%}")
+
+    Notes
+    -----
+    The decomposition identifies three types of comparisons:
+
+    1. **Earlier vs Later treated**: Units treated at time g₁ vs units
+       treated later at g₂ (g₁ < g₂). These are "good" comparisons.
+    2. **Later vs Earlier treated**: Units treated at g₂ vs already-treated
+       units at g₁. These are "forbidden" — they use treated units as
+       controls and can introduce negative weighting bias.
+    3. **Treated vs Never treated**: Always valid comparisons.
+
+    A large ``negative_weight_share`` signals that TWFE is unreliable
+    and a heterogeneity-robust estimator (C&S, Sun-Abraham) should be used.
+
+    See Goodman-Bacon (2021, *JEcon*), Theorem 1.
+    """
+    df = data.copy()
+
+    # Validate balanced panel
+    panel = df.pivot_table(index=id, columns=time, values=y, aggfunc='first')
+    n_units = len(panel)
+    time_periods = sorted(df[time].unique())
+    T = len(time_periods)
+
+    # Get treatment timing for each unit
+    # g_i = first period where treat == 1 (inf for never-treated)
+    unit_treat = df.groupby(id)[[treat, time]].apply(
+        lambda grp: grp.loc[grp[treat] == 1, time].min()
+        if (grp[treat] == 1).any() else np.inf
+    )
+
+    # Identify groups by treatment timing
+    timing_groups = sorted(unit_treat.unique())
+    never_treated = np.inf
+
+    # Overall TWFE estimate (for reference)
+    beta_twfe = _twfe_estimate(df, y, treat, time, id)
+
+    # Enumerate all 2×2 comparisons
+    comparisons = []
+
+    for i, g1 in enumerate(timing_groups):
+        for g2 in timing_groups[i + 1:]:
+            if g1 == never_treated:
+                continue
+
+            units_g1 = unit_treat[unit_treat == g1].index
+            units_g2 = unit_treat[unit_treat == g2].index
+
+            n1 = len(units_g1)
+            n2 = len(units_g2)
+
+            if n1 == 0 or n2 == 0:
+                continue
+
+            if g2 == never_treated:
+                # Type: Treated vs Never-treated
+                comp_type = 'Treated vs Never-treated'
+                est, wt = _pairwise_did(
+                    panel, units_g1, units_g2, g1, time_periods,
+                    n1, n2, n_units, T,
+                )
+                comparisons.append({
+                    'type': comp_type,
+                    'treated': g1,
+                    'control': 'Never',
+                    'estimate': est,
+                    'weight': wt,
+                })
+            else:
+                # Type 1: Earlier (g1) vs Later (g2) — "good"
+                est1, wt1 = _pairwise_did(
+                    panel, units_g1, units_g2, g1, time_periods,
+                    n1, n2, n_units, T,
+                )
+                comparisons.append({
+                    'type': 'Earlier vs Later treated',
+                    'treated': g1,
+                    'control': g2,
+                    'estimate': est1,
+                    'weight': wt1,
+                })
+
+                # Type 2: Later (g2) vs Earlier (g1) — "forbidden"
+                est2, wt2 = _pairwise_did(
+                    panel, units_g2, units_g1, g2, time_periods,
+                    n2, n1, n_units, T,
+                )
+                comparisons.append({
+                    'type': 'Later vs Already-treated',
+                    'treated': g2,
+                    'control': g1,
+                    'estimate': est2,
+                    'weight': wt2,
+                })
+
+    decomp = pd.DataFrame(comparisons)
+
+    if len(decomp) == 0:
+        return {
+            'beta_twfe': beta_twfe,
+            'decomposition': decomp,
+            'weighted_sum': 0.0,
+            'n_comparisons': 0,
+            'negative_weight_share': 0.0,
+        }
+
+    # Normalize weights to sum to 1
+    total_w = decomp['weight'].sum()
+    if total_w > 0:
+        decomp['weight'] = decomp['weight'] / total_w
+
+    weighted_sum = float((decomp['estimate'] * decomp['weight']).sum())
+
+    # Negative weight share (from "forbidden" comparisons)
+    forbidden = decomp['type'] == 'Later vs Already-treated'
+    neg_share = float(decomp.loc[forbidden, 'weight'].sum())
+
+    return {
+        'beta_twfe': beta_twfe,
+        'decomposition': decomp,
+        'weighted_sum': weighted_sum,
+        'n_comparisons': len(decomp),
+        'negative_weight_share': neg_share,
+    }
+
+
+def _twfe_estimate(df, y, treat, time, id_col):
+    """Standard TWFE DID regression: Y_it = α_i + γ_t + β·D_it + ε_it."""
+    # Demean by unit and time (within transformation)
+    panel = df.set_index([id_col, time])
+    Y = panel[y].unstack()
+    D = panel[treat].unstack()
+
+    # Double-demean
+    y_dm = Y.values - Y.values.mean(axis=1, keepdims=True) \
+           - Y.values.mean(axis=0, keepdims=True) + Y.values.mean()
+    d_dm = D.values - D.values.mean(axis=1, keepdims=True) \
+           - D.values.mean(axis=0, keepdims=True) + D.values.mean()
+
+    y_flat = y_dm.ravel()
+    d_flat = d_dm.ravel()
+    valid = np.isfinite(y_flat) & np.isfinite(d_flat)
+
+    if valid.sum() < 2 or np.var(d_flat[valid]) < 1e-12:
+        return 0.0
+
+    beta = float(np.sum(d_flat[valid] * y_flat[valid]) /
+                 np.sum(d_flat[valid] ** 2))
+    return beta
+
+
+def _pairwise_did(panel, units_treated, units_control, treat_time,
+                  time_periods, n_t, n_c, n_total, T):
+    """Compute 2×2 DID and Bacon weight for a pair of groups."""
+    # Pre/post periods relative to treat_time
+    pre_periods = [t for t in time_periods if t < treat_time]
+    post_periods = [t for t in time_periods if t >= treat_time]
+
+    if not pre_periods or not post_periods:
+        return 0.0, 0.0
+
+    # Outcome means
+    y_t_pre = panel.loc[units_treated, pre_periods].values.mean()
+    y_t_post = panel.loc[units_treated, post_periods].values.mean()
+    y_c_pre = panel.loc[units_control, pre_periods].values.mean()
+    y_c_post = panel.loc[units_control, post_periods].values.mean()
+
+    did_est = (y_t_post - y_t_pre) - (y_c_post - y_c_pre)
+
+    # Bacon weight ∝ n_k × n_l × V(D̃_kl)
+    n_k = n_t
+    n_l = n_c
+    s = len(post_periods) / T  # share of post-treatment periods
+    var_d = s * (1 - s)  # variance of demeaned treatment in this 2×2
+    weight = (n_k + n_l) * var_d
+
+    return float(did_est), float(weight)
+
+
+# Citation
+CausalResult._CITATIONS['bacon_decomposition'] = (
+    "@article{goodman2021difference,\n"
+    "  title={Difference-in-Differences with Variation in Treatment "
+    "Timing},\n"
+    "  author={Goodman-Bacon, Andrew},\n"
+    "  journal={Journal of Econometrics},\n"
+    "  volume={225},\n"
+    "  number={2},\n"
+    "  pages={254--277},\n"
+    "  year={2021},\n"
+    "  publisher={Elsevier}\n"
+    "}"
+)