data: pd.DataFrame,

column: str,

treated_values: Optional[Union[Any, List[Any]]] = None,

threshold: Optional[float] = None,

above_threshold: bool = True,

new_column: str = "treated",

) -> pd.DataFrame:

"""

df = data.copy()

if treated_values is not None and threshold is not None:

raise ValueError("Specify either 'treated_values' or 'threshold', not both.")

if treated_values is None and threshold is None:

raise ValueError("Must specify either 'treated_values' or 'threshold'.")

if column not in df.columns:

raise ValueError(f"Column '{column}' not found in DataFrame.")

if treated_values is not None:

# Convert single value to list

if not isinstance(treated_values, (list, tuple, set)):

treated_values = [treated_values]

df[new_column] = df[column].isin(treated_values).astype(int)

else:

# Use threshold

if above_threshold:

df[new_column] = (df[column] >= threshold).astype(int)

else:

df[new_column] = (df[column] <= threshold).astype(int)

data: pd.DataFrame,

time_column: str,

post_periods: Optional[Union[Any, List[Any]]] = None,

treatment_start: Optional[Any] = None,

new_column: str = "post",

) -> pd.DataFrame:

"""

df = data.copy()

if post_periods is not None and treatment_start is not None:

raise ValueError("Specify either 'post_periods' or 'treatment_start', not both.")

if post_periods is None and treatment_start is None:

raise ValueError("Must specify either 'post_periods' or 'treatment_start'.")

if time_column not in df.columns:

raise ValueError(f"Column '{time_column}' not found in DataFrame.")

if post_periods is not None:

# Convert single value to list

if not isinstance(post_periods, (list, tuple, set)):

post_periods = [post_periods]

df[new_column] = df[time_column].isin(post_periods).astype(int)

else:

# Use treatment_start - convert to same type as column if needed

col_dtype = df[time_column].dtype

if pd.api.types.is_datetime64_any_dtype(col_dtype):

treatment_start = pd.to_datetime(treatment_start)

df[new_column] = (df[time_column] >= treatment_start).astype(int)

data: pd.DataFrame,

value_columns: List[str],

id_column: str,

time_name: str = "period",

value_name: str = "value",

time_values: Optional[List[Any]] = None,

) -> pd.DataFrame:

"""

if not value_columns:

raise ValueError("value_columns cannot be empty.")

if id_column not in data.columns:

raise ValueError(f"Column '{id_column}' not found in DataFrame.")

for col in value_columns:

if col not in data.columns:

raise ValueError(f"Column '{col}' not found in DataFrame.")

if time_values is None:

time_values = list(range(len(value_columns)))

if len(time_values) != len(value_columns):

raise ValueError(

f"time_values length ({len(time_values)}) must match "

f"value_columns length ({len(value_columns)})."

)

# Get other columns to preserve (not id or value columns)

other_cols = [c for c in data.columns if c != id_column and c not in value_columns]

# Use pd.melt for better performance (vectorized)

long_df = pd.melt(

data,

id_vars=[id_column] + other_cols,

value_vars=value_columns,

var_name="_temp_var",

value_name=value_name,

)

# Map column names to time values

col_to_time = dict(zip(value_columns, time_values))

long_df[time_name] = long_df["_temp_var"].map(col_to_time)

long_df = long_df.drop("_temp_var", axis=1)

# Reorder columns and sort

cols = [id_column, time_name, value_name] + other_cols

data: pd.DataFrame,

unit_column: str,

time_column: str,

method: str = "inner",

fill_value: Optional[float] = None,

) -> pd.DataFrame:

"""

if unit_column not in data.columns:

raise ValueError(f"Column '{unit_column}' not found in DataFrame.")

if time_column not in data.columns:

raise ValueError(f"Column '{time_column}' not found in DataFrame.")

if method not in ["inner", "outer", "fill"]:

raise ValueError(f"method must be 'inner', 'outer', or 'fill', got '{method}'")

all_units = data[unit_column].unique()

all_periods = sorted(data[time_column].unique())

n_periods = len(all_periods)

if method == "inner":

# Keep only units that have all periods

unit_counts = data.groupby(unit_column)[time_column].nunique()

complete_units = unit_counts[unit_counts == n_periods].index

return data[data[unit_column].isin(complete_units)].copy()

elif method in ["outer", "fill"]:

# Create full grid of unit-period combinations

full_index = pd.MultiIndex.from_product(

[all_units, all_periods], names=[unit_column, time_column]

)

full_df = pd.DataFrame(index=full_index).reset_index()

