"""
Shared bootstrap utilities for multiplier bootstrap inference.
Provides weight generation, percentile CI, and p-value helpers used by
both CallawaySantAnna and ContinuousDiD estimators.
"""
import warnings
from typing import Optional, Tuple
import numpy as np
from diff_diff._backend import HAS_RUST_BACKEND, _rust_bootstrap_weights
__all__ = [
"generate_bootstrap_weights",
"generate_bootstrap_weights_batch",
"generate_bootstrap_weights_batch_numpy",
"generate_survey_multiplier_weights_batch",
"generate_rao_wu_weights",
"generate_rao_wu_weights_batch",
"compute_percentile_ci",
"compute_bootstrap_pvalue",
"compute_effect_bootstrap_stats",
"compute_effect_bootstrap_stats_batch",
"warn_bootstrap_failure_rate",
"stratified_bootstrap_indices",
]
def stratified_bootstrap_indices(
rng: np.random.Generator,
n_control: int,
n_treated: int,
n_bootstrap: int,
) -> Tuple[np.ndarray, np.ndarray]:
"""Generate stratified bootstrap sample indices.
Draws controls first, then treated, per replicate. A single ``rng``
advances through all replicates so Python and Rust consumers share an
identical bytestream under the same seed.
Parameters
----------
rng : np.random.Generator
Seeded numpy generator. Consumed positionally; caller owns seeding.
n_control : int
Size of the control pool. Indices drawn in ``[0, n_control)``.
n_treated : int
Size of the treated pool. Indices drawn in ``[0, n_treated)``.
n_bootstrap : int
Number of bootstrap replicates.
Returns
-------
(control_indices, treated_indices) : tuple of np.ndarray
Shapes ``(n_bootstrap, n_control)`` and ``(n_bootstrap, n_treated)``,
dtype ``int64``. Each row is one replicate's sample-with-replacement
indices. Callers map these back to unit identities via their own pool.
"""
control_idx = np.empty((n_bootstrap, n_control), dtype=np.int64)
treated_idx = np.empty((n_bootstrap, n_treated), dtype=np.int64)
for b in range(n_bootstrap):
if n_control > 0:
control_idx[b] = rng.choice(n_control, size=n_control, replace=True)
if n_treated > 0:
treated_idx[b] = rng.choice(n_treated, size=n_treated, replace=True)
return control_idx, treated_idx
def warn_bootstrap_failure_rate(
n_success: int,
n_attempted: int,
context: str,
threshold: float = 0.05,
stacklevel: int = 3,
) -> None:
"""Emit one proportional failure-rate warning after a replicate loop.
Replaces the hard-coded ``< N successes`` pattern that lets high-failure
runs (e.g. 11 of 200) pass silently. Does not emit when
``n_attempted == 0`` (callers handle that degenerate path explicitly);
when ``n_success == 0`` and ``n_attempted > 0``, the warning fires
describing the all-failed run.
Parameters
----------
n_success : int
Number of replicates that produced a finite estimate.
n_attempted : int
Total replicates attempted (``self.n_bootstrap``).
context : str
Short label for the caller (e.g. ``"TROP global bootstrap"``).
threshold : float, default=0.05
Failure-rate threshold above which a warning is emitted. ``0.05``
matches the existing SyntheticDiD bootstrap and placebo guards.
stacklevel : int, default=3
Passed to :func:`warnings.warn`.
"""
if n_attempted = n_attempted:
return
failure_rate = 1.0 - (n_success / n_attempted)
if failure_rate np.ndarray:
"""
Generate bootstrap weights for multiplier bootstrap.
Parameters
----------
n_units : int
Number of units (clusters) to generate weights for.
weight_type : str
Type of weights: "rademacher", "mammen", or "webb".
rng : np.random.Generator
Random number generator.
Returns
-------
np.ndarray
Array of bootstrap weights with shape (n_units,).
"""
if weight_type == "rademacher":
return rng.choice([-1.0, 1.0], size=n_units)
elif weight_type == "mammen":
sqrt5 = np.sqrt(5)
val1 = -(sqrt5 - 1) / 2
val2 = (sqrt5 + 1) / 2
p1 = (sqrt5 + 1) / (2 * sqrt5)
return rng.choice([val1, val2], size=n_units, p=[p1, 1 - p1])
elif weight_type == "webb":
values = np.array(
[
-np.sqrt(3 / 2),
-np.sqrt(2 / 2),
-np.sqrt(1 / 2),
np.sqrt(1 / 2),
np.sqrt(2 / 2),
np.sqrt(3 / 2),
]
)
return rng.choice(values, size=n_units)
else:
raise ValueError(
f"weight_type must be 'rademacher', 'mammen', or 'webb', " f"got '{weight_type}'"
)
def generate_bootstrap_weights_batch(
n_bootstrap: int,
n_units: int,
weight_type: str,
rng: np.random.Generator,
) -> np.ndarray:
"""
Generate all bootstrap weights at once (vectorized).
Uses Rust backend if available for parallel generation.
Parameters
----------
n_bootstrap : int
Number of bootstrap iterations.
n_units : int
Number of units (clusters) to generate weights for.
weight_type : str
Type of weights: "rademacher", "mammen", or "webb".
rng : np.random.Generator
Random number generator.
Returns
-------
np.ndarray
Array of bootstrap weights with shape (n_bootstrap, n_units).
"""
if HAS_RUST_BACKEND and _rust_bootstrap_weights is not None:
seed = rng.integers(0, 2**63 - 1)
return _rust_bootstrap_weights(n_bootstrap, n_units, weight_type, seed)
return generate_bootstrap_weights_batch_numpy(n_bootstrap, n_units, weight_type, rng)
def generate_bootstrap_weights_batch_numpy(
n_bootstrap: int,
n_units: int,
weight_type: str,
rng: np.random.Generator,
) -> np.ndarray:
"""
NumPy fallback implementation of :func:`generate_bootstrap_weights_batch`.
Parameters
----------
n_bootstrap : int
Number of bootstrap iterations.
n_units : int
Number of units (clusters) to generate weights for.
weight_type : str
Type of weights: "rademacher", "mammen", or "webb".
rng : np.random.Generator
Random number generator.
Returns
-------
np.ndarray
Array of bootstrap weights with shape (n_bootstrap, n_units).
"""
if weight_type == "rademacher":
return rng.choice([-1.0, 1.0], size=(n_bootstrap, n_units))
elif weight_type == "mammen":
sqrt5 = np.sqrt(5)
val1 = -(sqrt5 - 1) / 2
val2 = (sqrt5 + 1) / 2
p1 = (sqrt5 + 1) / (2 * sqrt5)
return rng.choice([val1, val2], size=(n_bootstrap, n_units), p=[p1, 1 - p1])
elif weight_type == "webb":
values = np.array(
[
-np.sqrt(3 / 2),
-np.sqrt(2 / 2),
-np.sqrt(1 / 2),
np.sqrt(1 / 2),
np.sqrt(2 / 2),
np.sqrt(3 / 2),
]
)
return rng.choice(values, size=(n_bootstrap, n_units))
else:
raise ValueError(
f"weight_type must be 'rademacher', 'mammen', or 'webb', " f"got '{weight_type}'"
)
def compute_percentile_ci(
boot_dist: np.ndarray,
alpha: float,
) -> Tuple[float, float]:
"""
Compute percentile confidence interval from bootstrap distribution.
Parameters
----------
boot_dist : np.ndarray
Bootstrap distribution (1-D array).
alpha : float
Significance level (e.g., 0.05 for 95% CI).
Returns
-------
tuple of float
``(lower, upper)`` confidence interval bounds.
"""
lower = float(np.percentile(boot_dist, alpha / 2 * 100))
upper = float(np.percentile(boot_dist, (1 - alpha / 2) * 100))
return (lower, upper)
def compute_bootstrap_pvalue(
original_effect: float,
boot_dist: np.ndarray,
n_valid: Optional[int] = None,
) -> float:
"""
Compute two-sided bootstrap p-value using the percentile method.
Parameters
----------
original_effect : float
Original point estimate.
boot_dist : np.ndarray
Bootstrap distribution of the effect.
n_valid : int, optional
Number of valid bootstrap samples for p-value floor.
If None, uses ``len(boot_dist)``.
Returns
-------
float
Two-sided bootstrap p-value.
"""
if original_effect >= 0:
p_one_sided = np.mean(boot_dist = 0)
p_value = min(2 * p_one_sided, 1.0)
n_for_floor = n_valid if n_valid is not None else len(boot_dist)
p_value = max(p_value, 1 / (n_for_floor + 1))
return float(p_value)
def compute_effect_bootstrap_stats(
original_effect: float,
boot_dist: np.ndarray,
alpha: float = 0.05,
context: str = "bootstrap distribution",
) -> Tuple[float, Tuple[float, float], float]:
"""
Compute bootstrap statistics for a single effect.
Filters non-finite samples, returning NaN for all statistics if
fewer than 50% of samples are valid.
Parameters
----------
original_effect : float
Original point estimate.
boot_dist : np.ndarray
Bootstrap distribution of the effect.
alpha : float, default=0.05
Significance level.
context : str, optional
Description for warning messages.
Returns
-------
se : float
Bootstrap standard error.
ci : tuple of float
Percentile confidence interval.
p_value : float
Bootstrap p-value.
"""
if not np.isfinite(original_effect):
return np.nan, (np.nan, np.nan), np.nan
finite_mask = np.isfinite(boot_dist)
n_valid = np.sum(finite_mask)
n_total = len(boot_dist)
if n_valid < n_total:
n_nonfinite = n_total - n_valid
warnings.warn(
f"Dropping {n_nonfinite}/{n_total} non-finite bootstrap samples "
f"in {context}. Bootstrap estimates based on remaining valid samples.",
RuntimeWarning,
stacklevel=3,
)
if n_valid < n_total * 0.5:
warnings.warn(
f"Too few valid bootstrap samples ({n_valid}/{n_total}) in {context}. "
"Returning NaN for SE/CI/p-value to signal invalid inference.",
RuntimeWarning,
stacklevel=3,
)
return np.nan, (np.nan, np.nan), np.nan
valid_dist = boot_dist[finite_mask]
se = float(np.std(valid_dist, ddof=1))
# Guard: if SE is not finite or zero, all inference fields must be NaN.
if not np.isfinite(se) or se tuple:
"""
Batch-compute bootstrap statistics for multiple effects at once.
Parameters
----------
original_effects : np.ndarray
Array of original point estimates, shape (n_effects,).
bootstrap_matrix : np.ndarray
Bootstrap distributions, shape (n_bootstrap, n_effects).
alpha : float, default=0.05
Significance level.
Returns
-------
ses : np.ndarray
Bootstrap SEs for each effect.
ci_lowers : np.ndarray
Lower CI bounds for each effect.
ci_uppers : np.ndarray
Upper CI bounds for each effect.
p_values : np.ndarray
Bootstrap p-values for each effect.
"""
n_bootstrap, n_effects = bootstrap_matrix.shape
ses = np.full(n_effects, np.nan)
ci_lowers = np.full(n_effects, np.nan)
ci_uppers = np.full(n_effects, np.nan)
p_values = np.full(n_effects, np.nan)
# Check for non-finite original effects
valid_effects = np.isfinite(original_effects)
if not np.any(valid_effects):
return ses, ci_lowers, ci_uppers, p_values
# Count valid bootstrap samples per effect
finite_mask = np.isfinite(bootstrap_matrix) # (n_bootstrap, n_effects)
n_valid = finite_mask.sum(axis=0) # (n_effects,)
# Determine which effects have enough valid samples
enough_valid = (n_valid >= n_bootstrap * 0.5) & valid_effects
if not np.any(enough_valid):
n_insufficient = int(np.sum(valid_effects))
if n_insufficient > 0:
warnings.warn(
f"{n_insufficient} effect(s) had too few valid bootstrap samples (= number of PSUs."
)
f = n_units_for_fpc / resolved_survey.fpc[0]
if f < 1.0:
weights = weights * np.sqrt(1.0 - f)
else:
weights = np.zeros_like(weights)
else:
# Stratified generate independently within strata
weights = np.empty((n_bootstrap, n_psu), dtype=np.float64)
# Build PSU column-index map
psu_to_col = {int(p): i for i, p in enumerate(psu_ids)}
unique_strata = np.unique(strata)
_singleton_cols = [] # For lonely_psu="adjust" pooling
for h in unique_strata:
mask_h = strata == h
if psu is not None:
psus_in_h = np.unique(psu[mask_h])
else:
psus_in_h = np.where(mask_h)[0]
n_h = len(psus_in_h)
cols = np.array([psu_to_col[int(p)] for p in psus_in_h])
if n_h < 2:
if _lonely_psu == "adjust":
# Collect for pooled pseudo-stratum processing
_singleton_cols.extend(cols.tolist())
else:
# remove / certainty zero weight
weights[:, cols] = 0.0
continue
# Generate weights for this stratum
stratum_weights = generate_bootstrap_weights_batch_numpy(
n_bootstrap, n_h, weight_type, rng
)
# FPC scaling
if resolved_survey.fpc is not None:
N_h = resolved_survey.fpc[mask_h][0]
if N_h < n_h:
raise ValueError(
f"FPC ({N_h}) is less than the number of PSUs "
f"({n_h}) in stratum {h}. FPC must be >= n_PSU."
)
f_h = n_h / N_h
if f_h < 1.0:
stratum_weights = stratum_weights * np.sqrt(1.0 - f_h)
else:
stratum_weights = np.zeros_like(stratum_weights)
weights[:, cols] = stratum_weights
# Pool singleton PSUs into a pseudo-stratum for "adjust"
if _singleton_cols:
n_pooled = len(_singleton_cols)
if n_pooled >= 2:
pooled_weights = generate_bootstrap_weights_batch_numpy(
n_bootstrap, n_pooled, weight_type, rng
)
# No FPC scaling for pooled singletons (conservative)
pooled_cols = np.array(_singleton_cols)
weights[:, pooled_cols] = pooled_weights
else:
# Single singleton cannot pool, zero weight (library-specific
# fallback; bootstrap adjust with one singleton = remove).
import warnings
warnings.warn(
"lonely_psu='adjust' with only 1 singleton stratum in "
"bootstrap: singleton PSU contributes zero variance "
"(same as 'remove'). At least 2 singleton strata are "
"needed for pooled pseudo-stratum bootstrap.",
UserWarning,
stacklevel=3,
)
weights[:, _singleton_cols[0]] = 0.0
return weights, psu_ids
def generate_rao_wu_weights(
resolved_survey: "ResolvedSurveyDesign",
rng: np.random.Generator,
) -> np.ndarray:
"""Generate one set of Rao-Wu (1988) rescaled observation weights.
Within each stratum *h* with *n_h* PSUs, draw ``m_h`` PSUs with
replacement and rescale observation weights by ``(n_h / m_h) * r_hi``
where ``r_hi`` is the count of PSU *i* being selected.
Without FPC: ``m_h = n_h - 1``.
With FPC: ``m_h = max(1, round((1 - f_h) * (n_h - 1)))``
(Rao, Wu & Yue 1992, Section 3).
For ``lonely_psu="adjust"``, singleton PSUs are pooled into a combined
pseudo-stratum and resampled together (no FPC scaling on pooled group).
Parameters
----------
resolved_survey : ResolvedSurveyDesign
Resolved survey design.
rng : np.random.Generator
Random number generator.
Returns
-------
np.ndarray
Rescaled observation weights, shape ``(n_obs,)``.
"""
n_obs = len(resolved_survey.weights)
base_weights = resolved_survey.weights
psu = resolved_survey.psu
strata = resolved_survey.strata
_lonely_psu_rw = resolved_survey.lonely_psu
rescaled = np.zeros(n_obs, dtype=np.float64)
if psu is None:
obs_psu = np.arange(n_obs)
else:
obs_psu = psu
if strata is None:
strata_masks = [np.ones(n_obs, dtype=bool)]
else:
unique_strata = np.unique(strata)
strata_masks = [strata == h for h in unique_strata]
# Collect singleton PSUs for "adjust" pooling
_singleton_info = [] # list of (mask_h, unique_psu_h) tuples
for mask_h in strata_masks:
psu_h = obs_psu[mask_h]
unique_psu_h = np.unique(psu_h)
n_h = len(unique_psu_h)
if n_h < 2:
if _lonely_psu_rw == "adjust":
_singleton_info.append((mask_h, unique_psu_h))
else:
# remove / certainty keep original weights (zero variance)
rescaled[mask_h] = base_weights[mask_h]
continue
# Compute resample size
if resolved_survey.fpc is not None:
N_h = resolved_survey.fpc[mask_h][0]
if N_h < n_h:
raise ValueError(
f"FPC ({N_h}) is less than the number of PSUs "
f"({n_h}). FPC must be >= number of PSUs."
)
f_h = n_h / N_h
if f_h >= 1.0:
# Census stratum keep original weights (zero variance)
rescaled[mask_h] = base_weights[mask_h]
continue
m_h = max(1, round((1.0 - f_h) * (n_h - 1)))
else:
m_h = n_h - 1
# Draw m_h PSUs with replacement
drawn_indices = rng.choice(n_h, size=m_h, replace=True)
counts = np.bincount(drawn_indices, minlength=n_h)
# Rescale factor per PSU: (n_h / m_h) * r_hi
scale_per_psu = (n_h / m_h) * counts.astype(np.float64)
# Map PSU local index for vectorized application
psu_to_local = {int(p): i for i, p in enumerate(unique_psu_h)}
obs_in_h = np.where(mask_h)[0]
local_indices = np.array([psu_to_local[int(obs_psu[idx])] for idx in obs_in_h])
rescaled[obs_in_h] = base_weights[obs_in_h] * scale_per_psu[local_indices]
# Pool singleton PSUs into a pseudo-stratum for "adjust"
if _singleton_info:
# Combine all singleton PSUs into one group
pooled_psus = np.concatenate([p for _, p in _singleton_info])
n_pooled = len(pooled_psus)
if n_pooled >= 2:
m_pooled = n_pooled - 1 # No FPC for pooled singletons
drawn = rng.choice(n_pooled, size=m_pooled, replace=True)
counts = np.bincount(drawn, minlength=n_pooled)
scale_per_psu = (n_pooled / m_pooled) * counts.astype(np.float64)
# Build PSU scale mapping and apply
psu_scale_map = {int(pooled_psus[i]): scale_per_psu[i] for i in range(n_pooled)}
for mask_h, _ in _singleton_info:
obs_in_h = np.where(mask_h)[0]
for idx in obs_in_h:
p = int(obs_psu[idx])
rescaled[idx] = base_weights[idx] * psu_scale_map.get(p, 1.0)
else:
# Single singleton cannot pool, keep base weights (library-specific
# fallback; bootstrap adjust with one singleton = remove).
import warnings
warnings.warn(
"lonely_psu='adjust' with only 1 singleton stratum in "
"bootstrap: singleton PSU contributes zero variance "
"(same as 'remove'). At least 2 singleton strata are "
"needed for pooled pseudo-stratum bootstrap.",
UserWarning,
stacklevel=2,
)
for mask_h, _ in _singleton_info:
rescaled[mask_h] = base_weights[mask_h]
return rescaled
def generate_rao_wu_weights_batch(
n_bootstrap: int,
resolved_survey: "ResolvedSurveyDesign",
rng: np.random.Generator,
) -> np.ndarray:
"""Generate multiple sets of Rao-Wu rescaled weights.
Parameters
----------
n_bootstrap : int
Number of bootstrap iterations.
resolved_survey : ResolvedSurveyDesign
Resolved survey design.
rng : np.random.Generator
Random number generator.
Returns
-------
np.ndarray
Rescaled weights, shape ``(n_bootstrap, n_obs)``.
"""
n_obs = len(resolved_survey.weights)
result = np.empty((n_bootstrap, n_obs), dtype=np.float64)
for b in range(n_bootstrap):
result[b] = generate_rao_wu_weights(resolved_survey, rng)
return result