diff --git a/CHANGELOG.md b/CHANGELOG.md index 605d9257..a687c869 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,6 +35,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `did_had_pretest_workflow`, ...) are unchanged in this release and removed separately. ### Fixed +- **`ImputationDiD` singular-variance fallback no longer densifies the normal matrix.** + The sparse-factorization fallback called `np.linalg.lstsq((A₀'[W]A₀).toarray(), …)` — + an `O((U+T+K)²)` dense materialization and OOM risk on large panels (triggered only + when the sparse factorization fails, e.g. a rank-deficient Ω₀). Both fallback sites now + solve via `scipy.sparse.linalg.lsmr` with no dense materialization. Solver choice + provably cannot change the estimator: least-squares solutions of the singular system + 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 plus a no-densify spy test through the full fit. Convergence is validated fail-closed: + an uncertified LSMR stop (condition-limit / max-iteration; machine-precision statuses + 4/5 count as certified per SciPy) gets one retry with an uncapped condition limit, + then raises internally and the variance boundary reports a NaN SE (full NaN inference + tuple) — raising rather than returning NaN matters because the missing-FE + `nan_to_num` in the psi product would launder a NaN vector into zeros and a finite, + wrong variance (locked by a fit-level NaN-inference regression test). Fallback warning text + updated ("sparse LSMR" instead of "dense lstsq"). The analogous `TwoStageDiD` dense + fallbacks are multi-RHS with coefficient-level consumers where the invariance argument + does not transfer; tracked separately in TODO. - **`HonestDiD` Δ^SD optimal-FLCI center parity with R (SE-audit B2b).** The optimal Fixed-Length CI optimizer was a flat Nelder-Mead over slope weights that landed on a different affine estimator than R `HonestDiD::findOptimalFLCI` at intermediate smoothness diff --git a/TODO.md b/TODO.md index db377ef5..1d76c9fb 100644 --- a/TODO.md +++ b/TODO.md @@ -45,7 +45,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | `EfficientDiD` conditional path: the largest remaining O(n) stage is the sieve/nuisance construction outside the tiled pass (~9s at 10k). (The `_ridge_solve_weights` Python-prep shave landed 2026-07-07 — the `omega_stack[rest]` fancy-index copy and tail scatter are skipped when no row is zero-masked, byte-identical outputs; the `zero_mask` abs scan itself remains, needed for correctness.) | `efficient_did_covariates.py` | CS-scaling | Mid | Low | | `CallawaySantAnna` repeated-cross-section ipw/dr paths (`_ipw_estimation_rc` / `_doubly_robust_rc`) re-solve the propensity logit for EVERY (g,t) cell with no cache — the panel paths dedup via `pscore_cache` keyed `(g, base_period[, t])`. The Phase 3 IRLS fast path already speeds each RC solve; the missing (g,t)-dedup cache is the remaining lever. Unprofiled surface — measure an RC covariate scenario first. | `staggered.py::_ipw_estimation_rc`, `_doubly_robust_rc` | CS-scaling | Mid | Low | | Rust `solve_ols` always runs a full thin SVD (equilibrated, gelsd-parity); at 40 covariates it is the top per-cell solver item on a CS dr fit (1.04s of 3.1s at 2M rows, 95 calls). A Cholesky/QR fast path with SVD fallback — mirroring the Phase 3 Python-side pattern (certify well-conditioned, fall back verbatim) — is the natural lever, but `solve_ols` is the universal OLS entry point (every estimator), so the blast radius needs the full backend-parity treatment (`TestSolveOLSSkipRankCheckParity` posture: fitted-values parity, not beta). | `rust/src/linalg.rs::solve_ols` | CS-scaling | Heavy | Low | -| `ImputationDiD` dense `(A0'A0).toarray()` scales `O((U+T+K)^2)` — OOM risk on large panels (only triggers when the sparse solver fails). Needs an alternative dense fallback or richer sparse strategy. | `imputation.py` | #141 | Heavy | Medium | +| `TwoStageDiD` dense `toarray()` fallbacks (`two_stage.py:297` unweighted + `:2922` GMM-sandwich weighted) share the OOM-risk pattern the ImputationDiD LSMR fix (2026-07) closed, but are MULTI-RHS solves whose `gamma_hat` feeds coefficient-level consumers directly — the null-space-invariance argument that made the imputation swap provably output-preserving does NOT transfer; needs its own analysis (does any consumer depend on the min-norm choice?) before an LSMR/column-loop swap. | `diff_diff/two_stage.py` | #141 | Mid | Low | | CR2 Bell-McCaffrey DOF uses a naive `O(n²k)` per-coefficient loop over cluster pairs; Pustejovsky-Tipton (2018) Appendix B has a scores-based formulation avoiding the full `n×n` `M`. Switch when a user hits a large-`n` cluster-robust design. | `linalg.py::_compute_cr2_bm` | Phase 1a | Heavy | Low | | Rust-backend CR2 Bell-McCaffrey: falls through to NumPy (the leverage/Satterthwaite-DOF path needs `return_dof` support, which the Rust vcov dispatch excludes). The one-way HC2 kernel landed 2026-07-07 (`compute_robust_vcov_hc2`, mirrors the NumPy hc2 branch at ~1e-15; near-singular hat-diagonal sentinel + Python-side warn-and-HC1-fallback). | `rust/src/linalg.rs` | Phase 1a | Mid | Low | | Wild cluster bootstrap CI inversion calls `_t_star(r)` ~O(100) times, each materializing a fresh `(B×n)` `y_star` + `(k×B)` refit + `(n×B)` residual arrays. Acceptable for the few-cluster regime; for large-`n`/large-`B`, chunk `_t_star` over draws or precompute the `r`-independent cluster-level pieces (restricted residuals are linear in `r`). | `utils.py::wild_bootstrap_se._t_star` | #543 | Mid | Low | diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index 9717378a..e1fe9f03 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -56,9 +56,9 @@ class _UntreatedProjection(NamedTuple): A_0: sparse.csr_matrix A_1: sparse.csr_matrix # solver(rhs) -> z; None when the factorization was exactly singular (the - # solve path then routes to a dense lstsq fallback). + # solve path then routes to the sparse LSMR least-squares fallback). solver: Optional[Callable[[np.ndarray], np.ndarray]] - A0tA0_csc: sparse.csc_matrix # retained for the dense-lstsq fallback + A0tA0_csc: sparse.csc_matrix # retained for the LSMR fallback survey_weights_0: Optional[np.ndarray] singular: bool @@ -68,6 +68,70 @@ class _UntreatedProjection(NamedTuple): # ============================================================================= +class _LSMRUnconvergedError(RuntimeError): + """LSMR failed to certify a solution on the singular-variance fallback. + + Raised (not returned as NaN) so the variance boundary can fail closed: + a NaN vector would be laundered into zeros by the missing-FE + ``nan_to_num`` in the psi product — producing a finite, WRONG variance — + whereas this exception is caught in ``_compute_conservative_variance`` + and converted to a NaN SE (the all-or-nothing NaN inference convention). + """ + + +def _lsmr_minnorm_normal_solve(A0tA0_csc, rhs: np.ndarray) -> np.ndarray: + """Least-squares solve of the (possibly singular) normal equations + ``(A_0'[W]A_0) z = rhs`` WITHOUT densifying the sparse matrix. + + Replaces the previous ``np.linalg.lstsq(A0tA0.toarray(), ...)`` fallback, + whose dense materialization scales ``O((U+T+K)^2)`` — an OOM risk on + large panels (the TODO row this resolves). ``scipy.sparse.linalg.lsmr`` + handles singular symmetric systems, converging to the minimum-norm + least-squares solution (the same solution family as ``lstsq``'s + pseudo-inverse solution). + + Solver choice cannot change the estimator output: any two least-squares + solutions differ by a ``null(A_0'[W]A_0) = null(sqrt(W) A_0)`` component, + which the downstream projection ``v_untreated = -[W_0] A_0 z`` + annihilates (unweighted: ``null = null(A_0)`` so ``A_0 z`` is invariant; + weighted: the weight multiplication zeroes exactly the rows where the + null component can be nonzero). Locked by the singular-system parity + test against a dense-lstsq oracle. + + CONVERGENCE IS VALIDATED (fail-closed): ``istop`` in ``{0, 1, 2, 4, 5}`` + means LSMR certified an (approximate) solution / least-squares solution + within ``atol``/``btol`` (4 and 5 are the machine-precision analogues of + 1 and 2 per SciPy's documentation); anything else (condition-limit stop, + max-iteration exhaustion) gets ONE retry with an uncapped condition + limit and a generous iteration budget, and if still uncertified raises + :class:`_LSMRUnconvergedError` — caught at the variance boundary and + converted to a NaN SE — rather than feeding a finite-but-unverified + solution into the Theorem 3 weights. + """ + import scipy.sparse.linalg as spla + + _certified = (0, 1, 2, 4, 5) + result = spla.lsmr(A0tA0_csc, rhs, atol=1e-14, btol=1e-14) + z, istop = result[0], int(result[1]) + if istop not in _certified or not np.all(np.isfinite(z)): + dim = A0tA0_csc.shape[0] + result = spla.lsmr( + A0tA0_csc, rhs, atol=1e-14, btol=1e-14, conlim=1e16, maxiter=max(50 * dim, 10_000) + ) + z, istop = result[0], int(result[1]) + if istop not in _certified or not np.all(np.isfinite(z)): + warnings.warn( + "ImputationDiD variance: the LSMR fallback solve of " + f"(A_0'[W]A_0) z = rhs did not converge (istop={istop}); " + "the affected variance is reported as NaN rather than from " + "an unverified solution.", + UserWarning, + stacklevel=3, + ) + raise _LSMRUnconvergedError(f"LSMR uncertified (istop={istop})") + return z + + class ImputationDiD(ImputationDiDBootstrapMixin): """ Borusyak-Jaravel-Spiess (2024) imputation DiD estimator. @@ -1460,25 +1524,31 @@ def _compute_conservative_variance( Standard error. """ sw_0 = survey_weights[omega_0_mask.values] if survey_weights is not None else None - cluster_psi_sums, _, ve_product = self._compute_cluster_psi_sums( - df=df, - outcome=outcome, - unit=unit, - time=time, - first_treat=first_treat, - covariates=covariates, - omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, - unit_fe=unit_fe, - time_fe=time_fe, - grand_mean=grand_mean, - delta_hat=delta_hat, - weights=weights, - cluster_var=cluster_var, - kept_cov_mask=kept_cov_mask, - survey_weights_0=sw_0, - proj_cache=proj_cache, - ) + try: + cluster_psi_sums, _, ve_product = self._compute_cluster_psi_sums( + df=df, + outcome=outcome, + unit=unit, + time=time, + first_treat=first_treat, + covariates=covariates, + omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, + unit_fe=unit_fe, + time_fe=time_fe, + grand_mean=grand_mean, + delta_hat=delta_hat, + weights=weights, + cluster_var=cluster_var, + kept_cov_mask=kept_cov_mask, + survey_weights_0=sw_0, + proj_cache=proj_cache, + ) + except _LSMRUnconvergedError: + # Solver failure is GLOBAL (the untreated projection is invalid), + # unlike per-observation missing-FE NaNs — fail the whole SE + # closed instead of letting nan_to_num launder it to zeros. + return np.nan if resolved_survey is not None: # Design-based variance with strata/PSU/FPC support @@ -1519,7 +1589,9 @@ def _build_untreated_projection( Uses scipy.sparse for FE dummy columns to reduce memory from O(N*(U+T)) to O(N) for the FE portion. An exactly singular ``A_0'[W]A_0`` makes ``sparse_factorized`` raise ``RuntimeError``; we emit a UserWarning (once - per fit) and record ``singular=True`` so the solve routes to dense lstsq. + per fit) and record ``singular=True`` so the solve routes to the sparse + LSMR least-squares fallback (no dense materialization; see + :func:`_lsmr_minnorm_normal_solve`). """ # Exclude rank-deficient covariates from design matrices if kept_cov_mask is not None and not np.all(kept_cov_mask): @@ -1587,21 +1659,22 @@ def _build_A_sparse(df_sub, unit_vals, time_vals): # Factorize once (factorize-once / solve-many). An exactly singular # matrix makes sparse_factorized raise RuntimeError -- the same condition # that previously surfaced as spsolve's MatrixRankWarning -> non-finite - # solution. Mirror the TwoStageDiD GMM-sandwich pattern: warn once and - # fall back to dense lstsq per target. (Bit-identical to the prior - # per-target spsolve for a single dense RHS -- both use the SuperLU - # simple driver with the same defaults.) + # solution. Warn once and fall back to the sparse LSMR least-squares + # solve per target (no dense materialization). (The factorized path is + # bit-identical to the prior per-target spsolve for a single dense + # RHS -- both use the SuperLU simple driver with the same defaults.) try: solver: Optional[Callable[[np.ndarray], np.ndarray]] = sparse_factorized(A0tA0_csc) singular = False except RuntimeError as exc: # Silent-failure audit axis C: emit a UserWarning on fallback instead - # of swallowing the error. Keep the "dense lstsq" substring (asserted + # of swallowing the error. Keep the "sparse LSMR" substring (asserted # by tests). warnings.warn( "ImputationDiD variance: sparse factorization of (A_0' [W] A_0) " - f"failed ({type(exc).__name__}); falling back to dense lstsq. This " - "may indicate a rank-deficient or near-singular normal-equations " + f"failed ({type(exc).__name__}); falling back to a sparse LSMR " + "least-squares solve (no dense materialization). This may " + "indicate a rank-deficient or near-singular normal-equations " "matrix and variance estimates may be less reliable.", UserWarning, stacklevel=2, @@ -1628,23 +1701,24 @@ def _solve_untreated_v(self, ctx: _UntreatedProjection, weights: np.ndarray) -> if ctx.singular: # Factorization was singular at build time (warned once already). - z, _, _, _ = np.linalg.lstsq(ctx.A0tA0_csc.toarray(), A1_w, rcond=None) + z = _lsmr_minnorm_normal_solve(ctx.A0tA0_csc, A1_w) else: assert ctx.solver is not None z = ctx.solver(A1_w) if not np.all(np.isfinite(z)): # Defensive, target-specific: a non-finite solve on an otherwise - # factorizable matrix routes this RHS to dense lstsq. Warn per + # factorizable matrix routes this RHS to the LSMR fallback. Warn per # target (silent-failure audit axis C) -- distinct from the # once-per-fit build-time singular warning. warnings.warn( "ImputationDiD variance: sparse solve of (A_0' [W] A_0) z = " - "A_1' w returned a non-finite solution; falling back to dense " - "lstsq for this target. Variance estimates may be less reliable.", + "A_1' w returned a non-finite solution; falling back to a " + "sparse LSMR least-squares solve for this target. Variance " + "estimates may be less reliable.", UserWarning, stacklevel=2, ) - z, _, _, _ = np.linalg.lstsq(ctx.A0tA0_csc.toarray(), A1_w, rcond=None) + z = _lsmr_minnorm_normal_solve(ctx.A0tA0_csc, A1_w) # v_untreated = -[W_0] A_0 z (WLS projection requires per-obs weight) v_untreated = -(ctx.A_0 @ z) diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index c99f4c1f..7892b1f3 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -1563,7 +1563,7 @@ Observation weights `v_it`: - For treated `(i,t) in Omega_1`: `v_it = w_it` (the aggregation weight) - For untreated `(i,t) in Omega_0` (FE-only **and** covariate cases): the exact imputation projection `v_untreated = -A_0 (A_0' A_0)^{-1} A_1' w_treated` (survey-weighted, with the left WLS weight factor `W_0`: `-W_0 A_0 (A_0' W_0 A_0)^{-1} A_1' w_treated`), where `A_0`, `A_1` are the two-way-FE (all unit dummies + time dummies dropping the first; plus any covariates) design matrices for untreated/treated observations. -**Note on v_it derivation:** The paper's Supplementary Proposition A3 gives the explicit `v_it^*` formula; it is not in the reviewed main-article PDF, so the projection is validated *empirically* against R `didimputation` (`tests/test_methodology_imputation.py::TestImputationDiDParityR`, SEs match to ~1e-10; the covariate branch — first stage `y ~ x | unit + time` on the untreated sample — is anchored separately by `TestImputationDiDCovariateParityR` on a time-varying-X panel, SEs ~2e-10). **Deviation note (superseded closed form):** the FE-only path previously used a closed form `-(w_i./n_{0,i} + w_.t/n_{0,t} - w../N_0)`, which is exact only for a *balanced* untreated set; because `Omega_0` is generically unbalanced in staggered designs (treated observations are removed), that form biased the SE (~27% on the parity panel) and was replaced by the exact projection above during the ImputationDiD methodology validation. A genuinely rank-deficient `A_0' A_0` (e.g. an unidentified period FE) routes to a dense least-squares fallback with a `UserWarning`. +**Note on v_it derivation:** The paper's Supplementary Proposition A3 gives the explicit `v_it^*` formula; it is not in the reviewed main-article PDF, so the projection is validated *empirically* against R `didimputation` (`tests/test_methodology_imputation.py::TestImputationDiDParityR`, SEs match to ~1e-10; the covariate branch — first stage `y ~ x | unit + time` on the untreated sample — is anchored separately by `TestImputationDiDCovariateParityR` on a time-varying-X panel, SEs ~2e-10). **Deviation note (superseded closed form):** the FE-only path previously used a closed form `-(w_i./n_{0,i} + w_.t/n_{0,t} - w../N_0)`, which is exact only for a *balanced* untreated set; because `Omega_0` is generically unbalanced in staggered designs (treated observations are removed), that form biased the SE (~27% on the parity panel) and was replaced by the exact projection above during the ImputationDiD methodology validation. A genuinely rank-deficient `A_0' A_0` (e.g. an unidentified period FE) routes to a sparse LSMR least-squares fallback with a `UserWarning` (no dense materialization; see the sparse-variance-solver Note below). Auxiliary model residuals (Equation 8): - Partition `Omega_1` into groups `G_g` (default: cohort × horizon) @@ -1611,7 +1611,7 @@ where `W_it(h) = 1[K_it = h]` are lead indicators, estimated on `Omega_0` only. - **Non-constant `first_treat` within a unit:** Emits `UserWarning` identifying the count and example unit. The estimator proceeds using the first observed value per unit (via `.first()` aggregation), but results may be unreliable. - **treatment_effects DataFrame weights:** `weight` column uses `1/n_valid` for finite tau_hat and 0 for NaN tau_hat, consistent with the ATT estimand (unweighted), or normalized survey weights `sw_i/sum(sw)` when `survey_design` is active. - **Rank-deficient covariates in variance:** Covariates with NaN coefficients (dropped for rank deficiency in Step 1) are excluded from the variance design matrices `A_0`/`A_1`. Only covariates with finite coefficients participate in the `v_it` projection. -- **Sparse variance solver:** the untreated projection `v_untreated = -A_0 (A_0'[W]A_0)^{-1} A_1'w` factorizes the normal-equations matrix `(A_0'[W]A_0)` once per `fit()` via `scipy.sparse.linalg.factorized` and reuses the factorization across every estimand target (overall ATT, each event-study horizon, each group, and the bootstrap precompute), solving only the target-specific RHS `A_1'w` per target -- factorize-once / solve-many (the design is target-invariant; only `weights` vary). This is **bit-identical** to the prior per-target `scipy.sparse.linalg.spsolve` for a single dense RHS (both use the SuperLU simple driver with the same defaults), built once instead of `O(targets)` times. Mirrors the TwoStageDiD GMM-sandwich `factorized` pattern. An exactly singular `(A_0'[W]A_0)` makes `factorized` raise `RuntimeError`; the build falls back to dense `lstsq` and emits a `UserWarning` once per fit (silent-failure audit axis C). A defensive per-target non-finite solve likewise routes to dense `lstsq` with a per-target `UserWarning`, so callers always know variance estimates came from the degraded path. The design is built/cached in `_build_untreated_projection` and solved per target in `_solve_untreated_v`. +- **Sparse variance solver:** the untreated projection `v_untreated = -A_0 (A_0'[W]A_0)^{-1} A_1'w` factorizes the normal-equations matrix `(A_0'[W]A_0)` once per `fit()` via `scipy.sparse.linalg.factorized` and reuses the factorization across every estimand target (overall ATT, each event-study horizon, each group, and the bootstrap precompute), solving only the target-specific RHS `A_1'w` per target -- factorize-once / solve-many (the design is target-invariant; only `weights` vary). This is **bit-identical** to the prior per-target `scipy.sparse.linalg.spsolve` for a single dense RHS (both use the SuperLU simple driver with the same defaults), built once instead of `O(targets)` times. Mirrors the TwoStageDiD GMM-sandwich `factorized` pattern. An exactly singular `(A_0'[W]A_0)` makes `factorized` raise `RuntimeError`; the build emits a `UserWarning` once per fit (silent-failure audit axis C) and the solve routes to a **sparse LSMR least-squares fallback** (`scipy.sparse.linalg.lsmr`, `atol=btol=1e-14`) — the previous dense `lstsq(toarray())` fallback materialized the `O((U+T+K)^2)` normal matrix, an OOM risk on large panels (2026-07). Solver choice cannot change the estimator output: least-squares solutions of the singular system differ only by `null(sqrt(W) A_0)` components, which the downstream projection `v = -[W_0] A_0 z` annihilates (dense-oracle parity test). Convergence is validated fail-closed: `istop` in `{0, 1, 2, 4, 5}` counts as certified (4/5 are SciPy's machine-precision analogues of 1/2); an uncertified stop gets one retry with an uncapped condition limit, then raises internally and the variance boundary reports a **full NaN inference tuple** — raising rather than returning NaN matters because the missing-FE `nan_to_num` in the psi product would otherwise launder a NaN vector into zeros and a finite, wrong variance. A defensive per-target non-finite solve likewise routes to the LSMR fallback with a per-target `UserWarning`, so callers always know variance estimates came from the degraded path. The design is built/cached in `_build_untreated_projection` and solved per target in `_solve_untreated_v`. - **Note:** Survey weights enter ImputationDiD via weighted iterative FE (Step 1), survey-weighted ATT aggregation (Step 3), and design-based variance via `compute_survey_if_variance()`. PSU clustering, stratification, and FPC are fully supported in the Theorem 3 variance path. When `resolved_survey` is present, the observation-level influence function (`v_it * epsilon_tilde_it`) is passed to `compute_survey_if_variance()` which applies the stratified PSU-level sandwich with FPC correction. Strata also enters survey df (n_PSU - n_strata) for t-distribution inference. Bootstrap + survey supported (Phase 6) via PSU-level multiplier weights. - **Bootstrap inference:** Uses multiplier bootstrap on the Theorem 3 influence function: `psi_i = sum_t v_it * epsilon_tilde_it`. Cluster-level psi sums are pre-computed for each aggregation target (overall, per-horizon, per-group), then perturbed with multiplier weights (Rademacher by default; configurable via `bootstrap_weights` parameter to use Mammen or Webb weights, matching CallawaySantAnna). This is a library extension (not in the paper) consistent with CallawaySantAnna/SunAbraham bootstrap patterns. - **Auxiliary residuals (Equation 8):** Implements the paper's *unit-clustered* Equation 8 aggregator, `tau_tilde_g = sum_i (sum_{t in G_g,i} v_it)(sum_{t in G_g,i} v_it * tau_hat_it) / sum_i (sum_{t in G_g,i} v_it)^2` (Borusyak-Jaravel-Spiess 2024, eq. 8, p. 3272; minimal-excess-variance derivation in Supplementary Appendix A.8): for each unit form the within-unit weight sum `a_{i,g}` and weighted-effect sum `b_{i,g}` over the unit's observations in group `g`, then combine across units. Groups partition `Omega_1` via `aux_partition` (default `"cohort_horizon"` = cohort × event-time; also `"cohort"` / `"horizon"`). Unimputable (NaN `tau_hat`) and off-target observations carry `v_it = 0` and are excluded from the aggregation — exact for finite `tau_hat` (a zero-weight row adds 0 to both `a` and `b`) and NaN-safe; a group with no contributing observations falls back to the unweighted group mean (a variance no-op, since `psi_g = sum_t v_it * eps_tilde_it = 0` there). diff --git a/tests/test_imputation.py b/tests/test_imputation.py index 2d7f65e2..5239ba62 100644 --- a/tests/test_imputation.py +++ b/tests/test_imputation.py @@ -908,7 +908,7 @@ def test_sparse_solver_dense_fallback_emits_warning(self): "diff_diff.imputation.sparse_factorized", side_effect=RuntimeError("test failure") ): with pytest.warns( - UserWarning, match="sparse factorization.*falling back to dense lstsq" + UserWarning, match="sparse factorization.*falling back to a sparse LSMR" ): est.fit( data, @@ -3056,3 +3056,149 @@ def test_main_fit_zero_weight_treated_unit_covariates(self): assert np.isnan(r.group_effects[2]["effect"]) # Cohort 3 is unaffected. assert np.isfinite(r.group_effects[3]["effect"]) + + +class TestLSMRFallbackParity: + """The sparse LSMR fallback replaces dense lstsq on the (possibly + singular) normal equations. Solver choice cannot change the estimator: + least-squares solutions differ only by null(A_0'[W]A_0) = null(sqrt(W)A_0) + components, which the projection v = -[W_0] A_0 z annihilates. Lock the + projection parity against a dense-lstsq oracle on a genuinely singular + system.""" + + def test_singular_system_projection_matches_dense_oracle(self): + import scipy.sparse as sp + + from diff_diff.imputation import _lsmr_minnorm_normal_solve + + rng = np.random.default_rng(3) + n, p = 200, 12 + A0_dense = rng.normal(size=(n, p)) + A0_dense[:, -1] = A0_dense[:, 0] # exact collinearity -> singular normal eqs + A_0 = sp.csr_matrix(A0_dense) + A0tA0 = sp.csc_matrix(A_0.T @ A_0) + rhs = rng.normal(size=p) + + z_lsmr = _lsmr_minnorm_normal_solve(A0tA0, rhs) + z_dense = np.linalg.lstsq(A0tA0.toarray(), rhs, rcond=None)[0] + assert np.all(np.isfinite(z_lsmr)) + # The z's may differ by a null-space component; the PROJECTION A_0 z + # (what the estimator consumes) must agree. + np.testing.assert_allclose(A_0 @ z_lsmr, A_0 @ z_dense, rtol=0, atol=1e-8) + + def test_weighted_singular_system_projection_matches_dense_oracle(self): + """Weighted variant (CI-review D1): the production path solves + (A_0'[W]A_0) z = rhs with survey weights W. Null-space components of + the weighted normal equations live in null(sqrt(W) A_0), so the + WEIGHTED projection W_0 A_0 z — what the weighted estimator + consumes — must agree across solvers even where the unweighted + projection A_0 z need not.""" + import scipy.sparse as sp + + from diff_diff.imputation import _lsmr_minnorm_normal_solve + + rng = np.random.default_rng(9) + n, p = 180, 10 + A0_dense = rng.normal(size=(n, p)) + A0_dense[:, -1] = 2.0 * A0_dense[:, 1] # exact collinearity + w = rng.uniform(0.2, 3.0, size=n) + w[:12] = 0.0 # zero-weight rows (subpopulation) stay inert + A_0 = sp.csr_matrix(A0_dense) + A0tWA0 = sp.csc_matrix((A_0.T.multiply(w)) @ A_0) + rhs = rng.normal(size=p) + + z_lsmr = _lsmr_minnorm_normal_solve(A0tWA0, rhs) + z_dense = np.linalg.lstsq(A0tWA0.toarray(), rhs, rcond=None)[0] + assert np.all(np.isfinite(z_lsmr)) + np.testing.assert_allclose(w * (A_0 @ z_lsmr), w * (A_0 @ z_dense), rtol=0, atol=1e-8) + + def test_no_dense_materialization_on_fallback(self, monkeypatch): + """The singular-build fallback path must never call .toarray() on the + normal matrix (the O((U+T+K)^2) OOM risk this closes).""" + import unittest.mock + + import diff_diff.imputation as imp + + data = generate_test_data(n_units=60, n_periods=6, seed=7) + + with unittest.mock.patch( + "diff_diff.imputation.sparse_factorized", side_effect=RuntimeError("forced") + ): + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + orig_lsmr = imp._lsmr_minnorm_normal_solve + calls = [] + + def _spy(mat, rhs): + calls.append(mat.shape) + mat.toarray = None # densifying would now raise + return orig_lsmr(mat, rhs) + + monkeypatch.setattr(imp, "_lsmr_minnorm_normal_solve", _spy) + res = ImputationDiD().fit( + data, outcome="outcome", unit="unit", time="time", first_treat="first_treat" + ) + assert calls, "fallback path did not route through the LSMR solver" + assert np.isfinite(res.overall_att) + assert any("sparse LSMR" in str(x.message) for x in w) + + def test_unconverged_lsmr_fails_closed_to_nan(self, monkeypatch): + """CI-review P1 regression: a finite-but-uncertified LSMR result + (istop outside {0,1,2} on both attempts) must NOT feed the variance; + the solve returns NaN so inference degrades to NaN.""" + import scipy.sparse as sp + + import diff_diff.imputation as imp + + def _fake_lsmr(A, b, **kwargs): + # finite vector, but istop=7 (max-iteration exhaustion) + return (np.ones(A.shape[0]), 7, 5, 1.0, 1.0, 1.0, 1.0, 1.0) + + monkeypatch.setattr("scipy.sparse.linalg.lsmr", _fake_lsmr) + A0tA0 = sp.csc_matrix(np.eye(4)) + with pytest.warns(UserWarning, match="did not converge"): + with pytest.raises(imp._LSMRUnconvergedError): + imp._lsmr_minnorm_normal_solve(A0tA0, np.ones(4)) + + def test_unconverged_lsmr_fit_level_nan_inference(self, monkeypatch): + """CI-review P0 regression: a globally failed solve must NOT be + laundered into finite inference by the missing-FE nan_to_num — the + full inference tuple degrades to NaN at the variance boundary.""" + import unittest.mock + + def _fake_lsmr(A, b, **kwargs): + return (np.ones(A.shape[0]), 7, 5, 1.0, 1.0, 1.0, 1.0, 1.0) + + data = generate_test_data(n_units=60, n_periods=6, seed=7) + monkeypatch.setattr("scipy.sparse.linalg.lsmr", _fake_lsmr) + with unittest.mock.patch( + "diff_diff.imputation.sparse_factorized", side_effect=RuntimeError("forced") + ): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = ImputationDiD().fit( + data, outcome="outcome", unit="unit", time="time", first_treat="first_treat" + ) + assert np.isfinite(res.overall_att) # point estimate unaffected + assert np.isnan(res.overall_se) + assert np.isnan(res.overall_t_stat) + assert np.isnan(res.overall_p_value) + assert np.all(np.isnan(np.asarray(res.overall_conf_int, dtype=float))) + + def test_machine_precision_istop_accepted(self, monkeypatch): + """CI-review P1 regression: istop 4/5 (machine-precision analogues of + 1/2 per SciPy) are certified — no retry, no failure handling.""" + import scipy.sparse as sp2 + + import diff_diff.imputation as imp2 + + calls = [] + + def _fake_lsmr(A, b, **kwargs): + calls.append(kwargs) + return (np.full(A.shape[0], 2.0), 4, 5, 1.0, 1.0, 1.0, 1.0, 1.0) + + monkeypatch.setattr("scipy.sparse.linalg.lsmr", _fake_lsmr) + z = imp2._lsmr_minnorm_normal_solve(sp2.csc_matrix(np.eye(3)), np.ones(3)) + assert len(calls) == 1 # accepted on the first attempt + np.testing.assert_array_equal(z, np.full(3, 2.0)) diff --git a/tests/test_methodology_imputation.py b/tests/test_methodology_imputation.py index 1fe9c13f..80d6a1cb 100644 --- a/tests/test_methodology_imputation.py +++ b/tests/test_methodology_imputation.py @@ -276,9 +276,9 @@ def test_singular_omega0_routes_to_dense_fallback(self) -> None: first_treat="first_treat", ) # The build-time RuntimeError on the singular factorization must trigger - # the dense-lstsq fallback (with a UserWarning carrying "dense lstsq"). - fallback_warnings = [w for w in caught if "dense lstsq" in str(w.message)] - assert fallback_warnings, "expected the dense-lstsq fallback under a singular Ω₀" + # the sparse LSMR fallback (with a UserWarning carrying "sparse LSMR"). + fallback_warnings = [w for w in caught if "sparse LSMR" in str(w.message)] + assert fallback_warnings, "expected the LSMR fallback under a singular Ω₀" # Factorize-once: the build-time singular warning fires a single time for # this single-target (overall-only) fit, not once per (g,t). assert len(fallback_warnings) == 1