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

[igerber]

Summary

• Adopts snap_absorbed_regressors (the two-stage FE-span snap + LSMR confirmation the absorb= estimators use) on ImputationDiD._compute_lead_coefficients — the change the demean-modernization PR deliberately deferred (TODO row). Lead indicators are the most plausible FE-spanned regressors: when a lead's calendar period contains only its cohort's rows on the untreated sample (e.g., never-treated units unobserved there), the lead collapses to a calendar-time dummy in the span of the absorbed time FE.
• Behavior: a spanned lead now snaps to exact zero → deterministic NaN coefficient for that horizon + a 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 (the junk direction then perturbs the identified lead coefficients; see the REGISTRY "Absorbed Fixed Effects" contract). Identified leads are untouched: the snap is a no-op there and the full imputation suites (incl. the R-parity anchors and the replicate/bit-equality grids) pass unmodified — 198 tests.
• New TestLeadSnapAbsorbed: (a) a spanned-lead design (never-treated span ends before a cohort pre-period) locks NaN + the cause-specific warning + finite untouched siblings; (b) a balanced design locks the no-op (no warning, leads finite). REGISTRY ImputationDiD note documents the guard.

Methodology references (required if estimator / math changes)

• Method name(s): Borusyak-Jaravel-Spiess (2024) pretrends Test 1 (Eq. 9 lead model) — identification-hygiene guard only; coefficients of identified leads are unchanged
• Paper / source link(s): REGISTRY.md § ImputationDiD (new "pretrends lead FE-span guard" Note) + § Absorbed Fixed Effects (snap contract)
• Any intentional deviations from the source (and why): None — an unidentified lead was never estimable; the guard makes its NaN deterministic and explained rather than convergence-dependent.

Validation

• Tests added/updated: tests/test_imputation.py::TestLeadSnapAbsorbed (2 behavioral tests); tests/test_imputation.py + tests/test_methodology_imputation.py + tests/test_replicate_weight_expansion.py all pass (198).
• Backtest / simulation / notebook evidence (if applicable): N/A

Security / privacy

• Confirm no secrets/PII in this PR: Yes

---

Generated with Claude Code

[github-actions]

Overall Assessment

✅ Looks good. No unmitigated P0/P1 findings.

Executive Summary

• Affected method: ImputationDiD BJS pretrends Test 1 / Eq. 9 lead model.
• The FE-span snap guard is documented in REGISTRY.md and matches the shared absorbed-FE contract.
• The implementation uses weighted pre-demean norms and passes the same survey weights into the snap and WLS path.
• Downstream rank-deficient OLS already returns NaN coefficients/VCV cleanly for dropped or rank-0 designs.
• I could not run the focused tests because pytest is not installed in this environment.

Methodology

• Severity: P3 informational
  Impact: The PR changes pretrend lead handling, but the behavior is explicitly documented as a guard for FE-spanned, unidentified leads. It routes lead indicators and covariates through snap_absorbed_regressors after within-transform, consistent with the absorbed-FE registry contract.
  Concrete fix: None required. See docs/methodology/REGISTRY.md:L1566 and diff_diff/imputation.py:L2297-L2325.

Code Quality

No findings.

Performance

No findings. The added LSMR confirmation is behind the existing two-stage snap screen, so it is scoped to near-absorbed candidate regressors rather than every column unconditionally; see diff_diff/utils.py:L3241-L3267.

Maintainability

No findings.

Tech Debt

No findings. The relevant deferred TODO row has been removed after implementation; remaining TODO entries do not track this specific gap. See TODO.md:L45-L58.

Security

No findings.

Documentation/Tests

• Severity: P3
  Impact: The spanned-lead test asserts effect and se are NaN, but it does not explicitly assert the full inference tuple (t_stat, p_value, conf_int) is NaN, despite the test comment describing full inference behavior. The implementation does use safe_inference, so this is a test-strength issue, not a correctness blocker.
  Concrete fix: In tests/test_imputation.py:L3103-L3105, also assert t_stat, p_value, and both CI endpoints are NaN, or reuse the existing assert_nan_inference helper.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 22a2b43fd0f9dcb0d2711c4040454d94039a6c4e

---

Overall Assessment

✅ Looks good. No unmitigated P0/P1 findings.

Executive Summary

