diff --git a/CHANGELOG.md b/CHANGELOG.md index 2f7c8dab..89ba7f13 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -454,6 +454,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 memory-traffic cleanup, not a headline speedup). The `zero_mask` abs scan is retained (correctness); the remaining conditional-path lever is the sieve/nuisance stage (see TODO). +- **`stute_test` / `stute_joint_pretest` unweighted bootstrap ~2x faster, bit-identical.** + The Appendix-D per-replicate OLS refit loops now hoist their loop invariants: Mammen + multiplier draws are batched through memory-bounded `rng.choice` calls (numpy fills + C-order, so the variate stream — and therefore every draw — is identical to the prior + per-iteration draws), and the d-moments + tie-safe CvM tie-block indices are precomputed + once per test instead of once per replicate. Each replicate still applies the same 1-D + refit/CvM operations on the same values, so `bootstrap_S`, `cvm_stat`, and `p_value` are + **bit-identical** to the literal per-iteration form (locked by a frozen byte-copy parity + test + per-helper equality tests). Measured 2.5x/2.1x/1.9x at G=1e3/1e4/1e5 (B=999). + This resolves the Phase-3 TODO row's ~2x target without its sketched O(G²) + `M = I - X(X'X)^{-1}X'` materialization (which would not scale in memory) and with + exact — not just "functional" — identity. The paper-faithful per-replicate refit form + is unchanged; the large-G advisory now cites measured post-hoist timings. ## [3.6.2] - 2026-07-03 diff --git a/TODO.md b/TODO.md index e55b68c0..4cfcc8cd 100644 --- a/TODO.md +++ b/TODO.md @@ -51,7 +51,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | 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 | | 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 | +| Multiplier-bootstrap weight chunking (CallawaySantAnna, EfficientDiD, and HAD — all wired through `diff_diff/bootstrap_chunking.py`) covers the **unstratified** survey-PSU generation (the default unit-level bootstrap — `cluster=None`, equivalently `cluster="unit"` — the large-`n_units` OOM case). Remaining gap: the **stratified** survey-PSU generator (`generate_survey_multiplier_weights_batch`, per-stratum + lonely-PSU pooling + FPC) still materializes the full `(n_bootstrap × n_psu)` matrix (consumed via sliced blocks). Stratified designs have few PSUs so this rarely OOMs; tile per-stratum generation over draws (each stratum's draws are independent → contiguous draw-blocks reproduce the stream bit-identically) if a large-PSU stratified design hits memory. | `diff_diff/bootstrap_chunking.py::iter_survey_multiplier_weight_blocks` | follow-up | 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 | | Adopt `snap_absorbed_regressors` (FE-spanned regressor two-stage snap + LSMR confirmation) on the ImputationDiD lead-indicator path — lead columns are the most plausible FE-spanned regressors, and the truncated-MAP-iterate exposure documented at `utils.py` applies; today only `solve_ols` rank detection guards it. Behavior change beyond the demean-modernization refactor, so deferred from that PR. | `imputation.py::_compute_lead_coefficients` | demean-modernization | Mid | Low | | Per-cell `treated_units`/`control_units` label arrays (`all_units[positions]`, ~O(n_control) alloc per (g,t) cell) are consumed only by the precomputed-None fallback of the combined-IF assembly, which no in-package caller reaches — build them lazily (or drop from the IF-info dict) to cut per-cell allocation at high cell counts. | `staggered.py::_compute_att_gt_fast` | CS-scaling | Mid | Low | diff --git a/diff_diff/had_pretests.py b/diff_diff/had_pretests.py index f9ed37aa..414bcce1 100644 --- a/diff_diff/had_pretests.py +++ b/diff_diff/had_pretests.py @@ -968,16 +968,29 @@ def _validate_1d_numeric(arr: np.ndarray, name: str) -> np.ndarray: return a -def _fit_ols_intercept_slope(d: np.ndarray, dy: np.ndarray) -> "tuple[float, float, np.ndarray]": +def _fit_ols_intercept_slope( + d: np.ndarray, + dy: np.ndarray, + d_moments: "Optional[tuple[float, np.ndarray, float]]" = None, +) -> "tuple[float, float, np.ndarray]": """Fit ``dy = a + b*d + eps`` via closed-form OLS. Returns ``(a_hat, b_hat, residuals)`` where ``residuals`` has the same length as ``d`` in the ORIGINAL input order (not sorted). + + ``d_moments`` optionally supplies precomputed ``(d_mean, d_dev, + var_d)`` for callers that refit against the SAME ``d`` many times + (the Stute bootstrap loops). The values are deterministic functions + of ``d`` computed with the identical expressions below, so passing + them is bit-identical to recomputing per call. """ - d_mean = d.mean() + if d_moments is not None: + d_mean, d_dev, var_d = d_moments + else: + d_mean = d.mean() + d_dev = d - d_mean + var_d = np.dot(d_dev, d_dev) dy_mean = dy.mean() - d_dev = d - d_mean - var_d = np.dot(d_dev, d_dev) if var_d <= 0.0: # Degenerate case: all dose values equal. Slope undefined. # Caller is responsible for gating before we reach here; if we @@ -1066,7 +1079,36 @@ def _fit_weighted_ols_intercept_only( return a_hat, 0.0, residuals -def _cvm_statistic(eps_sorted: np.ndarray, d_sorted: np.ndarray) -> float: +def _iter_mammen_rows( + n_bootstrap: int, + G: int, + rng: np.random.Generator, + max_batch_bytes: int = 64 * 1024 * 1024, +): + """Yield per-replicate Mammen weight rows drawn in memory-bounded batches. + + One ``rng.choice`` call per batch: numpy fills the output in C order, + so consecutive batched draws consume the generator's variate stream + identically to ``n_bootstrap`` sequential size-``G`` draws — the + bootstrap law is unchanged draw-for-draw — while peak memory stays + bounded (a single full ``(B, G)`` batch would be ``B*G*8`` bytes, + ~800 MB at G=100k, B=999). + """ + rows_per_batch = max(1, int(max_batch_bytes // (8 * max(G, 1)))) + drawn = 0 + while drawn < n_bootstrap: + take = min(rows_per_batch, n_bootstrap - drawn) + batch = _generate_mammen_weights((take, G), rng) + for r in range(take): + yield batch[r] + drawn += take + + +def _cvm_statistic( + eps_sorted: np.ndarray, + d_sorted: np.ndarray, + tie_blocks: "Optional[tuple[np.ndarray, np.ndarray]]" = None, +) -> float: """Compute the tie-safe Cramer-von Mises cusum statistic. Paper definition (Appendix D): @@ -1093,6 +1135,12 @@ def _cvm_statistic(eps_sorted: np.ndarray, d_sorted: np.ndarray) -> float: d_sorted : np.ndarray, shape (G,) Regressor values sorted ascending. Must be sorted consistently with ``eps_sorted``. + tie_blocks : tuple of (np.ndarray, np.ndarray), optional + Precomputed ``(tie_end_idx, counts)`` for callers that evaluate + the statistic against the SAME ``d_sorted`` many times (the Stute + bootstrap loops). Deterministic functions of ``d_sorted`` computed + with the identical expressions below, so passing them is + bit-identical to recomputing per call. Returns ------- @@ -1106,8 +1154,11 @@ def _cvm_statistic(eps_sorted: np.ndarray, d_sorted: np.ndarray) -> float: # already-sorted regressor gives per-unique-value counts; the last # index of each tie block is `cumsum(counts) - 1`, and np.repeat # expands that back to per-observation. - _, counts = np.unique(d_sorted, return_counts=True) - tie_end_idx = np.cumsum(counts) - 1 + if tie_blocks is not None: + tie_end_idx, counts = tie_blocks + else: + _, counts = np.unique(d_sorted, return_counts=True) + tie_end_idx = np.cumsum(counts) - 1 cumsum_tie_safe = np.repeat(cumsum[tie_end_idx], counts) return float(np.sum(cumsum_tie_safe * cumsum_tie_safe) / (G * G)) @@ -1691,9 +1742,11 @@ def stute_test( ones(G)``, weighted helpers reduce bit-exactly to the unweighted versions but bootstrap p-values diverge by Monte-Carlo noise (different RNG consumption between batched ``generate_survey_multiplier_weights_batch`` - and per-iteration ``_generate_mammen_weights``); use the - distribution-equivalence reduction test (large B) for trivial-pweight - parity, NOT numerical equivalence. + and the unweighted path's ``_generate_mammen_weights`` stream — the + latter is drawn in memory-bounded batches that consume the stream + identically to per-iteration draws); use the distribution-equivalence + reduction test (large B) for trivial-pweight parity, NOT numerical + equivalence. References ---------- @@ -1813,10 +1866,10 @@ def stute_test( if G > _STUTE_LARGE_G_THRESHOLD: warnings.warn( f"stute_test: G = {G} exceeds {_STUTE_LARGE_G_THRESHOLD}; the " - f"per-iteration refit is O(G) per iteration so the " - f"{n_bootstrap}-replication loop may take tens of seconds or " - f"more. Consider yatchew_hr_test() instead (paper Theorem 7 " - f"recommends Yatchew-HR at large G).", + f"per-replicate refit is O(G), so the {n_bootstrap}-replication " + f"bootstrap cost grows linearly in G (measured ~1.5s at G=1e5, " + f"B=999; ~10x that at G=1e6). Consider yatchew_hr_test() " + f"instead (paper Theorem 7 recommends Yatchew-HR at large G).", UserWarning, stacklevel=2, ) @@ -1936,12 +1989,25 @@ def stute_test( fitted = a_hat + b_hat * d_arr # baseline fitted values under H_0 if w_arr is None: - # Unweighted bit-exact path - identical to pre-PR code. - for b in range(n_bootstrap): - eta = _generate_mammen_weights(G, rng) + # Unweighted path: the Appendix-D per-replicate OLS refit with + # hoisted loop invariants. The eta draws are batched through + # memory-bounded rng.choice calls (identical variate stream as + # per-iteration size-G draws — see _iter_mammen_rows), and the + # d-moments and CvM tie-block indices are precomputed once (they + # are deterministic functions of d, recomputed identically inside + # the helpers). Each replicate applies the same 1-D refit/CvM + # operations on the same values as the literal per-iteration + # form, so bootstrap_S is bit-identical to pre-hoist code + # (verified on a seed x DGP grid; ~2.3x at G in [1e3, 5e4]). + d_mean_h = d_arr.mean() + d_dev_h = d_arr - d_mean_h + d_moments = (d_mean_h, d_dev_h, float(np.dot(d_dev_h, d_dev_h))) + _, tie_counts = np.unique(d_sorted, return_counts=True) + tie_blocks = (np.cumsum(tie_counts) - 1, tie_counts) + for b, eta in enumerate(_iter_mammen_rows(n_bootstrap, G, rng)): dy_b = fitted + eps * eta - _, _, eps_b = _fit_ols_intercept_slope(d_arr, dy_b) - bootstrap_S[b] = _cvm_statistic(eps_b[idx], d_sorted) + _, _, eps_b = _fit_ols_intercept_slope(d_arr, dy_b, d_moments=d_moments) + bootstrap_S[b] = _cvm_statistic(eps_b[idx], d_sorted, tie_blocks=tie_blocks) else: # Phase 4.5 C survey-aware path: PSU-level Mammen multipliers # (broadcast to per-obs perturbation), weighted OLS refit, weighted @@ -3274,18 +3340,24 @@ def stute_joint_pretest( bootstrap_S = np.empty(n_bootstrap, dtype=np.float64) if w_arr is None: - # Unweighted bit-exact path (stability invariant #1). - for b in range(n_bootstrap): - # SHARED eta across horizons - preserves unit-level dependence - # in the vector-valued empirical process. Independent-per-horizon - # draws would overstate precision. - eta = _generate_mammen_weights(G, rng) + # Unweighted bit-exact path (stability invariant #1). Same hoisted + # invariants as stute_test's unweighted loop: the eta draws come + # from memory-bounded batched rng.choice calls (identical variate + # stream as per-iteration draws — see _iter_mammen_rows) and the + # CvM tie-block indices are precomputed once — bootstrap_S is + # bit-identical to the per-iteration form. + _, tie_counts = np.unique(d_sorted, return_counts=True) + tie_blocks = (np.cumsum(tie_counts) - 1, tie_counts) + # SHARED eta across horizons - preserves unit-level dependence + # in the vector-valued empirical process. Independent-per-horizon + # draws would overstate precision. + for b, eta in enumerate(_iter_mammen_rows(n_bootstrap, G, rng)): S_b = 0.0 for k in horizon_labels: dy_b = fitted_arrays[k] + residuals_arrays[k] * eta beta_b = XtX_inv_Xt @ dy_b eps_b = dy_b - X @ beta_b - S_b += _cvm_statistic(eps_b[idx], d_sorted) + S_b += _cvm_statistic(eps_b[idx], d_sorted, tie_blocks=tie_blocks) bootstrap_S[b] = S_b else: # Phase 4.5 C survey-aware path: PSU-level Mammen multipliers diff --git a/diff_diff/utils.py b/diff_diff/utils.py index 4973c068..b2c8b3dc 100644 --- a/diff_diff/utils.py +++ b/diff_diff/utils.py @@ -5,7 +5,7 @@ import os import warnings from dataclasses import dataclass, field -from typing import Any, Dict, Iterable, List, Optional, Tuple +from typing import Any, Dict, Iterable, List, Optional, Tuple, Union import numpy as np import pandas as pd @@ -559,7 +559,9 @@ def _generate_webb_weights(n_clusters: int, rng: np.random.Generator) -> np.ndar return np.asarray(rng.choice(values, size=n_clusters)) -def _generate_mammen_weights(n_clusters: int, rng: np.random.Generator) -> np.ndarray: +def _generate_mammen_weights( + n_clusters: "Union[int, Tuple[int, ...]]", rng: np.random.Generator +) -> np.ndarray: """ Generate Mammen's two-point distribution weights. @@ -571,8 +573,12 @@ def _generate_mammen_weights(n_clusters: int, rng: np.random.Generator) -> np.nd Parameters ---------- - n_clusters : int - Number of clusters. + n_clusters : int or tuple of int + Number of clusters, or an output shape. A batched draw of shape + ``(B, G)`` consumes the generator's variate stream identically to + ``B`` sequential draws of size ``G`` (``rng.choice`` fills output + in C order), so batching replicate draws preserves the bootstrap + law draw-for-draw. rng : np.random.Generator Random number generator. diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 01e72b64..761f6c0c 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -3428,12 +3428,12 @@ Shipped in `diff_diff/had_pretests.py` as `stute_joint_pretest()` (residuals-in - **Note (Phase 2b baseline convention):** All event-time horizons use a uniform `F-1` anchor: `ΔY_{g,t} = Y_{g,t} - Y_{g,F-1}` for every `t`. This is consistent with the paper's Garrett-et-al. application (Section 5.1: "outcome `Y_{g,t} - Y_{g,2001}`" where `F = 2002`), simplifies event-time indexing (`e = t - F` so `e = -1` is the anchor, skipped), and keeps the implementation symmetric for pre- and post-period horizons. The paper review text's asymmetric "`Y_{g,t} - Y_{g,1}` for pre" / "`Y_{g,t} - Y_{g,F-1}` for post" phrasing is covered by the uniform convention since both give the same placebo interpretation under parallel trends (the paper's own applications use the uniform anchor). - **Note (Phase 2b result class):** `aggregate="event_study"` returns a new `HeterogeneousAdoptionDiDEventStudyResults` dataclass (distinct from the single-period `HeterogeneousAdoptionDiDResults`) with per-horizon arrays (`event_times`, `att`, `se`, `t_stat`, `p_value`, `conf_int_low`, `conf_int_high`, `n_obs_per_horizon`) and shared metadata. `to_dataframe()` returns a tidy per-horizon DataFrame; `to_dict()` returns a dict with list-of-per-horizon fields. The static return-type annotation on `fit()` is `Union[HeterogeneousAdoptionDiDResults, HeterogeneousAdoptionDiDEventStudyResults]`, matching the runtime polymorphism on `aggregate`; callers should narrow via `isinstance` (or via the aggregate they passed) when reading aggregate-specific fields. - **Note (Phase 3 pretest workflow):** Phase 3 ships the Section 4 pre-test diagnostics in a new module `diff_diff/had_pretests.py` (separate from the 2,800-line `had.py` estimator module). The three tests — `qug_test`, `stute_test`, `yatchew_hr_test` — accept raw `(d, dy)` arrays and return their own result dataclasses (`QUGTestResults`, `StuteTestResults`, `YatchewTestResults`); the composite `did_had_pretest_workflow` dispatches on `aggregate` (`"overall"` requires a balanced two-period panel; `"event_study"` accepts a multi-period panel with >= 3 periods) and returns `HADPretestReport` with a priority-ordered verdict string. The `data`-only workflow signature (no `result=`-consuming variant) avoids coupling pretests to the `BiasCorrectedFit` internal state, which does not expose residuals. Below-minimum sample sizes emit `UserWarning` and return all-NaN inference (no `ValueError` raise) for library consistency with the `safe_inference` convention; input-shape violations (non-1D, NaN-containing, mismatched length) still raise. No multiple-testing adjustment (Bonferroni/Holm) is applied — the paper does not prescribe one. **Partial-workflow semantic (path-dependent):** `aggregate="overall"` (default; two-period panel) runs paper steps 1 (QUG) + 3 (linearity via Stute + Yatchew-HR) ONLY; step 2 (Assumption 7 pre-trends test via Equation 18) is NOT covered on this path and a fail-to-reject verdict appends an explicit Assumption 7 gap flag (`"QUG and linearity diagnostics fail-to-reject; Assumption 7 pre-trends test NOT run (paper step 2 deferred)"`). `aggregate="event_study"` (multi-period panel, requires >=3 periods + earlier pre-period for joint pretrends) closes step 2 via joint Stute pre-trends (Equation 18) over pre-period horizons AND joint Stute homogeneity over post-period horizons; the verdict on a fail-to-reject across QUG + joint pretrends + joint homogeneity reads "TWFE admissible under Section 4 assumptions" without the Assumption 7 caveat. The `HADPretestReport` docstring and structural fields (`pretrends_joint`, `homogeneity_joint`) reflect the aggregate-dependent coverage. - - **Note (Phase 3 Stute bootstrap):** The Stute CvM statistic is implemented via the simplified-form `S = (1/G²) · Σ cumsum²`, which is algebraically equivalent to the paper's stated `Σ(g/G)² · ((1/g) Σ eps_{(h)})²` form. Bootstrap follows the paper's Appendix D Algorithm literally: each iteration refits OLS on `dy_b = a_hat + b_hat · d + eps · eta` (null DGP multiplied by Mammen weights). The Appendix-D vectorized matrix form (precomputing `M = I - X(X'X)⁻¹X'` once and applying it to `eps · eta` in each iteration) is functionally identical and ~2× faster; it is deferred as a performance follow-up to keep the Phase 3 reviewer surface small. Bootstrap reproducibility uses `np.random.default_rng(seed)` (library convention, matching `wild_bootstrap_se`); `seed=42` in tests for bitwise stability. + - **Note (Phase 3 Stute bootstrap):** The Stute CvM statistic is implemented via the simplified-form `S = (1/G²) · Σ cumsum²`, which is algebraically equivalent to the paper's stated `Σ(g/G)² · ((1/g) Σ eps_{(h)})²` form. Bootstrap follows the paper's Appendix D Algorithm literally: each iteration refits OLS on `dy_b = a_hat + b_hat · d + eps · eta` (null DGP multiplied by Mammen weights). The ~2× performance follow-up landed 2026-07-07 as **loop-invariant hoisting**, not the once-sketched `M = I - X(X'X)⁻¹X'` materialization (O(G²) memory — would not scale): Mammen draws are batched through memory-bounded `rng.choice` calls whose variate stream is identical to per-iteration draws, and the d-moments + CvM tie-block indices are precomputed once per test — each replicate still runs the paper-literal O(G) refit on the same values, so `bootstrap_S`/`cvm_stat`/`p_value` are bit-identical to the per-iteration form (frozen byte-copy parity test at `tests/test_had_pretests.py::TestStuteBootstrapHoistedInvariants`). Bootstrap reproducibility uses `np.random.default_rng(seed)` (library convention, matching `wild_bootstrap_se`); `seed=42` in tests for bitwise stability. - **Note (Phase 3 Yatchew normalizer):** `σ̂²_diff` in `yatchew_hr_test` divides by `2G` (paper-literal from Theorem 7), NOT by `2(G-1)`. A unit test hand-computes `σ̂²_diff` at `G=4` on deterministic inputs and asserts the `2G`-normalizer form at `atol=1e-12` so any later regression to the (asymptotically equivalent but finite-sample different) `2(G-1)` form would fail the test. `σ̂⁴_W` uses `np.mean(eps_{(g)}² · eps_{(g-1)}²)` which divides by `G-1` via `np.mean` length, matching the paper's Theorem 7 / Equation 29 normalization. - **Note (Phase 3 Yatchew tie policy):** `yatchew_hr_test` REJECTS duplicate dose values with a `UserWarning` + all-NaN result at the front door. The difference-based variance estimator `σ̂²_diff` and the heteroskedasticity-robust scale `σ̂⁴_W` both use adjacent differences (of sorted dy and of adjacent squared residuals, respectively); under tied doses the within-tie row ordering is arbitrary (stable sort falls back to input order), so the statistic becomes non-methodological and order-dependent. Callers with tied doses (mass-point designs, discretised dose registers) should use `stute_test` instead — its tie-safe Cramer-von Mises statistic collapses tie blocks to the post-tie cumulative sum and is provably order-invariant under within-tie permutations. A regression test on `stute_test` asserts bit-identical `cvm_stat` across tied-dose permutations at `atol=1e-14`; a matching regression on `yatchew_hr_test` asserts the NaN+warning behavior on duplicated and constant doses. This tie policy is a Phase 3 diff-diff choice (the paper describes Yatchew only as "sort by D" and does not specify a tie rule); an alternative implementation could document a justified tie-resolution convention and back it with permutation tests, but at the cost of a methodology surface the paper does not cover. **Workflow fallback:** `did_had_pretest_workflow` follows the paper's "Stute OR Yatchew" step-3 wording — a conclusive Stute (which is tie-safe) is sufficient to adjudicate linearity even when Yatchew returns NaN on tied-dose panels. The composite verdict then reads "QUG and linearity diagnostics fail-to-reject (Yatchew NaN - skipped); ..." rather than forcing the whole report to "inconclusive", and `all_pass` is `True` when QUG + at least one conclusive linearity test fail-to-reject. - **Note (Phase 3 exact-linear short-circuit):** Both `stute_test` and `yatchew_hr_test` detect numerically-exact linear fits (OLS residuals below machine precision relative to the signal) and short-circuit to `p=1.0, reject=False` without running the bootstrap / computing `T_hr`. The detection compares `sum(eps^2)` against the CENTERED total sum of squares `sum((dy - dybar)^2)` (ratio `<= 1e-24`), which is equivalent to `1 - R^2` and is TRANSLATION-INVARIANT under additive shifts in dy. Comparing against uncentered `sum(dy^2)` would NOT be translation-invariant: adding a large constant to dy inflates `sum(dy^2)` and can spuriously trip the short-circuit on genuinely noisy data. Regression tests pin (a) the exact-linear fail-to-reject behavior for `dy = a + b*d`, (b) translation-invariance under `dy -> dy + 1e12` (the short-circuit must not fire on noisy data regardless of the constant offset). - [x] Phase 3: `qug_test()` (`T = D_{2,(1)} / (D_{2,(2)} - D_{2,(1)})`, rejection `{T > 1/α - 1}`). One-sided asymptotic p-value `1 / (1 + T)` under the Exp(1)/Exp(1) limit law (Theorem 4). Zero-dose observations filtered upfront (with `UserWarning`); `D_{(1)} == D_{(2)}` tie returns all-NaN inference (conservative). Closed-form tight parity tested at `atol=1e-12`. -- [x] Phase 3: `stute_test()` Cramér-von Mises with Mammen wild bootstrap. Statistic `S = (1/G^2) Σ (cumsum_g)^2` (algebraically equivalent to paper's `Σ(g/G)^2 · ((1/g) Σ eps_{(h)})^2`). Bootstrap follows paper Appendix D Algorithm literal (per-iteration OLS refit). `n_bootstrap=999` default, `n_bootstrap >= 99` validated. `G < 10` returns NaN; `G > 100_000` emits a `UserWarning` pointing to `yatchew_hr_test`. Appendix-D vectorized matrix form deferred as a performance follow-up (tracked in `TODO.md`). +- [x] Phase 3: `stute_test()` Cramér-von Mises with Mammen wild bootstrap. Statistic `S = (1/G^2) Σ (cumsum_g)^2` (algebraically equivalent to paper's `Σ(g/G)^2 · ((1/g) Σ eps_{(h)})^2`). Bootstrap follows paper Appendix D Algorithm literal (per-iteration OLS refit). `n_bootstrap=999` default, `n_bootstrap >= 99` validated. `G < 10` returns NaN; `G > 100_000` emits a `UserWarning` pointing to `yatchew_hr_test`. The ~2× performance follow-up landed 2026-07-07 (loop-invariant hoisting, bit-identical — see the Phase 3 Stute bootstrap Note above). - [x] Phase 3: `yatchew_hr_test()` heteroskedasticity-robust specification test. Test statistic `T_hr = sqrt(G) · (σ̂²_lin - σ̂²_diff) / σ̂²_W` from paper Equation 29. Normalizer `σ̂²_diff` divides by `2G` (paper-literal Theorem 7), NOT `2(G-1)`; hand-computed tight parity asserted at `atol=1e-12`. One-sided standard-normal critical value. `G < 3` returns NaN. Phase 3 shipped only the linearity null (paper Theorem 7); the `null="mean_independence"` R-parity extension shipped post-PR #392 (see the algorithm-variant block above for the contract). - [x] Phase 3: `did_had_pretest_workflow()` composite helper. `data`-only entry point with `aggregate` dispatch: `aggregate="overall"` (default) requires a balanced two-period panel — multi-period panels are rejected at the front door by `_validate_had_panel` with a pointer to `aggregate="event_study"` — and runs steps 1 (QUG) + 3 (Stute + Yatchew-HR) only; `aggregate="event_study"` takes a multi-period panel (>=3 periods) and additionally runs step 2 (joint Stute pre-trends over pre-period horizons) + joint Stute homogeneity over post-period horizons, populating `pretrends_joint` / `homogeneity_joint`. `seed` forwards to all bootstrap-based tests (QUG and Yatchew are deterministic). Returns `HADPretestReport` with priority-ordered verdict string. On `aggregate="overall"` a fail-to-reject verdict explicitly flags the Assumption 7 gap rather than claiming unconditional TWFE safety: `"QUG and linearity diagnostics fail-to-reject; Assumption 7 pre-trends test NOT run (paper step 2 deferred)"`; on `aggregate="event_study"` a fail-to-reject across all three covered diagnostics reads `"TWFE admissible under Section 4 assumptions"` without the Assumption 7 caveat. Verdict priority follows the paper's one-way rule (TWFE admissible only if NO test rejects): **conclusive rejections are the primary verdict and are NEVER hidden by inconclusive status** — any unresolved-step note is appended via `"; additional steps unresolved: ..."` rather than replacing the rejection. The pure `"inconclusive - QUG NaN"` / `"inconclusive - both Stute and Yatchew linearity tests NaN"` forms only fire when NO conclusive test rejects AND a required step is unresolved. The partial-workflow fail-to-reject verdict may carry a `"(Yatchew NaN - skipped)"` (or Stute) suffix when one linearity test is NaN but the other is conclusive (step 3 resolved via the paper's "Stute OR Yatchew" wording). Bundled rejection-reason strings name each failed assumption in the conclusive-rejection case. `all_pass` is `True` iff QUG is conclusive AND at least one of Stute/Yatchew is conclusive AND no conclusive test rejects. **Non-negative-dose contract**: all three raw linearity helpers (`qug_test`, `stute_test`, `yatchew_hr_test`) raise a front-door `ValueError` on any `d < 0`, mirroring the `_validate_had_panel` guard (paper Section 2 HAD support restriction). On the `aggregate="overall"` path, the panel must already be exactly two periods (`_validate_had_panel` raises with a pointer to `aggregate="event_study"` otherwise); the first-difference helper computes `(t_post, t_pre)` per unit and feeds each raw helper directly. On the `aggregate="event_study"` path, joint Stute is dispatched across pre-period and post-period horizons directly (the joint Equation-18 form, no per-horizon pre-slicing). - [x] Phase 4: Pierce-Schott (2016) replication harness reproduces Figure 2 values. **Waived 2026-05-20:** see Deviations block above; the paper itself self-acknowledges that NP estimators are too noisy to be informative on the LBD-restricted PNTR panel (Section 5.2), and R parity at `atol=1e-8` via `tests/test_did_had_parity.py` is a strictly stronger correctness anchor than Figure-2 reproduction on a proxy panel. Tracked as Low-priority follow-up in `TODO.md`. diff --git a/tests/test_had_pretests.py b/tests/test_had_pretests.py index 4a56a81c..ee0778b2 100644 --- a/tests/test_had_pretests.py +++ b/tests/test_had_pretests.py @@ -5879,3 +5879,135 @@ def test_stratified_mc_power_under_alternative(self): f"Stratified Stute power under known alternative: " f"{power:.3f} (target > 0.50 at n_draws={n_draws})" ) + + +class TestStuteBootstrapHoistedInvariants: + """The 2026-07 loop-invariant hoist (batched Mammen draws, precomputed + d-moments, precomputed CvM tie blocks) must be bit-identical to the + literal per-iteration Appendix-D form. Locks (a) the RNG stream + identity of batched/chunked draws, (b) each helper's + precomputed-argument equality, and (c) end-to-end ``p_value``/ + ``cvm_stat`` equality against a frozen byte-copy of the pre-hoist + bootstrap loop.""" + + @staticmethod + def _dgp(G=400, seed=0): + rng = np.random.default_rng(seed) + d = np.round(rng.gamma(2.0, 1.0, G), 1) # rounding induces tie blocks + dy = 1.0 + 0.5 * d + 0.3 * d * d + rng.normal(0.0, 0.5, G) + return d, dy + + def test_mammen_stream_identity_batched_and_chunked(self): + """A (B, G) batched draw and the memory-bounded chunked iterator + both consume the generator's variate stream exactly like B + sequential size-G draws.""" + from diff_diff.had_pretests import _iter_mammen_rows + from diff_diff.utils import _generate_mammen_weights + + B, G = 13, 17 + rng_seq = np.random.default_rng(123) + seq = np.vstack([_generate_mammen_weights(G, rng_seq) for _ in range(B)]) + + rng_batch = np.random.default_rng(123) + batch = _generate_mammen_weights((B, G), rng_batch) + np.testing.assert_array_equal(seq, batch) + + # Force multiple chunks: room for only 3 rows per batch. + rng_chunk = np.random.default_rng(123) + chunked = np.vstack(list(_iter_mammen_rows(B, G, rng_chunk, max_batch_bytes=3 * G * 8))) + np.testing.assert_array_equal(seq, chunked) + + def test_fit_ols_d_moments_bit_identical(self): + from diff_diff.had_pretests import _fit_ols_intercept_slope + + d, dy = self._dgp() + a0, b0, e0 = _fit_ols_intercept_slope(d, dy) + d_mean = d.mean() + d_dev = d - d_mean + moments = (d_mean, d_dev, float(np.dot(d_dev, d_dev))) + a1, b1, e1 = _fit_ols_intercept_slope(d, dy, d_moments=moments) + assert a0 == a1 and b0 == b1 + np.testing.assert_array_equal(e0, e1) + + def test_cvm_tie_blocks_bit_identical(self): + from diff_diff.had_pretests import _cvm_statistic + + d, dy = self._dgp() + idx = np.argsort(d, kind="stable") + d_sorted = d[idx] + assert len(np.unique(d_sorted)) < len(d_sorted), "DGP must contain ties" + eps = dy - dy.mean() + s0 = _cvm_statistic(eps[idx], d_sorted) + _, counts = np.unique(d_sorted, return_counts=True) + tie_blocks = (np.cumsum(counts) - 1, counts) + s1 = _cvm_statistic(eps[idx], d_sorted, tie_blocks=tie_blocks) + assert s0 == s1 + + def test_stute_p_value_bit_identical_to_literal_loop(self): + """Frozen byte-copy of the pre-hoist unweighted bootstrap loop: + stute_test must reproduce its p_value/cvm_stat exactly.""" + from diff_diff.had_pretests import ( + _cvm_statistic, + _fit_ols_intercept_slope, + stute_test, + ) + from diff_diff.utils import _generate_mammen_weights + + for seed_dgp in (0, 1): + d, dy = self._dgp(seed=seed_dgp) + G = d.shape[0] + n_bootstrap, seed = 199, 42 + + r = stute_test(d, dy, n_bootstrap=n_bootstrap, seed=seed) + + # Frozen pre-hoist reference (literal per-iteration form). + a_hat, b_hat, eps = _fit_ols_intercept_slope(d, dy) + idx = np.argsort(d, kind="stable") + d_sorted = d[idx] + S = _cvm_statistic(eps[idx], d_sorted) + rng = np.random.default_rng(seed) + fitted = a_hat + b_hat * d + boot = np.empty(n_bootstrap) + for b in range(n_bootstrap): + eta = _generate_mammen_weights(G, rng) + dy_b = fitted + eps * eta + _, _, eps_b = _fit_ols_intercept_slope(d, dy_b) + boot[b] = _cvm_statistic(eps_b[idx], d_sorted) + p_ref = float((1.0 + float(np.sum(boot >= S))) / (n_bootstrap + 1.0)) + + assert r.cvm_stat == float(S) + assert r.p_value == p_ref + + def test_joint_p_value_invariant_to_per_iteration_draws(self, monkeypatch): + """End-to-end lock of the RNG-stream identity claim for the joint + loop: swapping the memory-bounded batched row iterator for literal + per-iteration draws must leave stute_joint_pretest's p_value and + joint statistic unchanged.""" + import diff_diff.had_pretests as hp + from diff_diff.had_pretests import _fit_ols_intercept_slope, stute_joint_pretest + from diff_diff.utils import _generate_mammen_weights + + d, _ = self._dgp(seed=3) + G = d.shape[0] + X = np.column_stack([np.ones(G), d]) + rng = np.random.default_rng(11) + residuals_by_horizon = {} + fitted_by_horizon = {} + for k in (0, 1): + dy_k = 1.0 + 0.4 * d + 0.2 * d * d + rng.normal(0.0, 0.5, G) + a_k, b_k, eps_k = _fit_ols_intercept_slope(d, dy_k) + residuals_by_horizon[k] = eps_k + fitted_by_horizon[k] = a_k + b_k * d + + kwargs = dict(alpha=0.05, n_bootstrap=199, seed=5) + r_batched = stute_joint_pretest(residuals_by_horizon, fitted_by_horizon, d, X, **kwargs) + + def per_iteration_rows(n_bootstrap, G, rng, max_batch_bytes=None): + for _ in range(n_bootstrap): + yield _generate_mammen_weights(G, rng) + + monkeypatch.setattr(hp, "_iter_mammen_rows", per_iteration_rows) + r_literal = stute_joint_pretest(residuals_by_horizon, fitted_by_horizon, d, X, **kwargs) + + assert r_batched.p_value == r_literal.p_value + assert r_batched.cvm_stat_joint == r_literal.cvm_stat_joint