"""
Power analysis tools for difference-in-differences study design.
This module provides power calculations and simulation-based power analysis
for DiD study design, helping practitioners answer questions like:
- "How many units do I need to detect an effect of size X?"
- "What is the minimum detectable effect given my sample size?"
- "What power do I have to detect a given effect?"
References
----------
Bloom, H. S. (1995). "Minimum Detectable Effects: A Simple Way to Report the
Statistical Power of Experimental Designs." Evaluation Review, 19(5), 547-556.
Burlig, F., Preonas, L., & Woerman, M. (2020). "Panel Data and Experimental Design."
Journal of Development Economics, 144, 102458.
Djimeu, E. W., & Houndolo, D.-G. (2016). "Power Calculation for Causal Inference
in Social Science: Sample Size and Minimum Detectable Effect Determination."
Journal of Development Effectiveness, 8(4), 508-527.
"""
import warnings
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Tuple
import numpy as np
import pandas as pd
from scipy import stats
# Maximum sample size returned when effect is too small to detect
# (e.g., zero effect or extremely small relative to noise)
MAX_SAMPLE_SIZE = 2**31 - 1
# ---------------------------------------------------------------------------
# Estimator registry maps estimator class names to DGP/fit/extract profiles
# ---------------------------------------------------------------------------
@dataclass
class _EstimatorProfile:
"""Internal profile describing how to run power simulations for an estimator."""
default_dgp: Callable
dgp_kwargs_builder: Callable
fit_kwargs_builder: Callable
result_extractor: Callable
min_n: int = 20
# ---------------------------------------------------------------------------
# SurveyPowerConfig carries DGP survey params for simulation power
# ---------------------------------------------------------------------------
@dataclass
class SurveyPowerConfig:
"""Configuration for survey-aware power simulations.
When passed to :func:`simulate_power`, :func:`simulate_mde`, or
:func:`simulate_sample_size`, the simulation loop generates data with
:func:`~diff_diff.prep.generate_survey_did_data` and automatically
injects a ``SurveyDesign`` into the estimator's ``fit()`` call.
Parameters
----------
n_strata : int, default=5
Number of geographic strata.
psu_per_stratum : int, default=8
Number of primary sampling units (PSUs) per stratum. Must be >= 2
for Taylor Series Linearization variance estimation.
fpc_per_stratum : float, default=200.0
Finite population correction (total PSUs per stratum).
weight_variation : str, default="moderate"
Sampling weight dispersion: ``"none"`` (all equal), ``"moderate"``
(range ~1-2), ``"high"`` (range ~1-4).
psu_re_sd : float, default=2.0
Standard deviation of PSU random effects. Controls intra-cluster
correlation and drives DEFF > 1.
psu_period_factor : float, default=0.5
Multiplier for PSU-period interaction shocks.
icc : float, optional
Target intra-class correlation (0 < icc < 1). Overrides
``psu_re_sd`` via variance decomposition.
weight_cv : float, optional
Target coefficient of variation for weights. Overrides
``weight_variation``.
informative_sampling : bool, default=False
If True, weights correlate with Y(0).
heterogeneous_te_by_strata : bool, default=False
If True, treatment effect varies by stratum.
include_replicate_weights : bool, default=False
If True, add JK1 delete-one-PSU replicate weight columns.
survey_design : SurveyDesign, optional
Override the auto-built SurveyDesign. When None, a default
``SurveyDesign(weights="weight", strata="stratum", psu="psu",
fpc="fpc")`` is used, matching ``generate_survey_did_data`` output.
Examples
--------
>>> from diff_diff import CallawaySantAnna, simulate_power, SurveyPowerConfig
>>> config = SurveyPowerConfig(n_strata=5, psu_per_stratum=8, icc=0.05)
>>> results = simulate_power(
... CallawaySantAnna(),
... n_units=200,
... treatment_effect=2.0,
... survey_config=config,
... n_simulations=100,
... seed=42,
... )
"""
n_strata: int = 5
psu_per_stratum: int = 8
fpc_per_stratum: float = 200.0
weight_variation: str = "moderate"
psu_re_sd: float = 2.0
psu_period_factor: float = 0.5
icc: Optional[float] = None
weight_cv: Optional[float] = None
informative_sampling: bool = False
heterogeneous_te_by_strata: bool = False
include_replicate_weights: bool = False
survey_design: Optional[Any] = None
def __post_init__(self) -> None:
if self.n_strata < 1:
raise ValueError(f"n_strata must be >= 1, got {self.n_strata}")
if self.psu_per_stratum < 2:
raise ValueError(
f"psu_per_stratum must be >= 2 for TSL variance estimation, "
f"got {self.psu_per_stratum}"
)
if self.weight_variation not in ("none", "moderate", "high"):
raise ValueError(
f"weight_variation must be 'none', 'moderate', or 'high', "
f"got '{self.weight_variation}'"
)
if not np.isfinite(self.psu_re_sd) or self.psu_re_sd < 0:
raise ValueError(f"psu_re_sd must be finite and >= 0, got {self.psu_re_sd}")
if not np.isfinite(self.fpc_per_stratum):
raise ValueError(f"fpc_per_stratum must be finite, got {self.fpc_per_stratum}")
if self.icc is not None and not (0 < self.icc < 1):
raise ValueError(f"icc must be between 0 and 1 (exclusive), got {self.icc}")
if self.icc is not None and self.psu_re_sd != 2.0:
raise ValueError(
"Cannot specify both icc and a non-default psu_re_sd. "
"icc overrides psu_re_sd via the ICC formula."
)
if self.weight_cv is not None:
if not np.isfinite(self.weight_cv) or self.weight_cv Any:
"""Return a SurveyDesign for this config.
Reflects the live ``self.survey_design`` value every call (no
caching). Finding #28 (axis J, silent-failures audit): the
previous ``_cached_survey_design`` was populated on first call
and never invalidated on mutation, so ``config.survey_design =
other_design`` silently kept returning the original. Since the
default ``SurveyDesign(...)`` construction is microseconds and
user-provided designs are just reference copies, there's no cache
cost worth keeping.
"""
if self.survey_design is not None:
return self.survey_design
from diff_diff.survey import SurveyDesign
return SurveyDesign(
weights="weight", strata="stratum", psu="psu", fpc="fpc"
)
@property
def min_viable_n(self) -> int:
"""Minimum n_units for a viable survey design (>= 2 units per PSU)."""
return self.n_strata * self.psu_per_stratum * 2
# -- DGP kwargs adapters -----------------------------------------------------
def _basic_dgp_kwargs(
n_units: int,
n_periods: int,
treatment_effect: float,
treatment_fraction: float,
treatment_period: int,
sigma: float,
) -> Dict[str, Any]:
return dict(
n_units=n_units,
n_periods=n_periods,
treatment_effect=treatment_effect,
treatment_fraction=treatment_fraction,
treatment_period=treatment_period,
noise_sd=sigma,
)
def _staggered_dgp_kwargs(
n_units: int,
n_periods: int,
treatment_effect: float,
treatment_fraction: float,
treatment_period: int,
sigma: float,
) -> Dict[str, Any]:
return dict(
n_units=n_units,
n_periods=n_periods,
treatment_effect=treatment_effect,
never_treated_frac=1 - treatment_fraction,
cohort_periods=[treatment_period],
dynamic_effects=False,
noise_sd=sigma,
)
def _factor_dgp_kwargs(
n_units: int,
n_periods: int,
treatment_effect: float,
treatment_fraction: float,
treatment_period: int,
sigma: float,
) -> Dict[str, Any]:
n_pre = treatment_period
n_post = n_periods - treatment_period
return dict(
n_units=n_units,
n_pre=n_pre,
n_post=n_post,
n_treated=max(1, int(n_units * treatment_fraction)),
treatment_effect=treatment_effect,
noise_sd=sigma,
)
def _ddd_dgp_kwargs(
n_units: int,
n_periods: int,
treatment_effect: float,
treatment_fraction: float,
treatment_period: int,
sigma: float,
) -> Dict[str, Any]:
return dict(
n_per_cell=max(2, n_units // 8),
treatment_effect=treatment_effect,
noise_sd=sigma,
)
# -- Fit kwargs builders ------------------------------------------------------
def _basic_fit_kwargs(
data: pd.DataFrame,
n_units: int,
n_periods: int,
treatment_period: int,
) -> Dict[str, Any]:
return dict(outcome="outcome", treatment="treated", time="post")
def _twfe_fit_kwargs(
data: pd.DataFrame,
n_units: int,
n_periods: int,
treatment_period: int,
) -> Dict[str, Any]:
return dict(outcome="outcome", treatment="treated", time="post", unit="unit")
def _multiperiod_fit_kwargs(
data: pd.DataFrame,
n_units: int,
n_periods: int,
treatment_period: int,
) -> Dict[str, Any]:
return dict(
outcome="outcome",
treatment="treated",
time="period",
post_periods=list(range(treatment_period, n_periods)),
)
def _staggered_fit_kwargs(
data: pd.DataFrame,
n_units: int,
n_periods: int,
treatment_period: int,
) -> Dict[str, Any]:
return dict(outcome="outcome", unit="unit", time="period", first_treat="first_treat")
def _ddd_fit_kwargs(
data: pd.DataFrame,
n_units: int,
n_periods: int,
treatment_period: int,
) -> Dict[str, Any]:
return dict(outcome="outcome", group="group", partition="partition", time="time")
def _trop_fit_kwargs(
data: pd.DataFrame,
n_units: int,
n_periods: int,
treatment_period: int,
) -> Dict[str, Any]:
return dict(outcome="outcome", treatment="treated", unit="unit", time="period")
def _sdid_fit_kwargs(
data: pd.DataFrame,
n_units: int,
n_periods: int,
treatment_period: int,
) -> Dict[str, Any]:
periods = sorted(data["period"].unique())
post_periods = [p for p in periods if p >= treatment_period]
return dict(
outcome="outcome",
treatment="treat",
unit="unit",
time="period",
post_periods=post_periods,
)
# -- Survey-aware DGP kwargs adapter ------------------------------------------
def _survey_dgp_kwargs(
n_units: int,
n_periods: int,
treatment_effect: float,
treatment_fraction: float,
treatment_period: int,
sigma: float,
survey_config: SurveyPowerConfig,
) -> Dict[str, Any]:
"""Build kwargs for generate_survey_did_data from simulate_power params."""
return dict(
n_units=n_units,
n_periods=n_periods,
treatment_effect=treatment_effect,
never_treated_frac=1 - treatment_fraction,
# 0-indexed treatment_period 1-indexed cohort_periods
cohort_periods=[treatment_period + 1],
noise_sd=sigma,
dynamic_effects=False,
n_strata=survey_config.n_strata,
psu_per_stratum=survey_config.psu_per_stratum,
fpc_per_stratum=survey_config.fpc_per_stratum,
weight_variation=survey_config.weight_variation,
psu_re_sd=survey_config.psu_re_sd,
psu_period_factor=survey_config.psu_period_factor,
icc=survey_config.icc,
weight_cv=survey_config.weight_cv,
informative_sampling=survey_config.informative_sampling,
heterogeneous_te_by_strata=survey_config.heterogeneous_te_by_strata,
include_replicate_weights=survey_config.include_replicate_weights,
return_true_population_att=True,
)
# -- Survey-aware fit kwargs builders -----------------------------------------
def _survey_basic_fit_kwargs(
data: pd.DataFrame,
n_units: int,
n_periods: int,
treatment_period: int,
survey_config: SurveyPowerConfig,
) -> Dict[str, Any]:
"""Fit kwargs for DifferenceInDifferences with survey design.
Uses ``ever_treated`` (time-invariant group indicator) rather than the
survey DGP's ``treated`` column (which is post-only: 1{g>0, t>=g}).
DifferenceInDifferences internally constructs ``treatment * time``,
so passing the post-only flag would make that interaction rank-deficient.
"""
return dict(
outcome="outcome",
treatment="ever_treated",
time="post",
survey_design=survey_config._build_survey_design(),
)
def _survey_twfe_fit_kwargs(
data: pd.DataFrame,
n_units: int,
n_periods: int,
treatment_period: int,
survey_config: SurveyPowerConfig,
) -> Dict[str, Any]:
"""Fit kwargs for TwoWayFixedEffects with survey design."""
return dict(
outcome="outcome",
treatment="ever_treated",
time="post",
unit="unit",
survey_design=survey_config._build_survey_design(),
)
def _survey_multiperiod_fit_kwargs(
data: pd.DataFrame,
n_units: int,
n_periods: int,
treatment_period: int,
survey_config: SurveyPowerConfig,
) -> Dict[str, Any]:
"""Fit kwargs for MultiPeriodDiD with survey design (1-indexed periods)."""
return dict(
outcome="outcome",
treatment="ever_treated",
unit="unit",
time="period",
# 1-indexed: post periods run from treatment_period+1 to n_periods
post_periods=list(range(treatment_period + 1, n_periods + 1)),
survey_design=survey_config._build_survey_design(),
)
def _survey_staggered_fit_kwargs(
data: pd.DataFrame,
n_units: int,
n_periods: int,
treatment_period: int,
survey_config: SurveyPowerConfig,
) -> Dict[str, Any]:
"""Fit kwargs for staggered estimators (CS, SA, etc.) with survey design."""
return dict(
outcome="outcome",
unit="unit",
time="period",
first_treat="first_treat",
survey_design=survey_config._build_survey_design(),
)
# -- Result extractors --------------------------------------------------------
def _extract_simple(result: Any) -> Tuple[float, float, float, Tuple[float, float]]:
return (result.att, result.se, result.p_value, result.conf_int)
def _extract_multiperiod(
result: Any,
) -> Tuple[float, float, float, Tuple[float, float]]:
return (result.avg_att, result.avg_se, result.avg_p_value, result.avg_conf_int)
def _extract_staggered(
result: Any,
) -> Tuple[float, float, float, Tuple[float, float]]:
_nan = float("nan")
_nan_ci = (_nan, _nan)
def _first(r: Any, *attrs: str, default: Any = _nan) -> Any:
for a in attrs:
v = getattr(r, a, None)
if v is not None:
return v
return default
return (
result.overall_att,
_first(result, "overall_se", "overall_att_se"),
_first(result, "overall_p_value", "overall_att_p_value"),
_first(result, "overall_conf_int", "overall_att_ci", default=_nan_ci),
)
# Keys derived from simulate_power() public params overriding these
# via data_generator_kwargs would desync the DGP from the result object.
_PROTECTED_DGP_KEYS = frozenset(
{
"treatment_effect", # true_effect in results / MDE search variable
"noise_sd", # sigma param
"n_units", # sample-size search variable
"n_periods", # n_periods param
"treatment_fraction", # treatment_fraction param
"treatment_period", # treatment_period param
"n_pre", # derived from treatment_period in factor-model DGPs
"n_post", # derived from n_periods - treatment_period in factor-model DGPs
}
)
# Keys managed by SurveyPowerConfig block in data_generator_kwargs when
# survey_config is active to prevent silent conflicts.
_SURVEY_CONFIG_KEYS = frozenset(
{
"n_strata",
"psu_per_stratum",
"fpc_per_stratum",
"weight_variation",
"psu_re_sd",
"psu_period_factor",
"icc",
"weight_cv",
"informative_sampling",
"heterogeneous_te_by_strata",
"include_replicate_weights",
"return_true_population_att",
"dynamic_effects",
"cohort_periods",
"never_treated_frac",
}
)
# -- Staggered DGP compatibility check ----------------------------------------
_STAGGERED_ESTIMATORS = frozenset(
{
"CallawaySantAnna",
"SunAbraham",
"ImputationDiD",
"TwoStageDiD",
"StackedDiD",
"EfficientDiD",
}
)
# Estimators that need a derived `post` column when using survey DGP
# (survey DGP produces `period`/`first_treat` but not `post`).
_SURVEY_POST_ESTIMATORS = frozenset({"DifferenceInDifferences", "TwoWayFixedEffects"})
# Survey fit kwargs builder lookup maps estimator name to builder function.
_SURVEY_FIT_BUILDERS: Dict[str, Callable] = {
"DifferenceInDifferences": _survey_basic_fit_kwargs,
"TwoWayFixedEffects": _survey_twfe_fit_kwargs,
"MultiPeriodDiD": _survey_multiperiod_fit_kwargs,
**{name: _survey_staggered_fit_kwargs for name in _STAGGERED_ESTIMATORS},
}
# Unsupported: factor-model and triple-diff estimators (survey DGP produces
# staggered cohort data, not factor-model or 2x2x2 data).
_SURVEY_UNSUPPORTED = frozenset({"TROP", "SyntheticDiD", "TripleDifference"})
def _check_staggered_dgp_compat(
estimator: Any,
data_generator_kwargs: Optional[Dict[str, Any]],
) -> None:
"""Warn if a staggered estimator's settings don't match the default DGP."""
name = type(estimator).__name__
if name not in _STAGGERED_ESTIMATORS:
return
dgp_overrides = data_generator_kwargs or {}
cohort_periods = dgp_overrides.get("cohort_periods")
has_multi_cohort = cohort_periods is not None and len(set(cohort_periods)) >= 2
issues: List[str] = []
# Check control_group="not_yet_treated" (CS, SA)
cg = getattr(estimator, "control_group", "never_treated")
if cg == "not_yet_treated" and not has_multi_cohort:
issues.append(
f' - {name} has control_group="not_yet_treated" but the default '
f"DGP generates a single treatment cohort with never-treated "
f"controls. Power may not reflect the intended not-yet-treated "
f"design.\n"
f" Fix: pass data_generator_kwargs="
f'{{"cohort_periods": [2, 4], "never_treated_frac": 0.0}} '
f"(or a custom data_generator)."
)
# Check anticipation > 0 (all staggered)
antic = getattr(estimator, "anticipation", 0)
if antic > 0:
issues.append(
f" - {name} has anticipation={antic} but the default DGP does "
f"not model anticipatory effects. The estimator will look for "
f"treatment effects {antic} period(s) before the DGP generates "
f"them, biasing power estimates.\n"
f" Fix: supply a custom data_generator that shifts the "
f"effect onset."
)
# Check clean_control on StackedDiD
if name == "StackedDiD":
cc = getattr(estimator, "clean_control", "not_yet_treated")
if cc == "strict" and not has_multi_cohort:
issues.append(
' - StackedDiD has clean_control="strict" but the default '
"single-cohort DGP makes strict controls equivalent to "
"never-treated controls.\n"
" Fix: pass data_generator_kwargs="
'{"cohort_periods": [2, 4]} '
"to test true strict clean-control behavior."
)
if issues:
msg = (
f"Staggered power DGP mismatch for {name}. The default "
f"single-cohort DGP may not match the estimator "
f"configuration:\n" + "\n".join(issues)
)
warnings.warn(msg, UserWarning, stacklevel=2)
def _ddd_effective_n(
n_units: int, data_generator_kwargs: Optional[Dict[str, Any]]
) -> Optional[int]:
"""Return effective DDD sample size, or None if no rounding occurred."""
overrides = data_generator_kwargs or {}
if "n_per_cell" in overrides:
eff = overrides["n_per_cell"] * 8
else:
eff = max(2, n_units // 8) * 8
return eff if eff != n_units else None
def _check_ddd_dgp_compat(
n_units: int,
n_periods: int,
treatment_fraction: float,
treatment_period: int,
data_generator_kwargs: Optional[Dict[str, Any]],
) -> None:
"""Warn when simulation inputs don't match DDD's fixed 222 design."""
issues: List[str] = []
# DDD is a fixed 2-period factorial; n_periods and treatment_period are ignored
if n_periods != 2:
issues.append(
f"n_periods={n_periods} is ignored (DDD uses a fixed " f"2-period design: pre/post)"
)
if treatment_period != 1:
issues.append(
f"treatment_period={treatment_period} is ignored (DDD "
f"always treats in the second period)"
)
# DDD's 222 factorial has inherent 50% treatment fraction
if treatment_fraction != 0.5:
issues.append(
f"treatment_fraction={treatment_fraction} is ignored "
f"(DDD uses a balanced 222 factorial where 50% of "
f"groups are treated)"
)
# n_units rounding: n_per_cell = max(2, n_units // 8)
eff_n = _ddd_effective_n(n_units, data_generator_kwargs)
if eff_n is not None:
eff_n_per_cell = eff_n // 8
issues.append(
f"effective sample size is {eff_n} "
f"(n_per_cell={eff_n_per_cell} 8 cells), "
f"not the requested n_units={n_units}"
)
if issues:
warnings.warn(
"TripleDifference uses a fixed 222 factorial DGP "
"(group partition time). "
+ "; ".join(issues)
+ ". Pass a custom data_generator for non-standard DDD designs.",
UserWarning,
stacklevel=2,
)
def _check_sdid_placebo_data(
data: pd.DataFrame,
estimator: Any,
est_kwargs: Dict[str, Any],
) -> None:
"""Check SyntheticDiD placebo feasibility on realized data.
This catches infeasible designs on the custom-DGP path where the
pre-generation check (which uses ``n_units * treatment_fraction``)
cannot run because treatment allocation is determined by the DGP.
"""
vm = getattr(estimator, "variance_method", "placebo")
if vm != "placebo":
return
treat_col = est_kwargs.get("treatment", "treat")
unit_col = est_kwargs.get("unit", "unit")
if treat_col not in data.columns or unit_col not in data.columns:
return # fit will fail with a more specific error
unit_treat = data.groupby(unit_col)[treat_col].first()
n_treated = int(unit_treat.sum())
n_control = len(unit_treat) - n_treated
if n_control Dict[str, _EstimatorProfile]:
"""Lazily build and return the estimator registry."""
global _ESTIMATOR_REGISTRY # noqa: PLW0603
if _ESTIMATOR_REGISTRY is not None:
return _ESTIMATOR_REGISTRY
from diff_diff.prep import (
generate_ddd_data,
generate_did_data,
generate_factor_data,
generate_staggered_data,
)
_ESTIMATOR_REGISTRY = {
# --- Basic DiD group ---
"DifferenceInDifferences": _EstimatorProfile(
default_dgp=generate_did_data,
dgp_kwargs_builder=_basic_dgp_kwargs,
fit_kwargs_builder=_basic_fit_kwargs,
result_extractor=_extract_simple,
min_n=20,
),
"TwoWayFixedEffects": _EstimatorProfile(
default_dgp=generate_did_data,
dgp_kwargs_builder=_basic_dgp_kwargs,
fit_kwargs_builder=_twfe_fit_kwargs,
result_extractor=_extract_simple,
min_n=20,
),
"MultiPeriodDiD": _EstimatorProfile(
default_dgp=generate_did_data,
dgp_kwargs_builder=_basic_dgp_kwargs,
fit_kwargs_builder=_multiperiod_fit_kwargs,
result_extractor=_extract_multiperiod,
min_n=20,
),
# --- Staggered group ---
"CallawaySantAnna": _EstimatorProfile(
default_dgp=generate_staggered_data,
dgp_kwargs_builder=_staggered_dgp_kwargs,
fit_kwargs_builder=_staggered_fit_kwargs,
result_extractor=_extract_staggered,
min_n=40,
),
"SunAbraham": _EstimatorProfile(
default_dgp=generate_staggered_data,
dgp_kwargs_builder=_staggered_dgp_kwargs,
fit_kwargs_builder=_staggered_fit_kwargs,
result_extractor=_extract_staggered,
min_n=40,
),
"ImputationDiD": _EstimatorProfile(
default_dgp=generate_staggered_data,
dgp_kwargs_builder=_staggered_dgp_kwargs,
fit_kwargs_builder=_staggered_fit_kwargs,
result_extractor=_extract_staggered,
min_n=40,
),
"TwoStageDiD": _EstimatorProfile(
default_dgp=generate_staggered_data,
dgp_kwargs_builder=_staggered_dgp_kwargs,
fit_kwargs_builder=_staggered_fit_kwargs,
result_extractor=_extract_staggered,
min_n=40,
),
"StackedDiD": _EstimatorProfile(
default_dgp=generate_staggered_data,
dgp_kwargs_builder=_staggered_dgp_kwargs,
fit_kwargs_builder=_staggered_fit_kwargs,
result_extractor=_extract_staggered,
min_n=40,
),
"EfficientDiD": _EstimatorProfile(
default_dgp=generate_staggered_data,
dgp_kwargs_builder=_staggered_dgp_kwargs,
fit_kwargs_builder=_staggered_fit_kwargs,
result_extractor=_extract_staggered,
min_n=40,
),
# --- Factor model group ---
"TROP": _EstimatorProfile(
default_dgp=generate_factor_data,
dgp_kwargs_builder=_factor_dgp_kwargs,
fit_kwargs_builder=_trop_fit_kwargs,
result_extractor=_extract_simple,
min_n=30,
),
"SyntheticDiD": _EstimatorProfile(
default_dgp=generate_factor_data,
dgp_kwargs_builder=_factor_dgp_kwargs,
fit_kwargs_builder=_sdid_fit_kwargs,
result_extractor=_extract_simple,
min_n=30,
),
# --- Triple difference ---
"TripleDifference": _EstimatorProfile(
default_dgp=generate_ddd_data,
dgp_kwargs_builder=_ddd_dgp_kwargs,
fit_kwargs_builder=_ddd_fit_kwargs,
result_extractor=_extract_simple,
min_n=64,
),
}
return _ESTIMATOR_REGISTRY
@dataclass
class PowerResults:
"""
Results from analytical power analysis.
Attributes
----------
power : float
Statistical power (probability of rejecting H0 when effect exists).
mde : float
Minimum detectable effect size.
required_n : int
Required total sample size (treated + control).
effect_size : float
Effect size used in calculation.
alpha : float
Significance level.
alternative : str
Alternative hypothesis ('two-sided', 'greater', 'less').
n_treated : int
Number of treated units.
n_control : int
Number of control units.
n_pre : int
Number of pre-treatment periods.
n_post : int
Number of post-treatment periods.
sigma : float
Residual standard deviation.
rho : float
Intra-cluster correlation (for panel data).
deff : float
Survey design effect (variance inflation factor).
design : str
Study design type ('basic_did', 'panel', 'staggered').
"""
power: float
mde: float
required_n: int
effect_size: float
alpha: float
alternative: str
n_treated: int
n_control: int
n_pre: int
n_post: int
sigma: float
rho: float = 0.0
deff: float = 1.0
design: str = "basic_did"
def __repr__(self) -> str:
"""Concise string representation."""
return (
f"PowerResults(power={self.power:.3f}, mde={self.mde:.4f}, "
f"required_n={self.required_n})"
)
def summary(self) -> str:
"""
Generate a formatted summary of power analysis results.
Returns
-------
str
Formatted summary table.
"""
lines = [
"=" * 60,
"Power Analysis for Difference-in-Differences".center(60),
"=" * 60,
"",
f"{'Design:':>> print(f"Required N: {results.required_n}")
Calculate power for given sample and effect:
>>> results = pa.power(effect_size=0.5, n_treated=50, n_control=50, sigma=1.0)
>>> print(f"Power: {results.power:.1%}")
Notes
-----
The power calculations are based on the variance of the DiD estimator:
For basic 2x2 DiD:
Var(ATT) = sigma^2 * (1/n_treated_post + 1/n_treated_pre
+ 1/n_control_post + 1/n_control_pre)
For panel DiD with T periods:
Var(ATT) = sigma^2 * (1/(N_treated * T) + 1/(N_control * T))
* (1 + (T-1)*rho) / (1 + (T-1)*rho)
Where rho is the intra-cluster correlation coefficient.
References
----------
Bloom, H. S. (1995). "Minimum Detectable Effects."
Burlig, F., Preonas, L., & Woerman, M. (2020). "Panel Data and Experimental Design."
"""
def __init__(
self,
alpha: float = 0.05,
power: float = 0.80,
alternative: str = "two-sided",
):
if not 0 < alpha < 1:
raise ValueError("alpha must be between 0 and 1")
if not 0 < power < 1:
raise ValueError("power must be between 0 and 1")
if alternative not in ("two-sided", "greater", "less"):
raise ValueError("alternative must be 'two-sided', 'greater', or 'less'")
self.alpha = alpha
self.target_power = power
self.alternative = alternative
@staticmethod
def _validate_deff(deff: float) -> None:
"""Validate deff parameter and warn if < 1."""
if not np.isfinite(deff) or deff Tuple[float, float]:
"""Get z critical values for alpha and power."""
if self.alternative == "two-sided":
z_alpha = stats.norm.ppf(1 - self.alpha / 2)
else:
z_alpha = stats.norm.ppf(1 - self.alpha)
z_beta = stats.norm.ppf(self.target_power)
return z_alpha, z_beta
def _compute_variance(
self,
n_treated: int,
n_control: int,
n_pre: int,
n_post: int,
sigma: float,
rho: float = 0.0,
deff: float = 1.0,
design: str = "basic_did",
) -> float:
"""
Compute variance of the DiD estimator.
Parameters
----------
n_treated : int
Number of treated units.
n_control : int
Number of control units.
n_pre : int
Number of pre-treatment periods.
n_post : int
Number of post-treatment periods.
sigma : float
Residual standard deviation.
rho : float
Intra-cluster correlation (for panel data).
deff : float
Survey design effect (variance inflation factor). Not redundant
with ``rho``: ``rho`` models within-unit serial correlation
(Moulton factor), ``deff`` models survey clustering/weighting.
design : str
Study design type.
Returns
-------
float
Variance of the DiD estimator.
"""
if design == "basic_did":
# For basic 2x2 DiD, each cell has n_treated/2 or n_control/2 obs
# assuming balanced design
n_t_pre = n_treated # treated units in pre-period
n_t_post = n_treated # treated units in post-period
n_c_pre = n_control
n_c_post = n_control
variance = sigma**2 * (1 / n_t_post + 1 / n_t_pre + 1 / n_c_post + 1 / n_c_pre)
elif design == "panel":
# Panel DiD with multiple periods
# Account for serial correlation via ICC
T = n_pre + n_post
# Design effect for clustering
design_effect = 1 + (T - 1) * rho
# Base variance (as if independent)
base_var = sigma**2 * (1 / n_treated + 1 / n_control)
# Adjust for clustering (Moulton factor)
variance = base_var * design_effect / T
else:
raise ValueError(f"Unknown design: {design}")
# Survey design effect (multiplicative variance inflation)
variance *= deff
return variance
def power(
self,
effect_size: float,
n_treated: int,
n_control: int,
sigma: float,
n_pre: int = 1,
n_post: int = 1,
rho: float = 0.0,
deff: float = 1.0,
) -> PowerResults:
"""
Calculate statistical power for given effect size and sample.
Parameters
----------
effect_size : float
Expected treatment effect size.
n_treated : int
Number of treated units.
n_control : int
Number of control units.
sigma : float
Residual standard deviation.
n_pre : int, default=1
Number of pre-treatment periods.
n_post : int, default=1
Number of post-treatment periods.
rho : float, default=0.0
Intra-cluster correlation for panel data.
deff : float, default=1.0
Survey design effect (variance inflation factor). Not redundant
with ``rho``: ``rho`` models within-unit serial correlation,
``deff`` models survey clustering/weighting.
Returns
-------
PowerResults
Power analysis results.
Examples
--------
>>> pa = PowerAnalysis()
>>> results = pa.power(effect_size=2.0, n_treated=50, n_control=50, sigma=5.0)
>>> print(f"Power: {results.power:.1%}")
"""
self._validate_deff(deff)
T = n_pre + n_post
design = "panel" if T > 2 else "basic_did"
variance = self._compute_variance(
n_treated, n_control, n_pre, n_post, sigma, rho, deff=deff, design=design
)
se = np.sqrt(variance)
# Calculate power
if self.alternative == "two-sided":
z_alpha = stats.norm.ppf(1 - self.alpha / 2)
# Power = P(reject | effect) = P(|Z| > z_alpha | effect)
power_val = (
1
- stats.norm.cdf(z_alpha - effect_size / se)
+ stats.norm.cdf(-z_alpha - effect_size / se)
)
elif self.alternative == "greater":
z_alpha = stats.norm.ppf(1 - self.alpha)
power_val = 1 - stats.norm.cdf(z_alpha - effect_size / se)
else: # less
z_alpha = stats.norm.ppf(1 - self.alpha)
power_val = stats.norm.cdf(-z_alpha - effect_size / se)
# Also compute MDE and required N for reference
mde = self._compute_mde_from_se(se)
required_n = self._compute_required_n(
effect_size,
sigma,
n_pre,
n_post,
rho,
design,
n_treated / (n_treated + n_control),
deff=deff,
)
return PowerResults(
power=power_val,
mde=mde,
required_n=required_n,
effect_size=effect_size,
alpha=self.alpha,
alternative=self.alternative,
n_treated=n_treated,
n_control=n_control,
n_pre=n_pre,
n_post=n_post,
sigma=sigma,
rho=rho,
deff=deff,
design=design,
)
def _compute_mde_from_se(self, se: float) -> float:
"""Compute MDE given standard error."""
z_alpha, z_beta = self._get_critical_values()
return (z_alpha + z_beta) * se
def mde(
self,
n_treated: int,
n_control: int,
sigma: float,
n_pre: int = 1,
n_post: int = 1,
rho: float = 0.0,
deff: float = 1.0,
) -> PowerResults:
"""
Calculate minimum detectable effect given sample size.
The MDE is the smallest effect size that can be detected with the
specified power and significance level.
Parameters
----------
n_treated : int
Number of treated units.
n_control : int
Number of control units.
sigma : float
Residual standard deviation.
n_pre : int, default=1
Number of pre-treatment periods.
n_post : int, default=1
Number of post-treatment periods.
rho : float, default=0.0
Intra-cluster correlation for panel data.
deff : float, default=1.0
Survey design effect (variance inflation factor).
Returns
-------
PowerResults
Power analysis results including MDE.
Examples
--------
>>> pa = PowerAnalysis(power=0.80)
>>> results = pa.mde(n_treated=100, n_control=100, sigma=10.0)
>>> print(f"MDE: {results.mde:.2f}")
"""
self._validate_deff(deff)
T = n_pre + n_post
design = "panel" if T > 2 else "basic_did"
variance = self._compute_variance(
n_treated, n_control, n_pre, n_post, sigma, rho, deff=deff, design=design
)
se = np.sqrt(variance)
mde = self._compute_mde_from_se(se)
return PowerResults(
power=self.target_power,
mde=mde,
required_n=n_treated + n_control,
effect_size=mde,
alpha=self.alpha,
alternative=self.alternative,
n_treated=n_treated,
n_control=n_control,
n_pre=n_pre,
n_post=n_post,
sigma=sigma,
rho=rho,
deff=deff,
design=design,
)
def _compute_required_n(
self,
effect_size: float,
sigma: float,
n_pre: int,
n_post: int,
rho: float,
design: str,
treat_frac: float = 0.5,
deff: float = 1.0,
) -> int:
"""Compute required sample size for given effect.
Note: this method has its own formula independent of _compute_variance,
so deff must be applied here separately (not double-counting).
"""
# Handle edge case of zero effect size
if effect_size == 0:
return MAX_SAMPLE_SIZE # Can't detect zero effect
z_alpha, z_beta = self._get_critical_values()
T = n_pre + n_post
if design == "basic_did":
n_total = (
2
* sigma**2
* (z_alpha + z_beta) ** 2
/ (effect_size**2 * treat_frac * (1 - treat_frac))
)
else: # panel
design_effect = 1 + (T - 1) * rho
n_total = (
2
* sigma**2
* (z_alpha + z_beta) ** 2
* design_effect
/ (effect_size**2 * treat_frac * (1 - treat_frac) * T)
)
# Survey design effect (multiplicative sample size inflation)
n_total *= deff
# Handle infinity case (extremely small effect)
if np.isinf(n_total):
return MAX_SAMPLE_SIZE
return max(4, int(np.ceil(n_total))) # At least 4 units
def sample_size(
self,
effect_size: float,
sigma: float,
n_pre: int = 1,
n_post: int = 1,
rho: float = 0.0,
treat_frac: float = 0.5,
deff: float = 1.0,
) -> PowerResults:
"""
Calculate required sample size to detect given effect.
Parameters
----------
effect_size : float
Treatment effect to detect.
sigma : float
Residual standard deviation.
n_pre : int, default=1
Number of pre-treatment periods.
n_post : int, default=1
Number of post-treatment periods.
rho : float, default=0.0
Intra-cluster correlation for panel data.
treat_frac : float, default=0.5
Fraction of units assigned to treatment.
deff : float, default=1.0
Survey design effect (variance inflation factor).
Returns
-------
PowerResults
Power analysis results including required sample size.
Examples
--------
>>> pa = PowerAnalysis(power=0.80)
>>> results = pa.sample_size(effect_size=5.0, sigma=10.0)
>>> print(f"Required N: {results.required_n}")
"""
self._validate_deff(deff)
T = n_pre + n_post
design = "panel" if T > 2 else "basic_did"
n_total = self._compute_required_n(
effect_size, sigma, n_pre, n_post, rho, design, treat_frac, deff=deff
)
n_treated = max(2, int(np.ceil(n_total * treat_frac)))
n_control = max(2, n_total - n_treated)
n_total = n_treated + n_control
# Compute actual power achieved
variance = self._compute_variance(
n_treated, n_control, n_pre, n_post, sigma, rho, deff=deff, design=design
)
se = np.sqrt(variance)
mde = self._compute_mde_from_se(se)
return PowerResults(
power=self.target_power,
mde=mde,
required_n=n_total,
effect_size=effect_size,
alpha=self.alpha,
alternative=self.alternative,
n_treated=n_treated,
n_control=n_control,
n_pre=n_pre,
n_post=n_post,
sigma=sigma,
rho=rho,
deff=deff,
design=design,
)
def power_curve(
self,
n_treated: int,
n_control: int,
sigma: float,
effect_sizes: Optional[List[float]] = None,
n_pre: int = 1,
n_post: int = 1,
rho: float = 0.0,
deff: float = 1.0,
) -> pd.DataFrame:
"""
Compute power for a range of effect sizes.
Parameters
----------
n_treated : int
Number of treated units.
n_control : int
Number of control units.
sigma : float
Residual standard deviation.
effect_sizes : list of float, optional
Effect sizes to evaluate. If None, uses a range from 0 to 3*MDE.
n_pre : int, default=1
Number of pre-treatment periods.
n_post : int, default=1
Number of post-treatment periods.
rho : float, default=0.0
Intra-cluster correlation.
deff : float, default=1.0
Survey design effect (variance inflation factor).
Returns
-------
pd.DataFrame
DataFrame with columns 'effect_size' and 'power'.
Examples
--------
>>> pa = PowerAnalysis()
>>> curve = pa.power_curve(n_treated=50, n_control=50, sigma=5.0)
>>> print(curve)
"""
# First get MDE to determine default range
mde_result = self.mde(n_treated, n_control, sigma, n_pre, n_post, rho, deff=deff)
if effect_sizes is None:
# Generate range from 0 to 2*MDE
effect_sizes = np.linspace(0, 2.5 * mde_result.mde, 50).tolist()
powers = []
for es in effect_sizes:
result = self.power(
effect_size=es,
n_treated=n_treated,
n_control=n_control,
sigma=sigma,
n_pre=n_pre,
n_post=n_post,
rho=rho,
deff=deff,
)
powers.append(result.power)
return pd.DataFrame({"effect_size": effect_sizes, "power": powers})
def sample_size_curve(
self,
effect_size: float,
sigma: float,
sample_sizes: Optional[List[int]] = None,
n_pre: int = 1,
n_post: int = 1,
rho: float = 0.0,
treat_frac: float = 0.5,
deff: float = 1.0,
) -> pd.DataFrame:
"""
Compute power for a range of sample sizes.
Parameters
----------
effect_size : float
Treatment effect size.
sigma : float
Residual standard deviation.
sample_sizes : list of int, optional
Total sample sizes to evaluate. If None, uses sensible range.
n_pre : int, default=1
Number of pre-treatment periods.
n_post : int, default=1
Number of post-treatment periods.
rho : float, default=0.0
Intra-cluster correlation.
treat_frac : float, default=0.5
Fraction assigned to treatment.
deff : float, default=1.0
Survey design effect (variance inflation factor).
Returns
-------
pd.DataFrame
DataFrame with columns 'sample_size' and 'power'.
"""
# Get required N to determine default range
required = self.sample_size(effect_size, sigma, n_pre, n_post, rho, treat_frac, deff=deff)
if sample_sizes is None:
min_n = max(10, required.required_n // 4)
max_n = required.required_n * 2
sample_sizes = list(range(min_n, max_n + 1, max(1, (max_n - min_n) // 50)))
powers = []
for n in sample_sizes:
n_treated = max(2, int(n * treat_frac))
n_control = max(2, n - n_treated)
result = self.power(
effect_size=effect_size,
n_treated=n_treated,
n_control=n_control,
sigma=sigma,
n_pre=n_pre,
n_post=n_post,
rho=rho,
deff=deff,
)
powers.append(result.power)
return pd.DataFrame({"sample_size": sample_sizes, "power": powers})
def simulate_power(
estimator: Any,
n_units: int = 100,
n_periods: int = 4,
treatment_effect: float = 5.0,
treatment_fraction: float = 0.5,
treatment_period: int = 2,
sigma: float = 1.0,
n_simulations: int = 500,
alpha: float = 0.05,
effect_sizes: Optional[List[float]] = None,
seed: Optional[int] = None,
data_generator: Optional[Callable] = None,
data_generator_kwargs: Optional[Dict[str, Any]] = None,
estimator_kwargs: Optional[Dict[str, Any]] = None,
result_extractor: Optional[Callable] = None,
progress: bool = True,
survey_config: Optional[SurveyPowerConfig] = None,
) -> SimulationPowerResults:
"""
Estimate power using Monte Carlo simulation.
This function simulates datasets with known treatment effects and estimates
power as the fraction of simulations where the null hypothesis is rejected.
Most built-in estimators are supported via an internal registry that selects
the appropriate data-generating process and fit signature automatically.
Parameters
----------
estimator : estimator object
DiD estimator to use (e.g., DifferenceInDifferences, CallawaySantAnna).
n_units : int, default=100
Number of units per simulation.
n_periods : int, default=4
Number of time periods.
treatment_effect : float, default=5.0
True treatment effect to simulate.
treatment_fraction : float, default=0.5
Fraction of units that are treated.
treatment_period : int, default=2
First post-treatment period (0-indexed).
sigma : float, default=1.0
Residual standard deviation (noise level).
n_simulations : int, default=500
Number of Monte Carlo simulations.
alpha : float, default=0.05
Significance level for hypothesis tests.
effect_sizes : list of float, optional
Multiple effect sizes to evaluate for power curve.
If None, uses only treatment_effect.
seed : int, optional
Random seed for reproducibility.
data_generator : callable, optional
Custom data generation function. When provided, bypasses the
registry DGP and calls this function with the standard kwargs
(n_units, n_periods, treatment_effect, etc.).
data_generator_kwargs : dict, optional
Additional keyword arguments for data generator.
estimator_kwargs : dict, optional
Additional keyword arguments for estimator.fit().
result_extractor : callable, optional
Custom function to extract results from the estimator output.
Takes the estimator result object and returns a tuple of
``(att, se, p_value, conf_int)``. Useful for unregistered
estimators with non-standard result schemas.
progress : bool, default=True
Whether to print progress updates.
survey_config : SurveyPowerConfig, optional
When provided, generates survey-structured data via
``generate_survey_did_data`` and injects ``SurveyDesign`` into
estimator ``fit()``. Mutually exclusive with ``data_generator``.
Supported estimators: DiD, TWFE, MultiPeriod, CS, SA, Imputation,
TwoStage, Stacked, Efficient. Unsupported: TROP, SyntheticDiD,
TripleDifference. ``heterogeneous_te_by_strata`` must be False.
Returns
-------
SimulationPowerResults
Simulation-based power analysis results.
Examples
--------
Basic power simulation:
>>> from diff_diff import DifferenceInDifferences, simulate_power
>>> did = DifferenceInDifferences()
>>> results = simulate_power(
... estimator=did,
... n_units=100,
... treatment_effect=5.0,
... sigma=5.0,
... n_simulations=500,
... seed=42
... )
>>> print(f"Power: {results.power:.1%}")
Power curve over multiple effect sizes:
>>> results = simulate_power(
... estimator=did,
... effect_sizes=[1.0, 2.0, 3.0, 5.0, 7.0],
... n_simulations=200,
... seed=42
... )
>>> print(results.power_curve_df())
With Callaway-Sant'Anna (auto-detected, no custom DGP needed):
>>> from diff_diff import CallawaySantAnna
>>> cs = CallawaySantAnna()
>>> results = simulate_power(cs, n_simulations=200, seed=42)
Notes
-----
The simulation approach:
1. Generate data with known treatment effect
2. Fit the estimator and record the p-value
3. Repeat n_simulations times
4. Power = fraction of simulations where p-value < alpha
References
----------
Burlig, F., Preonas, L., & Woerman, M. (2020). "Panel Data and Experimental Design."
"""
rng = np.random.default_rng(seed)
estimator_name = type(estimator).__name__
registry = _get_registry()
profile = registry.get(estimator_name)
# If no profile and no custom data_generator, raise
if profile is None and data_generator is None:
raise ValueError(
f"Estimator '{estimator_name}' not in registry. "
f"Provide a custom data_generator and estimator_kwargs "
f"(the full dict of keyword arguments for estimator.fit(), "
f"e.g. dict(outcome='y', treatment='treat', time='period'))."
)
# When a custom data_generator is provided, bypass registry DGP
use_custom_dgp = data_generator is not None
use_survey_dgp = survey_config is not None
# --- Survey config validation ---
if use_survey_dgp:
assert survey_config is not None # for type narrowing
if estimator_name in _SURVEY_UNSUPPORTED:
raise ValueError(
f"survey_config is not supported with {estimator_name}. "
f"generate_survey_did_data produces staggered cohort data "
f"incompatible with this estimator's DGP. Use the custom "
f"data_generator path for survey power with {estimator_name}."
)
if use_custom_dgp:
raise ValueError(
"survey_config and data_generator are mutually exclusive. "
"survey_config uses generate_survey_did_data internally."
)
if treatment_period < 1:
raise ValueError(
f"treatment_period must be >= 1 with survey_config "
f"(need at least one pre-treatment period), got {treatment_period}."
)
if estimator_name not in _SURVEY_FIT_BUILDERS:
raise ValueError(
f"No survey power profile for {estimator_name}. "
f"Supported: {sorted(_SURVEY_FIT_BUILDERS.keys())}."
)
if survey_config.heterogeneous_te_by_strata:
raise ValueError(
"heterogeneous_te_by_strata=True is not supported with "
"simulation power analysis. The DGP's population ATT diverges "
"from the input treatment_effect under heterogeneous effects, "
"which would make bias/coverage/RMSE metrics misleading."
)
data_gen_kwargs = data_generator_kwargs or {}
est_kwargs = estimator_kwargs or {}
# Block survey_design in estimator_kwargs when survey_config is active.
# Custom survey design overrides go through SurveyPowerConfig.survey_design.
if use_survey_dgp and "survey_design" in est_kwargs:
raise ValueError(
"estimator_kwargs cannot contain 'survey_design' when survey_config "
"is set. To override the auto-built SurveyDesign, pass it via "
"SurveyPowerConfig(survey_design=...)."
)
# Block survey-config-managed keys in data_generator_kwargs
if use_survey_dgp and data_gen_kwargs:
collisions = _SURVEY_CONFIG_KEYS & set(data_gen_kwargs)
if collisions:
raise ValueError(
f"data_generator_kwargs contains keys managed by survey_config: "
f"{sorted(collisions)}. Set these on SurveyPowerConfig instead."
)
# Block DGP params that make realized ATT diverge from scalar input,
# which would misstate bias/coverage/RMSE (same rationale as
# heterogeneous_te_by_strata rejection above).
te_interaction = data_gen_kwargs.get("te_covariate_interaction", 0.0)
if te_interaction != 0.0:
raise ValueError(
f"te_covariate_interaction={te_interaction} is not supported "
f"with survey_config. The DGP's population ATT diverges from "
f"the input treatment_effect under covariate-interaction "
f"heterogeneity, which would make bias/coverage/RMSE misleading."
)
# Enforce panel-mode alignment between DGP and estimator.
# Runs even with empty data_gen_kwargs to catch CS(panel=False) + default DGP.
if use_survey_dgp:
dgp_panel = data_gen_kwargs.get("panel", True)
est_panel = getattr(estimator, "panel", True)
if not dgp_panel:
if estimator_name != "CallawaySantAnna":
raise ValueError(
f"panel=False (repeated cross-sections) is not supported "
f"with {estimator_name} under survey_config. Only "
f"CallawaySantAnna supports repeated cross-sections."
)
if est_panel:
raise ValueError(
"data_generator_kwargs has panel=False but "
"CallawaySantAnna.panel=True. Use "
"CallawaySantAnna(panel=False) to match."
)
elif estimator_name == "CallawaySantAnna" and not est_panel:
raise ValueError(
"CallawaySantAnna(panel=False) requires "
"data_generator_kwargs={'panel': False} to generate "
"repeated cross-section data."
)
# Reject estimator settings that require a multi-cohort DGP.
# survey_config hard-codes a single-cohort DGP and blocks
# cohort_periods/never_treated_frac overrides.
control_group = getattr(estimator, "control_group", "never_treated")
clean_control = getattr(estimator, "clean_control", None)
if control_group in ("not_yet_treated", "last_cohort"):
raise ValueError(
f"survey_config does not support control_group='{control_group}' "
f"(requires multi-cohort DGP). Use the custom data_generator "
f"path for survey power with this control-group design."
)
if clean_control == "strict":
raise ValueError(
f"survey_config does not support clean_control='strict' "
f"(requires multi-cohort DGP). Use the custom data_generator "
f"path for survey power with strict clean controls."
)
# SyntheticDiD placebo variance requires n_control > n_treated.
# Check after merging data_generator_kwargs so overrides of n_treated
# are accounted for.
if estimator_name == "SyntheticDiD" and not use_custom_dgp:
vm = getattr(estimator, "variance_method", "placebo")
effective_n_treated = data_gen_kwargs.get(
"n_treated", max(1, int(n_units * treatment_fraction))
)
n_control = n_units - effective_n_treated
if vm == "placebo" and n_control 0:
pct = (sim + effect_idx * n_simulations) / (len(effect_sizes) * n_simulations)
print(f" Simulation progress: {pct:.0%}")
sim_seed = rng.integers(0, 2**31)
# --- Generate data ---
if use_survey_dgp:
assert survey_config is not None
assert _generate_survey_did_data is not None
dgp_kwargs = _survey_dgp_kwargs(
n_units=n_units,
n_periods=n_periods,
treatment_effect=effect,
treatment_fraction=treatment_fraction,
treatment_period=treatment_period,
sigma=sigma,
survey_config=survey_config,
)
dgp_kwargs.update(data_gen_kwargs)
dgp_kwargs.pop("seed", None)
data = _generate_survey_did_data(seed=sim_seed, **dgp_kwargs)
# Derive columns for non-staggered estimators.
# Survey DGP's `treated` is time-varying (1{g>0, t>=g}); basic/TWFE/
# MultiPeriod need a time-invariant group indicator (`ever_treated`).
if estimator_name not in _STAGGERED_ESTIMATORS:
data["ever_treated"] = (data["first_treat"] > 0).astype(int)
# Basic/TWFE also need a `post` period indicator.
if estimator_name in _SURVEY_POST_ESTIMATORS:
data["post"] = (data["period"] >= treatment_period + 1).astype(int)
# Collect DGP truth for metadata
dgp_truth = data.attrs.get("dgp_truth", {})
if dgp_truth:
kish = dgp_truth.get("deff_kish")
icc_r = dgp_truth.get("icc_realized")
if kish is not None:
deff_values.append(kish)
if icc_r is not None:
icc_values.append(icc_r)
elif use_custom_dgp:
assert data_generator is not None
data = data_generator(
n_units=n_units,
n_periods=n_periods,
treatment_effect=effect,
treatment_fraction=treatment_fraction,
treatment_period=treatment_period,
noise_sd=sigma,
seed=sim_seed,
**data_gen_kwargs,
)
else:
assert profile is not None
dgp_kwargs = profile.dgp_kwargs_builder(
n_units=n_units,
n_periods=n_periods,
treatment_effect=effect,
treatment_fraction=treatment_fraction,
treatment_period=treatment_period,
sigma=sigma,
)
dgp_kwargs.update(data_gen_kwargs)
dgp_kwargs.pop("seed", None)
data = profile.default_dgp(seed=sim_seed, **dgp_kwargs)
# Check SDID placebo feasibility on realized data (custom DGP path)
if effect_idx == 0 and sim == 0 and estimator_name == "SyntheticDiD":
_check_sdid_placebo_data(data, estimator, est_kwargs)
try:
# --- Fit estimator ---
if use_survey_dgp:
assert survey_config is not None
fit_builder = _SURVEY_FIT_BUILDERS[estimator_name]
fit_kwargs = fit_builder(
data, n_units, n_periods, treatment_period, survey_config
)
fit_kwargs.update(est_kwargs)
elif profile is not None and not use_custom_dgp:
fit_kwargs = profile.fit_kwargs_builder(
data, n_units, n_periods, treatment_period
)
fit_kwargs.update(est_kwargs)
else:
# Custom DGP fallback: use registry fit kwargs if available,
# otherwise use basic DiD signature
if profile is not None:
fit_kwargs = profile.fit_kwargs_builder(
data, n_units, n_periods, treatment_period
)
fit_kwargs.update(est_kwargs)
else:
fit_kwargs = dict(est_kwargs)
result = estimator.fit(data, **fit_kwargs)
# --- Extract results ---
if profile is not None:
att, se, p_val, ci = profile.result_extractor(result)
elif result_extractor is not None:
att, se, p_val, ci = result_extractor(result)
else:
att = result.att if hasattr(result, "att") else result.avg_att
se = result.se if hasattr(result, "se") else result.avg_se
p_val = result.p_value if hasattr(result, "p_value") else result.avg_p_value
ci = result.conf_int if hasattr(result, "conf_int") else result.avg_conf_int
# NaN p-value treat as non-rejection
rejected = bool(p_val < alpha) if not np.isnan(p_val) else False
estimates.append(att)
ses.append(se)
p_values.append(p_val)
rejections.append(rejected)
ci_contains_true.append(ci[0] str:
return (
f"SimulationMDEResults(mde={self.mde:.4f}, "
f"power_at_mde={self.power_at_mde:.3f}, "
f"n_steps={self.n_steps})"
)
def summary(self) -> str:
"""Generate a formatted summary."""
lines = [
"=" * 65,
"Simulation-Based MDE Results".center(65),
"=" * 65,
"",
f"{'Estimator:':