perf(linalg): scores-based CR2 Bell-McCaffrey DOF (PT2018 Appendix B)

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 <noreply@anthropic.com>

Author: igerber

CHANGELOG.md

@@ -481,6 +481,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   h=1 clamp-parity test). The symbol is imported independently (mixed-version safe — a
   stale extension degrades HC2 to NumPy without disabling older Rust accelerations).
   `return_dof` / weighted / CR2-BM requests stay on NumPy (CR2-BM tracked in TODO).
+- **CR2 Bell-McCaffrey Satterthwaite DOF: scores-based evaluation (PT2018 Appendix B).**
+  The unweighted per-contrast DOF previously materialized the dense `n×n` residual-maker
+  `M = I − X(X'X)⁻¹X'` and contracted it over all cluster pairs (`O(n²)` time per
+  contrast, `O(n²)` memory — 3.2 GB at n=20k). The pairwise matrix now collapses to
+  `B = diag(‖ω_g‖²) − P'M_U P` with `P = X'Ω` (disjoint cluster supports), costing
+  `O(nk + G²k)` per contrast with no `n×n` allocation: ~32x at n=5k/G=50 (0.57s→0.018s);
+  n=20k/G=100 completes in 0.12s where the old path would allocate 3.2 GB. Algebraically
+  identical — agreement with a frozen pair-loop oracle at rtol 1e-10 (balanced,
+  unbalanced, and compound-contrast designs); both NaN-reliability guards (noise floor,
+  cluster-count bound) unchanged. All 432 consumer tests pass unmodified.
 - **`CallawaySantAnna` per-(g,t) IF scatters converted from `np.add.at` to fancy `+=`**
   (`staggered.py::_cluster_robust_se_from_per_gt_if` — runs once per (g,t) cell when
   `cluster=` is set — and the general combined-IF assembly path in

TODO.md

@@ -46,7 +46,6 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m
 | `CallawaySantAnna` repeated-cross-section ipw/dr paths (`_ipw_estimation_rc` / `_doubly_robust_rc`) re-solve the propensity logit for EVERY (g,t) cell with no cache — the panel paths dedup via `pscore_cache` keyed `(g, base_period[, t])`. The Phase 3 IRLS fast path already speeds each RC solve; the missing (g,t)-dedup cache is the remaining lever. Unprofiled surface — measure an RC covariate scenario first. | `staggered.py::_ipw_estimation_rc`, `_doubly_robust_rc` | CS-scaling | Mid | Low |
 | Rust `solve_ols` always runs a full thin SVD (equilibrated, gelsd-parity); at 40 covariates it is the top per-cell solver item on a CS dr fit (1.04s of 3.1s at 2M rows, 95 calls). A Cholesky/QR fast path with SVD fallback — mirroring the Phase 3 Python-side pattern (certify well-conditioned, fall back verbatim) — is the natural lever, but `solve_ols` is the universal OLS entry point (every estimator), so the blast radius needs the full backend-parity treatment (`TestSolveOLSSkipRankCheckParity` posture: fitted-values parity, not beta). | `rust/src/linalg.rs::solve_ols` | CS-scaling | Heavy | Low |
 | `ImputationDiD` dense `(A0'A0).toarray()` scales `O((U+T+K)^2)` — OOM risk on large panels (only triggers when the sparse solver fails). Needs an alternative dense fallback or richer sparse strategy. | `imputation.py` | #141 | Heavy | Medium |
-| CR2 Bell-McCaffrey DOF uses a naive `O(n²k)` per-coefficient loop over cluster pairs; Pustejovsky-Tipton (2018) Appendix B has a scores-based formulation avoiding the full `n×n` `M`. Switch when a user hits a large-`n` cluster-robust design. | `linalg.py::_compute_cr2_bm` | Phase 1a | Heavy | Low |
 | Rust-backend CR2 Bell-McCaffrey: falls through to NumPy (the leverage/Satterthwaite-DOF path needs `return_dof` support, which the Rust vcov dispatch excludes). The one-way HC2 kernel landed 2026-07-07 (`compute_robust_vcov_hc2`, mirrors the NumPy hc2 branch at ~1e-15; near-singular hat-diagonal sentinel + Python-side warn-and-HC1-fallback). | `rust/src/linalg.rs` | Phase 1a | Mid | Low |
 | Wild cluster bootstrap CI inversion calls `_t_star(r)` ~O(100) times, each materializing a fresh `(B×n)` `y_star` + `(k×B)` refit + `(n×B)` residual arrays. Acceptable for the few-cluster regime; for large-`n`/large-`B`, chunk `_t_star` over draws or precompute the `r`-independent cluster-level pieces (restricted residuals are linear in `r`). | `utils.py::wild_bootstrap_se._t_star` | #543 | Mid | Low |
 | `SpilloverDiD` sparse cKDTree path for the staggered nearest-treated-distance helper (mirrors the static helper's sparse branch). `_compute_nearest_treated_distance_staggered` always builds dense `(n_units, n_treated_by_onset)` matrices per cohort; add a sparse branch gated on `n > _CONLEY_SPARSE_N_THRESHOLD`. | `spillover.py` | Wave B | Mid | Low |

diff_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,28 +2211,27 @@ 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`).
     max_abs_B_arr = np.zeros(m)
     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

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

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)