From 6c252d47e8f86d0ee3b64533b0711504589a9531 Mon Sep 17 00:00:00 2001 From: igerber Date: Tue, 7 Jul 2026 19:04:01 -0400 Subject: [PATCH 1/3] fix(imputation): FE-span snap guard on the pretrends lead model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Test-1 lead indicators + covariates now route through snap_absorbed_regressors after the within-transform (the same two-stage snap + LSMR confirmation the absorb= estimators use). A lead whose calendar period contains only its cohort's rows on Omega_0 collapses to a calendar-time dummy in the span of the absorbed time FE; it now snaps to exact zero — deterministic NaN coefficient + cause-specific warning naming lead[h] — instead of relying on the raw rank check alone, which the documented truncated-MAP-iterate exposure can defeat in slow-convergence regimes. Identified leads unchanged (full suites pass unmodified); behavioral tests lock the spanned-NaN contract and the no-op case. REGISTRY ImputationDiD note added. Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 11 +++++ TODO.md | 1 - diff_diff/imputation.py | 27 +++++++++++- docs/methodology/REGISTRY.md | 4 +- tests/test_imputation.py | 79 ++++++++++++++++++++++++++++++++++++ 5 files changed, 119 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 09b45f9c..cfa1726b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,6 +35,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `did_had_pretest_workflow`, ...) are unchanged in this release and removed separately. ### Fixed +- **`ImputationDiD` pretrends lead model gains the FE-span snap guard.** The Test-1 lead + indicators + covariates now route through `snap_absorbed_regressors` after the + within-transform (the same two-stage snap + LSMR confirmation the `absorb=` estimators + use). A lead whose calendar period contains only its cohort's rows on the untreated + sample collapses to a calendar-time dummy in the span of the absorbed time FE; it now + snaps to exact zero — deterministic NaN coefficient + cause-specific warning naming + `lead[h]` — instead of relying on the raw rank check alone, which the documented + truncated-MAP-iterate exposure can defeat in slow-convergence regimes (junk direction + perturbing the identified lead coefficients). Identified leads are unchanged (the full + imputation suites pass unmodified); behavioral tests lock both the spanned-NaN contract + and the no-op case. REGISTRY ImputationDiD note added. - **`HonestDiD` Δ^SD optimal-FLCI center parity with R (SE-audit B2b).** The optimal Fixed-Length CI optimizer was a flat Nelder-Mead over slope weights that landed on a different affine estimator than R `HonestDiD::findOptimalFLCI` at intermediate smoothness diff --git a/TODO.md b/TODO.md index 97444378..f8ac7eba 100644 --- a/TODO.md +++ b/TODO.md @@ -52,7 +52,6 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | `SpilloverDiD` sparse cKDTree path for the staggered nearest-treated-distance helper (mirrors the static helper's sparse branch). `_compute_nearest_treated_distance_staggered` always builds dense `(n_units, n_treated_by_onset)` matrices per cohort; add a sparse branch gated on `n > _CONLEY_SPARSE_N_THRESHOLD`. | `spillover.py` | Wave B | Mid | Low | | `HeterogeneousAdoptionDiD` Phase 3 Stute: Appendix-D vectorized form replaces the per-iteration OLS refit with a single precomputed `M = I - X(X'X)^{-1}X'` applied to `eps*eta` (~2× faster, functionally identical). Shipped the literal-refit form to match paper text. | `had_pretests.py::stute_test` | Phase 3 | Mid | Low | | Migrate `spillover._iterative_fe_subset` onto the shared `diff_diff.utils._iterative_fe_solve` (same Gauss-Seidel-on-codes recursion, specialized to a masked Butts subsample; ImputationDiD/TwoStageDiD already route through the shared helper — this takes the FE-solver copy count from 2 to 1). Preserve the SpilloverDiD positive-weight/NaN-FE REGISTRY contract and `_FE_ITER_MAX=100` budget (or align it to 10k with a REGISTRY note). | `spillover.py` | demean-modernization | Mid | Low | -| Adopt `snap_absorbed_regressors` (FE-spanned regressor two-stage snap + LSMR confirmation) on the ImputationDiD lead-indicator path — lead columns are the most plausible FE-spanned regressors, and the truncated-MAP-iterate exposure documented at `utils.py` applies; today only `solve_ols` rank detection guards it. Behavior change beyond the demean-modernization refactor, so deferred from that PR. | `imputation.py::_compute_lead_coefficients` | demean-modernization | Mid | Low | ### Testing / docs diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index 9717378a..e645d00d 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -37,7 +37,13 @@ ImputationDiDResults, ) from diff_diff.linalg import solve_ols -from diff_diff.utils import _iterative_fe_solve, demean_by_groups, safe_inference +from diff_diff.utils import ( + _iterative_fe_solve, + demean_by_groups, + pre_demean_norms, + safe_inference, + snap_absorbed_regressors, +) class _UntreatedProjection(NamedTuple): @@ -2289,6 +2295,7 @@ def _compute_lead_coefficients( # per-horizon n_obs counts below. within_transform pins [unit, time]; # [time, unit] here preserves the historical time-then-unit sweep order. narrow = df_0[[outcome, *all_x_cols, time, unit]].copy() + _pre_norms = pre_demean_norms(narrow, all_x_cols, weights=survey_weights_0) demeaned, _ = demean_by_groups( narrow, [outcome, *all_x_cols], @@ -2298,6 +2305,24 @@ def _compute_lead_coefficients( max_iter=10_000, tol=1e-10, ) + # FE-spanned regressors demean to numerical junk, not exact zero; + # snap them so rank handling drops them deterministically (NaN + # coefficient for that horizon) instead of the junk direction + # perturbing the identified lead coefficients. Lead indicators are + # the most plausible FE-spanned regressors here: with a single + # (balanced-restricted) cohort a lead h collapses to a calendar-time + # dummy on Omega_0, which lies exactly in the span of the absorbed + # time FE. + snap_absorbed_regressors( + demeaned, + all_x_cols, + _pre_norms, + absorbed_desc="unit and time fixed effects (pretrends lead model)", + group_vars=[time, unit], + rank_deficient_action=self.rank_deficient_action, + display_names={f"_lead_{h}": f"lead[{h}]" for h in pre_rel_times}, + weights=survey_weights_0, + ) y_dm = demeaned[outcome].to_numpy(dtype=np.float64) X_dm = demeaned[all_x_cols].to_numpy(dtype=np.float64) diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index c99f4c1f..d18454fd 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -1563,7 +1563,9 @@ Observation weights `v_it`: - For treated `(i,t) in Omega_1`: `v_it = w_it` (the aggregation weight) - For untreated `(i,t) in Omega_0` (FE-only **and** covariate cases): the exact imputation projection `v_untreated = -A_0 (A_0' A_0)^{-1} A_1' w_treated` (survey-weighted, with the left WLS weight factor `W_0`: `-W_0 A_0 (A_0' W_0 A_0)^{-1} A_1' w_treated`), where `A_0`, `A_1` are the two-way-FE (all unit dummies + time dummies dropping the first; plus any covariates) design matrices for untreated/treated observations. -**Note on v_it derivation:** The paper's Supplementary Proposition A3 gives the explicit `v_it^*` formula; it is not in the reviewed main-article PDF, so the projection is validated *empirically* against R `didimputation` (`tests/test_methodology_imputation.py::TestImputationDiDParityR`, SEs match to ~1e-10; the covariate branch — first stage `y ~ x | unit + time` on the untreated sample — is anchored separately by `TestImputationDiDCovariateParityR` on a time-varying-X panel, SEs ~2e-10). **Deviation note (superseded closed form):** the FE-only path previously used a closed form `-(w_i./n_{0,i} + w_.t/n_{0,t} - w../N_0)`, which is exact only for a *balanced* untreated set; because `Omega_0` is generically unbalanced in staggered designs (treated observations are removed), that form biased the SE (~27% on the parity panel) and was replaced by the exact projection above during the ImputationDiD methodology validation. A genuinely rank-deficient `A_0' A_0` (e.g. an unidentified period FE) routes to a dense least-squares fallback with a `UserWarning`. +**Note (pretrends lead FE-span guard, 2026-07):** the Test-1 lead model (`_compute_lead_coefficients`) routes its lead indicators + covariates through `snap_absorbed_regressors` after the within-transform (same two-stage snap + LSMR confirmation as the `absorb=` estimators; see "Absorbed Fixed Effects"). Lead indicators are the most plausible FE-spanned regressors: when a lead's calendar period contains only that cohort's rows on Omega_0 (e.g. never-treated units unobserved there), the lead collapses to a calendar-time dummy in the span of the absorbed time FE — it now snaps to exact zero and reports a deterministic NaN coefficient with a cause-specific warning (label `lead[h]`), instead of relying on the raw rank check alone (which the truncated-MAP-iterate exposure can defeat in slow-convergence regimes). Identified leads are unchanged (snap is a no-op; suite bit-stable). + +**Note on v_it derivation:** The paper's Supplementary Proposition A3 gives the explicit `v_it^*` formula; it is not in the reviewed main-article PDF, so the projection is validated *empirically* against R `didimputation` (`tests/test_methodology_imputation.py::TestImputationDiDParityR`, SEs match to ~1e-10). **Deviation note (superseded closed form):** the FE-only path previously used a closed form `-(w_i./n_{0,i} + w_.t/n_{0,t} - w../N_0)`, which is exact only for a *balanced* untreated set; because `Omega_0` is generically unbalanced in staggered designs (treated observations are removed), that form biased the SE (~27% on the parity panel) and was replaced by the exact projection above during the ImputationDiD methodology validation. A genuinely rank-deficient `A_0' A_0` (e.g. an unidentified period FE) routes to a dense least-squares fallback with a `UserWarning`. Auxiliary model residuals (Equation 8): - Partition `Omega_1` into groups `G_g` (default: cohort × horizon) diff --git a/tests/test_imputation.py b/tests/test_imputation.py index 2d7f65e2..1dc07615 100644 --- a/tests/test_imputation.py +++ b/tests/test_imputation.py @@ -3056,3 +3056,82 @@ def test_main_fit_zero_weight_treated_unit_covariates(self): assert np.isnan(r.group_effects[2]["effect"]) # Cohort 3 is unaffected. assert np.isfinite(r.group_effects[3]["effect"]) + + +class TestLeadSnapAbsorbed: + """FE-spanned lead indicators on the pretrends path are SNAPPED to exact + zero (deterministic NaN coefficient + cause-specific warning) instead of + reaching the solver as numerical junk — the snap_absorbed_regressors + adoption on _compute_lead_coefficients (TODO row: lead columns are the + most plausible FE-spanned regressors).""" + + @staticmethod + def _panel(never_treated_last_period): + rng = np.random.default_rng(5) + rows = [] + for i in range(30): + ft = 7 if i < 15 else 0 + periods = range(1, 8) if ft else range(1, never_treated_last_period + 1) + for t in periods: + y = ( + 1.0 + + 0.1 * i + + 0.2 * t + + (1.0 if (ft and t >= ft) else 0.0) + + rng.normal(0, 0.1) + ) + rows.append({"unit": i, "time": t, "first_treat": ft, "y": y}) + return pd.DataFrame(rows) + + def test_spanned_lead_snaps_to_nan_with_cause_warning(self): + # Never-treated units end at t=4, so Omega_0 at t=5 contains ONLY the + # g=7 cohort: lead[-2] == 1{t==5} on Omega_0 — exactly in the span of + # the absorbed time FE. + df = self._panel(never_treated_last_period=4) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + res = ImputationDiD(pretrends=True).fit( + df, + outcome="y", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + assert res.event_study_effects is not None + eff = res.event_study_effects + # The spanned lead is deterministically NaN (full inference tuple). + assert np.isnan(eff[-2]["effect"]) and np.isnan(eff[-2]["se"]) + # Identified leads stay finite (the junk direction never reached them). + for h in (-6, -5, -3): + assert np.isfinite(eff[h]["effect"]), f"h={h}" + assert np.isfinite(eff[h]["se"]) and eff[h]["se"] > 0, f"h={h}" + # Cause-specific snap warning names the display label, not the raw column. + snap_msgs = [ + str(x.message) + for x in w + if "collinear with the absorbed fixed effects" in str(x.message) + ] + assert any("lead[-2]" in m and "pretrends lead model" in m for m in snap_msgs), snap_msgs + + def test_identified_leads_unchanged_no_snap_warning(self): + # Balanced never-treated span: every lead period has never-treated + # rows in Omega_0 -> nothing is FE-spanned; the snap is a no-op and + # no cause-specific warning fires. + df = self._panel(never_treated_last_period=7) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + res = ImputationDiD(pretrends=True).fit( + df, + outcome="y", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + assert res.event_study_effects is not None + finite_leads = [ + h for h, e in res.event_study_effects.items() if h < -1 and np.isfinite(e["effect"]) + ] + assert len(finite_leads) >= 4 + assert not any("collinear with the absorbed fixed effects" in str(x.message) for x in w) From 372bf0a03017275cedc93f9af68eb7f6fc0e26d4 Mon Sep 17 00:00:00 2001 From: igerber Date: Wed, 8 Jul 2026 05:58:04 -0400 Subject: [PATCH 2/3] test: platform-robust lead-trap assertion (count-based, not horizon-specific) CI surfaced that WHICH redundant lead the rank handler drops after the snap (the leads-sum dummy trap column) is pivoted-QR/BLAS-order dependent: -4 on macOS/Accelerate, -3 on linux-arm py3.11, -6 on the pure-python CI backend. The snap contract itself (lead[-2] NaN + cause warning) is deterministic and unchanged; the sibling-finiteness assert now checks the count (exactly one trap drop) and survivor health instead of hardcoded horizons. Co-Authored-By: Claude Fable 5 --- tests/test_imputation.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/tests/test_imputation.py b/tests/test_imputation.py index 1dc07615..ec8b89d9 100644 --- a/tests/test_imputation.py +++ b/tests/test_imputation.py @@ -3100,11 +3100,18 @@ def test_spanned_lead_snaps_to_nan_with_cause_warning(self): ) assert res.event_study_effects is not None eff = res.event_study_effects - # The spanned lead is deterministically NaN (full inference tuple). + # The spanned lead is deterministically NaN (full inference tuple) — + # that is the snap's contract and is platform-independent. assert np.isnan(eff[-2]["effect"]) and np.isnan(eff[-2]["se"]) - # Identified leads stay finite (the junk direction never reached them). - for h in (-6, -5, -3): - assert np.isfinite(eff[h]["effect"]), f"h={h}" + # Of the remaining leads {-6,-5,-4,-3}, the leads-sum dummy trap costs + # exactly ONE more column — but WHICH one the rank handler drops is + # pivoted-QR/BLAS-order dependent (observed: -4 on macOS/Accelerate, + # -3 on linux-arm py3.11, -6 on the pure-python CI backend). Assert + # the count and the health of the survivors, not specific horizons. + others = [-6, -5, -4, -3] + finite = [h for h in others if np.isfinite(eff[h]["effect"])] + assert len(finite) == len(others) - 1, f"finite leads: {finite}" + for h in finite: assert np.isfinite(eff[h]["se"]) and eff[h]["se"] > 0, f"h={h}" # Cause-specific snap warning names the display label, not the raw column. snap_msgs = [ From 7763f71c2f23d0408411fb3e2606ecfe425c6024 Mon Sep 17 00:00:00 2001 From: igerber Date: Wed, 8 Jul 2026 07:17:25 -0400 Subject: [PATCH 3/3] test(imputation): assert the full NaN inference tuple on the spanned lead (review P3) Co-Authored-By: Claude Fable 5 --- tests/test_imputation.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_imputation.py b/tests/test_imputation.py index ec8b89d9..54de3a4a 100644 --- a/tests/test_imputation.py +++ b/tests/test_imputation.py @@ -3100,9 +3100,11 @@ def test_spanned_lead_snaps_to_nan_with_cause_warning(self): ) assert res.event_study_effects is not None eff = res.event_study_effects - # The spanned lead is deterministically NaN (full inference tuple) — - # that is the snap's contract and is platform-independent. + # The spanned lead is deterministically NaN — the FULL inference + # tuple (review P3: assert every field, not just effect/se). assert np.isnan(eff[-2]["effect"]) and np.isnan(eff[-2]["se"]) + assert np.isnan(eff[-2]["t_stat"]) and np.isnan(eff[-2]["p_value"]) + assert np.all(np.isnan(np.asarray(eff[-2]["conf_int"], dtype=float))) # Of the remaining leads {-6,-5,-4,-3}, the leads-sum dummy trap costs # exactly ONE more column — but WHICH one the rank handler drops is # pivoted-QR/BLAS-order dependent (observed: -4 on macOS/Accelerate,