Skip to content

Commit 27c4758

Browse files
igerberigerber
authored
perf(linalg): per-cell solver fast paths - IRLS Cholesky inner solve + rank-detection certification (#634)
1 parent a564bd7 commit 27c4758

6 files changed

Lines changed: 576 additions & 10 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,27 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
247247
supersedes it.
248248

249249
### Changed
250+
- **Per-cell solver fast paths for covariate fits (CallawaySantAnna and every
251+
estimator routing through the shared solvers).** Two pure-Python changes in
252+
`diff_diff/linalg.py`: (1) `solve_logit`'s IRLS inner step — previously a full
253+
SVD (`np.linalg.lstsq`) on the tall weighted design per iteration — now solves
254+
equilibrated normal equations by Cholesky under an explicit
255+
reciprocal-condition guard (LAPACK `dpocon`), falling back to the exact legacy
256+
solve for any iteration whose normal matrix cannot be certified
257+
well-conditioned; (2) `_detect_rank_deficiency` short-circuits the common
258+
full-rank case with a Gram/eigenvalue certification two orders stricter than
259+
the pivoted-QR boundary, so all rank counts, drop decisions, and pivot
260+
selections on deficient or uncertifiable designs are unchanged. Measured
261+
(3-rep medians, CallawaySantAnna, 100k units x 20 periods = 2M rows, 95 cells,
262+
`aggregate="all"`): dr fits 1.5s → 1.0s at 10 covariates (1.49x), 2.6s → 1.6s
263+
at 20 (1.64x), 5.7s → 3.1s at 40 (1.85x); ipw 1.8s → 0.9s (2.0x);
264+
survey-weighted dr 2.7s → 1.7s (1.6x); pure-Python backend gains are equal or
265+
larger. The solver stages themselves: IRLS solve 6.8x, rank detection 16.5x
266+
(at 40 covariates). Estimates are unchanged beyond machine precision (overall
267+
ATT/SE deltas exactly 0; per-cell max ~7e-15; the R-golden ipw SE parity at
268+
1e-6 abs and the dCDH bit-identity baselines hold unmodified), the Cholesky
269+
fallback fires on 0% of healthy fits (its near-separation trip path is locked
270+
by tests), and memory is flat.
250271
- **`CallawaySantAnna` no-covariate `estimation_method="dr"` per-cell SE is now
251272
influence-function-based** (`sqrt(sum(phi^2))`), matching the `reg`/`ipw` branches and R's
252273
`DRDID::drdid_panel`. It was the last per-cell SE on the ddof=1 plug-in

‎TODO.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,8 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m
4646
| Issue | Location | Origin | Effort | Priority |
4747
|-------|----------|--------|--------|----------|
4848
| `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 |
49+
| `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 |
50+
| 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 |
4951
| `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 |
5052
| 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 |
5153
| 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_diff/linalg.py‎

Lines changed: 106 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -39,8 +39,9 @@
3939
import numpy as np
4040
import pandas as pd
4141
from scipy import stats
42+
from scipy.linalg import cho_factor, cho_solve, qr
4243
from scipy.linalg import lstsq as scipy_lstsq
43-
from scipy.linalg import qr
44+
from scipy.linalg.lapack import dpocon
4445

4546
# Import Rust backend if available (from _backend to avoid circular imports)
4647
from diff_diff._backend import (
@@ -132,16 +133,54 @@ def _detect_rank_deficiency(
132133
Indices of columns that are linearly dependent (should be dropped).
133134
Empty if matrix is full rank.
134135
pivot : ndarray of int
135-
Column permutation from QR decomposition.
136+
Column permutation from QR decomposition. For a full-rank result the
137+
pivot carries no information (no caller consumes it; the stage-0
138+
certification below returns a trivial ``arange`` pivot).
136139
"""
137-
_, k = X.shape
140+
n, k = X.shape
138141
if k == 0:
139142
return 0, np.array([], dtype=int), np.array([], dtype=int)
140143

141144
# R's qr() uses tol = 1e-07 by default (sqrt(eps) ≈ 1.49e-08); we use 1e-07.
142145
if rcond is None:
143146
rcond = 1e-07
144147

148+
# Stage 0 — Gram full-rank CERTIFICATION (perf fast path; decisions never
149+
# change). The Gram is built directly from X (no equilibrated copy of the
150+
# tall matrix); diag(G) holds the squared column 2-norms, so equilibrating
151+
# G symmetrically by sqrt(diag(G)) is exactly column equilibration of X —
152+
# the `_rank_guarded_inv` convention. Certification threshold 1e-10 is the
153+
# documented Gram constant from `_rank_guarded_inv` (a Gram squares the
154+
# condition number of X, so 1e-10 on eigenvalues ~ cond(X_eq) < 1e5, two
155+
# orders STRICTER than the 1e-7 QR full-rank boundary): stage 0 never
156+
# reports full rank where the two-stage QR below would report a
157+
# deficiency, it only declines and falls through. (A pathological
158+
# Kahan-type matrix could in principle make pivoted-QR R-diagonals
159+
# undershoot the singular values enough to open a gap; real DiD designs —
160+
# dummies plus covariates — do not have that structure, and the
161+
# characterization test documents it.) Skipped when n < k (always
162+
# deficient — also keeps the sole pivot consumer, staggered.py's
163+
# underdetermined pair solve, structurally on the QR path) and when a
164+
# caller passes a LOOSER-than-default rcond (no caller does today; the
165+
# stricter-than-QR guarantee above assumes rcond <= 1e-7). Non-finite
166+
# entries poison diag(G), decline certification, and fall through so
167+
# scipy's qr raises ValueError on NaN/Inf exactly as before.
168+
if n >= k and rcond <= 1e-07:
169+
gram = X.T @ X
170+
diag = np.diag(gram)
171+
if np.all(np.isfinite(diag)) and np.all(diag > 0):
172+
scales = np.sqrt(diag)
173+
gram_eq = gram / scales[:, None] / scales[None, :]
174+
eigvals = np.linalg.eigvalsh(gram_eq)
175+
eig_min, eig_max = eigvals[0], eigvals[-1]
176+
if (
177+
np.isfinite(eig_min)
178+
and np.isfinite(eig_max)
179+
and eig_max > 0.0
180+
and eig_min > 1e-10 * eig_max
181+
):
182+
return k, np.array([], dtype=int), np.arange(k, dtype=int)
183+
145184
def _rank_and_pivot(M: np.ndarray) -> Tuple[int, np.ndarray]:
146185
# Pivoted QR: M @ P = Q @ R. The rank threshold is anchored to the
147186
# largest pivot diagonal |R[0,0]| (decreasing after pivoting).
@@ -2892,6 +2931,12 @@ def _compute_robust_vcov_numpy(
28922931
_LOGIT_SEPARATION_PROB_THRESHOLD = 1e-5
28932932
_DEFAULT_EPV_THRESHOLD = 10
28942933

2934+
# Reciprocal-condition guard for the IRLS normal-equations Cholesky fast path
2935+
# (solve_logit): cond(G) <= 1e6 bounds the Cholesky forward error at
2936+
# ~eps * cond ~ 2e-10 relative, consistent with the tol-bounded parity budget;
2937+
# anything worse falls back to the legacy tall-matrix lstsq for that iteration.
2938+
_IRLS_CHOL_RCOND_GUARD = 1e-6
2939+
28952940

28962941
def solve_logit(
28972942
X: np.ndarray,
@@ -2947,7 +2992,10 @@ def solve_logit(
29472992
users identify which logit estimation triggered the warning.
29482993
diagnostics_out : dict, optional
29492994
If provided, populated with EPV diagnostic info:
2950-
``{"epv": float, "n_events": int, "k": int, "is_low": bool}``.
2995+
``{"epv": float, "n_events": int, "k": int, "is_low": bool}``, plus
2996+
``"irls_chol_fallback_iters"`` - the number of IRLS iterations whose
2997+
normal-equations Cholesky fast path fell back to the legacy lstsq
2998+
solve (0 on well-conditioned fits).
29512999
29523000
Returns
29533001
-------
@@ -3097,7 +3145,31 @@ def solve_logit(
30973145
raise ValueError(msg)
30983146
warnings.warn(msg, UserWarning, stacklevel=2)
30993147

3100-
# IRLS (Fisher scoring)
3148+
# IRLS (Fisher scoring). Each weighted-least-squares step is solved via
3149+
# EQUILIBRATED normal equations + Cholesky with an explicit condition
3150+
# guard, falling back to the legacy tall-matrix lstsq (gelsd SVD) for any
3151+
# iteration whose normal matrix cannot be certified well-conditioned.
3152+
# Context: the OR path deliberately REMOVED a cho_solve(X'X) fast path
3153+
# because it was NOT scale-equilibrated (see the covariate-reg notes in
3154+
# staggered.py around `_equilibrated_lstsq`); this path differs on
3155+
# exactly that axis - (1) columns are equilibrated to unit 2-norm ONCE
3156+
# (a fixed reparameterization: X_eq @ beta_eq == X @ beta algebraically,
3157+
# so probabilities are unchanged and beta is unscaled per iteration);
3158+
# (2) cho_factor alone can SUCCEED with a garbage solution when cond(G)
3159+
# exceeds ~1e10, so the guard is a dpocon reciprocal-condition estimate,
3160+
# not just the factorization succeeding - working weights can crush a
3161+
# column's effective scale (a dummy on a near-separated subgroup has
3162+
# w_irls ~ 1e-10 on its support), so full column rank pre-loop does NOT
3163+
# imply a well-conditioned G; (3) the fallback reproduces the legacy
3164+
# computation for that iteration on the raw basis. IRLS state (beta,
3165+
# convergence tol, warnings) stays in the RAW basis: a scaled-basis tol
3166+
# would be ~sqrt(n)x tighter for every fit (the intercept column alone
3167+
# has 2-norm sqrt(n)), and the separation check below reads raw beta.
3168+
irls_col_norms = np.sqrt(np.einsum("ij,ij->j", X_solve, X_solve))
3169+
irls_safe_norms = np.where(irls_col_norms > 0, irls_col_norms, 1.0)
3170+
X_eq = X_solve / irls_safe_norms
3171+
chol_fallback_iters = 0
3172+
31013173
beta_solve = np.zeros(X_solve.shape[1])
31023174
converged = False
31033175

@@ -3120,9 +3192,33 @@ def solve_logit(
31203192

31213193
# Weighted least squares: solve (X'WX) beta = X'Wz
31223194
sqrt_w = np.sqrt(w_total)
3123-
Xw = X_solve * sqrt_w[:, None]
31243195
zw = z * sqrt_w
3125-
beta_new, _, _, _ = np.linalg.lstsq(Xw, zw, rcond=None)
3196+
beta_new = None
3197+
Xw_eq = X_eq * sqrt_w[:, None]
3198+
gram = Xw_eq.T @ Xw_eq
3199+
# 1-norm BEFORE factorization (dpocon contract); G is symmetric so
3200+
# the max absolute column sum is the 1-norm.
3201+
anorm = float(np.max(np.sum(np.abs(gram), axis=0)))
3202+
if np.isfinite(anorm):
3203+
try:
3204+
chol = cho_factor(gram)
3205+
except np.linalg.LinAlgError:
3206+
chol = None
3207+
if chol is not None:
3208+
# cho_factor default lower=False pairs with dpocon's default
3209+
# uplo='U' (dpocon has no `lower=` kwarg).
3210+
rcond_gram, pocon_info = dpocon(chol[0], anorm)
3211+
if (
3212+
pocon_info == 0
3213+
and np.isfinite(rcond_gram)
3214+
and rcond_gram > _IRLS_CHOL_RCOND_GUARD
3215+
):
3216+
beta_new = cho_solve(chol, Xw_eq.T @ zw) / irls_safe_norms
3217+
if beta_new is None:
3218+
# Guarded fallback: byte-identical to the pre-fast-path solve.
3219+
chol_fallback_iters += 1
3220+
Xw = X_solve * sqrt_w[:, None]
3221+
beta_new, _, _, _ = np.linalg.lstsq(Xw, zw, rcond=None)
31263222

31273223
# Check convergence
31283224
if np.max(np.abs(beta_new - beta_solve)) < tol:
@@ -3131,6 +3227,9 @@ def solve_logit(
31313227
break
31323228
beta_solve = beta_new
31333229

3230+
if diagnostics_out is not None:
3231+
diagnostics_out["irls_chol_fallback_iters"] = chol_fallback_iters
3232+
31343233
# Final predicted probabilities
31353234
eta_final = X_solve @ beta_solve
31363235
eta_final = np.clip(eta_final, -500, 500)

0 commit comments

Comments
 (0)