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

CHANGELOG.md

@@ -595,6 +595,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   `docs/methodology/REGISTRY.md` (both estimator sections + "Absorbed Fixed Effects").
 
 ### Performance
+- **`TwoStageDiD` / `SpilloverDiD` Stage-1 fallback no longer densifies the normal matrix.**
+  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. The TODO row's caution that the ImputationDiD
+  null-space-invariance argument "does NOT transfer" was refuted by the consumer trace:
+  every `gamma_hat` consumer is an `X_10`-range functional (`Psi = X_10 γ`; the GMM score
+  correction `c_g'γ` with `c_g ∈ rowspace(X_10)`), where `null(X'X) = null(X_10)` exactly
+  annihilates the min-norm ambiguity; the one `theta_exact` consumer outside that argument
+  (the bootstrap exact-residual helper's `X_1 @ theta` on treated rows) is covered by
+  min-norm agreement — both dense lstsq (SVD) and LSMR return the min-norm LS solution —
+  locked by a
+  dense-lstsq-oracle parity test on a singular Gram plus a no-densify guard through the
+  full fit. Uncertified LSMR (istop outside {0,1,2,4,5} after an uncapped retry) fails
+  closed: NaN vcov/SE on the analytical boundary, the established `None` degenerate on
+  the bootstrap boundary.
 - **Rust-backend HC2 vcov.** The Rust vcov path supported only HC1/CR1; one-way
   (unclustered, unweighted) HC2 now dispatches to a new `compute_robust_vcov_hc2` kernel
   mirroring the NumPy branch exactly — hat diagonals off the same bread,

TODO.md

@@ -43,7 +43,6 @@ 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 |
 | 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 |
-| `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 |
 | One-way (non-clustered) unweighted Bell-McCaffrey DOF still materializes the dense `n×n` residual-maker `M` in `_compute_bm_dof_from_contrasts` (`O(n²k)` hat build + `O(n²m)` per-contrast sums; practical for n < 10k). Switch to a scores-based evaluation like the clustered CR2-BM path now uses (singleton-cluster reduction of the same `B = diag(‖ω‖²) − P'M_U P` identity, row-chunked). | `diff_diff/linalg.py::_compute_bm_dof_from_contrasts` | #656 | Mid | Low |
 
 ### Testing / docs

diff_diff/spillover.py

@@ -40,7 +40,7 @@
 )
 from diff_diff.linalg import _rank_guarded_inv, solve_ols
 from diff_diff.results import SpilloverDiDResults
-from diff_diff.two_stage import _compute_gmm_corrected_meat
+from diff_diff.two_stage import _compute_gmm_corrected_meat, _LSMRUnconvergedError
 from diff_diff.utils import _iterative_fe_solve, safe_inference
 
 # Type alias mirroring diff_diff.conley.ConleyMetric so callers can supply
@@ -1589,17 +1589,16 @@ def _build_butts_fe_design_csr(
     -----
     Rank-deficient ``X_10' X_10`` (e.g. warn-and-drop units with no
     Omega_0 rows) is detected downstream by ``_compute_gmm_corrected_meat``
-    via ``sparse_factorized`` failure → ``np.linalg.lstsq`` fallback with
+    via ``sparse_factorized`` failure → certified sparse-LSMR fallback with
     a documented ``UserWarning``.
 
     **Re-factorization on entry:** when callers pass pre-mask integer
     codes that have had interior values dropped via ``finite_mask`` (a
     supported warn-and-drop fit), the input code arrays can be sparse —
     e.g. ``unit_codes = [0, 1, 3, 4]`` with code 2 dropped. Building
     ``X_10`` on the raw codes would materialize an all-zero FE column at
-    index 2, forcing ``sparse_factorized`` onto the dense
-    ``lstsq``/``XtX_10.toarray()`` fallback unnecessarily (large-memory
-    path on big panels). To avoid this, re-factorize via
+    index 2, forcing ``sparse_factorized`` onto the certified sparse-LSMR
+    fallback unnecessarily (a warned degraded path). To avoid this, re-factorize via
     :func:`pd.factorize` on entry to compact the code space to
     ``0..n_unique-1`` (no-op when codes are already contiguous; mirrors
     the column-space convention of ``TwoStageDiD._build_fe_design``).
@@ -3179,7 +3178,7 @@ def fit(
         # / time code; if the dropped unit sorts first, the fit-length and
         # full-length builds produce DIFFERENT column spaces (an all-zero
         # X_10 column for the dropped unit in the full-length build →
-        # rank-deficient `X_10' W X_10` → lstsq fallback → different
+        # rank-deficient `X_10' W X_10` → LSMR fallback → different
         # `gamma_hat`). The zero-pad invariant is preserved by zero-padding
         # the constructed Psi inside `_compute_gmm_corrected_meat` AFTER
         # the fit-sample gamma_hat / Psi build, NOT by rebuilding the FE
@@ -3333,25 +3332,30 @@ def fit(
         # arrays — survey-finite-mask subset of fit-sample inputs — plus
         # `X_*_sparse_fit` / `eps_10_fit` which are already built on
         # survey_finite_mask above).
-        meat_kept = _compute_gmm_corrected_meat(
-            X_1_sparse=X_1_sparse_fit,
-            X_10_sparse=X_10_sparse_fit,
-            eps_10=eps_10_fit,
-            X_2=X_2_kept_gamma,
-            eps_2=eps_2_fit_gamma,
-            vcov_type=_wave_d_vcov_mode,
-            cluster_ids=cluster_ids_for_meat,
-            conley_coords=conley_coords_for_meat,
-            conley_cutoff_km=_conley_cutoff_arg,
-            conley_metric=_conley_metric_arg,
-            conley_kernel="bartlett",
-            conley_time=conley_time_for_meat,
-            conley_unit=conley_unit_for_meat,
-            conley_lag_cutoff=_conley_lag_arg,
-            survey_weights=survey_weights_fit_gamma,
-            resolved_survey=resolved_survey_fit,
-            score_pad_mask=score_pad_mask_arg,
-        )
+        # An uncertified LSMR Stage-1 fallback solve inside the meat helper
+        # fails closed: NaN meat -> NaN SEs (the helper already warned).
+        try:
+            meat_kept = _compute_gmm_corrected_meat(
+                X_1_sparse=X_1_sparse_fit,
+                X_10_sparse=X_10_sparse_fit,
+                eps_10=eps_10_fit,
+                X_2=X_2_kept_gamma,
+                eps_2=eps_2_fit_gamma,
+                vcov_type=_wave_d_vcov_mode,
+                cluster_ids=cluster_ids_for_meat,
+                conley_coords=conley_coords_for_meat,
+                conley_cutoff_km=_conley_cutoff_arg,
+                conley_metric=_conley_metric_arg,
+                conley_kernel="bartlett",
+                conley_time=conley_time_for_meat,
+                conley_unit=conley_unit_for_meat,
+                conley_lag_cutoff=_conley_lag_arg,
+                survey_weights=survey_weights_fit_gamma,
+                resolved_survey=resolved_survey_fit,
+                score_pad_mask=score_pad_mask_arg,
+            )
+        except _LSMRUnconvergedError:
+            meat_kept = np.full((X_2_kept_gamma.shape[1], X_2_kept_gamma.shape[1]), np.nan)
 
         # Bread sandwich: A_22^{-1} = (X_2' W X_2)^{-1} via the shared rank-guarded
         # generalized inverse `_rank_guarded_inv` (column-drop on a near-singular