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.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 6b45c59ce540c03b0350ee2844bf2e309f09e36f

---

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 prior methodology-documentation drift is resolved: REGISTRY.md now documents sparse LSMR, certified istop handling, and fail-closed NaN variance behavior.
• The weighted parity test requested in the prior review has been added.
• No new inline inference anti-pattern was introduced in the changed analytical path; NaN SE flows through safe_inference().
• Static review only: I could not run tests because this environment lacks pytest and numpy.

Methodology

No blocking findings.

• Severity: None
• Impact: The LSMR replacement is a documented numerical implementation choice for the same projection contract, not an undocumented methodology deviation. The registry documents v_untreated = -[W_0]A_0(A_0'W_0A_0)^{-1}A_1'w, the LSMR fallback, and fail-closed NaN inference at docs/methodology/REGISTRY.md:L1562-L1614.
• Concrete fix: None required.

Code Quality

Finding C1 — P3 informational

• Location: tests/test_imputation.py:L868-L879, tests/test_imputation.py:L895-L898
• Impact: A test name/comment/docstring still says “dense fallback” / “dense lstsq,” even though the runtime path is now sparse LSMR. This is stale reviewer-facing text only.
• Concrete fix: Rename/update those test labels to “LSMR fallback.”

Performance

No findings.

• Severity: None
• Impact: The PR removes the A0tA0.toarray() materialization from the ImputationDiD fallback path.
• Concrete fix: None required.

Maintainability

No blocking findings.

• Severity: None
• Impact: _lsmr_minnorm_normal_solve() is localized, and _LSMRUnconvergedError is caught at the conservative variance boundary where full NaN inference can be preserved.
• Concrete fix: None required beyond the stale test wording above.

Tech Debt

Finding T1 — P3 informational, tracked

• Location: TODO.md:L48
• Impact: Analogous TwoStageDiD dense fallbacks remain, but the PR tracks them and correctly notes the ImputationDiD null-space invariance argument does not automatically transfer to multi-RHS coefficient consumers.
• Concrete fix: None required for this PR.

Security

No findings.

• Severity: None
• Impact: No secrets or unsafe execution/deserialization paths found in the changed files.
• Concrete fix: None required.

Documentation/Tests

Finding D1 — P3 informational, tracked

• Location: tests/test_csdid_ported.py:L950-L1026, TODO.md:L61
• Impact: Some Callaway-Sant’Anna golden assertions were loosened or removed in this PR. That reduces parity-test sensitivity, but it is test coverage only and is tracked as deferred tolerance-tightening work.
• Concrete fix: Optional: restore the tighter per-cell ATT/SE assertions, or make the TODO row explicitly name tests/test_csdid_ported.py and the affected C6-C8 checks.

Finding D2 — P3 informational

• Location: tests/test_imputation.py:L3061-L3204
• Impact: The new tests cover singular projection parity, weighted projection parity, no dense materialization, LSMR non-convergence raising, fit-level NaN inference, and istop=4/5 acceptance.
• Concrete fix: None required.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 1aa0424a892e9e52016182d36da87ac2a99cb70b

---

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 dense lstsq(toarray()) fallback is replaced with sparse lsmr in the singular/non-finite solve paths, and the registry documents the implementation choice and fail-closed behavior.
• The NaN/Inf edge case is handled correctly: uncertified LSMR raises internally, _compute_conservative_variance() returns NaN, and downstream inference uses safe_inference().
• Prior stale “dense fallback” test wording appears resolved.
• I could not run tests because this environment does not have pytest installed.

Methodology

No blocking findings.

• Severity: None
• Impact: The solver substitution is documented as an implementation choice for the same projection contract: v_untreated = -[W_0] A_0 (A_0' W_0 A_0)^{-1} A_1'w. The registry explicitly documents sparse LSMR, certified istop handling, and NaN fail-closed behavior at docs/methodology/REGISTRY.md:L1562-L1614.
• Concrete fix: None required.

Code Quality

No findings.

• Severity: None
• Impact: _lsmr_minnorm_normal_solve() is localized, both fallback call sites route through it, and _LSMRUnconvergedError is caught at the conservative variance boundary in diff_diff/imputation.py:L1527-L1551.
• Concrete fix: None required.

Performance

No findings.

• Severity: None
• Impact: The PR removes the A0tA0.toarray() materialization from the ImputationDiD singular fallback path; the remaining dense TwoStageDiD cases are tracked separately in TODO.md:L48.
• Concrete fix: None required.

Maintainability

No findings.

• Severity: None
• Impact: The fallback warning text, registry note, changelog, and tests now consistently refer to sparse LSMR rather than dense lstsq.
• Concrete fix: None required.

Tech Debt

No blocking findings.

• Severity: P3 informational, tracked
• Impact: Analogous TwoStageDiD dense fallback work remains, but it is explicitly tracked as separate analysis because the coefficient-level consumers do not inherit the same null-space-invariance argument. Location: TODO.md:L48.
• Concrete fix: None required for this PR.

Security

No findings.

• Severity: None
• Impact: No secrets, unsafe deserialization, shell execution, or user-controlled execution paths are introduced by the diff.
• Concrete fix: None required.

Documentation/Tests

Finding D1 — P3 informational

• Location: tests/test_imputation.py:L3145-L3148
• Impact: The new test docstring says uncertified means istop outside {0,1,2}, but the implemented and documented certified set is {0,1,2,4,5}. The test behavior is correct because it uses istop=7; this is stale wording only.
• Concrete fix: Update the docstring to say istop outside {0,1,2,4,5}.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: e537be45d2300488b1ca88357db6d554ee7593ca

---

Overall Assessment

✅ Looks good. No unmitigated P0/P1 findings.

Executive Summary

• Affected method: ImputationDiD, BJS (2024) Theorem 3 conservative variance untreated-projection fallback.
• The dense lstsq(toarray()) fallback is replaced by sparse LSMR without changing the documented projection contract.
• LSMR uncertified stops now fail closed: _LSMRUnconvergedError is caught at the variance boundary and downstream inference uses safe_inference().
• Prior docstring issue about certified istop values is resolved.
• I could not run tests because pytest is not installed in this environment.

Methodology

No findings.

• Severity: None
• Impact: The solver substitution preserves the registry’s projection formula v_untreated = -[W_0] A_0 (A_0' W_0 A_0)^-1 A_1'w; differing least-squares solutions only vary in the null space annihilated by the downstream projection. See docs/methodology/REGISTRY.md:L1562-L1614 and diff_diff/imputation.py:L82-L132.
• Concrete fix: None required.

Code Quality

No findings.

• Severity: None
• Impact: The LSMR helper is localized, the private exception is caught at _compute_conservative_variance, and public inference remains routed through safe_inference(). See diff_diff/imputation.py:L122-L131, diff_diff/imputation.py:L1527-L1551, and diff_diff/imputation.py:L718-L720.
• Concrete fix: None required.

Performance

No findings.

• Severity: None
• Impact: The modified ImputationDiD fallback no longer materializes A0tA0_csc.toarray() in production; both singular and non-finite solve fallbacks route through sparse LSMR. See diff_diff/imputation.py:L1702-L1721.
• Concrete fix: None required.

Maintainability

No findings.

• Severity: None
• Impact: Warning text, registry wording, changelog, and tests consistently describe the sparse LSMR fallback.
• Concrete fix: None required.

Tech Debt

• Severity: P3 informational, tracked
• Impact: Analogous dense TwoStageDiD fallbacks remain, but this is explicitly tracked in TODO.md:L48 with the reason the ImputationDiD null-space argument does not automatically transfer.
• Concrete fix: None required for this PR.

Security

No findings.

• Severity: None
• Impact: No secrets, unsafe deserialization, shell execution, or user-controlled execution paths are introduced by the diff.
• Concrete fix: None required.

Documentation/Tests

• Severity: P3 informational
• Impact: One new test docstring says the uncertified solve “returns NaN,” but the implementation intentionally raises _LSMRUnconvergedError; _compute_conservative_variance() returns NaN at the boundary. Behavior is correct. Location: tests/test_imputation.py:L3145-L3148.
• Concrete fix: Reword to “the solve raises and the variance boundary returns NaN.”

Validation not run: pytest is unavailable (pytest: command not found).

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 3ba44a6311b492bce5e5c6c408eecd2caa718da4

---

Overall Assessment

✅ Looks good. No unmitigated P0/P1 findings.

Executive Summary

• Affected method: ImputationDiD, BJS (2024) Theorem 3 conservative variance untreated-projection fallback.
• The dense lstsq(...toarray()) fallback is replaced with sparse LSMR in the two modified fallback sites.
• The methodology registry now documents the LSMR solver substitution, projection-invariance argument, convergence contract, and NaN fail-closed behavior.
• Prior concerns around uncertified LSMR stops and partial NaN inference are addressed: _LSMRUnconvergedError is caught at the variance boundary and routed through safe_inference().
• The analogous TwoStageDiD dense fallbacks are explicitly tracked in TODO.md, with the correct note that the imputation null-space argument does not automatically transfer.
• Validation not run: pytest is not installed in this environment.

Methodology

No findings.

• Severity: None
• Impact: The change preserves the documented projection contract for v_untreated = -[W_0] A_0 (A_0' W_0 A_0)^-1 A_1'w. The registry documents the solver substitution and why least-squares solution differences are annihilated by the downstream projection. See docs/methodology/REGISTRY.md:L1562-L1614 and diff_diff/imputation.py:L82-L132.
• Concrete fix: None required.

Code Quality

No findings.

• Severity: None
• Impact: The LSMR helper is localized, convergence is explicitly certified, and failure is represented by a private exception rather than a NaN vector that could be laundered by nan_to_num. See diff_diff/imputation.py:L71-L132, diff_diff/imputation.py:L1526-L1551, and diff_diff/imputation.py:L1702-L1727.
• Concrete fix: None required.

Performance

No findings.

• Severity: None
• Impact: The changed ImputationDiD fallback avoids production A0tA0_csc.toarray() materialization. Remaining dense TwoStageDiD fallbacks are outside this PR’s implementation scope and tracked in TODO.md:L48.
• Concrete fix: None required.

Maintainability

No findings.

• Severity: None
• Impact: Warning text, registry language, changelog entry, and tests consistently describe the sparse LSMR fallback and fail-closed behavior.
• Concrete fix: None required.

Tech Debt

• Severity: P3 informational, tracked
• Impact: The analogous TwoStageDiD dense fallback risk remains, but it is explicitly tracked with the required methodological caveat that multi-RHS coefficient consumers need separate analysis. Location: TODO.md:L48.
• Concrete fix: None required for this PR.

Security

No findings.

• Severity: None
• Impact: The diff does not introduce secrets, deserialization, shell execution, or user-controlled execution paths.
• Concrete fix: None required.

Documentation/Tests

No blocking findings.

• Severity: P3 informational
• Impact: New tests cover dense-oracle projection parity, weighted projection parity, no dense materialization, uncertified LSMR failure, fit-level NaN inference, and istop 4/5 acceptance. See tests/test_imputation.py:L3061-L3205 and tests/test_methodology_imputation.py:L269-L285.
• Concrete fix: None required.

Validation attempted but not run: PYTHONDONTWRITEBYTECODE=1 python -m pytest ... -q -p no:cacheprovider failed with No module named pytest.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 3db39205e81e80974c77872d88e2728d2957fafe

---

Overall Assessment

✅ Looks good — no unmitigated P0/P1 findings.

Executive Summary

• Affected method: ImputationDiD, BJS (2024) Theorem 3 conservative variance untreated-projection fallback.
• The dense lstsq(...toarray()) fallback is replaced with sparse LSMR in the two ImputationDiD fallback sites.
• The Methodology Registry documents the solver substitution, projection-invariance argument, convergence contract, and fail-closed NaN behavior.
• Prior LSMR concerns appear addressed: uncertified stops raise _LSMRUnconvergedError, _compute_conservative_variance() returns np.nan, and safe_inference() propagates full NaN inference.
• Validation not run: pytest is not installed. git diff --check passed.

Methodology

• Severity: P3 informational, documented implementation choice
• Impact: The solver substitution is documented in docs/methodology/REGISTRY.md:L1568 and docs/methodology/REGISTRY.md:L1616. The code preserves the documented projection contract v_untreated = -[W_0] A_0 z; least-squares solution differences are only null-space components annihilated by the downstream projection. See diff_diff/imputation.py:L88-L137 and diff_diff/imputation.py:L1700-L1733.
• Concrete fix: None required.

Code Quality

• Severity: None
• Impact: The LSMR helper is localized, checks certified istop values, retries once, warns, and raises a private exception instead of returning a NaN vector that could be zeroed by nan_to_num. See diff_diff/imputation.py:L120-L137 and diff_diff/imputation.py:L1478-L1481.
• Concrete fix: None required.

Performance

• Severity: None
• Impact: The changed fallback avoids production dense normal-matrix materialization. Remaining TwoStageDiD dense fallbacks are outside this PR’s implementation scope and tracked in TODO.md:L48.
• Concrete fix: None required.

Maintainability

• Severity: None
• Impact: Warning text, registry language, changelog text, and tests consistently refer to the sparse LSMR fallback.
• Concrete fix: None required.

Tech Debt

• Severity: P3 informational, tracked
• Impact: The analogous TwoStageDiD dense fallback risk remains, but it is explicitly tracked with the necessary methodology caveat that the ImputationDiD null-space argument does not transfer. See TODO.md:L48.
• Concrete fix: None required for this PR.

Security

• Severity: None
• Impact: No secrets, deserialization, shell execution, or user-controlled execution paths introduced.
• Concrete fix: None required.

Documentation/Tests

• Severity: None
• Impact: New tests cover singular dense-oracle projection parity, weighted parity, no dense materialization, uncertified LSMR failure, fit-level NaN inference, and acceptance of istop 4/5. See tests/test_imputation.py:L3147-L3290 and tests/test_methodology_imputation.py:L269-L285.
• Concrete fix: None required. Targeted pytest could not be run because pytest is unavailable in this environment.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 05c4b712d169882e2c0e92ff22657ff11ec86c35

---

Overall Assessment

✅ Looks good — no unmitigated P0/P1 findings.

Executive Summary

• Affected method: ImputationDiD, BJS (2024) Theorem 3 conservative variance untreated-projection fallback.
• Prior dense lstsq(...toarray()) fallback is replaced with sparse LSMR in both ImputationDiD fallback sites.
• The solver substitution, null-space projection invariance, convergence contract, and NaN fail-closed behavior are documented in REGISTRY.md.
• Prior re-review concerns remain addressed: uncertified LSMR stops raise _LSMRUnconvergedError, _compute_conservative_variance() returns np.nan, and safe_inference() makes t/p/CI all NaN.
• Validation: git diff --check passed. Targeted pytest could not run because pytest is not installed in this environment.

Methodology

• Severity: P3 informational, documented implementation choice
  Impact: The estimator’s consumed quantity remains v_untreated = -[W_0] A_0 z; the registry documents why least-squares solver choice is projection-invariant for singular (A_0'[W]A_0) and why uncertified solves fail closed. See diff_diff/imputation.py:L88-L137, diff_diff/imputation.py:L1700-L1733, and docs/methodology/REGISTRY.md:L1563-L1617.
  Concrete fix: None required.

Code Quality

• Severity: None
  Impact: The fallback helper is localized, retries uncertified LSMR once, validates finite output, and raises a private exception instead of returning a NaN vector that could be zeroed by nan_to_num. See diff_diff/imputation.py:L119-L137 and diff_diff/imputation.py:L1532-L1557.
  Concrete fix: None required.

Performance

• Severity: None
  Impact: The ImputationDiD production fallback no longer materializes the dense normal matrix. Remaining TwoStageDiD dense fallbacks are outside this PR’s changed implementation and are tracked with the necessary methodology caveat in TODO.md:L46.
  Concrete fix: None required.

Maintainability

• Severity: None
  Impact: Warning text, code comments, registry, changelog, and tests consistently describe the sparse LSMR fallback and fail-closed behavior.
  Concrete fix: None required.

Tech Debt

• Severity: P3 informational, tracked
  Impact: The analogous TwoStageDiD dense fallback risk remains, but it is explicitly tracked as deferred work and notes why the ImputationDiD invariance argument does not automatically transfer. See TODO.md:L46.
  Concrete fix: None required for this PR.

Security

• Severity: None
  Impact: No secrets, deserialization, shell execution, or user-controlled execution paths introduced.
  Concrete fix: None required.

Documentation/Tests

• Severity: None
  Impact: Tests cover dense-oracle projection parity, weighted parity, no dense materialization, uncertified LSMR failure, fit-level NaN inference, and accepted istop 4/5 statuses. See tests/test_imputation.py:L3147-L3291 and tests/test_methodology_imputation.py:L278-L285.
  Concrete fix: None required. Targeted pytest was not run locally because pytest is unavailable.