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

CHANGELOG.md

@@ -62,6 +62,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   NumPy's groupby order. Verified: 1 distinct value across 50 identical calls (was 3/8);
   NumPy parity unchanged (~1e-14, the normal cross-backend GEMM tolerance); bit-identity
   regression tests cover both contiguous and non-contiguous unsorted cluster ids.
+- **`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

@@ -49,7 +49,6 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m
 | CR2 Bell-McCaffrey DOF uses a naive `O(n²k)` per-coefficient loop over cluster pairs; Pustejovsky-Tipton (2018) Appendix B has a scores-based formulation avoiding the full `n×n` `M`. Switch when a user hits a large-`n` cluster-robust design. | `linalg.py::_compute_cr2_bm` | Phase 1a | Heavy | Low |
 | Rust-backend CR2 Bell-McCaffrey: falls through to NumPy (the leverage/Satterthwaite-DOF path needs `return_dof` support, which the Rust vcov dispatch excludes). The one-way HC2 kernel landed 2026-07-07 (`compute_robust_vcov_hc2`, mirrors the NumPy hc2 branch at ~1e-15; near-singular hat-diagonal sentinel + Python-side warn-and-HC1-fallback). | `rust/src/linalg.rs` | Phase 1a | Mid | Low |
 | `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 |
-| 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_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; 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`.
 
 Auxiliary model residuals (Equation 8):

tests/test_imputation.py

@@ -3056,3 +3056,91 @@ 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 — 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,
+        # -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 = [
+            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)