"""

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

from diff_diff.results_base import Diagnostic

# 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 <= 0:

raise ValueError(f"weight_cv must be finite and > 0, got {self.weight_cv}")

if self.weight_variation != "moderate":

raise ValueError(

"Cannot specify both weight_cv and a non-default "

"weight_variation. weight_cv overrides weight_variation."

)

if not np.isfinite(self.psu_period_factor) or self.psu_period_factor < 0:

raise ValueError(

f"psu_period_factor must be finite and >= 0, got {self.psu_period_factor}"

)

if self.fpc_per_stratum < self.psu_per_stratum:

raise ValueError(

f"fpc_per_stratum ({self.fpc_per_stratum}) must be >= "

f"psu_per_stratum ({self.psu_per_stratum})"

)

def _build_survey_design(self) -> 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,

)

def _ddd_panel_dgp_kwargs(

n_units: int,

n_periods: int,

treatment_effect: float,

treatment_fraction: float,

treatment_period: int,

sigma: float,

) -> Dict[str, Any]:

# Panel DDD DGP (n_periods > 2). `n_units` maps directly (no 8-cell //8

# rounding). `treatment_fraction` is intentionally NOT mapped — DDD is a

# balanced factorial, so group_frac/partition_frac are left at the DGP

# default 0.5. Omitting group_frac/partition_frac here also keeps them out

# of the _PROTECTED_DGP_KEYS collision check, so a user can still override

# the group/partition split via data_generator_kwargs.

return dict(

n_units=n_units,

n_periods=n_periods,

treatment_period=treatment_period,

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", post="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", post="time")

def _ddd_panel_fit_kwargs(

data: pd.DataFrame,

n_units: int,

n_periods: int,

treatment_period: int,

) -> Dict[str, Any]:

# Panel DDD: time="post" is generate_ddd_panel_data's derived binary

# pre/post indicator (vs the cross-sectional "time"). Clustering is NOT a

# fit kwarg — it resolves from the estimator's cluster="unit" attribute

# against the DGP's "unit" column.

return dict(outcome="outcome", group="group", partition="partition", post="post")

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 * post``,

so passing the post-only flag would make that interaction rank-deficient.

"""

