perf(two_stage): certified sparse-LSMR Stage-1 fallback — no dense normal-matrix materialization

[igerber]

Summary

• The RuntimeError fallback for the sparse Stage-1 factorization called np.linalg.lstsq(XtX_10.toarray(), ...) at four sites (analytical unweighted + weighted GMM sandwich, the SpilloverDiD Wave-D meat, and the bootstrap) — an O((U+T+K)²) dense materialization and OOM risk on large panels. All four now solve via certified per-column sparse LSMR (istop ∈ {0,1,2,4,5}; one uncapped-conlim retry; _LSMRUnconvergedError on failure).
• The TODO row's analysis deliverable: the row cautioned that ImputationDiD's null-space-invariance argument "does NOT transfer" because gamma_hat feeds coefficient-level consumers. The consumer trace refutes this for gamma_hat (every consumer is an X_10-range functional: Psi = X_10 γ; the GMM score correction c_g'γ with c_g ∈ rowspace(X_10); null(X'X) = null(X_10) exactly). The one consumer outside that argument — the bootstrap exact-residual helper's X_1 @ theta_exact on treated rows — is covered by min-norm agreement: both dense lstsq (SVD) and LSMR return the min-norm least-squares solution.
• Fail-closed boundaries: analytical vcov → NaN on _LSMRUnconvergedError (3 callers; the raise matters because the GMM-score nan_to_num would otherwise launder NaN scores into zeros and a finite, wrong variance); SpilloverDiD meat → NaN matrix; bootstrap → the established None degenerate contract.

Methodology references (required if estimator / math changes)

• Method name(s): Gardner (2022) two-stage DiD GMM sandwich / Butts (2021) SpilloverDiD Wave-D meat — numerical fallback-solver change only, no estimand/formula change.
• Paper / source link(s): Gardner (2022) arXiv:2207.05943; REGISTRY §TwoStageDiD Stage-1 fallback Note (rewritten with the two-pillar invariance argument).
• Any intentional deviations from the source (and why): None — solver choice is provably output-preserving (range-functional annihilation + min-norm agreement), locked by a helper-level dense-lstsq-oracle parity test on a singular Gram and a forced-fallback fit-level ATT/SE parity test.

Validation

• Tests added/updated: tests/test_two_stage.py (helper-level singular-Gram oracle parity incl. multi-RHS; forced-fallback fit-level LSMR-vs-dense-oracle ATT/SE parity; no-lstsq guard through the full fit; uncertified-istop fail-closed NaN; weighted/survey fallback with unit-constant weights); tests/test_spillover.py (meat LSMR fallback + fail-closed NaN through the spillover boundary).
• Backtest / simulation / notebook evidence (if applicable): N/A (fallback path only; happy path byte-unchanged).
• 451 two_stage + spillover tests pass; 3 clean local review rounds.

Security / privacy

• Confirm no secrets/PII in this PR: Yes

---

Generated with Claude Code

[github-actions]

Overall Assessment

⚠️ Needs changes — one P1 test-suite regression. I found no P0/P1 methodology correctness defect in the estimator change itself.

Executive Summary

• The affected methods are TwoStageDiD Gardner GMM sandwich/bootstrap and SpilloverDiD Wave-D GMM-corrected meat.
• The dense-to-LSMR fallback deviation is documented in docs/methodology/REGISTRY.md:L1725 with a labeled Note, so it is not a methodology defect.
• The analytical and Spillover fail-closed paths propagate NaN inference rather than finite wrong SEs.
• P1: tests/test_methodology_two_stage.py still asserts the old "lstsq" warning text, so the methodology suite will fail after this PR.
• Minor docs are stale in spillover.py, still describing the removed dense toarray() fallback.

Methodology

