perf(linalg,mpd): compute MPD cluster+hc2_bm CR2 precomputes once

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`

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