# Merge with original data

result = full_df.merge(data, on=[unit_column, time_column], how="left")

if method == "fill":

# Identify columns to fill (exclude unit and time columns)

cols_to_fill = [c for c in result.columns if c not in [unit_column, time_column]]

if fill_value is not None:

# Fill specified columns with fill_value

numeric_cols = result.select_dtypes(include=[np.number]).columns

for col in numeric_cols:

if col in cols_to_fill:

result[col] = result[col].fillna(fill_value)

else:

# Forward fill within each unit for non-key columns

result = result.sort_values([unit_column, time_column])

result[cols_to_fill] = result.groupby(unit_column)[cols_to_fill].ffill()

# Backward fill any remaining NaN at start

result[cols_to_fill] = result.groupby(unit_column)[cols_to_fill].bfill()

return result

data: pd.DataFrame,

outcome: str,

treatment: str,

time: str,

unit: Optional[str] = None,

raise_on_error: bool = True,

) -> Dict[str, Any]:

"""

errors = []

warnings = []

# Check columns exist

required_cols = [outcome, treatment, time]

if unit is not None:

required_cols.append(unit)

for col in required_cols:

if col not in data.columns:

errors.append(f"Required column '{col}' not found in DataFrame.")

if errors:

if raise_on_error:

raise ValueError("\n".join(errors))

return {"valid": False, "errors": errors, "warnings": warnings, "summary": {}}

# Check outcome is numeric

if not pd.api.types.is_numeric_dtype(data[outcome]):

errors.append(

f"Outcome column '{outcome}' must be numeric. " f"Got type: {data[outcome].dtype}"

)

# Check treatment is binary

treatment_vals = data[treatment].dropna().unique()

if not set(treatment_vals).issubset({0, 1}):

errors.append(

f"Treatment column '{treatment}' must be binary (0 or 1). "

f"Found values: {sorted(treatment_vals)}"

)

# Check time is binary for simple DiD

time_vals = data[time].dropna().unique()

if len(time_vals) == 2 and not set(time_vals).issubset({0, 1}):

warnings.append(

f"Time column '{time}' has 2 values but they are not 0 and 1: {sorted(time_vals)}. "

"For basic DiD, use 0 for pre-treatment and 1 for post-treatment."

)

# Check for missing values

for col in required_cols:

n_missing = data[col].isna().sum()

if n_missing > 0:

errors.append(

f"Column '{col}' has {n_missing} missing values. "

"Please handle missing data before fitting."

)

# Calculate summary statistics

summary = {}

if not errors:

summary["n_obs"] = len(data)

summary["n_treated"] = int((data[treatment] == 1).sum())

summary["n_control"] = int((data[treatment] == 0).sum())

summary["n_periods"] = len(time_vals)

if unit is not None:

summary["n_units"] = data[unit].nunique()

# Check for sufficient variation

if summary["n_treated"] == 0:

errors.append("No treated observations found (treatment column is all 0).")

if summary["n_control"] == 0:

errors.append("No control observations found (treatment column is all 1).")

# Check for each treatment-time combination

if len(time_vals) == 2:

# For 2-period DiD, check all four cells

for t_val in [0, 1]:

for p_val in time_vals:

count = len(data[(data[treatment] == t_val) & (data[time] == p_val)])

if count == 0:

errors.append(

f"No observations for treatment={t_val}, time={p_val}. "

"DiD requires observations in all treatment-time cells."

)

else:

# For multi-period, check that both treatment groups exist in multiple periods

for t_val in [0, 1]:

n_periods_with_obs = data[data[treatment] == t_val][time].nunique()

if n_periods_with_obs < 2:

group_name = "Treated" if t_val == 1 else "Control"

errors.append(

f"{group_name} group has observations in only {n_periods_with_obs} period(s). "

"DiD requires multiple periods per group."

)

# Panel-specific validation

if unit is not None and not errors:

# Check treatment is constant within units

unit_treatment_var = data.groupby(unit)[treatment].nunique()

units_with_varying_treatment = unit_treatment_var[unit_treatment_var > 1]

if len(units_with_varying_treatment) > 0:

warnings.append(

f"Treatment varies within {len(units_with_varying_treatment)} unit(s). "

"For standard DiD, treatment should be constant within units. "

"This may be intentional for staggered adoption designs."

)

# Check panel balance

periods_per_unit = data.groupby(unit)[time].nunique()

if periods_per_unit.min() != periods_per_unit.max():

warnings.append(

f"Unbalanced panel detected. Units have between "

f"{periods_per_unit.min()} and {periods_per_unit.max()} periods. "

"Consider using balance_panel() to balance the data."

)

valid = len(errors) == 0

if raise_on_error and not valid:

raise ValueError("Data validation failed:\n" + "\n".join(errors))

data: pd.DataFrame, outcome: str, treatment: str, time: str, unit: Optional[str] = None

) -> pd.DataFrame:

"""

