perf(linalg): compute hc2_bm CR2 sandwich once, not twice (#475)

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`

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_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:

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-