igerber · igerber · Jul 10, 2026 · Jul 9, 2026 · Jul 9, 2026 · Jul 9, 2026
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -619,6 +619,28 @@ 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 unweighted adjustment matrices: low-rank factored evaluation.**
+  After the scores-based DOF change below, ~85% of remaining CR2-BM runtime was the dense
+  per-cluster eigendecomposition `A_g = (I − H_gg)^{−1/2}` (`O(n_g³)` per cluster). In the
+  unweighted case `H_gg = X_g M_U X_g'` has rank ≤ k, so with `U_g = X_g M_U^{1/2}` and
+  `U_g'U_g = Q diag(λ) Q'`, the adjustment operator is exactly
+  `A_g = I + (U_g Q) diag(γ) (U_g Q)'` with `γ_i = ((1−λ_i)^{−1/2} − 1)/λ_i` (Moore-Penrose
+  zeroing `γ_i = −1/λ_i` when `1−λ_i ≤ 1e-10`, matching the dense path's pseudoinverse
+  convention for absorbed cluster FEs) — a k×k eigenproblem per cluster. Consumers only
+  ever apply `A_g` to skinny matrices (the residual vector for the meat; `X_g bread_inv`
+  for the DOF omegas), so the dense `(n_g, n_g)` `A_g` is never materialized: `O(n_g·k·min(n_g,k) + min(n_g,k)³)`
+  per cluster (the eigenproblem is solved on the smaller Gram side — k×k for large
+  clusters, the tiny dense `n_g×n_g` construction when `n_g ≤ k`, so small/singleton
+  clusters never regress vs the prior dense path) and `O(n k)` total memory, with `γ` evaluated via `expm1/log1p`
+  (stable as λ→0). End-to-end `_compute_cr2_bm`: 18→2.3 ms at n=5k/G=50/k=10;
+  119→5.3 ms at n=20k/G=100; 597→33 ms at n=100k/G=500; 4111→38 ms (~108x) at
+  n=100k/G=100/k=40. Algebraically identical — vcov/DOF match a frozen dense-eigh oracle
+  at rtol 1e-12/1e-10 with identical NaN-guard patterns (balanced, unbalanced, k=40,
+  leverage-1 absorbed-cluster-FE, singleton clusters, near-zero-leverage γ-stability;
+  `tests/test_linalg.py::TestCR2BMLowRankAdjustment`). The weighted (clubSandwich
+  WLS-CR2) path keeps the dense construction — its `G_g` carries the `S_W` bias term and
+  is not a rank-k identity perturbation. This obsoletes the planned Rust CR2-BM port
+  (the NumPy path is now BLAS-bound; see the re-scoped TODO row).
 - **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

diff --git a/TODO.md b/TODO.md
@@ -44,7 +44,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 |
 | `TwoStageDiD` dense `toarray()` fallbacks (`two_stage.py:297` unweighted + `:2922` GMM-sandwich weighted) share the OOM-risk pattern the ImputationDiD LSMR fix (2026-07) closed, but are MULTI-RHS solves whose `gamma_hat` feeds coefficient-level consumers directly — the null-space-invariance argument that made the imputation swap provably output-preserving does NOT transfer; needs its own analysis (does any consumer depend on the min-norm choice?) before an LSMR/column-loop swap. | `diff_diff/two_stage.py` | #141 | Mid | 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 |
 | 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 |
 
 ### Testing / docs
@@ -112,6 +111,7 @@ Doable in principle, but no current caller and/or explicitly out of paper scope.
 
 | Issue | Location | PR | Priority |
 |-------|----------|----|----------|
