fix(did): close remaining review follow-ups — H3/H4/H6/H7

Finishes the polish round so the DID module can be declared release-ready.

H3. _etwfe_never_only no longer mutates the caller's DataFrame.
    Removed the _ft_cache helper column; the cohort membership is now
    computed into a local Series and consumed inline. Verified that
    df.columns is unchanged after the call. Docstring also notes the
    per-cohort cluster SE adjustment so users aren't surprised when
    the aggregated SE differs from a full-sample fit.

H4. Slope indexing in _etwfe_with_xvar is now name-keyed.
    Builds an explicit coef_index dict ({column_name: position}) at
    design-matrix construction time; cohort ATT and per-xvar slopes
    are extracted by dict lookup, never by arithmetic offset.
    base_vcov likewise uses np.ix_ over the cohort-name keys. This
    makes future column-order refactors safe. A regression test
    verifies that reordering xvar=['x1','x2'] vs ['x2','x1'] yields
    identical per-cohort slopes (and that slope_x1 != slope_x2 on
    heterogeneous DGPs).

H6. Repeated-cross-section design warns on rank deficiency.
    Before the OLS fit, _etwfe_repeated_cs computes the design-matrix
    rank and, if deficient, emits a RuntimeWarning naming the deficit
    and suggesting concrete remedies. The pinv fallback in _ols_fit
    still runs so the regression doesn't crash, but the user now has
    a chance to notice a silently pathological specification.

H7. etwfe_emfx gains include_leads for full event-study output.
    wooldridge_did now builds event-time dummies for all rel_time
    values except rel = -1 (the reference category), so pre-trend
    leads are estimated alongside post-treatment lags. etwfe_emfx
    gains an include_leads=False default (backward-compat); set
    True to return the full event study including negative event
    times. SE uses the same vcov-based delta method (H1 fix).

Tests: 4 new regression cases (tests/test_did_summary.py), total 27
passing in 1.85s (was 23).

Combined with the prior review fixes (C1-C5, H1/H2/H5), all 8 issues
flagged by the code-review agent are now closed. The remaining M/L
items are genuine polish (markdown cosmetics, dead-code pruning, API
ergonomics) that don't affect correctness.

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

Author: brycewang-stanford

src/statspai/did/wooldridge_did.py

