"""

Data generation utilities for difference-in-differences analysis.

This module provides functions to generate synthetic datasets for testing

and validating DiD estimators, including basic 2x2 DiD, staggered adoption,

factor model data, triple difference, and event study designs.

"""

from typing import Dict, List, Optional

import numpy as np

import pandas as pd

def generate_did_data(

n_units: int = 100,

n_periods: int = 4,

treatment_effect: float = 5.0,

treatment_fraction: float = 0.5,

treatment_period: int = 2,

unit_fe_sd: float = 2.0,

time_trend: float = 0.5,

noise_sd: float = 1.0,

seed: Optional[int] = None,

) -> pd.DataFrame:

"""

Generate synthetic data for DiD analysis with known treatment effect.

Creates a balanced panel dataset with realistic features including

unit fixed effects, time trends, and a known treatment effect.

Parameters

----------

n_units : int, default=100

Number of units in the panel.

n_periods : int, default=4

Number of time periods.

treatment_effect : float, default=5.0

True average treatment effect on the treated.

treatment_fraction : float, default=0.5

Fraction of units that receive treatment.

treatment_period : int, default=2

First post-treatment period (0-indexed). Periods >= this are post.

unit_fe_sd : float, default=2.0

Standard deviation of unit fixed effects.

time_trend : float, default=0.5

Linear time trend coefficient.

noise_sd : float, default=1.0

Standard deviation of idiosyncratic noise.

seed : int, optional

Random seed for reproducibility.

Returns

-------

pd.DataFrame

Synthetic panel data with columns:

- unit: Unit identifier

- period: Time period

- treated: Treatment indicator (0/1)

- post: Post-treatment indicator (0/1)

- outcome: Outcome variable

- true_effect: The true treatment effect (for validation)

Examples

--------

Generate simple data for testing:

>>> data = generate_did_data(n_units=50, n_periods=4, treatment_effect=3.0, seed=42)

>>> len(data)

200

>>> data.columns.tolist()

['unit', 'period', 'treated', 'post', 'outcome', 'true_effect']

Verify treatment effect recovery:

>>> from diff_diff import DifferenceInDifferences

>>> did = DifferenceInDifferences()

>>> results = did.fit(data, outcome='outcome', treatment='treated', post='post')

>>> abs(results.att - 3.0) < 1.0 # Close to true effect

True

"""

rng = np.random.default_rng(seed)

# Determine treated units

n_treated = int(n_units * treatment_fraction)

treated_units = set(range(n_treated))

# Generate unit fixed effects

unit_fe = rng.normal(0, unit_fe_sd, n_units)

# Build data

records = []

for unit in range(n_units):

is_treated = unit in treated_units

for period in range(n_periods):

is_post = period >= treatment_period

# Base outcome

y = 10.0 # Baseline

y += unit_fe[unit] # Unit fixed effect

y += time_trend * period # Time trend

# Treatment effect (only for treated units in post-period)

effect = 0.0

if is_treated and is_post:

effect = treatment_effect

y += effect

# Add noise

y += rng.normal(0, noise_sd)

records.append(

{

"unit": unit,

"period": period,

"treated": int(is_treated),

"post": int(is_post),

"outcome": y,

"true_effect": effect,

}

)

return pd.DataFrame(records)

def generate_staggered_data(

n_units: int = 100,

n_periods: int = 10,

cohort_periods: Optional[List[int]] = None,

never_treated_frac: float = 0.3,

treatment_effect: float = 2.0,

dynamic_effects: bool = True,

effect_growth: float = 0.1,

unit_fe_sd: float = 2.0,

time_trend: float = 0.1,

noise_sd: float = 0.5,

seed: Optional[int] = None,

panel: bool = True,

) -> pd.DataFrame:

"""

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

-------

pd.DataFrame

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

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

"""

rng = np.random.default_rng(seed)

# Default cohort periods if not specified

if cohort_periods is None:

