fix(imputation): FE-span snap guard on the pretrends lead model

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

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_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)
 

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):

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)