From a1061c6e3250695d46b1cbe4c1208b370f23e2da Mon Sep 17 00:00:00 2001 From: igerber Date: Tue, 7 Jul 2026 06:59:17 -0400 Subject: [PATCH] perf(efficient_did): Rust batched-Cholesky ridge solve for conditional Omega* weights _ridge_solve_weights' batched np.linalg.solve (serial LAPACK LU over the (units x H x H) ridged Omega* stacks; the top stage at every scale after the kcov-table hoisting) now dispatches to a new Rust kernel batched_ridge_chol_solve_ones: hand-rolled unblocked Cholesky (the ridged Omega* is SPD by construction) in a reused per-thread scratch buffer, rayon-parallel over the unit axis, with an in-kernel faer LU fallback for non-SPD rows (NaN-poisoned on an exactly-zero pivot) and per-row legacy numpy recompute of any non-finite row, so pinv edge-case semantics are unchanged. Dispatch is on backend availability + float64 only - no batch-size cutoff (tile-invariance twins force one-unit batches through the same kernel). Measured (3-rep medians, 20p/5-cohort/5-cov, aggregate="all"): 2k 7.1 -> 4.3s (1.65x) | 10k 36.2 -> 22.7s (1.60x) 30k 129 -> 90s (1.43x) | 100k 17.8 -> 15.9 min | survey_1k 1.8x Ridge stage at 10k 17.1 -> 3.7s (kernel alone 9.6x); maxrss flat/lower. Deltas: post cells ~2e-12 rel, overall ATT ~1e-13; pure-python backend byte-identical; nocov path untouched; omega_ridge=0 never reaches the kernel. New tests: kernel parity vs numpy (well/ill-conditioned, NaN, exact singular, batch-split bit-identity, degenerate shapes) in test_rust_backend.py; dispatch contracts (kernel-called spy, in-process backend A/B, symbol-None degradation, bad-row-only exact recompute, exact-singular pinv arm, f32 declines) in test_efficient_did.py; 5 cargo tests. REGISTRY ridge Note extended (pure implementation change); CHANGELOG + performance-plan tables; TODO row swapped to the remaining Python-prep shave. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01X4AzrFUMqJxcUumSH31mSr --- CHANGELOG.md | 17 ++ TODO.md | 2 +- diff_diff/_backend.py | 13 + diff_diff/efficient_did_covariates.py | 59 ++++- docs/methodology/REGISTRY.md | 2 +- docs/performance-plan.md | 48 +++- rust/src/batched_solve.rs | 353 ++++++++++++++++++++++++++ rust/src/lib.rs | 7 + tests/test_efficient_did.py | 151 +++++++++++ tests/test_rust_backend.py | 138 ++++++++++ 10 files changed, 773 insertions(+), 17 deletions(-) create mode 100644 rust/src/batched_solve.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 1d243d81..209f0ce6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -258,6 +258,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 the no-covariates path is byte-identical and `omega_ridge=0` still routes the entire legacy path. A `pt_assumption="post"` covariate fit builds no tables at all (all cells are just-identified), keeping that path regression-free. +- **`EfficientDiD` conditional-path ridge solves dispatch to a Rust batched-Cholesky kernel** + (follow-up to the table hoisting above; Rust backend only). `_ridge_solve_weights`' batched + `np.linalg.solve` — a serial LAPACK LU sweep over the `(units × H × H)` ridged Omega* stacks + and the top stage at every scale post-hoisting — now routes to + `batched_ridge_chol_solve_ones` in the Rust backend: a hand-rolled unblocked Cholesky (the + ridged Omega* is SPD by construction) in a reused per-thread scratch buffer, parallelized + over the unit axis with rayon. Non-SPD rows (measured zero on realistic panels) fall back to + LU in-kernel, and any non-finite row is recomputed through the exact legacy numpy chain, so + the pseudoinverse edge-case semantics are unchanged. Measured (3-rep medians, 20 periods, + 5 cohorts, 5 covariates, `aggregate="all"`): fit 7.1s → 4.3s at 2k units (1.65x), + 36.2s → 22.7s at 10k (1.6x), 129s → 90s at 30k (1.4x), the 100k-unit fit 17.8 → 15.9 min at + slightly lower peak RSS; survey-weighted 3.0s → 1.7s at 1k (1.8x). The ridge-solve stage + itself is 4.7x faster at 10k (17.1s → 3.7s; the kernel alone is ~9.6x — the remainder is + shared Python-side prep, tracked in TODO.md). Results on the Rust backend move at + floating-point reassociation level only (post-treatment cells ~2e-12 relative, overall ATT + ~1e-13); the pure-Python backend is byte-identical, the no-covariates path is untouched, and + `omega_ridge=0` never reaches the kernel. - **CallawaySantAnna multiplier bootstrap rewritten as a fused, column-tiled scatter-GEMM** (EfficientDiD's bootstrap routes through the same kernel). The former loop sliced the `(block × n_units)` weight matrix twice per (g,t) cell per weight block — at a 40-period, diff --git a/TODO.md b/TODO.md index 6ca2f22d..034595fa 100644 --- a/TODO.md +++ b/TODO.md @@ -46,7 +46,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| -| `EfficientDiD` conditional path: post-hoisting (2026-07, 100k fit ~18 min) the top stage at every scale is `_ridge_solve_weights` — batched `np.linalg.solve` LU over `(n × H × H)` stacks (17.2s of a 35.7s fit at 10k units, ~170s at 100k; linear in n). The matrices are SPD after the ridge, so a batched Cholesky (potrf ~1/2 the LU flops) in the Rust backend — which could also thread across the stack — is the natural lever; numpy has no batched triangular solve, so pure-Python gains are limited. | `efficient_did_covariates.py::_ridge_solve_weights`, `rust/src/` | CS-scaling | Mid | Low | +| `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 | | `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/_backend.py b/diff_diff/_backend.py index 2c15de5f..2824fb90 100644 --- a/diff_diff/_backend.py +++ b/diff_diff/_backend.py @@ -73,6 +73,15 @@ except ImportError: _rust_demean_map = None +# Batched ridge-regularized SPD solve (EfficientDiD per-unit weights): +# imported independently for the same mixed-version reason as demean_map. +try: + from diff_diff._rust_backend import ( + batched_ridge_chol_solve_ones as _rust_batched_ridge_chol_solve, + ) +except ImportError: + _rust_batched_ridge_chol_solve = None + # Determine final backend based on environment variable and availability if _backend_env == "python": # Force pure Python mode - disable Rust even if available @@ -83,6 +92,8 @@ _rust_compute_robust_vcov = None # FE-absorption MAP demeaning kernel _rust_demean_map = None + # Batched ridge-regularized SPD solve + _rust_batched_ridge_chol_solve = None # TROP estimator acceleration (local method) _rust_unit_distance_matrix = None _rust_loocv_grid_search = None @@ -136,6 +147,8 @@ def rust_backend_info(): "_rust_compute_robust_vcov", # FE-absorption MAP demeaning kernel "_rust_demean_map", + # Batched ridge-regularized SPD solve (EfficientDiD per-unit weights) + "_rust_batched_ridge_chol_solve", # TROP estimator acceleration (local method) "_rust_unit_distance_matrix", "_rust_loocv_grid_search", diff --git a/diff_diff/efficient_did_covariates.py b/diff_diff/efficient_did_covariates.py index 42a35bf5..7623f479 100644 --- a/diff_diff/efficient_did_covariates.py +++ b/diff_diff/efficient_did_covariates.py @@ -26,6 +26,7 @@ import numpy as np from scipy.spatial.distance import cdist +from diff_diff._backend import HAS_RUST_BACKEND, _rust_batched_ridge_chol_solve from diff_diff.linalg import solve_ols # Default ridge for the Omega* inversion (see ``compute_per_unit_weights`` / @@ -1629,27 +1630,57 @@ def _ridge_solve_weights(omega_stack: np.ndarray, omega_ridge: float) -> np.ndar om = omega_stack[rest] trace = np.trace(om, axis1=1, axis2=2) - scale = np.maximum(trace / H, 0.0) - om_ridged = om + (omega_ridge * scale)[:, None, None] * np.eye(H)[None] + ridge = omega_ridge * np.maximum(trace / H, 0.0) + if ( + HAS_RUST_BACKEND + and _rust_batched_ridge_chol_solve is not None + and omega_stack.dtype == np.float64 + ): + # Batched Cholesky in Rust (ridge added in-kernel - no om_ridged + # temp). Dispatch is on availability + dtype ONLY: any batch-size + # cutoff would make forced-tiny-tile fits cross-algorithm vs the + # single-tile fit and break the tile-invariance twins. A non-finite + # row signals not-SPD (kernel LU fallback / NaN-poisoned exact + # singularity) - recompute exactly those rows via the legacy chain + # so their semantics (incl. the pinv arm) match the numpy path. + num = _rust_batched_ridge_chol_solve(om, ridge) + bad = np.flatnonzero(~np.isfinite(num).all(axis=1)) + if bad.size: + num[bad] = _ridge_solve_numpy(om[bad], ridge[bad]) + else: + num = _ridge_solve_numpy(om, ridge) + den = num @ ones + ok = np.abs(den) >= 1e-15 + solved = np.full((rest.size, H), 1.0 / H) + solved[ok] = num[ok] / den[ok, None] + weights[rest] = solved + return weights + + +def _ridge_solve_numpy(om: np.ndarray, ridge: np.ndarray) -> np.ndarray: + """Legacy numpy solve chain for ``(om_i + ridge_i * I) x = 1``. + + Batched LU solve; on an exactly-singular batch, per-unit solve with a + minimum-norm pseudoinverse backstop (unreachable for PSD Omega* with + lam > 0, kept so one pathological unit cannot fail the batch). This is + both the pure-Python path and the per-row recompute target for rows the + Rust kernel flags as non-finite. + """ + m_rows, H, _ = om.shape + om_ridged = om + ridge[:, None, None] * np.eye(H)[None] try: # gufunc solve treats a 2-D rhs as a matrix; a stack of vectors needs - # shape (n, H, 1) - num = np.linalg.solve(om_ridged, np.ones((rest.size, H, 1)))[..., 0] + # shape (m, H, 1) + return np.linalg.solve(om_ridged, np.ones((m_rows, H, 1)))[..., 0] except np.linalg.LinAlgError: - # Unreachable for PSD Omega* with lam > 0; kept as a per-unit - # minimum-norm fallback so one pathological unit cannot fail the batch. - num = np.empty((rest.size, H)) - for m in range(rest.size): + ones = np.ones(H) + num = np.empty((m_rows, H)) + for m in range(m_rows): try: num[m] = np.linalg.solve(om_ridged[m], ones) except np.linalg.LinAlgError: num[m] = np.linalg.pinv(om_ridged[m]) @ ones - den = num @ ones - ok = np.abs(den) >= 1e-15 - solved = np.full((rest.size, H), 1.0 / H) - solved[ok] = num[ok] / den[ok, None] - weights[rest] = solved - return weights + return num def compute_per_unit_weights( diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index d216ac96..6a89d00d 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -1254,7 +1254,7 @@ where `q_{g,e} = pi_g / sum_{g' in G_{trt,e}} pi_{g'}`. - **PT-Post regime (just-identified)**: Under PT-Post, EDiD automatically reduces to standard single-baseline estimator (Corollary 3.2). No downside to using EDiD -- it subsumes standard estimators - **Duplicate rows**: Duplicate `(unit, time)` entries are rejected with `ValueError`. The estimator requires exactly one observation per unit-period - **Note:** PT-All index set includes g'=∞ (never-treated) as a candidate comparison group and excludes period_1 for all g'. When g'=∞, the second and third Eq 3.9 terms telescope so all (∞, t_pre) moments produce the same 2x2 DiD value; these redundant moments are damped by the default Omega* ridge (see the Omega* ridge Note below; at `omega_ridge=0`, by the legacy pseudoinverse). When t_pre = period_1, the third term degenerates to E[Y_1 - Y_1 | G=g'] = 0 for any g', adding no information. Valid pairs require only t_pre < g' (pre-treatment for comparison group), not t_pre < g. Same-group pairs (g'=g) are valid and contribute overidentifying moments (Equation 3.9) — except the degenerate PRE-treatment self-pair (g'=g, t_pre=t), which the default ridge path excludes (see the ridge Note). -- **Note:** Omega* ridge regularization (`omega_ridge`, default 1e-6; v3.7) — a documented refinement of the Omega* inversion in Eq 3.5/3.13/4.3, in the space the paper leaves open (it assumes Omega* is invertible and does not prescribe finite-sample handling of singularity). Under PT-All the overidentified moment set makes the SAMPLE Omega* numerically singular by construction (the telescoping (∞, t_pre) moments above plus near-duplicate cross-cohort moments): measured cond 1e17–1e22 for 100% of units on realistic panels, with a spectrum of one exact-null direction (the degenerate self-pair, below) plus a cluster of statistically-null directions at relative eigenvalue 1e-5–1e-8, far below the ~1e-2 sampling noise of the covariance entries — and no clean spectral gap. The prior pseudoinverse fallback therefore sat on the rcond-cutoff cliff: ANY floating-point-level change (BLAS reordering, platform change, a 1-ulp data perturbation) redrew per-cell weights and moved per-cell ATT(g,t) at up to ~1e-2 relative (measured 1.2e-4 per-cell rel for a 1-ulp outcome perturbation on the pre-v3.7 code; overall ATT stable at ~1e-9 because the redraw averages out across units and cells). The ridge solves `(Omega* + omega_ridge * max(trace/H, 0) * I) x = 1` (trace-scaled: scale-equivariant, O(H), no SVD) — a smooth regularization with sensitivity bounded by 1/omega_ridge, not a cutoff. **Why the deviation is safe:** every moment individually identifies ATT(g,t), so any fixed weights summing to 1 keep the estimator consistent; the ridge trades a numerically ill-defined efficiency optimum for a stable one, and the plug-in EIF treatment of estimated weights (Remark 4.2) is unaffected. **Calibration evidence (2026-07):** default 1e-6 chosen by 1-ulp stability (per-cell rel <= 3e-9 vs 1.2e-4 legacy; candidates 1e-4/1e-6/1e-8 all pass the <=1e-6 target) with the HRS Table 6 anchors as a hard gate (all anchors within the published-value tolerance at every candidate; worst deviation 0.0257 SE, unchanged from legacy; shift <= 0.0001 SE); Monte Carlo on covariate-confounded DGPs shows bias/RMSE/SE-calibration/coverage statistically identical to legacy (ridge marginally better point metrics). **One-time value shift:** per-cell and event-study values on the covariate path move within the pre-existing indeterminacy band when upgrading (worst observed post-treatment cell shift ~0.6 of its own SE at n=500, shrinking with n: overall-ATT shift 1.6e-2 → 4e-3 → 1.1e-3 rel at n=500/1k/2k); the no-covariates path is essentially unchanged (~1e-7). `omega_ridge=0` restores the ENTIRE legacy code path bit-for-bit (both the omega constructions and the inv/pinv weights, including per-cell condition-number warnings and the legacy O(n^2 H^2) runtime). **Degenerate self-pair:** for PRE-treatment cells (t < g), the pair (g'=g, t_pre=t) telescopes to the identically-zero moment 0=0 (the exact-null Omega* direction). The legacy pseudoinverse truncated it, spreading weight over noisy moments (spurious pre-treatment placebos of ~5% of the effect size, pure noise amplification); a naive ridge would instead load all weight on the zero-variance moment, making placebos deterministically zero and silently disabling the pre-trend diagnostic. The default ridge path therefore drops this zero-information pair (fit-level filter; post-treatment cells never contain it), restoring honest data-driven placebos; `omega_ridge=0` keeps the legacy pair set. **Warnings/diagnostics:** the no-covariates path still computes per-cell condition numbers (cheap at (H,H)) for `results.omega_condition_numbers` and consolidates cells with cond > 1e12 into ONE fit-level warning (count + max cond) instead of the legacy per-cell pseudoinverse warnings; the covariate path intentionally computes NO per-unit condition numbers — they would cost exactly the per-unit SVDs the v3.7 rewrite removes, near-singularity there is structural (~always true under PT-All, so a warning would be always-on noise rather than signal), and the ridge handles it by design (the legacy covariate path likewise had no per-unit diagnostics). The scalability warning threshold moved from n > 5000 (legacy per-cell warning) to n > 50,000 (one fit-level warning; the kernel stage is still intrinsically O(n^2) but with a ~100x lower constant and tile-bounded memory). **Cross-cell table hoisting (2026-07 follow-up):** the tiled implementation now builds the kernel-covariance tables once per comparison group per unit-tile instead of per (g, t) cell — every Omega* term is `s_group * KCov(Y_u1 - Y_v1, Y_u2 - Y_v2 | group)`, keyed only by wide-outcome columns, so the per-cell H(H+1)/2-pair tables dedup to distinct product columns per group (~26x fewer kernel GEMM columns on a PT-All fit), and the kernel weight matrices are built and freed one group at a time (the tile memory budget is governed by the largest single group rather than the sum, giving proportionally fatter tiles at large n). Each cell's Omega* is then gathered from the group tables in the same per-entry operation order as the per-cell construction (value-exact, locked by test); this is a pure implementation change — measured results move only at floating-point reassociation level (post-treatment cells ~1e-12 relative, overall ATT ~1e-13; the no-covariates path is byte-identical), and `omega_ridge=0` still routes the entire legacy path. +- **Note:** Omega* ridge regularization (`omega_ridge`, default 1e-6; v3.7) — a documented refinement of the Omega* inversion in Eq 3.5/3.13/4.3, in the space the paper leaves open (it assumes Omega* is invertible and does not prescribe finite-sample handling of singularity). Under PT-All the overidentified moment set makes the SAMPLE Omega* numerically singular by construction (the telescoping (∞, t_pre) moments above plus near-duplicate cross-cohort moments): measured cond 1e17–1e22 for 100% of units on realistic panels, with a spectrum of one exact-null direction (the degenerate self-pair, below) plus a cluster of statistically-null directions at relative eigenvalue 1e-5–1e-8, far below the ~1e-2 sampling noise of the covariance entries — and no clean spectral gap. The prior pseudoinverse fallback therefore sat on the rcond-cutoff cliff: ANY floating-point-level change (BLAS reordering, platform change, a 1-ulp data perturbation) redrew per-cell weights and moved per-cell ATT(g,t) at up to ~1e-2 relative (measured 1.2e-4 per-cell rel for a 1-ulp outcome perturbation on the pre-v3.7 code; overall ATT stable at ~1e-9 because the redraw averages out across units and cells). The ridge solves `(Omega* + omega_ridge * max(trace/H, 0) * I) x = 1` (trace-scaled: scale-equivariant, O(H), no SVD) — a smooth regularization with sensitivity bounded by 1/omega_ridge, not a cutoff. **Why the deviation is safe:** every moment individually identifies ATT(g,t), so any fixed weights summing to 1 keep the estimator consistent; the ridge trades a numerically ill-defined efficiency optimum for a stable one, and the plug-in EIF treatment of estimated weights (Remark 4.2) is unaffected. **Calibration evidence (2026-07):** default 1e-6 chosen by 1-ulp stability (per-cell rel <= 3e-9 vs 1.2e-4 legacy; candidates 1e-4/1e-6/1e-8 all pass the <=1e-6 target) with the HRS Table 6 anchors as a hard gate (all anchors within the published-value tolerance at every candidate; worst deviation 0.0257 SE, unchanged from legacy; shift <= 0.0001 SE); Monte Carlo on covariate-confounded DGPs shows bias/RMSE/SE-calibration/coverage statistically identical to legacy (ridge marginally better point metrics). **One-time value shift:** per-cell and event-study values on the covariate path move within the pre-existing indeterminacy band when upgrading (worst observed post-treatment cell shift ~0.6 of its own SE at n=500, shrinking with n: overall-ATT shift 1.6e-2 → 4e-3 → 1.1e-3 rel at n=500/1k/2k); the no-covariates path is essentially unchanged (~1e-7). `omega_ridge=0` restores the ENTIRE legacy code path bit-for-bit (both the omega constructions and the inv/pinv weights, including per-cell condition-number warnings and the legacy O(n^2 H^2) runtime). **Degenerate self-pair:** for PRE-treatment cells (t < g), the pair (g'=g, t_pre=t) telescopes to the identically-zero moment 0=0 (the exact-null Omega* direction). The legacy pseudoinverse truncated it, spreading weight over noisy moments (spurious pre-treatment placebos of ~5% of the effect size, pure noise amplification); a naive ridge would instead load all weight on the zero-variance moment, making placebos deterministically zero and silently disabling the pre-trend diagnostic. The default ridge path therefore drops this zero-information pair (fit-level filter; post-treatment cells never contain it), restoring honest data-driven placebos; `omega_ridge=0` keeps the legacy pair set. **Warnings/diagnostics:** the no-covariates path still computes per-cell condition numbers (cheap at (H,H)) for `results.omega_condition_numbers` and consolidates cells with cond > 1e12 into ONE fit-level warning (count + max cond) instead of the legacy per-cell pseudoinverse warnings; the covariate path intentionally computes NO per-unit condition numbers — they would cost exactly the per-unit SVDs the v3.7 rewrite removes, near-singularity there is structural (~always true under PT-All, so a warning would be always-on noise rather than signal), and the ridge handles it by design (the legacy covariate path likewise had no per-unit diagnostics). The scalability warning threshold moved from n > 5000 (legacy per-cell warning) to n > 50,000 (one fit-level warning; the kernel stage is still intrinsically O(n^2) but with a ~100x lower constant and tile-bounded memory). **Cross-cell table hoisting (2026-07 follow-up):** the tiled implementation now builds the kernel-covariance tables once per comparison group per unit-tile instead of per (g, t) cell — every Omega* term is `s_group * KCov(Y_u1 - Y_v1, Y_u2 - Y_v2 | group)`, keyed only by wide-outcome columns, so the per-cell H(H+1)/2-pair tables dedup to distinct product columns per group (~26x fewer kernel GEMM columns on a PT-All fit), and the kernel weight matrices are built and freed one group at a time (the tile memory budget is governed by the largest single group rather than the sum, giving proportionally fatter tiles at large n). Each cell's Omega* is then gathered from the group tables in the same per-entry operation order as the per-cell construction (value-exact, locked by test); this is a pure implementation change — measured results move only at floating-point reassociation level (post-treatment cells ~1e-12 relative, overall ATT ~1e-13; the no-covariates path is byte-identical), and `omega_ridge=0` still routes the entire legacy path. **Rust-backend batched-Cholesky ridge solve (2026-07 follow-up):** on the Rust backend the batched ridge solve `(Omega* + lam * max(trace/H, 0) * I) x = 1` dispatches to a batched Cholesky kernel (the ridged Omega* is SPD by construction; a non-SPD row — measured zero on realistic panels — falls back to LU in-kernel, and any non-finite row is recomputed through the exact legacy numpy chain including its pseudoinverse backstop, so edge-case semantics are unchanged), parallelized over units. This too is a pure implementation change: Cholesky and LU solutions differ only at the condition-bounded floating-point level, and measured results move at reassociation level (post-treatment cells ~2e-12 relative, overall ATT ~1e-13); the pure-Python backend is byte-identical, and `omega_ridge=0` never reaches the kernel. - **Note:** Bootstrap aggregation uses fixed cohort-size weights for overall/event-study reaggregation, matching the CallawaySantAnna bootstrap pattern (staggered_bootstrap.py:281 computes `bootstrap_overall = bootstrap_atts_gt[:, post_indices] @ weights`; L297 uses the same fixed-weight pattern for event study). The analytical path includes a WIF correction; fixed-weight bootstrap captures the same sampling variability through per-cell EIF perturbation without re-estimating aggregation weights, consistent with both the library's CS implementation and the R `did` package. - **Overall ATT convention**: The library's `overall_att` uses cohort-size-weighted averaging of post-treatment (g,t) cells, matching the CallawaySantAnna simple aggregation. This differs from the paper's ES_avg (Eq 2.3), which uniformly averages over event-time horizons. ES_avg can be computed from event study output as `mean(event_study_effects[e]["effect"] for e >= 0)` diff --git a/docs/performance-plan.md b/docs/performance-plan.md index b6939e05..8ff210eb 100644 --- a/docs/performance-plan.md +++ b/docs/performance-plan.md @@ -326,7 +326,53 @@ The kernel-covariance stage itself: 21.2s -> 0.53s at 10k (40x). Deltas: post-treatment cells ~1e-12 rel (max 6e-12), overall ATT ~1e-13, nocov exactly 0. Post-hoist the top stages at every scale are `_ridge_solve_weights` (batched LAPACK LU over (n x H x H) stacks; 17.2s at 10k — Rust batched-Cholesky is the -natural follow-up, TODO row filed) and the nuisance sieves. +natural follow-up, resolved by the section below) and the nuisance sieves. + +### Rust batched-Cholesky ridge solve (2026-07 follow-up) + +Resolves the ridge-solve TODO row above. `_ridge_solve_weights`' batched +`np.linalg.solve` is a LAPACK dgesv gufunc — LU (2/3 H^3 flops) and SERIAL over +the batch — although the ridged Omega* is SPD (Cholesky = 1/3 H^3). Spike on +real captured batches (`.bench-local/edid_ridge_solve_spike.py`, cond_10k: 570 +calls, all H=59/60, m~1300-1740): numpy stage 15.6s extrapolated vs Rust kernel +1.6s = **9.6x kernel aggregate**; zero non-PD rows (min relative eigenvalue +1.7e-8 = the ridge floor), so the LU/NaN-poison fallback is pure defense. Raw +solve divergence Cholesky-vs-LU: max ~3.7e-7 per entry (cond ~6e7 x eps class) +— but fit-level deltas land at ~1e-12 because the weight perturbations live in +Omega*'s near-null directions (near-duplicate moments with near-identical +generated outcomes), exactly the insensitivity the ridge Note's "any fixed +weights summing to 1" argument predicts. + +Kernel-variant A/B (serial, per-H, m=1737): hand-rolled unblocked Cholesky in a +reused per-thread scratch vs faer `Llt::new` (which allocates a fresh Mat + +scratch per call): hand-rolled wins 10.3x at H=2 and 1.9x at H=8; faer wins +~1.4x at H=60 (SIMD factorization). Shipped hand-rolled (user-confirmed): zero +per-call allocation, wins where H is small, one oracle-locked code path; the +fit-level cost of not switching at H=60 is ~2% (the win is rayon parallelism +across the batch — Accelerate's serial dgesv is roughly at parity with the +serial hand-rolled Cholesky per solve). + +Measured arms (3-rep medians; BEFORE = pristine main 81d53f59 via +baseline-worktree PYTHONPATH; `.bench-local/edid_chol_arms_results.jsonl`): + +| scenario | BEFORE | AFTER | speedup | maxrss B->A | +|---|---|---|---|---| +| conditional n=2k | 7.08s | 4.29s | 1.65x | 0.59 -> 0.60 GB | +| conditional n=10k | 36.2s | 22.7s | 1.60x | 0.80 -> 0.75 GB | +| conditional n=30k | 129.0s | 90.1s | 1.43x | 1.09 -> 1.13 GB | +| conditional n=100k | 17.76 min (#629 measured, 1 rep) | 15.92 min (1 rep) | 1.12x | 4.02 -> 3.84 GB | +| nocov n=2k (guard) | 0.44s | 0.45s | 1.0x (path untouched) | flat | +| survey n=1k | 3.03s | 1.68s | 1.80x | flat | +| pure-python cond n=2k (guard) | 6.09s | 6.10s | byte-identical results | flat | + +Ridge stage at 10k: 17.1s -> 3.68s (4.65x; the aspirational >=5x stage gate was +narrowly missed — the kernel is 9.6x but ~2s of the remaining stage is shared +Python prep: the zero_mask full-stack abs scan + the `omega_stack[rest]` +fancy-index copy, ~28 GB extra memory traffic per 10k fit; TODO row filed for +the copy shave). Deltas: post cells max ~2e-12 rel, overall ATT ~1e-13, +event-study max ~5e-10, all cells incl. near-zero placebos <= 1.4e-11 abs; +pure-python arm byte-identical. After this the largest O(n) stage is the +sieve/nuisance construction outside the tiled pass (~9s at 10k). --- diff --git a/rust/src/batched_solve.rs b/rust/src/batched_solve.rs new file mode 100644 index 00000000..ad468df7 --- /dev/null +++ b/rust/src/batched_solve.rs @@ -0,0 +1,353 @@ +//! Batched ridge-regularized SPD solves for EfficientDiD per-unit weights. +//! +//! Solves `(A_i + ridge_i * I) x_i = 1` for a stack of small symmetric +//! matrices (H = 2..~60), one solve per unit per (g, t) cell. The matrices +//! are numerically PSD by construction (conditional Omega* covariances) and +//! the trace-scaled ridge makes them SPD, so an unblocked Cholesky +//! factorization (1/3 H^3 flops vs LU's 2/3 H^3) in a reused per-thread +//! scratch buffer is the primary path, parallelized over the batch with +//! rayon. Each matrix is solved independently in a fixed operation order, so +//! results are bit-deterministic across thread counts and batch splits (the +//! Python tile-invariance contract depends on this). +//! +//! The GIL is held throughout (PyReadonlyArray borrows cannot cross +//! `py.detach`; same pattern as `trop.rs`). Nothing in the parallel region +//! touches Python. + +use ndarray::{Array2, ArrayView2, ArrayViewMut1, Axis}; +use numpy::{PyArray2, PyReadonlyArray1, PyReadonlyArray3, ToPyArray}; +use pyo3::prelude::*; +use rayon::prelude::*; + +use faer::linalg::solvers::{PartialPivLu, Solve}; + +/// Solve `(A_i + ridge_i * I) x_i = 1` for a stack of symmetric matrices. +/// +/// # Arguments +/// * `a_stack` - (m, H, H) stack of symmetric (numerically PSD) matrices +/// * `ridge` - (m,) per-matrix ridge added to the diagonal (single addition +/// per diagonal entry - bit-identical to numpy's `a + r * I` there) +/// +/// # Returns +/// (m, H) solutions. Rows where the matrix is not SPD fall back to a faer +/// partial-pivot LU solve; a row whose LU has an exactly-zero pivot is +/// poisoned with NaN (mirroring LAPACK dgesv's exact-singularity signal so +/// the Python caller can route those rows through its legacy pinv path). +/// Non-finite values from a near-singular LU solve are returned as-is - the +/// Python caller treats any non-finite row as "recompute via legacy path". +/// +/// Degenerate shapes: m == 0 -> (0, H); H == 0 -> (m, 0). H == 1 solves via +/// the same Cholesky code (x = 1/sqrt(d)/sqrt(d), within 1 ulp of 1/d). +#[pyfunction] +#[pyo3(signature = (a_stack, ridge))] +pub fn batched_ridge_chol_solve_ones<'py>( + py: Python<'py>, + a_stack: PyReadonlyArray3<'py, f64>, + ridge: PyReadonlyArray1<'py, f64>, +) -> PyResult>> { + let a = a_stack.as_array(); + let r = ridge.as_array(); + let m = a.shape()[0]; + let h = a.shape()[1]; + + if a.shape()[2] != h { + return Err(PyErr::new::(format!( + "a_stack must be a stack of square matrices, got ({}, {}, {})", + m, + h, + a.shape()[2] + ))); + } + if r.len() != m { + return Err(PyErr::new::(format!( + "ridge length {} must match batch size {}", + r.len(), + m + ))); + } + + let mut out = Array2::::zeros((m, h)); + if m == 0 || h == 0 { + return Ok(out.to_pyarray(py)); + } + + // Bound per-task scheduling overhead: tiny matrices (H=2) get large + // chunks, big matrices (H=60, ~10us each) split freely. + let min_len = (4096 / (h * h)).max(1); + + out.axis_iter_mut(Axis(0)) + .into_par_iter() + .enumerate() + .zip(a.axis_iter(Axis(0)).into_par_iter()) + .with_min_len(min_len) + .for_each_init( + || (vec![0.0_f64; h * h], vec![0.0_f64; h]), + |(buf, x), ((i, mut out_row), a_i)| { + solve_one(a_i, r[i], buf, x, &mut out_row); + }, + ); + + Ok(out.to_pyarray(py)) +} + +/// Solve one (H, H) system into `out_row`, using the per-thread scratch. +fn solve_one( + a_i: ArrayView2<'_, f64>, + ridge_i: f64, + buf: &mut [f64], + x: &mut [f64], + out_row: &mut ArrayViewMut1<'_, f64>, +) { + let h = x.len(); + + // Copy into row-major scratch, adding the ridge on the diagonal. Every + // entry of buf is overwritten, so no state leaks between rows. + for row in 0..h { + for col in 0..h { + buf[row * h + col] = a_i[[row, col]]; + } + buf[row * h + row] += ridge_i; + } + + if cholesky_solve_ones_in_place(buf, x, h) { + for (o, v) in out_row.iter_mut().zip(x.iter()) { + *o = *v; + } + } else { + lu_fallback(a_i, ridge_i, out_row, h); + } +} + +/// Unblocked lower Cholesky in place on `buf` (row-major, lower triangle +/// written; upper triangle left as input values and never read), then +/// forward/back substitution against the implicit ones RHS into `x`. +/// Returns false without touching `x` if a pivot is not strictly positive. +// The negated comparison is the point: `!(d > 0.0)` is true for NaN while +// `d <= 0.0` is not, and NaN pivots must route to the LU fallback. +#[allow(clippy::neg_cmp_op_on_partial_ord)] +fn cholesky_solve_ones_in_place(buf: &mut [f64], x: &mut [f64], h: usize) -> bool { + for j in 0..h { + let mut d = buf[j * h + j]; + for k in 0..j { + d -= buf[j * h + k] * buf[j * h + k]; + } + // NaN-safe pivot check: must be `!(d > 0.0)`, NOT `d <= 0.0` - a NaN + // pivot compares false either way, and only this form routes NaN to + // the LU fallback instead of silently continuing with sqrt(NaN). + if !(d > 0.0) { + return false; + } + let l_jj = d.sqrt(); + buf[j * h + j] = l_jj; + for row in (j + 1)..h { + let mut v = buf[row * h + j]; + for k in 0..j { + v -= buf[row * h + k] * buf[j * h + k]; + } + buf[row * h + j] = v / l_jj; + } + } + + // Forward substitution: L y = 1. + for row in 0..h { + let mut v = 1.0_f64; + for k in 0..row { + v -= buf[row * h + k] * x[k]; + } + x[row] = v / buf[row * h + row]; + } + + // Back substitution: L^T z = y (reads L transposed). + for row in (0..h).rev() { + let mut v = x[row]; + for k in (row + 1)..h { + v -= buf[k * h + row] * x[k]; + } + x[row] = v / buf[row * h + row]; + } + + true +} + +/// Rare-path fallback for non-SPD matrices: faer partial-pivot LU. +/// +/// An exactly-zero U pivot poisons the row with NaN - LAPACK's dgesv raises +/// on exact zero pivots (numpy's LinAlgError -> legacy pinv route), while a +/// different elimination order could otherwise produce finite garbage here. +fn lu_fallback( + a_i: ArrayView2<'_, f64>, + ridge_i: f64, + out_row: &mut ArrayViewMut1<'_, f64>, + h: usize, +) { + let a_faer = faer::Mat::from_fn(h, h, |row, col| { + let v = a_i[[row, col]]; + if row == col { + v + ridge_i + } else { + v + } + }); + let lu = PartialPivLu::new(a_faer.as_ref()); + + let u = lu.U(); + for k in 0..h { + if u[(k, k)] == 0.0 { + for o in out_row.iter_mut() { + *o = f64::NAN; + } + return; + } + } + + let ones = faer::Mat::from_fn(h, 1, |_, _| 1.0); + let sol = lu.solve(&ones); + for (k, o) in out_row.iter_mut().enumerate() { + *o = sol[(k, 0)]; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use ndarray::{Array1, Array3}; + use rand::prelude::*; + use rand_xoshiro::Xoshiro256PlusPlus; + + /// Random SPD matrix A = B B^T + eps * I. + fn random_spd(h: usize, rng: &mut Xoshiro256PlusPlus, eps: f64) -> ndarray::Array2 { + let b = ndarray::Array2::from_shape_fn((h, h), |_| rng.random::() - 0.5); + let mut a = b.dot(&b.t()); + for j in 0..h { + a[[j, j]] += eps; + } + a + } + + /// Drive solve_one directly (the pyfunction wrapper needs Python). + fn solve_stack(a: &Array3, ridge: &Array1) -> Array2 { + let m = a.shape()[0]; + let h = a.shape()[1]; + let mut out = Array2::::zeros((m, h)); + let mut buf = vec![0.0_f64; h * h]; + let mut x = vec![0.0_f64; h]; + for (i, mut row) in out.axis_iter_mut(Axis(0)).enumerate() { + solve_one( + a.index_axis(Axis(0), i), + ridge[i], + &mut buf, + &mut x, + &mut row, + ); + } + out + } + + /// Oracle: faer LU on the same ridged matrix. + fn lu_oracle(a: &ndarray::Array2, ridge: f64) -> Vec { + let h = a.nrows(); + let a_faer = faer::Mat::from_fn( + h, + h, + |r, c| { + if r == c { + a[[r, c]] + ridge + } else { + a[[r, c]] + } + }, + ); + let lu = PartialPivLu::new(a_faer.as_ref()); + let ones = faer::Mat::from_fn(h, 1, |_, _| 1.0); + let sol = lu.solve(&ones); + (0..h).map(|k| sol[(k, 0)]).collect() + } + + #[test] + fn cholesky_matches_lu_oracle_on_random_spd() { + let mut rng = Xoshiro256PlusPlus::seed_from_u64(42); + for &h in &[1usize, 2, 3, 8, 30, 60] { + let m = 5; + let mut a = Array3::::zeros((m, h, h)); + for i in 0..m { + a.index_axis_mut(Axis(0), i) + .assign(&random_spd(h, &mut rng, 0.5)); + } + let ridge = Array1::from_elem(m, 1e-6); + let out = solve_stack(&a, &ridge); + for i in 0..m { + let oracle = lu_oracle(&a.index_axis(Axis(0), i).to_owned(), ridge[i]); + for k in 0..h { + let got = out[[i, k]]; + let want = oracle[k]; + let rel = (got - want).abs() / want.abs().max(1e-30); + assert!( + rel < 1e-10, + "h={} i={} k={} got={} want={} rel={}", + h, + i, + k, + got, + want, + rel + ); + } + } + } + } + + #[test] + fn nan_input_row_produces_non_finite_output() { + let mut a = Array3::::zeros((1, 3, 3)); + for j in 0..3 { + a[[0, j, j]] = 1.0; + } + a[[0, 1, 1]] = f64::NAN; + let ridge = Array1::from_elem(1, 1e-6); + let out = solve_stack(&a, &ridge); + assert!( + out.iter().any(|v| !v.is_finite()), + "NaN input must not produce an all-finite row: {:?}", + out + ); + } + + #[test] + fn exact_singular_zero_ridge_row_is_nan_poisoned() { + // diag(1, -1, 0): trace 0 -> zero ridge upstream; Cholesky fails at + // the -1 pivot; LU has an exactly-zero pivot -> poisoned row. + let mut a = Array3::::zeros((1, 3, 3)); + a[[0, 0, 0]] = 1.0; + a[[0, 1, 1]] = -1.0; + a[[0, 2, 2]] = 0.0; + let ridge = Array1::from_elem(1, 0.0); + let out = solve_stack(&a, &ridge); + assert!( + out.iter().all(|v| v.is_nan()), + "exact-singular row must be NaN-poisoned: {:?}", + out + ); + } + + #[test] + fn indefinite_full_rank_row_uses_lu_exactly() { + // diag(1, -1): Cholesky fails; LU gives exactly [1, -1]. + let mut a = Array3::::zeros((1, 2, 2)); + a[[0, 0, 0]] = 1.0; + a[[0, 1, 1]] = -1.0; + let ridge = Array1::from_elem(1, 0.0); + let out = solve_stack(&a, &ridge); + assert_eq!(out[[0, 0]], 1.0); + assert_eq!(out[[0, 1]], -1.0); + } + + #[test] + fn h_equals_one_matches_reciprocal() { + let mut a = Array3::::zeros((1, 1, 1)); + a[[0, 0, 0]] = 4.0; + let ridge = Array1::from_elem(1, 0.0); + let out = solve_stack(&a, &ridge); + let rel = (out[[0, 0]] - 0.25).abs() / 0.25; + assert!(rel < 1e-14, "got {}", out[[0, 0]]); + } +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 30b6312f..ceebb5b9 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -16,6 +16,7 @@ use std::collections::HashMap; #[cfg(feature = "alloc-profile")] mod alloc_profile; +mod batched_solve; mod bootstrap; mod demean; mod linalg; @@ -53,6 +54,12 @@ fn _rust_backend(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(linalg::solve_ols, m)?)?; m.add_function(wrap_pyfunction!(linalg::compute_robust_vcov, m)?)?; + // Batched ridge-regularized SPD solve (EfficientDiD per-unit weights) + m.add_function(wrap_pyfunction!( + batched_solve::batched_ridge_chol_solve_ones, + m + )?)?; + // TROP estimator acceleration (local method) m.add_function(wrap_pyfunction!(trop::compute_unit_distance_matrix, m)?)?; m.add_function(wrap_pyfunction!(trop::loocv_grid_search, m)?)?; diff --git a/tests/test_efficient_did.py b/tests/test_efficient_did.py index 91eb1cd1..d685f2a4 100644 --- a/tests/test_efficient_did.py +++ b/tests/test_efficient_did.py @@ -16,6 +16,8 @@ from edid_dgp import make_compustat_dgp from diff_diff import CallawaySantAnna, EDiD, EfficientDiD +from diff_diff._backend import HAS_RUST_BACKEND +from diff_diff._backend import _rust_batched_ridge_chol_solve as _ridge_chol_symbol from diff_diff.efficient_did_results import EfficientDiDResults from diff_diff.efficient_did_weights import ( enumerate_valid_triples, @@ -911,6 +913,155 @@ def test_one_ulp_stability_nocov(self): self._assert_close_surfaces(r_ulp, r_base, rtol=1e-6, atol=1e-9) +_requires_rust_chol = pytest.mark.skipif( + not HAS_RUST_BACKEND or _ridge_chol_symbol is None, + reason="Rust batched-Cholesky kernel not available", +) + + +class TestRidgeSolveRustDispatch: + """v3.8: `_ridge_solve_weights` dispatches to the Rust batched-Cholesky + kernel (`batched_ridge_chol_solve_ones`) on the rust backend. + + Dispatch is on availability + float64 dtype ONLY - no batch-size cutoff, + so the tile-invariance twins above (which force one-unit batches) stay + same-algorithm on both sides. Cholesky and LU legitimately differ at the + cond*eps level (~1e-7 on the near-singular ridged Omega*, cond ~1e6-1e8), + so cross-backend fit comparisons use loose tolerances while + row-level semantic contracts (legacy recompute of non-finite rows, + symbol-None degradation) are exact. + """ + + @staticmethod + def _spd_stack(m=40, H=6, seed=0): + rng = np.random.default_rng(seed) + b = rng.standard_normal((m, H, H)) + return b @ b.transpose(0, 2, 1) + 1.0 * np.eye(H) + + @staticmethod + def _legacy_weights(omega_stack, omega_ridge): + """`_ridge_solve_weights` with the rust symbol disabled.""" + import diff_diff.efficient_did_covariates as cov_mod + + orig = cov_mod._rust_batched_ridge_chol_solve + cov_mod._rust_batched_ridge_chol_solve = None + try: + return cov_mod._ridge_solve_weights(omega_stack, omega_ridge) + finally: + cov_mod._rust_batched_ridge_chol_solve = orig + + @_requires_rust_chol + def test_rust_kernel_called_on_conditional_fit(self, monkeypatch): + """The conditional covariate fit routes its per-cell solves through + the rust kernel when the backend is active.""" + import diff_diff.efficient_did_covariates as cov_mod + + calls = {"n": 0} + orig = cov_mod._rust_batched_ridge_chol_solve + + def spy(a_stack, ridge): + calls["n"] += 1 + return orig(a_stack, ridge) + + monkeypatch.setattr(cov_mod, "_rust_batched_ridge_chol_solve", spy) + df = _make_covariate_panel(n_units=80, n_periods=9, seed=42) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + EfficientDiD().fit( + df, + "y", + "unit", + "time", + "first_treat", + covariates=["x1", "x2"], + aggregate="all", + ) + assert calls["n"] > 0, "rust kernel never dispatched on the conditional fit" + + @_requires_rust_chol + def test_backend_ab_conditional_fit(self, monkeypatch): + """Rust-vs-legacy end-to-end fit A/B (in-process: the legacy arm + monkeypatches the dispatch symbol to None; DIFF_DIFF_BACKEND cannot + be flipped in-process). Tolerance is the Cholesky-vs-LU cond*eps + band on the ridged Omega* (~1e-7 raw, measured), NOT reassociation + noise - do not tighten toward the tile-twin 1e-10.""" + import diff_diff.efficient_did_covariates as cov_mod + + df = _make_covariate_panel(n_units=150, n_periods=9, seed=42) + fit_kwargs = dict(covariates=["x1", "x2"], aggregate="all") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r_rust = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", **fit_kwargs) + monkeypatch.setattr(cov_mod, "_rust_batched_ridge_chol_solve", None) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r_py = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", **fit_kwargs) + TestFusedConditionalPath._assert_close_surfaces(r_rust, r_py, rtol=1e-5, atol=1e-9) + + def test_symbol_none_degrades_to_legacy_exactly(self): + """With the dispatch symbol None (stale extension / pure-python + backend), `_ridge_solve_weights` must be the byte-identical legacy + chain - locks the mixed-version degradation contract.""" + om = self._spd_stack(m=25, H=5, seed=3) + got = self._legacy_weights(om, 1e-6) + + # hand-computed legacy chain (same expressions, same op order) + n, H, _ = om.shape + trace = np.trace(om, axis1=1, axis2=2) + ridge = 1e-6 * np.maximum(trace / H, 0.0) + om_ridged = om + ridge[:, None, None] * np.eye(H)[None] + num = np.linalg.solve(om_ridged, np.ones((n, H, 1)))[..., 0] + den = num @ np.ones(H) + expected = num / den[:, None] + np.testing.assert_array_equal(got, expected) + + @_requires_rust_chol + def test_bad_row_recompute_matches_legacy_exactly(self): + """A row the kernel flags non-finite (all-NaN omega) is recomputed + via the legacy numpy chain: THAT row must equal the legacy result + exactly. Good rows in the same batch keep Cholesky results and + legitimately differ from legacy at ~1e-15..1e-10 - they are only + checked loosely.""" + import diff_diff.efficient_did_covariates as cov_mod + + om = self._spd_stack(m=6, H=4, seed=7) + om[2] = np.nan # forces Cholesky fail -> LU NaN -> python recompute + w_rust = cov_mod._ridge_solve_weights(om, 1e-6) + w_py = self._legacy_weights(om, 1e-6) + + np.testing.assert_array_equal(w_rust[2], w_py[2]) + good = [0, 1, 3, 4, 5] + np.testing.assert_allclose(w_rust[good], w_py[good], rtol=1e-9, atol=1e-12) + + def test_exact_singular_pinv_arm(self): + """diag(1, -2, 0): trace < 0 -> zero ridge; exactly singular. Legacy: + batched solve raises -> per-row solve raises -> pinv gives + num = [1, -0.5, 0], den = 0.5 -> weights [2, -1, 0]. The rust path + must reach the identical pinv arm via NaN-poisoning + recompute + (a bare batched re-solve would crash instead).""" + import diff_diff.efficient_did_covariates as cov_mod + + om = np.zeros((1, 3, 3)) + om[0, 0, 0] = 1.0 + om[0, 1, 1] = -2.0 + w = cov_mod._ridge_solve_weights(om, 1e-6) + w_legacy = self._legacy_weights(om, 1e-6) + np.testing.assert_array_equal(w, w_legacy) + np.testing.assert_allclose(w[0], [2.0, -1.0, 0.0], atol=1e-12) + + def test_f32_stack_declines_dispatch(self): + """A float32 stack (reachable via the public compute_per_unit_weights) + must decline rust dispatch - no silent coercion - and match the + legacy path exactly.""" + from diff_diff.efficient_did_covariates import compute_per_unit_weights + + om64 = self._spd_stack(m=10, H=4, seed=11) + om32 = om64.astype(np.float32) + got = compute_per_unit_weights(om32, omega_ridge=1e-6) + expected = self._legacy_weights(om32, 1e-6) + np.testing.assert_array_equal(got, expected) + + class TestValidTriples: """Test enumerate_valid_triples with hand-worked examples.""" diff --git a/tests/test_rust_backend.py b/tests/test_rust_backend.py index 66337ed8..210ebb84 100644 --- a/tests/test_rust_backend.py +++ b/tests/test_rust_backend.py @@ -3369,3 +3369,141 @@ def test_invalid_values_fall_back_to_default(self, monkeypatch, bad): monkeypatch.setenv("DIFF_DIFF_DEMEAN_CHUNK_COLS", bad) assert utils_mod._resolve_demean_chunk_cols() is utils_mod._DEMEAN_MAP_CHUNK_COLS + + +from diff_diff._backend import ( # noqa: E402 + _rust_batched_ridge_chol_solve as _batched_chol_symbol, +) + + +@pytest.mark.skipif( + not HAS_RUST_BACKEND or _batched_chol_symbol is None, + reason="Rust backend or batched-Cholesky kernel not available", +) +class TestBatchedRidgeCholSolve: + """Rust `batched_ridge_chol_solve_ones` vs the numpy LU reference. + + Contract: solves (A_i + ridge_i * I) x = 1 per matrix via hand-rolled + Cholesky; non-SPD rows fall back to faer LU; an exactly-singular row is + NaN-poisoned (so the Python dispatch recomputes it via the legacy + chain). Cholesky-vs-LU parity is cond*eps-bounded, NOT bit-identical: + ~1e-12 on well-conditioned stacks, ~1e-5 budget on the cond~1e6-1e8 + ridged near-singular stacks production feeds it (measured ~4e-7 max on + real Omega* batches). Per-row op order is fixed, so results are + bit-deterministic across batch splits and thread counts. + """ + + @staticmethod + def _numpy_reference(a_stack, ridge): + m, h, _ = a_stack.shape + a_ridged = a_stack + ridge[:, None, None] * np.eye(h)[None] + return np.linalg.solve(a_ridged, np.ones((m, h, 1)))[..., 0] + + @staticmethod + def _spd_stack(m, h, seed, eps=1.0): + rng = np.random.default_rng(seed) + b = rng.standard_normal((m, h, h)) + return b @ b.transpose(0, 2, 1) + eps * np.eye(h) + + @pytest.mark.parametrize("h", [2, 3, 30, 60]) + @pytest.mark.parametrize("m", [1, 200]) + def test_well_conditioned_parity(self, h, m): + a = self._spd_stack(m, h, seed=h * 1000 + m) + ridge = 1e-6 * np.trace(a, axis1=1, axis2=2) / h + got = _batched_chol_symbol(a, ridge) + want = self._numpy_reference(a, ridge) + # atol covers near-zero solution entries: Cholesky-vs-LU error scales + # with the solution norm (cond*eps*||x||), not per-entry. + np.testing.assert_allclose(got, want, rtol=1e-12, atol=1e-13) + + def test_ill_conditioned_parity_cond_bounded(self): + """Production regime: numerically singular PSD + trace-scaled ridge + (floors relative eigenvalues at ~1e-8 -> cond ~1e8). Cholesky and LU + then differ at the cond*eps level; budget 1e-5 (~30x the measured + max on real Omega* stacks).""" + rng = np.random.default_rng(9) + m, h = 50, 20 + q, _ = np.linalg.qr(rng.standard_normal((h, h))) + eigs = np.logspace(0, -12, h) # exact-null tail beyond fp precision + a1 = (q * eigs) @ q.T + a = np.repeat(a1[None], m, axis=0) + 0.0 + # jitter each matrix a little to vary the batch (stay PSD) + jit = rng.standard_normal((m, h, h)) * 1e-9 + a = a + jit @ jit.transpose(0, 2, 1) + ridge = 1e-6 * np.trace(a, axis1=1, axis2=2) / h + got = _batched_chol_symbol(a, ridge) + want = self._numpy_reference(a, ridge) + assert np.isfinite(got).all() + np.testing.assert_allclose(got, want, rtol=1e-5, atol=1e-8) + + def test_nan_row_non_finite(self): + a = self._spd_stack(3, 4, seed=1) + a[1] = np.nan + ridge = np.full(3, 1e-6) + got = _batched_chol_symbol(a, ridge) + assert not np.isfinite(got[1]).all() + np.testing.assert_allclose(got[[0, 2]], self._numpy_reference(a, ridge)[[0, 2]], rtol=1e-12) + + def test_exact_singular_zero_ridge_nan_poisoned(self): + """diag(1, -2, 0) with zero ridge: Cholesky fails (negative pivot), + faer LU sees an exactly-zero U pivot -> whole row NaN (mirrors + LAPACK's exact-singularity signal; a finite-garbage row would + silently skip the dispatch's legacy pinv recompute).""" + a = np.zeros((1, 3, 3)) + a[0, 0, 0] = 1.0 + a[0, 1, 1] = -2.0 + got = _batched_chol_symbol(a, np.zeros(1)) + assert np.isnan(got).all() + + def test_indefinite_full_rank_lu_fallback(self): + """diag(1, -1) with zero ridge: not SPD but invertible - the LU + fallback returns the exact solution [1, -1].""" + a = np.zeros((1, 2, 2)) + a[0, 0, 0] = 1.0 + a[0, 1, 1] = -1.0 + got = _batched_chol_symbol(a, np.zeros(1)) + np.testing.assert_array_equal(got[0], [1.0, -1.0]) + + def test_batch_split_bit_identity(self): + """Full-stack result == concatenated sub-batch results, bitwise. + Locks the per-row-fixed-op-order determinism the EfficientDiD + tile-invariance twins depend on.""" + a = self._spd_stack(31, 12, seed=5) + ridge = 1e-6 * np.trace(a, axis1=1, axis2=2) / 12 + full = _batched_chol_symbol(a, ridge) + parts = np.vstack( + [ + _batched_chol_symbol(a[:7], ridge[:7]), + _batched_chol_symbol(a[7:20], ridge[7:20]), + _batched_chol_symbol(a[20:], ridge[20:]), + ] + ) + np.testing.assert_array_equal(full, parts) + + def test_degenerate_shapes(self): + """m=0 and H=0 are no-ops with the right shape; H=1 is the scalar + 1/(a+ridge) via the 1x1 factorization (within 1 ulp).""" + out_m0 = _batched_chol_symbol(np.zeros((0, 4, 4)), np.zeros(0)) + assert out_m0.shape == (0, 4) + out_h0 = _batched_chol_symbol(np.zeros((3, 0, 0)), np.zeros(3)) + assert out_h0.shape == (3, 0) + out_h1 = _batched_chol_symbol(np.full((1, 1, 1), 4.0), np.zeros(1)) + np.testing.assert_allclose(out_h1, [[0.25]], rtol=1e-14) + + def test_strided_input_defensive(self): + """Non-contiguous views produce identical results to a contiguous + copy (defensive: the live dispatch path's fancy-indexed stacks are + always C-contiguous).""" + a = self._spd_stack(20, 5, seed=8) + ridge = np.full(20, 1e-6) + strided = a[::2] + assert not strided.flags["C_CONTIGUOUS"] + got = _batched_chol_symbol(strided, ridge[::2]) + want = _batched_chol_symbol(np.ascontiguousarray(strided), np.ascontiguousarray(ridge[::2])) + np.testing.assert_array_equal(got, want) + + def test_shape_validation_errors(self): + with pytest.raises(ValueError, match="square"): + _batched_chol_symbol(np.zeros((2, 3, 4)), np.zeros(2)) + with pytest.raises(ValueError, match="ridge length"): + _batched_chol_symbol(np.zeros((2, 3, 3)), np.zeros(5))