diff --git a/CHANGELOG.md b/CHANGELOG.md index 48c11da5..ed8ab360 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -127,6 +127,22 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 supersedes it. ### Changed +- **CallawaySantAnna multiplier bootstrap rewritten as a fused, column-tiled scatter-GEMM** + (EfficientDiD's bootstrap routes through the same kernel). The former loop sliced the + `(block × n_units)` weight matrix twice per (g,t) cell per weight block — at a 40-period, + 10-cohort, 100k-unit panel (390 cells, 999 draws) that was ~95% of bootstrap wall time as + pure memory traffic. All perturbation columns (per-cell IFs, overall combined IF, + per-event-time combined IFs) are now scattered into a column-tiled influence matrix + (byte-capped tiles; each tile replays the bit-identical weight stream via RNG state + snapshot/restore) and consumed by one BLAS GEMM per weight block. Measured (5-rep medians, + rust backend + Accelerate): bootstrap fits 24.5s → 4.3s (5.7x) at the 4M-row flagship, + 5.6s → 0.31s (18.3x) no-covariate and 6.5s → 1.2s (5.5x) dr at 2M rows, 4.2s → 1.5s (2.7x) + survey-PSU; peak memory flat or lower. Point estimates are bit-identical (the bootstrap + never feeds them); bootstrap SEs/CIs/p-values match the previous loop to BLAS + reassociation (≤ ~1e-15 relative — far below bootstrap Monte-Carlo error), the same + numerics posture as the existing draw-axis chunking. EfficientDiD's unweighted path folds + its `1/n` prefactor into the influence column (`W @ (eif/n)` instead of `(W @ eif)/n`) — + the same reassociation-level equivalence. - **CallawaySantAnna aggregation SE assembly rewritten to O(n_units)** (also inherited by `StaggeredTripleDifference` aggregation). The combined influence-function assembly behind simple/event-study/group aggregated SEs — previously ~56-85% of analytical fit time at diff --git a/TODO.md b/TODO.md index f0c2ebca..09ef8439 100644 --- a/TODO.md +++ b/TODO.md @@ -47,6 +47,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| +| `EfficientDiD` analytical stage has a pathological scaling hot loop: 339s for ONE fit at just 2k units × 20 periods with 5 covariates (15+ CPU-min at 10k units; memory pressure at 100k), while CS handles 100k units in ~2s. Process sampling attributes it to a broadcasted NumPy subtract/multiply inside a Python-level loop (EIF/weights construction). Profile properly and vectorize; until then EfficientDiD is impractical beyond a few thousand units. Found while benchmarking the Phase 2 bootstrap rewrite (its bootstrap stage is 0.018s at 2k units — the analytical stage is the whole problem). | `efficient_did.py` / `efficient_did_weights.py` | CS-scaling | Heavy | Medium | | `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 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 | diff --git a/diff_diff/bootstrap_chunking.py b/diff_diff/bootstrap_chunking.py index 5cd597dc..0b873ed1 100644 --- a/diff_diff/bootstrap_chunking.py +++ b/diff_diff/bootstrap_chunking.py @@ -23,11 +23,27 @@ - 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. + +The *influence* side can be tiled too. When a fit needs many perturbation +columns (per-cell IFs, the overall combined IF, per-event-time combined IFs), +computing them one column at a time degenerates into per-column slicing/GEMV +over the weight blocks -- memory-bandwidth-bound, not FLOP-bound. Instead, +:func:`tiled_if_matmul` scatters the columns into an ``(n_units, tile_cols)`` +buffer and runs one BLAS GEMM per (weight block, column tile). The column axis +is tiled under a byte cap, and each column tile makes its own full pass over +the weight stream via :class:`ReplayableWeightStream`, which snapshots the +generator state at construction and restores it per pass -- so every pass sees +the bit-identical weight stream on both backends (Rust re-draws the same +``base_seed`` from the restored state; the NumPy fallback replays the stream +itself; the survey stratified branch regenerates its full PSU matrix from the +restored state). GEMM results match the per-column GEMV path to within BLAS +reassociation, same as the draw-axis chunking above. """ from __future__ import annotations -from typing import Iterator, Optional, Tuple +from collections import abc +from typing import Callable, Iterable, Iterator, List, Optional, Sequence, Tuple, Union import numpy as np @@ -183,3 +199,134 @@ def _generated() -> Iterator[Tuple[int, np.ndarray]]: yield chunk_start, block return psu_ids, _generated() + + +# Byte ceiling for the ``(n_units, tile_cols)`` float64 scatter buffer used by +# ``tiled_if_matmul``. Sibling of ``_TARGET_BLOCK_BYTES``; together they bound +# the kernel's live intermediates at ~384 MB regardless of n_units or the +# number of perturbation columns. Read at CALL time inside ``tiled_if_matmul`` +# (never bound as a def-time default) so tests can monkeypatch it. +_TARGET_TILE_BYTES = 128 * 1024 * 1024 + +# One perturbation column of ``tiled_if_matmul``'s influence matrix: a list of +# (index, values) contributions. ``index=None`` assigns the full-width dense +# ``values``; otherwise ``values`` is scattered at ``index``. A column may also +# be a zero-argument callable returning such a list, materialized only when its +# tile is filled (keeps peak memory at one O(n_units) temporary for callers +# whose columns are computed on the fly). +IFColumn = List[Tuple[Optional[np.ndarray], np.ndarray]] +IFColumnSpec = Union[IFColumn, Callable[[], IFColumn]] + + +class ReplayableWeightStream: + """Multiplier-weight block stream that can be re-iterated bit-identically. + + Wraps a weight-block iterator factory together with the ``Generator`` that + feeds it, snapshotting the generator state at construction. Every + ``__iter__`` restores the snapshot before invoking the factory, so each + pass yields the exact same ``(chunk_start, block)`` sequence: + + - **Rust backend**: ``iter_weight_blocks`` draws its per-run ``base_seed`` + from the rng lazily at the first ``next()``; the restored state + reproduces the same base seed, and rows are seeded absolutely from it. + - **NumPy fallback**: generation consumes the rng stream directly; the + restored state replays the identical stream. + - **Survey stratified branch**: ``iter_survey_multiplier_weight_blocks`` + generates its full PSU matrix eagerly at factory-call time -- from the + restored state, so per-pass recreation is also bit-identical. + + The snapshot MUST precede any rng-consuming call: construct this object + before probing generators that draw at call time (the survey stratified + branch draws eagerly; resolve ``psu_ids`` from the design's ``psu`` array + directly instead of via a probe call). After the final pass the rng ends + exactly where a single pass would leave it, so downstream consumers of the + same generator are unaffected by the number of passes. + """ + + def __init__( + self, + make_iter: Callable[[np.random.Generator], Iterator[Tuple[int, np.ndarray]]], + rng: np.random.Generator, + ) -> None: + self._make_iter = make_iter + self._rng = rng + self._state = rng.bit_generator.state + + def __iter__(self) -> Iterator[Tuple[int, np.ndarray]]: + self._rng.bit_generator.state = self._state + return self._make_iter(self._rng) + + +def tiled_if_matmul( + weight_stream: Iterable[Tuple[int, np.ndarray]], + n_bootstrap: int, + n_units: int, + columns: Sequence[IFColumnSpec], + *, + tile_bytes: Optional[int] = None, +) -> np.ndarray: + """Compute ``W @ M`` for the multiplier bootstrap without materializing + ``W`` (streamed in row blocks) or ``M`` (scattered in column tiles). + + ``W`` is the ``(n_bootstrap, n_units)`` weight matrix yielded by + ``weight_stream`` as ``(chunk_start, block)`` row blocks; ``M`` is the + ``(n_units, len(columns))`` influence matrix described by ``columns`` + (see :data:`IFColumnSpec`). Returns the dense ``(n_bootstrap, n_cols)`` + product. Columns are tiled under ``tile_bytes`` (default: module attribute + ``_TARGET_TILE_BYTES``, resolved at call time); each tile is scattered + once into an F-order buffer and then consumed by one BLAS GEMM per weight + block. With more than one tile, ``weight_stream`` is iterated once per + tile and therefore must be re-iterable with an identical stream -- + :class:`ReplayableWeightStream` provides exactly that. + + Scatter contract: contributions are written by plain assignment, so index + arrays must be unique within each contribution and the contributions of a + single column must be mutually disjoint (duplicate or overlapping indices + keep the last write instead of summing like ``W[:, idx] @ values`` would). + All in-repo callers satisfy this by construction (treated/control unit + sets are disjoint; dense columns are single-contribution). + + A non-finite influence value poisons only its own output column (every + output element is an independent dot product), matching the per-column + GEMV behavior this kernel replaces. Results match the per-column GEMV + path to within BLAS reassociation (~1 ULP), not bit-for-bit. + """ + n_cols = len(columns) + out = np.empty((n_bootstrap, n_cols)) + if n_cols == 0 or n_bootstrap == 0: + return out + if tile_bytes is None: + tile_bytes = _TARGET_TILE_BYTES + tile_cols = int(max(1, min(n_cols, tile_bytes // (max(1, n_units) * 8)))) + + # A single-pass iterator silently exhausts after the first tile, leaving + # uninitialized rows in `out` -- fail loudly instead. isinstance (not an + # iter() probe) so re-iterable streams see no extra pass. + if tile_cols < n_cols and isinstance(weight_stream, abc.Iterator): + raise ValueError( + "tiled_if_matmul needs more than one column tile but weight_stream " + "is a single-pass iterator; pass a re-iterable stream " + "(ReplayableWeightStream) so each tile can replay the weights." + ) + + buf = np.zeros((n_units, tile_cols), order="F") + for lo in range(0, n_cols, tile_cols): + hi = min(lo + tile_cols, n_cols) + width = hi - lo + if lo > 0: + buf[:, :width] = 0.0 + for c in range(lo, hi): + spec = columns[c] + if callable(spec): + spec = spec() + col = buf[:, c - lo] + for idx, values in spec: + if idx is None: + np.copyto(col, values) + else: + col[idx] = values + tile = buf[:, :width] + with np.errstate(divide="ignore", invalid="ignore", over="ignore"): + for chunk_start, w_block in weight_stream: + out[chunk_start : chunk_start + w_block.shape[0], lo:hi] = w_block @ tile + return out diff --git a/diff_diff/efficient_did_bootstrap.py b/diff_diff/efficient_did_bootstrap.py index 6aed6a12..5e4cca75 100644 --- a/diff_diff/efficient_did_bootstrap.py +++ b/diff_diff/efficient_did_bootstrap.py @@ -13,9 +13,11 @@ import numpy as np from diff_diff.bootstrap_chunking import ( + ReplayableWeightStream, compute_block_size, iter_survey_multiplier_weight_blocks, iter_weight_blocks, + tiled_if_matmul, ) from diff_diff.bootstrap_utils import ( compute_effect_bootstrap_stats as _compute_effect_bootstrap_stats_func, @@ -97,7 +99,6 @@ def _run_multiplier_bootstrap( rng = np.random.default_rng(self.seed) gt_pairs = list(group_time_effects.keys()) - n_gt = len(gt_pairs) # Original ATTs (independent of the draws; referenced per block below). original_atts = np.array([group_time_effects[gt]["effect"] for gt in gt_pairs]) @@ -120,13 +121,15 @@ def _run_multiplier_bootstrap( # at a time (unstratified designs tile the generation; stratified # designs have few PSUs and fall back to full generation + slicing). _block_size = compute_block_size(n_units, self.n_bootstrap) - psu_ids, _psu_blocks = iter_survey_multiplier_weight_blocks( - self.n_bootstrap, - resolved_survey, - self.bootstrap_weights, - rng, - block_size=_block_size, - ) + # Resolve psu_ids WITHOUT calling the generator: the stratified + # branch draws from the rng eagerly at call time, and the + # replayable stream below must snapshot the rng state before any + # draw. Duplicates the rng-free resolution both generator branches + # use (np.unique / np.arange). + if resolved_survey.psu is not None: + psu_ids = np.unique(resolved_survey.psu) + else: + psu_ids = np.arange(len(resolved_survey.weights)) # Single-cluster (G<2) survey-PSU multiplier bootstrap collapses # to constant multiplier draws → BLAS roundoff produces ≈0 # variance (NOT NaN). Downstream zero-SE guards check exact 0 and @@ -163,29 +166,51 @@ def _run_multiplier_bootstrap( np.array_equal(unit_to_psu_col, np.arange(n_units)) ) - 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] + # Factory recreating the PSU generation + unit expansion per pass. + def _make_weight_iter( + rng_: np.random.Generator, + ) -> Iterator[Tuple[int, np.ndarray]]: + _, _psu_blocks = iter_survey_multiplier_weight_blocks( + self.n_bootstrap, + resolved_survey, + self.bootstrap_weights, + rng_, + block_size=_block_size, + ) + + def _expanded() -> 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] + + return _expanded() - weight_blocks: Iterator[Tuple[int, np.ndarray]] = _weight_blocks() elif cluster_indices is not None and n_clusters is not None: # Cluster-level weights, expanded to unit level per block via the # helper's expand_index (block[:, cluster_indices]). - weight_blocks = iter_weight_blocks( - self.n_bootstrap, - n_clusters, - self.bootstrap_weights, - rng, - expand_index=cluster_indices, - ) + def _make_weight_iter( + rng_: np.random.Generator, + ) -> Iterator[Tuple[int, np.ndarray]]: + return iter_weight_blocks( + self.n_bootstrap, + n_clusters, + self.bootstrap_weights, + rng_, + expand_index=cluster_indices, + ) + else: # Standard unit-level weights, generated one row-block at a time. - weight_blocks = iter_weight_blocks( - self.n_bootstrap, n_units, self.bootstrap_weights, rng - ) + def _make_weight_iter( + rng_: np.random.Generator, + ) -> Iterator[Tuple[int, np.ndarray]]: + return iter_weight_blocks(self.n_bootstrap, n_units, self.bootstrap_weights, rng_) + + # Re-iterable stream: each column tile of the fused perturbation GEMM + # below makes its own full pass over the bit-identical weight stream. + weight_stream = ReplayableWeightStream(_make_weight_iter, rng) # eif SCALING is a SEPARATE axis from the weight PATH: it is keyed on # unit_level_weights (set whenever a SurveyDesign was passed — including a @@ -196,24 +221,32 @@ def _weight_blocks() -> Iterator[Tuple[int, np.ndarray]]: _has_unit_weights = unit_level_weights is not None _total_w = float(np.sum(unit_level_weights)) if _has_unit_weights else 1.0 - # Pre-allocate the small (n_bootstrap, n_gt) output; only this persists. - # Each weight block fills its rows of every column, then is discarded — - # peak memory is capped at O(block x n_units). The per-(g,t) scaled EIF is - # recomputed per block as a single O(n_units) temporary (not cached across - # all n_gt cells), so the perturbation adds no O(n_gt x n_units) allocation - # that would erode the memory win on weighted panels. The aggregations - # below (overall, event study, group) re-aggregate these columns and never - # touch the weight matrix. - bootstrap_atts = np.empty((self.n_bootstrap, n_gt)) - for _chunk_start, _w_block in weight_blocks: - _rows = slice(_chunk_start, _chunk_start + _w_block.shape[0]) - for j, gt in enumerate(gt_pairs): - with np.errstate(divide="ignore", invalid="ignore", over="ignore"): - if _has_unit_weights: - perturbation = _w_block @ (unit_level_weights * eif_by_gt[gt] / _total_w) - else: - perturbation = (_w_block @ eif_by_gt[gt]) / n_units - bootstrap_atts[_rows, j] = original_atts[j] + perturbation + # Fused perturbation GEMM over the replayable weight stream, one column + # per (g,t). Columns are LAZY callables materializing each scaled EIF + # only when its tile is filled — one O(n_units) temporary at a time, so + # the perturbation still adds no O(n_gt x n_units) allocation that + # would erode the memory win on weighted panels (tiles are capped by + # _TARGET_TILE_BYTES). The unweighted path folds its 1/n prefactor into + # the column (W @ (eif/n) instead of (W @ eif)/n) — a pure BLAS + # reassociation-level change. The aggregations below (overall, event + # study, group) re-aggregate these columns and never touch the weight + # matrix. + def _scaled_eif_column(gt: Tuple[Any, Any]): + def _make() -> List[Tuple[Optional[np.ndarray], np.ndarray]]: + if _has_unit_weights: + return [(None, unit_level_weights * eif_by_gt[gt] / _total_w)] + return [(None, eif_by_gt[gt] / n_units)] + + return _make + + perturbations = tiled_if_matmul( + weight_stream, + self.n_bootstrap, + n_units, + [_scaled_eif_column(gt) for gt in gt_pairs], + ) + with np.errstate(divide="ignore", invalid="ignore", over="ignore"): + bootstrap_atts = original_atts[None, :] + perturbations # Post-treatment mask — also exclude NaN effects post_mask = np.array( diff --git a/diff_diff/staggered_bootstrap.py b/diff_diff/staggered_bootstrap.py index 783b97c6..2c3ce962 100644 --- a/diff_diff/staggered_bootstrap.py +++ b/diff_diff/staggered_bootstrap.py @@ -13,9 +13,11 @@ import numpy as np from diff_diff.bootstrap_chunking import ( + ReplayableWeightStream, compute_block_size, iter_survey_multiplier_weight_blocks, iter_weight_blocks, + tiled_if_matmul, ) from diff_diff.bootstrap_utils import ( compute_bootstrap_pvalue as _compute_bootstrap_pvalue_func, @@ -354,13 +356,17 @@ def _run_multiplier_bootstrap( # 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, - ) + # Resolve psu_ids WITHOUT calling the generator: the stratified + # branch of iter_survey_multiplier_weight_blocks draws from the rng + # eagerly at call time, and the replayable stream below must + # snapshot the rng state before any draw. This duplicates the + # rng-free resolution both generator branches use (np.unique / + # np.arange), so the column order of the generated PSU matrix + # matches unit_to_psu_col. + if resolved_survey_unit.psu is not None: + psu_ids = np.unique(resolved_survey_unit.psu) + else: + psu_ids = np.arange(len(resolved_survey_unit.weights)) if len(psu_ids) < 2: import warnings as _warnings @@ -395,22 +401,40 @@ def _run_multiplier_bootstrap( np.array_equal(unit_to_psu_col, np.arange(n_units)) ) - # 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] + # Factory recreating the PSU generation + unit-level expansion per + # pass; the full (n_bootstrap, n_units) expansion is never + # materialized at once. + def _make_weight_iter( + rng_: np.random.Generator, + ) -> Iterator[Tuple[int, np.ndarray]]: + _, _psu_blocks = iter_survey_multiplier_weight_blocks( + self.n_bootstrap, + resolved_survey_unit, + self.bootstrap_weights, + rng_, + block_size=_block_size, + ) + + def _expanded() -> 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] + + return _expanded() - 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 - ) + def _make_weight_iter( + rng_: np.random.Generator, + ) -> Iterator[Tuple[int, np.ndarray]]: + return iter_weight_blocks(self.n_bootstrap, n_units, self.bootstrap_weights, rng_) + + # Re-iterable stream: each column tile of the fused perturbation GEMM + # below makes its own full pass over the bit-identical weight stream. + weight_stream = ReplayableWeightStream(_make_weight_iter, rng) # Pre-compute the overall combined IF once (reused across every block). # None exactly when the overall aggregation is skipped. @@ -433,68 +457,65 @@ def _weight_blocks() -> Iterator[Tuple[int, np.ndarray]]: n_global_units=n_units, ) - # 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 = {e: np.empty(self.n_bootstrap) for e in rel_periods} 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 = {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_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"] - ) + # Fused perturbation columns: [per-cell IFs | overall combined IF | + # per-event-time combined IFs]. One column-tiled GEMM over the + # replayable weight stream replaces the former per-cell + # ``W[:, idx] @ inf`` slicing loop, which was memory-bandwidth-bound + # (two fancy-index copies of the weight block per cell). The weight + # stream is bit-identical to the per-block path; the BLAS reductions + # may reassociate, so statistics match to within ~1 ULP (far below + # bootstrap Monte-Carlo error), not bit-for-bit. Treated/control index + # arrays are disjoint per cell, satisfying the kernel's + # assignment-scatter contract. + columns: List[Any] = [ + [ + (gt_treated_indices[j], gt_treated_inf[j]), + (gt_control_indices[j], gt_control_inf[j]), + ] + for j in range(n_gt) + ] + overall_col = -1 + if overall_combined_if is not None: + overall_col = len(columns) + columns.append([(None, overall_combined_if)]) + es_col0 = len(columns) + for e in rel_periods: + columns.append([(None, event_study_info[e]["combined_if"])]) + + perturbations = tiled_if_matmul(weight_stream, self.n_bootstrap, n_units, columns) + + # Reconstruct the bootstrap draws (small, n_bootstrap-sized arrays). + with np.errstate(divide="ignore", invalid="ignore", over="ignore"): + bootstrap_atts_gt = original_atts[None, :] + perturbations[:, :n_gt] + if skip_overall_aggregation: + bootstrap_overall = np.full(self.n_bootstrap, np.nan) + else: + bootstrap_overall = original_overall + perturbations[:, overall_col] + + bootstrap_event_study: Optional[Dict[int, np.ndarray]] = None + if event_study_info is not None: + bootstrap_event_study = { + e: event_study_info[e]["effect"] + perturbations[:, es_col0 + k] + for k, e in enumerate(rel_periods) + } + + # Group aggregation: fixed-weight re-aggregation of the completed + # perturbed cell draws (matches at reassociation level). + bootstrap_group: Optional[Dict[Any, np.ndarray]] = None + if group_agg_info is not None: + bootstrap_group = { + g: bootstrap_atts_gt[:, group_agg_info[g]["gt_indices"]] + @ group_agg_info[g]["weights"] + for g in group_list + } # 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 b9fcfbcb..b6dba61f 100644 --- a/docs/doc-deps.yaml +++ b/docs/doc-deps.yaml @@ -949,7 +949,7 @@ sources: - 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." + note: "Memory-bounded multiplier-bootstrap machinery: draw-block chunking of weight generation/consumption (shared by CallawaySantAnna, EfficientDiD, and HAD), plus the fused column-tiled scatter-GEMM perturbation kernel (tiled_if_matmul + ReplayableWeightStream, v3.7 Phase 2 — used by CallawaySantAnna and EfficientDiD; HAD consumes whole-matrix GEMMs directly and does not use the tiled kernel). Implementation detail — numerically equivalent (~1 ULP, BLAS reassociation) to the per-column GEMV path; weight stream bit-identical, replayable per column tile." diff_diff/prep.py: drift_risk: low diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index a7d990bd..39a80873 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -468,7 +468,7 @@ methodology is introduced. - `classical`, `hc2`, `hc2_bm`, `conley` — REJECTED at `__init__`. The rejection is **library-architectural, not paper-prescribed**: analytical-sandwich variance families (`classical`, `hc2`, `hc2_bm`) are defined on a single regression's hat matrix, and CS's per-(g,t) doubly-robust / IPW / outcome-regression structure has no equivalent single design matrix to compute hat-matrix leverage or Bell-McCaffrey Satterthwaite DOF on. Spatial-HAC (`conley`) likewise has no defined composition with per-unit IF aggregation today. See ["IF-based variance estimators vs analytical-sandwich estimators"](#if-based-variance-estimators-vs-analytical-sandwich-estimators) above for the structural taxonomy. *Cluster wiring:* -Prior to the bare-`cluster=` wiring fix, `CallawaySantAnna(cluster="X")` was a silent no-op — the parameter was stored at `__init__` but never consumed in the fit / aggregator / bootstrap pipeline (users got per-unit IF variance silently, even when they explicitly set `cluster="state"`). The fix synthesizes a minimal `SurveyDesign(psu=X, weight_type="pweight")` when bare `cluster=` is set without an explicit survey design, threading the synthesized PSU through the existing `_compute_stratified_psu_meat` aggregator (`staggered_aggregation.py:735-749`) and PSU-level multiplier bootstrap (`staggered_bootstrap.py:323-347`). Three-branch wiring at `staggered.py:~1500`: +Prior to the bare-`cluster=` wiring fix, `CallawaySantAnna(cluster="X")` was a silent no-op — the parameter was stored at `__init__` but never consumed in the fit / aggregator / bootstrap pipeline (users got per-unit IF variance silently, even when they explicitly set `cluster="state"`). The fix synthesizes a minimal `SurveyDesign(psu=X, weight_type="pweight")` when bare `cluster=` is set without an explicit survey design, threading the synthesized PSU through the existing `_compute_stratified_psu_meat` aggregator (`staggered_aggregation.py:735-749`) and PSU-level multiplier bootstrap (`staggered_bootstrap.py:334-437`). Three-branch wiring at `staggered.py:~1500`: 1. Bare `cluster=X` + no `survey_design` → synthesize `SurveyDesign(psu=X)`; refit `_resolve_survey_for_fit` on synthetic; `effective_survey_design = synthetic` so `_validate_unit_constant_survey` runs on it (preventing first-value-wins collapse for movers on panel data). 2. `survey_design` without PSU + `cluster=X` → call `_inject_cluster_as_psu(resolved_survey, cluster_ids)`. 3. `survey_design` with PSU + `cluster=X` → PSU wins; `_resolve_effective_cluster` emits `UserWarning` if partitions differ. diff --git a/docs/performance-plan.md b/docs/performance-plan.md index ad8eb7fb..202655e4 100644 --- a/docs/performance-plan.md +++ b/docs/performance-plan.md @@ -183,6 +183,54 @@ precision; analytical SEs to 5 significant digits; bootstrap SEs within ~1% (Mon --- +## CS multiplier-bootstrap fused tiled-GEMM rewrite (v3.7 Phase 2, 2026-07) + +Stage-level instrumentation of `_run_multiplier_bootstrap` at the 40p × 10c × 100k-unit +flagship (390 cells, 72 event times, 999 draws; 20.9s of a 24.4s fit) split the loop as: + +| stage | time | share | +|---|---:|---:| +| per-cell fancy-index copies `W[:, idx_t]`, `W[:, idx_c]` | 20.45s | **95%** | +| per-cell GEMVs | 0.73s | 3% | +| event-study GEMVs (72 full-width) | 0.27s | 1% | +| weight generation (rust, 1e8 draws) | 0.04s | 0.2% | + +Not FLOPs, not RNG — ~100GB of cache-hostile gather traffic from slicing the +`(block × n_units)` weight matrix twice per cell (mean cell density 0.33: never-treated +controls appear in every cell). The fix: scatter every perturbation column (per-cell IFs, +overall combined IF, per-event-time combined IFs) into a column-tiled influence matrix +and run one BLAS GEMM per (weight block, column tile) — `bootstrap_chunking.tiled_if_matmul` ++ `ReplayableWeightStream` (each column tile replays the bit-identical weight stream via +RNG state snapshot/restore; byte-capped tiles keep the kernel's live intermediates bounded +regardless of n_units or column count). A scipy-sparse cell-column variant was measured and +rejected (27x slower than dense GEMM at 33% density). EfficientDiD's per-cell full-width +GEMV bootstrap routes through the same kernel via lazily materialized scaled-EIF columns. + +Measured (5-rep medians, M4 Max, rust backend + Accelerate, otherwise-idle machine; +BEFORE = pristine `origin/main` worktree via PYTHONPATH, AFTER = frozen branch): + +| Scenario | Before | After | Speedup | Peak RSS | +|---|---:|---:|---:|---| +| 40p × 10c, 5-cov dr, boot-999 (4M rows, 390 cells) | 24.49s | 4.28s | **5.7x** | 2.9→2.9 GB | +| no-cov boot-999 (2M rows) | 5.64s | 0.31s | **18.3x** | 2.0→1.6 GB | +| 5-cov dr boot-999 (2M rows) | 6.49s | 1.19s | **5.5x** | 2.3→1.9 GB | +| survey-PSU 5-cov dr boot-999 (500 PSUs, non-identity map) | 4.22s | 1.54s | **2.7x** | 2.0→1.9 GB | +| RCS 5-cov dr boot-999 (2M obs) | 42.15s | 40.07s | 1.1x | 3.3→3.2 GB | + +ATT deltas exactly 0 on every scenario (the bootstrap never feeds point estimates); +bootstrap SE relative deltas 3e-16 to 1.2e-15 (BLAS reassociation only). The RCS scenario +barely moves because its fit time is the analytical repeated-cross-sections path, not the +bootstrap. EfficientDiD's bootstrap stage went 0.018s → 0.012s but is invisible inside its +fits: EfficientDiD's **analytical** stage has a pre-existing pathological hot loop +(339s at just 2k units × 20p with 5 covariates; sampled to a broadcasted subtract/multiply +in a Python-level loop) — tracked in TODO.md as its own item, out of scope here. + +Remaining CS bottlenecks (updated forecast): `_precompute_structures` pandas +groupby+pivot on analytical no-cov fits, and per-cell DR/IPW nuisance solvers at high +covariate counts (IRLS `lstsq`, rank-guard QR — Phase 3 candidate). + +--- + ## Memory scaling at the millions-of-units tail (v3.x, June 2026) At the scale where the dense working arrays - not compute - are the binding constraint diff --git a/tests/test_bootstrap_chunking.py b/tests/test_bootstrap_chunking.py index 6cc0009e..db5316fd 100644 --- a/tests/test_bootstrap_chunking.py +++ b/tests/test_bootstrap_chunking.py @@ -19,9 +19,11 @@ from diff_diff import CallawaySantAnna, EfficientDiD, HeterogeneousAdoptionDiD from diff_diff.bootstrap_chunking import ( + ReplayableWeightStream, compute_block_size, iter_survey_multiplier_weight_blocks, iter_weight_blocks, + tiled_if_matmul, ) from diff_diff.bootstrap_utils import ( generate_bootstrap_weights_batch, @@ -120,6 +122,186 @@ def test_single_block_matches_legacy_generator(self, weight_type): np.testing.assert_array_equal(chunked, legacy) +class _CountingStream: + """Re-iterable wrapper that counts how many passes were made.""" + + def __init__(self, stream): + self.stream = stream + self.n_passes = 0 + + def __iter__(self): + self.n_passes += 1 + return iter(self.stream) + + +class TestTiledIFMatmul: + """Oracle + replay + tiling contracts for the fused scatter-GEMM kernel. + + ``tiled_if_matmul`` replaces per-column ``W[:, idx] @ values`` GEMVs with + one GEMM per (weight block, column tile). The oracle below computes the OLD + per-column semantics on the materialized weight matrix and requires the + kernel to match at reassociation tolerance. Replay determinism of + ``ReplayableWeightStream`` is what makes multi-tile passes sound. + """ + + N_BOOT, N_UNITS, BLOCK = 61, 220, 17 + + def _stream(self, weight_type="rademacher", seed=42): + rng = np.random.default_rng(seed) + return ReplayableWeightStream( + lambda r: iter_weight_blocks( + self.N_BOOT, self.N_UNITS, weight_type, r, block_size=self.BLOCK + ), + rng, + ) + + def _weight_matrix(self, weight_type="rademacher", seed=42): + _, mat = _stack(self.N_BOOT, self.N_UNITS, weight_type, seed, block_size=self.BLOCK) + return mat + + def _columns(self): + rng = np.random.default_rng(0) + n = self.N_UNITS + sparse_idx = rng.choice(n, size=40, replace=False) + sparse_vals = rng.standard_normal(40) + dense_vals = rng.standard_normal(n) + # two mutually disjoint contributions in one column (the CS cell shape) + treated_idx = np.arange(0, 30) + control_idx = np.arange(50, 130) + treated_vals = rng.standard_normal(30) + control_vals = rng.standard_normal(80) + lazy_vals = rng.standard_normal(n) + nan_idx = np.array([3, 7]) + columns = [ + [(sparse_idx, sparse_vals)], + [(None, dense_vals)], + [(treated_idx, treated_vals), (control_idx, control_vals)], + lambda: [(None, lazy_vals)], + [(nan_idx, np.array([np.nan, np.nan]))], + [], # all-zero column + ] + return columns + + @staticmethod + def _old_semantics(W, columns): + """The pre-rewrite per-column GEMV semantics: sum of W[:, idx] @ vals.""" + out = np.zeros((W.shape[0], len(columns))) + for c, spec in enumerate(columns): + if callable(spec): + spec = spec() + for idx, values in spec: + if idx is None: + out[:, c] += W @ values + else: + out[:, c] += W[:, idx] @ values + return out + + @pytest.mark.parametrize("weight_type", WEIGHT_TYPES) + def test_oracle_matches_old_gemv_semantics(self, weight_type): + columns = self._columns() + W = self._weight_matrix(weight_type) + ref = self._old_semantics(W, columns) + out = tiled_if_matmul(self._stream(weight_type), self.N_BOOT, self.N_UNITS, columns) + finite_cols = [0, 1, 2, 3, 5] + np.testing.assert_allclose(out[:, finite_cols], ref[:, finite_cols], rtol=1e-12, atol=1e-14) + # NaN influence values poison exactly their own column + assert np.all(np.isnan(out[:, 4])) + assert not np.any(np.isnan(out[:, finite_cols])) + # the empty column is exactly zero + np.testing.assert_array_equal(out[:, 5], np.zeros(self.N_BOOT)) + + def test_single_tile_equals_multi_tile(self): + columns = self._columns() + one = tiled_if_matmul(self._stream(), self.N_BOOT, self.N_UNITS, columns, tile_bytes=2**62) + # 2 columns per tile -> 3 tiles + tiny = tiled_if_matmul( + self._stream(), + self.N_BOOT, + self.N_UNITS, + columns, + tile_bytes=2 * self.N_UNITS * 8, + ) + np.testing.assert_allclose(tiny, one, rtol=1e-12, atol=1e-14, equal_nan=True) + + def test_multi_tile_actually_replays_the_stream(self): + # Guards against a silent single-tile no-op: the multi-tile arm must + # iterate the weight stream once per tile, and still match the oracle. + columns = self._columns() + counting = _CountingStream(self._stream()) + out = tiled_if_matmul( + counting, + self.N_BOOT, + self.N_UNITS, + columns, + tile_bytes=2 * self.N_UNITS * 8, + ) + assert counting.n_passes == 3 # 6 columns, 2 per tile + ref = self._old_semantics(self._weight_matrix(), columns) + np.testing.assert_allclose(out[:, :4], ref[:, :4], rtol=1e-12, atol=1e-14) + + def test_single_pass_iterator_raises_when_tiled(self): + # A plain generator would silently exhaust after the first tile. + gen = iter_weight_blocks(self.N_BOOT, self.N_UNITS, "rademacher", np.random.default_rng(1)) + with pytest.raises(ValueError, match="single-pass"): + tiled_if_matmul( + gen, + self.N_BOOT, + self.N_UNITS, + self._columns(), + tile_bytes=2 * self.N_UNITS * 8, + ) + + def test_empty_inputs(self): + out = tiled_if_matmul(self._stream(), self.N_BOOT, self.N_UNITS, []) + assert out.shape == (self.N_BOOT, 0) + + @pytest.mark.parametrize("weight_type", WEIGHT_TYPES) + def test_replay_determinism(self, weight_type): + # Two passes over a ReplayableWeightStream yield byte-identical blocks + # on whichever backend is active, and the rng ends exactly where one + # plain pass would leave it. + rng = np.random.default_rng(7) + stream = ReplayableWeightStream( + lambda r: iter_weight_blocks( + self.N_BOOT, self.N_UNITS, weight_type, r, block_size=self.BLOCK + ), + rng, + ) + first = [(cs, b.copy()) for cs, b in stream] + second = [(cs, b.copy()) for cs, b in stream] + assert [cs for cs, _ in first] == [cs for cs, _ in second] + for (_, a), (_, b) in zip(first, second): + np.testing.assert_array_equal(a, b) + + # single plain pass from an identically-seeded rng: same end state + rng_ref = np.random.default_rng(7) + for _ in iter_weight_blocks( + self.N_BOOT, self.N_UNITS, weight_type, rng_ref, block_size=self.BLOCK + ): + pass + assert rng.bit_generator.state == rng_ref.bit_generator.state + + def test_survey_stratified_replay_determinism(self): + # The stratified survey branch generates its FULL PSU matrix eagerly at + # factory-call time -- the riskiest replay mechanism. Recreating the + # factory from restored state must reproduce it bit-identically. + psu = np.arange(6) + strata = np.array([0, 0, 0, 1, 1, 1]) + design = _design(psu=psu, strata=strata, weights=np.ones(6)) + rng = np.random.default_rng(11) + + def make_iter(r): + _, blocks = iter_survey_multiplier_weight_blocks( + 50, design, "rademacher", r, block_size=9 + ) + return blocks + + stream = ReplayableWeightStream(make_iter, rng) + first = np.vstack([b.copy() for _, b in stream]) + second = np.vstack([b.copy() for _, b in stream]) + np.testing.assert_array_equal(first, second) + + class TestCSBootstrapChunkInvariance: """CallawaySantAnna bootstrap output is invariant to the chunk size. @@ -270,6 +452,162 @@ def test_nonidentity_cluster_chunks_match_single(self, monkeypatch): np.testing.assert_allclose(t[col].to_numpy(), b[col].to_numpy(), rtol=1e-10, atol=1e-12) +class TestCSBootstrapTileInvariance: + """CallawaySantAnna bootstrap output is invariant to the COLUMN tile size. + + The fused perturbation GEMM tiles the influence-column axis under + ``_TARGET_TILE_BYTES``; each tile replays the weight stream via + ``ReplayableWeightStream``. Forcing one column per tile makes every fit + replay the stream n_cols times -- results must match the default + single-tile fit at reassociation tolerance on every weight path (unit, + non-identity cluster, survey-PSU, and STRATIFIED survey, whose per-pass + factory recreation regenerates the full PSU matrix from restored state). + """ + + @staticmethod + def _force_one_column_tiles(monkeypatch): + # tiled_if_matmul resolves the cap at call time, so patching the module + # attribute is effective (no def-time default binding). + monkeypatch.setattr("diff_diff.bootstrap_chunking._TARGET_TILE_BYTES", 1) + + def test_unit_path_tiles_match_single_tile(self, monkeypatch): + fit = TestCSBootstrapChunkInvariance()._fit + base = fit() + self._force_one_column_tiles(monkeypatch) + tiled = fit() + + assert tiled.overall_se == pytest.approx(base.overall_se, rel=1e-10, abs=1e-12) + assert tiled.cband_crit_value == pytest.approx(base.cband_crit_value, rel=1e-10, abs=1e-12) + assert tiled.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 = tiled.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) + for level in ("event_study", "group"): + bl = base.to_dataframe(level=level).reset_index(drop=True) + tl = tiled.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"] + 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_path_tiles_match_single_tile(self, monkeypatch): + # Non-identity PSU fan-out: the per-pass expansion copy must replay too. + fit = TestCSBootstrapChunkInvariance()._fit_clustered + base = fit() + self._force_one_column_tiles(monkeypatch) + tiled = fit() + + assert tiled.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 = tiled.to_dataframe().sort_values(["group", "time"]).reset_index(drop=True) + np.testing.assert_allclose(t["se"].to_numpy(), b["se"].to_numpy(), rtol=1e-10, atol=1e-12) + + @staticmethod + def _survey_panel(stratified): + rng = np.random.default_rng(3) + n_psu, units_per_psu, nt = 12, 4, 8 + nu = n_psu * units_per_psu + 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) + psu = np.repeat(np.repeat(np.arange(n_psu), units_per_psu), nt) + w = np.repeat(np.exp(rng.normal(0, 0.3, nu)), nt) + df = pd.DataFrame( + {"unit": units, "period": periods, "y": y, "first_treat": ft, "psu": psu, "w": w} + ) + if stratified: + # 3 strata x 4 PSUs -> routes the STRATIFIED (eager full-generation) + # branch of iter_survey_multiplier_weight_blocks + df["stratum"] = df["psu"] % 3 + return df + + def _fit_survey(self, stratified): + from diff_diff import SurveyDesign + + df = self._survey_panel(stratified) + design = ( + SurveyDesign(weights="w", strata="stratum", psu="psu") + if stratified + else SurveyDesign(weights="w", psu="psu") + ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return CallawaySantAnna( + control_group="never_treated", + estimation_method="dr", + n_bootstrap=200, + seed=42, + ).fit( + df, + outcome="y", + unit="unit", + time="period", + first_treat="first_treat", + aggregate="all", + survey_design=design, + ) + + @pytest.mark.parametrize("stratified", [False, True], ids=["psu", "stratified"]) + def test_survey_path_tiles_match_single_tile(self, monkeypatch, stratified): + base = self._fit_survey(stratified) + self._force_one_column_tiles(monkeypatch) + tiled = self._fit_survey(stratified) + + assert tiled.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 = tiled.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) + + @staticmethod + def _rcs_frame(): + # Repeated cross sections: every row one observation with a unique unit + # id; treatment timing assigned at the observation level (panel=False + # contract). The bootstrap then perturbs observation-level IFs through + # the same tiled kernel. + rng = np.random.default_rng(5) + n_obs, n_periods = 600, 6 + t = rng.integers(1, n_periods + 1, size=n_obs) + never = rng.random(n_obs) < 0.4 + g = np.where(never, 0, rng.choice([3, 5], size=n_obs)) + treated = (g > 0) & (t >= g) + y = 0.2 * t + 2.0 * treated + rng.standard_normal(n_obs) + return pd.DataFrame({"unit": np.arange(n_obs), "period": t, "y": y, "first_treat": g}) + + def test_rcs_path_tiles_match_single_tile(self, monkeypatch): + # panel=False routes observation-level IFs through the same kernel. + def fit(): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return CallawaySantAnna( + estimation_method="reg", panel=False, n_bootstrap=200, seed=42 + ).fit( + self._rcs_frame(), + "y", + "unit", + "period", + "first_treat", + aggregate="all", + ) + + base = fit() + self._force_one_column_tiles(monkeypatch) + tiled = fit() + + assert tiled.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 = tiled.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) @@ -418,6 +756,46 @@ def test_weights_only_survey_path(self, monkeypatch): self._run(monkeypatch, survey_design=SurveyDesign(weights="w")) +class TestEfficientDiDBootstrapTileInvariance: + """EfficientDiD bootstrap output is invariant to the COLUMN tile size. + + Tile-forced twins of all four ``TestEfficientDiDBootstrapChunkInvariance`` + weight paths: one column per tile makes the fused perturbation GEMM replay + the weight stream once per (g,t) column (unit, cluster/expand_index, + survey-PSU, and weights-only survey -- the last exercising the LAZY + scaled-EIF columns under multi-pass replay). + """ + + def _run(self, monkeypatch, **fit_kwargs): + harness = TestEfficientDiDBootstrapChunkInvariance() + base = harness._fit(**fit_kwargs) + base_ses = harness._ses(base) + assert np.isfinite(base_ses[0]) and np.isfinite(base_ses[1:]).any() + # tiled_if_matmul resolves the cap at call time -> patching the module + # attribute forces one column per tile. + monkeypatch.setattr("diff_diff.bootstrap_chunking._TARGET_TILE_BYTES", 1) + tiled = harness._fit(**fit_kwargs) + np.testing.assert_allclose( + harness._ses(tiled), base_ses, rtol=1e-9, atol=1e-12, equal_nan=True + ) + + def test_unit_path(self, monkeypatch): + self._run(monkeypatch) + + def test_cluster_path(self, monkeypatch): + self._run(monkeypatch, cluster="state") + + def test_survey_psu_path(self, monkeypatch): + from diff_diff.survey import SurveyDesign + + self._run(monkeypatch, survey_design=SurveyDesign(psu="state", weights="w")) + + def test_weights_only_survey_path(self, monkeypatch): + from diff_diff.survey import SurveyDesign + + self._run(monkeypatch, survey_design=SurveyDesign(weights="w")) + + class TestHADBootstrapChunkInvariance: """HAD event-study sup-t bootstrap is invariant to the chunk size. diff --git a/tests/test_staggered.py b/tests/test_staggered.py index 26f8c72a..fbe4cc69 100644 --- a/tests/test_staggered.py +++ b/tests/test_staggered.py @@ -1970,6 +1970,36 @@ def _failing_logit(*args, **kwargs): np.testing.assert_allclose(cell["effect"], ref["effect"], rtol=1e-12) np.testing.assert_allclose(cell["se"], ref["se"], rtol=1e-12) + def test_bootstrap_nan_cell_if_poisons_only_its_own_cell(self, monkeypatch): + """A cell whose stored IF is non-finite (e.g. the #619 rank-0 [1,X] + bread semantics on ipw/dr) must NaN only its OWN bootstrap SE: in the + fused perturbation GEMM each cell column is an independent dot + product, so neighbor cells stay finite. (The overall/aggregate SEs + legitimately consume the poisoned cell and are not asserted here.)""" + from diff_diff.staggered_bootstrap import CallawaySantAnnaBootstrapMixin + + orig = CallawaySantAnnaBootstrapMixin._run_multiplier_bootstrap + poisoned = {} + + def poisoning(self, group_time_effects, influence_func_info, *args, **kwargs): + gt = sorted(influence_func_info)[0] + poisoned["gt"] = gt + info = influence_func_info[gt] + info["treated_inf"] = np.full(len(np.asarray(info["treated_inf"])), np.nan) + return orig(self, group_time_effects, influence_func_info, *args, **kwargs) + + monkeypatch.setattr(CallawaySantAnnaBootstrapMixin, "_run_multiplier_bootstrap", poisoning) + data = generate_staggered_data(seed=42) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + result = CallawaySantAnna(n_bootstrap=99, seed=5).fit( + data, "outcome", "unit", "time", "first_treat" + ) + gt = poisoned["gt"] + assert np.isnan(result.group_time_effects[gt]["se"]) + neighbor_ses = [cell["se"] for key, cell in result.group_time_effects.items() if key != gt] + assert neighbor_ses and np.isfinite(neighbor_ses).all() + def test_underdetermined_control_cell_reg_no_crash(self): """Cells with fewer controls than covariate columns (n_c < k+1) fit a reduced design; the rank-guarded IF bread column-drops and the