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:

"""

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,

}

)

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:

"""

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,

}

)

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:

"""

rng = np.random.default_rng(seed)

n_control = n_units - n_treated

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,

}

)

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:

"""

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

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:

"""

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,

}

)

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:

"""

rng = np.random.default_rng(seed)

n_periods = n_pre + n_post

treatment_period = n_pre

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

event_time = period - treatment_period

# Common time trend

time_effect = period * 0.5

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,

"event_time": event_time,

"true_effect": effect,

}

)

n_units: int = 500,

n_periods: int = 4,

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

never_treated_frac: float = 0.3,

dose_distribution: str = "lognormal",

dose_params: Optional[Dict] = None,

att_function: str = "linear",

att_slope: float = 2.0,

att_intercept: float = 1.0,

unit_fe_sd: float = 2.0,

time_trend: float = 0.5,

noise_sd: float = 1.0,

seed: Optional[int] = None,

) -> pd.DataFrame:

"""

rng = np.random.default_rng(seed)

if cohort_periods is None:

cohort_periods = [2]

# Assign units to cohorts

n_never = int(n_units * never_treated_frac)

n_treated_total = n_units - n_never

n_per_cohort = n_treated_total // len(cohort_periods)

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

idx = 0

for i, g in enumerate(cohort_periods):

n_this = n_per_cohort if i < len(cohort_periods) - 1 else n_treated_total - idx

cohort_assignments[n_never + idx : n_never + idx + n_this] = g

idx += n_this

# Generate doses

default_params = {

"lognormal": {"mean": 0.5, "sigma": 0.5},

"uniform": {"low": 0.5, "high": 5.0},

"exponential": {"scale": 2.0},

}

params = dose_params or default_params.get(dose_distribution, {})

dose_per_unit = np.zeros(n_units)

treated_mask = cohort_assignments > 0

n_treated_actual = int(np.sum(treated_mask))

if dose_distribution == "lognormal":

dose_per_unit[treated_mask] = rng.lognormal(

mean=params.get("mean", 0.5),

sigma=params.get("sigma", 0.5),

size=n_treated_actual,

)

elif dose_distribution == "uniform":

dose_per_unit[treated_mask] = rng.uniform(

low=params.get("low", 0.5),

high=params.get("high", 5.0),

size=n_treated_actual,

)

elif dose_distribution == "exponential":

dose_per_unit[treated_mask] = rng.exponential(

scale=params.get("scale", 2.0),

size=n_treated_actual,

)

else:

raise ValueError(

f"dose_distribution must be 'lognormal', 'uniform', or 'exponential', "

f"got '{dose_distribution}'"

)

# ATT function

def _att_func(d):

if att_function == "linear":

return att_intercept + att_slope * d

elif att_function == "quadratic":

return att_intercept + att_slope * d**2

elif att_function == "log":

return att_intercept + att_slope * np.log1p(d)

else:

raise ValueError(

f"att_function must be 'linear', 'quadratic', or 'log', " f"got '{att_function}'"

)

# Unit fixed effects

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

# Build panel

periods = np.arange(1, n_periods + 1)

records = []

for i in range(n_units):

g_i = cohort_assignments[i]

d_i = dose_per_unit[i]

for t in periods:

# Potential outcome without treatment

y0 = unit_fe[i] + time_trend * t + rng.normal(0, noise_sd)

# Treatment effect

if g_i > 0 and t >= g_i:

att_d = _att_func(d_i)

else:

att_d = 0.0

records.append(

{

"unit": i,

"period": int(t),

"outcome": y0 + att_d,

"first_treat": int(g_i) if g_i > 0 else 0,

"dose": d_i,

"true_att": att_d,

}

)

n_units: int = 200,