diff --git a/CHANGELOG.md b/CHANGELOG.md index ae86a0c9..2b63a1cb 100644 --- a/CHANGELOG.md +++ b/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, diff --git a/TODO.md b/TODO.md index 460dc5e6..b4400c94 100644 --- a/TODO.md +++ b/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 --git a/diff_diff/spillover.py b/diff_diff/spillover.py index 2c676b9f..96f27f9c 100644 --- a/diff_diff/spillover.py +++ b/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,7 +1589,7 @@ 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 @@ -1597,9 +1597,8 @@ def _build_butts_fe_design_csr( 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 diff --git a/diff_diff/two_stage.py b/diff_diff/two_stage.py index 4335eab0..0fb92ae8 100644 --- a/diff_diff/two_stage.py +++ b/diff_diff/two_stage.py @@ -55,6 +55,83 @@ # trades throughput for bounded memory. Keep in sync with two_stage_bootstrap.py. _SPARSE_DENSE_THRESHOLD = 10_000_000 + +class _LSMRUnconvergedError(RuntimeError): + """LSMR could not certify the Stage-1 normal-equation solve; the + variance boundary converts this to NaN inference (fail-closed).""" + + +def _lsmr_certified_normal_solve( + gram_csc, rhs: np.ndarray, context: str = "TwoStageDiD GMM sandwich" +) -> np.ndarray: + """Least-squares solve of the (possibly singular) sparse Stage-1 Gram + system ``gram @ out = rhs`` via per-column LSMR — no dense + materialization of the ``(p_1, p_1)`` normal matrix (`O((U+T+K)^2)` + OOM risk on large panels; the pattern the ImputationDiD LSMR fix + closed, ported here after the consumer-invariance analysis). + + OUTPUT-PRESERVING despite the min-norm ambiguity on singular systems: + least-squares solutions differ only by a ``null(X'X) = null(X_10)`` + component (weighted: ``null(X'WX) = null(W^{1/2}X_10)`` — zero-weight + rows are inert in every weighted consumer because Psi/score/residual + contributions carry the same ``W`` factor), and EVERY ``gamma_hat`` + consumer is an ``X_10``-range functional. One ``theta_exact`` consumer + (the bootstrap exact-residual helper's ``X_1_sparse @ theta_exact``) + evaluates theta on TREATED rows where a ``null(X_10)`` component would + NOT annihilate — parity there holds for a second reason: both dense + ``lstsq`` (SVD) and LSMR return the MIN-NORM least-squares solution, so + the two solvers agree on the whole vector (to iterative tolerance), not + just on range functionals; the fit-level singular-design parity test + locks this at the V/SE level. The remaining consumers are the + ``X_10``-range functionals — ``Psi_stage1 = X_10 @ gamma_hat``, the GMM + score correction ``c_g' gamma_hat`` with ``c_g = X_{10,g}' eps_{10,g}`` + in ``rowspace(X_10)``, and Stage-1 residuals ``y - X_10 theta`` — so + every null component annihilates. Locked by the singular-system parity + test against a dense-lstsq oracle. + + CONVERGENCE IS VALIDATED (fail-closed): ``istop`` in ``{0, 1, 2, 4, 5}`` + certifies a solution / least-squares solution within tolerance (4/5 are + the machine-precision analogues of 1/2 per SciPy); anything else gets + ONE retry with an uncapped condition limit, then raises + :class:`_LSMRUnconvergedError` — converted to NaN inference at the + variance boundary rather than feeding an unverified solution into the + GMM sandwich. + """ + import scipy.sparse.linalg as spla + + _certified = (0, 1, 2, 4, 5) + rhs_2d = np.atleast_2d(np.asarray(rhs, dtype=np.float64)) + if rhs_2d.shape[0] == 1 and np.asarray(rhs).ndim == 1: + rhs_2d = rhs_2d.T + dim = gram_csc.shape[0] + out = np.empty((dim, rhs_2d.shape[1])) + for j in range(rhs_2d.shape[1]): + result = spla.lsmr(gram_csc, rhs_2d[:, j], 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)): + result = spla.lsmr( + gram_csc, + rhs_2d[:, j], + 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( + f"{context}: the LSMR fallback solve of the " + f"Stage-1 normal equations 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})") + out[:, j] = z + return out + + # ============================================================================= # Wave D — Gardner GMM-corrected meat for SpilloverDiD # ============================================================================= @@ -153,8 +230,9 @@ def _compute_gmm_corrected_meat( **`gamma_hat` solve** (mirror of `TwoStageDiD._compute_gmm_variance` pattern at `two_stage.py:1886-1917`): factorize ``X_10' W X_10`` via - ``sparse_factorized`` (fast path); fall back to dense ``lstsq`` with - UserWarning when factorization fails. ``W`` is the diagonal of + ``sparse_factorized`` (fast path); fall back to the certified sparse + LSMR solve (:func:`_lsmr_certified_normal_solve`) with UserWarning when + factorization fails. ``W`` is the diagonal of ``survey_weights`` when provided (Hájek-normalized, ``sum_i w_i = n``); identity otherwise. ``gamma_hat`` has shape ``(p_1, p_2)``. @@ -269,7 +347,7 @@ def _compute_gmm_corrected_meat( # 1. gamma_hat = (X_10' W X_10)^{-1} (X_1' W X_2). Mirror the existing # TwoStageDiD method at two_stage.py:1886-1917 — sparse_factorized - # fast path with dense lstsq fallback + UserWarning on singular. + # fast path with certified sparse-LSMR fallback + UserWarning on singular. # When survey_weights is provided, X_10/X_1 cross-products use W. if survey_weights is not None: XtX_10 = X_10_sparse.T @ X_10_sparse.multiply(survey_weights[:, None]) @@ -288,15 +366,15 @@ def _compute_gmm_corrected_meat( warnings.warn( "SpilloverDiD Wave D GMM sandwich: sparse factorization of " f"(X_10' X_10) failed ({type(exc).__name__}); falling back to " - "dense lstsq. This may indicate a rank-deficient or " + "sparse LSMR. This may indicate a rank-deficient or " "near-singular Stage 1 design and SE estimates may be less " "reliable.", UserWarning, stacklevel=2, ) - gamma_hat = np.linalg.lstsq(XtX_10.toarray(), Xt1_X2, rcond=None)[0] - if gamma_hat.ndim == 1: - gamma_hat = gamma_hat.reshape(-1, 1) + gamma_hat = _lsmr_certified_normal_solve( + XtX_10.tocsc(), Xt1_X2, context="SpilloverDiD Wave-D GMM meat" + ) # 2. Psi = (X_10 @ gamma_hat) * eps_10[:, None] - X_2 * eps_2[:, None]. # Under Wave E.1 survey path, weights enter via element-wise eps @@ -2372,23 +2450,28 @@ def _stage2_static( att = float(coef[0]) # GMM sandwich variance - V = self._compute_gmm_variance( - df=df, - unit=unit, - time=time, - covariates=covariates, - omega_0_mask=omega_0_mask, - unit_fe=unit_fe, - time_fe=time_fe, - delta_hat=delta_hat, - kept_cov_mask=kept_cov_mask, - X_2=X_2, - cluster_ids=df[cluster_var].values, - survey_weights=survey_weights, - resolved_survey=resolved_survey, - score_pad_mask=score_pad_mask, - cluster_ids_full=cluster_ids_full, - ) + # An uncertified LSMR Stage-1 fallback solve fails closed: + # NaN vcov -> NaN SE/t/p/CI (the helper already warned). + try: + V = self._compute_gmm_variance( + df=df, + unit=unit, + time=time, + covariates=covariates, + omega_0_mask=omega_0_mask, + unit_fe=unit_fe, + time_fe=time_fe, + delta_hat=delta_hat, + kept_cov_mask=kept_cov_mask, + X_2=X_2, + cluster_ids=df[cluster_var].values, + survey_weights=survey_weights, + resolved_survey=resolved_survey, + score_pad_mask=score_pad_mask, + cluster_ids_full=cluster_ids_full, + ) + except _LSMRUnconvergedError: + V = np.full((X_2.shape[1], X_2.shape[1]), np.nan) se = float(np.sqrt(max(V[0, 0], 0.0))) return att, se @@ -2543,23 +2626,28 @@ def _stage2_event_study( ) # GMM variance for full coefficient vector - V = self._compute_gmm_variance( - df=df, - unit=unit, - time=time, - covariates=covariates, - omega_0_mask=omega_0_mask, - unit_fe=unit_fe, - time_fe=time_fe, - delta_hat=delta_hat, - kept_cov_mask=kept_cov_mask, - X_2=X_2, - cluster_ids=df[cluster_var].values, - survey_weights=survey_weights, - resolved_survey=resolved_survey, - score_pad_mask=score_pad_mask, - cluster_ids_full=cluster_ids_full, - ) + # An uncertified LSMR Stage-1 fallback solve fails closed: + # NaN vcov -> NaN SE/t/p/CI (the helper already warned). + try: + V = self._compute_gmm_variance( + df=df, + unit=unit, + time=time, + covariates=covariates, + omega_0_mask=omega_0_mask, + unit_fe=unit_fe, + time_fe=time_fe, + delta_hat=delta_hat, + kept_cov_mask=kept_cov_mask, + X_2=X_2, + cluster_ids=df[cluster_var].values, + survey_weights=survey_weights, + resolved_survey=resolved_survey, + score_pad_mask=score_pad_mask, + cluster_ids_full=cluster_ids_full, + ) + except _LSMRUnconvergedError: + V = np.full((X_2.shape[1], X_2.shape[1]), np.nan) # Build results dict event_study_effects: Dict[int, Dict[str, Any]] = {} @@ -2668,23 +2756,28 @@ def _stage2_group( ) # GMM variance - V = self._compute_gmm_variance( - df=df, - unit=unit, - time=time, - covariates=covariates, - omega_0_mask=omega_0_mask, - unit_fe=unit_fe, - time_fe=time_fe, - delta_hat=delta_hat, - kept_cov_mask=kept_cov_mask, - X_2=X_2, - cluster_ids=df[cluster_var].values, - survey_weights=survey_weights, - resolved_survey=resolved_survey, - score_pad_mask=score_pad_mask, - cluster_ids_full=cluster_ids_full, - ) + # An uncertified LSMR Stage-1 fallback solve fails closed: + # NaN vcov -> NaN SE/t/p/CI (the helper already warned). + try: + V = self._compute_gmm_variance( + df=df, + unit=unit, + time=time, + covariates=covariates, + omega_0_mask=omega_0_mask, + unit_fe=unit_fe, + time_fe=time_fe, + delta_hat=delta_hat, + kept_cov_mask=kept_cov_mask, + X_2=X_2, + cluster_ids=df[cluster_var].values, + survey_weights=survey_weights, + resolved_survey=resolved_survey, + score_pad_mask=score_pad_mask, + cluster_ids_full=cluster_ids_full, + ) + except _LSMRUnconvergedError: + V = np.full((X_2.shape[1], X_2.shape[1]), np.nan) group_effects: Dict[Any, Dict[str, Any]] = {} for g in treatment_groups: @@ -2908,22 +3001,22 @@ def _compute_gmm_variance( ) theta_exact = np.asarray(solve_XtX(np.asarray(rhs_fe).ravel())).ravel() except RuntimeError as exc: - # Singular matrix — fall back to dense least-squares. Silent-failure + # Singular matrix — fall back to certified sparse LSMR. Silent-failure # audit axis C: emit a UserWarning on fallback instead of swallowing. warnings.warn( "TwoStageDiD GMM sandwich: sparse factorization of " f"(X'_{{10}} W X_{{10}}) failed ({type(exc).__name__}); falling " - "back to dense lstsq. This may indicate a rank-deficient or " + "back to sparse LSMR. This may indicate a rank-deficient or " "near-singular Stage 1 design matrix and SE estimates may be " "less reliable.", UserWarning, stacklevel=2, ) - XtWX_10_dense = XtWX_10.toarray() - gamma_hat = np.linalg.lstsq(XtWX_10_dense, Xt1_WX2, rcond=None)[0] - if gamma_hat.ndim == 1: - gamma_hat = gamma_hat.reshape(-1, 1) - theta_exact = np.linalg.lstsq(XtWX_10_dense, np.asarray(rhs_fe).ravel(), rcond=None)[0] + XtWX_10_csc = XtWX_10.tocsc() + gamma_hat = _lsmr_certified_normal_solve(XtWX_10_csc, Xt1_WX2) + theta_exact = _lsmr_certified_normal_solve( + XtWX_10_csc, np.asarray(rhs_fe).ravel() + ).ravel() # Exact Stage-1 / Stage-2 residuals. The point-estimate path uses the # iterative alternating-projection FE solver (`_iterative_fe`), which diff --git a/diff_diff/two_stage_bootstrap.py b/diff_diff/two_stage_bootstrap.py index 4a2a0b2d..7039cc66 100644 --- a/diff_diff/two_stage_bootstrap.py +++ b/diff_diff/two_stage_bootstrap.py @@ -212,17 +212,21 @@ def _compute_cluster_S_scores( # of swallowing the error. warnings.warn( "TwoStageDiD bootstrap: sparse factorization of X_10' X_10 " - f"failed ({type(exc).__name__}); falling back to dense lstsq. " + f"failed ({type(exc).__name__}); falling back to sparse LSMR. " "This may indicate a rank-deficient or near-singular Stage 1 " "design matrix and bootstrap SE estimates may be less reliable.", UserWarning, stacklevel=2, ) - XtX_10_dense = XtX_10.toarray() - gamma_hat = np.linalg.lstsq(XtX_10_dense, Xt1_X2, rcond=None)[0] - if gamma_hat.ndim == 1: - gamma_hat = gamma_hat.reshape(-1, 1) - theta_exact = np.linalg.lstsq(XtX_10_dense, np.asarray(rhs_fe).ravel(), rcond=None)[0] + from diff_diff.two_stage import _lsmr_certified_normal_solve + + XtX_10_csc = XtX_10.tocsc() + gamma_hat = _lsmr_certified_normal_solve( + XtX_10_csc, Xt1_X2, context="TwoStageDiD bootstrap" + ) + theta_exact = _lsmr_certified_normal_solve( + XtX_10_csc, np.asarray(rhs_fe).ravel(), context="TwoStageDiD bootstrap" + ).ravel() # Exact Stage-1 / Stage-2 residuals (shared with the analytical variance) so # the bootstrap influence function uses the same exact residuals as @@ -376,6 +380,8 @@ def _run_bootstrap( resolved_survey: Optional[Any] = None, ) -> Optional[TwoStageBootstrapResults]: """Run multiplier bootstrap on GMM influence function.""" + from diff_diff.two_stage import _LSMRUnconvergedError as _TS_LSMRUnconverged + if self.n_bootstrap < 50: warnings.warn( f"n_bootstrap={self.n_bootstrap} is low. Consider n_bootstrap >= 199 " @@ -410,20 +416,25 @@ def _run_bootstrap( X_2_static = D.reshape(-1, 1) - S_static, bread_static, unique_clusters, _ = self._compute_cluster_S_scores( - df=df, - unit=unit, - time=time, - covariates=covariates, - omega_0_mask=omega_0_mask, - unit_fe=unit_fe, - time_fe=time_fe, - delta_hat=delta_hat, - kept_cov_mask=kept_cov_mask, - X_2=X_2_static, - cluster_ids=cluster_ids, - survey_weights=survey_weights, - ) + # Uncertified LSMR Stage-1 fallback -> degenerate (None/NaN) + # bootstrap contract rather than unverified scores. + try: + S_static, bread_static, unique_clusters, _ = self._compute_cluster_S_scores( + df=df, + unit=unit, + time=time, + covariates=covariates, + omega_0_mask=omega_0_mask, + unit_fe=unit_fe, + time_fe=time_fe, + delta_hat=delta_hat, + kept_cov_mask=kept_cov_mask, + X_2=X_2_static, + cluster_ids=cluster_ids, + survey_weights=survey_weights, + ) + except _TS_LSMRUnconverged: + return None n_clusters = len(unique_clusters) @@ -559,20 +570,25 @@ def _run_bootstrap( if h_int in horizon_to_col: X_2_es[i, horizon_to_col[h_int]] = 1.0 - S_es, bread_es, _, dropped_es = self._compute_cluster_S_scores( - df=df, - unit=unit, - time=time, - covariates=covariates, - omega_0_mask=omega_0_mask, - unit_fe=unit_fe, - time_fe=time_fe, - delta_hat=delta_hat, - kept_cov_mask=kept_cov_mask, - X_2=X_2_es, - cluster_ids=cluster_ids, - survey_weights=survey_weights, - ) + # Uncertified LSMR Stage-1 fallback -> degenerate (None/NaN) + # bootstrap contract rather than unverified scores. + try: + S_es, bread_es, _, dropped_es = self._compute_cluster_S_scores( + df=df, + unit=unit, + time=time, + covariates=covariates, + omega_0_mask=omega_0_mask, + unit_fe=unit_fe, + time_fe=time_fe, + delta_hat=delta_hat, + kept_cov_mask=kept_cov_mask, + X_2=X_2_es, + cluster_ids=cluster_ids, + survey_weights=survey_weights, + ) + except _TS_LSMRUnconverged: + return None # boot_coef_es: (B, k_es) boot_coef_es = np.dot(np.dot(all_weights, S_es), bread_es.T) @@ -626,20 +642,25 @@ def _run_bootstrap( if g in group_to_col: X_2_grp[i, group_to_col[g]] = 1.0 - S_grp, bread_grp, _, dropped_grp = self._compute_cluster_S_scores( - df=df, - unit=unit, - time=time, - covariates=covariates, - omega_0_mask=omega_0_mask, - unit_fe=unit_fe, - time_fe=time_fe, - delta_hat=delta_hat, - kept_cov_mask=kept_cov_mask, - X_2=X_2_grp, - cluster_ids=cluster_ids, - survey_weights=survey_weights, - ) + # Uncertified LSMR Stage-1 fallback -> degenerate (None/NaN) + # bootstrap contract rather than unverified scores. + try: + S_grp, bread_grp, _, dropped_grp = self._compute_cluster_S_scores( + df=df, + unit=unit, + time=time, + covariates=covariates, + omega_0_mask=omega_0_mask, + unit_fe=unit_fe, + time_fe=time_fe, + delta_hat=delta_hat, + kept_cov_mask=kept_cov_mask, + X_2=X_2_grp, + cluster_ids=cluster_ids, + survey_weights=survey_weights, + ) + except _TS_LSMRUnconverged: + return None boot_coef_grp = np.dot(np.dot(all_weights, S_grp), bread_grp.T) # NaN any dropped (unidentified) group coefficient (via the explicit diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 743315a1..21471e51 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -1722,7 +1722,7 @@ Our implementation uses multiplier bootstrap on the GMM influence function: clus - **Note:** The Stage-1 iterative FE solver (`_iterative_fe`) routes through the shared bincount Gauss-Seidel helper `diff_diff.utils._iterative_fe_solve`, and the covariate within-transformation routes through the shared MAP engine `diff_diff.utils.demean_by_groups` (factorize-once + `np.bincount`, optional Rust kernel; group order `[time, unit]` preserving the historical time-then-unit sweep) — the same convergence contract, accumulation-order numerics (~1e-10 vs the pre-3.7 pandas loops, not bit-for-bit), and `max_iter=10_000` budget documented under "Absorbed Fixed Effects with Survey Weights". Both surfaces emit `UserWarning` via `diff_diff.utils.warn_if_not_converged` when `max_iter` exhausts without reaching `tol` (the demean warning now carries the shared-engine label naming the affected variables rather than the estimator name). Silent return of the current iterate was classified as a silent failure under the Phase 2 audit and replaced with an explicit signal to match the logistic/Poisson IRLS pattern in `linalg.py`. - **Note:** Zero-total-weight groups (e.g. whole PSUs zeroed by JK1/BRR replicate weights, which reach Stage 1 unmasked — `keep_mask` only drops always-treated units): a unit/period whose observations ALL carry zero weight surfaces as `NaN` FE (key retained for the rank-condition membership check; matches the SpilloverDiD `_iterative_fe_subset` REGISTRY contract — never a silent finite `0.0`), and the shared demean engine's inert-row guard leaves those rows un-demeaned instead of NaN-poisoning the column. Before 3.7 the pandas loops divided 0/0 there: the covariate replicate path NaN-poisoned `y_dm`/`X_dm`, failed EVERY replicate refit inside `solve_ols(check_finite=True)`, and returned NaN SEs after a non-convergence warning storm. It now produces finite replicate SEs. `_mask_nan_ytilde`'s "non-finite imputed outcomes" `UserWarning` is suppressed (via `warn_nan=False`) ONLY inside the replicate-refit closures, where NaN FE for zeroed PSUs is expected mechanics — the main-fit warning is unchanged. - **Note:** When the Stage-2 bread `X'_2 W X_2` is singular, both the analytical TSL variance (`two_stage.py`) and the multiplier-bootstrap bread (`two_stage_bootstrap.py`) now emit a `UserWarning` before falling back to `np.linalg.lstsq`. Previously this fallback was silent. Sibling of axis-A finding #17 in the Phase 2 silent-failures audit; surfaced by the repo-wide lstsq-fallback pattern grep that accompanied the StaggeredTripleDifference fix. -- **Note:** The GMM sandwich and bootstrap paths both use `scipy.sparse.linalg.factorized` for the Stage 1 normal-equations solve `(X'_{10} W X_{10}) gamma = X'_1 W X_2` and fall back to dense `lstsq` when the sparse factorization raises `RuntimeError` on a near-singular matrix. Both fallback sites emit a `UserWarning` (silent-failure audit axis C) so callers know SE estimates came from the degraded path rather than the fast sparse path. +- **Note:** The GMM sandwich and bootstrap paths both use `scipy.sparse.linalg.factorized` for the Stage 1 normal-equations solve `(X'_{10} W X_{10}) gamma = X'_1 W X_2` and fall back to a **certified sparse LSMR** solve when the sparse factorization raises `RuntimeError` on a near-singular matrix (2026-07; previously dense `lstsq(toarray())`, an `O((U+T+K)²)` materialization and OOM risk on large panels — the pattern the ImputationDiD LSMR fix closed). Solver choice cannot change the output: least-squares solutions differ only by `null(X'X) = null(X_10)` components (weighted: `null(W^{1/2}X_10)`; zero-weight rows are inert in every weighted consumer), and every `gamma_hat` consumer is an `X_10`-range functional — `Psi = X_10 gamma`, the GMM score correction `c_g' gamma` with `c_g = X_{10,g}' eps_{10,g}` in `rowspace(X_10)` — so its null component annihilates. The bootstrap's `X_1_sparse @ theta_exact` consumer evaluates theta on treated rows where annihilation does not apply; parity there holds because both dense `lstsq` (SVD) and LSMR return the MIN-NORM least-squares solution, so the solvers agree on the whole vector to iterative tolerance (helper-level dense-lstsq-oracle parity test on a singular Gram + fit-level singular-design V/SE parity test). Convergence is validated fail-closed: `istop` in `{0,1,2,4,5}` certified, one uncapped-conlim retry, then `_LSMRUnconvergedError` — the analytical variance boundary reports NaN vcov/SE (raising matters: the GMM score `nan_to_num` would otherwise launder NaN scores into zeros and a finite, wrong variance) and the bootstrap boundary degenerates to the established `None` contract. All fallback sites (analytical unweighted + weighted, SpilloverDiD Wave-D meat, bootstrap) emit a `UserWarning` (silent-failure audit axis C). - **Note:** The GMM sandwich re-solves the Stage-1 unit+time fixed effects **exactly** (sparse OLS reusing the `scipy.sparse.linalg.factorized` factorization of `(X'_{10} W X_{10})` already computed for `gamma_hat`), rather than reusing the iterative alternating-projection FE (`_iterative_fe`) that produces the point estimate. The iterative solver converges only to ~1e-7 on unbalanced *untreated* panels — negligible for the ATT, but enough to perturb the variance by ~1% relative to the analytical GMM sandwich. The exact re-solve makes the analytical GMM SE match R `did2s` to ~1e-7 (`tests/test_methodology_two_stage.py::TestTwoStageDiDParityR`), mirroring ImputationDiD's exact-sparse variance path (`imputation.py` `_build_A_sparse`). It also required adding an **intercept column** to `_build_fe_design` so the first-stage column space spans the constant (the grand mean): the prior intercept-free `[unit_1.., time_1..]` layout (drop first unit + first time, no intercept) silently omitted the grand mean, which the exact residual is first-order sensitive to (the iterative point-estimate solver absorbs the grand mean into its mean-based FE, so the point estimate was unaffected). Obs whose unit or time FE are unidentified (NaN; rank-deficient / Proposition-5) fall back to the iterative residual, so those edge cases are unchanged. The reported `overall_att` still uses the iterative FE (preserving point-estimate equivalence with ImputationDiD at 1e-10); only the variance uses the exact residuals. - **Note:** `vcov_type` is permanently narrow to `{"hc1"}` (Phase 1b threading). TwoStageDiD's variance is the Gardner (2022) two-stage GMM cluster-sandwich `V = (X'_2 W X_2)^{-1} (S' S) (X'_2 W X_2)^{-1}` with the per-cluster GMM-corrected score `S_g = gamma_hat' c_g - X'_{2g} eps_{2g}`. Analytical-sandwich families `{classical, hc2, hc2_bm}` are rejected at `__init__`/`fit()`: the GMM-corrected meat folds first-stage FE estimation uncertainty into the score via the `gamma_hat' c_g` term, so there is no single hat matrix spanning both stages on which HC2 leverage or Bell-McCaffrey Satterthwaite DOF can be defined, and the Gardner first-stage correction has not been derived for the leverage-corrected or homoskedastic meat structures (no reference implementation — `clubSandwich` covers single-equation WLS/OLS CR2, not two-stage GMM; mirrors the SpilloverDiD `vcov_type="classical"` rejection). `cluster=