From c599aeb27bf18b8e0af16576b159aa5c28c0e75a Mon Sep 17 00:00:00 2001 From: igerber Date: Tue, 7 Jul 2026 10:30:01 -0400 Subject: [PATCH] perf(linalg): per-cell solver fast paths - IRLS Cholesky inner solve + rank-detection certification Two pure-Python fast paths in the shared solvers (CallawaySantAnna covariate fits and every estimator routing through them): 1. solve_logit's IRLS inner step - previously a full gelsd SVD on the tall weighted design per iteration - now solves equilibrated normal equations by Cholesky under an explicit dpocon reciprocal-condition guard (1e-6), falling back to the exact legacy lstsq line for any iteration whose normal matrix cannot be certified well-conditioned (working weights can crush a column's effective scale on near-separated subgroups, so full column rank pre-loop does not imply a chol-safe normal matrix). Columns equilibrated ONCE; IRLS state, tol, and warnings stay in the raw basis. Fallback count exposed via diagnostics_out["irls_chol_fallback_iters"]. 2. _detect_rank_deficiency stage-0 certification - the always-run pivoted QR on the tall design is short-circuited by a Gram/eigvalsh certification at the _rank_guarded_inv 1e-10 Gram threshold (two orders stricter than the 1e-7 QR boundary, so it never reports full rank where the QR path would report a deficiency); n < k, non-finite, looser-rcond, and uncertifiable designs fall through to the existing two-stage QR verbatim (all drop decisions, pivots, and rank counts unchanged - incl. dCDH residual-df and the rust solve_ols routing boolean, proven by arm instrumentation). Measured (3-rep medians, CS 100k units x 20p = 2M rows, 95 cells): dr cov10 1.53->1.03s (1.49x) | cov20 2.64->1.61s (1.64x) dr cov40 5.69->3.07s (1.85x) | ipw cov20 1.79->0.90s (1.98x) survey dr cov20 2.73->1.68s (1.63x); python-backend pairs equal or better. Stages at cov40: IRLS solve 6.8x, rank detection 16.5x. Deltas: overall ATT/SE exactly 0, per-cell max ~7e-15; golden ipw R-parity and dCDH bit-identity baselines green unmodified; fallback 0% / certification 100% on all bench scenarios; maxrss flat. New tests: TestIRLSCholeskyFastPath (property parity vs legacy IRLS at the tol-bounded 1e-8 gate, always-raise byte-identity lock, natural guard-trip warning equivalence, convergence-iteration semantics) and TestRankDetectionStage0Certification (zero-QR cert spy, declined paths incl. the cond~1e6 cert-vs-QR boundary, Kahan characterization, rust solve_ols dispatch composition) in tests/test_linalg.py. REGISTRY notes for both fast paths; CHANGELOG + performance-plan tables; TODO rows for the RC pscore cache and rust solve_ols follow-ups. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01X4AzrFUMqJxcUumSH31mSr --- CHANGELOG.md | 21 ++ TODO.md | 2 + diff_diff/linalg.py | 113 ++++++++++- docs/methodology/REGISTRY.md | 18 +- docs/performance-plan.md | 54 +++++ tests/test_linalg.py | 378 +++++++++++++++++++++++++++++++++++ 6 files changed, 576 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9d2161da..7dbc1d46 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -247,6 +247,27 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 supersedes it. ### Changed +- **Per-cell solver fast paths for covariate fits (CallawaySantAnna and every + estimator routing through the shared solvers).** Two pure-Python changes in + `diff_diff/linalg.py`: (1) `solve_logit`'s IRLS inner step — previously a full + SVD (`np.linalg.lstsq`) on the tall weighted design per iteration — now solves + equilibrated normal equations by Cholesky under an explicit + reciprocal-condition guard (LAPACK `dpocon`), falling back to the exact legacy + solve for any iteration whose normal matrix cannot be certified + well-conditioned; (2) `_detect_rank_deficiency` short-circuits the common + full-rank case with a Gram/eigenvalue certification two orders stricter than + the pivoted-QR boundary, so all rank counts, drop decisions, and pivot + selections on deficient or uncertifiable designs are unchanged. Measured + (3-rep medians, CallawaySantAnna, 100k units x 20 periods = 2M rows, 95 cells, + `aggregate="all"`): dr fits 1.5s → 1.0s at 10 covariates (1.49x), 2.6s → 1.6s + at 20 (1.64x), 5.7s → 3.1s at 40 (1.85x); ipw 1.8s → 0.9s (2.0x); + survey-weighted dr 2.7s → 1.7s (1.6x); pure-Python backend gains are equal or + larger. The solver stages themselves: IRLS solve 6.8x, rank detection 16.5x + (at 40 covariates). Estimates are unchanged beyond machine precision (overall + ATT/SE deltas exactly 0; per-cell max ~7e-15; the R-golden ipw SE parity at + 1e-6 abs and the dCDH bit-identity baselines hold unmodified), the Cholesky + fallback fires on 0% of healthy fits (its near-separation trip path is locked + by tests), and memory is flat. - **`CallawaySantAnna` no-covariate `estimation_method="dr"` per-cell SE is now influence-function-based** (`sqrt(sum(phi^2))`), matching the `reg`/`ipw` branches and R's `DRDID::drdid_panel`. It was the last per-cell SE on the ddof=1 plug-in diff --git a/TODO.md b/TODO.md index 80ed68d4..b2c34a3e 100644 --- a/TODO.md +++ b/TODO.md @@ -46,6 +46,8 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| | `EfficientDiD` conditional path: post-Cholesky-dispatch (2026-07, 10k fit 22.7s) the ridge-solve stage is 3.7s, of which the Rust kernel is only ~1.6s — the remainder is shared Python prep inside `_ridge_solve_weights`: the `zero_mask` full-stack abs scan plus the `omega_stack[rest]` fancy-index copy (~28 GB of extra memory traffic per 10k fit). The copy is avoidable when no rows are zero-masked (`rest.size == n`, the common case) — value-identical on both backends since `om` is never mutated. After that, the largest O(n) stage is the sieve/nuisance construction outside the tiled pass (~9s at 10k). | `efficient_did_covariates.py::_ridge_solve_weights` | CS-scaling | Quick | 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 | | 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 HC2: the Rust path only supports HC1; HC2 and CR2 Bell-McCaffrey fall through to NumPy. Noticeable for large-`n` fits. | `rust/src/linalg.rs` | Phase 1a | Mid | Low | diff --git a/diff_diff/linalg.py b/diff_diff/linalg.py index 6647d370..c779caf5 100644 --- a/diff_diff/linalg.py +++ b/diff_diff/linalg.py @@ -39,8 +39,9 @@ import numpy as np import pandas as pd from scipy import stats +from scipy.linalg import cho_factor, cho_solve, qr from scipy.linalg import lstsq as scipy_lstsq -from scipy.linalg import qr +from scipy.linalg.lapack import dpocon # Import Rust backend if available (from _backend to avoid circular imports) from diff_diff._backend import ( @@ -132,9 +133,11 @@ def _detect_rank_deficiency( Indices of columns that are linearly dependent (should be dropped). Empty if matrix is full rank. pivot : ndarray of int - Column permutation from QR decomposition. + Column permutation from QR decomposition. For a full-rank result the + pivot carries no information (no caller consumes it; the stage-0 + certification below returns a trivial ``arange`` pivot). """ - _, k = X.shape + n, k = X.shape if k == 0: return 0, np.array([], dtype=int), np.array([], dtype=int) @@ -142,6 +145,42 @@ def _detect_rank_deficiency( if rcond is None: rcond = 1e-07 + # Stage 0 — Gram full-rank CERTIFICATION (perf fast path; decisions never + # change). The Gram is built directly from X (no equilibrated copy of the + # tall matrix); diag(G) holds the squared column 2-norms, so equilibrating + # G symmetrically by sqrt(diag(G)) is exactly column equilibration of X — + # the `_rank_guarded_inv` convention. Certification threshold 1e-10 is the + # documented Gram constant from `_rank_guarded_inv` (a Gram squares the + # condition number of X, so 1e-10 on eigenvalues ~ cond(X_eq) < 1e5, two + # orders STRICTER than the 1e-7 QR full-rank boundary): stage 0 never + # reports full rank where the two-stage QR below would report a + # deficiency, it only declines and falls through. (A pathological + # Kahan-type matrix could in principle make pivoted-QR R-diagonals + # undershoot the singular values enough to open a gap; real DiD designs — + # dummies plus covariates — do not have that structure, and the + # characterization test documents it.) Skipped when n < k (always + # deficient — also keeps the sole pivot consumer, staggered.py's + # underdetermined pair solve, structurally on the QR path) and when a + # caller passes a LOOSER-than-default rcond (no caller does today; the + # stricter-than-QR guarantee above assumes rcond <= 1e-7). Non-finite + # entries poison diag(G), decline certification, and fall through so + # scipy's qr raises ValueError on NaN/Inf exactly as before. + if n >= k and rcond <= 1e-07: + gram = X.T @ X + diag = np.diag(gram) + if np.all(np.isfinite(diag)) and np.all(diag > 0): + scales = np.sqrt(diag) + gram_eq = gram / scales[:, None] / scales[None, :] + eigvals = np.linalg.eigvalsh(gram_eq) + eig_min, eig_max = eigvals[0], eigvals[-1] + if ( + np.isfinite(eig_min) + and np.isfinite(eig_max) + and eig_max > 0.0 + and eig_min > 1e-10 * eig_max + ): + return k, np.array([], dtype=int), np.arange(k, dtype=int) + def _rank_and_pivot(M: np.ndarray) -> Tuple[int, np.ndarray]: # Pivoted QR: M @ P = Q @ R. The rank threshold is anchored to the # largest pivot diagonal |R[0,0]| (decreasing after pivoting). @@ -2892,6 +2931,12 @@ def _compute_robust_vcov_numpy( _LOGIT_SEPARATION_PROB_THRESHOLD = 1e-5 _DEFAULT_EPV_THRESHOLD = 10 +# Reciprocal-condition guard for the IRLS normal-equations Cholesky fast path +# (solve_logit): cond(G) <= 1e6 bounds the Cholesky forward error at +# ~eps * cond ~ 2e-10 relative, consistent with the tol-bounded parity budget; +# anything worse falls back to the legacy tall-matrix lstsq for that iteration. +_IRLS_CHOL_RCOND_GUARD = 1e-6 + def solve_logit( X: np.ndarray, @@ -2947,7 +2992,10 @@ def solve_logit( users identify which logit estimation triggered the warning. diagnostics_out : dict, optional If provided, populated with EPV diagnostic info: - ``{"epv": float, "n_events": int, "k": int, "is_low": bool}``. + ``{"epv": float, "n_events": int, "k": int, "is_low": bool}``, plus + ``"irls_chol_fallback_iters"`` - the number of IRLS iterations whose + normal-equations Cholesky fast path fell back to the legacy lstsq + solve (0 on well-conditioned fits). Returns ------- @@ -3097,7 +3145,31 @@ def solve_logit( raise ValueError(msg) warnings.warn(msg, UserWarning, stacklevel=2) - # IRLS (Fisher scoring) + # IRLS (Fisher scoring). Each weighted-least-squares step is solved via + # EQUILIBRATED normal equations + Cholesky with an explicit condition + # guard, falling back to the legacy tall-matrix lstsq (gelsd SVD) for any + # iteration whose normal matrix cannot be certified well-conditioned. + # Context: the OR path deliberately REMOVED a cho_solve(X'X) fast path + # because it was NOT scale-equilibrated (see the covariate-reg notes in + # staggered.py around `_equilibrated_lstsq`); this path differs on + # exactly that axis - (1) columns are equilibrated to unit 2-norm ONCE + # (a fixed reparameterization: X_eq @ beta_eq == X @ beta algebraically, + # so probabilities are unchanged and beta is unscaled per iteration); + # (2) cho_factor alone can SUCCEED with a garbage solution when cond(G) + # exceeds ~1e10, so the guard is a dpocon reciprocal-condition estimate, + # not just the factorization succeeding - working weights can crush a + # column's effective scale (a dummy on a near-separated subgroup has + # w_irls ~ 1e-10 on its support), so full column rank pre-loop does NOT + # imply a well-conditioned G; (3) the fallback reproduces the legacy + # computation for that iteration on the raw basis. IRLS state (beta, + # convergence tol, warnings) stays in the RAW basis: a scaled-basis tol + # would be ~sqrt(n)x tighter for every fit (the intercept column alone + # has 2-norm sqrt(n)), and the separation check below reads raw beta. + irls_col_norms = np.sqrt(np.einsum("ij,ij->j", X_solve, X_solve)) + irls_safe_norms = np.where(irls_col_norms > 0, irls_col_norms, 1.0) + X_eq = X_solve / irls_safe_norms + chol_fallback_iters = 0 + beta_solve = np.zeros(X_solve.shape[1]) converged = False @@ -3120,9 +3192,33 @@ def solve_logit( # Weighted least squares: solve (X'WX) beta = X'Wz sqrt_w = np.sqrt(w_total) - Xw = X_solve * sqrt_w[:, None] zw = z * sqrt_w - beta_new, _, _, _ = np.linalg.lstsq(Xw, zw, rcond=None) + beta_new = None + Xw_eq = X_eq * sqrt_w[:, None] + gram = Xw_eq.T @ Xw_eq + # 1-norm BEFORE factorization (dpocon contract); G is symmetric so + # the max absolute column sum is the 1-norm. + anorm = float(np.max(np.sum(np.abs(gram), axis=0))) + if np.isfinite(anorm): + try: + chol = cho_factor(gram) + except np.linalg.LinAlgError: + chol = None + if chol is not None: + # cho_factor default lower=False pairs with dpocon's default + # uplo='U' (dpocon has no `lower=` kwarg). + rcond_gram, pocon_info = dpocon(chol[0], anorm) + if ( + pocon_info == 0 + and np.isfinite(rcond_gram) + and rcond_gram > _IRLS_CHOL_RCOND_GUARD + ): + beta_new = cho_solve(chol, Xw_eq.T @ zw) / irls_safe_norms + if beta_new is None: + # Guarded fallback: byte-identical to the pre-fast-path solve. + chol_fallback_iters += 1 + Xw = X_solve * sqrt_w[:, None] + beta_new, _, _, _ = np.linalg.lstsq(Xw, zw, rcond=None) # Check convergence if np.max(np.abs(beta_new - beta_solve)) < tol: @@ -3131,6 +3227,9 @@ def solve_logit( break beta_solve = beta_new + if diagnostics_out is not None: + diagnostics_out["irls_chol_fallback_iters"] = chol_fallback_iters + # Final predicted probabilities eta_final = X_solve @ beta_solve eta_final = np.clip(eta_final, -500, 500) diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 84461972..5839af27 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -78,7 +78,7 @@ where τ is the ATT. - Rank-deficient design matrix (collinearity): warns and sets NA for dropped coefficients (R-style, matches `lm()`) - Tolerance: `1e-07` (matches R's `qr()` default), relative to largest diagonal element of R in QR decomposition - Controllable via `rank_deficient_action` parameter: "warn" (default), "error", or "silent" - - **Note (scale invariance):** the shared `diff_diff/linalg.py` rank check re-detects on column-equilibrated (unit 2-norm) columns when the raw pass reports a deficiency (adopting the higher equilibrated rank and its scale-corrected pivot selection only when a large-scale column inflated the threshold; genuine collinearity keeps the raw selection), and the least-squares solve equilibrates columns then unscales the coefficients — so rank detection and the fit are invariant to per-column scaling; for a well-scaled collinear design the dropped column is unchanged, while a scale-induced under-count adopts the scale-corrected equilibrated selection (which may differ from the raw choice but is guaranteed to retain an identified subset). This covers every covariate outcome-regression fit routed through `solve_ols` — DiD, TwoWayFixedEffects, MultiPeriodDiD, ImputationDiD, TwoStageDiD, TripleDifference, and (as of the OR scale-equilibration change) CallawaySantAnna's `_compute_all_att_gt_covariate_reg` / `_doubly_robust` and StaggeredTripleDifference's `_compute_or` — see the CallawaySantAnna rank-deficiency Note's scope statement. + - **Note (scale invariance):** the shared `diff_diff/linalg.py` rank check re-detects on column-equilibrated (unit 2-norm) columns when the raw pass reports a deficiency (adopting the higher equilibrated rank and its scale-corrected pivot selection only when a large-scale column inflated the threshold; genuine collinearity keeps the raw selection), and the least-squares solve equilibrates columns then unscales the coefficients — so rank detection and the fit are invariant to per-column scaling; for a well-scaled collinear design the dropped column is unchanged, while a scale-induced under-count adopts the scale-corrected equilibrated selection (which may differ from the raw choice but is guaranteed to retain an identified subset). This covers every covariate outcome-regression fit routed through `solve_ols` — DiD, TwoWayFixedEffects, MultiPeriodDiD, ImputationDiD, TwoStageDiD, TripleDifference, and (as of the OR scale-equilibration change) CallawaySantAnna's `_compute_all_att_gt_covariate_reg` / `_doubly_robust` and StaggeredTripleDifference's `_compute_or` — see the CallawaySantAnna rank-deficiency Note's scope statement. **Certification fast path (2026-07):** the common full-rank case is short-circuited by a Gram/eigenvalue certification on the equilibrated Gram at the documented `_rank_guarded_inv` 1e-10 Gram threshold — two orders stricter than the QR full-rank boundary, so it never reports full rank where the pivoted-QR path would report a deficiency; deficient or uncertifiable designs (including all n < k and non-finite inputs) fall through to the pivoted-QR path verbatim, leaving every drop decision, pivot selection, and rank count unchanged. Certification never SELECTS columns (it only confirms none need dropping), so the `_rank_guarded_inv` note about equilibrated-Gram column selection being reserved for the IF inverse still holds. - **Note (covariate-name collision guard):** `DifferenceInDifferences`, `MultiPeriodDiD`, and `TwoWayFixedEffects` build their `coefficients` dict by zipping a variable-name list (structural terms + user covariate column names, appended verbatim) with the coefficient vector. A covariate whose name equals a reserved structural name — `const`, the treatment/time column names, the `{treatment}:{time}` interaction (DiD), the `period_{p}` dummies / `{treatment}:period_{p}` interactions (MultiPeriodDiD), `ATT` (TwoWayFixedEffects), any fixed-effect / unit / time dummy name, or an internal `_`-prefixed working column (`_treat_time`, `_did_treatment`, `_treatment_post`) — would silently overwrite that structural coefficient (dict last-write-wins) with no error. `fit()` now calls `utils.validate_covariate_names()` before building the design, raising `ValueError` on such a collision (case-sensitive, so `Const` is allowed) and on duplicate covariate names. Reserved fixed-effect/unit/time dummy names are taken from the same `pd.get_dummies(..., drop_first=True)` call used to build them (exact match, including for `Categorical` columns). The TwoWayFixedEffects guard fires on both variance paths — the within-transform path exposes only `{"ATT": att}`, but a `_treatment_post`-named covariate would still clobber the internal interaction column. The influence-function estimators (CallawaySantAnna, etc.) key results by `(g, t)` and are unaffected. **Reference implementation(s):** @@ -272,7 +272,7 @@ where V is the VCV sub-matrix for post-treatment δ_e coefficients. not observed at all event times contribute to the periods they are present for. - Rank-deficient design matrix (collinearity): warns and sets NA for dropped coefficients (R-style, matches `lm()`) - - **Note (scale invariance):** shared `diff_diff/linalg.py` behavior — rank detection re-checks on column-equilibrated columns and the solve equilibrates/unscales, so detection and fit are invariant to per-column scaling. For a well-scaled collinear design the dropped column is unchanged; a scale-induced under-count adopts the scale-corrected equilibrated selection (which may differ from the raw choice but retains an identified subset). See the CallawaySantAnna rank-deficiency Note. + - **Note (scale invariance):** shared `diff_diff/linalg.py` behavior — rank detection re-checks on column-equilibrated columns and the solve equilibrates/unscales, so detection and fit are invariant to per-column scaling. For a well-scaled collinear design the dropped column is unchanged; a scale-induced under-count adopts the scale-corrected equilibrated selection (which may differ from the raw choice but retains an identified subset). The common full-rank case is short-circuited by a strictly-stricter Gram certification (2026-07) that leaves all drop decisions unchanged. See the CallawaySantAnna rank-deficiency Note. - **Note (covariate-name collision guard):** a covariate named like a reserved structural term (`const`, the treatment column, a `period_{p}` dummy, a `{treatment}:period_{p}` interaction, a fixed-effect dummy, or an internal `_did_*` column) — or a duplicate covariate name — raises `ValueError` (would otherwise silently overwrite a structural coefficient in the coef dict). See the DifferenceInDifferences "covariate-name collision guard" Note. - Average ATT (`avg_att`) is NA if any post-period effect is unidentified (R-style NA propagation) @@ -365,7 +365,7 @@ This matches the behavior of R's `fixest::feols()` with absorbed FE. - Treatment perfectly collinear with FE raises error with informative message listing dropped columns - Covariate collinearity emits warning but estimation continues (ATT still identified) - Rank-deficient design matrix: warns and sets NA for dropped coefficients (R-style, matches `lm()`) - - **Note (scale invariance):** shared `diff_diff/linalg.py` behavior — rank detection re-checks on column-equilibrated columns and the solve equilibrates/unscales, so detection and fit are invariant to per-column scaling. For a well-scaled collinear design the dropped column is unchanged; a scale-induced under-count adopts the scale-corrected equilibrated selection (which may differ from the raw choice but retains an identified subset). See the CallawaySantAnna rank-deficiency Note. + - **Note (scale invariance):** shared `diff_diff/linalg.py` behavior — rank detection re-checks on column-equilibrated columns and the solve equilibrates/unscales, so detection and fit are invariant to per-column scaling. For a well-scaled collinear design the dropped column is unchanged; a scale-induced under-count adopts the scale-corrected equilibrated selection (which may differ from the raw choice but retains an identified subset). The common full-rank case is short-circuited by a strictly-stricter Gram certification (2026-07) that leaves all drop decisions unchanged. See the CallawaySantAnna rank-deficiency Note. - **Note (covariate-name collision guard):** a covariate named `const`, `ATT`, `_treatment_post`, or a unit/time fixed-effect dummy name — or a duplicate covariate name — raises `ValueError` on both variance paths (would otherwise silently overwrite a structural coefficient on the full-dummy HC2/HC2-BM path). See the DifferenceInDifferences "covariate-name collision guard" Note. - Unbalanced panels handled via proper demeaning - **Note (iterative demeaning):** the two-way within transformation @@ -729,6 +729,18 @@ The multiplier bootstrap uses random weights w_i with E[w]=0 and Var(w)=1: - Algorithm: IRLS (Fisher scoring), matching R's `glm(family=binomial)` default - **Note:** Uses IRLS (Fisher scoring) for propensity score estimation, consistent with R's `did::att_gt()` which uses `glm(family=binomial)` internally + - **Note (IRLS inner solver, 2026-07):** each IRLS iteration's weighted + least-squares step is solved via equilibrated normal equations + Cholesky with + an explicit reciprocal-condition guard (LAPACK `dpocon`; guard 1e-6), falling + back to the exact legacy tall-matrix `lstsq` solve for any iteration whose + normal matrix cannot be certified well-conditioned (near-separation fits, where + working weights crush a column's effective scale). The IRLS ALGORITHM — working + weights `mu*(1-mu)`, working response, `tol=1e-8` on the raw-basis coefficient + change, `max_iter=25` — is unchanged; both inner solvers converge to the same + MLE, and measured coefficients move at the ~1e-15 level (tol-bounded worst case + ~1e-8 if an iteration count ever shifts by one; the R-golden ipw SE pin at + 1e-6 abs is unaffected). The per-fit fallback count is exposed via + `diagnostics_out["irls_chol_fallback_iters"]` (0 on well-conditioned fits). - Near-separation detection: Warns when predicted probabilities are within 1e-5 of 0 or 1, or when IRLS fails to converge - Trimming: Propensity scores clipped to `[pscore_trim, 1-pscore_trim]` (default diff --git a/docs/performance-plan.md b/docs/performance-plan.md index 8ff210eb..43a11a40 100644 --- a/docs/performance-plan.md +++ b/docs/performance-plan.md @@ -183,6 +183,60 @@ precision; analytical SEs to 5 significant digits; bootstrap SEs within ~1% (Mon --- +## CS per-cell solver fast paths (Phase 3, 2026-07) + +Fresh profile on main 49bbde68 (CallawaySantAnna dr, 100k units x 20 periods = +2M rows, 95 cells) after Phase 1 removed the aggregation dominance: + +| item | cov10 (fit 1.56s) | cov20 (2.71s) | cov40 (5.68s) | +|---|---|---|---| +| `np.linalg.lstsq` in solve_logit IRLS (184 calls) | 0.39s | 0.81s | 1.99s | +| `_detect_rank_deficiency` pivoted QR (141 calls) | 0.15s | 0.31s | 0.84s | +| rust solve_ols SVD (out of scope, TODO row) | 0.23s | 0.46s | 1.03s | + +Two pure-Python fast paths in `diff_diff/linalg.py`: + +1. **IRLS inner solve**: per-iteration gelsd SVD on the tall weighted design + -> equilibrated normal equations + Cholesky, guarded by a LAPACK `dpocon` + reciprocal-condition estimate (guard 1e-6; cho_factor alone can succeed + with garbage on cond(G) ~1e10-1e16, and working weights can crush a + column's effective scale on near-separated subgroups). Any uncertified + iteration falls back to the exact legacy lstsq line. Columns are + equilibrated ONCE (the repo removed a prior un-equilibrated cho_solve OR + fast path for scale sensitivity - this addresses that removal reason); + IRLS state stays in the raw basis (a scaled-basis tol would be ~sqrt(n)x + tighter for every fit). Fallback count exposed via + `diagnostics_out["irls_chol_fallback_iters"]`. +2. **Rank-detection stage-0 certification**: the always-run pivoted QR on the + tall design is short-circuited by a Gram/eigvalsh certification on the + equilibrated Gram at the `_rank_guarded_inv` 1e-10 threshold (two orders + stricter than the QR boundary; never certifies what QR would call + deficient); n < k, non-finite, looser-rcond, and uncertifiable designs + fall through to the existing two-stage QR verbatim. + +Measured arms (3-rep medians; BEFORE = pristine main 49bbde68 via +baseline-worktree PYTHONPATH, .so rebuilt at that SHA; BLAS threads pinned; +`.bench-local/cs_solver_arms_results.jsonl`): + +| scenario (@100k x 20p) | rust BEFORE -> AFTER | speedup | python pair | +|---|---|---|---| +| cov10_dr | 1.53 -> 1.03s | 1.49x | 1.49x | +| cov20_dr | 2.64 -> 1.61s | 1.64x | 1.65x | +| cov40_dr | 5.69 -> 3.07s | 1.85x | 1.95x | +| cov20_ipw | 1.79 -> 0.90s | 1.98x | 2.03x | +| cov20_survey_dr | 2.73 -> 1.68s | 1.63x | 1.60x | + +Component stages at cov40: solve_logit 2.54 -> 0.37s (6.8x), +_detect_rank_deficiency 0.84 -> 0.05s (16.5x). Instrumented during arms: +Cholesky fallback 0 iterations on every scenario/backend/arm; certification +rate 100% (141 -> 0 pivoted QRs per dr fit, 46 -> 0 per ipw fit); +rust-dispatch counts identical between arms (the solve_ols routing boolean is +provably unchanged). Deltas: overall ATT/SE exactly 0, per-cell max ~7e-15, +maxrss flat. After this the top per-cell item at cov40 is the rust solve_ols +SVD (1.04s; follow-up TODO row), then pandas frame prep. + +--- + ## CS multiplier-bootstrap fused tiled-GEMM rewrite (v3.7 Phase 2, 2026-07) Stage-level instrumentation of `_run_multiplier_bootstrap` at the 40p × 10c × 100k-unit diff --git a/tests/test_linalg.py b/tests/test_linalg.py index c07c90d3..c19cebed 100644 --- a/tests/test_linalg.py +++ b/tests/test_linalg.py @@ -1,15 +1,19 @@ """Tests for the unified linear algebra backend.""" +import warnings + import numpy as np import pandas as pd import pytest +from diff_diff import HAS_RUST_BACKEND from diff_diff.linalg import ( InferenceResult, LinearRegression, _rank_guarded_inv, compute_r_squared, compute_robust_vcov, + solve_logit, solve_ols, solve_poisson, ) @@ -1671,6 +1675,184 @@ def test_equilibrated_lstsq_f_order_overwrite_matches_reference(self): np.testing.assert_array_equal(coef, ref) +class TestRankDetectionStage0Certification: + """Stage-0 Gram/eigvalsh full-rank certification in _detect_rank_deficiency. + + Contract: certification is a pure fast path - it either certifies full + rank (returning (k, empty, arange)) on designs the two-stage QR would + ALSO call full rank, or declines and falls through to the existing QR + path verbatim. The cert threshold 1e-10 on equilibrated-Gram eigenvalues + (~ cond(X_eq) < 1e5) is two orders stricter than the QR full-rank + boundary (rcond=1e-7 on R-diagonals), so decisions on deficient designs + never change. + """ + + @staticmethod + def _qr_call_counter(monkeypatch): + import diff_diff.linalg as lmod + + calls = {"n": 0} + orig = lmod.qr + + def counting(*args, **kwargs): + calls["n"] += 1 + return orig(*args, **kwargs) + + monkeypatch.setattr(lmod, "qr", counting) + return calls + + def test_certified_full_rank_runs_zero_qr(self, monkeypatch): + """Well-conditioned full-rank design - including a 1e8-scale + independent column (equilibration makes it benign) - certifies with + ZERO pivoted-QR calls (the perf lock) and the trivial-pivot + contract.""" + from diff_diff.linalg import _detect_rank_deficiency + + rng = np.random.default_rng(42) + X = rng.standard_normal((500, 6)) + X[:, 3] *= 1e8 + calls = self._qr_call_counter(monkeypatch) + rank, dropped, pivot = _detect_rank_deficiency(X) + assert calls["n"] == 0 + assert rank == 6 + assert dropped.shape == (0,) and np.issubdtype(dropped.dtype, np.integer) + np.testing.assert_array_equal(pivot, np.arange(6)) + assert np.issubdtype(pivot.dtype, np.integer) + + def test_declined_collinear_matches_legacy_qr_selection(self, monkeypatch): + """A genuinely collinear design declines certification (singular + equilibrated Gram) and the (rank, dropped, pivot) triple is the raw + pivoted-QR answer, exactly as before stage-0 existed.""" + from scipy.linalg import qr as scipy_qr + + from diff_diff.linalg import _detect_rank_deficiency + + rng = np.random.default_rng(7) + X = rng.standard_normal((120, 4)) + X[:, 3] = X[:, 0] + X[:, 1] + calls = self._qr_call_counter(monkeypatch) + rank, dropped, pivot = _detect_rank_deficiency(X) + assert calls["n"] >= 1 # fell through to the QR path + assert rank == 3 + # reference: the raw pivoted-QR drop selection (existing contract) + r_ref, piv_ref = scipy_qr(X, mode="r", pivoting=True) + np.testing.assert_array_equal(dropped, np.sort(piv_ref[3:])) + np.testing.assert_array_equal(pivot, piv_ref) + + def test_declined_n_less_than_k_structural(self, monkeypatch): + """n < k skips certification structurally (always deficient; the + sole pivot consumer in staggered.py lives behind this shape).""" + from diff_diff.linalg import _detect_rank_deficiency + + rng = np.random.default_rng(3) + X = rng.standard_normal((3, 5)) + calls = self._qr_call_counter(monkeypatch) + rank, dropped, _pivot = _detect_rank_deficiency(X) + assert calls["n"] >= 1 + assert rank == 3 and len(dropped) == 2 + + def test_declined_zero_column(self): + """A zero column zeroes its Gram diagonal - certification declines + and the QR path drops it as before.""" + from diff_diff.linalg import _detect_rank_deficiency + + rng = np.random.default_rng(5) + X = rng.standard_normal((60, 3)) + X[:, 1] = 0.0 + rank, dropped, _ = _detect_rank_deficiency(X) + assert rank == 2 and 1 in dropped + + def test_nan_still_raises_value_error(self): + """Non-finite entries decline certification (they poison diag(G)) + and scipy's qr raises ValueError exactly as before.""" + from diff_diff.linalg import _detect_rank_deficiency + + rng = np.random.default_rng(9) + X = rng.standard_normal((50, 3)) + X[0, 0] = np.nan + with pytest.raises(ValueError): + _detect_rank_deficiency(X) + + def test_boundary_declines_cert_but_qr_full_rank(self, monkeypatch): + """cond(X_eq) ~ 1e6 sits BETWEEN the cert threshold (1e5) and the QR + full-rank boundary (~1e7): certification declines, stage-1 QR still + returns full rank - locks cert-strictly-stricter-than-QR ordering.""" + from diff_diff.linalg import _detect_rank_deficiency + + rng = np.random.default_rng(11) + X = rng.standard_normal((200, 3)) + # near-dependence BETWEEN columns (a small column NORM would be + # repaired by equilibration and correctly certified): sv ratio of the + # equilibrated design ~2e-6 -> Gram eig ratio ~4e-12 < 1e-10 declines + # cert, while the QR R-diagonal ratio ~2e-6 > 1e-7 stays full rank. + X[:, 2] = X[:, 0] + 3e-6 * rng.standard_normal(200) + calls = self._qr_call_counter(monkeypatch) + rank, dropped, _ = _detect_rank_deficiency(X) + assert calls["n"] >= 1 # cert declined + assert rank == 3 and len(dropped) == 0 + + def test_looser_rcond_skips_certification(self, monkeypatch): + """A caller-supplied rcond looser than 1e-7 disables stage-0 (the + stricter-than-QR guarantee only holds for rcond <= 1e-7); the QR + path answers with the caller's threshold as before.""" + from diff_diff.linalg import _detect_rank_deficiency + + rng = np.random.default_rng(13) + X = rng.standard_normal((100, 3)) + calls = self._qr_call_counter(monkeypatch) + rank, dropped, _ = _detect_rank_deficiency(X, rcond=1e-4) + assert calls["n"] >= 1 + assert rank == 3 and len(dropped) == 0 + + def test_kahan_characterization_cert_never_wrong(self): + """Kahan-type matrices are the theoretical gap between pivoted-QR + R-diagonals and singular values (the R-diagonal can undershoot + sigma_min by up to 2^(k-1) - reachable pathology, not bounded to + large k). Characterization, not a behavior pin: WHENEVER stage-0 + certifies, the SVD-based numerical rank agrees it is full rank at + the QR threshold - i.e. certification can disagree with legacy QR + only by being MORE correct (repairing a QR false-drop), never by + certifying a genuinely deficient design.""" + from diff_diff.linalg import _detect_rank_deficiency + + for k, theta in ((8, 0.6), (20, 0.35), (30, 0.28)): + s, c = np.sin(theta), np.cos(theta) + K = np.zeros((k, k)) + for i in range(k): + K[i, i] = s**i + K[i, i + 1 :] = -c * s**i + rank, dropped, _ = _detect_rank_deficiency(K) + if len(dropped) == 0 and rank == k: + sv = np.linalg.svd(K / np.sqrt(np.einsum("ij,ij->j", K, K)), compute_uv=False) + assert sv[-1] > 1e-7 * sv[0], ( + f"cert claimed full rank on a genuinely deficient Kahan " + f"matrix (k={k}, theta={theta})" + ) + + @pytest.mark.skipif(not HAS_RUST_BACKEND, reason="Rust backend not available") + def test_certified_design_still_dispatches_solve_ols_to_rust(self, monkeypatch): + """Composition lock: the stage-0-certified full-rank answer feeds the + solve_ols routing boolean, which must still dispatch to the Rust + solver (unweighted hc1 path).""" + import diff_diff.linalg as lmod + from diff_diff.linalg import solve_ols + + calls = {"n": 0} + orig = lmod._solve_ols_rust + + def counting(*args, **kwargs): + calls["n"] += 1 + return orig(*args, **kwargs) + + monkeypatch.setattr(lmod, "_solve_ols_rust", counting) + rng = np.random.default_rng(17) + X = rng.standard_normal((300, 5)) + y = X @ np.arange(1.0, 6.0) + rng.standard_normal(300) + coeffs, _, _ = solve_ols(X, y) + assert calls["n"] == 1 + assert np.all(np.isfinite(coeffs)) + + class TestEstimatorIntegration: """Integration tests verifying estimators produce correct results.""" @@ -1913,6 +2095,202 @@ def test_rank_deficient_action_error(self): solve_logit(X, y, rank_deficient_action="error") +def _legacy_irls_reference(X, y, weights=None, max_iter=25, tol=1e-8): + """Test-local reimplementation of the pre-fast-path IRLS inner loop: + per-iteration tall-matrix `np.linalg.lstsq(Xw, zw, rcond=None)` with the + identical working weights/response/convergence semantics. Shared by the + fast-path parity and convergence-semantics tests below. Assumes a + full-rank design (no rank/EPV handling - callers use clean fixtures). + Returns (beta_with_intercept, converged).""" + n = X.shape[0] + Xi = np.column_stack([np.ones(n), X]) + beta = np.zeros(Xi.shape[1]) + for _ in range(max_iter): + eta = np.clip(Xi @ beta, -500, 500) + mu = np.clip(1.0 / (1.0 + np.exp(-eta)), 1e-10, 1 - 1e-10) + w_irls = mu * (1.0 - mu) + z = eta + (y - mu) / w_irls + w_total = weights * w_irls if weights is not None else w_irls + sqrt_w = np.sqrt(w_total) + beta_new, _, _, _ = np.linalg.lstsq(Xi * sqrt_w[:, None], z * sqrt_w, rcond=None) + if np.max(np.abs(beta_new - beta)) < tol: + return beta_new, True + beta = beta_new + return beta, False + + +class TestIRLSCholeskyFastPath: + """solve_logit's equilibrated normal-equations Cholesky inner solve. + + Parity vs the legacy per-iteration lstsq is TOL-BOUNDED (atol 1e-8), + not bit-level: both solvers converge to the same MLE, but the iteration + at which the max|delta-beta| < tol check first crosses can legally shift + by one, moving the final beta by up to tol. Observed parity on + well-conditioned fits is ~1e-10..1e-12 (iteration counts match; the + quadratically-decaying final step dominates the difference). + """ + + @staticmethod + def _make_logit(n, k, seed, scale_col=None, weights_kind=None, sep=0.0): + rng = np.random.default_rng(seed) + X = rng.standard_normal((n, k)) + if scale_col is not None: + X[:, scale_col % k] *= 1e4 + beta = rng.standard_normal(k) * 0.5 + eta = X @ beta + sep * X[:, 0] + p = 1.0 / (1.0 + np.exp(-np.clip(eta, -30, 30))) + y = (rng.random(n) < p).astype(float) + w = None + if weights_kind == "positive": + w = np.exp(rng.normal(0, 0.5, n)) + elif weights_kind == "zeros": + w = np.exp(rng.normal(0, 0.5, n)) + w[rng.random(n) < 0.1] = 0.0 + elif weights_kind == "tiny": + w = np.exp(rng.normal(0, 0.5, n)) + w[rng.random(n) < 0.1] = 1e-12 + return X, y, w + + def test_property_parity_vs_legacy_lstsq(self): + """~20 random GLM datasets (varied n/k, scale disparity, positive / + exact-zero / tiny-positive weights, mild separation pressure): the + fast-path beta matches the legacy per-iteration lstsq reimplementation + at atol 1e-8 (tol-bounded; see class docstring for why not tighter), + with zero Cholesky fallbacks on these well-conditioned fits.""" + cases = [] + for i in range(20): + cases.append( + dict( + n=200 + 137 * i, + k=2 + (i % 9), + seed=100 + i, + scale_col=i % 3 if i % 4 == 0 else None, + weights_kind=[None, "positive", "zeros", "tiny"][i % 4], + sep=0.8 if i % 5 == 0 else 0.0, + ) + ) + worst = 0.0 + for case in cases: + X, y, w = self._make_logit(**case) + diag = {} + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + beta, probs = solve_logit(X, y, weights=w, diagnostics_out=diag) + ref, _ = _legacy_irls_reference(X, y, weights=w) + delta = float(np.nanmax(np.abs(beta - ref))) + worst = max(worst, delta) + assert delta < 1e-8, (case, delta) + # saturated fits legally produce probs of exactly 0.0/1.0 in + # float (eta clipped at +-500; exp underflows) - same as legacy + assert np.all((probs >= 0) & (probs <= 1)) + assert diag["irls_chol_fallback_iters"] == 0, (case, diag) + # typical parity is far below the gate; record it in the assert + assert worst < 1e-8 + + def test_forced_fallback_bit_identical_to_legacy(self, monkeypatch): + """With cho_factor monkeypatched to ALWAYS raise, every iteration + takes the guarded fallback, which must be byte-identical to the + legacy computation (the fallback reconstructs the exact pre-fast-path + lstsq line on the raw basis).""" + import diff_diff.linalg as lmod + + def always_raise(*args, **kwargs): + raise np.linalg.LinAlgError("forced") + + X, y, _ = self._make_logit(400, 5, seed=7) + diag = {} + monkeypatch.setattr(lmod, "cho_factor", always_raise) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + beta_fb, probs_fb = solve_logit(X, y, diagnostics_out=diag) + monkeypatch.undo() + ref, _ = _legacy_irls_reference(X, y) + assert diag["irls_chol_fallback_iters"] > 0 + np.testing.assert_array_equal(beta_fb, ref) + + def test_natural_guard_trip_warning_equivalence(self, monkeypatch): + """A separation-pressure dataset trips the dpocon guard naturally in + late IRLS iterations (working weights collapse, G ill-conditioned). + The fit is tol-bounded vs legacy - NOT byte-identical, because early + well-conditioned iterations take the Cholesky path - so the lock is + warning-set EQUIVALENCE plus a positive fallback count.""" + import diff_diff.linalg as lmod + + rng = np.random.default_rng(31) + n = 600 + # Quasi-separation on a DUMMY subgroup: y == 1 on every dummy row, + # so the dummy coefficient diverges and the working weights collapse + # on exactly that column's support (G's dummy diagonal shrinks ~ + # mu*(1-mu) -> 1e-10 RELATIVE to the healthy columns - a uniform + # collapse across all rows would rescale G without changing its + # conditioning and never trip the guard). + dummy = np.zeros(n) + dummy[:60] = 1.0 + x1 = rng.standard_normal(n) + X = np.column_stack([dummy, x1]) + y = (rng.random(n) < 1.0 / (1.0 + np.exp(-x1))).astype(float) + y[:60] = 1.0 + + diag = {} + with warnings.catch_warnings(record=True) as caught_new: + warnings.simplefilter("always") + solve_logit(X, y, diagnostics_out=diag) + assert ( + diag["irls_chol_fallback_iters"] > 0 + ), "fixture no longer trips the dpocon guard - regenerate it" + + def always_raise(*args, **kwargs): + raise np.linalg.LinAlgError("forced") + + monkeypatch.setattr(lmod, "cho_factor", always_raise) + with warnings.catch_warnings(record=True) as caught_legacy: + warnings.simplefilter("always") + solve_logit(X, y) + monkeypatch.undo() + + def warning_set(records): + return {(r.category.__name__, str(r.message)[:40]) for r in records} + + assert warning_set(caught_new) == warning_set(caught_legacy) + + def test_convergence_iteration_semantics_preserved(self, monkeypatch): + """The fast path must not change WHEN the IRLS loop converges on a + well-conditioned fit: find the minimal converging max_iter N under + the legacy solver (via the always-raise monkeypatch), then assert + the fast path converges (no warning) at N and warns at N-1. Parity + at ~1e-12 makes an iteration-count shift essentially impossible on + this fixture; if a platform ever shifts it by one, relax with an + in-test justification comment.""" + import diff_diff.linalg as lmod + + X, y, _ = self._make_logit(500, 4, seed=19) + + def always_raise(*args, **kwargs): + raise np.linalg.LinAlgError("forced") + + def converges(max_iter, force_legacy): + if force_legacy: + monkeypatch.setattr(lmod, "cho_factor", always_raise) + try: + with warnings.catch_warnings(record=True) as rec: + warnings.simplefilter("always") + solve_logit(X, y, max_iter=max_iter, check_separation=False) + return not any("did not converge" in str(r.message) for r in rec) + finally: + if force_legacy: + monkeypatch.undo() + + n_min = None + for m in range(1, 26): + if converges(m, force_legacy=True): + n_min = m + break + assert n_min is not None and n_min >= 2, n_min + + assert converges(n_min, force_legacy=False) + assert not converges(n_min - 1, force_legacy=False) + + class TestCheckPropensityDiagnostics: """Tests for propensity score diagnostic warnings."""