cohort_periods = [3, 5, 7] if n_periods >= 8 else [n_periods // 3, 2 * n_periods // 3]

# Validate cohort periods

for cp in cohort_periods:

if cp < 1 or cp >= n_periods:

raise ValueError(f"Cohort period {cp} must be between 1 and {n_periods - 1}")

# Determine number of never-treated and treated units

n_never = int(n_units * never_treated_frac)

n_treated = n_units - n_never

if not panel:

# --- Repeated cross-section mode ---

# Each period draws n_units independent observations with unique IDs.

# Cohorts are assigned from the same distribution as panel.

records = []

for period in range(n_periods):

# For each period, draw fresh cohort assignments

ft_period = np.zeros(n_units, dtype=int)

if n_treated > 0:

cohort_assignments = rng.choice(len(cohort_periods), size=n_treated)

ft_period[n_never:] = [cohort_periods[c] for c in cohort_assignments]

# Unique unit IDs per period

for i in range(n_units):

uid = f"u{period}_{i}"

unit_first_treat = ft_period[i]

is_ever_treated = unit_first_treat > 0

is_treated = is_ever_treated and period >= unit_first_treat

# Outcome: unit_fe_proxy (drawn fresh) + time trend + treatment + noise

unit_fe_proxy = rng.normal(0, unit_fe_sd)

y = 10.0 + unit_fe_proxy + time_trend * period

effect = 0.0

if is_treated:

time_since_treatment = period - unit_first_treat

if dynamic_effects:

effect = treatment_effect * (1 + effect_growth * time_since_treatment)

else:

effect = treatment_effect

y += effect

y += rng.normal(0, noise_sd)

records.append(

{

"unit": uid,

"period": period,

"outcome": y,

"first_treat": unit_first_treat,

"treated": int(is_treated),

"treat": int(is_ever_treated),

"true_effect": effect,

}

)

return pd.DataFrame(records)

# --- Panel mode (default) ---

# Assign treatment cohorts

first_treat = np.zeros(n_units, dtype=int)

if n_treated > 0:

cohort_assignments = rng.choice(len(cohort_periods), size=n_treated)

first_treat[n_never:] = [cohort_periods[c] for c in cohort_assignments]

# Generate unit fixed effects

unit_fe = rng.normal(0, unit_fe_sd, n_units)

# Build data

records = []

for unit in range(n_units):

unit_first_treat = first_treat[unit]

is_ever_treated = unit_first_treat > 0

for period in range(n_periods):

# Check if treated at this observation

is_treated = is_ever_treated and period >= unit_first_treat

# Base outcome: unit FE + time trend

y = 10.0 + unit_fe[unit] + time_trend * period

# Treatment effect

effect = 0.0

if is_treated:

time_since_treatment = period - unit_first_treat

if dynamic_effects:

effect = treatment_effect * (1 + effect_growth * time_since_treatment)

else:

effect = treatment_effect

y += effect

# Add noise

y += rng.normal(0, noise_sd)

records.append(

{

"unit": unit,

"period": period,

"outcome": y,

"first_treat": unit_first_treat,

"treated": int(is_treated),

"treat": int(is_ever_treated),

"true_effect": effect,

}

)

return pd.DataFrame(records)

def generate_factor_data(

n_units: int = 50,

n_pre: int = 10,

n_post: int = 5,

n_treated: int = 10,

n_factors: int = 2,

treatment_effect: float = 2.0,

factor_strength: float = 1.0,

treated_loading_shift: float = 0.5,

unit_fe_sd: float = 1.0,

noise_sd: float = 0.5,

seed: Optional[int] = None,

) -> pd.DataFrame:

"""

Generate synthetic panel data with interactive fixed effects (factor model).

Creates data following the DGP:

Y_it = mu + alpha_i + beta_t + Lambda_i'F_t + tau*D_it + eps_it

where Lambda_i'F_t is the interactive fixed effects component. Useful for

testing TROP (Triply Robust Panel) and comparing with SyntheticDiD.

Parameters

----------

n_units : int, default=50

Total number of units in the panel.

n_pre : int, default=10

Number of pre-treatment periods.

n_post : int, default=5

Number of post-treatment periods.

n_treated : int, default=10

Number of treated units (assigned to first n_treated unit IDs).

n_factors : int, default=2

Number of latent factors in the interactive fixed effects.

treatment_effect : float, default=2.0

True average treatment effect on the treated.

factor_strength : float, default=1.0

Scaling factor for interactive fixed effects.

treated_loading_shift : float, default=0.5

Shift in factor loadings for treated units (creates confounding).

unit_fe_sd : float, default=1.0

Standard deviation of unit fixed effects.

noise_sd : float, default=0.5

Standard deviation of idiosyncratic noise.

seed : int, optional

Random seed for reproducibility.

Returns

-------

pd.DataFrame

Synthetic factor model data with columns:

- unit: Unit identifier

- period: Time period

- outcome: Outcome variable

- 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

Examples

--------

Generate data with factor structure:

>>> data = generate_factor_data(n_units=50, n_factors=2, seed=42)

>>> data.shape

(750, 6)

Use with TROP estimator:

>>> from diff_diff import TROP

>>> trop = TROP(n_bootstrap=50, seed=42)

>>> results = trop.fit(data, outcome='outcome', treatment='treated',

... unit='unit', time='period',

... post_periods=list(range(10, 15)))

Notes

-----

The treated units have systematically different factor loadings

(shifted by `treated_loading_shift`), which creates confounding

that standard DiD cannot address but TROP can handle.

"""

rng = np.random.default_rng(seed)

n_periods = n_pre + n_post

if n_treated > n_units:

raise ValueError(f"n_treated ({n_treated}) cannot exceed n_units ({n_units})")

if n_treated < 1:

raise ValueError("n_treated must be at least 1")

# Generate factors F: (n_periods, n_factors)

F = rng.normal(0, 1, (n_periods, n_factors))

# Generate loadings Lambda: (n_factors, n_units)

# Treated units have shifted loadings (creates confounding)

Lambda = rng.normal(0, 1, (n_factors, n_units))

Lambda[:, :n_treated] += treated_loading_shift

# Unit fixed effects (treated units have higher baseline)

alpha = rng.normal(0, unit_fe_sd, n_units)

alpha[:n_treated] += 1.0

# Time fixed effects (linear trend)

beta = np.linspace(0, 2, n_periods)

# Generate outcomes

records = []

for i in range(n_units):

is_ever_treated = i < n_treated

for t in range(n_periods):

post = t >= n_pre

# Base outcome

y = 10.0 + alpha[i] + beta[t]

# Interactive fixed effects: Lambda_i' F_t

y += factor_strength * (Lambda[:, i] @ F[t, :])

# Treatment effect

effect = 0.0

if is_ever_treated and post:

effect = treatment_effect

y += effect

# Add noise

y += rng.normal(0, noise_sd)

records.append(

{

"unit": i,

"period": t,

"outcome": y,

"treated": int(is_ever_treated and post),

"treat": int(is_ever_treated),

"true_effect": effect,

}

)

return pd.DataFrame(records)

def generate_synthetic_control_data(

n_donors: int = 20,

n_pre: int = 24,

n_post: int = 6,

n_factors: int = 3,

n_predictors: int = 3,

treatment_effect: float = 5.0,

effect_type: str = "ramp",

effect_growth: float = 1.0,

factor_strength: float = 1.0,

factor_persistence: float = 0.9,

n_convex_donors: int = 4,

baseline: float = 50.0,

unit_baseline_sd: float = 5.0,

time_trend: float = 0.5,

predictor_noise_sd: float = 0.5,

noise_sd: float = 1.0,

seed: Optional[int] = None,

) -> pd.DataFrame:

"""

Generate a single-treated-unit panel for synthetic control (ADH) demos.

Creates a factor-model panel with ONE treated unit (``unit == 0``) and

``n_donors`` never-treated donors, designed so the treated unit's underlying

(noiseless) trajectory lies in the span of the donor trajectories — so a synthetic

control can reproduce it closely. The DGP:

Y_it = b_i + beta_t + factor_strength * (Lambda_i . F_t) + tau_it + eps_it

where ``F_t`` is a persistent (AR(1)) latent factor path, ``Lambda_i`` are

nonnegative factor loadings, ``b_i`` is a unit baseline, ``beta_t`` is a common

calendar time effect, ``tau_it`` is the treatment effect (treated unit, post

periods only), and ``eps_it ~ N(0, noise_sd)``.

The treated unit's loadings and baseline are an exact convex combination of

``n_convex_donors`` donors (Dirichlet weights), so in the **noiseless** limit the

treated unit's outcome path is itself that convex combination of donor paths — it

lies in the donor convex hull and a synthetic control reproduces it exactly. With

nonzero ``noise_sd`` / ``predictor_noise_sd`` the observed data is only

*approximately* in-hull, so a fitted synthetic control achieves a small (not

exactly zero) pre-period RMSPE. Note that the generating weights are NOT identified

by the estimator: many donor combinations reproduce the treated pre-period path

equally well, so a fitted synthetic control recovers the counterfactual outcome path

(and ATT), not this specific weight vector.

Unlike :func:`generate_factor_data` (which *shifts* treated loadings to create

confounding for TROP), this generator keeps the treated unit reproducible from

donors, which is what classic synthetic control requires.

Parameters

----------

n_donors : int, default=20

Number of never-treated donor units. The in-space placebo p-value floor is

``1 / (n_donors + 1)``; larger pools give finer permutation p-values.

n_pre : int, default=24

Number of pre-treatment periods.

n_post : int, default=6

Number of post-treatment periods.

n_factors : int, default=3

Number of latent factors driving the common time structure.

n_predictors : int, default=3

Number of time-invariant predictor covariates (each a noisy linear map of

the unit's loading vector). Pass these as ``predictors=`` to

:class:`~diff_diff.SyntheticControl` to exercise the V-search; otherwise

the fit defaults to pre-period outcomes and these columns are unused.

treatment_effect : float, default=5.0

Base treatment effect on the treated unit in post periods. Under

``effect_type="ramp"`` this is the first post-period effect.

effect_type : str, default="ramp"

``"ramp"`` (effect grows by ``effect_growth`` each post period) or

``"constant"`` (effect fixed at ``treatment_effect``).

effect_growth : float, default=1.0

Per-post-period increment under ``effect_type="ramp"`` (ignored otherwise).

factor_strength : float, default=1.0

Scaling on the interactive (loading . factor) component.

factor_persistence : float, default=0.9

AR(1) coefficient on the latent factors in ``[0, 1]``: ``0`` gives i.i.d.

factors, ``1`` a random walk. Serial dependence is what makes the

moving-block conformal scheme meaningful.

n_convex_donors : int, default=4

Number of donors whose convex combination defines the treated unit (must be

in ``[1, n_donors]``).

baseline : float, default=50.0

Mean unit baseline (e.g. a per-capita outcome level).

unit_baseline_sd : float, default=5.0

Standard deviation of donor baselines.

time_trend : float, default=0.5

Slope of the common calendar time effect ``beta_t = time_trend * t`` (a

common trend is matched by any convex combination, so it preserves in-hull

existence). Set to ``0.0`` for no trend.

predictor_noise_sd : float, default=0.5

Noise added to each predictor covariate.

noise_sd : float, default=1.0

Standard deviation of idiosyncratic outcome noise.

seed : int, optional

Random seed for reproducibility.

Returns

-------

pd.DataFrame

Long-format balanced panel with columns:

- ``unit``: unit identifier (``0`` is the treated unit; ``1..n_donors`` donors)

- ``period``: time period (``0..n_pre-1`` pre, ``n_pre..`` post)

- ``outcome``: outcome variable

- ``treatment``: absorbing 0/1 indicator (1 only for the treated unit in

post periods) — :meth:`SyntheticControl.fit` infers the treated unit and

post periods from this column

- ``treat``: unit-level ever-treated flag

- ``true_effect``: the true treatment effect for this observation

- ``x1 .. x{n_predictors}``: predictor covariates

Examples

--------

Generate a ramped-effect panel and fit a synthetic control:

>>> from diff_diff import generate_synthetic_control_data, SyntheticControl

>>> data = generate_synthetic_control_data(seed=42)

>>> res = SyntheticControl(seed=42).fit(

... data, outcome="outcome", treatment="treatment",

... unit="unit", time="period", predictors=["x1", "x2", "x3"])

>>> round(res.att, 1) > 0 # post-period gap recovers the injected effect

True

"""

if n_donors < 2:

raise ValueError(f"n_donors ({n_donors}) must be at least 2")

if n_pre < 1:

raise ValueError(f"n_pre ({n_pre}) must be at least 1")

if n_post < 1:

raise ValueError(f"n_post ({n_post}) must be at least 1")

if n_factors < 1:

raise ValueError(f"n_factors ({n_factors}) must be at least 1")

if n_predictors < 0:

raise ValueError(f"n_predictors ({n_predictors}) must be non-negative")

if n_convex_donors < 1 or n_convex_donors > n_donors:

raise ValueError(f"n_convex_donors ({n_convex_donors}) must be in [1, n_donors={n_donors}]")

if effect_type not in ("ramp", "constant"):

raise ValueError(f"effect_type ({effect_type!r}) must be 'ramp' or 'constant'")

if not (0.0 <= factor_persistence <= 1.0):

raise ValueError(f"factor_persistence ({factor_persistence}) must be in [0, 1]")

rng = np.random.default_rng(seed)

n_periods = n_pre + n_post

n_units = n_donors + 1 # unit 0 is treated, 1..n_donors are donors

# Persistent latent factors F: (n_periods, n_factors). AR(1): rho=1 -> random

# walk, rho=0 -> i.i.d. The synthetic control removes this common structure, so

# the post-fit residuals stay stationary regardless of persistence.

F = np.empty((n_periods, n_factors))

F[0] = rng.normal(0, 1, n_factors)

for t in range(1, n_periods):

F[t] = factor_persistence * F[t - 1] + rng.normal(0, 1, n_factors)

# Donor loadings (nonnegative so a convex combination is well-defined) and

# baselines.

lambda_donor = np.abs(rng.normal(0, 1, (n_factors, n_donors)))

b_donor = rng.normal(baseline, unit_baseline_sd, n_donors)

# Treated unit = exact convex combination of n_convex_donors donors (in-hull).

convex_idx = rng.choice(n_donors, size=n_convex_donors, replace=False)

w_star = np.zeros(n_donors)

w_star[convex_idx] = rng.dirichlet(np.ones(n_convex_donors))

lambda_treated = lambda_donor @ w_star

b_treated = float(b_donor @ w_star)

# Stack: column 0 = treated, columns 1.. = donors.

loadings = np.column_stack([lambda_treated, lambda_donor]) # (n_factors, n_units)

baselines = np.concatenate([[b_treated], b_donor]) # (n_units,)

# Time-invariant predictors: each PRIMARILY proxies one distinct latent factor

# (diagonal-dominant map) plus a small cross-factor mix and noise. Distinct

# proxies give each predictor complementary signal, so the V-search spreads

# weight across them instead of leaning on a single redundant column.

pred_map = 0.25 * rng.normal(0, 1, (n_predictors, n_factors))

for h in range(n_predictors):

pred_map[h, h % n_factors] += 1.0

predictors = pred_map @ loadings + rng.normal(

0, predictor_noise_sd, (n_predictors, n_units)

) # (n_predictors, n_units)

# Common calendar time effect (matched by any convex combination since sum w=1).

beta = time_trend * np.arange(n_periods)

records = []

for i in range(n_units):

is_treated_unit = i == 0

for t in range(n_periods):

post = t >= n_pre

y = baselines[i] + beta[t] + factor_strength * (loadings[:, i] @ F[t])

effect = 0.0

if is_treated_unit and post:

k = t - n_pre # post-period index, 0-based

if effect_type == "ramp":

effect = treatment_effect + effect_growth * k

else:

effect = treatment_effect

y += effect

y += rng.normal(0, noise_sd)

row = {

"unit": i,

"period": t,

"outcome": y,

"treatment": int(is_treated_unit and post),

"treat": int(is_treated_unit),

"true_effect": effect,

}

for h in range(n_predictors):

row[f"x{h + 1}"] = predictors[h, i]

records.append(row)

return pd.DataFrame(records)

def generate_ddd_data(

n_per_cell: int = 100,

treatment_effect: float = 2.0,

group_effect: float = 2.0,

partition_effect: float = 1.0,

time_effect: float = 0.5,

noise_sd: float = 1.0,

add_covariates: bool = False,

seed: Optional[int] = None,

) -> pd.DataFrame:

"""

Generate synthetic data for Triple Difference (DDD) analysis.

Creates data following the DGP:

Y = mu + G + P + T + G*P + G*T + P*T + tau*G*P*T + eps

where G=group, P=partition, T=time. The treatment effect (tau) only

applies to units that are in the treated group (G=1), eligible

partition (P=1), and post-treatment period (T=1).

Parameters

----------

n_per_cell : int, default=100

Number of observations per cell (8 cells total: 2x2x2).

treatment_effect : float, default=2.0

True average treatment effect on the treated (G=1, P=1, T=1).

group_effect : float, default=2.0

Main effect of being in treated group.

partition_effect : float, default=1.0

Main effect of being in eligible partition.

time_effect : float, default=0.5

Main effect of post-treatment period.

noise_sd : float, default=1.0

Standard deviation of idiosyncratic noise.

add_covariates : bool, default=False

If True, adds age and education covariates that affect outcome.

seed : int, optional

Random seed for reproducibility.

Returns

-------

pd.DataFrame

Synthetic DDD data with columns:

- outcome: Outcome variable

- group: Group indicator (0=control, 1=treated)

- partition: Partition indicator (0=ineligible, 1=eligible)

- time: Time indicator (0=pre, 1=post)

- unit_id: Unique unit identifier

- true_effect: The true treatment effect for this observation

- age: Age covariate (if add_covariates=True)

- education: Education covariate (if add_covariates=True)

Examples

--------

Generate DDD data:

>>> data = generate_ddd_data(n_per_cell=100, treatment_effect=3.0, seed=42)

>>> data.shape

(800, 6)

>>> data.groupby(['group', 'partition', 'time']).size()

group partition time

0 0 0 100

1 100

1 0 100

1 100

1 0 0 100

1 100

1 0 100

1 100

dtype: int64

Use with TripleDifference estimator:

>>> from diff_diff import TripleDifference

>>> ddd = TripleDifference()

>>> results = ddd.fit(data, outcome='outcome', group='group',

... partition='partition', post='time')

>>> abs(results.att - 3.0) < 1.0

True

"""

rng = np.random.default_rng(seed)

records = []

unit_id = 0

for g in [0, 1]: # group (0=control state, 1=treated state)

for p in [0, 1]: # partition (0=ineligible, 1=eligible)

for t in [0, 1]: # time (0=pre, 1=post)

for _ in range(n_per_cell):

# Base outcome with main effects

y = 50 + group_effect * g + partition_effect * p + time_effect * t

# Second-order interactions (non-treatment)

y += 1.5 * g * p # group-partition interaction

y += 1.0 * g * t # group-time interaction (diff trends)

y += 0.5 * p * t # partition-time interaction

# Treatment effect: ONLY for G=1, P=1, T=1

effect = 0.0

if g == 1 and p == 1 and t == 1:

effect = treatment_effect

y += effect

# Covariates (always generated for consistency)

age = rng.normal(40, 10)

education = rng.choice([12, 14, 16, 18], p=[0.3, 0.3, 0.25, 0.15])

if add_covariates:

y += 0.1 * age + 0.5 * education

# Add noise

y += rng.normal(0, noise_sd)

record = {

"outcome": y,

"group": g,

"partition": p,

"time": t,

"unit_id": unit_id,

"true_effect": effect,

}

if add_covariates:

record["age"] = age

record["education"] = education

records.append(record)

unit_id += 1

return pd.DataFrame(records)

def generate_panel_data(

n_units: int = 100,

n_periods: int = 8,

treatment_period: int = 4,

treatment_fraction: float = 0.5,

treatment_effect: float = 5.0,

parallel_trends: bool = True,

trend_violation: float = 1.0,

unit_fe_sd: float = 2.0,

noise_sd: float = 0.5,

seed: Optional[int] = None,

) -> pd.DataFrame:

"""

Generate synthetic panel data for parallel trends testing.

Creates panel data with optional violation of parallel trends, useful

for testing parallel trends diagnostics, placebo tests, and sensitivity

analysis methods.

Parameters

----------

n_units : int, default=100

Total number of units in the panel.

n_periods : int, default=8

Number of time periods.

treatment_period : int, default=4

First post-treatment period (0-indexed).

treatment_fraction : float, default=0.5

Fraction of units that receive treatment.

treatment_effect : float, default=5.0

True average treatment effect on the treated.

parallel_trends : bool, default=True

If True, treated and control groups have parallel pre-treatment trends.

If False, treated group has a steeper pre-treatment trend.

trend_violation : float, default=1.0

Size of the differential trend for treated group when parallel_trends=False.

Treated units have trend = common_trend + trend_violation.

unit_fe_sd : float, default=2.0

Standard deviation of unit fixed effects.

noise_sd : float, default=0.5

Standard deviation of idiosyncratic noise.

seed : int, optional

Random seed for reproducibility.

Returns

-------

pd.DataFrame

Synthetic panel data with columns:

- unit: Unit identifier

- period: Time period

- treated: Binary unit-level treatment indicator

- post: Binary post-treatment indicator

- outcome: Outcome variable

- true_effect: The true treatment effect for this observation

Examples

--------

Generate data with parallel trends:

>>> data_parallel = generate_panel_data(parallel_trends=True, seed=42)

>>> from diff_diff.utils import check_parallel_trends

>>> result = check_parallel_trends(data_parallel, outcome='outcome',

... time='period', treatment_group='treated',

... pre_periods=[0, 1, 2, 3])

>>> result['parallel_trends_plausible']

True

Generate data with trend violation:

>>> data_violation = generate_panel_data(parallel_trends=False, seed=42)

>>> result = check_parallel_trends(data_violation, outcome='outcome',

... time='period', treatment_group='treated',

... pre_periods=[0, 1, 2, 3])

>>> result['parallel_trends_plausible']

False

"""

rng = np.random.default_rng(seed)

if treatment_period < 1:

raise ValueError("treatment_period must be at least 1")

if treatment_period >= n_periods:

raise ValueError(f"treatment_period must be less than n_periods ({n_periods})")

n_treated = int(n_units * treatment_fraction)

records = []

for unit in range(n_units):

is_treated = unit < n_treated

unit_fe = rng.normal(0, unit_fe_sd)

for period in range(n_periods):

post = period >= treatment_period

# Base time effect (common trend)

if parallel_trends:

time_effect = period * 1.0

else:

# Different trends: treated has steeper pre-treatment trend

if is_treated:

time_effect = period * (1.0 + trend_violation)

else:

time_effect = period * 1.0

y = 10.0 + unit_fe + time_effect

# Treatment effect (only for treated in post-period)

effect = 0.0

if is_treated and post:

effect = treatment_effect

y += effect

# Add noise

y += rng.normal(0, noise_sd)

records.append(

{

"unit": unit,

"period": period,

"treated": int(is_treated),

"post": int(post),

"outcome": y,

"true_effect": effect,

}

)

return pd.DataFrame(records)

def generate_event_study_data(

n_units: int = 300,

n_pre: int = 5,

n_post: int = 5,

treatment_fraction: float = 0.5,

treatment_effect: float = 5.0,

unit_fe_sd: float = 2.0,

noise_sd: float = 2.0,

seed: Optional[int] = None,

) -> pd.DataFrame:

"""

Generate synthetic data for event study analysis.

Creates panel data with simultaneous treatment at period n_pre.

Useful for testing MultiPeriodDiD, pre-trends power analysis,

and HonestDiD sensitivity analysis.

Parameters

----------

n_units : int, default=300

Total number of units in the panel.

n_pre : int, default=5

Number of pre-treatment periods.

n_post : int, default=5

Number of post-treatment periods.

treatment_fraction : float, default=0.5

Fraction of units that receive treatment.

treatment_effect : float, default=5.0

True average treatment effect on the treated.

unit_fe_sd : float, default=2.0

Standard deviation of unit fixed effects.

noise_sd : float, default=2.0

Standard deviation of idiosyncratic noise.

seed : int, optional

Random seed for reproducibility.

Returns

-------

pd.DataFrame

Synthetic event study data with columns:

- unit: Unit identifier

- period: Time period

- treated: Binary unit-level treatment indicator

- post: Binary post-treatment indicator

- outcome: Outcome variable

- event_time: Time relative to treatment (negative=pre, 0+=post)

- true_effect: The true treatment effect for this observation

Examples

--------

Generate event study data:

>>> data = generate_event_study_data(n_units=300, n_pre=5, n_post=5, seed=42)

>>> data['event_time'].unique()

array([-5, -4, -3, -2, -1, 0, 1, 2, 3, 4])