docs+test(imputation): CI-review P3s — LSMR fallback doc sync + weighted parity oracle

M1: REGISTRY ImputationDiD passages (v_it derivation note + sparse
variance solver note) now describe the sparse LSMR fallback (certified
istop set {0,1,2,4,5}, uncertified -> full NaN inference tuple, dense
lstsq removed) instead of the pre-#657 dense-lstsq contract.
C1: five stale 'dense lstsq' comments/docstrings in imputation.py
updated to the LSMR contract.
D1: weighted singular-system parity test added — null-space components
of (A_0'[W]A_0) live in null(sqrt(W) A_0), so the WEIGHTED projection
W_0 A_0 z (what the weighted estimator consumes) must match the dense
oracle even where the unweighted projection need not; includes
zero-weight (subpopulation) rows.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: igerber

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
 
@@ -1589,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):
@@ -1657,16 +1659,16 @@ 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) "
@@ -1705,7 +1707,7 @@ def _solve_untreated_v(self, ctx: _UntreatedProjection, weights: np.ndarray) ->
             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(

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

tests/test_imputation.py

@@ -3086,6 +3086,32 @@ def test_singular_system_projection_matches_dense_oracle(self):
         # (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)."""