From 4823b037775188331739e41f8179095df45c46eb Mon Sep 17 00:00:00 2001 From: igerber Date: Wed, 8 Jul 2026 01:14:08 -0400 Subject: [PATCH 1/2] perf(linalg): scores-based CR2 Bell-McCaffrey DOF (PT2018 Appendix B) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The unweighted per-contrast Satterthwaite DOF materialized the dense n x n residual-maker M = I - X(X'X)^-1 X' and contracted it over all cluster pairs (O(n^2) time per contrast, O(n^2) memory — 3.2 GB at n=20k). With Omega the (n, G) matrix stacking omega_g on disjoint cluster supports, B = Omega' M Omega collapses to diag(||omega_g||^2) - P' M_U P with P = X' Omega: O(nk + G^2 k) per contrast, no n x n allocation. ~32x at n=5k/G=50 (0.57s -> 0.018s); n=20k/G=100 in 0.12s. Algebraically identical — frozen pair-loop oracle parity at rtol 1e-10 (balanced, unbalanced, compound contrasts); both NaN-reliability guards unchanged; 432 consumer tests pass unmodified. REGISTRY notes updated (per-fit setup wording + Appendix-B evaluation sentence). Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 10 ++++ TODO.md | 1 - diff_diff/linalg.py | 61 ++++++++++++++---------- docs/methodology/REGISTRY.md | 6 +-- tests/test_linalg.py | 92 ++++++++++++++++++++++++++++++++++++ 5 files changed, 142 insertions(+), 28 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 09b45f9c..61798c13 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -489,6 +489,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 h=1 clamp-parity test). The symbol is imported independently (mixed-version safe — a stale extension degrades HC2 to NumPy without disabling older Rust accelerations). `return_dof` / weighted / CR2-BM requests stay on NumPy (CR2-BM tracked in TODO). +- **CR2 Bell-McCaffrey Satterthwaite DOF: scores-based evaluation (PT2018 Appendix B).** + The unweighted per-contrast DOF previously materialized the dense `n×n` residual-maker + `M = I − X(X'X)⁻¹X'` and contracted it over all cluster pairs (`O(n²)` time per + contrast, `O(n²)` memory — 3.2 GB at n=20k). The pairwise matrix now collapses to + `B = diag(‖ω_g‖²) − P'M_U P` with `P = X'Ω` (disjoint cluster supports), costing + `O(nk + G²k)` per contrast with no `n×n` allocation: ~32x at n=5k/G=50 (0.57s→0.018s); + n=20k/G=100 completes in 0.12s where the old path would allocate 3.2 GB. Algebraically + identical — agreement with a frozen pair-loop oracle at rtol 1e-10 (balanced, + unbalanced, and compound-contrast designs); both NaN-reliability guards (noise floor, + cluster-count bound) unchanged. All 432 consumer tests pass unmodified. - **`CallawaySantAnna` per-(g,t) IF scatters converted from `np.add.at` to fancy `+=`** (`staggered.py::_cluster_robust_se_from_per_gt_if` — runs once per (g,t) cell when `cluster=` is set — and the general combined-IF assembly path in diff --git a/TODO.md b/TODO.md index 97444378..4727aa8d 100644 --- a/TODO.md +++ b/TODO.md @@ -46,7 +46,6 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | `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 CR2 Bell-McCaffrey: falls through to NumPy (the leverage/Satterthwaite-DOF path needs `return_dof` support, which the Rust vcov dispatch excludes). The one-way HC2 kernel landed 2026-07-07 (`compute_robust_vcov_hc2`, mirrors the NumPy hc2 branch at ~1e-15; near-singular hat-diagonal sentinel + Python-side warn-and-HC1-fallback). | `rust/src/linalg.rs` | Phase 1a | Mid | Low | | Wild cluster bootstrap CI inversion calls `_t_star(r)` ~O(100) times, each materializing a fresh `(B×n)` `y_star` + `(k×B)` refit + `(n×B)` residual arrays. Acceptable for the few-cluster regime; for large-`n`/large-`B`, chunk `_t_star` over draws or precompute the `r`-independent cluster-level pieces (restricted residuals are linear in `r`). | `utils.py::wild_bootstrap_se._t_star` | #543 | Mid | Low | | `SpilloverDiD` sparse cKDTree path for the staggered nearest-treated-distance helper (mirrors the static helper's sparse branch). `_compute_nearest_treated_distance_staggered` always builds dense `(n_units, n_treated_by_onset)` matrices per cohort; add a sparse branch gated on `n > _CONLEY_SPARSE_N_THRESHOLD`. | `spillover.py` | Wave B | Mid | Low | diff --git a/diff_diff/linalg.py b/diff_diff/linalg.py index b0b70e09..267ca592 100644 --- a/diff_diff/linalg.py +++ b/diff_diff/linalg.py @@ -2068,10 +2068,7 @@ def _compute_cr2_bm_vcov_and_dof( # `(tr B)² / tr(B²)` form (bit-equal to prior); weighted uses the full # clubSandwich P_array construction. if weights is None: - # Build the symmetric residual-maker M = I - H for the simple formula. - H = X @ M_U @ X.T - M = np.eye(n) - H - dof_vec = _cr2_bm_dof_inner(X, M, A_g_matrices, cluster_idx, M_U, contrasts) + dof_vec = _cr2_bm_dof_inner(X, A_g_matrices, cluster_idx, M_U, contrasts) else: dof_vec = _cr2_bm_dof_inner_weighted( X, @@ -2154,7 +2151,6 @@ def _compute_cr2_bm( def _cr2_bm_dof_inner( X: np.ndarray, - M: np.ndarray, A_g_matrices: Dict[Any, np.ndarray], cluster_idx: Dict[Any, np.ndarray], bread_inv: np.ndarray, @@ -2163,10 +2159,10 @@ def _cr2_bm_dof_inner( """Inner DOF loop, parameterized by an arbitrary contrast matrix. Computes the CR2 Bell-McCaffrey Satterthwaite DOF for each column of - ``contrasts`` (shape ``(k, m)``), using the precomputed residual-maker - ``M``, per-cluster adjustment matrices ``A_g_matrices``, cluster index - map ``cluster_idx``, and ``bread_inv``. The per-coefficient case is - recovered with ``contrasts=np.eye(k)``; compound contrasts (e.g., a + ``contrasts`` (shape ``(k, m)``), using the per-cluster adjustment + matrices ``A_g_matrices``, cluster index map ``cluster_idx``, and + ``bread_inv``. The per-coefficient case is recovered with + ``contrasts=np.eye(k)``; compound contrasts (e.g., a post-period-average ATT) are handled by the same algebra without duplication. @@ -2178,6 +2174,23 @@ def _cr2_bm_dof_inner( trace_B2 = sum_{g, h} (omega_g' M_{g, h} omega_h)**2 DOF(c) = trace_B**2 / trace_B2 + SCORES-BASED EVALUATION (Pustejovsky-Tipton 2018 Appendix B): the + cluster-pair contraction is never evaluated against an explicit + residual-maker. With ``Omega`` the ``(n, G)`` matrix stacking the + ``omega_g`` on their (disjoint) cluster supports and + ``M = I - X bread_inv X'``, the pairwise matrix + ``B[g, h] = omega_g' M_{g, h} omega_h`` collapses to + + B = Omega' M Omega = diag(||omega_g||^2) - P' bread_inv P, + P = X' Omega (k, G; column g is X_g' omega_g) + + so ``trace_B2 = ||B||_F^2`` costs ``O(n k + G^2 k)`` per contrast with + ``O(G k)`` memory — the previous form materialized the dense ``n x n`` + ``M`` and looped cluster pairs at ``O(n^2)`` per contrast, the exact + large-``n`` blowup the TODO row tracked. The two evaluations are + algebraically identical; floating-point agreement is ~1e-12 relative + (different accumulation order), locked by the frozen-oracle parity test. + Returns ------- dof_vec : ndarray of shape (m,) @@ -2186,6 +2199,7 @@ def _cr2_bm_dof_inner( """ m = contrasts.shape[1] unique_clusters = list(cluster_idx.keys()) + n_g_clusters = len(unique_clusters) # Precompute once: q-matrix (n, m) and A_g_Xbi (n_g, k) per cluster. # For unit-contrast inputs (contrasts=I_k), this matches the prior # inline implementation exactly: q[:, j] == X_bi[:, j] == X @ bread_inv @ e_j. @@ -2197,6 +2211,15 @@ def _cr2_bm_dof_inner( # Omega per cluster per contrast: (n_g, m) = A_g_Xbi[g] @ contrasts omega_all = {g: A_g_Xbi[g] @ contrasts for g in unique_clusters} + # Scores-based precomputes (see docstring): per cluster g, + # normsq[g, j] = ||omega_g^{(j)}||^2 and P_all[g, :, j] = X_g' omega_g^{(j)}. + normsq = np.zeros((n_g_clusters, m)) + P_all = np.zeros((n_g_clusters, X.shape[1], m)) + for gi, g in enumerate(unique_clusters): + om = omega_all[g] # (n_g, m) + normsq[gi] = np.einsum("ij,ij->j", om, om) + P_all[gi] = X[cluster_idx[g]].T @ om + dof_vec = np.empty(m) # Retain max|B_{g,h}| per contrast so we can NaN-guard noise-floor # degeneracies in a second pass (mirrors `_cr2_bm_dof_inner_weighted`). @@ -2204,21 +2227,11 @@ def _cr2_bm_dof_inner( for j in range(m): q = Q[:, j] trace_B = float(np.sum(q * q)) - trace_B2 = 0.0 - max_abs_B = 0.0 - omega_cache = {g: omega_all[g][:, j] for g in unique_clusters} - for g in unique_clusters: - idx_g = cluster_idx[g] - omega_g = omega_cache[g] - for h in unique_clusters: - idx_h = cluster_idx[h] - omega_h = omega_cache[h] - M_gh = M[np.ix_(idx_g, idx_h)] - val = float(omega_g @ M_gh @ omega_h) - trace_B2 += val * val - if abs(val) > max_abs_B: - max_abs_B = abs(val) - max_abs_B_arr[j] = max_abs_B + P_j = P_all[:, :, j] # (G, k) + B_j = -(P_j @ bread_inv @ P_j.T) + B_j[np.diag_indices_from(B_j)] += normsq[:, j] + trace_B2 = float(np.sum(B_j * B_j)) + max_abs_B_arr[j] = float(np.max(np.abs(B_j))) if B_j.size else 0.0 dof_vec[j] = (trace_B * trace_B) / trace_B2 if trace_B2 > 0 else np.nan # Noise-floor NaN-guard (unweighted analogue of the guard in diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index c99f4c1f..41c4df3e 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -238,12 +238,12 @@ where V is the VCV sub-matrix for post-treatment δ_e coefficients. per-cluster adjustment matrices, the sandwich vcov, AND the per-coefficient + post-period-average contrast DOF together from a single shared precompute build (`_compute_cr2_bm_vcov_and_dof` in `diff_diff/linalg.py`); the fit bypasses - `solve_ols`'s separate vcov dispatch on this path so the O(n²) CR2 setup is built - once per fit rather than twice. `_compute_cr2_bm` (per-coefficient vcov + DOF) and + `solve_ols`'s separate vcov dispatch on this path so the CR2 setup (per-cluster + adjustment blocks) is built once per fit rather than twice. `_compute_cr2_bm` (per-coefficient vcov + DOF) and `_compute_cr2_bm_contrast_dof` (DOF-only for arbitrary contrasts) are thin wrappers over that shared core, so every CR2 caller routes through one implementation. The consolidation is bit-identical to the prior two-call path (proven at atol=0). -- **Note (unweighted per-coef DOF guard):** the unweighted clustered CR2-BM per-coefficient DOF (`_cr2_bm_dof_inner`, the simple `(tr B)²/tr(B²)` form) carries the same two-part reliability guard as the weighted P-array path. (1) **Noise floor:** for a high-leverage FE-dummy / collinear nuisance column `trace_B2 = Σ B_{g,h}²` collapses to float64 accumulation noise while `trace_B` stays O(1), inflating the ratio to a non-physical DOF (observed ~1e61 on the absorbed-FE golden); a contrast whose `max|B_{g,h}|` sits below the batch-relative (`1e-10×max`, computed on the scale-normalized `max|B|/‖c‖²` since `B ∝ ‖c‖²` while the DOF is scale-invariant) or absolute (`(EPS·n·k·bread_scale)²`) floor is NaN'd. (2) **Cluster-count bound:** the Bell-McCaffrey Satterthwaite DOF is `(tr B)²/tr(B²)` with `B` PSD and cluster-structured, so it is bounded by `rank(B) ≤ G` (number of clusters); the simple unweighted form is numerically less faithful than clubSandwich's P-array form on high-leverage columns and can return a finite-but-inflated DOF above `G` (observed ~32.7 and ~16.3 vs R's 6 and 3, `G=8`), which is NaN'd as non-physical. The well-conditioned contrasts estimators consume — the treatment effect, event-study coefficients, and the compound post-period-average — are unaffected and match R clubSandwich; only the non-user-facing high-leverage nuisance DOFs are suppressed (exact P-array reproduction of those is deferred). A `UserWarning` fires per fit. Regression: `tests/test_estimators_vcov_type.py::TestDiDAbsorbedFERParity::test_unweighted_cr2_bm_per_coef_dof_no_nonphysical`. +- **Note (unweighted per-coef DOF guard):** the unweighted clustered CR2-BM per-coefficient DOF (`_cr2_bm_dof_inner`, the simple `(tr B)²/tr(B²)` form) carries the same two-part reliability guard as the weighted P-array path. (1) **Noise floor:** for a high-leverage FE-dummy / collinear nuisance column `trace_B2 = Σ B_{g,h}²` collapses to float64 accumulation noise while `trace_B` stays O(1), inflating the ratio to a non-physical DOF (observed ~1e61 on the absorbed-FE golden); a contrast whose `max|B_{g,h}|` sits below the batch-relative (`1e-10×max`, computed on the scale-normalized `max|B|/‖c‖²` since `B ∝ ‖c‖²` while the DOF is scale-invariant) or absolute (`(EPS·n·k·bread_scale)²`) floor is NaN'd. (2) **Cluster-count bound:** the Bell-McCaffrey Satterthwaite DOF is `(tr B)²/tr(B²)` with `B` PSD and cluster-structured, so it is bounded by `rank(B) ≤ G` (number of clusters); the simple unweighted form is numerically less faithful than clubSandwich's P-array form on high-leverage columns and can return a finite-but-inflated DOF above `G` (observed ~32.7 and ~16.3 vs R's 6 and 3, `G=8`), which is NaN'd as non-physical. The well-conditioned contrasts estimators consume — the treatment effect, event-study coefficients, and the compound post-period-average — are unaffected and match R clubSandwich; only the non-user-facing high-leverage nuisance DOFs are suppressed (exact P-array reproduction of those is deferred). A `UserWarning` fires per fit. Regression: `tests/test_estimators_vcov_type.py::TestDiDAbsorbedFERParity::test_unweighted_cr2_bm_per_coef_dof_no_nonphysical`. **Evaluation (2026-07):** the pairwise `B` matrix is computed scores-based per Pustejovsky-Tipton (2018) Appendix B — `B = diag(‖ω_g‖²) − P' M_U P` with `P = X'Ω` — instead of contracting against an explicit dense `n×n` residual-maker (`O(n k + G² k)` per contrast, no `n×n` allocation; algebraically identical, ~1e-12 float agreement locked by a frozen-oracle parity test). - **Note:** `LinearRegression.get_se()` / `get_inference()` clamp the vcov diagonal at 0 before `sqrt`. A high-leverage / degenerate coefficient (an absorbed-FE dummy near-collinear with the treatment, whose Satterthwaite DOF already hits the noise-floor diff --git a/tests/test_linalg.py b/tests/test_linalg.py index c19cebed..40aaf7b0 100644 --- a/tests/test_linalg.py +++ b/tests/test_linalg.py @@ -2714,3 +2714,95 @@ def test_return_dropped_mask(self): # Rank 0 -> every coordinate dropped. _, nd0, _, dropped0 = _rank_guarded_inv(np.zeros((3, 3)), return_dropped=True) assert nd0 == 3 and dropped0.all() + + +class TestCR2BMScoresBasedDOF: + """Frozen-oracle parity for the scores-based Satterthwaite DOF evaluation + (PT2018 Appendix B): B = diag(||omega_g||^2) - P' M_U P must agree with + the previous explicit-residual-maker cluster-pair contraction to float64 + accumulation tolerance, without materializing the n x n M.""" + + @staticmethod + def _oracle_pairloop_dof(X, cluster_ids, bread_matrix, contrasts): + """The PREVIOUS implementation, frozen: dense M = I - H + explicit + cluster-pair loop (no NaN guards — raw trace ratio).""" + from diff_diff.linalg import _cr2_adjustment_matrix + + n, k = X.shape + bread_inv = np.linalg.solve(bread_matrix, np.eye(k)) + unique = list(np.unique(cluster_ids)) + idx = {g: np.where(cluster_ids == g)[0] for g in unique} + S_W = bread_matrix + MUWTWUM = bread_inv @ S_W @ bread_inv + A = {} + for g in unique: + X_g = X[idx[g]] + H_gg = X_g @ bread_inv @ X_g.T + G_g = np.eye(len(idx[g])) - H_gg - H_gg.T + X_g @ MUWTWUM @ X_g.T + A[g] = _cr2_adjustment_matrix(G_g) + M = np.eye(n) - X @ bread_inv @ X.T + m = contrasts.shape[1] + out = np.empty(m) + for j in range(m): + c = contrasts[:, j] + q = X @ bread_inv @ c + tr_B = float(np.sum(q * q)) + omega = {g: A[g] @ X[idx[g]] @ bread_inv @ c for g in unique} + tr_B2 = 0.0 + for g in unique: + for h in unique: + val = float(omega[g] @ M[np.ix_(idx[g], idx[h])] @ omega[h]) + tr_B2 += val * val + out[j] = (tr_B * tr_B) / tr_B2 if tr_B2 > 0 else np.nan + return out + + def test_scores_based_matches_pairloop_oracle(self): + from diff_diff.linalg import _compute_cr2_bm + + rng = np.random.default_rng(11) + n, k, G = 240, 4, 12 + X = np.column_stack([np.ones(n), rng.normal(size=(n, k - 1))]) + cl = np.repeat(np.arange(G), n // G) + beta = rng.normal(size=k) + y = X @ beta + rng.normal(size=n) * (1 + 0.3 * np.abs(X[:, 1])) + resid = y - X @ np.linalg.lstsq(X, y, rcond=None)[0] + bread = X.T @ X + + vcov, dof = _compute_cr2_bm(X, resid, cl, bread) + oracle = self._oracle_pairloop_dof(X, cl, bread, np.eye(k)) + np.testing.assert_allclose(dof, oracle, rtol=1e-10) + assert np.all(np.isfinite(dof)) and np.all(dof > 0) and np.all(dof <= G) + + def test_compound_contrast_matches_oracle(self): + from diff_diff.linalg import _compute_cr2_bm_contrast_dof + + rng = np.random.default_rng(13) + n, k, G = 180, 5, 9 + X = np.column_stack([np.ones(n), rng.normal(size=(n, k - 1))]) + cl = np.repeat(np.arange(G), n // G) + bread = X.T @ X + # Compound averaging contrast (post-period-average style) + a unit vector. + c1 = np.zeros(k) + c1[2:] = 1.0 / (k - 2) + c2 = np.zeros(k) + c2[1] = 1.0 + contrasts = np.column_stack([c1, c2]) + dof = _compute_cr2_bm_contrast_dof(X, cl, bread, contrasts) + oracle = self._oracle_pairloop_dof(X, cl, bread, contrasts) + np.testing.assert_allclose(dof, oracle, rtol=1e-10) + + def test_unbalanced_clusters_match_oracle(self): + from diff_diff.linalg import _compute_cr2_bm + + rng = np.random.default_rng(17) + sizes = [5, 40, 12, 3, 25, 60, 8, 18] + cl = np.concatenate([np.full(sz, g) for g, sz in enumerate(sizes)]) + n = len(cl) + k = 3 + X = np.column_stack([np.ones(n), rng.normal(size=(n, k - 1))]) + y = X @ rng.normal(size=k) + rng.normal(size=n) + resid = y - X @ np.linalg.lstsq(X, y, rcond=None)[0] + bread = X.T @ X + _, dof = _compute_cr2_bm(X, resid, cl, bread) + oracle = self._oracle_pairloop_dof(X, cl, bread, np.eye(k)) + np.testing.assert_allclose(dof, oracle, rtol=1e-10) From 47d22ddf5dd39948e8843f0e460e4ef6286ff9b6 Mon Sep 17 00:00:00 2001 From: igerber Date: Wed, 8 Jul 2026 07:20:16 -0400 Subject: [PATCH 2/2] perf+docs(linalg): contrast-chunk the CR2-BM product buffers (review P2) + doc syncs (P3s) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P2: P_all was (G, k, m) — O(G*k^2) on the batched per-coefficient sweep (contrasts=eye(k)), contradicting the stated O(Gk) bound. The per-cluster product buffers are now contrast-chunked under a 64 MB module-constant cap (monkeypatchable); each contrast's B is computed independently so chunking is bit-identical (locked by a forced one-contrast-chunk test with assert_array_equal). P3s: stale residual-maker-M mention removed from the shared-core docstring; CHANGELOG + REGISTRY memory claims updated. Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 4 ++- diff_diff/linalg.py | 58 ++++++++++++++++++++++++------------ docs/methodology/REGISTRY.md | 2 +- tests/test_linalg.py | 20 +++++++++++++ 4 files changed, 63 insertions(+), 21 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 61798c13..3dae11f8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -494,7 +494,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `M = I − X(X'X)⁻¹X'` and contracted it over all cluster pairs (`O(n²)` time per contrast, `O(n²)` memory — 3.2 GB at n=20k). The pairwise matrix now collapses to `B = diag(‖ω_g‖²) − P'M_U P` with `P = X'Ω` (disjoint cluster supports), costing - `O(nk + G²k)` per contrast with no `n×n` allocation: ~32x at n=5k/G=50 (0.57s→0.018s); + `O(nk + G²k)` per contrast with no `n×n` allocation (per-cluster product buffers + contrast-chunked under a 64 MB cap, bit-identical across chunk counts, so the batched + per-coefficient sweep never allocates `O(G·k²)` at once): ~32x at n=5k/G=50 (0.57s→0.018s); n=20k/G=100 completes in 0.12s where the old path would allocate 3.2 GB. Algebraically identical — agreement with a frozen pair-loop oracle at rtol 1e-10 (balanced, unbalanced, and compound-contrast designs); both NaN-reliability guards (noise floor, diff --git a/diff_diff/linalg.py b/diff_diff/linalg.py index 267ca592..e0223ede 100644 --- a/diff_diff/linalg.py +++ b/diff_diff/linalg.py @@ -1917,8 +1917,8 @@ def _compute_cr2_bm_vcov_and_dof( Both :func:`_compute_cr2_bm` (per-coefficient vcov + DOF) and :func:`_compute_cr2_bm_contrast_dof` (DOF-only for arbitrary contrasts) are thin wrappers over this function, so the expensive precomputes - (``bread_inv``, ``S_W``, ``MUWTWUM``, the per-cluster ``A_g`` eigendecompositions, - and the unweighted residual-maker ``M``) are defined in exactly one place. + (``bread_inv``, ``S_W``, ``MUWTWUM``, and the per-cluster ``A_g`` + eigendecompositions) are defined in exactly one place. Consolidating the two formerly-duplicated precompute blocks lets a caller that needs both vcov and contrast DOF (e.g. :class:`MultiPeriodDiD` under ``cluster + hc2_bm``) build them once instead of twice. @@ -2149,6 +2149,16 @@ def _compute_cr2_bm( return vcov, dof_vec +# Contrast-chunk byte budget for the scores-based CR2-BM DOF pass: bounds the +# (G, k, chunk) per-cluster product buffer so a batched per-coefficient sweep +# (contrasts=eye(k), i.e. m == k) cannot allocate O(G*k*m) at once on +# full-dummy / absorbed-FE designs with many clusters and coefficients. +# Results are chunk-count invariant BIT-FOR-BIT: each contrast's B matrix is +# computed independently, so chunking never reassociates a contrast's sums. +# Module-level so tests can monkeypatch it to force the multi-chunk path. +_CR2_BM_CONTRAST_CHUNK_BYTES = 64 * 1024 * 1024 + + def _cr2_bm_dof_inner( X: np.ndarray, A_g_matrices: Dict[Any, np.ndarray], @@ -2185,7 +2195,8 @@ def _cr2_bm_dof_inner( P = X' Omega (k, G; column g is X_g' omega_g) so ``trace_B2 = ||B||_F^2`` costs ``O(n k + G^2 k)`` per contrast with - ``O(G k)`` memory — the previous form materialized the dense ``n x n`` + ``O(G k)`` memory per contrast (contrast-chunked buffers bounded by + ``_CR2_BM_CONTRAST_CHUNK_BYTES``, bit-identical across chunk counts) — the previous form materialized the dense ``n x n`` ``M`` and looped cluster pairs at ``O(n^2)`` per contrast, the exact large-``n`` blowup the TODO row tracked. The two evaluations are algebraically identical; floating-point agreement is ~1e-12 relative @@ -2213,26 +2224,35 @@ def _cr2_bm_dof_inner( # Scores-based precomputes (see docstring): per cluster g, # normsq[g, j] = ||omega_g^{(j)}||^2 and P_all[g, :, j] = X_g' omega_g^{(j)}. - normsq = np.zeros((n_g_clusters, m)) - P_all = np.zeros((n_g_clusters, X.shape[1], m)) - for gi, g in enumerate(unique_clusters): - om = omega_all[g] # (n_g, m) - normsq[gi] = np.einsum("ij,ij->j", om, om) - P_all[gi] = X[cluster_idx[g]].T @ om - + k_X = X.shape[1] dof_vec = np.empty(m) # Retain max|B_{g,h}| per contrast so we can NaN-guard noise-floor # degeneracies in a second pass (mirrors `_cr2_bm_dof_inner_weighted`). max_abs_B_arr = np.zeros(m) - for j in range(m): - q = Q[:, j] - trace_B = float(np.sum(q * q)) - P_j = P_all[:, :, j] # (G, k) - B_j = -(P_j @ bread_inv @ P_j.T) - B_j[np.diag_indices_from(B_j)] += normsq[:, j] - trace_B2 = float(np.sum(B_j * B_j)) - max_abs_B_arr[j] = float(np.max(np.abs(B_j))) if B_j.size else 0.0 - dof_vec[j] = (trace_B * trace_B) / trace_B2 if trace_B2 > 0 else np.nan + # Chunk the contrasts so the per-cluster product buffer is (G, k, c) with + # c bounded by _CR2_BM_CONTRAST_CHUNK_BYTES — a full-m buffer would be + # O(G*k*m), i.e. O(G*k^2) on the batched per-coefficient sweep. Each + # contrast's B is computed independently, so chunking is bit-identical. + chunk = max(1, int(_CR2_BM_CONTRAST_CHUNK_BYTES // max(n_g_clusters * k_X * 8, 1))) + for c0 in range(0, m, chunk): + c1 = min(c0 + chunk, m) + width = c1 - c0 + normsq = np.zeros((n_g_clusters, width)) + P_all = np.zeros((n_g_clusters, k_X, width)) + for gi, g in enumerate(unique_clusters): + om = omega_all[g][:, c0:c1] # (n_g, width) + normsq[gi] = np.einsum("ij,ij->j", om, om) + P_all[gi] = X[cluster_idx[g]].T @ om + for jj in range(width): + j = c0 + jj + q = Q[:, j] + trace_B = float(np.sum(q * q)) + P_j = P_all[:, :, jj] # (G, k) + B_j = -(P_j @ bread_inv @ P_j.T) + B_j[np.diag_indices_from(B_j)] += normsq[:, jj] + trace_B2 = float(np.sum(B_j * B_j)) + max_abs_B_arr[j] = float(np.max(np.abs(B_j))) if B_j.size else 0.0 + dof_vec[j] = (trace_B * trace_B) / trace_B2 if trace_B2 > 0 else np.nan # Noise-floor NaN-guard (unweighted analogue of the guard in # `_cr2_bm_dof_inner_weighted`). For a high-leverage FE-dummy / collinear diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 41c4df3e..89359262 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -243,7 +243,7 @@ where V is the VCV sub-matrix for post-treatment δ_e coefficients. `_compute_cr2_bm_contrast_dof` (DOF-only for arbitrary contrasts) are thin wrappers over that shared core, so every CR2 caller routes through one implementation. The consolidation is bit-identical to the prior two-call path (proven at atol=0). -- **Note (unweighted per-coef DOF guard):** the unweighted clustered CR2-BM per-coefficient DOF (`_cr2_bm_dof_inner`, the simple `(tr B)²/tr(B²)` form) carries the same two-part reliability guard as the weighted P-array path. (1) **Noise floor:** for a high-leverage FE-dummy / collinear nuisance column `trace_B2 = Σ B_{g,h}²` collapses to float64 accumulation noise while `trace_B` stays O(1), inflating the ratio to a non-physical DOF (observed ~1e61 on the absorbed-FE golden); a contrast whose `max|B_{g,h}|` sits below the batch-relative (`1e-10×max`, computed on the scale-normalized `max|B|/‖c‖²` since `B ∝ ‖c‖²` while the DOF is scale-invariant) or absolute (`(EPS·n·k·bread_scale)²`) floor is NaN'd. (2) **Cluster-count bound:** the Bell-McCaffrey Satterthwaite DOF is `(tr B)²/tr(B²)` with `B` PSD and cluster-structured, so it is bounded by `rank(B) ≤ G` (number of clusters); the simple unweighted form is numerically less faithful than clubSandwich's P-array form on high-leverage columns and can return a finite-but-inflated DOF above `G` (observed ~32.7 and ~16.3 vs R's 6 and 3, `G=8`), which is NaN'd as non-physical. The well-conditioned contrasts estimators consume — the treatment effect, event-study coefficients, and the compound post-period-average — are unaffected and match R clubSandwich; only the non-user-facing high-leverage nuisance DOFs are suppressed (exact P-array reproduction of those is deferred). A `UserWarning` fires per fit. Regression: `tests/test_estimators_vcov_type.py::TestDiDAbsorbedFERParity::test_unweighted_cr2_bm_per_coef_dof_no_nonphysical`. **Evaluation (2026-07):** the pairwise `B` matrix is computed scores-based per Pustejovsky-Tipton (2018) Appendix B — `B = diag(‖ω_g‖²) − P' M_U P` with `P = X'Ω` — instead of contracting against an explicit dense `n×n` residual-maker (`O(n k + G² k)` per contrast, no `n×n` allocation; algebraically identical, ~1e-12 float agreement locked by a frozen-oracle parity test). +- **Note (unweighted per-coef DOF guard):** the unweighted clustered CR2-BM per-coefficient DOF (`_cr2_bm_dof_inner`, the simple `(tr B)²/tr(B²)` form) carries the same two-part reliability guard as the weighted P-array path. (1) **Noise floor:** for a high-leverage FE-dummy / collinear nuisance column `trace_B2 = Σ B_{g,h}²` collapses to float64 accumulation noise while `trace_B` stays O(1), inflating the ratio to a non-physical DOF (observed ~1e61 on the absorbed-FE golden); a contrast whose `max|B_{g,h}|` sits below the batch-relative (`1e-10×max`, computed on the scale-normalized `max|B|/‖c‖²` since `B ∝ ‖c‖²` while the DOF is scale-invariant) or absolute (`(EPS·n·k·bread_scale)²`) floor is NaN'd. (2) **Cluster-count bound:** the Bell-McCaffrey Satterthwaite DOF is `(tr B)²/tr(B²)` with `B` PSD and cluster-structured, so it is bounded by `rank(B) ≤ G` (number of clusters); the simple unweighted form is numerically less faithful than clubSandwich's P-array form on high-leverage columns and can return a finite-but-inflated DOF above `G` (observed ~32.7 and ~16.3 vs R's 6 and 3, `G=8`), which is NaN'd as non-physical. The well-conditioned contrasts estimators consume — the treatment effect, event-study coefficients, and the compound post-period-average — are unaffected and match R clubSandwich; only the non-user-facing high-leverage nuisance DOFs are suppressed (exact P-array reproduction of those is deferred). A `UserWarning` fires per fit. Regression: `tests/test_estimators_vcov_type.py::TestDiDAbsorbedFERParity::test_unweighted_cr2_bm_per_coef_dof_no_nonphysical`. **Evaluation (2026-07):** the pairwise `B` matrix is computed scores-based per Pustejovsky-Tipton (2018) Appendix B — `B = diag(‖ω_g‖²) − P' M_U P` with `P = X'Ω` — instead of contracting against an explicit dense `n×n` residual-maker (`O(n k + G² k)` per contrast, no `n×n` allocation, per-cluster product buffers contrast-chunked under a byte cap (bit-identical across chunk counts); algebraically identical, ~1e-12 float agreement locked by a frozen-oracle parity test). - **Note:** `LinearRegression.get_se()` / `get_inference()` clamp the vcov diagonal at 0 before `sqrt`. A high-leverage / degenerate coefficient (an absorbed-FE dummy near-collinear with the treatment, whose Satterthwaite DOF already hits the noise-floor diff --git a/tests/test_linalg.py b/tests/test_linalg.py index 40aaf7b0..3c0aa183 100644 --- a/tests/test_linalg.py +++ b/tests/test_linalg.py @@ -2806,3 +2806,23 @@ def test_unbalanced_clusters_match_oracle(self): _, dof = _compute_cr2_bm(X, resid, cl, bread) oracle = self._oracle_pairloop_dof(X, cl, bread, np.eye(k)) np.testing.assert_allclose(dof, oracle, rtol=1e-10) + + def test_contrast_chunking_bit_identical(self, monkeypatch): + """CI-review P2: the per-cluster product buffer is contrast-chunked + (bounded by _CR2_BM_CONTRAST_CHUNK_BYTES) instead of O(G*k*m). + Forcing one-contrast chunks must reproduce the single-chunk DOF + bit-for-bit — each contrast's B is computed independently.""" + import diff_diff.linalg as la + + rng = np.random.default_rng(23) + n, k, G = 200, 6, 10 + X = np.column_stack([np.ones(n), rng.normal(size=(n, k - 1))]) + cl = np.repeat(np.arange(G), n // G) + y = X @ rng.normal(size=k) + rng.normal(size=n) + resid = y - X @ np.linalg.lstsq(X, y, rcond=None)[0] + bread = X.T @ X + + _, dof_one = la._compute_cr2_bm(X, resid, cl, bread) + monkeypatch.setattr(la, "_CR2_BM_CONTRAST_CHUNK_BYTES", G * k * 8) # 1 contrast/chunk + _, dof_many = la._compute_cr2_bm(X, resid, cl, bread) + np.testing.assert_array_equal(dof_many, dof_one)