+| Rust-backend CR2 Bell-McCaffrey port (`return_dof` in the Rust vcov dispatch + CR2 algebra) — **premise re-scoped 2026-07-09**: the scores-based DOF + low-rank factored `A_g` changes made the NumPy CR2-BM path BLAS-bound (`O(n_g k²)` per cluster; 4.1s→38ms at n=100k/k=40), so a Rust port buys ~nothing and adds a parity surface. Revisit only if profiling shows CR2-BM hot again. | `rust/src/linalg.rs` | — | Low |
 | CallawaySantAnna **unbalanced-panel R parity — LANDED** via `allow_unbalanced_panel=True` (matches R `did::att_gt(allow_unbalanced_panel=TRUE)` / `DRDID::reg_did_rc`: ATT bit-exact on cells AND dynamic aggregation via fixed unit-cohort-mass `pg` + a per-unit WIF; SE up to the documented CR1 `sqrt(G/(G-1))` factor). The earlier "weighting" framing was a mis-diagnosis — on unbalanced panels the dominant divergence from R is the *estimator* (within-cell differencing vs RC-on-pooled-obs), not only the weighting; both are resolved by the flag. The DEFAULT path keeps within-cell differencing as a documented design choice and now emits a `UserWarning` on unbalanced input (no-silent-failures). **Remaining deferred:** `survey_design=` × `allow_unbalanced_panel=` (per-obs vs per-unit weight resolution — currently fail-closed `NotImplementedError`); and covariate / ipw / dr × the flag R-parity verification (the RC path supports them; the committed golden covers `reg` no-cov). | `staggered.py`, `staggered_aggregation.py` | SE-audit D3 | Low |
 | CallawaySantAnna event-study bucket/weight construction is duplicated between the analytical aggregator (`staggered_aggregation.py::_aggregate_event_study`) and the multiplier bootstrap (`staggered_bootstrap.py`): both group (g,t) by `e = t - g`, apply the finite/NaN/reference masks, and read cohort weights. Both already consume the same source-materialized universal reference cells (so they agree), but the bucket logic is copy-pasted. Extract one shared helper returning per-event-time buckets (finite cells, NaN cells, reference flags, cohort weights, combined-IF inputs) used by both. Pure refactor; gate on byte-identical analytical + bootstrap output. | `staggered_aggregation.py`, `staggered_bootstrap.py` | SE-audit D3 | Low |
 | `StackedDiD` survey re-resolution intra-file dedup (raw-weight extraction ×3, compose-normalize ×3, resolve-on-stacked ×2). The cross-estimator ContinuousDiD/EfficientDiD panel-to-unit collapse consolidation LANDED (#226 shared helpers `ResolvedSurveyDesign.subset_to_units_by_row_idx` / `build_unit_first_row_index`); StackedDiD is deliberately NOT on that path (control units are duplicated across sub-experiments, so it re-resolves at stacked granularity rather than collapsing to one row per unit). The residual is stacked-specific, low value, and touches the numerically-sensitive composed-weight renormalization. Post-filter re-resolution / metadata-recompute unification across the three estimators was assessed and is not warranted — they use genuinely different mechanisms and already delegate to shared `_resolve_survey_for_fit` / `compute_survey_metadata`. | `stacked_did.py` | #226 | Low |

diff --git a/diff_diff/linalg.py b/diff_diff/linalg.py
@@ -1917,8 +1917,10 @@ 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``, and the per-cluster ``A_g``
-    eigendecompositions) are defined in exactly one place.
+    (``bread_inv``; the per-cluster low-rank ``A_g`` factors on the
+    unweighted path; ``S_W``, ``MUWTWUM``, and the per-cluster dense ``A_g``
+    eigendecompositions on the weighted path) 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.
@@ -2023,23 +2025,94 @@ def _compute_cr2_bm_vcov_and_dof(
     # Per-cluster indices
     cluster_idx = {g: np.where(cluster_ids_arr == g)[0] for g in unique_clusters}
 
+    vcov: Optional[np.ndarray]
+    if weights is None:
+        # LOW-RANK FACTORED A_g (2026-07, evaluation change — algebraically
+        # identical, parity ~1e-15): unweighted G_g = I - H_gg with
+        # H_gg = X_g M_U X_g' of rank <= k, so with U_g = X_g M_U^{1/2} and
+        # U_g' U_g = Q diag(lam) Q', the adjustment operator is
+        #
+        #   A_g = (I - U_g U_g')^{-1/2}
+        #       = I + (U_g Q) diag(gamma) (U_g Q)',
+        #   gamma_i = ((1 - lam_i)^{-1/2} - 1) / lam_i          (regular)
+        #   gamma_i = -1 / lam_i                                 (1 - lam_i <= tol:
+        #                                                        Moore-Penrose zeroing,
+        #                                                        e.g. absorbed cluster FE)
+        #
+        # matching `_cr2_adjustment_matrix`'s pseudoinverse convention (its
+        # tol=1e-10 acts on G_g eigenvalues, which on span(U_g) are exactly
+        # `1 - lam_i`; the orthogonal complement carries eigenvalue 1 and is
+        # untouched by construction). Consumers only ever need A_g applied to
+        # skinny matrices (the residual vector for the meat; X_g bread_inv
+        # for the DOF omegas), so the dense (n_g, n_g) A_g — previously an
+        # O(n_g^3) eigh per cluster, ~85% of CR2-BM runtime at n=100k — is
+        # never materialized: per-cluster work is O(n_g k min(n_g, k) +
+        # min(n_g, k)^3) — the eigenproblem is solved on the SMALLER Gram
+        # side (k x k when n_g > k; the tiny n_g x n_g dense construction
+        # when n_g <= k, so small/singleton clusters never regress vs the
+        # prior dense path). gamma is evaluated via expm1/log1p, stable as
+        # lam -> 0 (limit 1/2).
+        # (The unweighted S_W collapse also applies: S_W = X'X = bread_matrix
+        # and MUWTWUM = M_U, so the bias-term build is skipped entirely.)
+        wM, VM = np.linalg.eigh(0.5 * (M_U + M_U.T))
+        M_U_half = (VM * np.sqrt(np.maximum(wM, 0.0))) @ VM.T
+        cluster_scores = np.zeros((G, k)) if residuals is not None else None
+        A_g_Xbi: Dict[Any, np.ndarray] = {}
+        for gi, g in enumerate(unique_clusters):
+            idx_g = cluster_idx[g]
+            X_g = X[idx_g]
+            U_g = X_g @ M_U_half
+            B_g = X_g @ bread_inv
+            n_g = len(idx_g)
+            if n_g <= k:
+                # Small cluster (n_g <= k): the smaller Gram side is the
+                # n_g x n_g one, so the k x k eigenproblem would REGRESS vs
+                # the dense per-cluster construction (e.g. paired designs or
+                # singleton-heavy clusterings with many covariates). Build
+                # G_g = I - U_g U_g' directly — it is tiny — and reuse the
+                # dense pseudoinverse convention verbatim.
+                A_g_small = _cr2_adjustment_matrix(np.eye(n_g) - U_g @ U_g.T)
+                if residuals is not None:
+                    cluster_scores[gi] = X_g.T @ (A_g_small @ residuals[idx_g])
+                A_g_Xbi[g] = A_g_small @ B_g
+                continue
+            lam, Q_g = np.linalg.eigh(U_g.T @ U_g)
+            lam = np.maximum(lam, 0.0)
+            s_vals = 1.0 - lam
+            gamma = np.zeros(k)
+            regular = (lam > 0) & (s_vals > 1e-10)
+            gamma[regular] = np.expm1(-0.5 * np.log1p(-lam[regular])) / lam[regular]
+            pseudo = (lam > 0) & (s_vals <= 1e-10)
+            gamma[pseudo] = -1.0 / lam[pseudo]
+            UQ = U_g @ Q_g
+            if residuals is not None:
+                u_g = residuals[idx_g]
+                # s_g = X_g' A_g u_g = X_g'u_g + (X_g'UQ) diag(gamma) (UQ'u_g)
+                cluster_scores[gi] = X_g.T @ u_g + (X_g.T @ UQ) @ (gamma * (UQ.T @ u_g))
+            A_g_Xbi[g] = B_g + UQ @ (gamma[:, np.newaxis] * (UQ.T @ B_g))
+        if residuals is not None:
+            meat = cluster_scores.T @ cluster_scores
+            vcov = M_U @ meat @ M_U
+        else:
+            vcov = None
+        dof_vec = _cr2_bm_dof_inner(X, A_g_Xbi, cluster_idx, M_U, contrasts)
+        return vcov, dof_vec
+
+    # --- WEIGHTED PATH (clubSandwich WLS-CR2; dense per-cluster A_g) ---
     # S_W = sum_g X_g' diag(W_norm_g²) X_g (used for both vcov and DOF).
-    # For unweighted (W_norm=1): S_W = X'X = bread_matrix.
     S_W = np.zeros((k, k))
     for g in unique_clusters:
         idx_g = cluster_idx[g]
         X_g = X[idx_g]
         S_W += X_g.T @ (X_g * (W_norm[idx_g] ** 2)[:, np.newaxis])
     MUWTWUM = M_U @ S_W @ M_U
 
-    # Per-cluster adjustment matrices A_g via G_g (which collapses to
-    # `I - H_gg` in the unweighted special case).
+    # Per-cluster adjustment matrices A_g via G_g.
     A_g_matrices: Dict[Any, np.ndarray] = {}
     for g in unique_clusters:
         idx_g = cluster_idx[g]
         X_g = X[idx_g]
-        # Asymmetric weighted hat block. For unweighted (W_norm=1): H_gg is the
-        # standard symmetric `X_g @ M_U @ X_g.T`.
+        # Asymmetric weighted hat block.
         H_gg = (X_g @ M_U @ X_g.T) * W_norm[idx_g][np.newaxis, :]
         I_g = np.eye(len(idx_g))
         bias_term = X_g @ MUWTWUM @ X_g.T
@@ -2049,7 +2122,6 @@ def _compute_cr2_bm_vcov_and_dof(
     # --- VCOV (meat) --- only when residuals are supplied (DOF-only callers
     # pass residuals=None and skip this).
     # Per-cluster score: s_g = X_g' diag(W_norm_g) A_g u_g.
-    vcov: Optional[np.ndarray]
     if residuals is not None:
         cluster_scores = np.zeros((G, k))
         for gi, g in enumerate(unique_clusters):
@@ -2063,24 +2135,18 @@ def _compute_cr2_bm_vcov_and_dof(
     else:
         vcov = None
 
-    # --- Per-contrast Bell-McCaffrey cluster DOF ---
-    # The inner helper branches on `weights`: unweighted uses the simple
-    # `(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:
-        dof_vec = _cr2_bm_dof_inner_weighted(
-            X,
-            A_g_matrices,
-            cluster_idx,
-            M_U,
-            MUWTWUM,
-            W_norm,
-            contrasts,
-            w_scale=w_scale,
-        )
+    # --- Per-contrast Bell-McCaffrey cluster DOF (weighted: the full
+    # clubSandwich P_array construction) ---
+    dof_vec = _cr2_bm_dof_inner_weighted(
+        X,
+        A_g_matrices,
+        cluster_idx,
+        M_U,
+        MUWTWUM,
+        W_norm,
+        contrasts,
+        w_scale=w_scale,
+    )
 
     return vcov, dof_vec
 
@@ -2168,16 +2234,18 @@ def _compute_cr2_bm(
 
 def _cr2_bm_dof_inner(
     X: np.ndarray,
-    A_g_matrices: Dict[Any, np.ndarray],
+    A_g_Xbi: Dict[Any, np.ndarray],
     cluster_idx: Dict[Any, np.ndarray],
     bread_inv: np.ndarray,
     contrasts: np.ndarray,
 ) -> np.ndarray:
     """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 per-cluster adjustment
-    matrices ``A_g_matrices``, cluster index map ``cluster_idx``, and
+    ``contrasts`` (shape ``(k, m)``), using the per-cluster precomputed
+    ``A_g_Xbi[g] = A_g @ X_g @ bread_inv`` blocks (built by the shared core
+    via the low-rank factored ``A_g`` apply — the dense ``A_g`` is never
+    materialized), 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
@@ -2231,13 +2299,10 @@ 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.
+    # Precompute once: X_bi (the A_g_Xbi blocks arrive precomputed from the
+    # shared core's factored A_g apply). For unit-contrast inputs
+    # (contrasts=I_k): q[:, j] == X_bi[:, j] == X @ bread_inv @ e_j.
     X_bi = X @ bread_inv  # (n, k)
-    A_g_Xbi = {
-        g: A_g_matrices[g] @ X[cluster_idx[g]] @ bread_inv for g in unique_clusters
-    }  # each (n_g, k)
     # 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.