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

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 1dc75d4f4c57f3d47a7b74a5fda69bba2851f891

---

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 docs/methodology/REGISTRY.md:L1566 and aligns with the shared absorbed-FE contract at docs/methodology/REGISTRY.md:L4617-L4648.
• Snapped columns flow through solve_ols rank handling and safe_inference, preserving all-NaN inference for unidentified leads.
• The previous P3 documentation concern is resolved: the covariate-branch parity note is present at docs/methodology/REGISTRY.md:L1568.
• The previous P3 test-strength concern is resolved: the test asserts effect, SE, t-stat, p-value, and CI endpoints at tests/test_imputation.py:L3103-L3107.

Methodology

Severity: P3 informational
Impact: The new guard is a documented identification-hygiene change for FE-spanned lead/covariate regressors in the BJS Eq. 9 pretrends regression. The implementation captures pre-demean norms, applies the shared two-stage snap after the within transform, and then lets rank handling report unidentified coefficients as NaN (diff_diff/imputation.py:L2297-L2326). This matches the registry note and the absorbed-FE contract.
Concrete fix: None required.

Code Quality

No findings. The PR reuses pre_demean_norms() and snap_absorbed_regressors() instead of adding estimator-local collinearity logic.

Performance

No findings. The LSMR confirmation is gated by the shared screen_tol path in snap_absorbed_regressors() (diff_diff/utils.py:L3306-L3319), so ordinary identified columns stay on the fast norm check.

Maintainability

No findings. Both callers of _compute_lead_coefficients() inherit the shared helper path: event-study pretrend coefficients at diff_diff/imputation.py:L1964-L1980 and pretrend_test() at diff_diff/imputation.py:L2526-L2541.

Tech Debt

No findings. The deferred TODO for adopting snap_absorbed_regressors on the imputation lead path has been removed from TODO.md:L44-L53.

Security

No findings.

Documentation/Tests

Severity: P3 informational
Impact: Prior P3 items are resolved. The registry retains the covariate parity audit trail at docs/methodology/REGISTRY.md:L1568, and the new test checks the full NaN inference tuple for the snapped lead at tests/test_imputation.py:L3103-L3107.
Concrete fix: None required.

Validation not run: this environment lacks pytest and numpy (python -m pytest failed with No module named pytest; a direct smoke check failed with No module named numpy).

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 3f316515b4f36ffd89e021ba9be650d741ffa1d8

---

Overall Assessment

✅ Looks good. No unmitigated P0/P1 findings.

Executive Summary

• Affected method: ImputationDiD BJS pretrends Test 1 / Equation 9 lead model.
• The FE-span snap guard is documented in docs/methodology/REGISTRY.md:L1566 and aligns with the shared absorbed-FE contract in diff_diff/utils.py:L3215-L3338.
• Snapped/dropped lead columns flow through solve_ols and safe_inference, preserving all-NaN inference fields for unidentified leads at diff_diff/imputation.py:L2419-L2432.
• Prior P3 concerns remain resolved: registry documentation is present, and tests assert the full NaN tuple plus the no-snap case.
• Validation: git diff --check passed. Targeted pytest could not run because pytest is not installed.

Methodology

Severity: P3 informational
Impact: The PR changes identification hygiene for BJS pretrends lead coefficients, not the estimand for identified leads. BJS develops the staggered-adoption imputation framework with inference and tests for identifying assumptions; the local registry specifies the Test 1 / Equation 9 lead regression on Omega_0 and now explicitly documents the FE-span snap guard. diff_diff/imputation.py:L2297-L2325 applies the documented guard after within-transforming outcome, leads, and covariates; identified columns remain untouched. (arxiv.org)
Concrete fix: None required.

Code Quality

No findings.

Performance

No findings. The helper only invokes LSMR confirmation for screened near-zero within columns, per diff_diff/utils.py:L3307-L3325.

Maintainability

No findings. The implementation reuses pre_demean_norms() and snap_absorbed_regressors() instead of adding estimator-local rank logic.

Tech Debt

No findings. The deferred TODO for adopting snap_absorbed_regressors on this path has been removed, and the shipped behavior is documented in CHANGELOG.md:L65-L75.

Security

No findings.

Documentation/Tests

Severity: P3 informational
Impact: The registry note at docs/methodology/REGISTRY.md:L1566-L1600 documents the deviation/guard, so it is not a methodology defect. The new tests cover both the FE-spanned lead NaN contract and the no-op identified-lead case at tests/test_imputation.py:L3061-L3146.
Concrete fix: None required.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 09bfafb42d2c1ad048be46e507618752969d9e71

---

Overall Assessment

✅ Looks good. No unmitigated P0/P1 findings.

Executive Summary

• Affected method: ImputationDiD BJS pretrends Test 1 / Equation 9 lead model.
• Prior P1/P3 concerns appear resolved: the FE-span guard is implemented and documented.
• Snapped lead columns flow through existing rank-deficiency handling and safe_inference(), preserving all-NaN inference for unidentified leads.
• New tests cover the FE-spanned lead warning/NaN contract and the identified-leads no-op case.
• Static review only; I did not run pytest in this read-only environment.

Methodology

Severity: P3 informational
Impact: The PR changes identification hygiene for FE-spanned pretrend lead indicators, not the estimand for identified BJS leads. The behavior is explicitly documented in docs/methodology/REGISTRY.md:L1566 and matches the shared absorbed-FE contract at docs/methodology/REGISTRY.md:L4617-L4648. The implementation applies the guard after within-transforming the lead/covariate model in diff_diff/imputation.py:L2297-L2325.
Concrete fix: None required.

Code Quality

No findings. The implementation reuses pre_demean_norms() and snap_absorbed_regressors() rather than adding estimator-local collinearity logic: diff_diff/imputation.py:L2297-L2325.

Performance

No findings. The shared helper only invokes the LSMR confirmation for screened near-zero within columns, per diff_diff/utils.py:L3356-L3374.

Maintainability

No findings. Both event-study pretrends and pretrend_test() delegate through _compute_lead_coefficients, so the guard is centralized: diff_diff/imputation.py:L1964-L1981 and diff_diff/imputation.py:L2525-L2541.

Tech Debt

No findings. The relevant TODO row was removed from the actionable backlog, consistent with the guard now landing in diff_diff/imputation.py.

Security

No findings.

Documentation/Tests

Severity: P3 informational
Impact: Documentation was added in the registry and changelog. Tests assert the full NaN inference tuple for the snapped lead and the no-warning no-op case for identified leads: tests/test_imputation.py:L3061-L3146.
Concrete fix: None required.