Signal silent sparse -> dense lstsq fallback in ImputationDiD and TwoStageDiD

[igerber]

Summary

• Three sites where a sparse factorization failure silently fell back to dense lstsq now emit a UserWarning with the original exception type and a note about the degraded-path implication.
  • diff_diff/imputation.py:1516 — variance path, scipy.sparse.linalg.spsolve on (A_0' W A_0) z = A_1' w. Previous bare except Exception was swallowing the root cause.
  • diff_diff/two_stage.py:1647 — GMM sandwich Stage 1 normal-equations via sparse_factorized.
  • diff_diff/two_stage_bootstrap.py:134 — bootstrap path, same pattern.
• Each site is called at most a handful of times per .fit() (per aggregation level for the main variance path; per bootstrap preparation stage for bootstrap), so one warning per fallback event is appropriate and the axis-B aggregation wrapper pattern is not needed here.
• No numerical change: the returned gamma_hat / z is identical to the silent path; only the signal is new.

Methodology references

• Methods: Borusyak-Jaravel-Spiess (2024) imputation DiD (variance via Theorem 3 auxiliary-residual sparse solve); Gardner (2022) two-stage DiD (GMM sandwich via sparse factorization of Stage 1 normal equations).
• REGISTRY updated under both estimators. The ImputationDiD Sparse variance solver bullet now notes the warning-on-fallback; the TwoStageDiD section gains a new **Note:** describing the signal on both the GMM sandwich and bootstrap fallback paths.

Validation

• 3 new tests use unittest.mock.patch to force spsolve / sparse_factorized to raise RuntimeError, then run .fit() and assert the expected UserWarning fires:
  • tests/test_imputation.py::TestImputationVariance::test_sparse_solver_dense_fallback_emits_warning
  • tests/test_two_stage.py::TestTwoStageDiDVariance::test_sparse_factorized_dense_fallback_emits_warning
  • tests/test_two_stage.py::TestTwoStageDiDVariance::test_sparse_factorized_bootstrap_dense_fallback_emits_warning
• Existing pre-PR test test_sparse_solver_dense_fallback (which verifies the fallback produces finite SE) continues to pass.
• 14 tests pass in the two affected variance test classes.

Security / privacy

• Confirm no secrets/PII in this PR: Yes

[github-actions]

Overall Assessment

Looks good.

Executive Summary

• The PR touches only the sparse-to-dense fallback sites in ImputationDiD’s Theorem 3 variance path and TwoStageDiD’s analytical/bootstrap Stage 1 normal-equations solve; I did not find any change to the underlying estimand, identification assumptions, weighting, or variance formula. citeturn3search2turn2search0
• The new warning behavior is documented in the registry at docs/methodology/REGISTRY.md:L1082-L1082 and docs/methodology/REGISTRY.md:L1165-L1165.
• I found no unmitigated P0/P1 issues in methodology adherence, NaN/inference handling, control-group logic, or parameter propagation in the changed code.
• P3: the new TwoStage tests assert warning emission but do not also verify that the dense fallback still returns usable SEs/results.
• The pre-existing ImputationDiD dense-fallback memory risk remains tracked in TODO.md:L62-L62 and is non-blocking for this review.

Methodology

No findings. The affected methods are the ImputationDiD variance solve in diff_diff/imputation.py:L1516-L1531, the TwoStageDiD GMM sandwich solve in diff_diff/two_stage.py:L1647-L1668, and the TwoStage bootstrap solve in diff_diff/two_stage_bootstrap.py:L134-L154. Cross-checking those sites against the estimator docstrings and the cited Borusyak-Jaravel-Spiess and Gardner/did2s methodology indicates this PR only adds UserWarning signaling around an existing sparse-to-dense linear-algebra fallback. That is a documented implementation choice, not a methodology deviation. citeturn3search2turn2search0

Code Quality

No findings. The edits are limited to warnings.warn(...) calls immediately ahead of the existing dense lstsq fallback branches in the three touched sites.

Performance

No findings. The fallback still uses the same dense solve path as before, so the PR does not introduce a new runtime or asymptotic performance regression in the reviewed code.

Maintainability

No findings. The three warning sites are mirrored by corresponding registry updates and targeted tests, so the new behavior is documented where users and maintainers will look first.

Tech Debt

• Severity: P3 (tracked in TODO.md). Impact: the ImputationDiD fallback still materializes A0tA0_sparse.toarray(), so the existing degraded-path O((U+T+K)^2) memory risk remains. This is already tracked and is not approval-blocking for this PR. References: diff_diff/imputation.py:L1530-L1531, TODO.md:L62-L62. Concrete fix: none required for approval; address via the existing TODO item or a follow-up memory-bounded fallback.

Security

No findings. I did not see any new secret-handling, data-exposure, or unsafe input-surface changes in the diff.

Documentation/Tests

• Severity: P3. Impact: the new TwoStage tests in tests/test_two_stage.py:L493-L530 prove warning emission, but unlike the existing ImputationDiD fallback regression in tests/test_imputation.py:L862-L885, they do not assert that the dense fallback still produces finite, usable SEs/results. A future control-flow regression could therefore keep the warning while breaking the degraded path. Concrete fix: after each warned fit(), assert results.overall_se is finite and positive, or compare the warned run against an unpatched fit on a fixed dataset.
• Residual validation gap: I could not execute the tests in this environment because pytest is not available on PATH, and the default python interpreter here does not have numpy installed, so the test assessment is by inspection only.

[igerber]

/ai-review

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 28507ec7fd32a44cfecc87fdc628d0eca3e1158e

---

Overall Assessment

✅ Looks good

Executive Summary

• Affected methods: the ImputationDiD Theorem 3 variance projection fallback in diff_diff/imputation.py:L1516-L1531, the TwoStageDiD GMM sandwich Stage 1 normal-equations fallback in diff_diff/two_stage.py:L1647-L1669, and the analogous TwoStage bootstrap fallback in diff_diff/two_stage_bootstrap.py:L134-L155.
• Cross-checking those sites against the estimator docstrings and variance definitions in diff_diff/imputation.py:L1-L15, diff_diff/imputation.py:L1446-L1455, diff_diff/two_stage.py:L1-L22, and diff_diff/two_stage.py:L1561-L1581 shows no change to the estimand, weighting, variance formula, identification assumptions, or defaults.
• The new signaling behavior is documented in the Methodology Registry, so there is no undocumented methodology deviation: docs/methodology/REGISTRY.md:L1082-L1082 and docs/methodology/REGISTRY.md:L1165-L1165.
• The prior re-review item is resolved: the new TwoStage tests now assert both warning emission and finite usable overall_se after the dense fallback in both analytical and bootstrap paths tests/test_two_stage.py:L493-L548.
• The pre-existing ImputationDiD dense-fallback memory risk remains tracked and non-blocking in TODO.md:L62-L62.

Methodology

• No findings. The touched branches only add UserWarning emission around pre-existing sparse-to-dense fallback paths; they do not alter the Theorem 3 projection in ImputationDiD or the Butts/Gardner GMM sandwich construction in TwoStageDiD diff_diff/imputation.py:L1446-L1455 diff_diff/two_stage.py:L1561-L1581.

Code Quality

• No findings. The diff is narrowly scoped to warning on degraded-path entry and does not introduce any inline inference computation, partial NaN guards, or parameter-propagation issues in the changed code.

Performance

• No findings. Warning emission is negligible relative to the already-existing dense toarray()/lstsq fallback work in diff_diff/imputation.py:L1530-L1531, diff_diff/two_stage.py:L1667-L1669, and diff_diff/two_stage_bootstrap.py:L153-L155.

Maintainability

• No findings. The new behavior is documented where maintainers will look for it and is covered by targeted regression tests docs/methodology/REGISTRY.md:L1082-L1082 docs/methodology/REGISTRY.md:L1165-L1165 tests/test_imputation.py:L888-L911 tests/test_two_stage.py:L493-L548.

Tech Debt

• Severity: P3 (tracked in TODO.md). Impact: the pre-existing ImputationDiD degraded path still densifies (A_0' A_0) via .toarray(), so the large-panel O((U+T+K)^2) memory risk remains unchanged by this PR diff_diff/imputation.py:L1530-L1531 TODO.md:L62-L62. Concrete fix: none required for approval; retain the existing TODO or replace the dense fallback with a memory-bounded solver in a follow-up.

Security

• No findings. The diff adds warning emission and tests only; it introduces no new secret-handling, data-exposure, or unsafe execution surface.

Documentation/Tests

• No findings. The previous P3 test gap is resolved: tests/test_two_stage.py:L493-L548 now verifies both warning emission and finite usable SEs on the warned fallback paths, and Imputation still has both the existing finite-SE fallback regression and the new warning assertion in tests/test_imputation.py:L862-L911.
• Execution note: I could not run the tests in this environment because pytest is not installed (pytest and python -m pytest were both unavailable).