Finding: P3 informational
Impact: The solver substitution affects TwoStageDiD._compute_gmm_variance, TwoStageDiDBootstrapMixin._compute_cluster_S_scores, and SpilloverDiD Wave-D meat. The REGISTRY explicitly documents the certified sparse LSMR fallback, the invariance argument, accepted fail-closed behavior, and bootstrap None contract at docs/methodology/REGISTRY.md:L1725. The implementation follows that contract at diff_diff/two_stage.py:L2990-L3015, diff_diff/two_stage.py:L2451-L2470, diff_diff/two_stage_bootstrap.py:L210-L227, diff_diff/two_stage_bootstrap.py:L417-L435, and diff_diff/spillover.py:L3338-L3359.
Concrete fix: No methodology fix required.

Code Quality

Severity: None
Impact: No new inline inference anti-pattern or parameter-propagation issue found in the changed code.
Concrete fix: None.

Performance

Severity: None
Impact: The changed Stage-1 fallback no longer calls dense np.linalg.lstsq(XtX_10.toarray(), ...) at the modified sites. Remaining .toarray() calls in the reviewed files are for score aggregation paths, not the normal-matrix fallback.
Concrete fix: None.

Maintainability

Finding: P3
Impact: The new private LSMR helper is shared across analytical, bootstrap, and Spillover paths, which is good. The warning text inside _lsmr_certified_normal_solve is TwoStage-specific even when called through Spillover/bootstrap, which may confuse debugging but does not change behavior. See diff_diff/two_stage.py:L120-L128.
Concrete fix: Consider making the helper accept a short context string, or use estimator-neutral wording.

Tech Debt

Severity: None
Impact: Removing the old TODO row is justified because the PR adds the required REGISTRY note and implements the documented fallback.
Concrete fix: None.

Security

Severity: None
Impact: No secrets, unsafe IO, or security-sensitive changes found.
Concrete fix: None.

Documentation/Tests

Finding: P1
Impact: tests/test_methodology_two_stage.py:L408-L436 still documents and asserts the old dense fallback. The assertion requires a warning containing "lstsq", but the implementation now warns with "falling back to sparse LSMR" at diff_diff/two_stage.py:L3002-L3010. This will fail the methodology test suite and leaves the source-material regression test checking the wrong contract.
Concrete fix: Update the test docstring and assertion to expect the new sparse-LSMR warning, e.g. match "falling back to sparse LSMR".

Finding: P3
Impact: diff_diff/spillover.py:L1590-L1603 still says rank-deficient X_10'X_10 routes to np.linalg.lstsq / XtX_10.toarray(), which is now stale.
Concrete fix: Rewrite that note to describe the certified sparse LSMR fallback.

Path To Approval

1. Update tests/test_methodology_two_stage.py::test_singular_omega0_warns_and_falls_back to expect "sparse LSMR" instead of "lstsq" and adjust its docstring.
2. Update the stale Spillover docstring at diff_diff/spillover.py:L1590-L1603.

I did not run the suite; the local Python environment could not import numpy, so this is a static review.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: cb835e3abe815d86b7cd8ad20bbc1ec561d2ed18

---

Overall Assessment

✅ Looks good — no unmitigated P0/P1 findings. Prior P1 methodology-test regression is resolved.

Executive Summary

• Affected methods: TwoStageDiD Gardner GMM sandwich/bootstrap and SpilloverDiD Wave-D GMM-corrected meat.
• The dense Stage-1 normal-matrix fallback is now documented as certified sparse LSMR in docs/methodology/REGISTRY.md:L1725.
• The implementation fail-closes uncertified LSMR to NaN analytical inference or bootstrap None, matching the registry note.
• Previous P1 is fixed: tests/test_methodology_two_stage.py:L408-L436 now expects "falling back to sparse LSMR".
• Remaining issues are P3 documentation/maintainability polish only.
• I attempted targeted tests, but pytest is unavailable in this environment.

Methodology

