diff --git a/CHANGELOG.md b/CHANGELOG.md index d04bde5c..c50560c4 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 2927bb44..17d928c1 100644 --- a/TODO.md +++ b/TODO.md @@ -55,7 +55,6 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | `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 | | Multiplier-bootstrap weight chunking (CallawaySantAnna, EfficientDiD, and HAD — all wired through `diff_diff/bootstrap_chunking.py`) covers the **unstratified** survey-PSU generation (the default unit-level bootstrap — `cluster=None`, equivalently `cluster="unit"` — the large-`n_units` OOM case). Remaining gap: the **stratified** survey-PSU generator (`generate_survey_multiplier_weights_batch`, per-stratum + lonely-PSU pooling + FPC) still materializes the full `(n_bootstrap × n_psu)` matrix (consumed via sliced blocks). Stratified designs have few PSUs so this rarely OOMs; tile per-stratum generation over draws (each stratum's draws are independent → contiguous draw-blocks reproduce the stream bit-identically) if a large-PSU stratified design hits memory. | `diff_diff/bootstrap_chunking.py::iter_survey_multiplier_weight_blocks` | follow-up | 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 | | Per-cell `treated_units`/`control_units` label arrays (`all_units[positions]`, ~O(n_control) alloc per (g,t) cell) are consumed only by the precomputed-None fallback of the combined-IF assembly, which no in-package caller reaches — build them lazily (or drop from the IF-info dict) to cut per-cell allocation at high cell counts. | `staggered.py::_compute_att_gt_fast` | CS-scaling | 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 01e72b64..5cd755ca 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -1563,6 +1563,8 @@ 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 (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): 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)