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

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

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_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)