"""
Multiplier bootstrap inference for the Efficient DiD estimator.
Pattern follows CallawaySantAnnaBootstrapMixin (staggered_bootstrap.py).
Perturbs EIF values with random weights to obtain bootstrap distributions
of ATT(g,t) and aggregated parameters.
"""
import warnings
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple
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,
)
if TYPE_CHECKING:
from diff_diff.survey import ResolvedSurveyDesign
[docs]
@dataclass
class EDiDBootstrapResults:
"""Bootstrap inference results for EfficientDiD."""
n_bootstrap: int
weight_type: str
alpha: float
overall_att_se: float
overall_att_ci: Tuple[float, float]
overall_att_p_value: float
group_time_ses: Dict[Tuple[Any, Any], float]
group_time_cis: Dict[Tuple[Any, Any], Tuple[float, float]]
group_time_p_values: Dict[Tuple[Any, Any], float]
event_study_ses: Optional[Dict[int, float]] = None
event_study_cis: Optional[Dict[int, Tuple[float, float]]] = None
event_study_p_values: Optional[Dict[int, float]] = None
group_effect_ses: Optional[Dict[Any, float]] = None
group_effect_cis: Optional[Dict[Any, Tuple[float, float]]] = None
group_effect_p_values: Optional[Dict[Any, float]] = None
bootstrap_distribution: Optional[np.ndarray] = field(default=None, repr=False)
class EfficientDiDBootstrapMixin:
"""Mixin providing multiplier bootstrap for EfficientDiD."""
n_bootstrap: int
bootstrap_weights: str
alpha: float
seed: Optional[int]
anticipation: int
def _run_multiplier_bootstrap(
self,
group_time_effects: Dict[Tuple[Any, Any], Dict[str, Any]],
eif_by_gt: Dict[Tuple[Any, Any], np.ndarray],
n_units: int,
aggregate: Optional[str],
balance_e: Optional[int],
treatment_groups: List[Any],
cohort_fractions: Dict[float, float],
cluster_indices: Optional[np.ndarray] = None,
n_clusters: Optional[int] = None,
resolved_survey: Optional["ResolvedSurveyDesign"] = None,
unit_level_weights: Optional[np.ndarray] = None,
) -> EDiDBootstrapResults:
"""Run multiplier bootstrap on stored EIF values.
For each bootstrap draw *b*, perturb ATT(g,t) as::
ATT_b(g,t) = ATT(g,t) + (1/n) * xi_b @ eif_gt
where ``xi_b`` is an i.i.d. weight vector of length ``n_units``.
When ``cluster_indices`` is provided, weights are generated at the
cluster level and expanded to units.
Aggregations (overall, event study, group) are recomputed from
the perturbed ATT(g,t) values.
Note: Bootstrap aggregation uses fixed cohort-size weights, consistent
with the Callaway-Sant'Anna bootstrap pattern (staggered_bootstrap.py).
The analytical path includes a WIF correction for aggregated SEs, but
the bootstrap captures weight uncertainty through EIF perturbation.
This matches the R ``did`` package approach.
"""
if self.n_bootstrap < 50:
warnings.warn(
f"n_bootstrap={self.n_bootstrap} is low. Consider n_bootstrap >= 199 "
"for reliable inference.",
UserWarning,
stacklevel=3,
)
rng = np.random.default_rng(self.seed)
gt_pairs = list(group_time_effects.keys())
# 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:
# The flag definition above guarantees this (mypy can't track it).
assert resolved_survey is not None
# 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)
# 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
# miss this. EfficientDiD's cluster path is already protected by
# ``_validate_and_build_cluster_mapping`` (n_clusters≥2 at fit-time)
# and the unit path is protected by the balanced-panel validator;
# only the survey-PSU branch reaches the bootstrap with <2 PSUs.
if len(psu_ids) < 2:
warnings.warn(
f"Survey-PSU bootstrap with n_psu={len(psu_ids)} (<2 "
"independent PSUs) produces degenerate variance from BLAS "
"roundoff; returning NaN SE.",
UserWarning,
stacklevel=3,
)
return self._build_nan_bootstrap_results(
group_time_effects,
aggregate,
balance_e,
treatment_groups,
cohort_fractions,
)
# Build unit -> PSU column map
if resolved_survey.psu is not None:
psu_id_to_col = {int(p): c for c, p in enumerate(psu_ids)}
unit_to_psu_col = np.array(
[psu_id_to_col[int(resolved_survey.psu[i])] for i in range(n_units)]
)
else:
unit_to_psu_col = np.arange(n_units)
# 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))
)
# 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()
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]).
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.
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
# 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
# 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(
[
t >= g - self.anticipation and np.isfinite(original_atts[j])
for j, (g, t) in enumerate(gt_pairs)
]
)
post_indices = np.where(post_mask)[0]
# Overall ATT: fixed-weight re-aggregation of perturbed cell ATTs.
# This matches CallawaySantAnna._run_multiplier_bootstrap
# (staggered_bootstrap.py:281). The analytical path includes a WIF
# correction; bootstrap captures sampling variability through per-cell
# EIF perturbation without re-estimating weights — this is standard
# in both this library's CS implementation and the R did package.
skip_overall = len(post_indices) == 0
if skip_overall:
bootstrap_overall = np.full(self.n_bootstrap, np.nan)
original_overall = np.nan
else:
post_groups = [gt_pairs[i][0] for i in post_indices]
pg = np.array([cohort_fractions.get(g, 0.0) for g in post_groups])
agg_w = pg / pg.sum() if pg.sum() > 0 else np.ones(len(pg)) / len(pg)
original_overall = float(np.sum(agg_w * original_atts[post_mask]))
with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
bootstrap_overall = bootstrap_atts[:, post_indices] @ agg_w
# Event study: fixed-weight re-aggregation (same pattern as overall).
# See note above re: WIF — analytical WIF is not needed in bootstrap.
bootstrap_event_study = None
event_study_info = None
if aggregate in ("event_study", "all"):
event_study_info = self._prepare_es_agg_boot(
gt_pairs, original_atts, cohort_fractions, balance_e
)
bootstrap_event_study = {}
for e, info in event_study_info.items():
idx = info["gt_indices"]
w = info["weights"]
with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
bootstrap_event_study[e] = bootstrap_atts[:, idx] @ w
# Group aggregation
bootstrap_group = None
group_agg_info = None
if aggregate in ("group", "all"):
group_agg_info = self._prepare_group_agg_boot(gt_pairs, original_atts, treatment_groups)
bootstrap_group = {}
for g, info in group_agg_info.items():
idx = info["gt_indices"]
w = info["weights"]
with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
bootstrap_group[g] = bootstrap_atts[:, idx] @ w
# Compute statistics
gt_ses: Dict[Tuple[Any, Any], float] = {}
gt_cis: Dict[Tuple[Any, Any], Tuple[float, float]] = {}
gt_pvals: Dict[Tuple[Any, Any], float] = {}
for j, gt in enumerate(gt_pairs):
se, ci, pv = _compute_effect_bootstrap_stats_func(
original_atts[j],
bootstrap_atts[:, j],
alpha=self.alpha,
context=f"ATT(g={gt[0]}, t={gt[1]})",
)
gt_ses[gt] = se
gt_cis[gt] = ci
gt_pvals[gt] = pv
if skip_overall:
ov_se, ov_ci, ov_pv = np.nan, (np.nan, np.nan), np.nan
else:
ov_se, ov_ci, ov_pv = _compute_effect_bootstrap_stats_func(
original_overall,
bootstrap_overall,
alpha=self.alpha,
context="overall ATT",
)
es_ses = es_cis = es_pvs = None
if bootstrap_event_study is not None and event_study_info is not None:
es_ses, es_cis, es_pvs = {}, {}, {}
for e in sorted(event_study_info.keys()):
se, ci, pv = _compute_effect_bootstrap_stats_func(
event_study_info[e]["effect"],
bootstrap_event_study[e],
alpha=self.alpha,
context=f"event study (e={e})",
)
es_ses[e] = se
es_cis[e] = ci
es_pvs[e] = pv
g_ses = g_cis = g_pvs = None
if bootstrap_group is not None and group_agg_info is not None:
g_ses, g_cis, g_pvs = {}, {}, {}
for g in sorted(group_agg_info.keys()):
se, ci, pv = _compute_effect_bootstrap_stats_func(
group_agg_info[g]["effect"],
bootstrap_group[g],
alpha=self.alpha,
context=f"group effect (g={g})",
)
g_ses[g] = se
g_cis[g] = ci
g_pvs[g] = pv
return EDiDBootstrapResults(
n_bootstrap=self.n_bootstrap,
weight_type=self.bootstrap_weights,
alpha=self.alpha,
overall_att_se=ov_se,
overall_att_ci=ov_ci,
overall_att_p_value=ov_pv,
group_time_ses=gt_ses,
group_time_cis=gt_cis,
group_time_p_values=gt_pvals,
event_study_ses=es_ses,
event_study_cis=es_cis,
event_study_p_values=es_pvs,
group_effect_ses=g_ses,
group_effect_cis=g_cis,
group_effect_p_values=g_pvs,
bootstrap_distribution=bootstrap_overall,
)
def _prepare_es_agg_boot(
self,
gt_pairs: List[Tuple[Any, Any]],
original_atts: np.ndarray,
cohort_fractions: Dict[float, float],
balance_e: Optional[int],
) -> Dict[int, Dict[str, Any]]:
"""Prepare event-study aggregation info for bootstrap."""
effects_by_e: Dict[int, List[Tuple[int, float, float]]] = {}
for j, (g, t) in enumerate(gt_pairs):
if not np.isfinite(original_atts[j]):
continue # Skip NaN cells
e = t - g
if e not in effects_by_e:
effects_by_e[e] = []
effects_by_e[e].append((j, original_atts[j], cohort_fractions.get(g, 0.0)))
if balance_e is not None:
groups_at_e = {
gt_pairs[j][0]
for j, (g, t) in enumerate(gt_pairs)
if t - g == balance_e and np.isfinite(original_atts[j])
}
balanced: Dict[int, List[Tuple[int, float, float]]] = {}
for j, (g, t) in enumerate(gt_pairs):
if g in groups_at_e:
if not np.isfinite(original_atts[j]):
continue # Skip NaN cells even in balanced set
e = t - g
if e not in balanced:
balanced[e] = []
balanced[e].append((j, original_atts[j], cohort_fractions.get(g, 0.0)))
effects_by_e = balanced
if balance_e is not None and not effects_by_e:
warnings.warn(
f"balance_e={balance_e}: no cohort has a finite effect at the "
"anchor horizon. Event study will be empty.",
UserWarning,
stacklevel=2,
)
result = {}
for e, elist in effects_by_e.items():
indices = np.array([x[0] for x in elist])
effs = np.array([x[1] for x in elist])
pgs = np.array([x[2] for x in elist])
w = pgs / pgs.sum() if pgs.sum() > 0 else np.ones(len(pgs)) / len(pgs)
result[e] = {
"gt_indices": indices,
"weights": w,
"effect": float(np.sum(w * effs)),
}
return result
def _prepare_group_agg_boot(
self,
gt_pairs: List[Tuple[Any, Any]],
original_atts: np.ndarray,
treatment_groups: List[Any],
) -> Dict[Any, Dict[str, Any]]:
"""Prepare group-level aggregation info for bootstrap."""
result = {}
for g in treatment_groups:
group_data = [
(j, original_atts[j])
for j, (gg, t) in enumerate(gt_pairs)
if gg == g and t >= g - self.anticipation and np.isfinite(original_atts[j])
]
if not group_data:
continue
indices = np.array([x[0] for x in group_data])
effs = np.array([x[1] for x in group_data])
w = np.ones(len(effs)) / len(effs)
result[g] = {
"gt_indices": indices,
"weights": w,
"effect": float(np.sum(w * effs)),
}
return result
def _build_nan_bootstrap_results(
self,
group_time_effects: Dict[Tuple[Any, Any], Dict[str, Any]],
aggregate: Optional[str],
balance_e: Optional[int],
treatment_groups: List[Any],
cohort_fractions: Dict[float, float],
) -> EDiDBootstrapResults:
"""Return an all-NaN ``EDiDBootstrapResults`` for degenerate bootstrap.
Used when survey-PSU bootstrap collapses to G<2 PSUs and would
otherwise produce ≈0 SE from BLAS roundoff. Each NaN dict is keyed
to the same (g,t)/event-time/group reductions the downstream
override loop at ``efficient_did.py:1078-1115`` expects, so the
override finds each key and overwrites analytical SE with NaN.
Setting these dicts to ``None`` instead would let the analytical
SE leak through, defeating the NaN-propagation contract; keying
an empty dict would silently no-op the override for every key.
``event_study_ses``/``group_effect_ses`` are ``None`` (not empty)
when ``aggregate`` does not request them, matching the
``is not None`` gates at ``efficient_did.py:1090, 1109``.
"""
gt_pairs = list(group_time_effects.keys())
gt_ses: Dict[Tuple[Any, Any], float] = {gt: np.nan for gt in gt_pairs}
gt_cis: Dict[Tuple[Any, Any], Tuple[float, float]] = {
gt: (np.nan, np.nan) for gt in gt_pairs
}
gt_pvs: Dict[Tuple[Any, Any], float] = {gt: np.nan for gt in gt_pairs}
original_atts = np.array([group_time_effects[gt]["effect"] for gt in gt_pairs])
es_ses = es_cis = es_pvs = None
if aggregate in ("event_study", "all"):
es_info = self._prepare_es_agg_boot(
gt_pairs, original_atts, cohort_fractions, balance_e
)
if es_info:
es_ses = {e: np.nan for e in es_info.keys()}
es_cis = {e: (np.nan, np.nan) for e in es_info.keys()}
es_pvs = {e: np.nan for e in es_info.keys()}
g_ses = g_cis = g_pvs = None
if aggregate in ("group", "all"):
g_info = self._prepare_group_agg_boot(gt_pairs, original_atts, treatment_groups)
if g_info:
g_ses = {g: np.nan for g in g_info.keys()}
g_cis = {g: (np.nan, np.nan) for g in g_info.keys()}
g_pvs = {g: np.nan for g in g_info.keys()}
return EDiDBootstrapResults(
n_bootstrap=self.n_bootstrap,
weight_type=self.bootstrap_weights,
alpha=self.alpha,
overall_att_se=np.nan,
overall_att_ci=(np.nan, np.nan),
overall_att_p_value=np.nan,
group_time_ses=gt_ses,
group_time_cis=gt_cis,
group_time_p_values=gt_pvs,
event_study_ses=es_ses,
event_study_cis=es_cis,
event_study_p_values=es_pvs,
group_effect_ses=g_ses,
group_effect_cis=g_cis,
group_effect_p_values=g_pvs,
bootstrap_distribution=None,
)