Finding: P3 informational
Impact: The solver substitution affects TwoStageDiD._compute_gmm_variance, TwoStageDiDBootstrapMixin._compute_cluster_S_scores, and SpilloverDiD Wave-D meat. This is documented with a labeled Note in docs/methodology/REGISTRY.md:L1725, including the invariance argument and fail-closed behavior. Per the review rules, this documented numerical implementation choice is not a defect.
Concrete fix: None required.

Code Quality

Severity: None
Impact: I found no new inline inference anti-pattern or parameter propagation issue in the changed code.
Concrete fix: None.

Performance

Severity: None
Impact: No remaining Stage-1 Gram XtX_10.toarray() / dense lstsq fallback path was found in the modified implementation. The remaining executable dense lstsq is the Stage-2 exact-residual re-solve at diff_diff/two_stage_bootstrap.py:L110-L123, which is separate from this PR’s Stage-1 normal-matrix materialization issue.
Concrete fix: None.

Maintainability

Finding: P3
Impact: _lsmr_certified_normal_solve() emits a TwoStage-specific warning even when called through Spillover or bootstrap, which can make debugging less clear. See diff_diff/two_stage.py:L120-L128.
Concrete fix: Optionally pass a short context string into the helper and use estimator-specific warning text.

Tech Debt

Severity: None
Impact: Removing the old TODO row is acceptable because the PR adds the methodology registry note and implements the documented fallback.
Concrete fix: None.

Security

Severity: None
Impact: No secrets, unsafe IO, or security-sensitive changes found in the diff.
Concrete fix: None.

Documentation/Tests

Finding: P3
Impact: One Spillover docstring still says rank-deficient X_10'X_10 falls back to np.linalg.lstsq, even though the implementation now uses certified sparse LSMR. See diff_diff/spillover.py:L1590-L1593. A nearby comment also still says lstsq fallback at diff_diff/spillover.py:L3180-L3182.
Concrete fix: Update those references to say certified sparse LSMR fallback.

Severity: None
Impact: The previous P1 test regression is fixed at tests/test_methodology_two_stage.py:L408-L436. New tests cover LSMR parity, no-densify behavior, fail-closed NaN inference, weighted fallback, bootstrap warning text, and Spillover fail-closed behavior.
Concrete fix: None.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 9c7cd1f9785c3a079386d555e8f34f1caa882698

---

Overall Assessment

✅ Looks good — no unmitigated P0/P1 findings.

Executive Summary

• Affected methods: TwoStageDiD Gardner GMM sandwich/bootstrap and SpilloverDiD Wave-D GMM-corrected meat.
• The sparse LSMR Stage-1 fallback is documented in docs/methodology/REGISTRY.md:L1725 with fail-closed behavior.
• Implementation matches the registry: uncertified LSMR becomes NaN analytical inference or bootstrap None.
• Prior P3 items are addressed: Spillover stale dense lstsq references were updated, and fallback warnings are now context-specific.
• I found one P3 documentation precision issue only.
• I could not run tests here: pytest and numpy are not installed.

Methodology

Finding: P3 informational
Impact: The solver substitution affects TwoStageDiD._compute_gmm_variance, TwoStageDiDBootstrapMixin._compute_cluster_S_scores, and SpilloverDiD Wave-D meat. This is documented as a Note in docs/methodology/REGISTRY.md:L1725, including the range-functional/min-norm argument and fail-closed boundary. Per review rules, this documented numerical implementation choice is not a defect.
Concrete fix: None required.

Code Quality

Severity: None
Impact: No new inline inference anti-pattern, public parameter propagation issue, or partial NaN guard found in the modified code. Analytical paths route through safe_inference() after NaN vcov/SE propagation.
Concrete fix: None.

Performance

Severity: None
Impact: The modified Stage-1 fallback sites no longer materialize XtX_10.toarray() / XtWX_10.toarray() for dense lstsq. Remaining toarray() calls in two_stage.py:L3060-L3064 are score-aggregation paths, not the Stage-1 normal-matrix fallback.
Concrete fix: None.

