Skip to content

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

Merged
igerber merged 1 commit into
mainfrom
perf/mpd-cluster-cr2-compute-once
Jun 28, 2026
Merged

perf(linalg,mpd): compute MPD cluster+hc2_bm CR2 precomputes once#566
igerber merged 1 commit into
mainfrom
perf/mpd-cluster-cr2-compute-once

Conversation

@igerber

@igerber igerber commented Jun 28, 2026

Copy link
Copy Markdown
Owner

Summary

  • Extract a shared _compute_cr2_bm_vcov_and_dof core in linalg.py; make _compute_cr2_bm and _compute_cr2_bm_contrast_dof thin bit-identical wrappers over it. Removes the ~50-line precompute block that was duplicated between the two functions, so every CR2 caller (_compute_robust_vcov_numpy, StackedDiD, SunAbraham, WooldridgeDiD, the weighted singleton dispatch, MPD) routes through one implementation.
  • MultiPeriodDiD(cluster=..., vcov_type="hc2_bm") (non-survey, unweighted) previously built the expensive CR2 Bell-McCaffrey 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 + post-period-average ATT Satterthwaite DOF. The fit now bypasses solve_ols's vcov on this path and computes vcov and the combined DOF from a single precompute build, then expands the reduced vcov with NaN for dropped columns.
  • The shared filter guards the residuals subscript with residuals is not None, so DOF-only callers passing weights with zero rows do not crash. The survey_weights is None clause keeps the bypass byte-identical to solve_ols (any hypothetical weighted path falls back to the prior two-call behavior).
  • Removes the resolved MPD cluster+hc2_bm Performance row from TODO.md; completes the follow-up flagged in the perf(linalg): compute hc2_bm CR2 sandwich once, not twice (#475) #565 LinearRegression CHANGELOG entry.

Methodology references (required if estimator / math changes)

  • Method name(s): CR2 Bell-McCaffrey cluster-robust variance + Satterthwaite contrast DOF (Pustejovsky-Tipton 2018 / Imbens-Kolesar 2016; clubSandwich WLS-CR2 port).
  • Paper / source link(s): docs/methodology/REGISTRY.md → "Variance Estimation, Cluster-Robust SE"; clubSandwich::vcovCR(type="CR2") + coef_test(test="Satterthwaite").
  • Any intentional deviations from the source (and why): None. Pure performance dedup — the CR2 sandwich and Satterthwaite DOF formulas are unchanged, so no REGISTRY edit is required.

Validation

  • Tests added/updated: tests/test_estimators_vcov_type.py (new TestMPDClusterHC2BMSharedPrecompute: asserts the per-cluster adjustment matrix is built exactly G times — one precompute build, down from 2·G — on balanced and unbalanced designs, plus fit determinism); tests/test_linalg_hc2_bm.py (wrapper/core atol=0 equivalence for the per-coef and DOF-only paths, unweighted + weighted; zero-weight + residuals=None DOF-only no-crash regression).
  • Backtest / simulation / notebook evidence (if applicable): Bit-identical proven at atol=0 for vcov_, per-period DOF, avg-ATT DOF, p-values, and CIs across balanced, unbalanced, and rank-deficient designs (60 fields × 3 scenarios). The existing R/clubSandwich goldens (test_multi_period_cluster_hc2_bm_avg_att_uses_clubsandwich_dof, tests/test_linalg_hc2_bm.py) remain the independent absolute-value anchor and pass. Targeted suites green (test_estimators_vcov_type, test_linalg_hc2_bm, test_methodology_wls_cr2, test_estimators, plus the StackedDiD/SunAbraham/Wooldridge callers); black/ruff clean; mypy 0 new errors vs baseline.

Security / privacy

  • Confirm no secrets/PII in this PR: Yes

🤖 Generated with Claude Code

@github-actions

Copy link
Copy Markdown

Overall Assessment: ✅ Looks good

Executive Summary

  • The affected method is CR2 Bell-McCaffrey cluster-robust variance with Satterthwaite DOF, especially MultiPeriodDiD(cluster=..., vcov_type="hc2_bm").
  • The linalg refactor preserves the registered unweighted and weighted formulas; the shared core keeps the same CR2 adjustment, sandwich meat, and contrast DOF logic.
  • The MPD bypass correctly computes reduced vcov + per-coef/avg-ATT DOF once, then expands dropped columns with NaN.
  • No P0/P1 methodology, inference, NaN-guard, or parameter propagation defects found.
  • One P3 documentation drift: the methodology registry still says MPD cluster+hc2_bm recomputes CR2 setup twice.

Methodology

Finding 1 — P3 informational: stale registry performance note

  • Location: docs/methodology/REGISTRY.md:L236-L242
  • Impact: The registry still says clustered hc2_bm MPD fits “pay the O(n²) CR2 setup twice” and points to a TODO follow-up, but this PR removes that TODO row and implements the shared _compute_cr2_bm_vcov_and_dof path.
  • Concrete fix: Update or remove that note so the registry reflects that vcov and contrast DOF now share one CR2 precompute build.

No methodology defect found in the implementation. The formulas match the registry/docstrings: unweighted uses the PT-style trace formula, weighted routes through the documented clubSandwich WLS-CR2 P-array form.

Code Quality

No findings. The new shared core is a reasonable consolidation, and wrappers remain thin.

Performance

No findings. The MPD cluster+hc2_bm bypass targets the documented duplicate CR2 precompute cost and adds a mechanism test asserting one adjustment-matrix build per cluster.

Maintainability

No findings. Centralizing CR2 precomputes reduces duplicate variance/DOF logic.

Tech Debt

No findings. The removed TODO row matches the implemented optimization; only the registry note needs cleanup.

Security

No findings. No secrets or security-sensitive changes identified in the changed files.

Documentation/Tests

Finding 1 — P3 informational: registry/changelog mismatch

  • Location: docs/methodology/REGISTRY.md:L236-L242, TODO.md:L56-L60
  • Impact: Documentation now disagrees with the PR behavior and the deleted TODO row.
  • Concrete fix: Remove the stale “currently recomputes” note or replace it with a short note that _compute_cr2_bm_vcov_and_dof shares the precompute build.

Targeted tests could not be run in this environment because pytest is not installed. Static review of the added tests shows coverage for wrapper/core equivalence, zero-weight DOF-only handling, balanced/unbalanced MPD precompute count, and fit determinism.

@igerber igerber force-pushed the perf/mpd-cluster-cr2-compute-once branch from ac5d27d to 2e20408 Compare June 28, 2026 12:07
@github-actions

Copy link
Copy Markdown

🔁 AI review rerun (requested by @igerber)

Head SHA: 2e2040859dfcfa61b90b24a03558fb86c5039c27


Overall Assessment: ✅ Looks good

Executive Summary

  • Prior stale-registry finding is resolved: REGISTRY.md now documents the shared _compute_cr2_bm_vcov_and_dof path.
  • Affected method: CR2 Bell-McCaffrey cluster-robust variance with Satterthwaite contrast DOF for MultiPeriodDiD(cluster=..., vcov_type="hc2_bm").
  • No P0/P1 methodology, variance/SE, NaN-propagation, or inference anti-pattern defects found.
  • Tests added cover shared-core wrapper equivalence, zero-weight DOF-only handling, MPD single-precompute behavior, and determinism.
  • Targeted tests could not be executed here because pytest is not installed.

Methodology

No findings.

The shared core preserves the registered CR2 sandwich and Satterthwaite contrast-DOF formulas, and the registry now explicitly documents the one-precompute MPD path at docs/methodology/REGISTRY.md:L236-L244.

Code Quality

Finding 1 — P3 informational: dormant weighted fallback omits weights=survey_weights

  • Location: diff_diff/estimators.py:L2039-L2049
  • Impact: The fallback comment says it preserves the prior hypothetical weighted cluster hc2_bm behavior, but the new _compute_cr2_bm_contrast_dof(...) call no longer forwards weights=survey_weights. Current public MPD paths appear unaffected because any survey_weights come from survey_design, and ResolvedSurveyDesign.needs_survey_vcov routes those fits out before this block.
  • Concrete fix: Either pass weights=survey_weights in that fallback or remove the weighted-preservation wording if the branch is intentionally unreachable.

Performance

No findings. The mechanism test at tests/test_estimators_vcov_type.py:L2284-L2316 directly guards the intended reduction to one _cr2_adjustment_matrix build per cluster.

Maintainability

No findings beyond the P3 fallback note above. Centralizing CR2 precomputes reduces duplicated math paths.

Tech Debt

No findings. The removed TODO row matches the completed MPD cluster+hc2_bm precompute optimization.

Security

No findings. No secrets or security-sensitive changes identified.

Documentation/Tests

No findings.

Static review of the added tests shows relevant coverage in tests/test_linalg_hc2_bm.py:L774-L840 and tests/test_estimators_vcov_type.py:L2244-L2343. I could not run them locally: /bin/bash: pytest: command not found.

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) <noreply@anthropic.com>
@igerber igerber force-pushed the perf/mpd-cluster-cr2-compute-once branch from 2e20408 to 974455d Compare June 28, 2026 12:14
@github-actions