return dict(

outcome="outcome",

treatment="ever_treated",

post="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 control_group (pre-M-095: clean_control) on StackedDiD

if name == "StackedDiD":

cc = getattr(estimator, "control_group", "not_yet_treated")

if cc == "strict" and not has_multi_cohort:

issues.append(

' - StackedDiD has control_group="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 _ddd_panel_cells_populated(n: int, group_frac: float, partition_frac: float) -> bool:

"""Whether ``generate_ddd_panel_data`` would populate all 4 (group,partition)

cells at ``n`` units. Mirrors the rounded stratified allocation + non-empty

validation in ``prep_dgp.generate_ddd_panel_data`` (kept in lockstep)."""

if n < 4:

return False

n_g1 = int(round(n * group_frac))

n_g0 = n - n_g1

n_p1_g0 = int(round(n_g0 * partition_frac))

n_p1_g1 = int(round(n_g1 * partition_frac))

cells = (n_g0 - n_p1_g0, n_p1_g0, n_g1 - n_p1_g1, n_p1_g1)

return min(cells) >= 1

def _ddd_panel_viable_min_n(

group_frac: float, partition_frac: float, floor: int = 16, search_max: int = 100000

) -> int:

"""Smallest ``n_units`` for which ``generate_ddd_panel_data`` populates all

four (group,partition) cells under the given split, floored at ``floor``.

For the balanced default (0.5/0.5) this is 4, so the result is ``floor``;

skewed splits (e.g. 0.1/0.1) need more units before every cell is non-empty,

so the sample-size search must bracket above this value (the registry

documents group_frac/partition_frac as data_generator_kwargs overrides)."""

# Validate up front (matching generate_ddd_panel_data) so an out-of-range

# split raises the same clear message here — before the bracketing logic can

# surface a misleading "n_range below the minimum" error downstream.

if not (0.0 < group_frac < 1.0):

raise ValueError(f"group_frac must be in (0, 1); got {group_frac}.")

if not (0.0 < partition_frac < 1.0):

raise ValueError(f"partition_frac must be in (0, 1); got {partition_frac}.")

for n in range(4, search_max + 1):

if _ddd_panel_cells_populated(n, group_frac, partition_frac):

return max(floor, n)

raise ValueError(

f"No panel-DDD sample size <= {search_max} populates all four "

f"(group, partition) cells for group_frac={group_frac}, "

f"partition_frac={partition_frac}; move the split closer to 0.5."

)

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 2×2×2 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 2×2×2 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 2×2×2 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 2×2×2 factorial DGP "

"(group × partition × time). "

+ "; ".join(issues)

+ ". Pass a custom data_generator for non-standard DDD designs.",

UserWarning,

stacklevel=2,

)

def _check_ddd_panel_dgp_compat(

estimator: Any,

treatment_fraction: float,

data_generator_kwargs: Optional[Dict[str, Any]],

) -> None:

"""Compat checks for the panel DDD power path (``n_periods > 2``).

Unlike the cross-sectional ``_check_ddd_dgp_compat``, ``n_periods`` and

``treatment_period`` are honored here (no warning). ``treatment_fraction``

is still inert (the panel DGP is a balanced 2×2×2). The key addition is the

clustering caveat: ``generate_ddd_panel_data`` has within-unit serial

correlation, so unclustered SEs overstate power.

"""

overrides = data_generator_kwargs or {}

if "n_per_cell" in overrides:

raise ValueError(

"data_generator_kwargs contains 'n_per_cell', a cross-sectional "

"generate_ddd_data parameter. The panel DDD power path "

"(n_periods > 2) uses generate_ddd_panel_data, which sizes the panel "

"by n_units directly. Control the design via n_units and the "

"group_frac / partition_frac data_generator_kwargs instead."

)

if treatment_fraction != 0.5:

warnings.warn(

f"treatment_fraction={treatment_fraction} is ignored for "

f"TripleDifference power: generate_ddd_panel_data uses a balanced "

f"2×2×2 design (group_frac=partition_frac=0.5). Pass group_frac / "

f"partition_frac via data_generator_kwargs to vary the split.",

UserWarning,

stacklevel=2,

)

# The panel DGP hard-names its unit column "unit", and TripleDifference

# resolves clustering from `self.cluster` against a same-named data column,

# so on this auto-DGP path the only correct value is the literal "unit" —

# this is NOT a general clustering check.

if getattr(estimator, "cluster", None) != "unit":

warnings.warn(

"TripleDifference power on the panel DGP (n_periods > 2) has "

"within-unit serial correlation, so unclustered standard errors are "

"anti-conservative and overstate power. Construct the estimator as "

'TripleDifference(cluster="unit") so the reported power reflects '

"cluster-robust (Liang-Zeger CR1) inference.",

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 <= n_treated:

raise ValueError(

f"SyntheticDiD placebo variance requires more control than "

f"treated units, but the generated data has n_control={n_control}, "

f"n_treated={n_treated}. Either adjust your data_generator so that "

f"n_control > n_treated, or use "

f"SyntheticDiD(variance_method='bootstrap') (paper-faithful refit; "

f"~5-30x slower than placebo) or SyntheticDiD(variance_method='jackknife')."

)

# -- Registry construction (deferred to avoid import-time cost) ---------------

_ESTIMATOR_REGISTRY: Optional[Dict[str, _EstimatorProfile]] = None

def _get_registry() -> 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

def _ddd_panel_profile() -> "_EstimatorProfile":

"""Profile for panel DDD power simulations (n_periods > 2).

Routes TripleDifference power analysis to ``generate_ddd_panel_data``

(which honors ``n_periods``/``treatment_period``) instead of the

cross-sectional 2x2x2 ``generate_ddd_data``. See ``simulate_power``.

"""

from diff_diff.prep import generate_ddd_panel_data

# min_n=16 is intentionally lower than the cross-sectional DDD min_n=64

# (which counts 8 cells x n_per_cell); the panel DGP maps n_units directly

# and only requires n_units >= 4, so 16 gives >=1 unit per (group,partition)

# cell with margin.

return _EstimatorProfile(

default_dgp=generate_ddd_panel_data,

dgp_kwargs_builder=_ddd_panel_dgp_kwargs,

fit_kwargs_builder=_ddd_panel_fit_kwargs,

result_extractor=_extract_simple,

min_n=16,

)

@dataclass

class PowerResults(Diagnostic):

"""

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.