diff_diff.generate_staggered_data#

diff_diff.generate_staggered_data(n_units=100, n_periods=10, cohort_periods=None, never_treated_frac=0.3, treatment_effect=2.0, dynamic_effects=True, effect_growth=0.1, unit_fe_sd=2.0, time_trend=0.1, noise_sd=0.5, seed=None, panel=True)[source]

Generate synthetic data for staggered adoption DiD analysis.

Creates panel data where different units receive treatment at different times (staggered rollout). Useful for testing CallawaySantAnna, SunAbraham, and other staggered DiD estimators.

Parameters:

  • n_units (int, default=100) – Total number of units in the panel.

  • n_periods (int, default=10) – Number of time periods.

  • cohort_periods (list of int, optional) – Periods when treatment cohorts are first treated. If None, defaults to [3, 5, 7] for a 10-period panel.

  • never_treated_frac (float, default=0.3) – Fraction of units that are never treated (cohort 0).

  • treatment_effect (float, default=2.0) – Base treatment effect at time of treatment.

  • dynamic_effects (bool, default=True) – If True, treatment effects grow over time since treatment.

  • effect_growth (float, default=0.1) – Per-period growth in treatment effect (if dynamic_effects=True). Effect at time t since treatment: effect * (1 + effect_growth * t).

  • unit_fe_sd (float, default=2.0) – Standard deviation of unit fixed effects.

  • time_trend (float, default=0.1) – Linear time trend coefficient.

  • noise_sd (float, default=0.5) – Standard deviation of idiosyncratic noise.

  • seed (int, optional) – Random seed for reproducibility.

  • panel (bool, default=True) – If True (default), generate balanced panel data (same units across all periods). If False, generate repeated cross-section data where each period draws independent observations with globally unique IDs.

Returns:

Synthetic staggered adoption data with columns: - unit: Unit identifier - period: Time period - outcome: Outcome variable - first_treat: First treatment period (0 = never treated) - treated: Binary indicator (1 if treated at this observation) - treat: Binary unit-level ever-treated indicator - true_effect: The true treatment effect for this observation

Return type:

pd.DataFrame

Examples

Generate staggered adoption data:

>>> data = generate_staggered_data(n_units=100, n_periods=10, seed=42)
>>> data['first_treat'].value_counts().sort_index()
0     30
3     24
5     23
7     23
Name: first_treat, dtype: int64

Use with Callaway-Sant’Anna estimator:

>>> from diff_diff import CallawaySantAnna
>>> cs = CallawaySantAnna()
>>> results = cs.fit(data, outcome='outcome', unit='unit',
...                  time='period', first_treat='first_treat')
>>> results.overall_att > 0
True