@@ -183,15 +183,19 @@ def wooldridge_did(
         interaction_cols.append(col)
 
     # ── Also build cohort × relative-time dummies for event study ───
+    # H7 fix: build event dummies for BOTH leads and lags. Omit
+    # rel = -1 as the reference category so the design is identified.
+    # Post-only consumers can still filter via etwfe_emfx(include_leads=False).
     event_cols: List[str] = []
     rel_times = set()
     for g in cohorts:
         mask_g = df["_ft"] == g
         for t_val in periods:
             rel = int(t_val - g)
-            if rel < 0:
-                continue  # only post-treatment for basic spec
-            col = f"_coh{int(g)}_rel{rel}"
+            if rel == -1:
+                continue  # reference / omitted period
+            col = (f"_coh{int(g)}_rel{rel}" if rel >= 0
+                   else f"_coh{int(g)}_rel_neg{abs(rel)}")
             df[col] = ((mask_g) & (df[time] == t_val)).astype(float)
             event_cols.append(col)
             rel_times.add(rel)
@@ -308,9 +312,15 @@ def wooldridge_did(
 
             ev_rows = []
             for j, col in enumerate(event_cols):
-                parts = col.replace("_coh", "").replace("_rel", " ").split()
-                coh_val = int(parts[0])
-                rel_val = int(parts[1])
+                # Parse name: _coh{g}_rel{k}  or  _coh{g}_rel_neg{k}
+                stripped = col.replace("_coh", "")
+                if "_rel_neg" in stripped:
+                    coh_part, rel_part = stripped.split("_rel_neg")
+                    rel_val = -int(rel_part)
+                else:
+                    coh_part, rel_part = stripped.split("_rel")
+                    rel_val = int(rel_part)
+                coh_val = int(coh_part)
                 idx_j = j + 1
                 ev_rows.append({
                     "cohort": coh_val,
@@ -564,6 +574,13 @@ def _etwfe_with_xvar(
     all_slope = [c for v in slope_cols_by_cohort.values() for c in v]
     all_inter = base_cols + all_slope
 
+    # H4 fix: explicit name-to-index map so slope lookups never rely on
+    # implicit ordering. Column 0 is the constant; columns [1:] are
+    # X_cols = [f"{c}_dm" for c in all_inter] + ctrl_dm_cols in order.
+    coef_index: Dict[str, int] = {"_const": 0}
+    for i, col in enumerate(all_inter, start=1):
+        coef_index[col] = i
+
     # Demean interactions via FE projection
     for col in all_inter:
         u_m = df.groupby(group)[col].transform("mean")
@@ -597,11 +614,14 @@ def _etwfe_with_xvar(
     df_resid = max(n_obs - X.shape[1], 1)
 
     k = len(cohorts)
-    p_x = len(xvar)
+    p_x = len(xvar)  # retained for the single-xvar backward-compat block
 
     cohort_results = []
-    for i, g in enumerate(cohorts):
-        b_idx = 1 + i
+    for g in cohorts:
+        # H4 fix: look up baseline and slope indices by coefficient name
+        # rather than by arithmetic position, so future column-order
+        # changes cannot silently mis-attribute.
+        b_idx = coef_index[f"_coh{int(g)}_post"]
         att = float(beta[b_idx])
         att_se = float(se[b_idx])
         p = float(2 * (1 - stats.t.cdf(abs(att / att_se) if att_se > 0 else 0,
@@ -610,9 +630,8 @@ def _etwfe_with_xvar(
             "cohort": int(g),
             "att_at_xmean": att, "att_se": att_se, "att_pvalue": p,
         }
-        # Per-xvar slopes
-        for j, x in enumerate(xvar):
-            s_idx = 1 + k + i * p_x + j
+        for x in xvar:
+            s_idx = coef_index[f"_coh{int(g)}_post_x_{x}"]
             slope = float(beta[s_idx])
             slope_se = float(se[s_idx])
             p_s = float(2 * (1 - stats.t.cdf(
@@ -633,7 +652,9 @@ def _etwfe_with_xvar(
     sizes = detail["n_obs"].values.astype(float)
     weights_vec = sizes / sizes.sum() if sizes.sum() > 0 else np.ones(k) / k
     att_overall = float(weights_vec @ detail["att_at_xmean"].values)
-    base_vcov = vcov[1:1 + k, 1:1 + k]
+    # H4 fix: extract baseline-coefficient vcov by explicit index lookup
+    base_idx = np.array([coef_index[f"_coh{int(g)}_post"] for g in cohorts])
+    base_vcov = vcov[np.ix_(base_idx, base_idx)]
     att_se_overall = float(np.sqrt(weights_vec @ base_vcov @ weights_vec))
     t_stat = att_overall / att_se_overall if att_se_overall > 0 else np.nan
     p_overall = float(2 * (1 - stats.t.cdf(abs(t_stat), df_resid)))
@@ -751,6 +772,23 @@ def _etwfe_repeated_cs(
     y_vec = dfv["_y"].values
     X = np.column_stack([np.ones(len(y_vec))] + [dfv[c].values for c in design_cols])
     cl_arr = dfv[cluster].values if cluster else None
+
+    # H6 fix: repeated-CS design is more collinearity-prone than the
+    # within-demeaned panel path. Warn the user loudly when the design
+    # matrix is rank-deficient; results are still produced via pinv.
+    rank = int(np.linalg.matrix_rank(X))
+    if rank < X.shape[1]:
+        import warnings as _warnings
+        deficit = X.shape[1] - rank
+        _warnings.warn(
+            f"etwfe(panel=False): design matrix is rank-deficient by "
+            f"{deficit} column(s) (rank={rank}, ncol={X.shape[1]}). "
+            "Falling back to pseudoinverse; some coefficients are "
+            "arbitrary linear combinations. Consider dropping collinear "
+            "controls, shortening the event window, or using panel=True.",
+            RuntimeWarning, stacklevel=3,
+        )
+
     beta, se, vcov = _ols_fit(X, y_vec, cluster=cl_arr)
     n_obs = len(y_vec)
     df_resid = max(n_obs - X.shape[1], 1)
@@ -816,15 +854,25 @@ def _etwfe_never_only(
     Runs a separate ETWFE regression per cohort, each using only
     (units in cohort g) ∪ (never-treated units). Combines cohort ATTs
     with cohort-size weighting. Matches R ``etwfe(cgroup='never')``.
+
+    Notes
+    -----
+    Per-cohort regressions each run on a different subset (cohort g +
+    never-treated), so the cluster small-sample correction
+    ``(n_cl/(n_cl-1))`` is evaluated cohort-by-cohort. Cohort SEs may
+    therefore be slightly larger than a single full-sample regression
+    would produce. The aggregated SE assumes cross-cohort independence,
+    which is exact under this per-cohort design.
     """
+    # H3 fix: compute the first-treat series locally rather than
+    # writing a helper column back to the outer frame. Prevents
+    # accidental column leakage when callers re-use `data`.
     df = data.copy()
-    ft = df[first_treat].replace(0, np.nan)
-    df["_ft_cache"] = ft
-    cohorts = sorted(df.loc[df["_ft_cache"].notna(), "_ft_cache"].unique())
+    ft_local = df[first_treat].replace(0, np.nan)
+    cohorts = sorted(ft_local.dropna().unique())
     if len(cohorts) == 0:
         raise ValueError("No treated cohorts found. Check 'first_treat' column.")
-    never_mask = df["_ft_cache"].isna()
-    never_ids = df.loc[never_mask, group].unique()
+    never_ids = df.loc[ft_local.isna(), group].unique()
     if len(never_ids) == 0:
         raise ValueError(
             "cgroup='nevertreated' requires at least one never-treated "
@@ -834,8 +882,7 @@ def _etwfe_never_only(
     rows: List[Dict[str, Any]] = []
     ses: List[float] = []
     for g in cohorts:
-        # Keep cohort g + never-treated
-        coh_ids = df.loc[df["_ft_cache"] == g, group].unique()
+        coh_ids = df.loc[ft_local == g, group].unique()
         keep = np.concatenate([coh_ids, never_ids])
         sub = df.loc[df[group].isin(keep)].copy()
         if xvar:
@@ -1404,6 +1451,7 @@ def etwfe_emfx(
     result: CausalResult,
     type: str = "simple",
     alpha: float = 0.05,
+    include_leads: bool = False,
 ) -> CausalResult:
     """
     R ``etwfe::emfx``-style aggregated marginal effects for an ETWFE fit.
@@ -1429,6 +1477,15 @@ def etwfe_emfx(
         Aggregation type.
     alpha : float, default 0.05
         Significance level for confidence intervals.
+    include_leads : bool, default False
+        For ``type='event'`` and ``type='calendar'``, whether to include
+        pre-treatment relative times (``rel_time < 0``) in the output.
+        These coefficients identify pre-trends and are informative for
+        parallel-trends inspection. Default ``False`` for backward
+        compatibility with earlier versions; set ``True`` for full
+        event-study output matching the R ``etwfe::emfx(type='event')``
+        default. ``rel_time = -1`` is always the reference category
+        and is excluded.
 
     Returns
     -------
@@ -1554,6 +1611,12 @@ def etwfe_emfx(
     df_resid = max(result.n_obs - len(cohorts), 1)
     t_crit = stats.t.ppf(1 - alpha / 2, df_resid)
 
+    # H7 fix: default still post-only to preserve prior behaviour, but
+    # pre-treatment leads are available via include_leads=True for
+    # pre-trend inspection (standard event-study practice).
+    if not include_leads:
+        es = es.loc[es["rel_time"] >= 0].copy()
+
     if type == "event":
         key_col = "rel_time"; label_col = "event_time"
     else:

tests/test_did_summary.py

@@ -321,3 +321,64 @@ def test_did_summary_plot_rejects_non_did_summary_result():
     )
     with pytest.raises(ValueError, match="_did_summary_marker"):
         sp.did_summary_plot(bad)
+
+
+# ───────────────────────────────────────────────────────────────────────
+# 9. Follow-up fixes (H3 / H4 / H6 / H7)
+# ───────────────────────────────────────────────────────────────────────
+
+def test_etwfe_never_only_does_not_leak_columns(staggered_df):
+    df_before = staggered_df.copy()
+    cols_before = list(df_before.columns)
+    sp.etwfe(df_before, y="y", time="time", first_treat="first_treat",
+             group="unit", cgroup="nevertreated")
+    # No helper column should leak back into the caller's frame
+    assert list(df_before.columns) == cols_before
+
+
+def test_etwfe_xvar_order_invariant_and_distinct(staggered_df):
+    rng = np.random.default_rng(0)
+    df = staggered_df.copy()
+    df["x_a"] = rng.standard_normal(len(df))
+    df["x_b"] = rng.standard_normal(len(df)) * 2 + 5
+    r = sp.etwfe(df, y="y", time="time", first_treat="first_treat",
+                 group="unit", xvar=["x_a", "x_b"])
+    r_swap = sp.etwfe(df, y="y", time="time", first_treat="first_treat",
+                      group="unit", xvar=["x_b", "x_a"])
+    # Name-keyed indexing: swapping xvar order does not change slopes
+    assert np.allclose(r.detail["slope_x_a"].values,
+                       r_swap.detail["slope_x_a"].values)
+    assert np.allclose(r.detail["slope_x_b"].values,
+                       r_swap.detail["slope_x_b"].values)
+    # And slopes for different xvars must actually differ
+    for _, row in r.detail.iterrows():
+        assert abs(row["slope_x_a"] - row["slope_x_b"]) > 1e-8
+
+
+def test_etwfe_panel_false_rank_deficient_warns(staggered_df):
+    import warnings as _w
+    df = staggered_df.copy()
+    df["const_col"] = 1.0  # perfectly collinear with intercept
+    with _w.catch_warnings(record=True) as caught:
+        _w.simplefilter("always")
+        sp.etwfe(df, y="y", time="time", first_treat="first_treat",
+                 group="unit", panel=False, controls=["const_col"])
+        rank_warnings = [w for w in caught
+                         if issubclass(w.category, RuntimeWarning)
+                         and "rank-deficient" in str(w.message)]
+        assert len(rank_warnings) >= 1
+
+
+def test_etwfe_emfx_event_include_leads(staggered_df):
+    fit = sp.etwfe(staggered_df, y="y", time="time",
+                   first_treat="first_treat", group="unit")
+    # Default: post-only (backward compatibility)
+    ev_post = sp.etwfe_emfx(fit, type="event")
+    assert ev_post.detail["event_time"].min() == 0
+    # Explicit leads
+    ev_full = sp.etwfe_emfx(fit, type="event", include_leads=True)
+    assert ev_full.detail["event_time"].min() < 0
+    # rel_time = -1 is the reference category and must not appear
+    assert -1 not in ev_full.detail["event_time"].values
+    # Leads use the proper vcov-based SE, not the independence fallback
+    assert ev_full.model_info["se_method"].startswith("vcov-based")