Copy link
Copy Markdown

🔁 AI review rerun (requested by @igerber)

Head SHA: 974455d2b7f7c40ddfc91b6fcb6076863efe4a92


Overall Assessment: ✅ Looks good

Executive Summary

  • Affected method: CR2 Bell-McCaffrey cluster-robust variance with Satterthwaite contrast DOF, especially MultiPeriodDiD(cluster=..., vcov_type="hc2_bm").
  • No P0/P1 methodology, variance/SE, NaN-propagation, or inference anti-pattern issues found.
  • The registry now documents the shared _compute_cr2_bm_vcov_and_dof path and MPD one-precompute behavior at docs/methodology/REGISTRY.md:L226-L244.
  • Prior weighted-fallback comment concern is addressed: the current branch is explicitly documented as unreachable/defensive and preserves prior behavior at diff_diff/estimators.py:L2039-L2056.
  • Could not run tests locally because pytest and even numpy are not installed in this environment.

Methodology

No findings.

The shared core preserves the registered CR2 sandwich and Satterthwaite contrast-DOF formulas: per-cluster A_g, S_W, G_g, score meat, and contrast DOF are centralized in diff_diff/linalg.py:L1712-L1893. The MPD bypass computes vcov and per-coef/avg-ATT DOF from the reduced identified design, then expands with NaNs for dropped columns at diff_diff/estimators.py:L1969-L2038.