# Group by treatment and time

summary = (

data.groupby([treatment, time])[outcome]

.agg([("n", "count"), ("mean", "mean"), ("std", "std"), ("min", "min"), ("max", "max")])

.round(4)

)

# Calculate time values for labeling

time_vals = sorted(data[time].unique())

# Add group labels based on sorted time values (not literal 0/1)

if len(time_vals) == 2:

pre_val, post_val = time_vals[0], time_vals[1]

def format_label(x: tuple) -> str:

treatment_label = "Treated" if x[0] == 1 else "Control"

time_label = "Post" if x[1] == post_val else "Pre"

return f"{treatment_label} - {time_label}"

summary.index = summary.index.map(format_label)

# Calculate means for each cell

treated_pre = data[(data[treatment] == 1) & (data[time] == pre_val)][outcome].mean()

treated_post = data[(data[treatment] == 1) & (data[time] == post_val)][outcome].mean()

control_pre = data[(data[treatment] == 0) & (data[time] == pre_val)][outcome].mean()

control_post = data[(data[treatment] == 0) & (data[time] == post_val)][outcome].mean()

# Calculate DiD

treated_diff = treated_post - treated_pre

control_diff = control_post - control_pre

did_estimate = treated_diff - control_diff

# Add to summary as a new row

did_row = pd.DataFrame(

{"n": ["-"], "mean": [did_estimate], "std": ["-"], "min": ["-"], "max": ["-"]},

index=["DiD Estimate"],

)

summary = pd.concat([summary, did_row])

else:

summary.index = summary.index.map(

lambda x: f"{'Treated' if x[0] == 1 else 'Control'} - Period {x[1]}"

)

data: pd.DataFrame, time_column: str, treatment_time_column: str, new_column: str = "event_time"

) -> pd.DataFrame:

"""

df = data.copy()

if time_column not in df.columns:

raise ValueError(f"Column '{time_column}' not found in DataFrame.")

if treatment_time_column not in df.columns:

raise ValueError(f"Column '{treatment_time_column}' not found in DataFrame.")

# Calculate event time

df[new_column] = df[time_column] - df[treatment_time_column]

# Handle never-treated (inf or NaN in treatment time)

col = df[treatment_time_column]

if pd.api.types.is_numeric_dtype(col):

never_treated = col.isna() | np.isinf(col)

else:

never_treated = col.isna()

df.loc[never_treated, new_column] = np.nan

data: pd.DataFrame,

unit_column: str,

time_column: str,

treatment_column: str,

outcome: str,

covariates: Optional[List[str]] = None,

) -> pd.DataFrame:

"""

agg_cols = {outcome: "mean", unit_column: "nunique"}

if covariates:

for cov in covariates:

agg_cols[cov] = "mean"

cohort_data = data.groupby([treatment_column, time_column]).agg(agg_cols).reset_index()

# Rename columns

cohort_data = cohort_data.rename(columns={unit_column: "n_units", outcome: f"mean_{outcome}"})

data: pd.DataFrame,

unit_column: str,

time_column: str,

outcome_column: str,

treatment_column: Optional[str] = None,

treated_units: Optional[List[Any]] = None,

pre_periods: Optional[List[Any]] = None,

covariates: Optional[List[str]] = None,

outcome_weight: float = 0.7,

covariate_weight: float = 0.3,

exclude_units: Optional[List[Any]] = None,

require_units: Optional[List[Any]] = None,

n_top: Optional[int] = None,

suggest_treatment_candidates: bool = False,

n_treatment_candidates: int = 5,

lambda_reg: float = 0.0,

) -> pd.DataFrame:

