fix(did): gardner_did event-study reference-category contamination

The Stage-2 dummy regression pooled all non-event rows (never-treated +
treated units outside the event-study horizon) into a single baseline
category, dragging every event-time coefficient toward the mean of that
pool.  On a synthetic panel with true τ=2 and strict parallel trends,
this produced:

  before: pre-trends ≈ -0.30 (should be 0), post ≈ +1.72 (should be 2.0)
  after:  pre-trends ≈ +0.01 (≈0),           post ≈ +2.02 (≈2.0)

The fix replaces the Stage-2 dummy regression in event-study mode with a
direct Borusyak-Jaravel-Spiess-style within-(cohort × relative-time)
averaging of the imputed gap ỹ = Y - Ŷ(0).  The non-event-study path
(single ATT) already used the correct regression and is unchanged.

Cluster-robust SEs in event-study mode now use the standard cluster
sum-of-residuals formula on each bin separately.

Also tightens test_gardner_event_study_support to require np.mean(post)
> 1.6 and |pre_trend| < 0.3 (previously 1.0 / 0.8).

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

Author: brycewang-stanford

src/statspai/did/gardner_2s.py

@@ -220,7 +220,16 @@ def gardner_did(
     y_hat_0 = A_full @ coefs
     y_tilde = y_all_arr - y_hat_0
 
-    # ── Stage 2: regress residualised outcome on treatment ──────── #
+    # ── Stage 2: recover treatment effects from the imputed gap ──── #
+    # Overall ATT: clustered OLS of ỹ on the treatment indicator, with
+    # an intercept to absorb any mean residual in the untreated rows.
+    # Event study: direct within-(cohort × relative-time) averaging of ỹ
+    # — the Borusyak-Jaravel-Spiess style — to avoid the reference-
+    # category contamination bias that a Stage-2 dummy regression would
+    # introduce (the "baseline" in a dummy regression lumps never-treated
+    # units together with treated units outside the event-study horizon,
+    # pulling every coefficient toward the residual mean).
+    cl = df[cluster].to_numpy()
     if event_study:
         rel_time = np.where(
             np.isfinite(ft) & (ft > 0),
@@ -234,28 +243,51 @@ def gardner_did(
                 horizon.append(0)
             horizon = sorted(set(horizon))
 
-        reg_cols = {}
+        names, est_list, se_list = [], [], []
         for k in horizon:
-            reg_cols[f"D_k{int(k):+d}"] = (rel_time == k).astype(float)
-        X2 = np.column_stack(list(reg_cols.values()))
-        names = list(reg_cols.keys())
+            key = f"D_k{int(k):+d}"
+            names.append(key)
+            mask = rel_time == k
+            if mask.sum() == 0:
+                est_list.append(float("nan"))
+                se_list.append(float("nan"))
+                continue
+            y_k = y_tilde[mask]
+            coef_k = float(np.mean(y_k))
+            # Cluster-robust SE of the within-bin mean
+            cl_k = cl[mask]
+            uniq = np.unique(cl_k)
+            G = len(uniq)
+            if G > 1:
+                group_means = np.array([y_k[cl_k == g].mean() for g in uniq])
+                # SE of the unweighted mean over n rows, allowing cluster
+                # correlation: Var(mean) ≈ (1/n²) Σ_g (Σ_{i∈g} (y_ki - coef))²
+                sq = 0.0
+                for g in uniq:
+                    idx = cl_k == g
+                    sq += float(np.sum(y_k[idx] - coef_k)) ** 2
+                var_k = sq / (len(y_k) ** 2)
+                se_k = float(np.sqrt(max(var_k, 0.0)))
+            else:
+                se_k = float(np.std(y_k, ddof=1) / np.sqrt(len(y_k)))
+            est_list.append(coef_k)
+            se_list.append(se_k)
+        est = np.array(est_list)
+        se = np.array(se_list)
+        coef_dict = dict(zip(names, est))
+        se_dict = dict(zip(names, se))
     else:
         X2 = df["_D"].to_numpy(dtype=float).reshape(-1, 1)
         names = ["ATT"]
-
-    # Stage-2 includes intercept (absorbs any post-residualisation mean shift)
-    design2 = np.column_stack([np.ones(len(y_tilde)), X2])
-    cl = df[cluster].to_numpy()
-
-    coef2, *_ = np.linalg.lstsq(design2, y_tilde, rcond=None)
-    resid2 = y_tilde - design2 @ coef2
-    V = _cluster_vcov(design2, resid2, cl)
-    se_full = np.sqrt(np.clip(np.diag(V), 0, None))
-    # Slice off intercept
-    est = coef2[1:]
-    se = se_full[1:]
-    coef_dict = dict(zip(names, est))
-    se_dict = dict(zip(names, se))
+        design2 = np.column_stack([np.ones(len(y_tilde)), X2])
+        coef2, *_ = np.linalg.lstsq(design2, y_tilde, rcond=None)
+        resid2 = y_tilde - design2 @ coef2
+        V = _cluster_vcov(design2, resid2, cl)
+        se_full = np.sqrt(np.clip(np.diag(V), 0, None))
+        est = coef2[1:]
+        se = se_full[1:]
+        coef_dict = dict(zip(names, est))
+        se_dict = dict(zip(names, se))
 
     z = sp_stats.norm.ppf(1 - alpha / 2)
     ci = {

tests/test_gardner_2s.py

@@ -53,11 +53,11 @@ def test_gardner_event_study_support():
     assert set(es["horizon"]) == {"D_k-2", "D_k-1", "D_k+0", "D_k+1", "D_k+2"}
     # Post-treatment coefs should be ≈ 2.0, pre-treatment ≈ 0
     post = [es["coef"][k] for k in ("D_k+0", "D_k+1", "D_k+2")]
-    # Event-study coefs may shift because Stage-1 fit depends on which cohorts
-    # are still untreated at each horizon.  Demand only the right sign + scale.
-    assert np.mean(post) > 1.0
+    # After the v1.5.1 reference-category fix, pre-trend should be ~0 and
+    # post-treatment should closely track the truth.
+    assert np.mean(post) > 1.6
     for k in ("D_k-2", "D_k-1"):
-        assert abs(es["coef"][k]) < 0.8  # pre-trend small relative to ATT
+        assert abs(es["coef"][k]) < 0.3  # pre-trend tight relative to ATT
 
 
 def test_gardner_alias_did_2stage():