diff --git a/CHANGELOG.md b/CHANGELOG.md index 37966400..755ce2dc 100644 --- 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 index 25b53d1c..da84b936 100644 --- 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 index 85679a3a..84c4a873 100644 --- 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,8 +2025,81 @@ 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] @@ -2032,14 +2107,12 @@ def _compute_cr2_bm_vcov_and_dof( 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,7 +2234,7 @@ 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, @@ -2176,8 +2242,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 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. diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index a2c66615..e0891af5 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -244,7 +244,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. 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 (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). **Low-rank A_g factorization (2026-07):** the unweighted per-cluster adjustment operator `A_g = (I − H_gg)^{−1/2}` is evaluated from the k×k eigenproblem of `U_g'U_g` (`U_g = X_g M_U^{1/2}`; `A_g = I + (U_g Q) diag(γ) (U_g Q)'`, `γ = ((1−λ)^{−1/2}−1)/λ` via expm1/log1p, Moore-Penrose zeroing at `1−λ ≤ 1e-10` matching `_cr2_adjustment_matrix`'s convention) and only ever applied to skinny matrices — the dense `(n_g, n_g)` `A_g` (an `O(n_g³)` eigh per cluster) is never materialized for `n_g > k`; clusters with `n_g ≤ k` keep the (tiny) dense construction since it is the smaller Gram side. Algebraically identical; vcov/DOF match a frozen dense-eigh oracle at rtol 1e-12/1e-10 incl. the leverage-1 absorbed-cluster-FE and singleton-cluster regimes (`TestCR2BMLowRankAdjustment`). The weighted clubSandwich path keeps the dense construction (its `G_g` carries the `S_W` bias term). - **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_estimators_vcov_type.py b/tests/test_estimators_vcov_type.py index 00c0a86c..23be4857 100644 --- a/tests/test_estimators_vcov_type.py +++ b/tests/test_estimators_vcov_type.py @@ -2438,9 +2438,11 @@ class TestMPDClusterHC2BMSharedPrecompute: Mechanism guard for the perf dedup: vcov and the per-coefficient + post-period-average contrast DOF now come from a single `_compute_cr2_bm_vcov_and_dof` call, so the expensive per-cluster - adjustment matrices (`_cr2_adjustment_matrix`) are built exactly once per - cluster. Before the dedup, solve_ols's vcov path and the separate - contrast-DOF call each built them, i.e. `2 * G`. + precomputes are built exactly once. Before the dedup, solve_ols's vcov + path and the separate contrast-DOF call each built them. (The spy + counts shared-core calls rather than `_cr2_adjustment_matrix` calls + because the unweighted path no longer materializes dense per-cluster + adjustment matrices at all — the low-rank factored `A_g` apply.) (Absolute SE/DOF values are pinned independently by the R/clubSandwich goldens in `test_multi_period_cluster_hc2_bm_avg_att_uses_clubsandwich_dof` @@ -2473,31 +2475,31 @@ def _unbalanced_panel(): @pytest.mark.parametrize("which", ["balanced", "unbalanced"]) def test_cr2_precompute_built_once(self, which, monkeypatch): - """`_cr2_adjustment_matrix` is called exactly `G` times (one precompute - build), not `2 * G`, on the cluster+hc2_bm path.""" + """`_compute_cr2_bm_vcov_and_dof` (the shared core, whose body builds + the per-cluster precomputes exactly once) is entered exactly once on + the cluster+hc2_bm path — not once for vcov and again for DOF.""" import diff_diff.linalg as L data = self._balanced_panel() if which == "balanced" else self._unbalanced_panel() - n_clusters = data["unit"].nunique() - orig = L._cr2_adjustment_matrix + orig = L._compute_cr2_bm_vcov_and_dof calls = {"n": 0} def _counting(*args, **kwargs): calls["n"] += 1 return orig(*args, **kwargs) - monkeypatch.setattr(L, "_cr2_adjustment_matrix", _counting) + monkeypatch.setattr(L, "_compute_cr2_bm_vcov_and_dof", _counting) with warnings.catch_warnings(): warnings.simplefilter("ignore") res = MultiPeriodDiD(vcov_type="hc2_bm", cluster="unit").fit( data, outcome="y", treatment="treated", time="time", unit="unit" ) - # One build: exactly one adjustment matrix per cluster (was 2 * G when + # One build: the shared core entered exactly once (was twice when # solve_ols's vcov and the contrast-DOF call each built the precomputes). - assert calls["n"] == n_clusters, ( - f"Expected {n_clusters} _cr2_adjustment_matrix calls (one CR2 " + assert calls["n"] == 1, ( + f"Expected exactly 1 _compute_cr2_bm_vcov_and_dof call (one CR2 " f"precompute build), got {calls['n']} — the precompute is being " "built more than once." ) diff --git a/tests/test_linalg.py b/tests/test_linalg.py index 01c619cc..1f48eb20 100644 --- a/tests/test_linalg.py +++ b/tests/test_linalg.py @@ -2716,6 +2716,122 @@ def test_return_dropped_mask(self): assert nd0 == 3 and dropped0.all() +class TestCR2BMLowRankAdjustment: + """The low-rank factored A_g apply reproduces the dense per-cluster + eigendecomposition path (frozen here as the oracle) at ~1e-14: vcov, + per-coefficient DOF, and the NaN-guard pattern. The identity: unweighted + G_g = I - U_g U_g' (U_g = X_g M_U^{1/2}, rank <= k), so + A_g = I + (U_g Q) diag(gamma) (U_g Q)' from the k x k eigenproblem of + U_g'U_g — O(n_g k^2) per cluster instead of the O(n_g^3) dense eigh + (~85% of CR2-BM runtime at n=100k pre-change).""" + + @staticmethod + def _oracle_dense_cr2(X, residuals, cluster_ids, bread_matrix): + """Frozen pre-change dense evaluation: per-cluster G_g = I - H_gg, + A_g via _cr2_adjustment_matrix (dense eigh), meat from dense A_g @ u, + A_g_Xbi from dense A_g — fed to the production scores-based DOF.""" + from diff_diff.linalg import _cr2_adjustment_matrix, _cr2_bm_dof_inner + + n, k = X.shape + unique = np.unique(cluster_ids) + cluster_idx = {g: np.where(cluster_ids == g)[0] for g in unique} + bread_inv = np.linalg.solve(bread_matrix, np.eye(k)) + scores = np.zeros((len(unique), k)) + A_g_Xbi = {} + for gi, g in enumerate(unique): + idx = cluster_idx[g] + X_g = X[idx] + H = X_g @ bread_inv @ X_g.T + A_g = _cr2_adjustment_matrix(np.eye(len(idx)) - H) + scores[gi] = X_g.T @ (A_g @ residuals[idx]) + A_g_Xbi[g] = A_g @ X_g @ bread_inv + vcov = bread_inv @ (scores.T @ scores) @ bread_inv + dof = _cr2_bm_dof_inner(X, A_g_Xbi, cluster_idx, bread_inv, np.eye(k)) + return vcov, dof + + @staticmethod + def _design(n, G, k, singular_cluster=False, singletons=False, seed=7): + rng = np.random.default_rng(seed) + if singletons: + cl = np.arange(n) % G + else: + base = np.repeat(np.arange(G), n // G) + cl = np.concatenate([base, np.full(n - base.size, G - 1)]) + X = np.column_stack([np.ones(n), rng.normal(size=(n, k - 1))]) + if singular_cluster: + # absorbed cluster-0 FE: cluster 0 carries within-cluster + # leverage 1 -> G_g eigenvalue 0 -> Moore-Penrose zeroing branch + X[:, 1] = (cl == 0).astype(float) + y = X @ rng.normal(size=k) + rng.normal(size=n) + resid = y - X @ np.linalg.lstsq(X, y, rcond=None)[0] + return X, resid, cl, X.T @ X + + @pytest.mark.parametrize( + "kw", + [ + dict(n=2000, G=20, k=10), + dict(n=1777, G=13, k=8), + dict(n=4000, G=25, k=40), + dict(n=2000, G=20, k=10, singular_cluster=True), + dict(n=300, G=300, k=5, singletons=True), + dict(n=600, G=300, k=12), + ], + ids=[ + "balanced", + "unbalanced", + "k40", + "leverage1-absorbed-FE", + "singletons", + "tiny-clusters-k-gt-ng", + ], + ) + def test_matches_frozen_dense_oracle(self, kw): + from diff_diff.linalg import _compute_cr2_bm + + X, resid, cl, bread = self._design(**kw) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + vcov, dof = _compute_cr2_bm(X, resid, cl, bread) + vcov_o, dof_o = self._oracle_dense_cr2(X, resid, cl, bread) + # scale-aware atol: cross-evaluation ~1-ULP drift reads as large + # RELATIVE error on near-zero off-diagonals under other BLAS kernels + # (see test_gamma_stable_near_zero_leverage's tolerance note). + np.testing.assert_allclose(vcov, vcov_o, rtol=1e-12, atol=1e-13 * np.max(np.abs(vcov_o))) + assert np.array_equal(np.isnan(dof), np.isnan(dof_o)), "NaN-guard pattern diverged" + fin = ~np.isnan(dof_o) + if fin.any(): + np.testing.assert_allclose(dof[fin], dof_o[fin], rtol=1e-10) + + def test_gamma_stable_near_zero_leverage(self): + """A near-zero-leverage cluster (tiny lam) exercises the expm1/log1p + evaluation of gamma; naive (1/sqrt(1-lam)-1)/lam is catastrophically + cancellative there. The oracle comparison pins the stable branch. + + Tolerance note: rtol 1e-11, not 1e-12 — the low-rank and dense + evaluations reduce in different GEMM orders, and on this fixture's + huge 5000-row cluster the ~1-ULP absolute drift (measured 1.3e-17) + lands on small-magnitude off-diagonals where it reads as ~1.3e-12 + RELATIVE on OpenBLAS/linux-arm CI (exact-equal on Accelerate) — the + documented cross-evaluation BLAS-reassociation caveat.""" + from diff_diff.linalg import _compute_cr2_bm + + rng = np.random.default_rng(11) + # one huge cluster + many tiny ones -> tiny per-cluster leverage + cl = np.concatenate([np.zeros(5000, dtype=int), 1 + np.arange(200) % 40]) + n = cl.size + X = np.column_stack([np.ones(n), rng.normal(size=(n, 3))]) + y = X @ rng.normal(size=4) + rng.normal(size=n) + resid = y - X @ np.linalg.lstsq(X, y, rcond=None)[0] + bread = X.T @ X + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + vcov, dof = _compute_cr2_bm(X, resid, cl, bread) + vcov_o, dof_o = self._oracle_dense_cr2(X, resid, cl, bread) + np.testing.assert_allclose(vcov, vcov_o, rtol=1e-11) + fin = ~np.isnan(dof_o) + np.testing.assert_allclose(dof[fin], dof_o[fin], rtol=1e-10) + + class TestCR2BMScoresBasedDOF: """Frozen-oracle parity for the scores-based Satterthwaite DOF evaluation (algebraic identity; PT2018 §3.1 Satterthwaite DOF):