Maintainability

Severity: None
Impact: The prior warning-context concern is resolved by _lsmr_certified_normal_solve(..., context=...) and context-specific Spillover/bootstrap callers.
Concrete fix: None.

Tech Debt

Severity: None
Impact: Removing the old TODO row is acceptable because the registry now documents the solver contract and the implementation removes the tracked OOM-risk path.
Concrete fix: None.

Security

Severity: None
Impact: No secrets, unsafe IO, or security-sensitive changes found in the diff.
Concrete fix: None.

Documentation/Tests

Finding: P3
Impact: CHANGELOG.md:L603-L608 says every gamma_hat/theta_exact consumer is an X_10-range functional, but the shared exact-residual helper evaluates X_1_sparse @ theta_exact on treated rows (diff_diff/two_stage_bootstrap.py:L92-L99). The registry correctly documents the min-norm rationale for that non-range-functional consumer, so this is documentation precision only. Similar wording appears in tests/test_two_stage.py:L525-L532, while tests/test_two_stage.py:L650-L658 says the fit-level parity test covers the singular case even though the fixture is full-rank.
Concrete fix: Reword the changelog/test docstrings to match the registry: gamma_hat consumers are range functionals; theta_exact treated-row use is covered by min-norm agreement.

Severity: P3 verification limitation
Impact: I attempted the targeted LSMR fallback tests, but this environment lacks pytest and numpy, so runtime verification was not possible here.
Concrete fix: Ensure CI runs the targeted tests/test_two_stage.py, tests/test_spillover.py, and tests/test_methodology_two_stage.py cases.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 70d00c26623e25ab70e84ddc1ebf0ff678583473

---

Overall Assessment

✅ Looks good — no unmitigated P0/P1 findings.

Executive Summary

• Affected methods: TwoStageDiD Gardner GMM sandwich/bootstrap and SpilloverDiD Wave-D GMM-corrected meat.
• The sparse LSMR fallback is documented in docs/methodology/REGISTRY.md:L1725 as a numerical implementation choice with fail-closed behavior.
• Analytical LSMR failure propagates to NaN SE/inference; bootstrap failure follows the documented None boundary.
• Prior documentation precision issue is addressed in the new changelog/test wording.
• Remaining items are P3 only: one stale registry phrase and local verification was blocked by missing dependencies.

Methodology

Finding: P3 informational
Impact: The solver substitution affects Stage-1 fallback solves in diff_diff/two_stage.py:L3015-L3019, diff_diff/two_stage_bootstrap.py:L223-L229, and diff_diff/two_stage.py:L375-L377. docs/methodology/REGISTRY.md:L1725 documents the range-functional/min-norm invariance argument and fail-closed boundaries. Per review rules, this documented numerical implementation choice is not a defect.
Concrete fix: None required.

Code Quality

Severity: None
Impact: No new public parameter propagation issue, partial NaN guard, or new inline inference anti-pattern found in the changed paths. The analytical call sites route SEs through safe_inference() after NaN vcov propagation.
Concrete fix: None.

Performance

Severity: None
Impact: The changed Stage-1 fallback sites no longer call dense lstsq on XtX_10.toarray() / XtWX_10.toarray(). Remaining toarray() calls in the touched runtime code are score-aggregation paths, not the normal-matrix fallback.
Concrete fix: None.

Maintainability

Severity: None
Impact: The new helper centralizes LSMR certification and uses caller-specific warning context. Spillover and bootstrap boundaries handle _LSMRUnconvergedError explicitly.
Concrete fix: None.

Tech Debt

Severity: None
Impact: Removing the TODO row is acceptable because the OOM-risk fallback is implemented and the methodology note now documents the solver contract.
Concrete fix: None.

Security

Severity: None
Impact: No secrets, unsafe IO, or security-sensitive changes found in the diff.
Concrete fix: None.

Documentation/Tests