• Affected method: ImputationDiD BJS pretrends Test 1 / Eq. 9 lead model.
• The FE-span guard is documented in REGISTRY.md, so this is not an undocumented methodology deviation.
• The implementation applies the snap after within-transform and preserves survey-weighted scaling through pre_demean_norms(..., weights=...) and snap_absorbed_regressors(..., weights=...).
• Dropped lead columns flow through existing rank-deficient OLS expansion and safe_inference, so NaN inference behavior is structurally protected.
• Prior P3 test-strength comment remains: the spanned-lead test still does not assert the full inference tuple.
• I could not run tests here because pytest is not installed.

Methodology

• Severity: P3 informational
  Impact: The PR changes ImputationDiD._compute_lead_coefficients, affecting BJS pretrends Test 1 / Eq. 9 lead estimation. The behavior is explicitly documented as a pretrends lead FE-span guard in REGISTRY.md, and it matches the shared “Absorbed Fixed Effects” snap contract.
  Concrete fix: None required. See docs/methodology/REGISTRY.md:L1566-L1568, docs/methodology/REGISTRY.md:L4615-L4646, and diff_diff/imputation.py:L2297-L2325.

Code Quality

No findings. The changed code uses the shared utility rather than adding inline collinearity or inference logic.

Performance

No findings. The added LSMR confirmation is gated behind the existing snap screen in snap_absorbed_regressors, so it only applies to near-absorbed candidates. See diff_diff/utils.py:L3253-L3267.

Maintainability

No findings. The change reuses the common FE-span utility and passes display labels for lead warnings.

Tech Debt

No findings. The deferred TODO row for adopting snap_absorbed_regressors on the imputation lead path was removed as completed.

Security

No findings.

Documentation/Tests

• Severity: P3
  Impact: Prior P3 remains. test_spanned_lead_snaps_to_nan_with_cause_warning says the spanned lead is NaN for the “full inference tuple,” but only asserts effect and se. The implementation uses safe_inference, so this is a test-strength gap, not a correctness blocker.
  Concrete fix: Add assert_nan_inference(eff[-2]) or explicit assertions for t_stat, p_value, and both CI endpoints at tests/test_imputation.py:L3103-L3105.

Validation not run: pytest / python -m pytest are unavailable in this environment (No module named pytest).

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 7763f71c2f23d0408411fb3e2606ecfe425c6024

---

Overall Assessment

✅ Looks good. No unmitigated P0/P1 findings.

Executive Summary

• Affected method: ImputationDiD BJS pretrends Test 1 / Eq. 9 lead model.
• The FE-span snap guard is documented in REGISTRY.md and matches the shared absorbed-FE contract.
• Dropped/spanned leads flow through solve_ols rank handling and safe_inference, so NaN inference is joint rather than partial.
• The prior P3 test-strength finding is addressed: the new test now asserts effect, se, t_stat, p_value, and CI endpoints.
• One P3 documentation item remains: the registry rewrite dropped an adjacent covariate-branch parity note.

Methodology

No P0/P1 findings.

Severity: P3 informational
Impact: The methodology change is documented and scoped to unidentified lead regressors in the BJS Eq. 9 pretrends model. diff_diff/imputation.py:L2297-L2325 captures pre-demean norms, runs the weighted within transform, and applies snap_absorbed_regressors; docs/methodology/REGISTRY.md:L1566 documents the guard, and docs/methodology/REGISTRY.md:L4615-L4646 documents the shared FE-spanned regressor contract.
Concrete fix: None required.

Code Quality

No findings. The PR reuses pre_demean_norms / snap_absorbed_regressors rather than adding local rank or inference logic.

Performance

No findings. The added LSMR confirmation is behind the shared screen_tol gate in snap_absorbed_regressors, so normal identified columns take the fast norm path.

Maintainability

No findings. The guard is applied at both call sites of _compute_lead_coefficients through the shared helper path: diff_diff/imputation.py:L1964-L1980 and diff_diff/imputation.py:L2526-L2541.

Tech Debt

No findings. The TODO item for adopting snap_absorbed_regressors on the imputation lead path was removed as completed in TODO.md:L45-L55.

Security

No findings.

Documentation/Tests

Severity: P3
Impact: The registry hunk rewrites the neighboring v_it derivation note and drops the prior mention that the covariate branch is separately anchored by tests/test_methodology_imputation.py::TestImputationDiDCovariateParityR, which still exists at tests/test_methodology_imputation.py:L1193. This weakens the methodology audit trail but does not change behavior.
Concrete fix: Restore the prior covariate-branch parity clause in docs/methodology/REGISTRY.md:L1568.

Prior P3 resolved: tests/test_imputation.py:L3103-L3107 now asserts the full NaN inference tuple for the snapped lead.

Validation not run: python -m pytest is unavailable in this environment (No module named pytest).