From a3c4d365f1d0909614fb6f56254494068f5ec065 Mon Sep 17 00:00:00 2001 From: igerber Date: Wed, 8 Jul 2026 01:14:08 -0400 Subject: [PATCH 1/9] 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 cb6bedde..9c7382f5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -590,6 +590,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 consumer (ring membership, `S_it`, the far-away check, the event-study `d_bar` trigger) compares against thresholds at or below that cutoff. Helper- and fit-level equality tests pin the sparse arm against the dense path (atol 1e-12 end-to-end). +- **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 02a27138..8013996e 100644 --- a/TODO.md +++ b/TODO.md @@ -46,7 +46,6 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | `EfficientDiD` conditional path: the largest remaining O(n) stage is the sieve/nuisance construction outside the tiled pass (~9s at 10k). (The `_ridge_solve_weights` Python-prep shave landed 2026-07-07 — the `omega_stack[rest]` fancy-index copy and tail scatter are skipped when no row is zero-masked, byte-identical outputs; the `zero_mask` abs scan itself remains, needed for correctness.) | `efficient_did_covariates.py` | 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 | | Adopt `snap_absorbed_regressors` (FE-spanned regressor two-stage snap + LSMR confirmation) on the ImputationDiD lead-indicator path — lead columns are the most plausible FE-spanned regressors, and the truncated-MAP-iterate exposure documented at `utils.py` applies; today only `solve_ols` rank detection guards it. Behavior change beyond the demean-modernization refactor, so deferred from that PR. | `imputation.py::_compute_lead_coefficients` | demean-modernization | 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 4470aa40..a0a90133 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 a22e7b09ea7148274b183fdc3288ab04797a4e64 Mon Sep 17 00:00:00 2001 From: igerber Date: Wed, 8 Jul 2026 07:20:16 -0400 Subject: [PATCH 2/9] 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 9c7382f5..99aa798f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -595,7 +595,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 a0a90133..8c460514 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) From cc04ff132b8638296b45c1c6f268f5c0ada9dea9 Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 9 Jul 2026 06:26:53 -0400 Subject: [PATCH 3/9] docs(linalg): correct the scores-based-evaluation citation (review P2) The accessible PT2018 arXiv text labels Appendix B as simulation-study details; the Satterthwaite DOF is Section 3.1 / Eq. 13. The scores identity B = Omega'M Omega = diag(||omega_g||^2) - P'M_U P is presented as an algebraic identity rather than paper-attributed, across the linalg docstring, oracle-test docstring, CHANGELOG entry, REGISTRY note, and the PR title. Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 2 +- diff_diff/linalg.py | 3 ++- docs/methodology/REGISTRY.md | 2 +- tests/test_linalg.py | 3 ++- 4 files changed, 6 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 99aa798f..37ecc658 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -590,7 +590,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 consumer (ring membership, `S_it`, the far-away check, the event-study `d_bar` trigger) compares against thresholds at or below that cutoff. Helper- and fit-level equality tests pin the sparse arm against the dense path (atol 1e-12 end-to-end). -- **CR2 Bell-McCaffrey Satterthwaite DOF: scores-based evaluation (PT2018 Appendix B).** +- **CR2 Bell-McCaffrey Satterthwaite DOF: scores-based evaluation (algebraic identity; PT2018 §3.1 / Eq. 13 Satterthwaite DOF).** 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 diff --git a/diff_diff/linalg.py b/diff_diff/linalg.py index e0223ede..c3a09894 100644 --- a/diff_diff/linalg.py +++ b/diff_diff/linalg.py @@ -2184,7 +2184,8 @@ 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 + SCORES-BASED EVALUATION (algebraic identity; the Satterthwaite DOF + itself is Pustejovsky-Tipton 2018 §3.1 / Eq. 13): 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 diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 8c460514..a5006940 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, 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 (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 via the algebraic identity `B = Ω'MΩ = diag(‖ω_g‖²) − P' M_U P` with `P = X'Ω` (the Satterthwaite DOF itself is Pustejovsky-Tipton 2018 §3.1 / Eq. 13) — 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 3c0aa183..51a4b8d2 100644 --- a/tests/test_linalg.py +++ b/tests/test_linalg.py @@ -2718,7 +2718,8 @@ def test_return_dropped_mask(self): 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 + (algebraic identity; PT2018 §3.1 Satterthwaite DOF): + 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.""" From cb2e2889b9fcc1aab107f97545c61e1f0e728a92 Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 9 Jul 2026 10:04:33 -0400 Subject: [PATCH 4/9] test+docs: chunk-count invariance is ~1 ULP, not bit-for-bit (CI failure fix) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI (linux-arm + Windows) failed test_contrast_chunking_bit_identical at 1-ULP differences: each contrast's B is computed independently, but the per-cluster GEMM X_g' omega_g[:, c0:c1] runs over a width-c slice and BLAS kernels (GEMV vs GEMM, platform-dependent) may accumulate a column differently at width 1 vs width m — exact on Accelerate, 1-ULP drift elsewhere. The documented chunking-reassociation caveat applies: test renamed to test_contrast_chunking_invariant with assert_allclose(rtol=1e-13), and the bit-for-bit claims corrected in the module-constant comment, docstring, CHANGELOG, and REGISTRY. Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 3 ++- diff_diff/linalg.py | 15 +++++++++++---- docs/methodology/REGISTRY.md | 2 +- tests/test_linalg.py | 12 ++++++++---- 4 files changed, 22 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 37ecc658..159d6493 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -596,7 +596,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 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 (per-cluster product buffers - contrast-chunked under a 64 MB cap, bit-identical across chunk counts, so the batched + contrast-chunked under a 64 MB cap, chunk-count invariant to ~1 ULP — BLAS kernels + can accumulate a GEMM column differently at different slice widths — 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, diff --git a/diff_diff/linalg.py b/diff_diff/linalg.py index c3a09894..ce055403 100644 --- a/diff_diff/linalg.py +++ b/diff_diff/linalg.py @@ -2153,8 +2153,12 @@ def _compute_cr2_bm( # (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. +# Each contrast's B matrix is computed independently, so chunking never +# reassociates a contrast's OWN sums — but the per-cluster GEMM +# `X_g' omega_g[:, c0:c1]` runs over a width-c slice, and BLAS kernels +# (GEMV vs GEMM, platform-dependent) may accumulate a column differently at +# different widths: chunk-count invariance holds to ~1 ULP (observed exact +# on Accelerate, 1-ULP drift on OpenBLAS/arm + Windows), NOT bit-for-bit. # Module-level so tests can monkeypatch it to force the multi-chunk path. _CR2_BM_CONTRAST_CHUNK_BYTES = 64 * 1024 * 1024 @@ -2197,7 +2201,9 @@ def _cr2_bm_dof_inner( so ``trace_B2 = ||B||_F^2`` costs ``O(n k + G^2 k)`` per contrast with ``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`` + ``_CR2_BM_CONTRAST_CHUNK_BYTES``, chunk-count invariant to ~1 ULP — BLAS + kernels may accumulate a GEMM column differently at different slice + widths) — 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 @@ -2233,7 +2239,8 @@ def _cr2_bm_dof_inner( # 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. + # contrast's B is computed independently; chunk-count invariance holds + # to ~1 ULP (BLAS width-dependent column accumulation), not bit-for-bit. 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) diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index a5006940..e3e2be05 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 via the algebraic identity `B = Ω'MΩ = diag(‖ω_g‖²) − P' M_U P` with `P = X'Ω` (the Satterthwaite DOF itself is Pustejovsky-Tipton 2018 §3.1 / Eq. 13) — 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 (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 via the algebraic identity `B = Ω'MΩ = diag(‖ω_g‖²) − P' M_U P` with `P = X'Ω` (the Satterthwaite DOF itself is Pustejovsky-Tipton 2018 §3.1 / Eq. 13) — 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 (chunk-count invariant to ~1 ULP — BLAS width-dependent column accumulation); 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 51a4b8d2..9f5e264b 100644 --- a/tests/test_linalg.py +++ b/tests/test_linalg.py @@ -2808,11 +2808,15 @@ def test_unbalanced_clusters_match_oracle(self): 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): + def test_contrast_chunking_invariant(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.""" + Forcing one-contrast chunks reproduces the single-chunk DOF to + ~1 ULP: each contrast's B is computed independently, but the + per-cluster GEMM runs over a width-c slice and BLAS kernels may + accumulate a column differently at width 1 vs width m (observed + exact on Accelerate, 1-ULP drift on OpenBLAS/arm + Windows CI — + the documented chunking-reassociation caveat).""" import diff_diff.linalg as la rng = np.random.default_rng(23) @@ -2826,4 +2830,4 @@ def test_contrast_chunking_bit_identical(self, monkeypatch): _, 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) + np.testing.assert_allclose(dof_many, dof_one, rtol=1e-13) From 156563ecfad1f6b95827802feff487ad841d2042 Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 9 Jul 2026 10:14:01 -0400 Subject: [PATCH 5/9] docs(linalg): bit-equal wording + citation precision (review optional P3s) Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 2 +- diff_diff/linalg.py | 21 +++++++++++++-------- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 159d6493..50fa56de 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -590,7 +590,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 consumer (ring membership, `S_it`, the far-away check, the event-study `d_bar` trigger) compares against thresholds at or below that cutoff. Helper- and fit-level equality tests pin the sparse arm against the dense path (atol 1e-12 end-to-end). -- **CR2 Bell-McCaffrey Satterthwaite DOF: scores-based evaluation (algebraic identity; PT2018 §3.1 / Eq. 13 Satterthwaite DOF).** +- **CR2 Bell-McCaffrey Satterthwaite DOF: scores-based evaluation (algebraic identity; PT2018 scalar Satterthwaite t-test DOF — the one-row HTZ case, §3.1).** 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 diff --git a/diff_diff/linalg.py b/diff_diff/linalg.py index ce055403..92f73147 100644 --- a/diff_diff/linalg.py +++ b/diff_diff/linalg.py @@ -2065,8 +2065,9 @@ def _compute_cr2_bm_vcov_and_dof( # --- Per-contrast Bell-McCaffrey cluster DOF --- # The inner helper branches on `weights`: unweighted uses the simple - # `(tr B)² / tr(B²)` form (bit-equal to prior); weighted uses the full - # clubSandwich P_array construction. + # `(tr B)² / tr(B²)` form (algebraically identical to the prior + # pair-loop evaluation; oracle parity within floating-point tolerance); + # weighted uses the full clubSandwich P_array construction. if weights is None: dof_vec = _cr2_bm_dof_inner(X, A_g_matrices, cluster_idx, M_U, contrasts) else: @@ -2188,8 +2189,9 @@ 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 (algebraic identity; the Satterthwaite DOF - itself is Pustejovsky-Tipton 2018 §3.1 / Eq. 13): the + SCORES-BASED EVALUATION (algebraic identity; the DOF itself is + Pustejovsky-Tipton 2018's scalar Satterthwaite t-test — the one-row + case of their HTZ small-sample correction, §3.1): 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 @@ -2552,8 +2554,10 @@ def _compute_cr2_bm_contrast_dof( ``contrasts=np.eye(k)``. weights : ndarray of shape (n,), optional Original (un-normalized) weights. ``None`` for unweighted; routes - through the bit-equal simple ``(tr B)² / tr(B²)`` formula. When - provided, routes through the clubSandwich WLS-CR2 P_array form. + through the simple ``(tr B)² / tr(B²)`` formula (algebraically + identical to the prior evaluation, parity within floating-point + tolerance). When provided, routes through the clubSandwich WLS-CR2 + P_array form. Returns ------- @@ -2646,8 +2650,9 @@ def _compute_bm_dof_from_contrasts( X, cluster_ids_singleton, bread_matrix, contrasts, weights=weights ) - # Unweighted: keep the simple (tr B)² / tr(B²) formula for bit-equal - # backward compatibility with prior unweighted Bell-McCaffrey output. + # Unweighted: keep the simple (tr B)² / tr(B²) formula — algebraically + # identical backward compatibility with prior unweighted Bell-McCaffrey + # output (scores-based evaluation; floating-point-tolerance parity). try: bread_inv_c = np.linalg.solve(bread_matrix, contrasts) except np.linalg.LinAlgError as e: From c7c63392f76de62e72fa59e6cbbaf763d5b83b1a Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 9 Jul 2026 10:27:05 -0400 Subject: [PATCH 6/9] perf(linalg): row-chunk the (G,G) pairwise CR2-BM matrix under the byte cap CI-review P2: B_j was materialized dense (G,G) per contrast, an O(G^2) memory bound the docs understated as O(Gk). trace_B2 (Frobenius sum) and max|B| are row-separable, so B_j now streams in row blocks bounded by _CR2_BM_CONTRAST_CHUNK_BYTES; neither the n x n residual-maker, the O(G*k*m) product buffers, nor O(G^2) pairwise entries are ever held at once. CHANGELOG/REGISTRY memory claims updated to the honest bound; the forced-tiny-cap invariance test now exercises both chunk axes (row_chunk = 6 < G = 10). Also reworded the stale "bit-equal to prior" REGISTRY sentence (review P3) to floating-point-tolerance parity. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Lbd6nqWmg4snvvBmegwqiw --- CHANGELOG.md | 10 ++++++---- diff_diff/linalg.py | 32 ++++++++++++++++++++++++-------- docs/methodology/REGISTRY.md | 4 ++-- tests/test_linalg.py | 5 ++++- 4 files changed, 36 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 50fa56de..95196674 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -595,10 +595,12 @@ 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 (per-cluster product buffers - contrast-chunked under a 64 MB cap, chunk-count invariant to ~1 ULP — BLAS kernels - can accumulate a GEMM column differently at different slice widths — so the batched - per-coefficient sweep never allocates `O(G·k²)` at once): ~32x at n=5k/G=50 (0.57s→0.018s); + `O(nk + G²k)` per contrast with memory bounded by a 64 MB cap — the per-cluster + product buffers are contrast-chunked and the `(G, G)` pairwise matrix is row-chunked + (its Frobenius sum and max are row-separable), so neither the `n×n` residual-maker, + `O(G·k·m)` product buffers, nor `O(G²)` pairwise entries are ever held at once + (chunk-count invariant to ~1 ULP — BLAS kernels can accumulate a GEMM column + differently at different slice widths): ~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 92f73147..045b64a6 100644 --- a/diff_diff/linalg.py +++ b/diff_diff/linalg.py @@ -2202,10 +2202,13 @@ 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 per contrast (contrast-chunked buffers bounded by - ``_CR2_BM_CONTRAST_CHUNK_BYTES``, chunk-count invariant to ~1 ULP — BLAS - kernels may accumulate a GEMM column differently at different slice - widths) — the previous form materialized the dense ``n x n`` + memory bounded by ``_CR2_BM_CONTRAST_CHUNK_BYTES`` (the per-cluster + product buffers are contrast-chunked AND the ``(G, G)`` pairwise matrix + is row-chunked — its Frobenius sum and max are row-separable — so + neither ``O(G k m)`` nor ``O(G^2)`` is ever held at once; chunk-count + invariant to ~1 ULP, BLAS kernels may accumulate a GEMM column + differently at different slice widths) — 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 @@ -2253,15 +2256,28 @@ def _cr2_bm_dof_inner( 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 + # Row-chunk the (G, G) pairwise matrix B_j under the same byte cap: + # a full B_j is O(G^2) per contrast, which can dominate on + # many-cluster designs (G in the tens of thousands). trace_B2 is a + # row-separable Frobenius sum and max|B| a row-separable max, so + # row blocks accumulate both without ever holding all of B_j. + row_chunk = max(1, int(_CR2_BM_CONTRAST_CHUNK_BYTES // max(n_g_clusters * 8, 1))) 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 + PB = P_j @ bread_inv # (G, k) + trace_B2 = 0.0 + max_abs_B = 0.0 + for r0 in range(0, n_g_clusters, row_chunk): + r1 = min(r0 + row_chunk, n_g_clusters) + B_rows = -(PB[r0:r1] @ P_j.T) # (rows, G) + B_rows[np.arange(r0, r1) - r0, np.arange(r0, r1)] += normsq[r0:r1, jj] + trace_B2 += float(np.sum(B_rows * B_rows)) + if B_rows.size: + max_abs_B = max(max_abs_B, float(np.max(np.abs(B_rows)))) + max_abs_B_arr[j] = max_abs_B 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 e3e2be05..a10d0789 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 via the algebraic identity `B = Ω'MΩ = diag(‖ω_g‖²) − P' M_U P` with `P = X'Ω` (the Satterthwaite DOF itself is Pustejovsky-Tipton 2018 §3.1 / Eq. 13) — 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 (chunk-count invariant to ~1 ULP — BLAS width-dependent column accumulation); 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 via the algebraic identity `B = Ω'MΩ = diag(‖ω_g‖²) − P' M_U P` with `P = X'Ω` (the Satterthwaite DOF itself is Pustejovsky-Tipton 2018 §3.1 / Eq. 13) — instead of contracting against an explicit dense `n×n` residual-maker (`O(n k + G² k)` per contrast; memory bounded by a byte cap — per-cluster product buffers contrast-chunked and the `(G, G)` pairwise matrix row-chunked (Frobenius sum and max are row-separable), so neither the `n×n` residual-maker nor `O(G²)` pairwise entries are ever held at once; chunk-count invariant to ~1 ULP — BLAS width-dependent column accumulation; 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 @@ -3398,7 +3398,7 @@ Shipped in `diff_diff/had_pretests.py` as `stute_joint_pretest()` (residuals-in **Requirements checklist (tracks implementation phase completion):** - [x] Phase 1a: Epanechnikov / triangular / uniform kernels with closed-form `κ_k` constants (`diff_diff/local_linear.py`). - [x] Phase 1a: Univariate local-linear regression at a boundary (`local_linear_fit` in `diff_diff/local_linear.py`). -- [x] Phase 1a: HC2 + Bell-McCaffrey DOF correction in `diff_diff/linalg.py` via `vcov_type="hc2_bm"` enum (both one-way and CR2 cluster-robust with Imbens-Kolesar / Pustejovsky-Tipton Satterthwaite DOF). **Weighted Bell-McCaffrey (`hc2_bm + weights`, both one-way and cluster) is now supported on the analytical linalg surface (`compute_robust_vcov` / `solve_ols` / `LinearRegression`) via the clubSandwich WLS-CR2 port.** The port matches clubSandwich's `pweight` (sampling-weight) convention only — `weight_type="aweight"` and `weight_type="fweight"` raise `NotImplementedError` on `hc2_bm + weights` (separate methodology task; CR1 / `vcov_type="hc1"` continues to support all three weight types). Estimator-level `survey_design=` paths continue to use the Taylor-series linearization (TSL) survey variance, which takes precedence over the analytical sandwich (no change in this PR). **Known precision limit on high-leverage FE-dummy coefficients**: the Satterthwaite DOF formula `(tr P)² / sum(P²)` is at the float64 noise floor for contrasts whose projection onto the design is essentially zero (typically FE-dummy coefficients where the dummy fires in a single cluster). In this regime, BLAS reduction-order differences between NumPy and R's BLAS produce 15-30% DOF discrepancies despite vcov matching at machine precision. `_cr2_bm_dof_inner_weighted` detects this regime via two criteria applied union-wise — **(a) batch-relative**: per-contrast max|P| below `1e-10 ×` the largest contrast's max|P| (catches noise-floor coefficients in a per-coefficient sweep where a non-noise reference is available); **(b) absolute single-contrast safe**: per-contrast max|P| below `(EPS × n × k × max(bread_inv_scale, 1))²` (catches single-contrast calls like MPD avg_att where no batch reference exists) — and returns NaN with a `UserWarning` rather than ship arbitrarily-different DOF on the user-facing surface. The coefficient SEs remain valid; only the DOF (and any t-test / CI that depends on it) is suppressed. The diff-diff implementation matches clubSandwich's specific algebra (R source: `CR-adjustments.R::CR2`, `clubSandwich.R::vcov_CR`, `coef_test.R::Satterthwaite_df`), **not** a textbook Pustejovsky-Tipton (2018) §3.3 transform-once derivation. clubSandwich uses W (not √W) in the hat matrix `H_gg = X_g M_U X_g' W_g`, W² in the bias-correction term `S_W = Σ_g X_g' W_g² X_g`, and unweighted residuals in the score construction `s_g = X_g' W_g A_g e_g`. The Satterthwaite DOF uses the full H1/H2/H3 array construction (`get_arrays.R::get_GH`), not the simplified `(tr B)² / tr(B²)` form (which works for the unweighted case but diverges from clubSandwich by 0.5-30% on weighted designs — see `feedback_wls_cr2_clubsandwich_parity`). Pinned at atol=1e-10 against `clubSandwich::vcovCR(..., type="CR2") + coef_test(test="Satterthwaite")$df_Satt` and `Wald_test(test="HTZ")$df_denom` on six weighted scenarios in `benchmarks/data/clubsandwich_cr2_golden.json` (vcov + non-noise-floor per-coefficient DOF + compound-contrast DOF; high-leverage FE-dummy coefficients are suppressed to NaN per the precision-limit note below) (`tests/test_methodology_wls_cr2.py`). Unweighted CR2-BM behavior is bit-equal to prior (regression-safe via `TestUnweightedRegressionStillBitEqual` + `TestDOFFormulaDualPathEquivalence`). The one-way weighted HC2-BM path uses clubSandwich's singleton-cluster CR2 reduction (each obs is its own cluster), routed via `_compute_bm_dof_from_contrasts` when `weights is not None`. clubSandwich pin: `≥ 0.7.0`. +- [x] Phase 1a: HC2 + Bell-McCaffrey DOF correction in `diff_diff/linalg.py` via `vcov_type="hc2_bm"` enum (both one-way and CR2 cluster-robust with Imbens-Kolesar / Pustejovsky-Tipton Satterthwaite DOF). **Weighted Bell-McCaffrey (`hc2_bm + weights`, both one-way and cluster) is now supported on the analytical linalg surface (`compute_robust_vcov` / `solve_ols` / `LinearRegression`) via the clubSandwich WLS-CR2 port.** The port matches clubSandwich's `pweight` (sampling-weight) convention only — `weight_type="aweight"` and `weight_type="fweight"` raise `NotImplementedError` on `hc2_bm + weights` (separate methodology task; CR1 / `vcov_type="hc1"` continues to support all three weight types). Estimator-level `survey_design=` paths continue to use the Taylor-series linearization (TSL) survey variance, which takes precedence over the analytical sandwich (no change in this PR). **Known precision limit on high-leverage FE-dummy coefficients**: the Satterthwaite DOF formula `(tr P)² / sum(P²)` is at the float64 noise floor for contrasts whose projection onto the design is essentially zero (typically FE-dummy coefficients where the dummy fires in a single cluster). In this regime, BLAS reduction-order differences between NumPy and R's BLAS produce 15-30% DOF discrepancies despite vcov matching at machine precision. `_cr2_bm_dof_inner_weighted` detects this regime via two criteria applied union-wise — **(a) batch-relative**: per-contrast max|P| below `1e-10 ×` the largest contrast's max|P| (catches noise-floor coefficients in a per-coefficient sweep where a non-noise reference is available); **(b) absolute single-contrast safe**: per-contrast max|P| below `(EPS × n × k × max(bread_inv_scale, 1))²` (catches single-contrast calls like MPD avg_att where no batch reference exists) — and returns NaN with a `UserWarning` rather than ship arbitrarily-different DOF on the user-facing surface. The coefficient SEs remain valid; only the DOF (and any t-test / CI that depends on it) is suppressed. The diff-diff implementation matches clubSandwich's specific algebra (R source: `CR-adjustments.R::CR2`, `clubSandwich.R::vcov_CR`, `coef_test.R::Satterthwaite_df`), **not** a textbook Pustejovsky-Tipton (2018) §3.3 transform-once derivation. clubSandwich uses W (not √W) in the hat matrix `H_gg = X_g M_U X_g' W_g`, W² in the bias-correction term `S_W = Σ_g X_g' W_g² X_g`, and unweighted residuals in the score construction `s_g = X_g' W_g A_g e_g`. The Satterthwaite DOF uses the full H1/H2/H3 array construction (`get_arrays.R::get_GH`), not the simplified `(tr B)² / tr(B²)` form (which works for the unweighted case but diverges from clubSandwich by 0.5-30% on weighted designs — see `feedback_wls_cr2_clubsandwich_parity`). Pinned at atol=1e-10 against `clubSandwich::vcovCR(..., type="CR2") + coef_test(test="Satterthwaite")$df_Satt` and `Wald_test(test="HTZ")$df_denom` on six weighted scenarios in `benchmarks/data/clubsandwich_cr2_golden.json` (vcov + non-noise-floor per-coefficient DOF + compound-contrast DOF; high-leverage FE-dummy coefficients are suppressed to NaN per the precision-limit note below) (`tests/test_methodology_wls_cr2.py`). Unweighted CR2-BM behavior is algebraically identical to prior (parity within floating-point tolerance after the 2026-07 scores-based evaluation change; regression-safe via `TestUnweightedRegressionStillBitEqual` + `TestDOFFormulaDualPathEquivalence`). The one-way weighted HC2-BM path uses clubSandwich's singleton-cluster CR2 reduction (each obs is its own cluster), routed via `_compute_bm_dof_from_contrasts` when `weights is not None`. clubSandwich pin: `≥ 0.7.0`. - **Note (scope limitation on absorbed FE):** HC2 and HC2 + Bell-McCaffrey on within-transformed designs still depend on the FULL FE hat matrix because FWL preserves coefficients and residuals but NOT the hat matrix: `h_ii = x_i' (X'X)^{-1} x_i` on the reduced design is not the diagonal of the full FE projection, and CR2's block adjustment `A_g = (I - H_gg)^{-1/2}` likewise depends on the full cluster-block hat matrix. The status across the three estimators that previously rejected this combination: - **`DifferenceInDifferences(absorb=..., vcov_type in {"hc2","hc2_bm"})` — SUPPORTED (auto-route).** When the user pairs `absorb=` with HC2 / HC2-BM, `DiD.fit()` internally promotes the absorb columns to `fixed_effects=` so the existing full-dummy code path computes the algebraically correct vcov from the full FE projection. Verified at ~1e-10 vs `lm() + sandwich::vcovHC(type="HC2")` and `lm() + clubSandwich::vcovCR(cluster=..., type="CR2")` (singleton-cluster CR2 trick for one-way HC2-BM Satterthwaite DOF; PT2018 §3.3 unweighted CR2 algebra). **User-visible surface change**: under the auto-route, the entire `DiDResults` (coefficients, vcov, residuals, fitted_values, r_squared) reflect the full-dummy fit rather than the within-transformed fit — the FE-dummy entries are included in `result.coefficients` / `result.vcov`, `r_squared` is computed on the un-demeaned outcome, and `residuals` / `fitted_values` are on the original scale. `result.att` is unaffected (FWL-equivalent). HC1/CR1 paths on `absorb=` are unchanged (no leverage term). **Survey-design scope**: when `survey_design=` is supplied, the existing survey variance path (Taylor-series linearization / replicate weights) takes precedence over the analytical HC2/HC2-BM sandwich; the auto-route only changes the FE handling (removing the prior reject) and does not redirect to the analytical small-sample sandwich on survey fits. - **`TwoWayFixedEffects(vcov_type in {"hc2","hc2_bm"})` — SUPPORTED (inline full-dummy build).** TWFE has no `absorb=` / `fixed_effects=` parameter (the unit + time FE are baked into the estimator's identity), so the same parameter-swap auto-route used for DiD-absorb / MPD-absorb is not directly applicable. Instead, `TwoWayFixedEffects.fit()` bypasses the within-transform when `vcov_type in {"hc2","hc2_bm"}` and builds the full-dummy design `[intercept, treated×post, covariates, factor(unit), factor(time)]` explicitly, then runs OLS through the standard `solve_ols` path so the leverage correction and BM DOF compute on the full FE projection. Verified at atol=1e-10 vs `lm(y ~ treat_post + factor(unit) + factor(post)) + sandwich::vcovHC(type="HC2")` for HC2 and vs `clubSandwich::vcovCR(cluster=seq_len(n), type="CR2")` for the singleton-cluster one-way HC2-BM Satterthwaite DOF; vs `vcovCR(cluster=unit, type="CR2")` for the auto-cluster CR2-BM path. **Auto-cluster default (non-survey analytical path):** TWFE's unit auto-cluster is preserved on `hc2_bm` (routes to CR2-BM at unit) and on `hc2 + wild_bootstrap` (the bootstrap consumes the cluster structure for resampling regardless of the analytical sandwich choice); dropped on explicit `hc2 + analytical` to match the one-way contract (the linalg validator rejects `hc2 + cluster_ids`). `hc2_bm + analytical` with no explicit cluster yields the auto-cluster CR2-BM path. **Survey-design exception:** under `survey_design=` with no explicit `cluster=`, TWFE intentionally keeps the documented implicit-PSU path (the auto-cluster is NOT injected into the survey PSU structure); users who want unit-level PSU injection under a survey design must pass explicit `cluster="unit"` or set `survey_design.psu` directly. **User-visible surface change** (matches the DiD-absorb / MPD-absorb disclosure above): under HC2 / HC2-BM, `result.coefficients`, `result.vcov`, `result.residuals`, `result.fitted_values`, and `result.r_squared` reflect the full-dummy fit rather than the within-transformed reduced fit (FE-dummy entries are included, `r_squared` is computed on the un-demeaned outcome, residuals/fitted are on the original scale). `result.att`, its SE, and analytical inference are unchanged (FWL-equivalent). HC1 / CR1 / Conley / classical paths remain on the within-transform (no leverage term in those vcov families). **Survey-design scope** (mirrors the DiD-absorb auto-route contract above): when `survey_design=` is supplied, the existing survey variance path (Taylor-series linearization for analytical-weight designs, or replicate-weight variance for BRR/Fay/JK1/JKn/SDR) takes precedence over the analytical HC2/HC2-BM sandwich; the full-dummy build only changes the FE handling (removing the prior reject) and does not redirect to the analytical small-sample sandwich on survey fits. **Replicate-weight survey designs** are blocked at the estimator level: `vcov_type in {"hc2","hc2_bm"}` + replicate weights raises `NotImplementedError` because the replicate refit path re-demeans per replicate, which doesn't compose with the full-dummy build (would require per-replicate full-dummy refit); workaround: use `vcov_type="hc1"` for replicate-weight CR1. `hc2_bm + weights` (Gates 4-5) is now lifted via the clubSandwich WLS-CR2 port; the survey-design rejection here is a separate estimator-level gate (analytical sandwich vs survey TSL precedence), independent of the linalg validator. diff --git a/tests/test_linalg.py b/tests/test_linalg.py index 9f5e264b..01c619cc 100644 --- a/tests/test_linalg.py +++ b/tests/test_linalg.py @@ -2816,7 +2816,10 @@ def test_contrast_chunking_invariant(self, monkeypatch): per-cluster GEMM runs over a width-c slice and BLAS kernels may accumulate a column differently at width 1 vs width m (observed exact on Accelerate, 1-ULP drift on OpenBLAS/arm + Windows CI — - the documented chunking-reassociation caveat).""" + the documented chunking-reassociation caveat). The same tiny cap + also forces the (G, G) pairwise matrix into multiple row blocks + (row_chunk = cap // (G*8) = 6 < G = 10 here), so this run covers + both chunk axes against the unchunked single-pass result.""" import diff_diff.linalg as la rng = np.random.default_rng(23) From 93d5d24f8d997873dee7ae27ef78c987d9e8b3e0 Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 9 Jul 2026 10:36:18 -0400 Subject: [PATCH 7/9] perf(linalg): contrast-chunk the q/omega score arrays; scope remaining docstring claims MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups on the row-chunk commit: - P2: "memory bounded by the cap" was still too broad — Q (n,m) and the per-cluster omega_all dict (n,m total) were full-width precomputes. Both now compute inside the contrast chunk (same total FLOPs), leaving X_bi and A_g_Xbi (both O(nk), input-scale) as the only full-size arrays; docstring/CHANGELOG/REGISTRY state exactly which arrays are bounded. - P3: one-way BM DOF's "scores-based formulation (tracked in TODO.md)" pointed at the CR2-BM row this PR resolves — added a specific TODO row for the one-way dense-M path and repointed the note. - P3: scoped _compute_cr2_bm's "bit-equal to prior unweighted behavior" claim to the vcov; the DOF evaluation is scores-based (floating-point- tolerance parity). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Lbd6nqWmg4snvvBmegwqiw --- CHANGELOG.md | 11 +++++----- TODO.md | 1 + diff_diff/linalg.py | 39 ++++++++++++++++++++++-------------- docs/methodology/REGISTRY.md | 2 +- 4 files changed, 32 insertions(+), 21 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 95196674..8cca2271 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -595,12 +595,13 @@ 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 memory bounded by a 64 MB cap — the per-cluster + `O(nk + G²k)` per contrast; peak memory = two `O(nk)` input-scale score precomputes + plus working buffers bounded by a 64 MB cap — q vectors, per-cluster omegas, and product buffers are contrast-chunked and the `(G, G)` pairwise matrix is row-chunked - (its Frobenius sum and max are row-separable), so neither the `n×n` residual-maker, - `O(G·k·m)` product buffers, nor `O(G²)` pairwise entries are ever held at once - (chunk-count invariant to ~1 ULP — BLAS kernels can accumulate a GEMM column - differently at different slice widths): ~32x at n=5k/G=50 (0.57s→0.018s); + (its Frobenius sum and max are row-separable), so none of the `n×n` residual-maker, + `O(n·m)` score arrays, `O(G·k·m)` product buffers, or `O(G²)` pairwise entries is + ever held at once (chunk-count invariant to ~1 ULP — BLAS kernels can accumulate a + GEMM column differently at different slice widths): ~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/TODO.md b/TODO.md index 8013996e..1a6af85f 100644 --- a/TODO.md +++ b/TODO.md @@ -47,6 +47,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | 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 | | 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 | +| One-way (non-clustered) unweighted Bell-McCaffrey DOF still materializes the dense `n×n` residual-maker `M` in `_compute_bm_dof_from_contrasts` (`O(n²k)` hat build + `O(n²m)` per-contrast sums; practical for n < 10k). Switch to a scores-based evaluation like the clustered CR2-BM path now uses (singleton-cluster reduction of the same `B = diag(‖ω‖²) − P'M_U P` identity, row-chunked). | `diff_diff/linalg.py::_compute_bm_dof_from_contrasts` | #656 | Mid | Low | | Adopt `snap_absorbed_regressors` (FE-spanned regressor two-stage snap + LSMR confirmation) on the ImputationDiD lead-indicator path — lead columns are the most plausible FE-spanned regressors, and the truncated-MAP-iterate exposure documented at `utils.py` applies; today only `solve_ols` rank detection guards it. Behavior change beyond the demean-modernization refactor, so deferred from that PR. | `imputation.py::_compute_lead_coefficients` | demean-modernization | Mid | Low | ### Testing / docs diff --git a/diff_diff/linalg.py b/diff_diff/linalg.py index 045b64a6..c063271a 100644 --- a/diff_diff/linalg.py +++ b/diff_diff/linalg.py @@ -2110,8 +2110,10 @@ def _compute_cr2_bm( Unweighted special case (``weights=None``): ``W_norm=1``, ``S_W=X'X``, ``M_U @ S_W @ M_U = M_U``, so ``G_g`` collapses to ``I - H_gg`` (the - symmetric form). Bit-equal to the prior unweighted behavior at machine - precision (atol=1e-14 regression-safety). + symmetric form). The vcov is bit-equal to the prior unweighted behavior + at machine precision (atol=1e-14 regression-safety); the per-contrast + DOF uses the scores-based evaluation (algebraically identical, parity + within floating-point tolerance — see `_cr2_bm_dof_inner`). Meat = ``sum_g s_g s_g'``; VCOV = ``M_U meat M_U`` (where ``M_U`` is the normalized bread inverse; ``w_scale`` cancels in the final vcov). @@ -2201,13 +2203,16 @@ def _cr2_bm_dof_inner( 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 - memory bounded by ``_CR2_BM_CONTRAST_CHUNK_BYTES`` (the per-cluster - product buffers are contrast-chunked AND the ``(G, G)`` pairwise matrix - is row-chunked — its Frobenius sum and max are row-separable — so - neither ``O(G k m)`` nor ``O(G^2)`` is ever held at once; chunk-count - invariant to ~1 ULP, BLAS kernels may accumulate a GEMM column - differently at different slice widths) — the previous form + so ``trace_B2 = ||B||_F^2`` costs ``O(n k + G^2 k)`` per contrast. Peak + memory = two ``O(n k)`` score precomputes (``X_bi`` and the per-cluster + ``A_g_Xbi`` blocks — input-scale, same order as ``X`` itself) plus + working buffers bounded by ``_CR2_BM_CONTRAST_CHUNK_BYTES``: the q + vectors, per-cluster omegas, and product buffers are contrast-chunked + and the ``(G, G)`` pairwise matrix is row-chunked (its Frobenius sum + and max are row-separable), so none of ``O(n m)``, ``O(G k m)``, or + ``O(G^2)`` is ever held at once (chunk-count invariant to ~1 ULP, BLAS + kernels may accumulate a GEMM column differently at different slice + widths) — 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 @@ -2227,12 +2232,12 @@ def _cr2_bm_dof_inner( # For unit-contrast inputs (contrasts=I_k), this matches the prior # inline implementation exactly: q[:, j] == X_bi[:, j] == X @ bread_inv @ e_j. X_bi = X @ bread_inv # (n, k) - Q = X_bi @ contrasts # (n, m) — q vectors as columns A_g_Xbi = { g: A_g_matrices[g] @ X[cluster_idx[g]] @ bread_inv for g in unique_clusters } # each (n_g, k) - # 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} + # The q vectors (X_bi @ contrasts) and per-cluster omegas + # (A_g_Xbi[g] @ contrasts) are computed per contrast chunk below, never + # at full width m, so no O(n*m) array is ever held. # 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)}. @@ -2250,10 +2255,12 @@ def _cr2_bm_dof_inner( for c0 in range(0, m, chunk): c1 = min(c0 + chunk, m) width = c1 - c0 + contrasts_c = contrasts[:, c0:c1] + Q_chunk = X_bi @ contrasts_c # (n, width) — q vectors as columns 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) + om = A_g_Xbi[g] @ contrasts_c # (n_g, width) normsq[gi] = np.einsum("ij,ij->j", om, om) P_all[gi] = X[cluster_idx[g]].T @ om # Row-chunk the (G, G) pairwise matrix B_j under the same byte cap: @@ -2264,7 +2271,7 @@ def _cr2_bm_dof_inner( row_chunk = max(1, int(_CR2_BM_CONTRAST_CHUNK_BYTES // max(n_g_clusters * 8, 1))) for jj in range(width): j = c0 + jj - q = Q[:, j] + q = Q_chunk[:, jj] trace_B = float(np.sum(q * q)) P_j = P_all[:, :, jj] # (G, k) PB = P_j @ bread_inv # (G, k) @@ -2620,7 +2627,9 @@ def _compute_bm_dof_from_contrasts( matches the numerator. Allocates an ``(n, n)`` temporary for ``M`` so the cost is ``O(n^2 k)`` for the hat build plus ``O(n^2 m)`` for the per- contrast sums. Practical for ``n < 10_000``; larger designs should switch - to a scores-based formulation (tracked in TODO.md). + to a scores-based formulation like the clustered CR2-BM path + (`_cr2_bm_dof_inner`) now uses — tracked as its own TODO.md row (the + original CR2-BM row this note pointed at is resolved). **Weighted** (``weights is not None``): dispatches to the clubSandwich singleton-cluster CR2 reduction (each observation is its own cluster) diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index a10d0789..e1214a0e 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 via the algebraic identity `B = Ω'MΩ = diag(‖ω_g‖²) − P' M_U P` with `P = X'Ω` (the Satterthwaite DOF itself is Pustejovsky-Tipton 2018 §3.1 / Eq. 13) — instead of contracting against an explicit dense `n×n` residual-maker (`O(n k + G² k)` per contrast; memory bounded by a byte cap — per-cluster product buffers contrast-chunked and the `(G, G)` pairwise matrix row-chunked (Frobenius sum and max are row-separable), so neither the `n×n` residual-maker nor `O(G²)` pairwise entries are ever held at once; chunk-count invariant to ~1 ULP — BLAS width-dependent column accumulation; 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 via the algebraic identity `B = Ω'MΩ = diag(‖ω_g‖²) − P' M_U P` with `P = X'Ω` (the Satterthwaite DOF itself is Pustejovsky-Tipton 2018 §3.1 / Eq. 13) — instead of contracting against an explicit dense `n×n` residual-maker (`O(n k + G² k)` per contrast; peak memory = two `O(n k)` input-scale score precomputes plus working buffers bounded by a byte cap — q vectors, per-cluster omegas, and product buffers contrast-chunked and the `(G, G)` pairwise matrix row-chunked (Frobenius sum and max are row-separable), so none of the `n×n` residual-maker, `O(n m)` score arrays, or `O(G²)` pairwise entries is ever held at once; chunk-count invariant to ~1 ULP — BLAS width-dependent column accumulation; 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 From aaed82dbaa5e4ed40aa7ca29b1340ab3f6d48564 Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 9 Jul 2026 10:44:53 -0400 Subject: [PATCH 8/9] perf(linalg): count every width-scaled buffer in the CR2-BM chunk denominator Review round 3 P2: the contrast-chunk denominator only counted the P_all slab (G*k), so Q_chunk (n per width) escaped the cap when n dominates, and the docs' "bounded by the cap" claim ignored the intrinsic one-contrast floor. The denominator now counts all width-scaled buffers (n + n_g_max + G*k + G, in bytes); docstring/CHANGELOG/REGISTRY state the cap is subject to a one-contrast lower bound (a single contrast needs the O(n) q vector and O(G*k) P_j/PB buffers by construction). Default 64 MB cap still yields a single chunk at benchmark sizes. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Lbd6nqWmg4snvvBmegwqiw --- CHANGELOG.md | 6 ++++-- diff_diff/linalg.py | 32 +++++++++++++++++++++----------- docs/methodology/REGISTRY.md | 2 +- 3 files changed, 26 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8cca2271..6cb1acf8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -596,8 +596,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 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; peak memory = two `O(nk)` input-scale score precomputes - plus working buffers bounded by a 64 MB cap — q vectors, per-cluster omegas, and - product buffers are contrast-chunked and the `(G, G)` pairwise matrix is row-chunked + plus working buffers capped at 64 MB subject to a one-contrast lower bound (a single + contrast intrinsically needs `O(n)` + `O(G·k)` buffers) — q vectors, per-cluster + omegas, and product buffers are contrast-chunked with every width-scaled buffer + counted in the chunk denominator, and the `(G, G)` pairwise matrix is row-chunked (its Frobenius sum and max are row-separable), so none of the `n×n` residual-maker, `O(n·m)` score arrays, `O(G·k·m)` product buffers, or `O(G²)` pairwise entries is ever held at once (chunk-count invariant to ~1 ULP — BLAS kernels can accumulate a diff --git a/diff_diff/linalg.py b/diff_diff/linalg.py index c063271a..b26f289d 100644 --- a/diff_diff/linalg.py +++ b/diff_diff/linalg.py @@ -2206,11 +2206,14 @@ def _cr2_bm_dof_inner( so ``trace_B2 = ||B||_F^2`` costs ``O(n k + G^2 k)`` per contrast. Peak memory = two ``O(n k)`` score precomputes (``X_bi`` and the per-cluster ``A_g_Xbi`` blocks — input-scale, same order as ``X`` itself) plus - working buffers bounded by ``_CR2_BM_CONTRAST_CHUNK_BYTES``: the q - vectors, per-cluster omegas, and product buffers are contrast-chunked - and the ``(G, G)`` pairwise matrix is row-chunked (its Frobenius sum - and max are row-separable), so none of ``O(n m)``, ``O(G k m)``, or - ``O(G^2)`` is ever held at once (chunk-count invariant to ~1 ULP, BLAS + working buffers capped by ``_CR2_BM_CONTRAST_CHUNK_BYTES`` subject to a + one-contrast lower bound (a single contrast intrinsically needs the + ``O(n)`` q vector and ``O(G k)`` P_j/PB buffers): the q vectors, + per-cluster omegas, and product buffers are contrast-chunked with all + width-scaled buffers counted in the chunk denominator, and the + ``(G, G)`` pairwise matrix is row-chunked (its Frobenius sum and max + are row-separable), so none of ``O(n m)``, ``O(G k m)``, or ``O(G^2)`` + is ever held at once (chunk-count invariant to ~1 ULP, BLAS kernels may accumulate a GEMM column differently at different slice widths) — the previous form materialized the dense ``n x n`` @@ -2246,12 +2249,19 @@ def _cr2_bm_dof_inner( # 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) - # 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; chunk-count invariance holds - # to ~1 ULP (BLAS width-dependent column accumulation), not bit-for-bit. - chunk = max(1, int(_CR2_BM_CONTRAST_CHUNK_BYTES // max(n_g_clusters * k_X * 8, 1))) + # Chunk the contrasts so every width-scaled working buffer stays under + # _CR2_BM_CONTRAST_CHUNK_BYTES: per unit of contrast width we allocate a + # Q_chunk column (n), a transient per-cluster omega (largest cluster + # n_g_max), a P_all slab (G*k), and a normsq row (G) — a full-m sweep + # would be O((n + G*k)*m). The cap is subject to a one-contrast lower + # bound: a single contrast intrinsically needs the O(n) q vector and the + # O(G*k) P_j / PB buffers. Each contrast's B is computed independently; + # chunk-count invariance holds to ~1 ULP (BLAS width-dependent column + # accumulation), not bit-for-bit. + n = X.shape[0] + n_g_max = max((idx.size for idx in cluster_idx.values()), default=0) + per_width_bytes = (n + n_g_max + n_g_clusters * k_X + n_g_clusters) * 8 + chunk = max(1, int(_CR2_BM_CONTRAST_CHUNK_BYTES // max(per_width_bytes, 1))) for c0 in range(0, m, chunk): c1 = min(c0 + chunk, m) width = c1 - c0 diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index e1214a0e..43a5d62f 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 via the algebraic identity `B = Ω'MΩ = diag(‖ω_g‖²) − P' M_U P` with `P = X'Ω` (the Satterthwaite DOF itself is Pustejovsky-Tipton 2018 §3.1 / Eq. 13) — instead of contracting against an explicit dense `n×n` residual-maker (`O(n k + G² k)` per contrast; peak memory = two `O(n k)` input-scale score precomputes plus working buffers bounded by a byte cap — q vectors, per-cluster omegas, and product buffers contrast-chunked and the `(G, G)` pairwise matrix row-chunked (Frobenius sum and max are row-separable), so none of the `n×n` residual-maker, `O(n m)` score arrays, or `O(G²)` pairwise entries is ever held at once; chunk-count invariant to ~1 ULP — BLAS width-dependent column accumulation; 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 via the algebraic identity `B = Ω'MΩ = diag(‖ω_g‖²) − P' M_U P` with `P = X'Ω` (the Satterthwaite DOF itself is Pustejovsky-Tipton 2018 §3.1 / Eq. 13) — instead of contracting against an explicit dense `n×n` residual-maker (`O(n k + G² k)` per contrast; peak memory = two `O(n k)` input-scale score precomputes plus working buffers capped by a byte cap subject to a one-contrast lower bound (a single contrast intrinsically needs `O(n)` + `O(G k)` buffers) — q vectors, per-cluster omegas, and product buffers contrast-chunked with every width-scaled buffer counted in the chunk denominator, and the `(G, G)` pairwise matrix row-chunked (Frobenius sum and max are row-separable), so none of the `n×n` residual-maker, `O(n m)` score arrays, or `O(G²)` pairwise entries is ever held at once; chunk-count invariant to ~1 ULP — BLAS width-dependent column accumulation; 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 From be4e5d28c92047d09fc1ed8cbbb017c60694acea Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 9 Jul 2026 10:51:58 -0400 Subject: [PATCH 9/9] =?UTF-8?q?docs(registry):=20PT2018=20Satterthwaite=20?= =?UTF-8?q?DOF=20is=20=C2=A73.1=20Eq.=2011=20(q=3D1=20reduction=20of=20AHT?= =?UTF-8?q?=20Eq.=2013)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review round 4 P3: the scalar t-test Satterthwaite degrees of freedom are Eq. 11 in PT2018 §3.1; Eq. 13 is the multi-parameter AHT expression that reduces to Eq. 11 at q=1. linalg.py's own citation names only §3.1 (no equation number) and needs no change. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Lbd6nqWmg4snvvBmegwqiw --- docs/methodology/REGISTRY.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 43a5d62f..63346c18 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 via the algebraic identity `B = Ω'MΩ = diag(‖ω_g‖²) − P' M_U P` with `P = X'Ω` (the Satterthwaite DOF itself is Pustejovsky-Tipton 2018 §3.1 / Eq. 13) — instead of contracting against an explicit dense `n×n` residual-maker (`O(n k + G² k)` per contrast; peak memory = two `O(n k)` input-scale score precomputes plus working buffers capped by a byte cap subject to a one-contrast lower bound (a single contrast intrinsically needs `O(n)` + `O(G k)` buffers) — q vectors, per-cluster omegas, and product buffers contrast-chunked with every width-scaled buffer counted in the chunk denominator, and the `(G, G)` pairwise matrix row-chunked (Frobenius sum and max are row-separable), so none of the `n×n` residual-maker, `O(n m)` score arrays, or `O(G²)` pairwise entries is ever held at once; chunk-count invariant to ~1 ULP — BLAS width-dependent column accumulation; 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 via the algebraic identity `B = Ω'MΩ = diag(‖ω_g‖²) − P' M_U P` with `P = X'Ω` (the Satterthwaite DOF itself is Pustejovsky-Tipton 2018 §3.1 Eq. 11; equivalently the `q=1` reduction of the AHT Eq. 13) — instead of contracting against an explicit dense `n×n` residual-maker (`O(n k + G² k)` per contrast; peak memory = two `O(n k)` input-scale score precomputes plus working buffers capped by a byte cap subject to a one-contrast lower bound (a single contrast intrinsically needs `O(n)` + `O(G k)` buffers) — q vectors, per-cluster omegas, and product buffers contrast-chunked with every width-scaled buffer counted in the chunk denominator, and the `(G, G)` pairwise matrix row-chunked (Frobenius sum and max are row-separable), so none of the `n×n` residual-maker, `O(n m)` score arrays, or `O(G²)` pairwise entries is ever held at once; chunk-count invariant to ~1 ULP — BLAS width-dependent column accumulation; 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