perf(efficient-did, had): chunk multiplier-bootstrap weight generation to bound peak memory

CHANGELOG.md

@@ -44,6 +44,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   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.
+- **EfficientDiD and HeterogeneousAdoptionDiD multiplier bootstraps now tile weight generation
+  over draws too, via the same `diff_diff/bootstrap_chunking.py` helper.** Both built the same
+  dense `(n_bootstrap × n_units)` multiplier-weight matrix CallawaySantAnna did — EfficientDiD
+  in its per-`(g,t)` EIF perturbation, HAD in its event-study sup-t band — the dominant
+  allocation at large `n_units` (~40 GB at 5M units × 999 reps). It is now generated and
+  consumed one draw-block at a time, capping peak memory at `O(block × n_units)`, so these
+  estimators reach the same millions-of-units scale as the chunked CallawaySantAnna path. The
+  weight *stream* is bit-identical on both backends; end-to-end bootstrap SEs and sup-t
+  critical values match the un-chunked path to within floating-point reassociation of the BLAS
+  reductions (~1 ULP, far below bootstrap Monte-Carlo error). As with CallawaySantAnna,
+  stratified survey designs (few PSUs) are unchanged — full generation + sliced blocks — with
+  the deferred per-stratum tiling tracked in TODO.md.
 - **`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

TODO.md

@@ -65,7 +65,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 |
+| 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 |
 
 ### Testing / docs
 

diff_diff/efficient_did_bootstrap.py

@@ -8,15 +8,17 @@
 
 import warnings
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, Iterator, List, Optional, Tuple
 
 import numpy as np
 
-from diff_diff.bootstrap_utils import (
-    compute_effect_bootstrap_stats as _compute_effect_bootstrap_stats_func,
+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 as _generate_bootstrap_weights_batch,
+    compute_effect_bootstrap_stats as _compute_effect_bootstrap_stats_func,
 )
 
 
@@ -97,21 +99,33 @@ def _run_multiplier_bootstrap(
         gt_pairs = list(group_time_effects.keys())
         n_gt = len(gt_pairs)
 
-        # Generate bootstrap weights — PSU-level when survey design is present,
-        # cluster-level if clustered, unit-level otherwise.
+        # Original ATTs (independent of the draws; referenced per block below).
+        original_atts = np.array([group_time_effects[gt]["effect"] for gt in gt_pairs])
+
+        # Bootstrap weights are generated AND consumed one draw-block at a time so
+        # the dense (n_bootstrap, n_units) weight matrix is never materialized in
+        # full — the dominant allocation at large n_units. Weight source per path:
+        # PSU-level under a survey design, cluster-level if clustered, unit-level
+        # otherwise. The weight stream is bit-identical to the un-chunked path; the
+        # BLAS weights @ eif reductions may reassociate, so SEs match to within
+        # ~1 ULP (far below bootstrap Monte-Carlo error), not bit-for-bit.
         _use_survey_bootstrap = resolved_survey is not None and (
             resolved_survey.strata is not None
             or resolved_survey.psu is not None
             or resolved_survey.fpc is not None
         )
 
         if _use_survey_bootstrap:
-            from diff_diff.bootstrap_utils import (
-                generate_survey_multiplier_weights_batch as _gen_survey_weights,
-            )
-
-            psu_weights, psu_ids = _gen_survey_weights(
-                self.n_bootstrap, resolved_survey, self.bootstrap_weights, rng
+            # PSU-level multiplier weights, generated and expanded one draw-block
+            # 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,
             )
             # Single-cluster (G<2) survey-PSU multiplier bootstrap collapses
             # to constant multiplier draws → BLAS roundoff produces ≈0
@@ -143,35 +157,63 @@ def _run_multiplier_bootstrap(
                 )
             else:
                 unit_to_psu_col = np.arange(n_units)
-            all_weights = psu_weights[:, unit_to_psu_col]
+            # When each unit is its own PSU the expansion is an identity
+            # permutation — skip the needless full-block copy (CS parity).
+            _psu_is_identity = len(psu_ids) == n_units and bool(
+                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]
+
+            weight_blocks: Iterator[Tuple[int, np.ndarray]] = _weight_blocks()
         elif cluster_indices is not None and n_clusters is not None:
-            cluster_weights = _generate_bootstrap_weights_batch(
-                self.n_bootstrap, n_clusters, self.bootstrap_weights, rng
+            # 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,
             )
-            # Expand cluster weights to unit level
-            all_weights = cluster_weights[:, cluster_indices]
         else:
-            all_weights = _generate_bootstrap_weights_batch(
+            # 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
             )
 
-        # Original ATTs
-        original_atts = np.array([group_time_effects[gt]["effect"] for gt in gt_pairs])
-
-        # Perturbed ATTs: (n_bootstrap, n_gt)
-        # Under survey design, perturb survey-score object w_i * eif_i / sum(w)
-        # to match the analytical variance convention (compute_survey_if_variance).
-        bootstrap_atts = np.zeros((self.n_bootstrap, n_gt))
-        for j, gt in enumerate(gt_pairs):
-            eif_gt = eif_by_gt[gt]  # shape (n_units,)
-            with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
-                if unit_level_weights is not None:
-                    total_w = float(np.sum(unit_level_weights))
-                    eif_scaled = unit_level_weights * eif_gt / total_w
-                    perturbation = all_weights @ eif_scaled
-                else:
-                    perturbation = (all_weights @ eif_gt) / n_units
-            bootstrap_atts[:, j] = original_atts[j] + perturbation
+        # 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
+        # weights-only design that takes the unit weight path above), NOT on
+        # _use_survey_bootstrap. With weights present we perturb the survey-score
+        # object w_i * eif_i / sum(w) (matches compute_survey_if_variance);
+        # otherwise the raw eif with a 1/n prefactor applied after the matmul.
+        _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
 
         # Post-treatment mask — also exclude NaN effects
         post_mask = np.array(

diff_diff/had.py

@@ -71,6 +71,11 @@
 import numpy as np
 import pandas as pd
 
+from diff_diff.bootstrap_chunking import (
+    compute_block_size,
+    iter_survey_multiplier_weight_blocks,
+    iter_weight_blocks,
+)
 from diff_diff.local_linear import (
     BandwidthResult,
     BiasCorrectedFit,
@@ -2096,11 +2101,7 @@ def _sup_t_multiplier_bootstrap(
         n_valid)`` when fewer than half the draws are finite — warns the
         caller.
     """
-    from diff_diff.bootstrap_utils import (
-        apply_stratum_centering,
-        generate_bootstrap_weights_batch,
-        generate_survey_multiplier_weights_batch,
-    )
+    from diff_diff.bootstrap_utils import apply_stratum_centering
 
     influence_matrix = np.asarray(influence_matrix, dtype=np.float64)
     att_per_horizon = np.asarray(att_per_horizon, dtype=np.float64)
@@ -2159,8 +2160,15 @@ def _sup_t_multiplier_bootstrap(
                     "contributions; matches the 'remove' analytical target) "
                     "or pass cband=False to skip the simultaneous band."
                 )
-        psu_weights, psu_ids = generate_survey_multiplier_weights_batch(
-            n_bootstrap, resolved_survey, bootstrap_weights, rng
+        # PSU-level multiplier weights, generated one draw-block at a time so the
+        # (n_bootstrap, n_psu) matrix is never materialized in full (n_psu ==
+        # n_units when each unit is its own PSU). Unstratified designs tile the
+        # generation; stratified / degenerate designs fall back to full
+        # generation + slicing. Weight stream bit-identical to the un-chunked
+        # generate_survey_multiplier_weights_batch.
+        _block_size = compute_block_size(n_units, n_bootstrap)
+        psu_ids, _psu_blocks = iter_survey_multiplier_weight_blocks(
+            n_bootstrap, resolved_survey, bootstrap_weights, rng, block_size=_block_size
         )
         # Aggregate Psi to PSU level, stratum-demean, and apply the
         # small-sample correction so Var_xi(xi @ Psi_psu_scaled) matches
@@ -2171,7 +2179,7 @@ def _sup_t_multiplier_bootstrap(
         # (1 - f_h) FPC factor into the multipliers, so we only need to
         # pre-process Psi at the PSU level (aggregate → stratum-demean →
         # sqrt(n_h / (n_h - 1)) rescale).
-        n_psu = int(psu_weights.shape[1])
+        n_psu = len(psu_ids)
         psu_id_to_col = {int(p): c for c, p in enumerate(psu_ids)}
         Psi_psu = np.zeros((n_psu, n_horizons), dtype=np.float64)
         if resolved_survey.psu is not None:
@@ -2191,21 +2199,23 @@ def _sup_t_multiplier_bootstrap(
         # § "Note (Stute stratified survey-bootstrap calibration)".
         apply_stratum_centering(Psi_psu, resolved_survey, psu_ids, psu_axis=0)
 
-        # PSU-level perturbations: (B, H) = (B, n_psu) @ (n_psu, H).
-        # No (1/n) prefactor — Psi_psu_scaled is already on the θ̂-scale
-        # matched to the analytical variance.
-        with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
-            perturbations = psu_weights @ Psi_psu  # (B, H)
+        # PSU-level perturbations: (B, H) = (B, n_psu) @ (n_psu, H), accumulated
+        # one draw-block at a time (output is the small (B, H), so only the weight
+        # block is large). No (1/n) prefactor — Psi_psu_scaled is already on the
+        # θ̂-scale matched to the analytical variance.
+        perturbations = np.empty((n_bootstrap, n_horizons), dtype=np.float64)
+        for _cs, _psu_block in _psu_blocks:
+            with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
+                perturbations[_cs : _cs + _psu_block.shape[0]] = _psu_block @ Psi_psu
     else:
-        all_bootstrap_weights = generate_bootstrap_weights_batch(
-            n_bootstrap, n_units, bootstrap_weights, rng
-        )  # (B, G)
-        # Unit-level iid multipliers: no stratum centering needed.
-        # Var(xi @ Psi) = sum_g psi_g² matches the trivial analytical
-        # variance from compute_survey_if_variance at the IF-scale-
-        # invariant tolerance (PR #359 convention).
-        with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
-            perturbations = all_bootstrap_weights @ influence_matrix  # (B, H)
+        # Unit-level iid multipliers, generated one draw-block at a time at unit
+        # width (no stratum centering). Var(xi @ Psi) = sum_g psi_g² matches the
+        # trivial analytical variance from compute_survey_if_variance at the
+        # IF-scale-invariant tolerance (PR #359 convention).
+        perturbations = np.empty((n_bootstrap, n_horizons), dtype=np.float64)
+        for _cs, _w_block in iter_weight_blocks(n_bootstrap, n_units, bootstrap_weights, rng):
+            with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
+                perturbations[_cs : _cs + _w_block.shape[0]] = _w_block @ influence_matrix
 
     # t-statistics via per-horizon analytical SE.
     with np.errstate(divide="ignore", invalid="ignore", over="ignore"):