From 1e82d8e0770540710ffda29aed4ef84b32e389f4 Mon Sep 17 00:00:00 2001 From: igerber Date: Tue, 7 Jul 2026 18:12:13 -0400 Subject: [PATCH 1/2] perf(utils): precompute r-independent pieces of the wild-bootstrap test inversion (~7x, bit-identical outputs) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The WCR bootstrap DGP is linear in the candidate null r (y*(r) = A + r*B), so beta_j*(r) is affine and the per-cluster score variance is the PSD quadratic qa + r*qb + r^2*qc — five precomputable (B,) vectors. The CI inversion's ~O(100) _t_star(r) evaluations (each previously materializing fresh (B,n)/(k,B)/(n,B) arrays) become O(B) arithmetic; the one-time precompute pass is chunked over draws (~256MB cap) so peak memory is bounded for large n/B — resolving both options the TODO row offered at once. Verified vs origin/main on a 5-seed few-cluster grid (backend pinned): SE, p-value, and inverted CI endpoints all bit-identical; 6.8x end-to-end. boottest parity suite (60 tests, pinned values) unchanged. Also parks the CS RC pscore (g,t)-dedup TODO row with findings: the RC propensity is fit on the pooled two-period sample whose period-t blocks change with every t, so no two cells share a logit — caching would change the estimator away from DRDID::*_rc. Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 14 +++++++ TODO.md | 7 ++-- diff_diff/utils.py | 72 +++++++++++++++++++++++++++++++----- tests/test_wild_bootstrap.py | 35 ++++++++++++++++++ 4 files changed, 114 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 09b45f9c..a5095086 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -489,6 +489,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 h=1 clamp-parity test). The symbol is imported independently (mixed-version safe — a stale extension degrades HC2 to NumPy without disabling older Rust accelerations). `return_dof` / weighted / CR2-BM requests stay on NumPy (CR2-BM tracked in TODO). +- **Wild cluster bootstrap test inversion ~7x faster with bounded memory.** The WCR + bootstrap DGP is linear in the candidate null `r` (`y*(r) = A + r·B` with r-independent + `A`, `B`), so the per-cluster score decomposition and the studentizing variance reduce + to five precomputed `(B,)` vectors (the variance is the PSD quadratic + `qa + r·qb + r²·qc`). The CI inversion's ~O(100) `_t_star(r)` evaluations — each of + which previously materialized fresh `(B×n)` bootstrap-outcome, `(k×B)` refit and `(n×B)` + residual arrays — now cost O(B) arithmetic each, and the ONE precompute pass is chunked + over draws (a conservative per-chunk byte budget sized against the peak count of live + `(Bc, n)` temporaries) so peak memory is bounded for large `n`/`B`. Verified against + origin/main on a 5-seed few-cluster grid: SE, p-value, and inverted CI endpoints all + **bit-identical** (backend pinned), 6.8x end-to-end; the boottest parity suite + (`tests/test_wild_bootstrap.py`, 60 tests incl. pinned values) passes unchanged. The + quadratic-form evaluation is a ~1-ULP reassociation of the per-call `sum(scores²)`; + the strict-inequality tie guard absorbs sub-1e-9 shifts by design. - **`CallawaySantAnna` per-(g,t) IF scatters converted from `np.add.at` to fancy `+=`** (`staggered.py::_cluster_robust_se_from_per_gt_if` — runs once per (g,t) cell when `cluster=` is set — and the general combined-IF assembly path in diff --git a/TODO.md b/TODO.md index 97444378..4591b9e5 100644 --- a/TODO.md +++ b/TODO.md @@ -30,6 +30,7 @@ The `Origin` column (Actionable tables) and the `PR` column (Deferred tables) bo |-------|----------|--------|--------|----------| | `SyntheticControl` conformal (CWZ 2021) extensions: (a) one-sided / signed-`t` variants (§7); (b) covariates in the conformal proxy (`X_jt`, eqs 4/6 — current proxy is outcomes-only); (c) AR / innovation-permutation path (Lemmas 5-7) for time-series proxies. The joint test, pointwise CIs, and average-effect CI have landed. | `conformal.py`, `synthetic_control_results.py` | CWZ-2021 | Heavy | Low | | `ContinuousDiD` CGBS-2024 extensions. (a) `covariates=` kwarg — **DONE (reg/dr)**; (b) discrete-treatment saturated regression (`treatment_type="discrete"`) — **DONE**; (c) lowest-dose-as-control per Remark 3.1 when `P(D=0)=0` (`control_group="lowest_dose"`) — **DONE** (discrete + continuous mass-point, single-cohort; estimand `ATT(d)−ATT(d_L)`; see REGISTRY Note #7). Remaining (all deferred `NotImplementedError`, documented): `estimation_method="ipw"` on the dose curve (scalar-adjustment / degenerate); `covariates=` × `survey_design=` (weighted OR + weighted nuisance IF); multi-cohort **heterogeneous-support** discrete aggregation (support-aware: average each dose only over the cohorts that observe it); **multi-cohort `lowest_dose`** (within-cohort `d_L` reference + support-aware cross-cohort aggregation); and **`covariates=` × `lowest_dose`** (conditional-PT-relative-to-`d_L` estimand). Single-cohort / 2-period / shared-support multi-cohort are supported. | `continuous_did.py` | CGBS-2024 | Heavy | Low | +| **Rust-backend `solve_ols(return_vcov=True)` is run-to-run NONDETERMINISTIC** (~1e-14 rel): 8 identical clustered-vcov calls in one process produced 3 distinct values (observed 2026-07-07, Apple Silicon M4, `maturin --features accelerate`; the pure-Python backend is bit-stable run-to-run). Same-seed reproducibility is a documented library expectation, and per the Python-canonical convention the Rust port should at minimum be deterministic. Likely rayon/threaded-BLAS reduction ordering in the faer SVD solve or the vcov meat; investigate and either force a deterministic reduction path or document the wobble as a platform contract with an explicit note. Repro: `solve_ols(X(400x3), y, cluster_ids=10x40, return_vcov=True)` in a loop, compare `vcov[1,1]` bit patterns. | `rust/src/linalg.rs::solve_ols`, `diff_diff/linalg.py` | wild-bootstrap follow-up | Mid | Medium | | TWFE's HC2/HC2-BM inline full-dummy build (`twfe.py:280-315`) duplicates the dummy-construction logic in `DifferenceInDifferences(fixed_effects=...)` (`estimators.py:478-486`). Extract a shared helper, or delegate TWFE's HC2/HC2-BM path to DiD's `fixed_effects=` branch (with TWFE-specific cluster-default threading), to reduce drift risk on FE naming / survey behavior / result-surface conventions. Substantive refactor — touches both estimators. | `twfe.py::fit`, `estimators.py::DifferenceInDifferences.fit` | follow-up | Heavy | Low | | `HonestDiD` Δ^SD optimal-FLCI center parity — **LANDED (SE-audit B2b)**: replaced the flat Nelder-Mead affine-estimator optimizer with a faithful port of R `HonestDiD::findOptimalFLCI`'s nested convex program (inner min-worst-case-bias at fixed estimator SD `h` via SLSQP QCQP; outer grid-zoom over `h`). Matches R's center + half-length + optimalVec to ~1e-3 (median ~1e-5) across a stress grid; the prior ~9% intermediate-M center drift is removed (widths always matched, coverage unaffected). Analytical folded-normal cv is more accurate than R's MC `.qfoldednormal`. Golden `honest_flci_golden.json` + `TestHonestFLCIParityR`. | `honest_did.py::_flci_solve` | SE-audit | Done | Low | @@ -43,12 +44,10 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| | `EfficientDiD` conditional path: the largest remaining O(n) stage is the sieve/nuisance construction outside the tiled pass (~9s at 10k). (The `_ridge_solve_weights` Python-prep shave landed 2026-07-07 — the `omega_stack[rest]` fancy-index copy and tail scatter are skipped when no row is zero-masked, byte-identical outputs; the `zero_mask` abs scan itself remains, needed for correctness.) | `efficient_did_covariates.py` | CS-scaling | Mid | Low | -| `CallawaySantAnna` repeated-cross-section ipw/dr paths (`_ipw_estimation_rc` / `_doubly_robust_rc`) re-solve the propensity logit for EVERY (g,t) cell with no cache — the panel paths dedup via `pscore_cache` keyed `(g, base_period[, t])`. The Phase 3 IRLS fast path already speeds each RC solve; the missing (g,t)-dedup cache is the remaining lever. Unprofiled surface — measure an RC covariate scenario first. | `staggered.py::_ipw_estimation_rc`, `_doubly_robust_rc` | CS-scaling | Mid | Low | | Rust `solve_ols` always runs a full thin SVD (equilibrated, gelsd-parity); at 40 covariates it is the top per-cell solver item on a CS dr fit (1.04s of 3.1s at 2M rows, 95 calls). A Cholesky/QR fast path with SVD fallback — mirroring the Phase 3 Python-side pattern (certify well-conditioned, fall back verbatim) — is the natural lever, but `solve_ols` is the universal OLS entry point (every estimator), so the blast radius needs the full backend-parity treatment (`TestSolveOLSSkipRankCheckParity` posture: fitted-values parity, not beta). | `rust/src/linalg.rs::solve_ols` | CS-scaling | Heavy | Low | | `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 | | 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 CR2 Bell-McCaffrey: falls through to NumPy (the leverage/Satterthwaite-DOF path needs `return_dof` support, which the Rust vcov dispatch excludes). The one-way HC2 kernel landed 2026-07-07 (`compute_robust_vcov_hc2`, mirrors the NumPy hc2 branch at ~1e-15; near-singular hat-diagonal sentinel + Python-side warn-and-HC1-fallback). | `rust/src/linalg.rs` | Phase 1a | Mid | Low | -| Wild cluster bootstrap CI inversion calls `_t_star(r)` ~O(100) times, each materializing a fresh `(B×n)` `y_star` + `(k×B)` refit + `(n×B)` residual arrays. Acceptable for the few-cluster regime; for large-`n`/large-`B`, chunk `_t_star` over draws or precompute the `r`-independent cluster-level pieces (restricted residuals are linear in `r`). | `utils.py::wild_bootstrap_se._t_star` | #543 | Mid | Low | | `SpilloverDiD` sparse cKDTree path for the staggered nearest-treated-distance helper (mirrors the static helper's sparse branch). `_compute_nearest_treated_distance_staggered` always builds dense `(n_units, n_treated_by_onset)` matrices per cohort; add a sparse branch gated on `n > _CONLEY_SPARSE_N_THRESHOLD`. | `spillover.py` | Wave B | Mid | Low | | `HeterogeneousAdoptionDiD` Phase 3 Stute: Appendix-D vectorized form replaces the per-iteration OLS refit with a single precomputed `M = I - X(X'X)^{-1}X'` applied to `eps*eta` (~2× faster, functionally identical). Shipped the literal-refit form to match paper text. | `had_pretests.py::stute_test` | Phase 3 | Mid | Low | | Migrate `spillover._iterative_fe_subset` onto the shared `diff_diff.utils._iterative_fe_solve` (same Gauss-Seidel-on-codes recursion, specialized to a masked Butts subsample; ImputationDiD/TwoStageDiD already route through the shared helper — this takes the FE-solver copy count from 2 to 1). Preserve the SpilloverDiD positive-weight/NaN-FE REGISTRY contract and `_FE_ITER_MAX=100` budget (or align it to 10k with a REGISTRY note). | `spillover.py` | demean-modernization | Mid | Low | @@ -73,6 +72,8 @@ Not currently actionable. Retained for provenance + AI-review deviation-document | Issue | Location | PR | Priority | |-------|----------|----|----------| +| `SyntheticControl` fit-snapshot residency (`_SyntheticControlFitSnapshot`) — **investigated 2026-07-07, parked**: the snapshot ALIASES the fit's own working pivots (zero extra construction cost); the retained residency implements the documented freeze contract (post-fit mutation of estimator inputs must not change `in_space_placebo()` / `leave_one_out()` / conformal output on an already-returned results object, and `__getstate__` already excludes it from pickles). A compact array representation saves only pandas overhead (the float panel dominates); releasing residency needs new API surface (`release`/opt-out flag) or a freeze-contract change. Revisit on user demand for very large donor panels. | `synthetic_control.py`, `synthetic_control_results.py` | follow-up | Low | +| Stratified survey-PSU multiplier-weight draw-tiling — **investigated 2026-07-07, parked**: the stratified generator (`generate_survey_multiplier_weights_batch`) consumes ONE sequential rng stream stratum-major (`rng.choice(size=(n_bootstrap, n_h))` per stratum, then lonely-PSU pooling), so draw-chunked assembly CANNOT reproduce the stream bit-identically (contra the old row's parenthetical) — it would need per-stratum generator state skipping (PCG64.advance + per-weight-type variate accounting; fragile) or a stream-layout change (MC-level SE changes → baseline/golden recapture + REGISTRY note). Stratified designs have few PSUs, so the full `(n_bootstrap × n_psu)` matrix rarely matters; unstratified (the large-`n_units` case) is already tiled. Revisit only if a large-PSU stratified design hits memory, as a documented stream change. | `diff_diff/bootstrap_chunking.py::iter_survey_multiplier_weight_blocks` | follow-up | Low | | CBWSDID covariate balancing (`StackedDiD(balance="entropy")`) v1 supports only balanced event windows + `weighting="aggregate"`; unbalanced/ragged panels fail closed (unit-count vs observation-count corrector convention unresolved off balanced panels). Matching-based balancing and the repeated `0→1`/`1→0` episode extension are also deferred. Documented in REGISTRY StackedDiD "Covariate balancing (CBWSDID)" Notes. | `stacked_did.py`, `balancing.py`, REGISTRY | follow-up | Low | | dCDH: Phase-1 per-period placebo `DID_M^pl` has NaN SE (no IF derivation for the per-period aggregation path). Multi-horizon placebos (`L_max ≥ 1`) have valid SE. | `chaisemartin_dhaultfoeuille.py` | #294 | Low | | dCDH: survey cell-period allocator's post-period attribution is a library convention, not derived from the observation-level survey linearization. MC coverage is empirically close to nominal; a formal derivation (or covariance-aware two-cell alternative) is deferred. Documented in REGISTRY survey IF expansion Note. | `chaisemartin_dhaultfoeuille.py`, REGISTRY | #408 | Medium | @@ -120,8 +121,6 @@ Doable in principle, but no current caller and/or explicitly out of paper scope. | Issue | Location | PR | Priority | |-------|----------|----|----------| -| `SyntheticControl` fit-snapshot residency (`_SyntheticControlFitSnapshot`) — **investigated 2026-07-07, parked**: the snapshot ALIASES the fit's own working pivots (zero extra construction cost); the retained residency implements the documented freeze contract (post-fit mutation of estimator inputs must not change `in_space_placebo()` / `leave_one_out()` / conformal output on an already-returned results object, and `__getstate__` already excludes it from pickles). A compact array representation saves only pandas overhead (the float panel dominates); releasing residency needs new API surface (`release`/opt-out flag) or a freeze-contract change. Revisit on user demand for very large donor panels. | `synthetic_control.py`, `synthetic_control_results.py` | follow-up | Low | -| Stratified survey-PSU multiplier-weight draw-tiling — **investigated 2026-07-07, parked**: the stratified generator (`generate_survey_multiplier_weights_batch`) consumes ONE sequential rng stream stratum-major (`rng.choice(size=(n_bootstrap, n_h))` per stratum, then lonely-PSU pooling), so draw-chunked assembly CANNOT reproduce the stream bit-identically (contra the old row's parenthetical) — it would need per-stratum generator state skipping (PCG64.advance + per-weight-type variate accounting; fragile) or a stream-layout change (MC-level SE changes → baseline/golden recapture + REGISTRY note). Stratified designs have few PSUs, so the full `(n_bootstrap × n_psu)` matrix rarely matters; unstratified (the large-`n_units` case) is already tiled. Revisit only if a large-PSU stratified design hits memory, as a documented stream change. | `diff_diff/bootstrap_chunking.py::iter_survey_multiplier_weight_blocks` | follow-up | Low | | CallawaySantAnna **unbalanced-panel R parity — LANDED** via `allow_unbalanced_panel=True` (matches R `did::att_gt(allow_unbalanced_panel=TRUE)` / `DRDID::reg_did_rc`: ATT bit-exact on cells AND dynamic aggregation via fixed unit-cohort-mass `pg` + a per-unit WIF; SE up to the documented CR1 `sqrt(G/(G-1))` factor). The earlier "weighting" framing was a mis-diagnosis — on unbalanced panels the dominant divergence from R is the *estimator* (within-cell differencing vs RC-on-pooled-obs), not only the weighting; both are resolved by the flag. The DEFAULT path keeps within-cell differencing as a documented design choice and now emits a `UserWarning` on unbalanced input (no-silent-failures). **Remaining deferred:** `survey_design=` × `allow_unbalanced_panel=` (per-obs vs per-unit weight resolution — currently fail-closed `NotImplementedError`); and covariate / ipw / dr × the flag R-parity verification (the RC path supports them; the committed golden covers `reg` no-cov). | `staggered.py`, `staggered_aggregation.py` | SE-audit D3 | Low | | CallawaySantAnna event-study bucket/weight construction is duplicated between the analytical aggregator (`staggered_aggregation.py::_aggregate_event_study`) and the multiplier bootstrap (`staggered_bootstrap.py`): both group (g,t) by `e = t - g`, apply the finite/NaN/reference masks, and read cohort weights. Both already consume the same source-materialized universal reference cells (so they agree), but the bucket logic is copy-pasted. Extract one shared helper returning per-event-time buckets (finite cells, NaN cells, reference flags, cohort weights, combined-IF inputs) used by both. Pure refactor; gate on byte-identical analytical + bootstrap output. | `staggered_aggregation.py`, `staggered_bootstrap.py` | SE-audit D3 | Low | | `StackedDiD` survey re-resolution intra-file dedup (raw-weight extraction ×3, compose-normalize ×3, resolve-on-stacked ×2). The cross-estimator ContinuousDiD/EfficientDiD panel-to-unit collapse consolidation LANDED (#226 shared helpers `ResolvedSurveyDesign.subset_to_units_by_row_idx` / `build_unit_first_row_index`); StackedDiD is deliberately NOT on that path (control units are duplicated across sub-experiments, so it re-resolves at stacked granularity rather than collapsing to one row per unit). The residual is stacked-specific, low value, and touches the numerically-sensitive composed-weight renormalization. Post-filter re-resolution / metadata-recompute unification across the three estimators was assessed and is not warranted — they use genuinely different mechanisms and already delegate to shared `_resolve_survey_for_fit` / `compute_survey_metadata`. | `stacked_did.py` | #226 | Low | diff --git a/diff_diff/utils.py b/diff_diff/utils.py index 4973c068..52406411 100644 --- a/diff_diff/utils.py +++ b/diff_diff/utils.py @@ -636,6 +636,12 @@ def _wild_weight_matrix( return weights +# Per-chunk byte budget for wild_bootstrap_se's r-independent precompute +# pass (divided by a conservative x8 peak-temporary multiplier to size the +# draw-rows per chunk). Read at CALL time so tests can monkeypatch it. +_WILD_PRECOMPUTE_CHUNK_BYTES = 256 * 1024 * 1024 + + def wild_bootstrap_se( X: np.ndarray, y: np.ndarray, @@ -868,20 +874,66 @@ def _degenerate() -> WildBootstrapResults: # that p(r) is a stable (monotone, step) function amenable to root-finding. weights = _wild_weight_matrix(n_clusters, n_bootstrap, weight_type, rng) n_boot_eff = int(weights.shape[0]) - weights_obs = weights[:, cl_idx] # (B, n) + + # ---- r-independent precompute for the test inversion -------------------- + # The WCR DGP is LINEAR in the candidate null r: + # u_r = m_y - r * m_xj + # y*(r) = y - u_r * (1 - w_obs) = A + r * B, + # with A = y - m_y*(1 - w_obs) and B = m_xj*(1 - w_obs) (both (B, n) and + # r-independent). Everything downstream is affine in y*, so + # beta_j*(r) = alpha1 + r * beta1 ((B,) each) + # scores(r) = SA + r * SB ((G, B)) + # se*(r)^2 / corr = qa + r * qb + r^2 * qc ((B,) each, + # the PSD quadratic form sum_g (SA + r SB)^2). The CI inversion calls + # _t_star ~O(100) times; precomputing these five (B,) vectors turns each + # call from three (B x n)/(n x B) GEMM passes into O(B) arithmetic, and + # the ONE precompute pass is chunked over draws so peak memory stays + # bounded (the old code materialized fresh (B, n) intermediates on every + # call). Not bit-identical to the per-call form (the variance is evaluated + # through the expanded quadratic instead of sum(scores^2) at each r — + # ~1 ULP reassociation class); the strict-inequality tie guard below + # absorbs sub-1e-9 shifts, so the discrete outputs (p-value, CI + # endpoints) are unchanged on the regression grid. + # The cluster-level weights matrix is only (B, G); the (B, n) + # observation-level expansion is built PER CHUNK below so no full + # (B, n) array is ever materialized. + alpha1 = np.empty(n_boot_eff, dtype=np.float64) + beta1 = np.empty(n_boot_eff, dtype=np.float64) + qa = np.empty(n_boot_eff, dtype=np.float64) + qb = np.empty(n_boot_eff, dtype=np.float64) + qc = np.empty(n_boot_eff, dtype=np.float64) + # Conservative sizing: up to ~8 (Bc, n) float64 arrays/temporaries can + # be live around the residual + score construction (one_minus_w_blk, + # A_blk, B_blk, RA_blk, RB_blk, plus elementwise/GEMM temporaries), so + # the divisor uses that peak multiplier against the byte cap below + # (module constant, read at call time so tests can monkeypatch it to + # force the multi-chunk branch). + _rows_per_chunk = max(1, int(_WILD_PRECOMPUTE_CHUNK_BYTES // (8 * 8 * max(n, 1)))) + for _cs in range(0, n_boot_eff, _rows_per_chunk): + blk = slice(_cs, min(_cs + _rows_per_chunk, n_boot_eff)) + one_minus_w_blk = 1.0 - weights[blk][:, cl_idx] # (Bc, n) + A_blk = y[None, :] - m_y[None, :] * one_minus_w_blk # (Bc, n) + B_blk = m_xj[None, :] * one_minus_w_blk # (Bc, n) + alpha1[blk] = A_blk @ a_vec + beta1[blk] = B_blk @ a_vec + # Residual-maker applied to each draw row: R = Z - (Z @ proj.T) @ X_eff.T + RA_blk = A_blk - (A_blk @ proj.T) @ X_eff.T # (Bc, n) + RB_blk = B_blk - (B_blk @ proj.T) @ X_eff.T + SA_blk = (a_vec[None, :] * RA_blk) @ cluster_indicator.T # (Bc, G) + SB_blk = (a_vec[None, :] * RB_blk) @ cluster_indicator.T + qa[blk] = np.sum(SA_blk * SA_blk, axis=1) + qb[blk] = 2.0 * np.sum(SA_blk * SB_blk, axis=1) + qc[blk] = np.sum(SB_blk * SB_blk, axis=1) def _t_star(r: float) -> np.ndarray: """Studentized bootstrap statistics t*(r) under H0: beta_j = r.""" - u_r = m_y - r * m_xj # restricted residuals at r (n,) - # WCR DGP: y* = fitted_restricted + u_r * w = (y - u_r) + u_r * w_obs. - y_star = y[None, :] - u_r[None, :] * (1.0 - weights_obs) # (B, n) - beta_j_star = y_star @ a_vec # (B,) - coef_full = proj @ y_star.T # (k_eff, B) - resid_star = y_star.T - X_eff @ coef_full # (n, B) bootstrap residuals - scores = cluster_indicator @ (a_vec[:, None] * resid_star) # (G, B) per-cluster scores - se_star = np.sqrt(corr * np.sum(scores**2, axis=0)) # (B,) + # PSD quadratic form; roundoff can dip microscopically negative at an + # interior root — clamp so sqrt is defined, then the se>0 guard NaNs + # the degenerate draws exactly as the per-call form did. + var_star = corr * np.maximum(qa + r * qb + r * r * qc, 0.0) + se_star = np.sqrt(var_star) # (B,) with np.errstate(divide="ignore", invalid="ignore"): - t = (beta_j_star - r) / se_star + t = (alpha1 + r * beta1 - r) / se_star t[~(se_star > 0)] = np.nan return t diff --git a/tests/test_wild_bootstrap.py b/tests/test_wild_bootstrap.py index ba8499a2..12d2650b 100644 --- a/tests/test_wild_bootstrap.py +++ b/tests/test_wild_bootstrap.py @@ -1480,3 +1480,38 @@ def test_wild_bootstrap_rank_deficient_storage_vcov_does_not_crash(): assert res.vcov is not None assert np.any(np.isnan(res.vcov)) assert not np.any(np.isinf(res.vcov)) + + +class TestPrecomputeChunking: + """The r-independent precompute pass must be chunk-count invariant: forcing + many draw-chunks (tiny byte budget) reproduces the single-chunk outputs + bit-for-bit (each chunk computes its own rows independently).""" + + def test_multi_chunk_bit_identical(self, monkeypatch): + import diff_diff.utils as du + + rng = np.random.default_rng(3) + G, n_per = 9, 30 + n = G * n_per + cl = np.repeat(np.arange(G), n_per) + x = rng.normal(size=n) + d = (np.arange(n) % 2).astype(float) + X = np.column_stack([np.ones(n), d, x]) + y = 1.0 + 0.4 * d + 0.2 * x + rng.normal(size=n) + rng.normal(size=G)[cl] + resid = y - X @ np.linalg.lstsq(X, y, rcond=None)[0] + + one = du.wild_bootstrap_se(X, y, resid, cl, 1, n_bootstrap=499, seed=11) + # ~3 rows per chunk -> hundreds of chunks. + monkeypatch.setattr(du, "_WILD_PRECOMPUTE_CHUNK_BYTES", 3 * 8 * 8 * n) + many = du.wild_bootstrap_se(X, y, resid, cl, 1, n_bootstrap=499, seed=11) + + # p-value: strict-count statistic with a 1e-9 relative tie guard — + # exact equality expected. se / CI endpoints tolerate the ambient + # Rust-backend run-to-run vcov wobble (~1e-14 rel; see the TODO row + # on rust solve_ols nondeterminism — the pure-Python backend is + # bit-stable and the chunked precompute itself is deterministic + # numpy). + assert many.p_value == one.p_value + np.testing.assert_allclose(many.se, one.se, rtol=1e-12) + np.testing.assert_allclose(many.ci_lower, one.ci_lower, rtol=1e-6) + np.testing.assert_allclose(many.ci_upper, one.ci_upper, rtol=1e-6) From e7a866b2a04233b2b0eeb7bd2bcbe7c319e75399 Mon Sep 17 00:00:00 2001 From: igerber Date: Wed, 8 Jul 2026 07:13:57 -0400 Subject: [PATCH 2/2] docs+test(wild): stale test-count + bit-for-bit wording (review P3s) Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 3 ++- tests/test_wild_bootstrap.py | 5 +++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a5095086..6f45dbf4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -500,7 +500,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `(Bc, n)` temporaries) so peak memory is bounded for large `n`/`B`. Verified against origin/main on a 5-seed few-cluster grid: SE, p-value, and inverted CI endpoints all **bit-identical** (backend pinned), 6.8x end-to-end; the boottest parity suite - (`tests/test_wild_bootstrap.py`, 60 tests incl. pinned values) passes unchanged. The + (`tests/test_wild_bootstrap.py`, incl. pinned values) passes unchanged, plus a new + chunk-count-invariance regression class (`TestPrecomputeChunking`). The quadratic-form evaluation is a ~1-ULP reassociation of the per-call `sum(scores²)`; the strict-inequality tie guard absorbs sub-1e-9 shifts by design. - **`CallawaySantAnna` per-(g,t) IF scatters converted from `np.add.at` to fancy `+=`** diff --git a/tests/test_wild_bootstrap.py b/tests/test_wild_bootstrap.py index 12d2650b..ba5851cd 100644 --- a/tests/test_wild_bootstrap.py +++ b/tests/test_wild_bootstrap.py @@ -1485,9 +1485,10 @@ def test_wild_bootstrap_rank_deficient_storage_vcov_does_not_crash(): class TestPrecomputeChunking: """The r-independent precompute pass must be chunk-count invariant: forcing many draw-chunks (tiny byte budget) reproduces the single-chunk outputs - bit-for-bit (each chunk computes its own rows independently).""" + (each chunk computes its own rows independently, so the p-value is exact; + se/CI carry documented ambient-backend tolerances, see in-test notes).""" - def test_multi_chunk_bit_identical(self, monkeypatch): + def test_multi_chunk_count_invariant(self, monkeypatch): import diff_diff.utils as du rng = np.random.default_rng(3)