""" Utility functions for causal inference methods. This module provides common utility functions used across different causal inference methods. """ from typing import Dict, List, Set, Optional, Union, Any, Tuple import numpy as np import pandas as pd import scipy.stats as stats from sklearn.preprocessing import StandardScaler import matplotlib.pyplot as plt import seaborn as sns from statsmodels.stats.outliers_influence import variance_inflation_factor from sklearn.linear_model import LogisticRegression import logging # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) def check_binary_treatment(treatment_series: pd.Series) -> bool: """ Check if treatment variable is binary. Args: treatment_series: Series containing treatment variable Returns: Boolean indicating if treatment is binary """ unique_values = set(treatment_series.unique()) # Remove NaN values if present unique_values = {x for x in unique_values if pd.notna(x)} # Check if there are exactly 2 unique values if len(unique_values) != 2: return False # Check if values are 0/1 or similar binary encoding sorted_vals = sorted(unique_values) # Check common binary encodings: 0/1, False/True, etc. binary_pairs = [ (0, 1), (False, True), ("0", "1"), ("no", "yes"), ("false", "true") ] # Convert to strings for comparison if needed if not all(isinstance(v, (int, float, bool)) for v in sorted_vals): # Convert to lowercase strings for comparison str_vals = [str(v).lower() for v in sorted_vals] for pair in binary_pairs: str_pair = [str(v).lower() for v in pair] if str_vals == str_pair: return True return False # For numeric values, check if they're 0/1 or can be easily mapped to 0/1 if sorted_vals == [0, 1]: return True # Check if there are only two values that could be easily mapped return len(unique_values) == 2 def calculate_standardized_differences(df: pd.DataFrame, treatment: str, covariates: List[str]) -> Dict[str, float]: """ Calculate standardized differences between treated and control groups. Args: df: DataFrame containing the data treatment: Name of treatment variable covariates: List of covariate variable names Returns: Dictionary with standardized differences for each covariate """ treated = df[df[treatment] == 1] control = df[df[treatment] == 0] std_diffs = {} for cov in covariates: # Skip if covariate has missing values if df[cov].isna().any(): std_diffs[cov] = np.nan continue t_mean = treated[cov].mean() c_mean = control[cov].mean() t_var = treated[cov].var() c_var = control[cov].var() # Pooled standard deviation pooled_std = np.sqrt((t_var + c_var) / 2) # Avoid division by zero if pooled_std == 0: std_diffs[cov] = 0 else: std_diffs[cov] = (t_mean - c_mean) / pooled_std return std_diffs def check_overlap(df: pd.DataFrame, treatment: str, propensity_scores: np.ndarray, threshold: float = 0.5) -> Dict[str, Any]: """ Check overlap in propensity scores between treated and control groups. Args: df: DataFrame containing the data treatment: Name of treatment variable propensity_scores: Array of propensity scores threshold: Threshold for sufficient overlap (proportion of range) Returns: Dictionary with overlap statistics """ df_copy = df.copy() df_copy['propensity_score'] = propensity_scores treated = df_copy[df_copy[treatment] == 1]['propensity_score'] control = df_copy[df_copy[treatment] == 0]['propensity_score'] min_treated = treated.min() max_treated = treated.max() min_control = control.min() max_control = control.max() overall_min = min(min_treated, min_control) overall_max = max(max_treated, max_control) # Range of overlap overlap_min = max(min_treated, min_control) overlap_max = min(max_treated, max_control) # Check if there is any overlap if overlap_max < overlap_min: overlap_proportion = 0 sufficient_overlap = False else: # Calculate proportion of overall range that has overlap overall_range = overall_max - overall_min if overall_range == 0: # All values are the same overlap_proportion = 1.0 sufficient_overlap = True else: overlap_proportion = (overlap_max - overlap_min) / overall_range sufficient_overlap = overlap_proportion >= threshold return { "treated_range": (float(min_treated), float(max_treated)), "control_range": (float(min_control), float(max_control)), "overlap_range": (float(overlap_min), float(overlap_max)), "overlap_proportion": float(overlap_proportion), "sufficient_overlap": sufficient_overlap } def plot_propensity_overlap(df: pd.DataFrame, treatment: str, propensity_scores: np.ndarray, save_path: Optional[str] = None) -> None: """ Plot overlap in propensity scores. Args: df: DataFrame containing the data treatment: Name of treatment variable propensity_scores: Array of propensity scores save_path: Optional path to save the plot """ df_copy = df.copy() df_copy['propensity_score'] = propensity_scores plt.figure(figsize=(10, 6)) # Plot histograms sns.histplot(df_copy.loc[df_copy[treatment] == 1, 'propensity_score'], bins=20, alpha=0.5, label='Treated', color='blue', kde=True) sns.histplot(df_copy.loc[df_copy[treatment] == 0, 'propensity_score'], bins=20, alpha=0.5, label='Control', color='red', kde=True) plt.title('Propensity Score Distributions') plt.xlabel('Propensity Score') plt.ylabel('Count') plt.legend() if save_path: plt.savefig(save_path, dpi=300, bbox_inches='tight') plt.show() def plot_covariate_balance(standardized_diffs: Dict[str, float], threshold: float = 0.1, save_path: Optional[str] = None) -> None: """ Plot standardized differences for covariates before and after matching. Args: standardized_diffs: Dictionary with standardized differences threshold: Threshold for acceptable balance save_path: Optional path to save the plot """ # Convert to DataFrame for plotting df = pd.DataFrame({ 'Covariate': list(standardized_diffs.keys()), 'Standardized Difference': list(standardized_diffs.values()) }) # Sort by absolute standardized difference df['Absolute Difference'] = np.abs(df['Standardized Difference']) df = df.sort_values('Absolute Difference', ascending=False) plt.figure(figsize=(12, len(standardized_diffs) * 0.4 + 2)) # Plot horizontal bars ax = sns.barplot(x='Standardized Difference', y='Covariate', data=df, palette=['red' if abs(x) > threshold else 'green' for x in df['Standardized Difference']]) # Add vertical lines for thresholds plt.axvline(x=threshold, color='red', linestyle='--', alpha=0.7) plt.axvline(x=-threshold, color='red', linestyle='--', alpha=0.7) plt.axvline(x=0, color='black', linestyle='-', alpha=0.7) plt.title('Covariate Balance: Standardized Differences') plt.xlabel('Standardized Difference') plt.tight_layout() if save_path: plt.savefig(save_path, dpi=300, bbox_inches='tight') plt.show() def check_temporal_structure(df: pd.DataFrame) -> Dict[str, Any]: """ Check if dataset has temporal structure. Args: df: DataFrame to check Returns: Dictionary with temporal structure information """ # Check for date/time columns date_cols = [] for col in df.columns: # Check if column has date in name if any(date_term in col.lower() for date_term in ['date', 'time', 'year', 'month', 'day', 'period']): date_cols.append(col) # Check if column can be converted to datetime if df[col].dtype == 'object': try: pd.to_datetime(df[col], errors='raise') date_cols.append(col) except: pass # Check for panel structure - look for ID columns id_cols = [] for col in df.columns: # Check if column has ID in name if any(id_term in col.lower() for id_term in ['id', 'identifier', 'key', 'code']): unique_count = df[col].nunique() # If column has multiple values but fewer than 10% of rows, likely an ID if 1 < unique_count < len(df) * 0.1: id_cols.append(col) # Check if there are multiple observations per unit is_panel = False panel_units = None if id_cols and date_cols: # For each ID column, check if there are multiple time periods for id_col in id_cols: obs_per_id = df.groupby(id_col).size() if (obs_per_id > 1).any(): is_panel = True panel_units = id_col break return { "has_temporal_structure": len(date_cols) > 0, "temporal_columns": date_cols, "potential_id_columns": id_cols, "is_panel_data": is_panel, "panel_units": panel_units } def check_for_discontinuities(df: pd.DataFrame, outcome: str, threshold_zscore: float = 3.0) -> Dict[str, Any]: """ Check for potential discontinuities in continuous variables. Args: df: DataFrame to check outcome: Name of outcome variable threshold_zscore: Z-score threshold for detecting discontinuities Returns: Dictionary with discontinuity information """ potential_running_vars = [] # Check only numeric columns that aren't the outcome numeric_cols = df.select_dtypes(include=[np.number]).columns numeric_cols = [col for col in numeric_cols if col != outcome] for col in numeric_cols: # Skip if too many unique values (unlikely to be a running variable) if df[col].nunique() > 100: continue # Sort values and calculate differences sorted_vals = np.sort(df[col].unique()) if len(sorted_vals) <= 1: continue diffs = np.diff(sorted_vals) mean_diff = np.mean(diffs) std_diff = np.std(diffs) # Skip if all differences are the same if std_diff == 0: continue # Calculate z-scores of differences zscores = (diffs - mean_diff) / std_diff # Check if any z-score exceeds threshold if np.any(np.abs(zscores) > threshold_zscore): # Potential discontinuity found max_idx = np.argmax(np.abs(zscores)) threshold = (sorted_vals[max_idx] + sorted_vals[max_idx + 1]) / 2 # Check if outcome means differ across threshold below_mean = df[df[col] < threshold][outcome].mean() above_mean = df[df[col] >= threshold][outcome].mean() # Only include if outcome means differ substantially if abs(above_mean - below_mean) > 0.1 * df[outcome].std(): potential_running_vars.append({ "variable": col, "threshold": float(threshold), "z_score": float(zscores[max_idx]), "outcome_diff": float(above_mean - below_mean) }) return { "has_discontinuities": len(potential_running_vars) > 0, "potential_running_variables": potential_running_vars } def find_potential_instruments(df: pd.DataFrame, treatment: str, outcome: str, correlation_threshold: float = 0.3) -> Dict[str, Any]: """ Find potential instrumental variables. Args: df: DataFrame to check treatment: Name of treatment variable outcome: Name of outcome variable correlation_threshold: Threshold for correlation with treatment Returns: Dictionary with potential instruments information """ # Get numeric columns that aren't treatment or outcome numeric_cols = df.select_dtypes(include=[np.number]).columns potential_ivs = [col for col in numeric_cols if col != treatment and col != outcome] iv_results = [] for col in potential_ivs: # Skip if column has too many missing values if df[col].isna().mean() > 0.1: continue # Check correlation with treatment corr_treatment = df[[col, treatment]].corr().iloc[0, 1] # Check correlation with outcome corr_outcome = df[[col, outcome]].corr().iloc[0, 1] # Potential IV should be correlated with treatment but not directly with outcome if abs(corr_treatment) > correlation_threshold and abs(corr_outcome) < correlation_threshold/2: iv_results.append({ "variable": col, "correlation_with_treatment": float(corr_treatment), "correlation_with_outcome": float(corr_outcome), "strength": "Strong" if abs(corr_treatment) > 0.5 else "Moderate" }) return { "has_potential_instruments": len(iv_results) > 0, "potential_instruments": iv_results } def test_parallel_trends(df: pd.DataFrame, treatment: str, outcome: str, time_var: str, unit_var: str) -> Dict[str, Any]: """ Test for parallel trends assumption in difference-in-differences. Args: df: DataFrame to check treatment: Name of treatment variable outcome: Name of outcome variable time_var: Name of time variable unit_var: Name of unit variable Returns: Dictionary with parallel trends test results """ # Ensure time_var is properly formatted df = df.copy() if df[time_var].dtype != 'int64': # Try to convert to datetime and then to period try: df[time_var] = pd.to_datetime(df[time_var]) # Get unique periods and map to integers periods = df[time_var].dt.to_period('M').unique() period_dict = {p: i for i, p in enumerate(sorted(periods))} df['time_period'] = df[time_var].dt.to_period('M').map(period_dict) time_var = 'time_period' except: # If conversion fails, try to map unique values to integers unique_times = df[time_var].unique() time_dict = {t: i for i, t in enumerate(sorted(unique_times))} df['time_period'] = df[time_var].map(time_dict) time_var = 'time_period' # Identify treatment and control groups # Treatment indicator should be 0 or 1 for each unit (not time-varying) unit_treatment = df.groupby(unit_var)[treatment].max() treatment_units = unit_treatment[unit_treatment == 1].index control_units = unit_treatment[unit_treatment == 0].index # Find time of treatment implementation if len(treatment_units) > 0: treatment_time = df[df[unit_var].isin(treatment_units) & (df[treatment] == 1)][time_var].min() else: # No treated units found return { "parallel_trends": False, "reason": "No treated units found", "pre_trend_correlation": None, "pre_trend_p_value": None } # Select pre-treatment periods pre_treatment = df[df[time_var] < treatment_time] # Calculate average outcome by time and group treated_means = pre_treatment[pre_treatment[unit_var].isin(treatment_units)].groupby(time_var)[outcome].mean() control_means = pre_treatment[pre_treatment[unit_var].isin(control_units)].groupby(time_var)[outcome].mean() # Need enough pre-treatment periods to test if len(treated_means) < 3: return { "parallel_trends": None, "reason": "Insufficient pre-treatment periods", "pre_trend_correlation": None, "pre_trend_p_value": None } # Align indices and calculate trends common_periods = sorted(set(treated_means.index).intersection(set(control_means.index))) if len(common_periods) < 3: return { "parallel_trends": None, "reason": "Insufficient common pre-treatment periods", "pre_trend_correlation": None, "pre_trend_p_value": None } treated_trends = np.diff(treated_means[common_periods]) control_trends = np.diff(control_means[common_periods]) # Calculate correlation between trends correlation, p_value = stats.pearsonr(treated_trends, control_trends) # Test if trends are parallel (high correlation, not significantly different) parallel_trends = correlation > 0.7 and p_value < 0.05 return { "parallel_trends": parallel_trends, "reason": "Trends are parallel" if parallel_trends else "Trends are not parallel", "pre_trend_correlation": float(correlation), "pre_trend_p_value": float(p_value) } def preprocess_data(df: pd.DataFrame, treatment_var: str, outcome_var: str, covariates: List[str], verbose: bool = True) -> pd.DataFrame: """ Preprocess the dataset to handle missing values and encode categorical variables. Args: df (pd.DataFrame): The dataset treatment_var (str): The treatment variable name outcome_var (str): The outcome variable name covariates (list): List of covariate variable names verbose (bool): Whether to print verbose output Returns: Tuple[pd.DataFrame, str, str, List[str], Dict[str, Any]]: Preprocessed dataset, updated treatment var name, updated outcome var name, updated covariates list, and column mappings. """ df_processed = df.copy() column_mappings: Dict[str, Any] = {} # Store original dtypes for mapping original_dtypes = {col: str(df_processed[col].dtype) for col in df_processed.columns} # Report missing values all_vars = [treatment_var, outcome_var] + covariates missing_data = df_processed[all_vars].isnull().sum() total_missing = missing_data.sum() if total_missing > 0: if verbose: logger.info(f"Dataset contains {total_missing} missing values:") for col in missing_data[missing_data > 0].index: percent = (missing_data[col] / len(df_processed)) * 100 if verbose: logger.info(f" - {col}: {missing_data[col]} missing values ({percent:.2f}%)") else: if verbose: logger.info("No missing values found in relevant columns.") # return df_processed # No preprocessing needed if no missing values # Handle missing values in treatment variable if df_processed[treatment_var].isnull().sum() > 0: if verbose: logger.info(f"Filling missing values in treatment variable '{treatment_var}' with mode") # For treatment, use mode (most common value) mode_val = df_processed[treatment_var].mode()[0] if not df_processed[treatment_var].mode().empty else 0 df_processed[treatment_var] = df_processed[treatment_var].fillna(mode_val) # Handle missing values in outcome variable if df_processed[outcome_var].isnull().sum() > 0: if verbose: logger.info(f"Filling missing values in outcome variable '{outcome_var}' with mean") # For outcome, use mean mean_val = df_processed[outcome_var].mean() df_processed[outcome_var] = df_processed[outcome_var].fillna(mean_val) # Handle missing values in covariates for col in covariates: if df_processed[col].isnull().sum() > 0: if pd.api.types.is_numeric_dtype(df_processed[col]): # For numeric covariates, use mean if verbose: logger.info(f"Filling missing values in numeric covariate '{col}' with mean") mean_val = df_processed[col].mean() df_processed[col] = df_processed[col].fillna(mean_val) elif pd.api.types.is_categorical_dtype(df_processed[col]) or df_processed[col].dtype == 'object': # For categorical covariates, use mode mode_val = df_processed[col].mode()[0] if not df_processed[col].mode().empty else "Missing" if verbose: logger.info(f"Filling missing values in categorical covariate '{col}' with mode ('{mode_val}')") df_processed[col] = df_processed[col].fillna(mode_val) else: # For other types, create a "Missing" category if verbose: logger.info(f"Filling missing values in covariate '{col}' of type {df_processed[col].dtype} with 'Missing' category") # Ensure the column is of object type before filling with string if df_processed[col].dtype != 'object': try: df_processed[col] = df_processed[col].astype(object) except Exception as e: logger.warning(f"Could not convert column {col} to object type to fill NAs: {e}. Skipping fill.") continue df_processed[col] = df_processed[col].fillna("Missing") # --- Categorical Encoding --- updated_treatment_var = treatment_var updated_outcome_var = outcome_var # Helper function for label encoding binary categoricals def label_encode_binary(series: pd.Series, var_name: str) -> Tuple[pd.Series, Dict[int, Any]]: uniques = series.dropna().unique() mapping = {} if len(uniques) == 2: # Try to map to 0 and 1 consistently, e.g., sort and assign # Or if boolean, map True to 1, False to 0 if series.dtype == 'bool': mapping = {0: False, 1: True} return series.astype(int), mapping # For non-boolean, sort to ensure consistent mapping # However, direct replacement is safer to control which becomes 0 and 1 # For simplicity here, we'll make a simple map. # A more robust approach might involve explicit mapping rules or user input. sorted_uniques = sorted(uniques, key=lambda x: str(x)) # sort to make it deterministic map_dict = {sorted_uniques[0]: 0, sorted_uniques[1]: 1} mapping = {v: k for k, v in map_dict.items()} # Inverse map for column_mappings if verbose: logger.info(f"Label encoding binary variable '{var_name}': {map_dict}") return series.map(map_dict), mapping elif len(uniques) == 1: # Single unique value, treat as constant (encode as 0) if verbose: logger.info(f"Binary variable '{var_name}' has only one unique value '{uniques[0]}'. Encoding as 0.") map_dict = {uniques[0]:0} mapping = {0: uniques[0]} return series.map(map_dict), mapping return series, mapping # No change if not binary # Encode Treatment Variable if df_processed[treatment_var].dtype == 'object' or df_processed[treatment_var].dtype == 'category' or df_processed[treatment_var].dtype == 'bool': original_series = df_processed[treatment_var].copy() df_processed[treatment_var], value_map = label_encode_binary(df_processed[treatment_var], treatment_var) if value_map: # If encoding happened column_mappings[treatment_var] = { 'original_dtype': original_dtypes[treatment_var], 'transformed_as': 'label_encoded_binary', 'new_column_name': treatment_var, # Name doesn't change 'value_map': value_map } if verbose: logger.info(f"Encoded treatment variable '{treatment_var}' to numeric.") # Encode Outcome Variable if df_processed[outcome_var].dtype == 'object' or df_processed[outcome_var].dtype == 'category' or df_processed[outcome_var].dtype == 'bool': original_series = df_processed[outcome_var].copy() df_processed[outcome_var], value_map = label_encode_binary(df_processed[outcome_var], outcome_var) if value_map: # If encoding happened column_mappings[outcome_var] = { 'original_dtype': original_dtypes[outcome_var], 'transformed_as': 'label_encoded_binary', 'new_column_name': outcome_var, # Name doesn't change 'value_map': value_map } if verbose: logger.info(f"Encoded outcome variable '{outcome_var}' to numeric.") # Encode Covariates (One-Hot Encoding for non-numeric) updated_covariates = [] categorical_covariates_to_encode = [] for cov in covariates: if cov not in df_processed.columns: # If a covariate was dropped or is an instrument etc. if verbose: logger.warning(f"Covariate '{cov}' not found in DataFrame columns after initial processing. Skipping encoding for it.") continue if df_processed[cov].dtype == 'object' or df_processed[cov].dtype == 'category' or pd.api.types.is_bool_dtype(df_processed[cov]): # Check if it's binary - if so, can also label encode # However, for consistency with get_dummies and to handle multi-category, # we'll let get_dummies handle it, or apply label encoding for binary covariates too. # For simplicity, let's stick to one-hot for all categorical covariates. if len(df_processed[cov].dropna().unique()) > 1 : # Only encode if more than 1 unique value categorical_covariates_to_encode.append(cov) else: # If only one unique value or all NaNs (already handled), it's constant-like if verbose: logger.info(f"Categorical covariate '{cov}' has <= 1 unique value after NA handling. Treating as constant-like, not one-hot encoding.") updated_covariates.append(cov) # Keep as is, will likely be numeric 0 or some constant else: # Already numeric updated_covariates.append(cov) if categorical_covariates_to_encode: if verbose: logger.info(f"One-hot encoding categorical covariates: {categorical_covariates_to_encode} using pd.get_dummies (drop_first=True)") # Store original columns before get_dummies to identify new ones original_df_columns = set(df_processed.columns) df_processed = pd.get_dummies(df_processed, columns=categorical_covariates_to_encode, prefix_sep='_', drop_first=True, dummy_na=False) # dummy_na=False since we handled NAs # Identify new columns created by get_dummies new_dummy_columns = list(set(df_processed.columns) - original_df_columns) updated_covariates.extend(new_dummy_columns) for original_cov_name in categorical_covariates_to_encode: # Find which dummy columns correspond to this original covariate related_dummies = [col for col in new_dummy_columns if col.startswith(original_cov_name + '_')] column_mappings[original_cov_name] = { 'original_dtype': original_dtypes[original_cov_name], 'transformed_as': 'one_hot_encoded', 'encoded_columns': related_dummies, # 'dropped_category': can be inferred if needed, but not explicitly stored for simplicity here } if verbose: logger.info(f" Original covariate '{original_cov_name}' resulted in dummy variables: {related_dummies}") if verbose: logger.info("Preprocessing complete.") if column_mappings: logger.info(f"Column mappings generated: {column_mappings}") else: logger.info("No column encodings were applied.") return df_processed, updated_treatment_var, updated_outcome_var, list(dict.fromkeys(updated_covariates)), column_mappings def check_collinearity(df: pd.DataFrame, covariates: List[str]) -> Optional[List[str]]: # Implementation of check_collinearity function # This function should return a list of collinear variables or None pass