"""

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

@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,

)