Code Quality

No findings.

Performance

No findings.

The new MPD mechanism test asserts one _cr2_adjustment_matrix build per cluster at tests/test_estimators_vcov_type.py:L2284-L2316, directly covering the intended dedup.

Maintainability

No findings.

Centralizing _compute_cr2_bm and _compute_cr2_bm_contrast_dof through _compute_cr2_bm_vcov_and_dof reduces duplicated methodology code.

Tech Debt

No findings.

The removed TODO row matches the completed MPD cluster+hc2_bm precompute optimization.

Security

No findings.

Documentation/Tests

No blocking findings.

Added tests cover wrapper/core equivalence, weighted zero-row DOF-only handling, MPD single-precompute behavior, and deterministic repeated fits at tests/test_linalg_hc2_bm.py:L774-L840 and tests/test_estimators_vcov_type.py:L2244-L2343.

Validation not executed here: /usr/bin/python: No module named pytest; a direct smoke run also failed because numpy is not installed.

@igerber igerber added the ready-for-ci Triggers CI test workflows label Jun 28, 2026
@igerber igerber merged commit 6892498 into main Jun 28, 2026
33 of 34 checks passed
@igerber igerber deleted the perf/mpd-cluster-cr2-compute-once branch June 28, 2026 14:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ready-for-ci Triggers CI test workflows

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant