diff --git a/CHANGELOG.md b/CHANGELOG.md index 2fa76a84..28455c94 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,6 +31,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 Korn-Graubard (1990), and Solon-Haider-Wooldridge (2015) to `docs/references.rst`. ### Changed +- **CallawaySantAnna multiplier bootstrap now tiles weight generation over draws, cutting + peak memory at large `n_units`.** The dense `(n_bootstrap × n_units)` multiplier-weight + matrix (the dominant allocation for the default unit-level bootstrap — `cluster=None`, + equivalently `cluster="unit"` — where each unit is its own + PSU) is generated and consumed one draw-block at a time via the new + `diff_diff/bootstrap_chunking.py` helper instead of being materialized in full. Measured peak + RSS at 999 bootstrap reps drops ~79% at 500k units (11.6 GB → 2.4 GB) and ~68% at 1M units + (10.8 GB → 3.4 GB); the previously out-of-reach millions-of-units × 999-rep regime now stays + near the fit's memory floor. The weight *stream* is bit-identical on both backends (Rust + absolute per-row seeding; NumPy in-order stream); end-to-end bootstrap SEs match to within + floating-point reassociation of the BLAS reductions (~1 ULP, far below bootstrap Monte-Carlo + error). Stratified survey designs (few PSUs) are unchanged (full generation + sliced blocks); + see TODO.md for the deferred per-stratum tiling. - **`run_placebo_test`'s `fake_group` path now filters ever-treated units by default.** The dispatcher threads its `treatment` column into `placebo_group_test`, so the fake-group placebo runs on never-treated units only (a more-correct placebo). Calling diff --git a/TODO.md b/TODO.md index 72b298c4..b4d3dee7 100644 --- a/TODO.md +++ b/TODO.md @@ -66,6 +66,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | `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 | | Rust faer SVD ndarray-to-faer conversion overhead (minimal vs SVD cost). | `rust/src/linalg.rs:67` | #115 | Quick | Low | +| CallawaySantAnna multiplier-bootstrap weight chunking covers the **unstratified** survey-PSU generation (the default unit-level bootstrap — `cluster=None`, equivalently `cluster="unit"` — the large-`n_units` OOM case). Two gaps remain: (1) EfficientDiD and HAD bootstraps still materialize the full `(n_bootstrap × n_units)` weight matrix — wire them through `diff_diff/bootstrap_chunking.py`; (2) 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`, `efficient_did_bootstrap.py`, `had.py` | follow-up | Mid | Low | ### Testing / docs diff --git a/diff_diff/bootstrap_chunking.py b/diff_diff/bootstrap_chunking.py new file mode 100644 index 00000000..5cd597dc --- /dev/null +++ b/diff_diff/bootstrap_chunking.py @@ -0,0 +1,185 @@ +"""Memory-bounded chunking for multiplier-bootstrap weight matrices. + +The multiplier bootstrap perturbs cached influence functions with a dense +``(n_bootstrap, n_units)`` weight matrix. At large ``n_units`` that matrix +dominates peak memory (e.g. ``999 x 5_000_000 x 8`` bytes is ~40 GB). Every +consumer is a left-multiply ``weights @ influence_vector`` whose result is small +(``(n_bootstrap,)`` or ``(n_bootstrap, n_gt)``), so the bootstrap can be tiled +over the *draw* dimension: generate and consume the weights in row-blocks of +``B``, capping the live intermediate at ``(B, n_units)``. FLOPs are identical to +the un-chunked path -- only the draw axis is tiled. The generated weight stream +is *bit-identical* to the un-chunked matrix (see below); the downstream +``weights @ influence`` matmuls go through BLAS, whose reduction order depends on +the operand row-count, so the resulting statistics match the un-chunked path to +within floating-point reassociation (typically <~1 ULP), far below bootstrap +Monte-Carlo error -- not bit-for-bit. + +Bit-identity of the weight *generation* is preserved on **both** backends: + +- **Rust** seeds each row absolutely as ``base_seed + row_index`` + (``rust/src/bootstrap.rs``), so calling the generator per block with base seed + ``base_seed + chunk_start`` reproduces the exact un-chunked rows. Exactly one + ``rng.integers`` draw is consumed, matching the un-chunked wrapper. +- The **NumPy** fallback draws the matrix row-major from the ``Generator`` + stream, so consuming it in contiguous, in-order blocks from the same generator + reproduces the identical sequence. +""" + +from __future__ import annotations + +from typing import Iterator, Optional, Tuple + +import numpy as np + +from diff_diff._backend import HAS_RUST_BACKEND, _rust_bootstrap_weights +from diff_diff.bootstrap_utils import generate_bootstrap_weights_batch_numpy + +# Byte ceiling for a single ``(B, n_units)`` float64 weight block. 256 MB keeps +# the live intermediate small at millions of units while staying large enough +# that the per-block matmuls remain BLAS-efficient and chunk overhead (a handful +# of extra Python iterations / FFI calls) is negligible. +_TARGET_BLOCK_BYTES = 256 * 1024 * 1024 + + +def compute_block_size( + n_units: int, n_bootstrap: int, target_bytes: int = _TARGET_BLOCK_BYTES +) -> int: + """Number of bootstrap rows per block so a ``(B, n_units)`` float64 block + stays under ``target_bytes``. Always in ``[1, n_bootstrap]``.""" + if n_units <= 0: + return max(1, n_bootstrap) + b = target_bytes // (n_units * 8) + return int(max(1, min(max(1, n_bootstrap), b))) + + +def iter_weight_blocks( + n_bootstrap: int, + n_gen: int, + weight_type: str, + rng: np.random.Generator, + *, + expand_index: Optional[np.ndarray] = None, + block_size: Optional[int] = None, +) -> Iterator[Tuple[int, np.ndarray]]: + """Yield ``(chunk_start, block)`` pairs covering all ``n_bootstrap`` draws. + + ``block`` has shape ``(B, width)`` where ``width = len(expand_index)`` when + ``expand_index`` is given, else ``n_gen``. Weights are generated at width + ``n_gen`` (unit / cluster / PSU level) and, when ``expand_index`` is given, + expanded to unit level via ``block[:, expand_index]`` (cluster->unit or + PSU->unit fan-out). The concatenation of all yielded blocks is bit-identical + to a single ``generate_bootstrap_weights_batch(n_bootstrap, n_gen, ...)`` + followed by the same expansion. + + Generation is in-order and stateful on ``rng`` (NumPy fallback) -- the caller + must consume the iterator sequentially, which the chunk loop does. + """ + width = n_gen if expand_index is None else int(len(expand_index)) + if block_size is None: + block_size = compute_block_size(width, n_bootstrap) + if block_size < 1: + raise ValueError(f"block_size must be >= 1, got {block_size}") + + rust_gen = ( + _rust_bootstrap_weights + if (HAS_RUST_BACKEND and _rust_bootstrap_weights is not None) + else None + ) + # Draw exactly one base seed (matching the un-chunked Rust wrapper); the + # NumPy fallback consumes the rng stream directly per block instead. + base_seed = int(rng.integers(0, 2**63 - 1)) if rust_gen is not None else 0 + + for chunk_start in range(0, n_bootstrap, block_size): + rows = min(block_size, n_bootstrap - chunk_start) + if rust_gen is not None: + block = rust_gen(rows, n_gen, weight_type, base_seed + chunk_start) + else: + block = generate_bootstrap_weights_batch_numpy(rows, n_gen, weight_type, rng) + if expand_index is not None: + block = block[:, expand_index] + yield chunk_start, block + + +def iter_survey_multiplier_weight_blocks( + n_bootstrap: int, + resolved_survey: object, + weight_type: str, + rng: np.random.Generator, + *, + block_size: int, +) -> Tuple[np.ndarray, Iterator[Tuple[int, np.ndarray]]]: + """Chunked PSU-level multiplier weights for the survey-aware bootstrap. + + Returns ``(psu_ids, blocks)`` where ``blocks`` yields + ``(chunk_start, (B, n_psu))`` PSU-weight blocks covering all draws. + + For UNSTRATIFIED designs (``strata is None``, ``n_psu >= 2``) the + ``(n_bootstrap, n_psu)`` matrix is generated one draw-block at a time via + :func:`iter_weight_blocks` plus the unstratified FPC scalar -- bit-identical + to the unstratified branch of + :func:`diff_diff.bootstrap_utils.generate_survey_multiplier_weights_batch`, + but the full matrix is never materialized. This is the path taken by + ``cluster="unit"`` (each unit its own PSU, ``n_psu == n_units``), the case + that otherwise dominates bootstrap memory at large n_units. + + Stratified designs (and the ``n_psu < 2`` degenerate case) fall back to full + generation + sliced blocks: per-stratum / lonely-PSU generation is not tiled + here, but stratified designs have few PSUs so the full matrix is small. + """ + from diff_diff.bootstrap_utils import generate_survey_multiplier_weights_batch + + if block_size < 1: + raise ValueError(f"block_size must be >= 1, got {block_size}") + + psu = getattr(resolved_survey, "psu", None) + strata = getattr(resolved_survey, "strata", None) + if psu is None: + n_psu = len(resolved_survey.weights) # type: ignore[attr-defined] + psu_ids = np.arange(n_psu) + else: + psu_ids = np.unique(psu) + n_psu = len(psu_ids) + + if strata is not None or n_psu < 2: + # Stratified or degenerate single-PSU: full generation (small here). + weights, psu_ids = generate_survey_multiplier_weights_batch( + n_bootstrap, resolved_survey, weight_type, rng + ) + + def _sliced() -> Iterator[Tuple[int, np.ndarray]]: + for chunk_start in range(0, n_bootstrap, block_size): + yield chunk_start, weights[chunk_start : chunk_start + block_size] + + return psu_ids, _sliced() + + # Unstratified, n_psu >= 2: tile the generation over draws. Mirror the + # unstratified FPC scaling from generate_survey_multiplier_weights_batch. + fpc = getattr(resolved_survey, "fpc", None) + fpc_scale = 1.0 + fpc_zero = False + if fpc is not None: + # psu=None already sets n_psu = len(weights), so n_units_for_fpc == n_psu + # on both branches of the original generator. + n_units_for_fpc = n_psu + if fpc[0] < n_units_for_fpc: + raise ValueError( + f"FPC ({fpc[0]}) is less than the number of PSUs " + f"({n_units_for_fpc}). FPC must be >= number of PSUs." + ) + f = n_units_for_fpc / fpc[0] + if f < 1.0: + fpc_scale = float(np.sqrt(1.0 - f)) + else: + fpc_zero = True + + def _generated() -> Iterator[Tuple[int, np.ndarray]]: + for chunk_start, block in iter_weight_blocks( + n_bootstrap, n_psu, weight_type, rng, block_size=block_size + ): + if fpc_zero: + block = np.zeros_like(block) + elif fpc_scale != 1.0: + block = block * fpc_scale + yield chunk_start, block + + return psu_ids, _generated() diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index 46ee15bf..13497b11 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -282,11 +282,16 @@ class CallawaySantAnna( Recommended: 999 or more for reliable inference. .. note:: Memory Usage - The bootstrap stores all weights in memory as a (n_bootstrap, n_units) - float64 array. For large datasets, this can be significant: - - 1K bootstrap × 10K units = ~80 MB - - 10K bootstrap × 100K units = ~8 GB - Consider reducing n_bootstrap if memory is constrained. + Bootstrap multiplier weights are generated and consumed one + draw-block at a time (see :mod:`diff_diff.bootstrap_chunking`), so the + full ``(n_bootstrap, n_units)`` weight matrix is never materialized. + The live weight intermediate is bounded by roughly + ``max(~256 MB, 8 * n_units)`` bytes -- a block holds at least one full + draw row -- independent of ``n_bootstrap``. Only the small bootstrap + *output* arrays (``(n_bootstrap, n_group_time)`` and ``(n_bootstrap,)`` + per aggregation) stay fully in memory. Stratified survey designs are + the current exception (the full PSU-weight matrix is built up front, + but PSUs are few). bootstrap_weights : str, default="rademacher" Type of weights for multiplier bootstrap: @@ -445,7 +450,6 @@ def __init__( pscore_fallback: str = "error", vcov_type: str = "hc1", ): - import warnings if control_group not in ["never_treated", "not_yet_treated"]: raise ValueError( diff --git a/diff_diff/staggered_bootstrap.py b/diff_diff/staggered_bootstrap.py index ba057b22..783b97c6 100644 --- a/diff_diff/staggered_bootstrap.py +++ b/diff_diff/staggered_bootstrap.py @@ -8,10 +8,15 @@ import warnings from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple import numpy as np +from diff_diff.bootstrap_chunking import ( + compute_block_size, + iter_survey_multiplier_weight_blocks, + iter_weight_blocks, +) from diff_diff.bootstrap_utils import ( compute_bootstrap_pvalue as _compute_bootstrap_pvalue_func, ) @@ -24,12 +29,6 @@ from diff_diff.bootstrap_utils import ( compute_percentile_ci as _compute_percentile_ci_func, ) -from diff_diff.bootstrap_utils import ( - generate_bootstrap_weights_batch as _generate_bootstrap_weights_batch, -) -from diff_diff.bootstrap_utils import ( - generate_survey_multiplier_weights_batch as _generate_survey_multiplier_weights_batch, -) if TYPE_CHECKING: import pandas as pd @@ -347,9 +346,20 @@ def _run_multiplier_bootstrap( _bootstrap_cluster_variance_unidentified = False if _use_survey_bootstrap: - # PSU-level multiplier weights - psu_weights, psu_ids = _generate_survey_multiplier_weights_batch( - self.n_bootstrap, resolved_survey_unit, self.bootstrap_weights, rng + # PSU-level multiplier weights, generated AND expanded one draw-block + # at a time so the (n_bootstrap, n_units) matrix is never built in + # full. This is the dominant allocation at large n_units, including + # the default unit-level bootstrap (cluster=None, equivalently + # cluster="unit": each unit its own PSU, n_psu == n_units). + # Unstratified designs tile the generation; stratified designs (few + # PSUs) fall back to full generation + sliced blocks. + _block_size = compute_block_size(n_units, self.n_bootstrap) + psu_ids, _psu_blocks = iter_survey_multiplier_weight_blocks( + self.n_bootstrap, + resolved_survey_unit, + self.bootstrap_weights, + rng, + block_size=_block_size, ) if len(psu_ids) < 2: import warnings as _warnings @@ -377,44 +387,35 @@ def _run_multiplier_bootstrap( # Each unit is its own PSU — identity mapping unit_to_psu_col = np.arange(n_units) - # Expand PSU weights to unit level for per-(g,t) perturbation - # Shape: (n_bootstrap, n_units) - all_bootstrap_weights = psu_weights[:, unit_to_psu_col] - else: - # Standard unit-level weights (no survey or weights-only) - all_bootstrap_weights = _generate_bootstrap_weights_batch( - self.n_bootstrap, n_units, self.bootstrap_weights, rng + # When each unit is its own PSU (e.g. cluster="unit"), the PSU block + # is already unit-aligned, so the fancy-index expansion is an + # identity permutation whose only effect is a needless full-block + # copy (doubling live block memory). Detect that once and skip it. + _psu_is_identity = len(psu_ids) == n_units and bool( + np.array_equal(unit_to_psu_col, np.arange(n_units)) ) - # Vectorized bootstrap ATT(g,t) computation - # Compute all bootstrap ATTs for all (g,t) pairs using matrix operations - bootstrap_atts_gt = np.zeros((self.n_bootstrap, n_gt)) + # Expand each PSU-weight block to unit level on demand; the full + # (n_bootstrap, n_units) expansion is never materialized at once. + def _weight_blocks() -> Iterator[Tuple[int, np.ndarray]]: + for _cs, _psu_block in _psu_blocks: + if _psu_is_identity: + yield _cs, _psu_block + else: + yield _cs, _psu_block[:, unit_to_psu_col] - for j in range(n_gt): - treated_idx = gt_treated_indices[j] - control_idx = gt_control_indices[j] - treated_inf = gt_treated_inf[j] - control_inf = gt_control_inf[j] - - # Extract weights for this (g,t)'s units across all bootstrap iterations - # Shape: (n_bootstrap, n_treated) and (n_bootstrap, n_control) - treated_weights = all_bootstrap_weights[:, treated_idx] - control_weights = all_bootstrap_weights[:, control_idx] - - # Vectorized perturbation: matrix-vector multiply - # Shape: (n_bootstrap,) - # Suppress RuntimeWarnings for edge cases (small samples, extreme weights) - with np.errstate(divide="ignore", invalid="ignore", over="ignore"): - perturbations = treated_weights @ treated_inf + control_weights @ control_inf - - # Let non-finite values propagate - they will be handled at statistics computation - bootstrap_atts_gt[:, j] = original_atts[j] + perturbations - - # Vectorized overall ATT using combined IF (includes WIF) - # Shape: (n_bootstrap,) - if skip_overall_aggregation: - bootstrap_overall = np.full(self.n_bootstrap, np.nan) + weight_blocks: Iterator[Tuple[int, np.ndarray]] = _weight_blocks() else: + # Standard unit-level weights (no survey or weights-only), generated + # one row-block at a time directly at unit width. + weight_blocks = iter_weight_blocks( + self.n_bootstrap, n_units, self.bootstrap_weights, rng + ) + + # Pre-compute the overall combined IF once (reused across every block). + # None exactly when the overall aggregation is skipped. + overall_combined_if: Optional[np.ndarray] = None + if not skip_overall_aggregation: # Use combined IF (standard IF + WIF) for proper bootstrap post_gt_pairs = [gt_pairs[i] for i in post_treatment_indices] post_groups = np.array([gt_pairs[i][0] for i in post_treatment_indices]) @@ -431,38 +432,69 @@ def _run_multiplier_bootstrap( global_unit_to_idx=unit_to_idx, n_global_units=n_units, ) - with np.errstate(divide="ignore", invalid="ignore", over="ignore"): - bootstrap_overall = original_overall + all_bootstrap_weights @ overall_combined_if - # Vectorized event study aggregation using combined IFs - # Non-finite values handled at statistics computation stage + # Pre-allocate the small bootstrap output arrays. Only these (sized in + # n_bootstrap, not n_units) persist; each weight block is discarded once + # its rows are written, capping peak memory at O(block x n_units). + bootstrap_atts_gt = np.empty((self.n_bootstrap, n_gt)) + if skip_overall_aggregation: + bootstrap_overall = np.full(self.n_bootstrap, np.nan) + else: + bootstrap_overall = np.empty(self.n_bootstrap) + rel_periods: List[int] = [] bootstrap_event_study: Optional[Dict[int, np.ndarray]] = None if event_study_info is not None: rel_periods = sorted(event_study_info.keys()) - bootstrap_event_study = {} - for e in rel_periods: - agg_info = event_study_info[e] - # Use combined IF (standard IF + WIF) for proper bootstrap - with np.errstate(divide="ignore", invalid="ignore", over="ignore"): - bootstrap_event_study[e] = ( - agg_info["effect"] + all_bootstrap_weights @ agg_info["combined_if"] - ) + bootstrap_event_study = {e: np.empty(self.n_bootstrap) for e in rel_periods} - # Vectorized group aggregation - # Non-finite values handled at statistics computation stage group_list: List[Any] = [] bootstrap_group: Optional[Dict[Any, np.ndarray]] = None if group_agg_info is not None: group_list = sorted(group_agg_info.keys()) - bootstrap_group = {} - for g in group_list: - agg_info = group_agg_info[g] - gt_indices = agg_info["gt_indices"] - weights = agg_info["weights"] - # Suppress RuntimeWarnings for edge cases + bootstrap_group = {g: np.empty(self.n_bootstrap) for g in group_list} + + # Consume the weights one row-block at a time. Each block fills its rows + # of every output array; only the draw axis is tiled. The weight stream + # is bit-identical to the un-chunked path; the BLAS weights @ influence + # reductions may reassociate, so statistics match to within ~1 ULP (far + # below bootstrap Monte-Carlo error), not bit-for-bit. + for _chunk_start, _w_block in weight_blocks: + _rows = slice(_chunk_start, _chunk_start + _w_block.shape[0]) + + # ATT(g,t) + for j in range(n_gt): + treated_weights = _w_block[:, gt_treated_indices[j]] + control_weights = _w_block[:, gt_control_indices[j]] + with np.errstate(divide="ignore", invalid="ignore", over="ignore"): + perturbations = ( + treated_weights @ gt_treated_inf[j] + control_weights @ gt_control_inf[j] + ) + bootstrap_atts_gt[_rows, j] = original_atts[j] + perturbations + + # Overall ATT (combined IF includes WIF); skipped when None. + if overall_combined_if is not None: with np.errstate(divide="ignore", invalid="ignore", over="ignore"): - bootstrap_group[g] = bootstrap_atts_gt[:, gt_indices] @ weights + bootstrap_overall[_rows] = original_overall + _w_block @ overall_combined_if + + # Event study aggregation (combined IFs) + if bootstrap_event_study is not None and event_study_info is not None: + for e in rel_periods: + agg_info = event_study_info[e] + with np.errstate(divide="ignore", invalid="ignore", over="ignore"): + bootstrap_event_study[e][_rows] = ( + agg_info["effect"] + _w_block @ agg_info["combined_if"] + ) + + # Group aggregation (reads this block's freshly written ATT(g,t) rows) + if bootstrap_group is not None and group_agg_info is not None: + for g in group_list: + agg_info = group_agg_info[g] + with np.errstate(divide="ignore", invalid="ignore", over="ignore"): + bootstrap_group[g][_rows] = ( + bootstrap_atts_gt[_rows][:, agg_info["gt_indices"]] + @ agg_info["weights"] + ) # Batch compute bootstrap statistics for ATT(g,t) batch_ses, batch_ci_lo, batch_ci_hi, batch_pv = _compute_effect_bootstrap_stats_batch_func( diff --git a/docs/doc-deps.yaml b/docs/doc-deps.yaml index 7fd25a50..c0408168 100644 --- a/docs/doc-deps.yaml +++ b/docs/doc-deps.yaml @@ -907,6 +907,14 @@ sources: type: roadmap note: "Phase 6 bootstrap+survey interaction" + diff_diff/bootstrap_chunking.py: + drift_risk: low + docs: + - path: docs/methodology/REGISTRY.md + section: "Bootstrap, Survey Bootstrap" + type: methodology + note: "Memory-bounded draw-block chunking of multiplier-bootstrap weight generation/consumption (shared by CallawaySantAnna; EfficientDiD/HAD adoption is follow-up). Implementation detail — numerically equivalent (~1 ULP, BLAS reassociation) to the un-chunked path; weight stream bit-identical." + diff_diff/prep.py: drift_risk: low docs: diff --git a/tests/test_bootstrap_chunking.py b/tests/test_bootstrap_chunking.py new file mode 100644 index 00000000..bfd69688 --- /dev/null +++ b/tests/test_bootstrap_chunking.py @@ -0,0 +1,323 @@ +"""Tests for memory-bounded multiplier-bootstrap weight chunking. + +The chunking in :mod:`diff_diff.bootstrap_chunking` tiles the bootstrap *draw* +dimension to cap peak memory at ``O(block x n_units)`` instead of +``O(n_bootstrap x n_units)``. Its load-bearing guarantee is that tiling +reproduces the un-chunked weight *stream* exactly (bit-identical), on whichever +backend is active (Rust absolute per-row seeding; NumPy in-order stream). These +tests lock the weight-stream bit-identity at the helper level and end-to-end +chunk-invariance (to floating-point reassociation) through CallawaySantAnna, +under whatever ``DIFF_DIFF_BACKEND`` the CI matrix selects. +""" + +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest + +from diff_diff import CallawaySantAnna +from diff_diff.bootstrap_chunking import ( + compute_block_size, + iter_survey_multiplier_weight_blocks, + iter_weight_blocks, +) +from diff_diff.bootstrap_utils import ( + generate_bootstrap_weights_batch, + generate_survey_multiplier_weights_batch, +) + +WEIGHT_TYPES = ["rademacher", "mammen", "webb"] + + +def _stack(n_bootstrap, n_gen, weight_type, seed, block_size, expand_index=None): + """Concatenate all weight blocks from a fresh, identically-seeded rng.""" + rng = np.random.default_rng(seed) + blocks = list( + iter_weight_blocks( + n_bootstrap, + n_gen, + weight_type, + rng, + expand_index=expand_index, + block_size=block_size, + ) + ) + starts = [cs for cs, _ in blocks] + mat = np.vstack([b for _, b in blocks]) + return starts, mat + + +class TestComputeBlockSize: + def test_always_in_bounds(self): + assert compute_block_size(1000, 200) <= 200 + assert compute_block_size(1000, 200) >= 1 + + def test_huge_n_units_floors_to_one_row(self): + assert compute_block_size(10**9, 200) == 1 + + def test_tiny_n_units_fits_all_in_one_block(self): + assert compute_block_size(1, 200) == 200 + + def test_respects_target_bytes(self): + # 100 units x 8 bytes = 800 B/row; an 8000 B budget -> 10 rows/block + assert compute_block_size(100, 500, target_bytes=8000) == 10 + + +class TestWeightStreamBitIdentity: + """Tiling the draw dimension reproduces the single-block stream exactly.""" + + @pytest.mark.parametrize("weight_type", WEIGHT_TYPES) + @pytest.mark.parametrize("block_size", [1, 7, 33, 198]) + def test_chunked_equals_single_block(self, weight_type, block_size): + n_bootstrap, n_gen, seed = 199, 53, 12345 + _, single = _stack(n_bootstrap, n_gen, weight_type, seed, block_size=n_bootstrap) + starts, chunked = _stack(n_bootstrap, n_gen, weight_type, seed, block_size=block_size) + assert single.shape == (n_bootstrap, n_gen) + assert chunked.shape == (n_bootstrap, n_gen) + # exact: the chunking promise is bit-identity, not approximate equality + np.testing.assert_array_equal(chunked, single) + # blocks cover every draw exactly once, in order + assert starts == list(range(0, n_bootstrap, block_size)) + + @pytest.mark.parametrize("weight_type", WEIGHT_TYPES) + def test_expand_index_is_chunk_invariant(self, weight_type): + # cluster/PSU fan-out: generate at n_gen, expand to unit width per block + n_bootstrap, n_clusters, n_units, seed = 100, 9, 40, 7 + expand = np.array([i % n_clusters for i in range(n_units)]) + _, single = _stack( + n_bootstrap, + n_clusters, + weight_type, + seed, + block_size=n_bootstrap, + expand_index=expand, + ) + _, chunked = _stack( + n_bootstrap, + n_clusters, + weight_type, + seed, + block_size=11, + expand_index=expand, + ) + assert single.shape == (n_bootstrap, n_units) + np.testing.assert_array_equal(chunked, single) + + @pytest.mark.parametrize("weight_type", WEIGHT_TYPES) + def test_single_block_matches_legacy_generator(self, weight_type): + # iter_weight_blocks in single-block mode must reproduce the legacy + # generate_bootstrap_weights_batch wrapper exactly (matched seeds), so the + # chunked path is anchored to the pre-existing generator, not just to its + # own single-block mode. + n_bootstrap, n_gen, seed = 199, 53, 999 + legacy = generate_bootstrap_weights_batch( + n_bootstrap, n_gen, weight_type, np.random.default_rng(seed) + ) + _, chunked = _stack(n_bootstrap, n_gen, weight_type, seed, block_size=n_bootstrap) + assert chunked.shape == legacy.shape + np.testing.assert_array_equal(chunked, legacy) + + +class TestCSBootstrapChunkInvariance: + """CallawaySantAnna bootstrap output is invariant to the chunk size. + + The generated weight stream is bit-identical across chunk sizes (locked by + ``TestWeightStreamBitIdentity``). The downstream ``weights @ influence`` + matmuls go through BLAS, whose reduction order depends on the operand + row-count, so the resulting statistics match to within floating-point + reassociation (~1 ULP) rather than bit-for-bit -- far below bootstrap + Monte-Carlo error. This mirrors the repo's assert_allclose convention for + float linalg. + """ + + @staticmethod + def _panel(): + rng = np.random.default_rng(0) + nu, nt = 120, 8 + units = np.repeat(np.arange(nu), nt) + periods = np.tile(np.arange(nt), nu) + n = nu * nt + cohort = rng.integers(0, 3, nu) + ft_unit = np.where(cohort == 0, 0, np.where(cohort == 1, 3, 5)) + ft = np.repeat(ft_unit, nt) + post = (periods >= ft) & (ft > 0) + y = rng.standard_normal(n) + 0.1 * periods + 2.0 * post + 0.5 * np.repeat(cohort, nt) + return pd.DataFrame({"unit": units, "period": periods, "y": y, "first_treat": ft}) + + def _fit(self): + return CallawaySantAnna( + control_group="never_treated", + estimation_method="dr", + cluster="unit", + n_bootstrap=200, + seed=42, + ).fit( + self._panel(), + outcome="y", + unit="unit", + time="period", + first_treat="first_treat", + aggregate="all", + ) + + def test_tiny_chunks_match_single_chunk(self, monkeypatch): + # Default path: this small panel fits in a single block. + base = self._fit() + # Force many tiny blocks on every weight path (unit + survey). + monkeypatch.setattr("diff_diff.bootstrap_chunking.compute_block_size", lambda *a, **k: 7) + monkeypatch.setattr("diff_diff.staggered_bootstrap.compute_block_size", lambda *a, **k: 7) + tiny = self._fit() + + # Continuous bootstrap statistics match to within BLAS reassociation. + assert tiny.overall_se == pytest.approx(base.overall_se, rel=1e-10, abs=1e-12) + assert tiny.cband_crit_value == pytest.approx(base.cband_crit_value, rel=1e-10, abs=1e-12) + # p-values are discrete proportions over draws; a borderline draw can + # flip under reassociation, shifting a p-value by O(1/n_bootstrap). + assert tiny.overall_p_value == pytest.approx(base.overall_p_value, abs=0.02) + + b = base.to_dataframe().sort_values(["group", "time"]).reset_index(drop=True) + t = tiny.to_dataframe().sort_values(["group", "time"]).reset_index(drop=True) + for col in ["se", "conf_int_lower", "conf_int_upper"]: + np.testing.assert_allclose(t[col].to_numpy(), b[col].to_numpy(), rtol=1e-10, atol=1e-12) + np.testing.assert_allclose(t["p_value"].to_numpy(), b["p_value"].to_numpy(), atol=0.02) + + # Event-study and group aggregate effects/SEs/CIs also match under chunking. + for level in ("event_study", "group"): + bl = base.to_dataframe(level=level).reset_index(drop=True) + tl = tiny.to_dataframe(level=level).reset_index(drop=True) + num_cols = [c for c in bl.columns if bl[c].dtype.kind in "fi" and c != "p_value"] + assert num_cols, f"no numeric columns to compare for level={level}" + for col in num_cols: + np.testing.assert_allclose( + tl[col].to_numpy(), bl[col].to_numpy(), rtol=1e-9, atol=1e-10 + ) + + def test_cluster_none_default_chunks_match_single(self, monkeypatch): + # cluster=None is the public default (auto-clusters at unit); confirm the + # default path is chunk-invariant end-to-end. + def fit(): + return CallawaySantAnna( + control_group="never_treated", + estimation_method="dr", + cluster=None, + n_bootstrap=200, + seed=42, + ).fit( + self._panel(), + outcome="y", + unit="unit", + time="period", + first_treat="first_treat", + aggregate="simple", + ) + + base = fit() + monkeypatch.setattr("diff_diff.bootstrap_chunking.compute_block_size", lambda *a, **k: 7) + monkeypatch.setattr("diff_diff.staggered_bootstrap.compute_block_size", lambda *a, **k: 7) + tiny = fit() + assert tiny.overall_se == pytest.approx(base.overall_se, rel=1e-10, abs=1e-12) + + @staticmethod + def _clustered_panel(): + # Units grouped into states (5 units/state) -> the unit->PSU map is a + # genuine many-units-to-one-PSU fan-out (non-identity expansion), unlike + # cluster="unit" above. + rng = np.random.default_rng(1) + n_states, units_per_state, nt = 10, 5, 8 + nu = n_states * units_per_state + units = np.repeat(np.arange(nu), nt) + periods = np.tile(np.arange(nt), nu) + n = nu * nt + cohort = rng.integers(0, 3, nu) + ft_unit = np.where(cohort == 0, 0, np.where(cohort == 1, 3, 5)) + ft = np.repeat(ft_unit, nt) + post = (periods >= ft) & (ft > 0) + y = rng.standard_normal(n) + 0.1 * periods + 2.0 * post + 0.5 * np.repeat(cohort, nt) + state = np.repeat(np.repeat(np.arange(n_states), units_per_state), nt) + return pd.DataFrame( + {"unit": units, "period": periods, "y": y, "first_treat": ft, "state": state} + ) + + def _fit_clustered(self): + return CallawaySantAnna( + control_group="never_treated", + estimation_method="dr", + cluster="state", + n_bootstrap=200, + seed=42, + ).fit( + self._clustered_panel(), + outcome="y", + unit="unit", + time="period", + first_treat="first_treat", + aggregate="all", + ) + + def test_nonidentity_cluster_chunks_match_single(self, monkeypatch): + # Exercises the non-identity PSU fan-out expansion under tiny chunks. + base = self._fit_clustered() + monkeypatch.setattr("diff_diff.bootstrap_chunking.compute_block_size", lambda *a, **k: 7) + monkeypatch.setattr("diff_diff.staggered_bootstrap.compute_block_size", lambda *a, **k: 7) + tiny = self._fit_clustered() + + assert tiny.overall_se == pytest.approx(base.overall_se, rel=1e-10, abs=1e-12) + b = base.to_dataframe().sort_values(["group", "time"]).reset_index(drop=True) + t = tiny.to_dataframe().sort_values(["group", "time"]).reset_index(drop=True) + for col in ["se", "conf_int_lower", "conf_int_upper"]: + np.testing.assert_allclose(t[col].to_numpy(), b[col].to_numpy(), rtol=1e-10, atol=1e-12) + + +def _design(psu=None, strata=None, fpc=None, weights=None, lonely_psu="adjust"): + """Minimal duck-typed ResolvedSurveyDesign for the survey weight generators.""" + return SimpleNamespace(psu=psu, strata=strata, fpc=fpc, weights=weights, lonely_psu=lonely_psu) + + +class TestSurveyWeightBlocks: + """`iter_survey_multiplier_weight_blocks` reproduces the full survey generator. + + The chunked survey path must yield the exact same PSU-weight matrix and + psu_ids as `generate_survey_multiplier_weights_batch`, so the + CallawaySantAnna survey/cluster bootstrap is bit-identical regardless of + chunking. Covers unstratified generation (tiled), `psu=None`, the FPC + scalar, and the stratified fallback (sliced). + """ + + def _assert_matches_full(self, design, weight_type, seed, block_size): + full_w, full_ids = generate_survey_multiplier_weights_batch( + 199, design, weight_type, np.random.default_rng(seed) + ) + ids, blocks = iter_survey_multiplier_weight_blocks( + 199, design, weight_type, np.random.default_rng(seed), block_size=block_size + ) + chunked = np.vstack([b for _, b in blocks]) + np.testing.assert_array_equal(ids, full_ids) + assert chunked.shape == full_w.shape + np.testing.assert_array_equal(chunked, full_w) + + @pytest.mark.parametrize("weight_type", WEIGHT_TYPES) + def test_unstratified_psu_tiled_matches_full(self, weight_type): + # 8 PSUs, 3 units each; tiled generation (block_size << n_bootstrap) + psu = np.repeat(np.arange(8), 3) + design = _design(psu=psu, weights=np.ones(len(psu))) + self._assert_matches_full(design, weight_type, seed=5, block_size=7) + + @pytest.mark.parametrize("weight_type", WEIGHT_TYPES) + def test_psu_none_matches_full(self, weight_type): + # psu=None -> each observation is its own PSU + design = _design(psu=None, weights=np.ones(20)) + self._assert_matches_full(design, weight_type, seed=9, block_size=11) + + def test_unstratified_fpc_scaling_matches_full(self): + # f = n_psu / fpc = 10/100 = 0.1 -> sqrt(0.9) scaling on every weight + psu = np.arange(10) + design = _design(psu=psu, fpc=np.full(10, 100.0), weights=np.ones(10)) + self._assert_matches_full(design, "rademacher", seed=3, block_size=13) + + def test_stratified_fallback_matches_full(self): + # 2 strata x 3 PSUs -> falls back to full generation + sliced blocks + psu = np.arange(6) + strata = np.array([0, 0, 0, 1, 1, 1]) + design = _design(psu=psu, strata=strata, weights=np.ones(6)) + self._assert_matches_full(design, "rademacher", seed=7, block_size=9) diff --git a/tests/test_methodology_callaway.py b/tests/test_methodology_callaway.py index ee882039..f17a6d10 100644 --- a/tests/test_methodology_callaway.py +++ b/tests/test_methodology_callaway.py @@ -21,7 +21,9 @@ from diff_diff import CallawaySantAnna from diff_diff.prep import generate_staggered_data -from diff_diff.staggered_bootstrap import _generate_bootstrap_weights_batch +from diff_diff.bootstrap_utils import ( + generate_bootstrap_weights_batch as _generate_bootstrap_weights_batch, +) # =============================================================================