From 974455d2b7f7c40ddfc91b6fcb6076863efe4a92 Mon Sep 17 00:00:00 2001 From: igerber Date: Sun, 28 Jun 2026 07:41:26 -0400 Subject: [PATCH] perf(linalg,mpd): compute MPD cluster+hc2_bm CR2 precomputes once MultiPeriodDiD(cluster=..., vcov_type="hc2_bm") built the expensive CR2 Bell-McCaffrey per-cluster precomputes (A_g eigendecompositions, S_W, the residual-maker M = I - H) twice per fit: once in solve_ols's vcov path (whose per-coefficient DOF was then discarded) and again in the separate _compute_cr2_bm_contrast_dof call for the per-coef + post-period-average ATT Satterthwaite DOF. Extract a shared core `_compute_cr2_bm_vcov_and_dof` (vcov + DOF for arbitrary contrasts from one precompute build); `_compute_cr2_bm` and _compute_cr2_bm_contrast_dof` become thin bit-identical wrappers over it (removing the formerly-duplicated ~50-line precompute block, so every CR2 caller routes through one implementation). The shared filter guards the residuals subscript with `residuals is not None` so DOF-only callers that pass weights with zero rows do not crash. At the MPD fit level, bypass solve_ols's vcov on the cluster+hc2_bm (non-survey, unweighted) path and compute vcov + the combined per-coef/avg-ATT DOF from one call, then expand with NaN for dropped columns; the `survey_weights is None` clause keeps the bypass byte-identical to solve_ols (else it falls back to the prior two-call behavior). Bit-identical vcov_, per-period DOF, avg-ATT DOF, p-values, and CIs (proven at atol=0 across balanced, unbalanced, and rank-deficient designs); a mechanism test asserts the per-cluster adjustment matrix is built exactly once per cluster (down from twice). No methodology, numerical, or public-API change. Updates the stale REGISTRY note that described the prior twice-built precompute. Resolves the MPD cluster+hc2_bm Performance row in TODO.md and completes the MultiPeriodDiD follow-up noted in the #565 LinearRegression entry. Co-Authored-By: Claude Opus 4.8 (1M context) --- CHANGELOG.md | 19 +- TODO.md | 1 - diff_diff/estimators.py | 71 ++++++-- diff_diff/linalg.py | 268 ++++++++++++++--------------- docs/methodology/REGISTRY.md | 16 +- tests/test_estimators_vcov_type.py | 130 +++++++++++++- tests/test_linalg_hc2_bm.py | 69 ++++++++ 7 files changed, 409 insertions(+), 165 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 652339058..6a878325c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -87,10 +87,25 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 DOF now come from a single combined call (the CR2 helper returns both together), halving the per-coefficient CR2 sandwich cost on every `hc2_bm` fit (weighted one-way WLS-CR2, weighted/unweighted clustered CR2-BM — i.e. `DiD`/`TWFE`/`MPD` with `vcov_type="hc2_bm"`). - (`MultiPeriodDiD`'s separate post-period-average contrast DOF still has its own recomputation, - tracked as an open follow-up in `TODO.md`.) **Bit-identical** SEs and DOF + (`MultiPeriodDiD`'s separate post-period-average contrast DOF recomputation is addressed by the + `MultiPeriodDiD` entry below.) **Bit-identical** SEs and DOF (proven at `atol=0`); no methodology, numerical, or public-API change. Minor side effect: the HC2→HC1 high-leverage-fallback `UserWarning` now fires once per fit instead of twice. +- **`MultiPeriodDiD(cluster=..., vcov_type="hc2_bm")`: build the CR2 precomputes once, not + twice.** The cluster + CR2 Bell-McCaffrey path built the expensive per-cluster precomputes + (the `A_g` adjustment-matrix eigendecompositions, `S_W`, the residual-maker `M = I − H`) twice + per fit — once in `solve_ols`'s vcov path (whose per-coefficient DOF was then discarded) and + again in the separate `_compute_cr2_bm_contrast_dof` call for the per-coefficient and + post-period-average ATT Satterthwaite DOF. The vcov **and** all DOF now come from a single + shared CR2 core (`_compute_cr2_bm_vcov_and_dof`) per fit; `_compute_cr2_bm` and + `_compute_cr2_bm_contrast_dof` become thin wrappers over it (removing the formerly-duplicated + ~50-line precompute block, so every CR2 caller routes through one implementation). The fit-level + call computes vcov and the per-coefficient + avg-ATT DOF together from one precompute build. + **Bit-identical** `vcov_`, per-period DOF, avg-ATT DOF, p-values, and CIs (proven at `atol=0` + across balanced, unbalanced, and rank-deficient designs; a mechanism test asserts the + per-cluster adjustment matrix is built exactly once per cluster, down from twice). Completes the + `MultiPeriodDiD` follow-up noted in the `LinearRegression` entry above. No methodology, + numerical, or public-API change. ### Fixed - **Corrected the Korn & Graubard (1990) citation venue** in `docs/methodology/REGISTRY.md` diff --git a/TODO.md b/TODO.md index 5eec47029..b3f63915e 100644 --- a/TODO.md +++ b/TODO.md @@ -56,7 +56,6 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| | `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 | -| MPD `cluster+hc2_bm` computes CR2 precomputes twice — `solve_ols → _compute_cr2_bm` for vcov+DOF, then `_compute_cr2_bm_contrast_dof` for the post-period-average contrast DOF. Both rebuild `H`, `M`, per-cluster `A_g`. Plumb the contrast DOF through the vcov path or share via a cached helper. | `linalg.py`, `estimators.py::MultiPeriodDiD.fit` | follow-up | Mid | Low | | 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 HC2: the Rust path only supports HC1; HC2 and CR2 Bell-McCaffrey fall through to NumPy. Noticeable for large-`n` fits. | `rust/src/linalg.rs` | Phase 1a | Mid | Low | | `SyntheticControl` retains a full `_SyntheticControlFitSnapshot` (pivoted panels) on EVERY fit for the opt-in `in_space_placebo()`, so callers who never run the placebo pay `O(units × periods × predictor-vars)` memory. Store a compact array/index representation, or build the snapshot lazily on first placebo call. | `synthetic_control.py`, `synthetic_control_results.py` | follow-up | Mid | Low | diff --git a/diff_diff/estimators.py b/diff_diff/estimators.py index 22d8d1acf..c6ae3b49e 100644 --- a/diff_diff/estimators.py +++ b/diff_diff/estimators.py @@ -1754,6 +1754,27 @@ def fit( # type: ignore[override] # Remap implicit "classical" + cluster to CR1 (legacy backward compat). _fit_vcov_type = self._resolve_effective_vcov_type(effective_cluster_ids) + # Cluster + CR2 Bell-McCaffrey (non-survey, unweighted) shares the SAME + # expensive CR2 precomputes (per-cluster A_g eigendecompositions, S_W, + # the residual-maker M) between the vcov and the per-coef/avg-ATT + # contrast DOF. Rather than let `solve_ols` build them for the vcov and + # then rebuild them for the contrast DOF below, skip solve_ols's vcov on + # this path and compute vcov + DOF together via one + # `_compute_cr2_bm_vcov_and_dof` call (see the hc2_bm DOF block). The + # `survey_weights is None` clause keeps the bypass byte-identical to + # solve_ols: solve_ols would pass `weights=survey_weights`, and the + # one-call path passes `weights=None`, so they only agree when + # survey_weights is None (the documented contract on this path — survey + # designs route through the TSL/replicate paths). If a future weighted + # entry point set survey_weights here, the flag is False and the code + # falls back to the original two-call behavior. + _is_mpd_cr2_path = ( + _fit_vcov_type == "hc2_bm" + and effective_cluster_ids is not None + and not _use_survey_vcov + and survey_weights is None + ) + # Resolve Conley arrays from column names (init-time) plus the # estimator's `time` / `unit` columns. CRITICAL: read from the # ORIGINAL `data` frame, NOT `working_data` — if absorb is used @@ -1790,7 +1811,7 @@ def fit( # type: ignore[override] X, y, return_fitted=True, - return_vcov=not _use_survey_vcov, + return_vcov=(not _use_survey_vcov) and not _is_mpd_cr2_path, cluster_ids=effective_cluster_ids, column_names=var_names, rank_deficient_action=self.rank_deficient_action, @@ -1945,15 +1966,22 @@ def _refit_mp_absorb(w_r): # rather than the shared n-k fallback. _bm_dof_per_coef: Optional[np.ndarray] = None _bm_dof_avg: Optional[float] = None + # On the `_is_mpd_cr2_path` bypass, solve_ols did not compute vcov + # (return_vcov=False). If every coefficient was dropped, synthesize the + # all-NaN vcov solve_ols would have returned (linalg.py:1230/1019); the + # BM DOF block below is skipped (no identified coefficients). + if _is_mpd_cr2_path and vcov is None and np.all(np.isnan(coefficients)): + vcov = np.full((X.shape[1], X.shape[1]), np.nan) if ( self.vcov_type == "hc2_bm" and not _use_survey_vcov - and vcov is not None + and (vcov is not None or _is_mpd_cr2_path) and not np.all(np.isnan(coefficients)) ): from diff_diff.linalg import ( _compute_bm_dof_from_contrasts, _compute_cr2_bm_contrast_dof, + _compute_cr2_bm_vcov_and_dof, _compute_hat_diagonals, ) @@ -1989,16 +2017,37 @@ def _refit_mp_absorb(w_r): contrasts, weights=survey_weights, ) + elif _is_mpd_cr2_path: + # Cluster-aware CR2 BM: vcov AND the per-coefficient + + # post-period-average compound contrast (Gate 6 lift) DOF + # from a SINGLE precompute build — the perf dedup. solve_ols + # bypassed vcov on this path, so compute it here from the + # same (X_kept, residuals, cluster_ids, bread_kept) it would + # have used internally (→ byte-identical vcov), then expand + # with NaN for dropped columns. weights=None per the + # _is_mpd_cr2_path guard (survey_weights is None here; survey + # designs route through the TSL path). + _vcov_reduced, _dof_all = _compute_cr2_bm_vcov_and_dof( + X_kept, + effective_cluster_ids, + bread_kept, + contrasts, + residuals=residuals, + weights=None, + ) + vcov = _expand_vcov_with_nan(_vcov_reduced, X.shape[1], _kept) else: - # Cluster-aware CR2 BM Satterthwaite DOF for per-coefficient - # AND post-period-average compound contrast (Gate 6 lift). - # This branch is guarded above by `not _use_survey_vcov`, - # so when reached, `survey_weights` is None (survey designs - # always route through the TSL path). The clubSandwich - # WLS-CR2 port lifted `_compute_cr2_bm_contrast_dof` to - # accept `weights=`, but no MPD entry point currently - # passes non-None weights here — `weights=None` is the - # de facto contract on this code path today. + # Defensive fallback, currently UNREACHABLE: reaching the + # cluster sub-branch requires `not _use_survey_vcov`, and any + # non-None `survey_weights` comes from a SurveyDesign whose + # `needs_survey_vcov` is True — so `_use_survey_vcov` would be + # True and the whole hc2_bm block is skipped. Hence + # `_is_mpd_cr2_path` is always satisfied here in practice (it + # already requires `survey_weights is None`), and this branch + # only fails safe for a future weighted-cluster entry point. + # It mirrors the pre-refactor call EXACTLY (solve_ols's vcov is + # kept; DOF is computed weights-free, exactly as the prior code + # did) so it adds no new, untested weighted-CR2 DOF behavior. _dof_all = _compute_cr2_bm_contrast_dof( X_kept, effective_cluster_ids, diff --git a/diff_diff/linalg.py b/diff_diff/linalg.py index 69626709f..f43efed6b 100644 --- a/diff_diff/linalg.py +++ b/diff_diff/linalg.py @@ -1709,55 +1709,48 @@ def _cr2_adjustment_matrix(G_g: np.ndarray, tol: float = 1e-10) -> np.ndarray: return (eigvecs * inv_sqrt) @ eigvecs.T -def _compute_cr2_bm( +def _compute_cr2_bm_vcov_and_dof( X: np.ndarray, - residuals: np.ndarray, cluster_ids: np.ndarray, bread_matrix: np.ndarray, + contrasts: np.ndarray, + residuals: Optional[np.ndarray] = None, weights: Optional[np.ndarray] = None, -) -> Tuple[np.ndarray, np.ndarray]: - """CR2 Bell-McCaffrey cluster-robust variance with per-coefficient DOF. - - Implements ``clubSandwich::vcovCR(..., type="CR2") + coef_test(test= - "Satterthwaite")`` for both unweighted and weighted ``lm`` fits. The - weighted form uses clubSandwich's specific WLS-CR2 algebra (which is - NOT a textbook PT2018 §3.3 transform-once derivation; see - docs/methodology/REGISTRY.md for the algorithm details). - - For each cluster ``g`` (with normalized weights ``W_norm = w / mean(w)``): - - ``H_gg = X_g M_U X_g' W_g`` (asymmetric weighted hat; W not sqrt(W)) - - ``S_W = sum_g X_g' W_g² X_g`` (W² in the bias-correction term) - - ``G_g = I - H_gg - H_gg' + X_g M_U S_W M_U X_g'`` - - ``A_g = G_g^{-1/2}`` via symmetric eigendecomposition with - pseudoinverse handling (see :func:`_cr2_adjustment_matrix`). - - Per-cluster score ``s_g = X_g' W_g A_g u_g`` (u_g = raw residual) - - 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). - - 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). - - Per-coefficient Satterthwaite DOF: see :func:`_cr2_bm_dof_inner` for the - unweighted simple formula and the weighted full P_array construction. +) -> Tuple[Optional[np.ndarray], np.ndarray]: + """Shared CR2 Bell-McCaffrey core — build precomputes once, return ``(vcov, dof)``. + + Single source of truth for the CR2 sandwich and the Satterthwaite 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. + 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. Parameters ---------- X : ndarray of shape (n, k) - residuals : ndarray of shape (n,) - Raw residuals ``y - X beta_hat`` from the (weighted) fit. cluster_ids : ndarray of shape (n,) bread_matrix : ndarray of shape (k, k) ``X'WX`` if weighted, ``X'X`` if unweighted. + contrasts : ndarray of shape (k, m) + Each column is a contrast vector for the Satterthwaite DOF. The + per-coefficient case is recovered with ``contrasts=np.eye(k)``. + residuals : ndarray of shape (n,), optional + Raw residuals ``y - X beta_hat`` from the (weighted) fit. When ``None`` + the meat / vcov is skipped and ``vcov`` is returned as ``None`` + (DOF-only callers); the per-cluster precomputes and DOF are unaffected. weights : ndarray of shape (n,), optional Original (un-normalized) weights. ``None`` for unweighted. Returns ------- - vcov : ndarray of shape (k, k) - dof_vec : ndarray of shape (k,) + vcov : ndarray of shape (k, k) or None + ``None`` when ``residuals is None``. + dof_vec : ndarray of shape (m,) + Satterthwaite DOF per contrast column. """ n, k = X.shape cluster_ids_arr = np.asarray(cluster_ids) @@ -1777,7 +1770,11 @@ def _compute_cr2_bm( positive_mask = weights_arr > 0 if not np.all(positive_mask): X = X[positive_mask] - residuals = residuals[positive_mask] + # DOF-only callers pass `residuals=None`; guard the subscript so the + # shared filter does not blow up on them (e.g. StackedDiD's weighted + # contrast-DOF path and the weighted singleton-cluster dispatch). + if residuals is not None: + residuals = residuals[positive_mask] cluster_ids_arr = cluster_ids_arr[positive_mask] weights_arr = weights_arr[positive_mask] weights = weights_arr # Rebind for downstream w_scale/W_norm logic @@ -1801,6 +1798,8 @@ def _compute_cr2_bm( G = len(unique_clusters) if G < 2: raise ValueError(f"Need at least 2 clusters for cluster-robust SEs, got {G}") + if contrasts.ndim != 2 or contrasts.shape[0] != k: + raise ValueError(f"contrasts must have shape (k={k}, m); got {contrasts.shape}") try: bread_inv = np.linalg.solve(bread_matrix, np.eye(k)) @@ -1853,27 +1852,32 @@ def _compute_cr2_bm( G_g = I_g - H_gg - H_gg.T + bias_term A_g_matrices[g] = _cr2_adjustment_matrix(G_g) - # --- VCOV (meat) --- + # --- 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. - cluster_scores = np.zeros((G, k)) - for gi, g in enumerate(unique_clusters): - idx_g = cluster_idx[g] - u_g = residuals[idx_g] - A_g = A_g_matrices[g] - adjusted = A_g @ u_g - cluster_scores[gi] = X[idx_g].T @ (W_norm[idx_g] * adjusted) - meat = cluster_scores.T @ cluster_scores - vcov = M_U @ meat @ M_U - - # --- Per-coefficient Bell-McCaffrey cluster DOF --- - # Delegate to the contrast-aware helper. The 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. + vcov: Optional[np.ndarray] + if residuals is not None: + cluster_scores = np.zeros((G, k)) + for gi, g in enumerate(unique_clusters): + idx_g = cluster_idx[g] + u_g = residuals[idx_g] + A_g = A_g_matrices[g] + adjusted = A_g @ u_g + cluster_scores[gi] = X[idx_g].T @ (W_norm[idx_g] * adjusted) + meat = cluster_scores.T @ cluster_scores + vcov = M_U @ meat @ M_U + else: + vcov = None + + # --- 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. 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, np.eye(k)) + dof_vec = _cr2_bm_dof_inner(X, M, A_g_matrices, cluster_idx, M_U, contrasts) else: dof_vec = _cr2_bm_dof_inner_weighted( X, @@ -1882,13 +1886,78 @@ def _compute_cr2_bm( M_U, MUWTWUM, W_norm, - np.eye(k), + contrasts, w_scale=w_scale, ) return vcov, dof_vec +def _compute_cr2_bm( + X: np.ndarray, + residuals: np.ndarray, + cluster_ids: np.ndarray, + bread_matrix: np.ndarray, + weights: Optional[np.ndarray] = None, +) -> Tuple[np.ndarray, np.ndarray]: + """CR2 Bell-McCaffrey cluster-robust variance with per-coefficient DOF. + + Implements ``clubSandwich::vcovCR(..., type="CR2") + coef_test(test= + "Satterthwaite")`` for both unweighted and weighted ``lm`` fits. The + weighted form uses clubSandwich's specific WLS-CR2 algebra (which is + NOT a textbook PT2018 §3.3 transform-once derivation; see + docs/methodology/REGISTRY.md for the algorithm details). + + For each cluster ``g`` (with normalized weights ``W_norm = w / mean(w)``): + - ``H_gg = X_g M_U X_g' W_g`` (asymmetric weighted hat; W not sqrt(W)) + - ``S_W = sum_g X_g' W_g² X_g`` (W² in the bias-correction term) + - ``G_g = I - H_gg - H_gg' + X_g M_U S_W M_U X_g'`` + - ``A_g = G_g^{-1/2}`` via symmetric eigendecomposition with + pseudoinverse handling (see :func:`_cr2_adjustment_matrix`). + - Per-cluster score ``s_g = X_g' W_g A_g u_g`` (u_g = raw residual) + + 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). + + 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). + + Per-coefficient Satterthwaite DOF: see :func:`_cr2_bm_dof_inner` for the + unweighted simple formula and the weighted full P_array construction. + + Parameters + ---------- + X : ndarray of shape (n, k) + residuals : ndarray of shape (n,) + Raw residuals ``y - X beta_hat`` from the (weighted) fit. + cluster_ids : ndarray of shape (n,) + bread_matrix : ndarray of shape (k, k) + ``X'WX`` if weighted, ``X'X`` if unweighted. + weights : ndarray of shape (n,), optional + Original (un-normalized) weights. ``None`` for unweighted. + + Returns + ------- + vcov : ndarray of shape (k, k) + dof_vec : ndarray of shape (k,) + """ + # Thin wrapper: per-coefficient vcov + DOF is the shared core with + # `contrasts = I_k`. See :func:`_compute_cr2_bm_vcov_and_dof` for the + # single-source-of-truth implementation. + vcov, dof_vec = _compute_cr2_bm_vcov_and_dof( + X, + cluster_ids, + bread_matrix, + np.eye(X.shape[1]), + residuals=residuals, + weights=weights, + ) + assert vcov is not None # residuals provided ⇒ vcov computed + return vcov, dof_vec + + def _cr2_bm_dof_inner( X: np.ndarray, M: np.ndarray, @@ -2180,93 +2249,18 @@ def _compute_cr2_bm_contrast_dof( _compute_cr2_bm : per-coefficient DOF (calls this helper internally with ``contrasts=np.eye(k)``). """ - n, k = X.shape - cluster_ids_arr = np.asarray(cluster_ids) - unique_clusters = np.unique(cluster_ids_arr) - # Subpopulation invariance: physically filter `weights > 0` rows before - # building per-cluster matrices. Mirrors the same fix in `_compute_cr2_bm` - # (CI codex P0 on PR #475 round 2). Zero-weight rows must be inert; the - # caller's bread_matrix = X.T @ (X * w[:, None]) is invariant to their - # removal, so no bread rebuild is needed. - if weights is not None: - weights_arr_eff = np.asarray(weights, dtype=np.float64) - positive_mask = weights_arr_eff > 0 - if not np.all(positive_mask): - X = X[positive_mask] - cluster_ids_arr = cluster_ids_arr[positive_mask] - weights_arr_eff = weights_arr_eff[positive_mask] - weights = weights_arr_eff - n = X.shape[0] - unique_clusters = np.unique(cluster_ids_arr) - eff_clusters = np.array( - [g for g in unique_clusters if float(np.sum(weights_arr_eff[cluster_ids_arr == g])) > 0] - ) - if len(eff_clusters) < 2: - raise ValueError( - f"Need at least 2 clusters with positive total weight for " - f"cluster-robust SEs, got {len(eff_clusters)} effective " - f"clusters out of {len(unique_clusters)} unique." - ) - unique_clusters = eff_clusters - if len(unique_clusters) < 2: - raise ValueError( - f"Need at least 2 clusters for cluster-robust SEs, got " f"{len(unique_clusters)}" - ) - if contrasts.ndim != 2 or contrasts.shape[0] != k: - raise ValueError(f"contrasts must have shape (k={k}, m); got {contrasts.shape}") - - try: - bread_inv = np.linalg.solve(bread_matrix, np.eye(k)) - except np.linalg.LinAlgError as e: - if "Singular" in str(e): - raise ValueError( - "Design matrix is rank-deficient (singular X'X matrix). " - "Cannot compute CR2 Bell-McCaffrey variance." - ) from e - raise - - # Normalize weights (clubSandwich convention; w_scale cancels in DOF ratio). - # Use M_U = (X' W_norm X)^{-1} = w_scale * bread_inv throughout, matching - # clubSandwich's internal matrix used in the P_array construction. - if weights is not None: - weights_arr = np.asarray(weights, dtype=np.float64) - pos = weights_arr > 0 - w_scale = float(np.mean(weights_arr[pos])) if np.any(pos) else 1.0 - W_norm = weights_arr / w_scale - M_U = w_scale * bread_inv - else: - W_norm = np.ones(n, dtype=np.float64) - M_U = bread_inv - - cluster_idx = {g: np.where(cluster_ids_arr == g)[0] for g in unique_clusters} - - # S_W = sum_g X_g' diag(W_norm_g²) X_g. For unweighted: 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 - - 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 (collapses to symmetric when W_norm=1). - 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 - G_g = I_g - H_gg - H_gg.T + bias_term - A_g_matrices[g] = _cr2_adjustment_matrix(G_g) - - if weights is None: - # Bit-equal to prior unweighted path: simple (tr B)² / tr(B²) form. - H = X @ M_U @ X.T - M = np.eye(n) - H - return _cr2_bm_dof_inner(X, M, A_g_matrices, cluster_idx, M_U, contrasts) - return _cr2_bm_dof_inner_weighted( - X, A_g_matrices, cluster_idx, M_U, MUWTWUM, W_norm, contrasts, w_scale=w_scale + # Thin wrapper: DOF-only is the shared core with `residuals=None` (the meat + # / vcov is skipped). See :func:`_compute_cr2_bm_vcov_and_dof` for the + # single-source-of-truth implementation of the per-cluster precomputes. + _, dof_vec = _compute_cr2_bm_vcov_and_dof( + X, + cluster_ids, + bread_matrix, + contrasts, + residuals=None, + weights=weights, ) + return dof_vec def _compute_bm_dof_from_contrasts( diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index a5d36462d..2405db13a 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -233,13 +233,15 @@ where V is the VCV sub-matrix for post-treatment δ_e coefficients. `Wald_test(constraints=matrix(c, 1), test="HTZ")$df_denom`. Weighted CR2-BM (`survey_design=` paths) is still gated separately; see the rows in `TODO.md` under Methodology/Correctness. -- **Note:** the cluster-aware contrast-DOF path currently recomputes the CR2 hat - matrix and per-cluster adjustment matrices that `solve_ols` already built for the - vcov dispatch — clustered `hc2_bm` MPD fits pay the O(n²) CR2 setup twice in - exchange for a clean call-site contract. Acceptable for typical cluster-robust - DiD panel sizes (n ≤ few thousand); tracked in `TODO.md` under Performance for - a follow-up that plumbs the contrast DOF through the existing CR2 vcov path or - shares precomputes. +- **Note:** the cluster-aware MPD `hc2_bm` path computes the CR2 hat matrix, + 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 + `_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:** `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 a41f6a23f..b831e40f8 100644 --- a/tests/test_estimators_vcov_type.py +++ b/tests/test_estimators_vcov_type.py @@ -2181,14 +2181,22 @@ def test_collision_raises_on_all_paths(self, vcov_type, name): df = self._panel_with(name) with pytest.raises(ValueError, match="collide"): TwoWayFixedEffects(vcov_type=vcov_type).fit( - df, outcome="y", treatment="treated", time="time", unit="unit", + df, + outcome="y", + treatment="treated", + time="time", + unit="unit", covariates=[name], ) def test_hc2_full_dummy_noncolliding_preserves_coefs(self): df = self._panel_with("x1") r = TwoWayFixedEffects(vcov_type="hc2").fit( - df, outcome="y", treatment="treated", time="time", unit="unit", + df, + outcome="y", + treatment="treated", + time="time", + unit="unit", covariates=["x1"], ) ck = r.coefficients @@ -2199,7 +2207,11 @@ def test_hc2_full_dummy_noncolliding_preserves_coefs(self): def test_within_transform_noncolliding_returns_att_only(self): df = self._panel_with("x1") r = TwoWayFixedEffects().fit( # default hc1 -> within-transform - df, outcome="y", treatment="treated", time="time", unit="unit", + df, + outcome="y", + treatment="treated", + time="time", + unit="unit", covariates=["x1"], ) # The within-transform path exposes only the ATT coefficient by design; @@ -2215,13 +2227,117 @@ def test_within_path_does_not_materialize_fe_dummies(self, monkeypatch): df = self._panel_with("x1") def _boom(*args, **kwargs): - raise AssertionError( - "pd.get_dummies must not be called on the within-transform path" - ) + raise AssertionError("pd.get_dummies must not be called on the within-transform path") monkeypatch.setattr(pd, "get_dummies", _boom) r = TwoWayFixedEffects().fit( - df, outcome="y", treatment="treated", time="time", unit="unit", + df, + outcome="y", + treatment="treated", + time="time", + unit="unit", covariates=["x1"], ) assert set(r.coefficients.keys()) == {"ATT"} + + +class TestMPDClusterHC2BMSharedPrecompute: + """`MultiPeriodDiD(cluster=..., vcov_type='hc2_bm')` builds the CR2 + Bell-McCaffrey precomputes ONCE, not twice. + + 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`. + + (Absolute SE/DOF values are pinned independently by the R/clubSandwich + goldens in `test_multi_period_cluster_hc2_bm_avg_att_uses_clubsandwich_dof` + and `test_linalg_hc2_bm.py`; these tests guard the *new invariants* the + refactor introduces.) + """ + + @staticmethod + def _balanced_panel(): + rng = np.random.default_rng(2) + rows = [] + for i in range(20): + treated = int(i >= 10) + for t in range(3): + y = rng.normal(0.0, 1.0) + 0.5 * treated * (t >= 1) + rows.append({"unit": i, "time": t, "treated": treated, "y": y}) + return pd.DataFrame(rows) + + @staticmethod + def _unbalanced_panel(): + rng = np.random.default_rng(7) + rows = [] + for i in range(24): + treated = int(i >= 12) + periods = [0, 1, 2, 3] if (i % 3 != 0) else [0, 2, 3] + for t in periods: + y = rng.normal(0.0, 1.0) + 0.4 * treated * (t >= 2) + rows.append({"unit": i, "time": t, "treated": treated, "y": y}) + return pd.DataFrame(rows) + + @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.""" + 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 + calls = {"n": 0} + + def _counting(*args, **kwargs): + calls["n"] += 1 + return orig(*args, **kwargs) + + monkeypatch.setattr(L, "_cr2_adjustment_matrix", _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 + # 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 " + f"precompute build), got {calls['n']} — the precompute is being " + "built more than once." + ) + # Inference is still finite (the dedup did not break the path). + assert np.isfinite(res.avg_att) and np.isfinite(res.avg_se) + assert np.isfinite(res.avg_p_value) + + def test_fit_is_reproducible(self): + """Two independent fits (and a repeat fit of the same estimator) give + identical avg-ATT and per-period inference — determinism + the + fit-does-not-mutate-config contract on the bypass path.""" + data = self._balanced_panel() + + est = MultiPeriodDiD(vcov_type="hc2_bm", cluster="unit") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r1 = est.fit(data, outcome="y", treatment="treated", time="time", unit="unit") + # Repeat fit of the SAME object (config must be unmutated). + r2 = est.fit(data, outcome="y", treatment="treated", time="time", unit="unit") + # Fresh object, same config. + r3 = MultiPeriodDiD(vcov_type="hc2_bm", cluster="unit").fit( + data, outcome="y", treatment="treated", time="time", unit="unit" + ) + + for r in (r2, r3): + assert r.avg_att == r1.avg_att + assert r.avg_se == r1.avg_se + assert r.avg_t_stat == r1.avg_t_stat + assert r.avg_p_value == r1.avg_p_value + assert set(r.period_effects) == set(r1.period_effects) + for p in r1.period_effects: + assert r.period_effects[p].effect == r1.period_effects[p].effect + assert r.period_effects[p].se == r1.period_effects[p].se diff --git a/tests/test_linalg_hc2_bm.py b/tests/test_linalg_hc2_bm.py index 14a019e4e..ebb01d7d4 100644 --- a/tests/test_linalg_hc2_bm.py +++ b/tests/test_linalg_hc2_bm.py @@ -28,6 +28,7 @@ _compute_bm_dof_oneway, _compute_cr2_bm, _compute_cr2_bm_contrast_dof, + _compute_cr2_bm_vcov_and_dof, _compute_hat_diagonals, _cr2_adjustment_matrix, compute_robust_vcov, @@ -769,3 +770,71 @@ def test_too_few_clusters_raises(self): bread = X.T @ X with pytest.raises(ValueError, match=r"[Nn]eed at least 2 clusters"): _compute_cr2_bm_contrast_dof(X, cluster, bread, np.eye(k)) + + def test_wrappers_are_bit_identical_to_shared_core(self): + """`_compute_cr2_bm` and `_compute_cr2_bm_contrast_dof` are thin + wrappers over `_compute_cr2_bm_vcov_and_dof`, so they must reproduce + the core's output exactly (atol=0 / array_equal). + + This is the refactor's structural guard: the per-coefficient vcov+DOF + path (`contrasts=eye(k)`, residuals provided) and the DOF-only contrast + path (`residuals=None`) both route through the single core, unweighted + and weighted. + """ + rng = np.random.default_rng(20260628) + n, k = 40, 3 + X = np.column_stack([np.ones(n), rng.standard_normal(n), rng.standard_normal(n)]) + beta = np.array([0.5, 1.0, -0.7]) + y = X @ beta + rng.standard_normal(n) * 0.3 + cluster = np.repeat(np.arange(8), 5) + # Compound contrast: per-coef columns plus an average of the slopes. + C = np.column_stack([np.eye(k), np.array([0.0, 0.5, 0.5])]) + + for weights in (None, rng.uniform(0.5, 2.0, size=n)): + if weights is None: + bread = X.T @ X + else: + bread = X.T @ (X * weights[:, np.newaxis]) + coef = np.linalg.solve(bread, X.T @ (y if weights is None else y * weights)) + residuals = y - X @ coef + + # Per-coefficient vcov + DOF via wrapper vs core (eye(k)). + v_wrap, d_wrap = _compute_cr2_bm(X, residuals, cluster, bread, weights=weights) + v_core, d_core = _compute_cr2_bm_vcov_and_dof( + X, cluster, bread, np.eye(k), residuals=residuals, weights=weights + ) + assert np.array_equal(v_wrap, v_core) + assert np.array_equal(d_wrap, d_core, equal_nan=True) + + # DOF-only contrast path via wrapper vs core (residuals=None). + dof_wrap = _compute_cr2_bm_contrast_dof(X, cluster, bread, C, weights=weights) + vcov_core, dof_core = _compute_cr2_bm_vcov_and_dof( + X, cluster, bread, C, residuals=None, weights=weights + ) + assert vcov_core is None # no residuals -> meat/vcov skipped + assert np.array_equal(dof_wrap, dof_core, equal_nan=True) + + def test_dof_only_with_zero_weights_does_not_crash(self): + """DOF-only callers pass `weights=` AND `residuals=None`; the shared + core's zero-weight filter must guard the residuals subscript. + + Regression for the merge of `_compute_cr2_bm`'s filter (which subscripts + `residuals` unconditionally) into the shared core: without the + `residuals is not None` guard, a DOF-only call with zero weights would + raise `TypeError` (StackedDiD's weighted contrast-DOF path and the + weighted singleton-cluster dispatch hit exactly this). + """ + rng = np.random.default_rng(20260629) + n, k = 40, 3 + X = np.column_stack([np.ones(n), rng.standard_normal(n), rng.standard_normal(n)]) + cluster = np.repeat(np.arange(8), 5) + weights = np.ones(n) + # Zero out a few rows, keeping >=2 clusters with positive total weight. + weights[[0, 7, 13]] = 0.0 + bread = X.T @ (X * weights[:, np.newaxis]) + C = np.column_stack([np.eye(k), np.array([0.0, 0.5, 0.5])]) + + # Must not raise (the bug was an unconditional residuals[positive_mask]). + dof = _compute_cr2_bm_contrast_dof(X, cluster, bread, C, weights=weights) + assert dof.shape == (C.shape[1],) + assert np.all(np.isfinite(dof))