fix(imputation): LSMR singular-variance fallback — no dense normal-matrix materialization

[igerber]

Summary

• Closes the ImputationDiD dense-fallback TODO row (the Medium-priority OOM risk from Add Borusyak-Jaravel-Spiess (2024) Imputation DiD estimator #141): the singular-variance fallback called np.linalg.lstsq((A₀'[W]A₀).toarray(), …) — an O((U+T+K)²) dense materialization, triggered when the sparse factorization fails (e.g. rank-deficient Ω₀). Both fallback sites now solve via scipy.sparse.linalg.lsmr with no dense materialization (locked by a spy test that poisons .toarray through a full fit).
• Why the swap is provably output-preserving: least-squares solutions of the singular normal equations differ only by null(√W·A₀) components, which the downstream projection v = −[W₀]A₀z annihilates — locked by a dense-lstsq-oracle parity test on a genuinely singular system.
• Fail-closed convergence contract (two local-review rounds hardened this): istop ∈ {0,1,2,4,5} counts as certified (4/5 are SciPy's machine-precision analogues of 1/2 — round-3 P1); an uncertified stop gets one retry with an uncapped condition limit, then raises a private _LSMRUnconvergedError caught at the variance boundary and converted to a full NaN inference tuple. Raising rather than returning NaN matters: the missing-FE nan_to_num in the psi product would launder a NaN vector into zeros and a finite, wrong variance (round-2 P0) — a fit-level regression test forces istop=7 end-to-end and asserts NaN SE/t/p/CI with a finite point estimate.
• The analogous TwoStageDiD dense fallbacks are multi-RHS with coefficient-level consumers, where the null-space-invariance argument does not transfer — tracked as a new TODO row with the analysis needed, not blindly swapped.

Methodology references (required if estimator / math changes)

• Method name(s): BJS (2024) Theorem 3 conservative variance — solver substitution on the untreated projection only; certified solves are exact to atol=1e-14, uncertified solves fail closed to NaN
• Paper / source link(s): REGISTRY.md § ImputationDiD (exact-projection contract)
• Any intentional deviations from the source (and why): None — projection output parity oracle-locked; failure mode strictly safer (NaN vs potential finite-wrong).

Validation

• Tests added/updated: tests/test_imputation.py::TestLSMRFallbackParity (5 tests: singular-oracle parity, no-densify spy, unconverged raise, fit-level NaN inference, istop-4/5 acceptance); warning-text assertions updated in 2 files. Suites green: 179 imputation + methodology-imputation; two_stage untouched (126).
• 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, specifically BJS (2024) Theorem 3 conservative variance untreated-projection fallback.
• The estimator still computes the documented projection v_untreated = -W_0 A_0 (A_0' W_0 A_0)^-1 A_1'w; swapping dense lstsq for sparse lsmr is a valid numerical implementation choice, not a methodology defect.
• The uncertified-solve path fails closed to NaN SE and then flows through safe_inference(), preserving all-or-nothing NaN inference.
• No new inline inference anti-pattern was introduced in the changed analytical path.
• Static review only: I could not run tests locally because this review environment is missing scipy.

Methodology

Finding M1 — P3 informational

• Location: docs/methodology/REGISTRY.md:L1566-L1614, diff_diff/imputation.py:L82-L132
• Impact: The code now uses sparse LSMR for singular/non-finite fallback solves, but the registry still says rank-deficient A_0'A_0 routes to dense least-squares / dense lstsq. This is documentation drift, not a methodology blocker: the consumed object is the projected W_0 A_0 z, and least-squares solution differences in the null space are annihilated by that projection. SciPy documents lsmr as a least-squares solver and statuses 4/5 as machine-precision analogues of 1/2, matching the certified-status guard. (docs.scipy.org)
• Concrete fix: Update the ImputationDiD registry sparse-variance-solver note to say the singular fallback uses sparse LSMR, documents the certified istop set, and reports NaN SE on uncertified convergence.

Code Quality

Finding C1 — P3 informational

• Location: diff_diff/imputation.py:L59-L61, diff_diff/imputation.py:L1590-L1592, diff_diff/imputation.py:L1657-L1663, diff_diff/imputation.py:L1668-L1670, diff_diff/imputation.py:L1706-L1709
• Impact: Several comments/docstrings still refer to “dense lstsq” or old warning assertions. Runtime behavior is correct, but future reviewers may misread the fallback contract.
• Concrete fix: Replace the stale comment text with “sparse LSMR least-squares fallback” and remove references to preserving the old dense-warning substring.

Performance

No findings.

• Severity: None
• Impact: The PR removes the dense normal-matrix materialization in the ImputationDiD fallback path.
• Concrete fix: None required.

Maintainability

No blocking findings.

• Severity: None
• Impact: The new helper is localized and the exception boundary is placed at _compute_conservative_variance(), which is the right abstraction for failing the whole SE closed.
• Concrete fix: None required beyond the P3 stale-comment cleanup above.

Tech Debt

Finding T1 — P3 informational, tracked

• Location: TODO.md:L49
• Impact: The analogous TwoStageDiD dense fallback risk remains, but the PR explicitly tracks it and correctly notes that ImputationDiD’s null-space projection argument does not automatically transfer to coefficient-level multi-RHS consumers.
• Concrete fix: None required for this PR.

Security

No findings.

• Severity: None
• Impact: No secrets, unsafe deserialization, filesystem/network expansion, or user-controlled execution paths introduced.
• Concrete fix: None required.

Documentation/Tests

Finding D1 — P3 informational

• Location: tests/test_imputation.py:L3062-L3177, tests/test_methodology_imputation.py:L274-L285
• Impact: The new tests cover singular parity, no .toarray() fallback, uncertified LSMR failure, fit-level NaN inference, and accepted machine-precision stop statuses. The only notable gap is that the singular parity oracle is unweighted while the production path also supports survey-weighted A_0'W_0A_0.
• Concrete fix: Optional: add a weighted singular-system parity test asserting W_0 A_0 z_lsmr matches the dense oracle projection.