From 6ca2c26ec13bd9059803601a37c31f1c892c6dd0 Mon Sep 17 00:00:00 2001 From: igerber Date: Sun, 12 Apr 2026 12:22:09 -0400 Subject: [PATCH 01/12] feat: add survey-aware power analysis (SurveyPowerConfig + deff) Connect survey infrastructure to power module: - SurveyPowerConfig dataclass for simulation-based survey power - survey_config param on simulate_power/mde/sample_size swaps DGP to generate_survey_did_data and injects SurveyDesign into fit() - deff param on closed-form PowerAnalysis methods and convenience functions - 9 estimators supported, 3 blocked with clear error (factor model DGPs) - DGP truth (Kish DEFF, realized ICC) reported in simulation results Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/__init__.py | 2 + diff_diff/power.py | 498 +++++++++++++++++++++++++++++++++++++++--- tests/test_power.py | 306 ++++++++++++++++++++++++++ 3 files changed, 780 insertions(+), 26 deletions(-) diff --git a/diff_diff/__init__.py b/diff_diff/__init__.py index 03a07e6e6..c6cc6ee2f 100644 --- a/diff_diff/__init__.py +++ b/diff_diff/__init__.py @@ -63,6 +63,7 @@ SimulationMDEResults, SimulationPowerResults, SimulationSampleSizeResults, + SurveyPowerConfig, compute_mde, compute_power, compute_sample_size, @@ -368,6 +369,7 @@ "SimulationMDEResults", "SimulationPowerResults", "SimulationSampleSizeResults", + "SurveyPowerConfig", "compute_mde", "compute_power", "compute_sample_size", diff --git a/diff_diff/power.py b/diff_diff/power.py index 8ad20ec3f..2ca1ce514 100644 --- a/diff_diff/power.py +++ b/diff_diff/power.py @@ -49,6 +49,118 @@ class _EstimatorProfile: 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 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.weight_cv is not None and self.weight_cv <= 0: + raise ValueError(f"weight_cv must be > 0, got {self.weight_cv}") + 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 user-supplied SurveyDesign or auto-build from DGP column names.""" + 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 ----------------------------------------------------- @@ -203,6 +315,114 @@ def _sdid_fit_kwargs( ) +# -- 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.""" + return dict( + outcome="outcome", + treatment="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="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="treated", + 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 -------------------------------------------------------- @@ -252,6 +472,28 @@ def _first(r: Any, *attrs: str, default: Any = _nan) -> Any: } ) +# 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 ---------------------------------------- @@ -266,6 +508,22 @@ def _first(r: Any, *attrs: str, default: Any = _nan) -> Any: } ) +# 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, @@ -567,6 +825,8 @@ class PowerResults: 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'). """ @@ -583,6 +843,7 @@ class PowerResults: n_post: int sigma: float rho: float = 0.0 + deff: float = 1.0 design: str = "basic_did" def __repr__(self) -> str: @@ -624,6 +885,7 @@ def summary(self) -> str: "-" * 60, f"{'Residual SD (sigma):':<30} {self.sigma:>10.4f}", f"{'Intra-cluster correlation:':<30} {self.rho:>10.4f}", + *([f"{'Design effect (DEFF):':<30} {self.deff:>10.4f}"] if self.deff != 1.0 else []), "", "-" * 60, "Power Analysis Results".center(60), @@ -662,6 +924,7 @@ def to_dict(self) -> Dict[str, Any]: "n_post": self.n_post, "sigma": self.sigma, "rho": self.rho, + "deff": self.deff, "design": self.design, } @@ -735,6 +998,9 @@ class SimulationPowerResults: rmse: float = field(init=False) simulation_results: Optional[List[Dict[str, Any]]] = field(default=None, repr=False) effective_n_units: Optional[int] = None + survey_config: Optional[Any] = field(default=None, repr=False) + mean_deff: Optional[float] = None + mean_icc_realized: Optional[float] = None def __post_init__(self): """Compute derived statistics.""" @@ -789,6 +1055,21 @@ def summary(self) -> str: lines.append( f"{'Effective sample size:':<35} {self.effective_n_units}" f" (DDD grid-rounded)" ) + if self.survey_config is not None: + lines.extend( + [ + "", + "-" * 65, + "Survey Design".center(65), + "-" * 65, + f"{'Strata:':<35} {self.survey_config.n_strata}", + f"{'PSUs per stratum:':<35} {self.survey_config.psu_per_stratum}", + ] + ) + if self.mean_deff is not None: + lines.append(f"{'Mean Kish DEFF:':<35} {self.mean_deff:.4f}") + if self.mean_icc_realized is not None: + lines.append(f"{'Mean realized ICC:':<35} {self.mean_icc_realized:.4f}") lines.append("=" * 65) return "\n".join(lines) @@ -822,6 +1103,8 @@ def to_dict(self) -> Dict[str, Any]: "alpha": self.alpha, "estimator_name": self.estimator_name, "effective_n_units": self.effective_n_units, + "mean_deff": self.mean_deff, + "mean_icc_realized": self.mean_icc_realized, } return d @@ -921,6 +1204,18 @@ def __init__( self.target_power = power self.alternative = alternative + @staticmethod + def _validate_deff(deff: float) -> None: + """Validate deff parameter and warn if < 1.""" + if deff <= 0: + raise ValueError(f"deff must be > 0, got {deff}") + if deff < 1.0: + warnings.warn( + f"deff={deff:.4f} < 1.0 implies net variance reduction " + f"(e.g., from stratification). This is valid but unusual.", + stacklevel=3, + ) + def _get_critical_values(self) -> Tuple[float, float]: """Get z critical values for alpha and power.""" if self.alternative == "two-sided": @@ -938,6 +1233,7 @@ def _compute_variance( n_post: int, sigma: float, rho: float = 0.0, + deff: float = 1.0, design: str = "basic_did", ) -> float: """ @@ -957,6 +1253,10 @@ def _compute_variance( 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. @@ -990,6 +1290,9 @@ def _compute_variance( else: raise ValueError(f"Unknown design: {design}") + # Survey design effect (multiplicative variance inflation) + variance *= deff + return variance def power( @@ -1001,6 +1304,7 @@ def power( 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. @@ -1021,6 +1325,10 @@ def power( 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 ------- @@ -1033,10 +1341,13 @@ def power( >>> 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, design) + variance = self._compute_variance( + n_treated, n_control, n_pre, n_post, sigma, rho, deff=deff, design=design + ) se = np.sqrt(variance) # Calculate power @@ -1058,7 +1369,14 @@ def power( # 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) + effect_size, + sigma, + n_pre, + n_post, + rho, + design, + n_treated / (n_treated + n_control), + deff=deff, ) return PowerResults( @@ -1074,6 +1392,7 @@ def power( n_post=n_post, sigma=sigma, rho=rho, + deff=deff, design=design, ) @@ -1090,6 +1409,7 @@ def mde( n_pre: int = 1, n_post: int = 1, rho: float = 0.0, + deff: float = 1.0, ) -> PowerResults: """ Calculate minimum detectable effect given sample size. @@ -1111,6 +1431,8 @@ def mde( 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 ------- @@ -1123,10 +1445,13 @@ def mde( >>> 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, design) + 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) @@ -1144,6 +1469,7 @@ def mde( n_post=n_post, sigma=sigma, rho=rho, + deff=deff, design=design, ) @@ -1156,8 +1482,13 @@ def _compute_required_n( rho: float, design: str, treat_frac: float = 0.5, + deff: float = 1.0, ) -> int: - """Compute required sample size for given effect.""" + """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 @@ -1167,16 +1498,6 @@ def _compute_required_n( T = n_pre + n_post if design == "basic_did": - # Var = sigma^2 * (1/n_t + 1/n_t + 1/n_c + 1/n_c) = sigma^2 * (2/n_t + 2/n_c) - # For balanced: Var = sigma^2 * 4/n where n = n_t = n_c - # SE = sqrt(Var), effect_size = (z_alpha + z_beta) * SE - # n = 4 * sigma^2 * (z_alpha + z_beta)^2 / effect_size^2 - - # For general allocation with treat_frac: - # Var = sigma^2 * 2 * (1/(N*p) + 1/(N*(1-p))) - # = 2 * sigma^2 / N * (1/p + 1/(1-p)) - # = 2 * sigma^2 / N * (1/(p*(1-p))) - n_total = ( 2 * sigma**2 @@ -1186,9 +1507,6 @@ def _compute_required_n( else: # panel design_effect = 1 + (T - 1) * rho - # Var = sigma^2 * (1/n_t + 1/n_c) * design_effect / T - # For balanced: Var = 2 * sigma^2 / N * design_effect / T - n_total = ( 2 * sigma**2 @@ -1197,6 +1515,9 @@ def _compute_required_n( / (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 @@ -1211,6 +1532,7 @@ def sample_size( 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. @@ -1229,6 +1551,8 @@ def sample_size( 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 ------- @@ -1241,11 +1565,12 @@ def sample_size( >>> 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 + effect_size, sigma, n_pre, n_post, rho, design, treat_frac, deff=deff ) n_treated = max(2, int(np.ceil(n_total * treat_frac))) @@ -1253,7 +1578,9 @@ def sample_size( n_total = n_treated + n_control # Compute actual power achieved - variance = self._compute_variance(n_treated, n_control, n_pre, n_post, sigma, rho, design) + 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) @@ -1270,6 +1597,7 @@ def sample_size( n_post=n_post, sigma=sigma, rho=rho, + deff=deff, design=design, ) @@ -1282,6 +1610,7 @@ def power_curve( 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. @@ -1302,6 +1631,8 @@ def power_curve( 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 ------- @@ -1315,7 +1646,7 @@ def power_curve( >>> print(curve) """ # First get MDE to determine default range - mde_result = self.mde(n_treated, n_control, sigma, n_pre, n_post, rho) + 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 @@ -1331,6 +1662,7 @@ def power_curve( n_pre=n_pre, n_post=n_post, rho=rho, + deff=deff, ) powers.append(result.power) @@ -1345,6 +1677,7 @@ def sample_size_curve( 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. @@ -1365,6 +1698,8 @@ def sample_size_curve( 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 ------- @@ -1372,7 +1707,7 @@ def sample_size_curve( 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) + 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) @@ -1391,6 +1726,7 @@ def sample_size_curve( n_pre=n_pre, n_post=n_post, rho=rho, + deff=deff, ) powers.append(result.power) @@ -1414,6 +1750,7 @@ def simulate_power( 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. @@ -1530,10 +1867,46 @@ def simulate_power( # 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())}." + ) data_gen_kwargs = data_generator_kwargs or {} est_kwargs = estimator_kwargs or {} + # 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." + ) + # SyntheticDiD placebo variance requires n_control > n_treated. # Check after merging data_generator_kwargs so overrides of n_treated # are accounted for. @@ -1618,6 +1991,16 @@ def simulate_power( primary_rejections: List[bool] = [] primary_ci_contains: List[bool] = [] + # Survey DGP truth accumulation (DEFF/ICC are DGP properties, + # independent of effect size, so averaging across all sims is correct) + deff_values: List[float] = [] + icc_values: List[float] = [] + + # Lazy import for survey DGP (mirrors registry's lazy import pattern) + _generate_survey_did_data: Optional[Callable] = None + if use_survey_dgp: + from diff_diff.prep import generate_survey_did_data as _generate_survey_did_data + for effect_idx, effect in enumerate(effect_sizes): is_primary = effect_idx == primary_idx @@ -1636,7 +2019,37 @@ def simulate_power( sim_seed = rng.integers(0, 2**31) # --- Generate data --- - if use_custom_dgp: + 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 `post` column for basic/TWFE estimators + 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, @@ -1668,7 +2081,14 @@ def simulate_power( try: # --- Fit estimator --- - if profile is not None and not use_custom_dgp: + 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 ) @@ -1775,6 +2195,9 @@ def simulate_power( ) ], effective_n_units=effective_n_units, + survey_config=survey_config, + mean_deff=float(np.nanmean(deff_values)) if deff_values else None, + mean_icc_realized=float(np.nanmean(icc_values)) if icc_values else None, ) @@ -1823,6 +2246,7 @@ class SimulationMDEResults: search_path: List[Dict[str, float]] estimator_name: str effective_n_units: Optional[int] = None + survey_config: Optional[Any] = field(default=None, repr=False) def __repr__(self) -> str: return ( @@ -1920,6 +2344,7 @@ class SimulationSampleSizeResults: search_path: List[Dict[str, float]] estimator_name: str effective_n_units: Optional[int] = None + survey_config: Optional[Any] = field(default=None, repr=False) def __repr__(self) -> str: return ( @@ -1993,6 +2418,7 @@ def simulate_mde( estimator_kwargs: Optional[Dict[str, Any]] = None, result_extractor: Optional[Callable] = None, progress: bool = True, + survey_config: Optional[SurveyPowerConfig] = None, ) -> SimulationMDEResults: """ Find the minimum detectable effect via simulation-based bisection search. @@ -2075,6 +2501,7 @@ def simulate_mde( estimator_kwargs=estimator_kwargs, result_extractor=result_extractor, progress=False, + survey_config=survey_config, ) def _power_at(effect: float) -> float: @@ -2108,6 +2535,7 @@ def _power_at(effect: float) -> float: search_path=search_path, estimator_name=estimator_name, effective_n_units=effective_n_units, + survey_config=survey_config, ) if power_hi < power: warnings.warn( @@ -2135,6 +2563,7 @@ def _power_at(effect: float) -> float: n_steps=len(search_path), search_path=search_path, estimator_name=estimator_name, + survey_config=survey_config, effective_n_units=effective_n_units, ) @@ -2180,6 +2609,7 @@ def _power_at(effect: float) -> float: search_path=search_path, estimator_name=estimator_name, effective_n_units=effective_n_units, + survey_config=survey_config, ) @@ -2201,6 +2631,7 @@ def simulate_sample_size( estimator_kwargs: Optional[Dict[str, Any]] = None, result_extractor: Optional[Callable] = None, progress: bool = True, + survey_config: Optional[SurveyPowerConfig] = None, ) -> SimulationSampleSizeResults: """ Find the required sample size via simulation-based bisection search. @@ -2304,6 +2735,7 @@ def _snap_n(n: int, direction: str = "down", floor: Optional[int] = None) -> int estimator_kwargs=estimator_kwargs, result_extractor=result_extractor, progress=False, + survey_config=survey_config, ) def _power_at_n(n: int) -> float: @@ -2317,6 +2749,8 @@ def _power_at_n(n: int) -> float: # --- Bracket --- abs_min = 16 if is_ddd_grid else 4 + if survey_config is not None: + abs_min = max(abs_min, survey_config.min_viable_n) if n_range is not None: lo, hi = _snap_n(n_range[0], "up", floor=abs_min), _snap_n( n_range[1], "down", floor=abs_min @@ -2340,6 +2774,7 @@ def _power_at_n(n: int) -> float: n_steps=len(search_path), search_path=search_path, estimator_name=estimator_name, + survey_config=survey_config, ) power_hi = _power_at_n(hi) if power_hi < power: @@ -2389,6 +2824,7 @@ def _power_at_n(n: int) -> float: n_steps=len(search_path), search_path=search_path, estimator_name=estimator_name, + survey_config=survey_config, ) # Fall through to bisection with lo..hi bracket else: @@ -2444,6 +2880,7 @@ def _power_at_n(n: int) -> float: n_steps=len(search_path), search_path=search_path, estimator_name=estimator_name, + survey_config=survey_config, ) @@ -2456,6 +2893,7 @@ def compute_mde( n_pre: int = 1, n_post: int = 1, rho: float = 0.0, + deff: float = 1.0, ) -> float: """ Convenience function to compute minimum detectable effect. @@ -2478,6 +2916,8 @@ def compute_mde( 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 ------- @@ -2490,7 +2930,7 @@ def compute_mde( >>> print(f"MDE: {mde:.2f}") """ pa = PowerAnalysis(alpha=alpha, power=power) - result = pa.mde(n_treated, n_control, sigma, n_pre, n_post, rho) + result = pa.mde(n_treated, n_control, sigma, n_pre, n_post, rho, deff=deff) return result.mde @@ -2503,6 +2943,7 @@ def compute_power( n_pre: int = 1, n_post: int = 1, rho: float = 0.0, + deff: float = 1.0, ) -> float: """ Convenience function to compute power for given effect and sample. @@ -2525,6 +2966,8 @@ def compute_power( 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 ------- @@ -2537,7 +2980,7 @@ def compute_power( >>> print(f"Power: {power:.1%}") """ pa = PowerAnalysis(alpha=alpha) - result = pa.power(effect_size, n_treated, n_control, sigma, n_pre, n_post, rho) + result = pa.power(effect_size, n_treated, n_control, sigma, n_pre, n_post, rho, deff=deff) return result.power @@ -2550,6 +2993,7 @@ def compute_sample_size( n_post: int = 1, rho: float = 0.0, treat_frac: float = 0.5, + deff: float = 1.0, ) -> int: """ Convenience function to compute required sample size. @@ -2572,6 +3016,8 @@ def compute_sample_size( 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 ------- @@ -2584,5 +3030,5 @@ def compute_sample_size( >>> print(f"Required N: {n}") """ pa = PowerAnalysis(alpha=alpha, power=power) - result = pa.sample_size(effect_size, sigma, n_pre, n_post, rho, treat_frac) + result = pa.sample_size(effect_size, sigma, n_pre, n_post, rho, treat_frac, deff=deff) return result.required_n diff --git a/tests/test_power.py b/tests/test_power.py index 116a97a50..8ea876d8e 100644 --- a/tests/test_power.py +++ b/tests/test_power.py @@ -20,6 +20,7 @@ SimulationSampleSizeResults, StackedDiD, SunAbraham, + SurveyPowerConfig, SyntheticDiD, TripleDifference, TwoStageDiD, @@ -2121,3 +2122,308 @@ def test_sdid_placebo_custom_dgp_sample_size_raises(self): seed=42, progress=False, ) + + +# --------------------------------------------------------------------------- +# Survey-aware power tests +# --------------------------------------------------------------------------- + +# Small survey config for fast tests +_SURVEY_CFG = SurveyPowerConfig(n_strata=3, psu_per_stratum=4, icc=0.05) +_SIM_KW = dict(n_units=100, n_periods=4, treatment_period=2, sigma=1.0, progress=False) + + +class TestSurveyPower: + """Tests for survey-aware power analysis (SurveyPowerConfig + deff).""" + + # -- Simulation path: smoke tests for each estimator group -- + + def test_survey_simulate_power_cs(self): + """CallawaySantAnna with survey_config runs and returns valid power.""" + result = simulate_power( + CallawaySantAnna(), + treatment_effect=3.0, + n_simulations=20, + seed=42, + survey_config=_SURVEY_CFG, + **_SIM_KW, + ) + assert 0 <= result.power <= 1 + assert result.survey_config is not None + assert isinstance(result, SimulationPowerResults) + + def test_survey_simulate_power_basic_did(self): + """DifferenceInDifferences with survey_config works.""" + result = simulate_power( + DifferenceInDifferences(), + treatment_effect=3.0, + n_simulations=20, + seed=42, + survey_config=_SURVEY_CFG, + **_SIM_KW, + ) + assert 0 <= result.power <= 1 + + def test_survey_simulate_power_twfe(self): + """TwoWayFixedEffects with survey_config works.""" + result = simulate_power( + TwoWayFixedEffects(), + treatment_effect=3.0, + n_simulations=20, + seed=42, + survey_config=_SURVEY_CFG, + **_SIM_KW, + ) + assert 0 <= result.power <= 1 + + def test_survey_simulate_power_multiperiod(self): + """MultiPeriodDiD with survey_config works.""" + result = simulate_power( + MultiPeriodDiD(), + treatment_effect=3.0, + n_simulations=20, + seed=42, + survey_config=_SURVEY_CFG, + **_SIM_KW, + ) + assert 0 <= result.power <= 1 + + @pytest.mark.parametrize( + "estimator_cls", + [SunAbraham, ImputationDiD, TwoStageDiD, StackedDiD, EfficientDiD], + ) + def test_survey_staggered_estimators(self, estimator_cls): + """All staggered estimators work with survey_config.""" + result = simulate_power( + estimator_cls(), + treatment_effect=3.0, + n_simulations=10, + seed=42, + survey_config=_SURVEY_CFG, + **_SIM_KW, + ) + assert 0 <= result.power <= 1 + + # -- Validation: unsupported estimators -- + + def test_survey_rejects_trop(self): + with pytest.raises(ValueError, match="not supported with TROP"): + simulate_power(TROP(), n_simulations=1, seed=42, survey_config=_SURVEY_CFG, **_SIM_KW) + + def test_survey_rejects_sdid(self): + with pytest.raises(ValueError, match="not supported with SyntheticDiD"): + simulate_power( + SyntheticDiD(), + n_simulations=1, + seed=42, + survey_config=_SURVEY_CFG, + **_SIM_KW, + ) + + def test_survey_rejects_ddd(self): + with pytest.raises(ValueError, match="not supported with TripleDifference"): + simulate_power( + TripleDifference(), + n_simulations=1, + seed=42, + survey_config=_SURVEY_CFG, + **_SIM_KW, + ) + + def test_survey_rejects_custom_dgp(self): + with pytest.raises(ValueError, match="mutually exclusive"): + simulate_power( + CallawaySantAnna(), + data_generator=lambda **kw: None, + survey_config=_SURVEY_CFG, + n_simulations=1, + seed=42, + **_SIM_KW, + ) + + # -- Metadata -- + + def test_survey_metadata(self): + """mean_deff and mean_icc_realized populated in results.""" + result = simulate_power( + CallawaySantAnna(), + treatment_effect=3.0, + n_simulations=10, + seed=42, + survey_config=_SURVEY_CFG, + **_SIM_KW, + ) + assert result.mean_deff is not None + assert result.mean_deff > 1.0 + assert result.mean_icc_realized is not None + assert result.mean_icc_realized > 0 + + def test_survey_power_curve(self): + """survey_config works with effect_sizes (power curve).""" + result = simulate_power( + CallawaySantAnna(), + effect_sizes=[1.0, 3.0], + n_simulations=10, + seed=42, + survey_config=_SURVEY_CFG, + **_SIM_KW, + ) + assert len(result.powers) == 2 + assert result.powers[1] >= result.powers[0] + + # -- Validation: data_generator_kwargs conflicts -- + + def test_survey_data_gen_kwargs_blocked(self): + """Survey-config-managed keys in data_generator_kwargs raise ValueError.""" + with pytest.raises(ValueError, match="managed by survey_config"): + simulate_power( + CallawaySantAnna(), + data_generator_kwargs={"n_strata": 10}, + survey_config=_SURVEY_CFG, + n_simulations=1, + seed=42, + **_SIM_KW, + ) + + def test_survey_data_gen_kwargs_passthrough(self): + """Allowed keys (unit_fe_sd) pass through.""" + result = simulate_power( + CallawaySantAnna(), + treatment_effect=3.0, + data_generator_kwargs={"unit_fe_sd": 0.5}, + survey_config=_SURVEY_CFG, + n_simulations=10, + seed=42, + **_SIM_KW, + ) + assert 0 <= result.power <= 1 + + def test_survey_treatment_period_validation(self): + """treatment_period=0 raises with survey_config.""" + with pytest.raises(ValueError, match="treatment_period must be >= 1"): + simulate_power( + CallawaySantAnna(), + treatment_period=0, + survey_config=_SURVEY_CFG, + n_simulations=1, + seed=42, + n_units=100, + n_periods=4, + sigma=1.0, + progress=False, + ) + + # -- simulate_mde and simulate_sample_size -- + + def test_survey_mde(self): + """simulate_mde with survey_config.""" + result = simulate_mde( + CallawaySantAnna(), + n_simulations=10, + max_steps=3, + seed=42, + survey_config=_SURVEY_CFG, + **_SIM_KW, + ) + assert isinstance(result, SimulationMDEResults) + assert result.mde > 0 + assert result.survey_config is not None + + def test_survey_sample_size_min_n_floor(self): + """simulate_sample_size with survey_config respects min_viable_n.""" + cfg = SurveyPowerConfig(n_strata=5, psu_per_stratum=8) + # Use small effect so bisection needs larger N + result = simulate_sample_size( + CallawaySantAnna(), + treatment_effect=0.5, + sigma=3.0, + n_simulations=10, + max_steps=3, + seed=42, + survey_config=cfg, + n_periods=4, + treatment_period=2, + progress=False, + ) + assert isinstance(result, SimulationSampleSizeResults) + assert result.required_n >= cfg.min_viable_n + assert result.survey_config is not None + + # -- Closed-form deff tests -- + + def test_closed_form_deff_default(self): + """deff=1.0 preserves existing behavior exactly.""" + p1 = compute_power(effect_size=5.0, n_treated=50, n_control=50, sigma=10.0) + p2 = compute_power(effect_size=5.0, n_treated=50, n_control=50, sigma=10.0, deff=1.0) + assert p1 == p2 + + def test_closed_form_deff_increases_mde(self): + """deff > 1 increases MDE.""" + mde1 = compute_mde(n_treated=50, n_control=50, sigma=10.0) + mde2 = compute_mde(n_treated=50, n_control=50, sigma=10.0, deff=2.0) + assert mde2 > mde1 + + def test_closed_form_deff_increases_required_n(self): + """deff > 1 increases required N.""" + n1 = compute_sample_size(effect_size=5.0, sigma=10.0) + n2 = compute_sample_size(effect_size=5.0, sigma=10.0, deff=2.0) + assert n2 > n1 + + def test_closed_form_deff_and_rho(self): + """Both deff and rho can be set simultaneously.""" + pa = PowerAnalysis() + result = pa.power( + effect_size=5.0, + n_treated=50, + n_control=50, + sigma=10.0, + n_pre=2, + n_post=2, + rho=0.3, + deff=2.0, + ) + assert 0 < result.power < 1 + assert result.deff == 2.0 + assert result.rho == 0.3 + + def test_closed_form_deff_in_results(self): + """deff appears in PowerResults.to_dict().""" + pa = PowerAnalysis() + result = pa.power(effect_size=5.0, n_treated=50, n_control=50, sigma=10.0, deff=1.5) + d = result.to_dict() + assert "deff" in d + assert d["deff"] == 1.5 + + def test_closed_form_deff_warning(self): + """deff < 1.0 emits warning.""" + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + compute_power(effect_size=5.0, n_treated=50, n_control=50, sigma=10.0, deff=0.8) + assert any("deff=0.8000 < 1.0" in str(x.message) for x in w) + + def test_closed_form_deff_invalid(self): + """deff <= 0 raises ValueError.""" + with pytest.raises(ValueError, match="deff must be > 0"): + compute_power(effect_size=5.0, n_treated=50, n_control=50, sigma=10.0, deff=0.0) + + # -- SurveyPowerConfig validation -- + + def test_survey_config_validation_psu(self): + with pytest.raises(ValueError, match="psu_per_stratum"): + SurveyPowerConfig(psu_per_stratum=1) + + def test_survey_config_validation_strata(self): + with pytest.raises(ValueError, match="n_strata"): + SurveyPowerConfig(n_strata=0) + + def test_survey_config_validation_weight_variation(self): + with pytest.raises(ValueError, match="weight_variation"): + SurveyPowerConfig(weight_variation="extreme") + + def test_survey_config_validation_icc(self): + with pytest.raises(ValueError, match="icc"): + SurveyPowerConfig(icc=1.5) + + def test_survey_config_validation_fpc(self): + with pytest.raises(ValueError, match="fpc_per_stratum"): + SurveyPowerConfig(fpc_per_stratum=3, psu_per_stratum=8) From 56b07ea7d4d11a20e339fdee3ef99a7ce88d8f9b Mon Sep 17 00:00:00 2001 From: igerber Date: Sun, 12 Apr 2026 12:32:00 -0400 Subject: [PATCH 02/12] fix: address AI review - document deff in REGISTRY.md, cache SurveyDesign - Add REGISTRY.md notes for analytical deff parameter (variance/sample-size inflation formulas, deff vs rho distinction) and survey_config simulation path (supported estimators, mutual exclusivity, protected keys) - Cache SurveyDesign in SurveyPowerConfig._build_survey_design() to avoid rebuilding per simulation iteration Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/power.py | 15 ++++++++++----- docs/methodology/REGISTRY.md | 2 ++ 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/diff_diff/power.py b/diff_diff/power.py index 2ca1ce514..a899584a2 100644 --- a/diff_diff/power.py +++ b/diff_diff/power.py @@ -148,12 +148,17 @@ def __post_init__(self) -> None: ) def _build_survey_design(self) -> Any: - """Return user-supplied SurveyDesign or auto-build from DGP column names.""" - if self.survey_design is not None: - return self.survey_design - from diff_diff.survey import SurveyDesign + """Return cached SurveyDesign (built once, reused across simulations).""" + if not hasattr(self, "_cached_survey_design"): + if self.survey_design is not None: + self._cached_survey_design = self.survey_design + else: + from diff_diff.survey import SurveyDesign - return SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc") + self._cached_survey_design = SurveyDesign( + weights="weight", strata="stratum", psu="psu", fpc="fpc" + ) + return self._cached_survey_design @property def min_viable_n(self) -> int: diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 9da6214bf..ed0467b3c 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -2254,6 +2254,8 @@ n = 2(t_{α/2} + t_{1-κ})² σ² / MDE² - **Note:** `simulate_sample_size()` rejects `n_per_cell` in `data_generator_kwargs` for `TripleDifference` because `n_per_cell` is derived from `n_units` (the search variable). A fixed override would freeze the effective sample size across bisection iterations, making the search degenerate. Use `simulate_power()` with a fixed `n_per_cell` override instead, or pass a custom `data_generator`. - **Note:** The simulation-based power registry (`simulate_power`, `simulate_mde`, `simulate_sample_size`) uses a single-cohort staggered DGP by default. Estimators configured with `control_group="not_yet_treated"`, `clean_control="strict"`, or `anticipation>0` will receive a `UserWarning` because the default DGP does not match their identification strategy. Users must supply `data_generator_kwargs` (e.g., `cohort_periods=[2, 4]`, `never_treated_frac=0.0`) or a custom `data_generator` to match the estimator design. - **Note:** The `TripleDifference` registry adapter uses `generate_ddd_data`, a fixed 2×2×2 factorial DGP (group × partition × time). The `n_periods`, `treatment_period`, and `treatment_fraction` parameters are ignored — DDD always simulates 2 periods with balanced groups. `n_units` is mapped to `n_per_cell = max(2, n_units // 8)` (effective total N = `n_per_cell × 8`), so non-multiples of 8 are rounded down and values below 16 are clamped to 16. A `UserWarning` is emitted when simulation inputs differ from the effective DDD design. When rounding occurs, all result objects (`SimulationPowerResults`, `SimulationMDEResults`, `SimulationSampleSizeResults`) set `effective_n_units` to the actual sample size used; it is `None` when no rounding occurred. `simulate_sample_size()` snaps bisection candidates to multiples of 8 so that `required_n` is always a realizable DDD sample size. Passing `n_per_cell` in `data_generator_kwargs` suppresses the effective-N rounding warning but not warnings for ignored parameters (`n_periods`, `treatment_period`, `treatment_fraction`). +- **Note:** The analytical power methods (`PowerAnalysis.power/mde/sample_size` and the `compute_power/compute_mde/compute_sample_size` convenience functions) accept a `deff` parameter (survey design effect, default 1.0). This inflates variance multiplicatively: `Var(ATT) *= deff`, and inflates required sample size: `n_total *= deff`. The `deff` parameter is **not redundant** with `rho` (intra-cluster correlation): `rho` models within-unit serial correlation in panel data via the Moulton factor `1 + (T-1)*rho`, while `deff` models the survey design effect from stratified multi-stage sampling (clustering + unequal weighting). A survey panel study may need both. Values `deff > 0` are accepted; `deff < 1.0` (net variance reduction, e.g., from stratification gain) emits a warning. +- **Note:** The simulation-based power functions (`simulate_power/simulate_mde/simulate_sample_size`) accept a `survey_config` parameter (`SurveyPowerConfig` dataclass). When set, the simulation loop uses `generate_survey_did_data` instead of the default registry DGP, and automatically injects `SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc")` into the estimator's `fit()` call. Supported estimators: DifferenceInDifferences, TwoWayFixedEffects, MultiPeriodDiD, CallawaySantAnna, SunAbraham, ImputationDiD, TwoStageDiD, StackedDiD, EfficientDiD. Unsupported (raises `ValueError`): TROP, SyntheticDiD, TripleDifference (generate_survey_did_data produces staggered cohort data incompatible with factor-model/DDD DGPs). `survey_config` and `data_generator` are mutually exclusive. `data_generator_kwargs` may not contain keys managed by `SurveyPowerConfig` (n_strata, psu_per_stratum, etc.) but may contain passthrough DGP params (unit_fe_sd, add_covariates, panel, strata_sizes). `simulate_sample_size` raises the bisection floor to `n_strata * psu_per_stratum * 2` to ensure viable survey structure. **Reference implementation(s):** - R: `pwr` package (general), `DeclareDesign` (simulation-based) From 40acf5b13b9a28080981e6e7cda16c458b6e7ce3 Mon Sep 17 00:00:00 2001 From: igerber Date: Sun, 12 Apr 2026 12:54:54 -0400 Subject: [PATCH 03/12] fix: address CI review - ever_treated variable, min_viable_n floor, stronger tests - Fix P1: basic/TWFE/MultiPeriod survey adapters now derive `ever_treated` (time-invariant group indicator) instead of using the DGP's post-only `treated` column, which caused rank-deficient design matrices - Fix P1: MultiPeriodDiD adapter now passes `unit="unit"` for proper time-varying-treatment validation - Fix P1: simulate_sample_size auto-bracketing starts from max(min_n, abs_min) and clamps early-return to abs_min, enforcing min_viable_n contract - Fix P2: strengthen tests to assert finite mean_estimate/mean_se (not just 0 <= power <= 1), add large-effect floor regression test Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/power.py | 30 ++++++++++++++++++++++-------- tests/test_power.py | 34 +++++++++++++++++++++++++++++++--- 2 files changed, 53 insertions(+), 11 deletions(-) diff --git a/diff_diff/power.py b/diff_diff/power.py index a899584a2..30f84a57a 100644 --- a/diff_diff/power.py +++ b/diff_diff/power.py @@ -367,10 +367,16 @@ def _survey_basic_fit_kwargs( treatment_period: int, survey_config: SurveyPowerConfig, ) -> Dict[str, Any]: - """Fit kwargs for DifferenceInDifferences with survey design.""" + """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="treated", + treatment="ever_treated", time="post", survey_design=survey_config._build_survey_design(), ) @@ -386,7 +392,7 @@ def _survey_twfe_fit_kwargs( """Fit kwargs for TwoWayFixedEffects with survey design.""" return dict( outcome="outcome", - treatment="treated", + treatment="ever_treated", time="post", unit="unit", survey_design=survey_config._build_survey_design(), @@ -403,7 +409,8 @@ def _survey_multiperiod_fit_kwargs( """Fit kwargs for MultiPeriodDiD with survey design (1-indexed periods).""" return dict( outcome="outcome", - treatment="treated", + 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)), @@ -2040,7 +2047,12 @@ def simulate_power( dgp_kwargs.pop("seed", None) data = _generate_survey_did_data(seed=sim_seed, **dgp_kwargs) - # Derive `post` column for basic/TWFE estimators + # 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) @@ -2789,7 +2801,7 @@ def _power_at_n(n: int) -> float: UserWarning, ) else: - lo = min_n + lo = max(min_n, abs_min) power_lo = _power_at_n(lo) if power_lo >= power: # Floor achieves target — search downward for true minimum @@ -2812,15 +2824,17 @@ def _power_at_n(n: int) -> float: (s for s in search_path if s["power"] >= power), key=lambda s: s["n_units"], ) + # Clamp to abs_min (enforces survey min_viable_n contract) + best_n = max(int(best["n_units"]), abs_min) warnings.warn( - f"Power at n={int(best['n_units'])} is " + f"Power at n={best_n} is " f"{best['power']:.2f} >= target {power}. Could not " f"find a smaller N below target power. Pass " f"n_range=(lo, hi) to refine.", UserWarning, ) return SimulationSampleSizeResults( - required_n=int(best["n_units"]), + required_n=best_n, power_at_n=best["power"], target_power=power, alpha=alpha, diff --git a/tests/test_power.py b/tests/test_power.py index 8ea876d8e..5fc128ca6 100644 --- a/tests/test_power.py +++ b/tests/test_power.py @@ -2153,7 +2153,7 @@ def test_survey_simulate_power_cs(self): assert isinstance(result, SimulationPowerResults) def test_survey_simulate_power_basic_did(self): - """DifferenceInDifferences with survey_config works.""" + """DifferenceInDifferences with survey_config produces finite estimates.""" result = simulate_power( DifferenceInDifferences(), treatment_effect=3.0, @@ -2163,9 +2163,13 @@ def test_survey_simulate_power_basic_did(self): **_SIM_KW, ) assert 0 <= result.power <= 1 + # Verify non-degenerate: finite mean estimate and SE (not rank-deficient) + assert np.isfinite(result.mean_estimate) + assert np.isfinite(result.mean_se) + assert result.mean_se > 0 def test_survey_simulate_power_twfe(self): - """TwoWayFixedEffects with survey_config works.""" + """TwoWayFixedEffects with survey_config produces finite estimates.""" result = simulate_power( TwoWayFixedEffects(), treatment_effect=3.0, @@ -2175,9 +2179,12 @@ def test_survey_simulate_power_twfe(self): **_SIM_KW, ) assert 0 <= result.power <= 1 + assert np.isfinite(result.mean_estimate) + assert np.isfinite(result.mean_se) + assert result.mean_se > 0 def test_survey_simulate_power_multiperiod(self): - """MultiPeriodDiD with survey_config works.""" + """MultiPeriodDiD with survey_config produces finite estimates.""" result = simulate_power( MultiPeriodDiD(), treatment_effect=3.0, @@ -2187,6 +2194,9 @@ def test_survey_simulate_power_multiperiod(self): **_SIM_KW, ) assert 0 <= result.power <= 1 + assert np.isfinite(result.mean_estimate) + assert np.isfinite(result.mean_se) + assert result.mean_se > 0 @pytest.mark.parametrize( "estimator_cls", @@ -2349,6 +2359,24 @@ def test_survey_sample_size_min_n_floor(self): assert result.required_n >= cfg.min_viable_n assert result.survey_config is not None + def test_survey_sample_size_large_effect_floor(self): + """Large effect early-return still respects min_viable_n.""" + cfg = SurveyPowerConfig(n_strata=5, psu_per_stratum=8) + # Large effect - first probe likely achieves target power immediately + result = simulate_sample_size( + CallawaySantAnna(), + treatment_effect=10.0, + sigma=1.0, + n_simulations=10, + max_steps=3, + seed=42, + survey_config=cfg, + n_periods=4, + treatment_period=2, + progress=False, + ) + assert result.required_n >= cfg.min_viable_n # must be >= 80 + # -- Closed-form deff tests -- def test_closed_form_deff_default(self): From 685c92609d243503b8b267d56bd950412bfc3746 Mon Sep 17 00:00:00 2001 From: igerber Date: Sun, 12 Apr 2026 13:19:38 -0400 Subject: [PATCH 04/12] fix: address CI review round 2 - block heterogeneous TE, fix hi-bracket, docstrings - P0: reject heterogeneous_te_by_strata=True in survey power validation (DGP population_att diverges from input treatment_effect, making bias/coverage/RMSE metrics misleading) - P1: fix simulate_sample_size auto-bracketing hi to respect abs_min (hi = max(2*lo, abs_min, 100) prevents probing below survey floor) - P2: add survey_config parameter to simulate_power/mde/sample_size docstrings with constraints and supported estimators - P2: add regression tests for heterogeneous_te rejection and large-floor (n_strata=10, psu_per_stratum=10, min_viable_n=200) auto-bracketing branch Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/power.py | 24 +++++++++++++++++++++++- tests/test_power.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/diff_diff/power.py b/diff_diff/power.py index 30f84a57a..0f12b46d0 100644 --- a/diff_diff/power.py +++ b/diff_diff/power.py @@ -1812,6 +1812,13 @@ def simulate_power( 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 ------- @@ -1906,6 +1913,13 @@ def simulate_power( 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 {} @@ -2482,6 +2496,9 @@ def simulate_mde( Forwarded to ``simulate_power()``. progress : bool, default=True Whether to print progress updates. + survey_config : SurveyPowerConfig, optional + Survey-aware simulation config. Forwarded to ``simulate_power()``. + See :func:`simulate_power` for details and constraints. Returns ------- @@ -2693,6 +2710,11 @@ def simulate_sample_size( Forwarded to ``simulate_power()``. progress : bool, default=True Whether to print progress updates. + survey_config : SurveyPowerConfig, optional + Survey-aware simulation config. Forwarded to ``simulate_power()``. + When set, the bisection floor is raised to + ``survey_config.min_viable_n`` to ensure viable survey structure. + See :func:`simulate_power` for details and constraints. Returns ------- @@ -2847,7 +2869,7 @@ def _power_at_n(n: int) -> float: ) # Fall through to bisection with lo..hi bracket else: - hi = max(100, 2 * min_n) + hi = max(2 * lo, abs_min, 100) for _ in range(10): if _power_at_n(hi) >= power: break diff --git a/tests/test_power.py b/tests/test_power.py index 5fc128ca6..bb46c1054 100644 --- a/tests/test_power.py +++ b/tests/test_power.py @@ -2377,6 +2377,36 @@ def test_survey_sample_size_large_effect_floor(self): ) assert result.required_n >= cfg.min_viable_n # must be >= 80 + def test_survey_sample_size_large_floor_auto_bracket(self): + """Auto-bracketing hi respects min_viable_n > 100.""" + cfg = SurveyPowerConfig(n_strata=10, psu_per_stratum=10) + # min_viable_n = 10 * 10 * 2 = 200, which exceeds default hi=100 + result = simulate_sample_size( + CallawaySantAnna(), + treatment_effect=0.5, + sigma=3.0, + n_simulations=10, + max_steps=3, + seed=42, + survey_config=cfg, + n_periods=4, + treatment_period=2, + progress=False, + ) + assert result.required_n >= cfg.min_viable_n # must be >= 200 + + def test_survey_rejects_heterogeneous_te(self): + """heterogeneous_te_by_strata=True rejected with simulation power.""" + cfg = SurveyPowerConfig(heterogeneous_te_by_strata=True) + with pytest.raises(ValueError, match="heterogeneous_te_by_strata"): + simulate_power( + CallawaySantAnna(), + n_simulations=1, + seed=42, + survey_config=cfg, + **_SIM_KW, + ) + # -- Closed-form deff tests -- def test_closed_form_deff_default(self): From a1d6199933eebde01f4362bea2754bf761604eb7 Mon Sep 17 00:00:00 2001 From: igerber Date: Sun, 12 Apr 2026 13:36:37 -0400 Subject: [PATCH 05/12] fix: block te_covariate_interaction and covariate_effects with survey_config P0: reject data_generator_kwargs with te_covariate_interaction != 0.0 or covariate_effects when survey_config is active. These DGP params make the realized population ATT diverge from the scalar treatment_effect input, which would misstate bias/coverage/RMSE metrics in simulation power. Add regression tests for both rejection paths. Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/power.py | 17 +++++++++++++++++ tests/test_power.py | 27 +++++++++++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/diff_diff/power.py b/diff_diff/power.py index 0f12b46d0..deea01ee4 100644 --- a/diff_diff/power.py +++ b/diff_diff/power.py @@ -1932,6 +1932,23 @@ def simulate_power( 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." + ) + if "covariate_effects" in data_gen_kwargs: + raise ValueError( + "covariate_effects is not supported with survey_config. " + "Custom covariate effects can shift the realized population " + "ATT away from the input treatment_effect." + ) # SyntheticDiD placebo variance requires n_control > n_treated. # Check after merging data_generator_kwargs so overrides of n_treated diff --git a/tests/test_power.py b/tests/test_power.py index bb46c1054..c2bb42227 100644 --- a/tests/test_power.py +++ b/tests/test_power.py @@ -2407,6 +2407,33 @@ def test_survey_rejects_heterogeneous_te(self): **_SIM_KW, ) + def test_survey_rejects_te_covariate_interaction(self): + """te_covariate_interaction != 0 rejected (diverges population ATT).""" + with pytest.raises(ValueError, match="te_covariate_interaction"): + simulate_power( + CallawaySantAnna(), + data_generator_kwargs={ + "add_covariates": True, + "te_covariate_interaction": 1.0, + }, + survey_config=_SURVEY_CFG, + n_simulations=1, + seed=42, + **_SIM_KW, + ) + + def test_survey_rejects_covariate_effects(self): + """covariate_effects rejected (can shift population ATT).""" + with pytest.raises(ValueError, match="covariate_effects"): + simulate_power( + CallawaySantAnna(), + data_generator_kwargs={"covariate_effects": (0.5, 0.3)}, + survey_config=_SURVEY_CFG, + n_simulations=1, + seed=42, + **_SIM_KW, + ) + # -- Closed-form deff tests -- def test_closed_form_deff_default(self): From a39a96f36030b81e32eda29525720f239f9cbf52 Mon Sep 17 00:00:00 2001 From: igerber Date: Sun, 12 Apr 2026 14:03:59 -0400 Subject: [PATCH 06/12] fix: address CI review round 4 - panel=False guard, drop covariate_effects block - P1: reject panel=False in data_generator_kwargs for panel-only estimators (TWFE, SA, Imputation, TwoStage, Stacked, Efficient); only CS supports repeated cross-sections under survey_config - P2: remove over-broad covariate_effects rejection (it affects baseline outcomes and ICC calibration, not realized population_att) - Add regression tests: panel=False rejected for TWFE, allowed for CS(panel=False) Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/power.py | 16 ++++++++++------ tests/test_power.py | 23 ++++++++++++++++++----- 2 files changed, 28 insertions(+), 11 deletions(-) diff --git a/diff_diff/power.py b/diff_diff/power.py index deea01ee4..2761d3bbe 100644 --- a/diff_diff/power.py +++ b/diff_diff/power.py @@ -1943,12 +1943,16 @@ def simulate_power( f"the input treatment_effect under covariate-interaction " f"heterogeneity, which would make bias/coverage/RMSE misleading." ) - if "covariate_effects" in data_gen_kwargs: - raise ValueError( - "covariate_effects is not supported with survey_config. " - "Custom covariate effects can shift the realized population " - "ATT away from the input treatment_effect." - ) + # Block panel=False for panel-only estimators. Only + # CallawaySantAnna supports repeated cross-sections. + if not data_gen_kwargs.get("panel", True): + _RCS_SUPPORTED = frozenset({"CallawaySantAnna"}) + if estimator_name not in _RCS_SUPPORTED: + raise ValueError( + f"panel=False (repeated cross-sections) is not supported " + f"with {estimator_name} under survey_config. Only " + f"{sorted(_RCS_SUPPORTED)} support repeated cross-sections." + ) # SyntheticDiD placebo variance requires n_control > n_treated. # Check after merging data_generator_kwargs so overrides of n_treated diff --git a/tests/test_power.py b/tests/test_power.py index c2bb42227..021300894 100644 --- a/tests/test_power.py +++ b/tests/test_power.py @@ -2422,18 +2422,31 @@ def test_survey_rejects_te_covariate_interaction(self): **_SIM_KW, ) - def test_survey_rejects_covariate_effects(self): - """covariate_effects rejected (can shift population ATT).""" - with pytest.raises(ValueError, match="covariate_effects"): + def test_survey_rejects_panel_false_for_panel_only(self): + """panel=False rejected for panel-only estimators (e.g., TWFE).""" + with pytest.raises(ValueError, match="panel=False.*not supported"): simulate_power( - CallawaySantAnna(), - data_generator_kwargs={"covariate_effects": (0.5, 0.3)}, + TwoWayFixedEffects(), + data_generator_kwargs={"panel": False}, survey_config=_SURVEY_CFG, n_simulations=1, seed=42, **_SIM_KW, ) + def test_survey_allows_panel_false_for_cs(self): + """panel=False allowed for CallawaySantAnna (supports RCS).""" + result = simulate_power( + CallawaySantAnna(panel=False), + treatment_effect=3.0, + data_generator_kwargs={"panel": False}, + survey_config=_SURVEY_CFG, + n_simulations=10, + seed=42, + **_SIM_KW, + ) + assert 0 <= result.power <= 1 + # -- Closed-form deff tests -- def test_closed_form_deff_default(self): From 904a4d77286b1be2df717b561811c4e950d8f73f Mon Sep 17 00:00:00 2001 From: igerber Date: Sun, 12 Apr 2026 14:40:07 -0400 Subject: [PATCH 07/12] fix: enforce CS panel-mode alignment between estimator and DGP P1: validate that CallawaySantAnna.panel matches the survey DGP's panel flag in both directions: - CS(panel=True) + DGP panel=False -> rejected - CS(panel=False) + default DGP (panel=True) -> rejected - Panel check now runs even with empty data_generator_kwargs Add regression tests for both misaligned CS combinations. Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/power.py | 27 +++++++++++++++++++++------ tests/test_power.py | 25 ++++++++++++++++++++++++- 2 files changed, 45 insertions(+), 7 deletions(-) diff --git a/diff_diff/power.py b/diff_diff/power.py index 2761d3bbe..4808b870f 100644 --- a/diff_diff/power.py +++ b/diff_diff/power.py @@ -1943,16 +1943,31 @@ def simulate_power( f"the input treatment_effect under covariate-interaction " f"heterogeneity, which would make bias/coverage/RMSE misleading." ) - # Block panel=False for panel-only estimators. Only - # CallawaySantAnna supports repeated cross-sections. - if not data_gen_kwargs.get("panel", True): - _RCS_SUPPORTED = frozenset({"CallawaySantAnna"}) - if estimator_name not in _RCS_SUPPORTED: + + # 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"{sorted(_RCS_SUPPORTED)} support repeated cross-sections." + 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." + ) # SyntheticDiD placebo variance requires n_control > n_treated. # Check after merging data_generator_kwargs so overrides of n_treated diff --git a/tests/test_power.py b/tests/test_power.py index 021300894..e0edc7701 100644 --- a/tests/test_power.py +++ b/tests/test_power.py @@ -2435,7 +2435,7 @@ def test_survey_rejects_panel_false_for_panel_only(self): ) def test_survey_allows_panel_false_for_cs(self): - """panel=False allowed for CallawaySantAnna (supports RCS).""" + """panel=False allowed for CallawaySantAnna(panel=False) (supports RCS).""" result = simulate_power( CallawaySantAnna(panel=False), treatment_effect=3.0, @@ -2447,6 +2447,29 @@ def test_survey_allows_panel_false_for_cs(self): ) assert 0 <= result.power <= 1 + def test_survey_rejects_cs_panel_mismatch_dgp_rcs(self): + """CS(panel=True) + DGP panel=False rejected.""" + with pytest.raises(ValueError, match="CallawaySantAnna.panel=True"): + simulate_power( + CallawaySantAnna(), # panel=True by default + data_generator_kwargs={"panel": False}, + survey_config=_SURVEY_CFG, + n_simulations=1, + seed=42, + **_SIM_KW, + ) + + def test_survey_rejects_cs_panel_mismatch_est_rcs(self): + """CS(panel=False) + default DGP (panel=True) rejected.""" + with pytest.raises(ValueError, match="panel=False.*requires"): + simulate_power( + CallawaySantAnna(panel=False), + survey_config=_SURVEY_CFG, + n_simulations=1, + seed=42, + **_SIM_KW, + ) + # -- Closed-form deff tests -- def test_closed_form_deff_default(self): From ce910c7fb9c7edf9d59b6157170596429f34155d Mon Sep 17 00:00:00 2001 From: igerber Date: Sun, 12 Apr 2026 15:09:14 -0400 Subject: [PATCH 08/12] fix: address CI review round 6 - block survey_design override, config conflicts, docs - P1: reject estimator_kwargs["survey_design"] when survey_config is set (prevents silently overwriting injected design; use SurveyPowerConfig.survey_design) - P2: add icc/psu_re_sd and weight_cv/weight_variation mutual-exclusion validation to SurveyPowerConfig.__post_init__ (mirrors DGP validation) - P3: update REGISTRY.md survey_config note to document panel=False coupling with CallawaySantAnna and estimator_kwargs survey_design restriction - Add regression tests for all three guards Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/power.py | 19 +++++++++++++++++++ docs/methodology/REGISTRY.md | 2 +- tests/test_power.py | 20 ++++++++++++++++++++ 3 files changed, 40 insertions(+), 1 deletion(-) diff --git a/diff_diff/power.py b/diff_diff/power.py index 4808b870f..bfe0ce009 100644 --- a/diff_diff/power.py +++ b/diff_diff/power.py @@ -139,8 +139,18 @@ def __post_init__(self) -> None: ) 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 and self.weight_cv <= 0: raise ValueError(f"weight_cv must be > 0, got {self.weight_cv}") + if self.weight_cv is not None and self.weight_variation != "moderate": + raise ValueError( + "Cannot specify both weight_cv and a non-default " + "weight_variation. weight_cv overrides weight_variation." + ) if self.fpc_per_stratum < self.psu_per_stratum: raise ValueError( f"fpc_per_stratum ({self.fpc_per_stratum}) must be >= " @@ -1924,6 +1934,15 @@ def simulate_power( 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) diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index ed0467b3c..a28361aa4 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -2255,7 +2255,7 @@ n = 2(t_{α/2} + t_{1-κ})² σ² / MDE² - **Note:** The simulation-based power registry (`simulate_power`, `simulate_mde`, `simulate_sample_size`) uses a single-cohort staggered DGP by default. Estimators configured with `control_group="not_yet_treated"`, `clean_control="strict"`, or `anticipation>0` will receive a `UserWarning` because the default DGP does not match their identification strategy. Users must supply `data_generator_kwargs` (e.g., `cohort_periods=[2, 4]`, `never_treated_frac=0.0`) or a custom `data_generator` to match the estimator design. - **Note:** The `TripleDifference` registry adapter uses `generate_ddd_data`, a fixed 2×2×2 factorial DGP (group × partition × time). The `n_periods`, `treatment_period`, and `treatment_fraction` parameters are ignored — DDD always simulates 2 periods with balanced groups. `n_units` is mapped to `n_per_cell = max(2, n_units // 8)` (effective total N = `n_per_cell × 8`), so non-multiples of 8 are rounded down and values below 16 are clamped to 16. A `UserWarning` is emitted when simulation inputs differ from the effective DDD design. When rounding occurs, all result objects (`SimulationPowerResults`, `SimulationMDEResults`, `SimulationSampleSizeResults`) set `effective_n_units` to the actual sample size used; it is `None` when no rounding occurred. `simulate_sample_size()` snaps bisection candidates to multiples of 8 so that `required_n` is always a realizable DDD sample size. Passing `n_per_cell` in `data_generator_kwargs` suppresses the effective-N rounding warning but not warnings for ignored parameters (`n_periods`, `treatment_period`, `treatment_fraction`). - **Note:** The analytical power methods (`PowerAnalysis.power/mde/sample_size` and the `compute_power/compute_mde/compute_sample_size` convenience functions) accept a `deff` parameter (survey design effect, default 1.0). This inflates variance multiplicatively: `Var(ATT) *= deff`, and inflates required sample size: `n_total *= deff`. The `deff` parameter is **not redundant** with `rho` (intra-cluster correlation): `rho` models within-unit serial correlation in panel data via the Moulton factor `1 + (T-1)*rho`, while `deff` models the survey design effect from stratified multi-stage sampling (clustering + unequal weighting). A survey panel study may need both. Values `deff > 0` are accepted; `deff < 1.0` (net variance reduction, e.g., from stratification gain) emits a warning. -- **Note:** The simulation-based power functions (`simulate_power/simulate_mde/simulate_sample_size`) accept a `survey_config` parameter (`SurveyPowerConfig` dataclass). When set, the simulation loop uses `generate_survey_did_data` instead of the default registry DGP, and automatically injects `SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc")` into the estimator's `fit()` call. Supported estimators: DifferenceInDifferences, TwoWayFixedEffects, MultiPeriodDiD, CallawaySantAnna, SunAbraham, ImputationDiD, TwoStageDiD, StackedDiD, EfficientDiD. Unsupported (raises `ValueError`): TROP, SyntheticDiD, TripleDifference (generate_survey_did_data produces staggered cohort data incompatible with factor-model/DDD DGPs). `survey_config` and `data_generator` are mutually exclusive. `data_generator_kwargs` may not contain keys managed by `SurveyPowerConfig` (n_strata, psu_per_stratum, etc.) but may contain passthrough DGP params (unit_fe_sd, add_covariates, panel, strata_sizes). `simulate_sample_size` raises the bisection floor to `n_strata * psu_per_stratum * 2` to ensure viable survey structure. +- **Note:** The simulation-based power functions (`simulate_power/simulate_mde/simulate_sample_size`) accept a `survey_config` parameter (`SurveyPowerConfig` dataclass). When set, the simulation loop uses `generate_survey_did_data` instead of the default registry DGP, and automatically injects `SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc")` into the estimator's `fit()` call. Supported estimators: DifferenceInDifferences, TwoWayFixedEffects, MultiPeriodDiD, CallawaySantAnna, SunAbraham, ImputationDiD, TwoStageDiD, StackedDiD, EfficientDiD. Unsupported (raises `ValueError`): TROP, SyntheticDiD, TripleDifference (generate_survey_did_data produces staggered cohort data incompatible with factor-model/DDD DGPs). `survey_config` and `data_generator` are mutually exclusive. `data_generator_kwargs` may not contain keys managed by `SurveyPowerConfig` (n_strata, psu_per_stratum, etc.) but may contain passthrough DGP params (unit_fe_sd, add_covariates, strata_sizes). Repeated cross-section survey power (`panel=False`) is only supported for `CallawaySantAnna(panel=False)` with a matching `data_generator_kwargs={"panel": False}`; both mismatch directions are rejected. `estimator_kwargs` may not contain `survey_design` when `survey_config` is set (use `SurveyPowerConfig(survey_design=...)` instead). `simulate_sample_size` raises the bisection floor to `n_strata * psu_per_stratum * 2` to ensure viable survey structure. **Reference implementation(s):** - R: `pwr` package (general), `DeclareDesign` (simulation-based) diff --git a/tests/test_power.py b/tests/test_power.py index e0edc7701..0927dd789 100644 --- a/tests/test_power.py +++ b/tests/test_power.py @@ -2548,3 +2548,23 @@ def test_survey_config_validation_icc(self): def test_survey_config_validation_fpc(self): with pytest.raises(ValueError, match="fpc_per_stratum"): SurveyPowerConfig(fpc_per_stratum=3, psu_per_stratum=8) + + def test_survey_config_validation_icc_psu_re_sd_conflict(self): + with pytest.raises(ValueError, match="icc.*psu_re_sd"): + SurveyPowerConfig(icc=0.05, psu_re_sd=3.0) + + def test_survey_config_validation_weight_cv_variation_conflict(self): + with pytest.raises(ValueError, match="weight_cv.*weight_variation"): + SurveyPowerConfig(weight_cv=0.5, weight_variation="high") + + def test_survey_rejects_estimator_kwargs_survey_design(self): + """estimator_kwargs cannot contain survey_design when survey_config set.""" + with pytest.raises(ValueError, match="estimator_kwargs.*survey_design"): + simulate_power( + CallawaySantAnna(), + estimator_kwargs={"survey_design": None}, + survey_config=_SURVEY_CFG, + n_simulations=1, + seed=42, + **_SIM_KW, + ) From 4322648a3878cc7f18a7e223e43a5b86d5aded0c Mon Sep 17 00:00:00 2001 From: igerber Date: Sun, 12 Apr 2026 15:24:05 -0400 Subject: [PATCH 09/12] fix: address CI review round 7 P2 nits - finiteness validation - Reject non-finite deff (NaN, inf) in _validate_deff() - Add weight_cv finiteness and psu_period_factor non-negativity checks to SurveyPowerConfig.__post_init__ (mirrors DGP validation) - Add regression tests for all new edge cases Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/power.py | 19 ++++++++++++------- tests/test_power.py | 24 +++++++++++++++++++++++- 2 files changed, 35 insertions(+), 8 deletions(-) diff --git a/diff_diff/power.py b/diff_diff/power.py index bfe0ce009..3f63e8760 100644 --- a/diff_diff/power.py +++ b/diff_diff/power.py @@ -144,12 +144,17 @@ def __post_init__(self) -> None: "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 and self.weight_cv <= 0: - raise ValueError(f"weight_cv must be > 0, got {self.weight_cv}") - if self.weight_cv is not None and self.weight_variation != "moderate": + 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( - "Cannot specify both weight_cv and a non-default " - "weight_variation. weight_cv overrides weight_variation." + 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( @@ -1229,8 +1234,8 @@ def __init__( @staticmethod def _validate_deff(deff: float) -> None: """Validate deff parameter and warn if < 1.""" - if deff <= 0: - raise ValueError(f"deff must be > 0, got {deff}") + if not np.isfinite(deff) or deff <= 0: + raise ValueError(f"deff must be finite and > 0, got {deff}") if deff < 1.0: warnings.warn( f"deff={deff:.4f} < 1.0 implies net variance reduction " diff --git a/tests/test_power.py b/tests/test_power.py index 0927dd789..4ca54aaf3 100644 --- a/tests/test_power.py +++ b/tests/test_power.py @@ -2522,9 +2522,19 @@ def test_closed_form_deff_warning(self): compute_power(effect_size=5.0, n_treated=50, n_control=50, sigma=10.0, deff=0.8) assert any("deff=0.8000 < 1.0" in str(x.message) for x in w) + def test_closed_form_deff_nan(self): + """deff=NaN raises ValueError.""" + with pytest.raises(ValueError, match="deff must be finite"): + compute_power(effect_size=5.0, n_treated=50, n_control=50, sigma=10.0, deff=np.nan) + + def test_closed_form_deff_inf(self): + """deff=inf raises ValueError.""" + with pytest.raises(ValueError, match="deff must be finite"): + compute_power(effect_size=5.0, n_treated=50, n_control=50, sigma=10.0, deff=np.inf) + def test_closed_form_deff_invalid(self): """deff <= 0 raises ValueError.""" - with pytest.raises(ValueError, match="deff must be > 0"): + with pytest.raises(ValueError, match="deff must be finite"): compute_power(effect_size=5.0, n_treated=50, n_control=50, sigma=10.0, deff=0.0) # -- SurveyPowerConfig validation -- @@ -2557,6 +2567,18 @@ def test_survey_config_validation_weight_cv_variation_conflict(self): with pytest.raises(ValueError, match="weight_cv.*weight_variation"): SurveyPowerConfig(weight_cv=0.5, weight_variation="high") + def test_survey_config_validation_weight_cv_nonfinite(self): + with pytest.raises(ValueError, match="weight_cv must be finite"): + SurveyPowerConfig(weight_cv=np.inf) + + def test_survey_config_validation_psu_period_factor_nonfinite(self): + with pytest.raises(ValueError, match="psu_period_factor must be finite"): + SurveyPowerConfig(psu_period_factor=np.nan) + + def test_survey_config_validation_psu_period_factor_negative(self): + with pytest.raises(ValueError, match="psu_period_factor must be finite"): + SurveyPowerConfig(psu_period_factor=-1.0) + def test_survey_rejects_estimator_kwargs_survey_design(self): """estimator_kwargs cannot contain survey_design when survey_config set.""" with pytest.raises(ValueError, match="estimator_kwargs.*survey_design"): From c5f9b3c589a6f56bb3ce9405cd62a35e02139037 Mon Sep 17 00:00:00 2001 From: igerber Date: Sun, 12 Apr 2026 15:44:55 -0400 Subject: [PATCH 10/12] fix: address CI review round 8 - multi-cohort guard, strata_sizes guard - P1: reject control_group='not_yet_treated' and clean_control='strict' with survey_config (require multi-cohort DGP that survey path doesn't support) - P1: reject strata_sizes in data_generator_kwargs for simulate_sample_size (sum(strata_sizes) == n_units but n_units varies during bisection) - Update REGISTRY.md survey-power note to document these restrictions - Add regression tests for all three rejection paths Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/power.py | 29 +++++++++++++++++++++++++++ docs/methodology/REGISTRY.md | 2 +- tests/test_power.py | 39 ++++++++++++++++++++++++++++++++++++ 3 files changed, 69 insertions(+), 1 deletion(-) diff --git a/diff_diff/power.py b/diff_diff/power.py index 3f63e8760..7e3fd1f76 100644 --- a/diff_diff/power.py +++ b/diff_diff/power.py @@ -1992,6 +1992,23 @@ def simulate_power( "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 == "not_yet_treated": + raise ValueError( + f"survey_config does not support control_group='not_yet_treated' " + f"(requires multi-cohort DGP). Use the custom data_generator " + f"path for survey power with not-yet-treated controls." + ) + 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 @@ -2846,6 +2863,18 @@ def _power_at_n(n: int) -> float: print(f" Sample size search: n={n}, power={pwr:.3f}") return pwr + # Block strata_sizes in sample-size search (same class as n_per_cell for DDD): + # strata_sizes requires sum(strata_sizes) == n_units, but n_units varies + # during bisection so a fixed strata_sizes would fail mid-search. + if survey_config is not None and data_generator_kwargs: + if "strata_sizes" in data_generator_kwargs: + raise ValueError( + "strata_sizes in data_generator_kwargs is not supported with " + "simulate_sample_size() because n_units varies during the " + "bisection search. Use simulate_power() with a fixed n_units " + "and strata_sizes instead." + ) + # --- Bracket --- abs_min = 16 if is_ddd_grid else 4 if survey_config is not None: diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index a28361aa4..8563e7583 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -2255,7 +2255,7 @@ n = 2(t_{α/2} + t_{1-κ})² σ² / MDE² - **Note:** The simulation-based power registry (`simulate_power`, `simulate_mde`, `simulate_sample_size`) uses a single-cohort staggered DGP by default. Estimators configured with `control_group="not_yet_treated"`, `clean_control="strict"`, or `anticipation>0` will receive a `UserWarning` because the default DGP does not match their identification strategy. Users must supply `data_generator_kwargs` (e.g., `cohort_periods=[2, 4]`, `never_treated_frac=0.0`) or a custom `data_generator` to match the estimator design. - **Note:** The `TripleDifference` registry adapter uses `generate_ddd_data`, a fixed 2×2×2 factorial DGP (group × partition × time). The `n_periods`, `treatment_period`, and `treatment_fraction` parameters are ignored — DDD always simulates 2 periods with balanced groups. `n_units` is mapped to `n_per_cell = max(2, n_units // 8)` (effective total N = `n_per_cell × 8`), so non-multiples of 8 are rounded down and values below 16 are clamped to 16. A `UserWarning` is emitted when simulation inputs differ from the effective DDD design. When rounding occurs, all result objects (`SimulationPowerResults`, `SimulationMDEResults`, `SimulationSampleSizeResults`) set `effective_n_units` to the actual sample size used; it is `None` when no rounding occurred. `simulate_sample_size()` snaps bisection candidates to multiples of 8 so that `required_n` is always a realizable DDD sample size. Passing `n_per_cell` in `data_generator_kwargs` suppresses the effective-N rounding warning but not warnings for ignored parameters (`n_periods`, `treatment_period`, `treatment_fraction`). - **Note:** The analytical power methods (`PowerAnalysis.power/mde/sample_size` and the `compute_power/compute_mde/compute_sample_size` convenience functions) accept a `deff` parameter (survey design effect, default 1.0). This inflates variance multiplicatively: `Var(ATT) *= deff`, and inflates required sample size: `n_total *= deff`. The `deff` parameter is **not redundant** with `rho` (intra-cluster correlation): `rho` models within-unit serial correlation in panel data via the Moulton factor `1 + (T-1)*rho`, while `deff` models the survey design effect from stratified multi-stage sampling (clustering + unequal weighting). A survey panel study may need both. Values `deff > 0` are accepted; `deff < 1.0` (net variance reduction, e.g., from stratification gain) emits a warning. -- **Note:** The simulation-based power functions (`simulate_power/simulate_mde/simulate_sample_size`) accept a `survey_config` parameter (`SurveyPowerConfig` dataclass). When set, the simulation loop uses `generate_survey_did_data` instead of the default registry DGP, and automatically injects `SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc")` into the estimator's `fit()` call. Supported estimators: DifferenceInDifferences, TwoWayFixedEffects, MultiPeriodDiD, CallawaySantAnna, SunAbraham, ImputationDiD, TwoStageDiD, StackedDiD, EfficientDiD. Unsupported (raises `ValueError`): TROP, SyntheticDiD, TripleDifference (generate_survey_did_data produces staggered cohort data incompatible with factor-model/DDD DGPs). `survey_config` and `data_generator` are mutually exclusive. `data_generator_kwargs` may not contain keys managed by `SurveyPowerConfig` (n_strata, psu_per_stratum, etc.) but may contain passthrough DGP params (unit_fe_sd, add_covariates, strata_sizes). Repeated cross-section survey power (`panel=False`) is only supported for `CallawaySantAnna(panel=False)` with a matching `data_generator_kwargs={"panel": False}`; both mismatch directions are rejected. `estimator_kwargs` may not contain `survey_design` when `survey_config` is set (use `SurveyPowerConfig(survey_design=...)` instead). `simulate_sample_size` raises the bisection floor to `n_strata * psu_per_stratum * 2` to ensure viable survey structure. +- **Note:** The simulation-based power functions (`simulate_power/simulate_mde/simulate_sample_size`) accept a `survey_config` parameter (`SurveyPowerConfig` dataclass). When set, the simulation loop uses `generate_survey_did_data` instead of the default registry DGP, and automatically injects `SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc")` into the estimator's `fit()` call. Supported estimators: DifferenceInDifferences, TwoWayFixedEffects, MultiPeriodDiD, CallawaySantAnna, SunAbraham, ImputationDiD, TwoStageDiD, StackedDiD, EfficientDiD. Unsupported (raises `ValueError`): TROP, SyntheticDiD, TripleDifference (generate_survey_did_data produces staggered cohort data incompatible with factor-model/DDD DGPs). `survey_config` and `data_generator` are mutually exclusive. `data_generator_kwargs` may not contain keys managed by `SurveyPowerConfig` (n_strata, psu_per_stratum, etc.) but may contain passthrough DGP params (unit_fe_sd, add_covariates, strata_sizes). Repeated cross-section survey power (`panel=False`) is only supported for `CallawaySantAnna(panel=False)` with a matching `data_generator_kwargs={"panel": False}`; both mismatch directions are rejected. `estimator_kwargs` may not contain `survey_design` when `survey_config` is set (use `SurveyPowerConfig(survey_design=...)` instead). Estimator settings that require a multi-cohort DGP (`control_group="not_yet_treated"`, `clean_control="strict"`) are rejected because the survey DGP uses a single cohort; use the custom `data_generator` path for these configurations. `simulate_sample_size` raises the bisection floor to `n_strata * psu_per_stratum * 2` to ensure viable survey structure and rejects `strata_sizes` in `data_generator_kwargs` (it depends on `n_units` which varies during bisection). **Reference implementation(s):** - R: `pwr` package (general), `DeclareDesign` (simulation-based) diff --git a/tests/test_power.py b/tests/test_power.py index 4ca54aaf3..abfd5b6c8 100644 --- a/tests/test_power.py +++ b/tests/test_power.py @@ -2590,3 +2590,42 @@ def test_survey_rejects_estimator_kwargs_survey_design(self): seed=42, **_SIM_KW, ) + + def test_survey_rejects_not_yet_treated(self): + """control_group='not_yet_treated' rejected (needs multi-cohort DGP).""" + with pytest.raises(ValueError, match="not_yet_treated"): + simulate_power( + CallawaySantAnna(control_group="not_yet_treated"), + survey_config=_SURVEY_CFG, + n_simulations=1, + seed=42, + **_SIM_KW, + ) + + def test_survey_rejects_clean_control_strict(self): + """clean_control='strict' rejected (needs multi-cohort DGP).""" + with pytest.raises(ValueError, match="clean_control.*strict"): + simulate_power( + StackedDiD(clean_control="strict"), + survey_config=_SURVEY_CFG, + n_simulations=1, + seed=42, + **_SIM_KW, + ) + + def test_survey_sample_size_rejects_strata_sizes(self): + """strata_sizes in data_generator_kwargs rejected for sample_size search.""" + with pytest.raises(ValueError, match="strata_sizes.*not supported"): + simulate_sample_size( + CallawaySantAnna(), + treatment_effect=3.0, + data_generator_kwargs={"strata_sizes": [20, 20, 20]}, + survey_config=SurveyPowerConfig(n_strata=3, psu_per_stratum=4), + n_simulations=10, + max_steps=3, + seed=42, + n_periods=4, + treatment_period=2, + sigma=1.0, + progress=False, + ) From b91a48549f932aa8cf8be02b80be7c3b07626a55 Mon Sep 17 00:00:00 2001 From: igerber Date: Sun, 12 Apr 2026 16:04:49 -0400 Subject: [PATCH 11/12] fix: address CI review round 9 - last_cohort guard, scalar validation - P1: reject control_group='last_cohort' (EfficientDiD) with survey_config (needs multi-cohort DGP, same as not_yet_treated) - P2: add psu_re_sd and fpc_per_stratum finiteness validation to SurveyPowerConfig.__post_init__ - Update REGISTRY.md to list last_cohort alongside other restrictions - Add regression tests for last_cohort rejection and scalar validation Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/power.py | 10 +++++++--- docs/methodology/REGISTRY.md | 2 +- tests/test_power.py | 23 +++++++++++++++++++++++ 3 files changed, 31 insertions(+), 4 deletions(-) diff --git a/diff_diff/power.py b/diff_diff/power.py index 7e3fd1f76..51f12441a 100644 --- a/diff_diff/power.py +++ b/diff_diff/power.py @@ -137,6 +137,10 @@ def __post_init__(self) -> None: 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: @@ -1997,11 +2001,11 @@ def simulate_power( # cohort_periods/never_treated_frac overrides. control_group = getattr(estimator, "control_group", "never_treated") clean_control = getattr(estimator, "clean_control", None) - if control_group == "not_yet_treated": + if control_group in ("not_yet_treated", "last_cohort"): raise ValueError( - f"survey_config does not support control_group='not_yet_treated' " + 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 not-yet-treated controls." + f"path for survey power with this control-group design." ) if clean_control == "strict": raise ValueError( diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 8563e7583..0a0622c1c 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -2255,7 +2255,7 @@ n = 2(t_{α/2} + t_{1-κ})² σ² / MDE² - **Note:** The simulation-based power registry (`simulate_power`, `simulate_mde`, `simulate_sample_size`) uses a single-cohort staggered DGP by default. Estimators configured with `control_group="not_yet_treated"`, `clean_control="strict"`, or `anticipation>0` will receive a `UserWarning` because the default DGP does not match their identification strategy. Users must supply `data_generator_kwargs` (e.g., `cohort_periods=[2, 4]`, `never_treated_frac=0.0`) or a custom `data_generator` to match the estimator design. - **Note:** The `TripleDifference` registry adapter uses `generate_ddd_data`, a fixed 2×2×2 factorial DGP (group × partition × time). The `n_periods`, `treatment_period`, and `treatment_fraction` parameters are ignored — DDD always simulates 2 periods with balanced groups. `n_units` is mapped to `n_per_cell = max(2, n_units // 8)` (effective total N = `n_per_cell × 8`), so non-multiples of 8 are rounded down and values below 16 are clamped to 16. A `UserWarning` is emitted when simulation inputs differ from the effective DDD design. When rounding occurs, all result objects (`SimulationPowerResults`, `SimulationMDEResults`, `SimulationSampleSizeResults`) set `effective_n_units` to the actual sample size used; it is `None` when no rounding occurred. `simulate_sample_size()` snaps bisection candidates to multiples of 8 so that `required_n` is always a realizable DDD sample size. Passing `n_per_cell` in `data_generator_kwargs` suppresses the effective-N rounding warning but not warnings for ignored parameters (`n_periods`, `treatment_period`, `treatment_fraction`). - **Note:** The analytical power methods (`PowerAnalysis.power/mde/sample_size` and the `compute_power/compute_mde/compute_sample_size` convenience functions) accept a `deff` parameter (survey design effect, default 1.0). This inflates variance multiplicatively: `Var(ATT) *= deff`, and inflates required sample size: `n_total *= deff`. The `deff` parameter is **not redundant** with `rho` (intra-cluster correlation): `rho` models within-unit serial correlation in panel data via the Moulton factor `1 + (T-1)*rho`, while `deff` models the survey design effect from stratified multi-stage sampling (clustering + unequal weighting). A survey panel study may need both. Values `deff > 0` are accepted; `deff < 1.0` (net variance reduction, e.g., from stratification gain) emits a warning. -- **Note:** The simulation-based power functions (`simulate_power/simulate_mde/simulate_sample_size`) accept a `survey_config` parameter (`SurveyPowerConfig` dataclass). When set, the simulation loop uses `generate_survey_did_data` instead of the default registry DGP, and automatically injects `SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc")` into the estimator's `fit()` call. Supported estimators: DifferenceInDifferences, TwoWayFixedEffects, MultiPeriodDiD, CallawaySantAnna, SunAbraham, ImputationDiD, TwoStageDiD, StackedDiD, EfficientDiD. Unsupported (raises `ValueError`): TROP, SyntheticDiD, TripleDifference (generate_survey_did_data produces staggered cohort data incompatible with factor-model/DDD DGPs). `survey_config` and `data_generator` are mutually exclusive. `data_generator_kwargs` may not contain keys managed by `SurveyPowerConfig` (n_strata, psu_per_stratum, etc.) but may contain passthrough DGP params (unit_fe_sd, add_covariates, strata_sizes). Repeated cross-section survey power (`panel=False`) is only supported for `CallawaySantAnna(panel=False)` with a matching `data_generator_kwargs={"panel": False}`; both mismatch directions are rejected. `estimator_kwargs` may not contain `survey_design` when `survey_config` is set (use `SurveyPowerConfig(survey_design=...)` instead). Estimator settings that require a multi-cohort DGP (`control_group="not_yet_treated"`, `clean_control="strict"`) are rejected because the survey DGP uses a single cohort; use the custom `data_generator` path for these configurations. `simulate_sample_size` raises the bisection floor to `n_strata * psu_per_stratum * 2` to ensure viable survey structure and rejects `strata_sizes` in `data_generator_kwargs` (it depends on `n_units` which varies during bisection). +- **Note:** The simulation-based power functions (`simulate_power/simulate_mde/simulate_sample_size`) accept a `survey_config` parameter (`SurveyPowerConfig` dataclass). When set, the simulation loop uses `generate_survey_did_data` instead of the default registry DGP, and automatically injects `SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc")` into the estimator's `fit()` call. Supported estimators: DifferenceInDifferences, TwoWayFixedEffects, MultiPeriodDiD, CallawaySantAnna, SunAbraham, ImputationDiD, TwoStageDiD, StackedDiD, EfficientDiD. Unsupported (raises `ValueError`): TROP, SyntheticDiD, TripleDifference (generate_survey_did_data produces staggered cohort data incompatible with factor-model/DDD DGPs). `survey_config` and `data_generator` are mutually exclusive. `data_generator_kwargs` may not contain keys managed by `SurveyPowerConfig` (n_strata, psu_per_stratum, etc.) but may contain passthrough DGP params (unit_fe_sd, add_covariates, strata_sizes). Repeated cross-section survey power (`panel=False`) is only supported for `CallawaySantAnna(panel=False)` with a matching `data_generator_kwargs={"panel": False}`; both mismatch directions are rejected. `estimator_kwargs` may not contain `survey_design` when `survey_config` is set (use `SurveyPowerConfig(survey_design=...)` instead). Estimator settings that require a multi-cohort DGP (`control_group="not_yet_treated"`, `control_group="last_cohort"`, `clean_control="strict"`) are rejected because the survey DGP uses a single cohort; use the custom `data_generator` path for these configurations. `simulate_sample_size` raises the bisection floor to `n_strata * psu_per_stratum * 2` to ensure viable survey structure and rejects `strata_sizes` in `data_generator_kwargs` (it depends on `n_units` which varies during bisection). **Reference implementation(s):** - R: `pwr` package (general), `DeclareDesign` (simulation-based) diff --git a/tests/test_power.py b/tests/test_power.py index abfd5b6c8..7a8c74be4 100644 --- a/tests/test_power.py +++ b/tests/test_power.py @@ -2602,6 +2602,17 @@ def test_survey_rejects_not_yet_treated(self): **_SIM_KW, ) + def test_survey_rejects_last_cohort(self): + """control_group='last_cohort' rejected (needs multi-cohort DGP).""" + with pytest.raises(ValueError, match="last_cohort"): + simulate_power( + EfficientDiD(control_group="last_cohort"), + survey_config=_SURVEY_CFG, + n_simulations=1, + seed=42, + **_SIM_KW, + ) + def test_survey_rejects_clean_control_strict(self): """clean_control='strict' rejected (needs multi-cohort DGP).""" with pytest.raises(ValueError, match="clean_control.*strict"): @@ -2629,3 +2640,15 @@ def test_survey_sample_size_rejects_strata_sizes(self): sigma=1.0, progress=False, ) + + def test_survey_config_validation_psu_re_sd_negative(self): + with pytest.raises(ValueError, match="psu_re_sd"): + SurveyPowerConfig(psu_re_sd=-1.0) + + def test_survey_config_validation_psu_re_sd_nan(self): + with pytest.raises(ValueError, match="psu_re_sd"): + SurveyPowerConfig(psu_re_sd=np.nan) + + def test_survey_config_validation_fpc_nan(self): + with pytest.raises(ValueError, match="fpc_per_stratum must be finite"): + SurveyPowerConfig(fpc_per_stratum=np.inf) From 87661d7943bdcc1b1e9db77bb4308f46891d8ce4 Mon Sep 17 00:00:00 2001 From: igerber Date: Sun, 12 Apr 2026 16:24:39 -0400 Subject: [PATCH 12/12] fix: address CI review round 10 - _snap_n floor enforcement for n_range - P1: fix _snap_n() to honor floor even when grid_step==1 (non-DDD), so explicit n_range with survey_config is clamped to min_viable_n - Add regression test for n_range=(10, 200) with min_viable_n=80 Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/power.py | 4 ++-- tests/test_power.py | 19 +++++++++++++++++++ 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/diff_diff/power.py b/diff_diff/power.py index 51f12441a..87959b0a8 100644 --- a/diff_diff/power.py +++ b/diff_diff/power.py @@ -2834,9 +2834,9 @@ def simulate_sample_size( ) def _snap_n(n: int, direction: str = "down", floor: Optional[int] = None) -> int: - if grid_step == 1: - return n actual_floor = floor if floor is not None else min_n + if grid_step == 1: + return max(actual_floor, n) if direction == "up": return max(actual_floor, ((n + grid_step - 1) // grid_step) * grid_step) return max(actual_floor, (n // grid_step) * grid_step) diff --git a/tests/test_power.py b/tests/test_power.py index 7a8c74be4..793f56c98 100644 --- a/tests/test_power.py +++ b/tests/test_power.py @@ -2395,6 +2395,25 @@ def test_survey_sample_size_large_floor_auto_bracket(self): ) assert result.required_n >= cfg.min_viable_n # must be >= 200 + def test_survey_sample_size_explicit_n_range_clamped(self): + """Explicit n_range below survey floor is clamped to min_viable_n.""" + cfg = SurveyPowerConfig(n_strata=5, psu_per_stratum=8) + # n_range=(10, 200) but min_viable_n=80, so lo should be clamped to 80 + result = simulate_sample_size( + CallawaySantAnna(), + treatment_effect=3.0, + sigma=1.0, + n_range=(10, 200), + n_simulations=10, + max_steps=3, + seed=42, + survey_config=cfg, + n_periods=4, + treatment_period=2, + progress=False, + ) + assert result.required_n >= cfg.min_viable_n # must be >= 80 + def test_survey_rejects_heterogeneous_te(self): """heterogeneous_te_by_strata=True rejected with simulation power.""" cfg = SurveyPowerConfig(heterogeneous_te_by_strata=True)