"""

# -------------------------------------------------------------------------

# Input validation

# -------------------------------------------------------------------------

for col in [unit_column, time_column, outcome_column]:

if col not in data.columns:

raise ValueError(f"Column '{col}' not found in DataFrame.")

if treatment_column is not None and treatment_column not in data.columns:

raise ValueError(f"Treatment column '{treatment_column}' not found in DataFrame.")

if covariates:

for cov in covariates:

if cov not in data.columns:

raise ValueError(f"Covariate column '{cov}' not found in DataFrame.")

if not 0 <= outcome_weight <= 1:

raise ValueError("outcome_weight must be between 0 and 1")

if not 0 <= covariate_weight <= 1:

raise ValueError("covariate_weight must be between 0 and 1")

if treated_units is not None and treatment_column is not None:

raise ValueError("Specify either 'treated_units' or 'treatment_column', not both.")

if require_units and exclude_units:

invalid_required = [u for u in require_units if u in exclude_units]

if invalid_required:

raise ValueError(f"Units cannot be both required and excluded: {invalid_required}")

# -------------------------------------------------------------------------

# Determine pre-treatment periods

# -------------------------------------------------------------------------

all_periods = sorted(data[time_column].unique())

if pre_periods is None:

mid_point = len(all_periods) // 2

pre_periods = all_periods[:mid_point]

else:

pre_periods = list(pre_periods)

if len(pre_periods) == 0:

raise ValueError("No pre-treatment periods specified or inferred.")

# -------------------------------------------------------------------------

# Identify treated and control units

# -------------------------------------------------------------------------

all_units = list(data[unit_column].unique())

if treated_units is not None:

treated_set = set(treated_units)

elif treatment_column is not None:

unit_treatment = data.groupby(unit_column)[treatment_column].first()

treated_set = set(unit_treatment[unit_treatment == 1].index)

elif suggest_treatment_candidates:

# Treatment candidate discovery mode - no treated units

treated_set = set()

else:

raise ValueError(

"Must specify treated_units, treatment_column, or set "

"suggest_treatment_candidates=True"

)

# -------------------------------------------------------------------------

# Treatment candidate discovery mode

# -------------------------------------------------------------------------

if suggest_treatment_candidates and len(treated_set) == 0:

return _suggest_treatment_candidates(

data, unit_column, time_column, outcome_column, pre_periods, n_treatment_candidates

)

if len(treated_set) == 0:

raise ValueError("No treated units found.")

# Determine control candidates

control_candidates = [u for u in all_units if u not in treated_set]

if exclude_units:

control_candidates = [u for u in control_candidates if u not in exclude_units]

if len(control_candidates) == 0:

raise ValueError("No control units available after exclusions.")

# -------------------------------------------------------------------------

# Create outcome matrices (pre-treatment)

# -------------------------------------------------------------------------

pre_data = data[data[time_column].isin(pre_periods)]

pivot = pre_data.pivot(index=time_column, columns=unit_column, values=outcome_column)

# Filter to pre_periods that exist in data

valid_pre_periods = [p for p in pre_periods if p in pivot.index]

if len(valid_pre_periods) == 0:

raise ValueError("No data found for specified pre-treatment periods.")

# Filter control_candidates to those present in pivot (handles unbalanced panels)

control_candidates = [c for c in control_candidates if c in pivot.columns]

if len(control_candidates) == 0:

raise ValueError("No control units found in pre-treatment data.")

# Control outcomes: shape (n_pre_periods, n_control_candidates)

Y_control = pivot.loc[valid_pre_periods, control_candidates].values.astype(float)

# Treated outcomes mean: shape (n_pre_periods,)

treated_list = [u for u in treated_set if u in pivot.columns]

if len(treated_list) == 0:

raise ValueError("Treated units not found in pre-treatment data.")

Y_treated_mean = pivot.loc[valid_pre_periods, treated_list].mean(axis=1).values.astype(float)

# -------------------------------------------------------------------------

# Compute outcome trend scores

# -------------------------------------------------------------------------

# Synthetic weights (higher = better match)

synthetic_weights = compute_synthetic_weights(Y_control, Y_treated_mean, lambda_reg=lambda_reg)

# RMSE for each control vs treated mean (use nanmean to handle missing data)

rmse_scores = []

for j in range(len(control_candidates)):

y_c = Y_control[:, j]

rmse = np.sqrt(np.nanmean((y_c - Y_treated_mean) ** 2))

rmse_scores.append(rmse)

# Convert RMSE to similarity score (lower RMSE = higher score)

max_rmse = max(rmse_scores) if rmse_scores else 1.0

min_rmse = min(rmse_scores) if rmse_scores else 0.0

rmse_range = max_rmse - min_rmse

if rmse_range < 1e-10:

# All controls have identical/similar pre-trends (includes single control case)

outcome_trend_scores = [1.0] * len(rmse_scores)

else:

# Normalize so best control gets 1.0, worst gets 0.0