From e5e7362977e647eda3086a0444d2c44c8a2acf26 Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 27 Jun 2026 20:03:11 -0400 Subject: [PATCH] perf(linalg): compute hc2_bm CR2 sandwich once, not twice LinearRegression(vcov_type="hc2_bm").fit computed the CR2/HC2-BM sandwich twice: once inside solve_ols for the vcov, then again via compute_robust_vcov(return_dof=True) just to extract the per-coefficient Bell-McCaffrey DOF. The CR2 helper returns (vcov, dof) together, so the second sandwich (per-cluster A_g + the O(n^2 k) Satterthwaite loop) was pure waste (#475). Skip solve_ols's vcov on the hc2_bm-not-survey path (return_vcov=False via a new _is_bm_path flag) and source BOTH vcov and dof from one compute_robust_vcov(return_dof=True) call, reusing the existing _expand_vcov_with_nan rank-deficient expansion. solve_ols (and its ~50 callers) is left untouched. Affects every hc2_bm fit: weighted one-way WLS-CR2, weighted/unweighted clustered CR2-BM (DiD/TWFE/MPD with vcov_type="hc2_bm"). Bit-identical SEs + per-coef DOF (proven at atol=0 across weighted/ unweighted x one-way/clustered + rank-deficient; both vcov computations route through the same _compute_robust_vcov_numpy). No methodology, numerical, or public-API change. Minor: the HC2->HC1 high-leverage fallback UserWarning now fires once per fit instead of twice. Co-Authored-By: Claude Opus 4.8 (1M context) --- CHANGELOG.md | 11 ++++ TODO.md | 1 - diff_diff/linalg.py | 76 ++++++++++++++----------- tests/test_methodology_wls_cr2.py | 95 +++++++++++++++++++++++++++++++ 4 files changed, 148 insertions(+), 35 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 68c12688..d9b5fdfd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -68,6 +68,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 #groups)` factorizations to one. **Bit-identical** to the prior per-target `spsolve` (proven at `atol=0` across FE-only, covariate, survey-weighted, and bootstrap paths); no methodology, numerical, or public-API change. +- **`LinearRegression(vcov_type="hc2_bm")`: compute the CR2 sandwich once, not twice.** The fit + path previously computed the CR2/HC2-BM vcov inside `solve_ols`, then recomputed the entire + Bell-McCaffrey sandwich (per-cluster `A_g`, the `O(n²k)` Satterthwaite loop) a second time via + `compute_robust_vcov(return_dof=True)` just to extract the per-coefficient DOF. The vcov **and** + 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 + (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. ### Fixed - **Corrected the Korn & Graubard (1990) citation venue** in `docs/methodology/REGISTRY.md` diff --git a/TODO.md b/TODO.md index e6523cfe..2c3b95af 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 | -| `LinearRegression.fit()` pays the CR2 cost twice on the weighted `hc2_bm` path: once in `solve_ols(..., return_vcov=True)` and again via `compute_robust_vcov(..., return_dof=True)` for `_bm_dof`. Fix: thread `return_dof` through `solve_ols`, or cache the per-cluster `A_g` / `MUWTWUM` precomputes. (CI codex P3 on #475.) | `linalg.py` | #475 | Mid | Low | | 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 | diff --git a/diff_diff/linalg.py b/diff_diff/linalg.py index 8f08ba04..69626709 100644 --- a/diff_diff/linalg.py +++ b/diff_diff/linalg.py @@ -3542,14 +3542,20 @@ def fit( ) if _fit_vcov_type != "classical" or effective_cluster_ids is not None: - # Use solve_ols with robust/cluster SEs - # When survey vcov will be used, skip standard vcov computation + # Use solve_ols with robust/cluster SEs. + # When survey vcov will be used, skip standard vcov computation. + # For hc2_bm (non-survey), ALSO skip solve_ols's vcov: the CR2 + # sandwich produces (vcov, dof) together, so the block below gets + # BOTH from a single compute_robust_vcov(return_dof=True) call + # instead of computing the vcov here and recomputing the whole + # sandwich a second time just for the dof (#475). + _is_bm_path = _fit_vcov_type == "hc2_bm" and not _use_survey_vcov coefficients, residuals, fitted, vcov = solve_ols( X, y, cluster_ids=effective_cluster_ids, return_fitted=True, - return_vcov=not _use_survey_vcov, + return_vcov=(not _use_survey_vcov) and not _is_bm_path, rank_deficient_action=self.rank_deficient_action, weights=_fit_weights, weight_type=_fit_weight_type, @@ -3562,22 +3568,24 @@ def fit( conley_unit=self.conley_unit, conley_lag_cutoff=self.conley_lag_cutoff, ) - # For hc2_bm, compute per-coefficient Bell-McCaffrey DOF. Both - # the one-way HC2+BM case and the cluster CR2 case are supported, - # including the weighted clustered CR2 path via the clubSandwich - # WLS-CR2 port. The dispatcher already rejects non-pweight weight - # types for hc2_bm + weights, so reaching this block guarantees - # `_compute_cr2_bm` returns a finite per-coefficient DOF. - if ( - _fit_vcov_type == "hc2_bm" - and not _use_survey_vcov - and vcov is not None - and not np.all(np.isnan(coefficients)) - ): - # Identified columns for DOF (rank-deficient case sets NaN coefs). + # For hc2_bm (non-survey), compute the CR2 vcov AND per-coefficient + # Bell-McCaffrey DOF together in a SINGLE compute_robust_vcov call + # (solve_ols skipped its vcov above via `_is_bm_path`). Both the + # one-way HC2+BM case and the (weighted) clustered CR2 case route + # through the same `_compute_robust_vcov_numpy`/`_compute_cr2_bm`, so + # this is bit-identical to the prior two-call form while computing the + # O(n^2 k) sandwich once instead of twice (#475). The dispatcher + # already rejects non-pweight weight types for hc2_bm + weights. + if _is_bm_path: + # Rank-deficient solves set NaN coefficients for dropped columns. nan_mask = np.isnan(coefficients) - if not np.any(nan_mask): - _, self._bm_dof = compute_robust_vcov( + if np.all(nan_mask): + # All columns dropped: NaN vcov (matches solve_ols's + # rank-deficient all-drop) and no BM DOF (n-k fallback). + vcov = np.full((X.shape[1], X.shape[1]), np.nan) + self._bm_dof = None + elif not np.any(nan_mask): + vcov, self._bm_dof = compute_robust_vcov( X, residuals, cluster_ids=effective_cluster_ids, @@ -3587,23 +3595,23 @@ def fit( return_dof=True, ) else: - # Per-coef DOF only for identified coefficients; set NaN for dropped. + # Rank-deficient: compute on identified columns only, then + # expand BOTH vcov and per-coef DOF with NaN for dropped + # columns (mirrors solve_ols's `_expand_vcov_with_nan` path). kept = np.where(~nan_mask)[0] - if len(kept) > 0: - _, dof_kept = compute_robust_vcov( - X[:, kept], - residuals, - cluster_ids=effective_cluster_ids, - weights=_fit_weights, - weight_type=_fit_weight_type, - vcov_type="hc2_bm", - return_dof=True, - ) - full = np.full(X.shape[1], np.nan) - full[kept] = dof_kept - self._bm_dof = full - else: - self._bm_dof = np.full(X.shape[1], np.nan) + vcov_reduced, dof_kept = compute_robust_vcov( + X[:, kept], + residuals, + cluster_ids=effective_cluster_ids, + weights=_fit_weights, + weight_type=_fit_weight_type, + vcov_type="hc2_bm", + return_dof=True, + ) + vcov = _expand_vcov_with_nan(vcov_reduced, X.shape[1], kept) + full = np.full(X.shape[1], np.nan) + full[kept] = dof_kept + self._bm_dof = full else: self._bm_dof = None else: diff --git a/tests/test_methodology_wls_cr2.py b/tests/test_methodology_wls_cr2.py index 29feaf02..d507d534 100644 --- a/tests/test_methodology_wls_cr2.py +++ b/tests/test_methodology_wls_cr2.py @@ -461,6 +461,101 @@ def test_lr_weighted_oneway_hc2_bm_df_matches_helper(self): err_msg="One-way LinearRegression _bm_dof must match helper Satterthwaite DOF", ) + def test_lr_hc2_bm_vcov_value_matches_helper(self): + """#475 dedup: the fit-level ``self.vcov_`` is now sourced from the same + combined ``compute_robust_vcov(return_dof=True)`` call as ``_bm_dof``. It + must match an independently-computed CR2 vcov, and each coefficient's SE + must equal its sqrt-diagonal. Pins the vcov VALUE directly (the other LR + tests only pin ``_bm_dof`` end-to-end).""" + from diff_diff.linalg import LinearRegression, compute_robust_vcov + + rng = np.random.default_rng(4751) + n = 36 + X = np.column_stack([np.ones(n), rng.normal(size=n), rng.normal(size=n)]) + y = X @ np.array([1.0, 0.5, -0.3]) + rng.normal(0, 0.5, n) + w = rng.uniform(0.5, 2.0, n) + cl = np.arange(n) % 6 + + for kw in ( + {"weights": w, "weight_type": "pweight", "cluster_ids": cl}, + {"weights": w, "weight_type": "pweight"}, + {}, # unweighted one-way + ): + if "weights" in kw: + coef = np.linalg.solve(X.T @ (X * w[:, None]), X.T @ (w * y)) + else: + coef = np.linalg.solve(X.T @ X, X.T @ y) + resid = y - X @ coef + helper_vcov, _ = compute_robust_vcov( + X, + resid, + cluster_ids=kw.get("cluster_ids"), + weights=kw.get("weights"), + weight_type=kw.get("weight_type", "pweight"), + vcov_type="hc2_bm", + return_dof=True, + ) + + lr = LinearRegression(include_intercept=False, vcov_type="hc2_bm", **kw) + lr.fit(X, y) + np.testing.assert_allclose( + lr.vcov_, + helper_vcov, + atol=1e-9, + rtol=1e-9, + err_msg=f"LinearRegression.vcov_ must match the CR2 helper vcov (kw={kw})", + ) + for i in range(X.shape[1]): + inf_i = lr.get_inference(index=i) + np.testing.assert_allclose( + inf_i.se, + np.sqrt(lr.vcov_[i, i]), + rtol=1e-12, + err_msg=f"SE must be sqrt(vcov_[{i},{i}]) (kw={kw})", + ) + + def test_lr_hc2_bm_cr2_sandwich_computed_once(self): + """#475 dedup mechanism proof: the CR2 sandwich + (``_compute_robust_vcov_numpy``) is computed ONCE per hc2_bm fit, not + twice (previously once inside ``solve_ols`` for the vcov + once via + ``compute_robust_vcov`` for the dof). Patching ``_compute_robust_vcov_numpy`` + covers BOTH the ``_compute_cr2_bm`` (weighted / clustered) and the + unweighted-one-way ``_compute_bm_dof_oneway`` sub-paths.""" + import diff_diff.linalg as _linalg + from diff_diff.linalg import LinearRegression + + rng = np.random.default_rng(4752) + n = 36 + X = np.column_stack([np.ones(n), rng.normal(size=n), rng.normal(size=n)]) + y = X @ np.array([1.0, 0.5, -0.3]) + rng.normal(0, 0.5, n) + w = rng.uniform(0.5, 2.0, n) + cl = np.arange(n) % 6 + + orig = _linalg._compute_robust_vcov_numpy + for kw in ( + {"weights": w, "weight_type": "pweight"}, # weighted one-way + {"weights": w, "weight_type": "pweight", "cluster_ids": cl}, # weighted cluster + {}, # unweighted one-way (-> _compute_bm_dof_oneway) + {"cluster_ids": cl}, # unweighted cluster + ): + calls = {"n": 0} + + def _counting(*a, _orig=orig, _calls=calls, **k): + _calls["n"] += 1 + return _orig(*a, **k) + + _linalg._compute_robust_vcov_numpy = _counting + try: + lr = LinearRegression(include_intercept=False, vcov_type="hc2_bm", **kw) + lr.fit(X, y) + finally: + _linalg._compute_robust_vcov_numpy = orig + assert lr._bm_dof is not None, f"_bm_dof must be set (kw={kw})" + assert calls["n"] == 1, ( + f"CR2 sandwich must run once, not {calls['n']}x (kw={kw}); the #475 " + "dedup collapses solve_ols's vcov + the separate dof recompute." + ) + class TestWLSCR2FEDoFNoiseGuard: """Regression for R6 P1: weighted CR2-BM Satterthwaite DOF for high-