Finding: P3
Impact: One older registry sentence still says a rank-deficient Spillover survey path leads to “lstsq fallback” at docs/methodology/REGISTRY.md:L4177. The authoritative new note at docs/methodology/REGISTRY.md:L1725 is correct, so this is documentation precision only.
Concrete fix: Reword that phrase to “LSMR fallback” or “fallback path.”

Finding: P3 verification limitation
Impact: I could not run the targeted tests because this environment lacks both pytest and numpy.
Concrete fix: Ensure CI runs the targeted tests/test_two_stage.py, tests/test_spillover.py, and tests/test_methodology_two_stage.py cases.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 087679094ac8dad0b95e517c4f60c8e587049943

---

Overall Assessment

✅ Looks good — no unmitigated P0/P1 findings.

Executive Summary

• Affected methods: TwoStageDiD Gardner/Butts-Gardner GMM sandwich + multiplier bootstrap, and SpilloverDiD Wave-D GMM-corrected meat.
• The solver change is documented in docs/methodology/REGISTRY.md:L1725 as a certified sparse LSMR fallback, with fail-closed behavior.
• Analytical variance paths convert uncertified LSMR to NaN SE/inference; Spillover converts to NaN meat/SE.
• The prior stale Spillover registry wording is fixed at docs/methodology/REGISTRY.md:L4177.
• I could not run tests locally because pytest and numpy are not installed in this environment.

Methodology

Finding: P3 informational
Impact: The PR changes the numerical fallback for Stage-1 normal-equation solves, not the estimand or sandwich formula. The affected fallback sites are centralized through _lsmr_certified_normal_solve() at diff_diff/two_stage.py:L64-L130, used by the Spillover meat at diff_diff/two_stage.py:L365-L377, analytical TwoStage variance at diff_diff/two_stage.py:L3003-L3019, and bootstrap scores at diff_diff/two_stage_bootstrap.py:L210-L229. The Registry documents this as a valid implementation choice with range/min-norm invariance and fail-closed boundaries at docs/methodology/REGISTRY.md:L1725.
Concrete fix: None required.

Code Quality

Severity: None
Impact: No new inline inference anti-pattern or partial NaN guard found in the changed paths. Analytical callers catch _LSMRUnconvergedError and route through existing safe_inference() propagation via NaN variance, e.g. diff_diff/two_stage.py:L2455-L2477, diff_diff/two_stage.py:L2631-L2650, and diff_diff/two_stage.py:L2761-L2780.
Concrete fix: None.

Performance

Severity: None
Impact: The changed Stage-1 fallback sites no longer materialize the dense normal matrix via XtX_10.toarray() / XtWX_10.toarray(). Remaining .toarray() calls in the touched runtime files are score aggregation paths, not the replaced normal-matrix fallback.
Concrete fix: None.

Maintainability

Severity: None
Impact: The new helper centralizes LSMR stop-code certification and warning behavior. Call-site boundaries are explicit for analytical TwoStage, Spillover meat, and bootstrap paths.
Concrete fix: None.

Tech Debt

Severity: None
Impact: The removed TODO row was addressed by implementation plus Registry documentation. No new untracked deferral was introduced.
Concrete fix: None.

Security

Severity: None
Impact: No secrets, unsafe IO, or security-sensitive behavior found in the diff.
Concrete fix: None.

Documentation/Tests

Finding: P3 verification limitation
Impact: Added tests cover sparse-LSMR fallback parity, no dense lstsq fallback, weighted survey fallback, uncertified analytical NaN propagation, and Spillover fail-closed behavior at tests/test_two_stage.py:L525-L727, tests/test_spillover.py:L676-L707, and tests/test_methodology_two_stage.py:L408-L437. I could not execute them locally because pytest is unavailable and importing numpy fails.
Concrete fix: Ensure CI runs the targeted tests/test_two_stage.py, tests/test_spillover.py, and tests/test_methodology_two_stage.py cases.