perf(linalg): per-cell solver fast paths - IRLS Cholesky inner solve + rank-detection certification (CS covariate fits 1.5-2x)

[igerber]

Summary

• Two pure-Python fast paths in the shared solvers (diff_diff/linalg.py), hit by CallawaySantAnna ipw/dr covariate fits and every estimator routing through them (StaggeredTripleDifference, TripleDifference, ContinuousDiD, WooldridgeDiD logit; solve_ols pre-check, solve_poisson, dCDH residual-df).
• IRLS inner solve: solve_logit's per-iteration tall-matrix np.linalg.lstsq (gelsd SVD) becomes equilibrated normal equations + Cholesky under an explicit LAPACK dpocon reciprocal-condition guard (1e-6), with a per-iteration fallback that reproduces the legacy solve byte-for-byte when the normal matrix cannot be certified well-conditioned (working weights can crush a dummy column's effective scale on near-separated subgroups, so pre-loop full rank does not imply a chol-safe normal matrix; cho_factor alone can succeed with garbage at cond ~1e10-1e16). Columns are equilibrated once; IRLS state, tol, and warnings stay in the raw basis. The prior removal of an un-equilibrated cho_solve(X'X) OR fast path is cited in-code, and this path differs on exactly that axis.
• Rank-detection certification: _detect_rank_deficiency short-circuits the common full-rank case with a Gram/eigvalsh certification at the documented _rank_guarded_inv 1e-10 Gram threshold - two orders stricter than the pivoted-QR 1e-7 boundary, so it never reports full rank where the QR path would report a deficiency. All n < k, non-finite, looser-rcond, and uncertifiable designs fall through to the existing two-stage QR verbatim: every drop decision, pivot selection, and rank count is unchanged (including dCDH's residual-df and the Rust solve_ols routing boolean, proven by arm instrumentation).
• Measured (3-rep medians, CS @100k units x 20 periods = 2M rows, 95 cells): dr 1.49x/1.64x/1.85x at 10/20/40 covariates, ipw 1.98x, survey-weighted dr 1.63x; pure-Python backend pairs equal or better (up to 2.03x). Component stages at 40 covariates: IRLS solve 6.8x, rank detection 16.5x. Estimates unchanged beyond machine precision: overall ATT/SE deltas exactly 0, per-cell max ~7e-15; Cholesky fallback 0% and certification 100% on all bench scenarios; maxrss flat.

Methodology references (required if estimator / math changes)

• Method name(s): CallawaySantAnna propensity-score estimation (IRLS logistic regression) and the shared rank-deficiency detection
• Paper / source link(s): R glm(family=binomial) IRLS convention and R qr()/lm() rank handling per docs/methodology/REGISTRY.md (IRLS inner-solver Note and rank-detection scale-invariance Notes extended in this PR)
• Any intentional deviations from the source (and why): None - the IRLS algorithm (working weights, working response, tol=1e-8, max_iter=25) and all rank/drop decisions are unchanged; only the inner linear-algebra implementations differ, at the condition-bounded floating-point level (measured ~1e-15; tol-bounded worst case ~1e-8).

Validation

• Tests added/updated: tests/test_linalg.py - new TestIRLSCholeskyFastPath (property parity vs a legacy-IRLS reimplementation at the tol-bounded 1e-8 gate over 20 varied datasets incl. zero/tiny weights and separation pressure; monkeypatched always-raise fallback byte-identity lock; natural dpocon-trip warning-set equivalence + diagnostics_out["irls_chol_fallback_iters"] counter; convergence-iteration semantics lock) and TestRankDetectionStage0Certification (zero-QR certification spy incl. a 1e8-scale column; declined paths - collinear, n<k structural, zero column, NaN still raises, looser rcond, and a between-columns cond~1e6 boundary locking cert-strictly-stricter-than-QR; Kahan characterization asserting certification can only disagree with QR by being more correct; Rust solve_ols dispatch composition lock). All existing rank-selection, logit, golden ipw/dr (R-parity SE abs<1e-6), survey self-consistency (atol=1e-10), and dCDH bit-identity baselines pass UNMODIFIED on both backends (1714 rust / 1594 python targeted sweep; full suite 8567 passed with only the known pre-existing order-dependent dCDH failure, which passes in isolation).
• Backtest / simulation / notebook evidence (if applicable): perf arms in docs/performance-plan.md (BEFORE = pristine main via baseline worktree, per-backend arm-pairs, BLAS threads pinned, instrumented fallback/certification/dispatch counters).

Security / privacy

• Confirm no secrets/PII in this PR: Yes

---

Generated with Claude Code

[github-actions]

PR Review Report

Overall Assessment

✅ Looks good — I found no unmitigated P0/P1 issues in the PR diff.

Executive Summary

• The affected methods are solve_logit() propensity-score IRLS and shared _detect_rank_deficiency() rank handling.
• Both numerical changes are documented in docs/methodology/REGISTRY.md with explicit Notes, so they are implementation choices rather than undocumented methodology deviations.
• The IRLS algorithm contract remains unchanged: same working weights, working response, raw-basis tolerance, max iterations, and fallback to legacy lstsq.
• The rank fast path is full-rank-only and preserves the documented QR fallback for deficient, non-finite, n < k, and looser-rcond designs.
• Added tests cover Cholesky fallback, warning equivalence, rank-certification decline paths, non-finite handling, and Rust dispatch composition.
• I could not run the test suite in this review environment because the bare Python interpreter lacks NumPy/SciPy.

Methodology

Finding: P3 informational — Documented solver fast paths

Severity: P3
Impact: CallawaySantAnna IPW/DR propensity-score fits and any estimator using shared OLS/rank helpers now use faster numerical paths, but the registry documents both changes as implementation choices. The IRLS note at docs/methodology/REGISTRY.md:L728-L743 states that the IRLS algorithm is unchanged and only the inner WLS solve changes. The rank note at docs/methodology/REGISTRY.md:L78-L81 documents the Gram/eigenvalue full-rank certification and QR fallback contract.
Concrete fix: No action required.

Code Quality

No P0/P1 findings.

The Cholesky path in diff_diff/linalg.py:L3193-L3221 is guarded by dpocon and falls back to the pre-existing raw-basis np.linalg.lstsq line when uncertified. The fallback counter is recorded via diagnostics_out at diff_diff/linalg.py:L3230-L3231.

Performance

Finding: P3 informational — Remaining performance work is tracked

Severity: P3
Impact: The PR adds TODO entries for repeated-cross-section propensity-score caching and Rust solve_ols fast paths in TODO.md:L44-L47. These are performance follow-ups, not correctness blockers.
Concrete fix: No action required for this PR.

Maintainability

No findings.

The new tests in tests/test_linalg.py:L1678-L1845 and tests/test_linalg.py:L2098-L2291 are focused on the new contracts and make the fast-path/fallback boundaries explicit.

Tech Debt

No untracked blocking tech debt found. New deferrable performance items are tracked in TODO.md, which satisfies the project’s deferred-work policy.

Security

No findings. The diff only changes numerical solver code, tests, and documentation; I saw no secrets or sensitive-data handling changes.

Documentation/Tests

Finding: P3 informational — Local tests not run in this review environment

Severity: P3
Impact: I could not execute the tests because the available Python interpreter has no NumPy/SciPy installed. The PR does add targeted tests for the new linalg behavior.
Concrete fix: Ensure CI runs the targeted tests/test_linalg.py suite and existing estimator parity tests.