Spaces:
Running
Running
| # Utility functions for Difference-in-Differences | |
| import pandas as pd | |
| import logging | |
| logger = logging.getLogger(__name__) | |
| def create_post_indicator(df: pd.DataFrame, time_var: str, treatment_period_start: any) -> pd.Series: | |
| """Creates the post-treatment indicator variable. | |
| Checks if time_var is already a 0/1 indicator; otherwise, compares to treatment_period_start. | |
| """ | |
| try: | |
| time_var_series = df[time_var] | |
| # Ensure numeric for checks and direct comparison | |
| if pd.api.types.is_bool_dtype(time_var_series): | |
| time_var_series = time_var_series.astype(int) | |
| # Check if it's already a binary 0/1 indicator | |
| if pd.api.types.is_numeric_dtype(time_var_series): | |
| unique_vals = set(time_var_series.dropna().unique()) | |
| if unique_vals == {0, 1}: | |
| logger.info(f"Time variable '{time_var}' is already a binary 0/1 indicator. Using it directly as post indicator.") | |
| return time_var_series.astype(int) | |
| else: | |
| # Numeric, but not 0/1, so compare with treatment_period_start | |
| logger.info(f"Time variable '{time_var}' is numeric. Comparing with treatment_period_start: {treatment_period_start}") | |
| return (time_var_series >= treatment_period_start).astype(int) | |
| else: | |
| # Non-numeric and not boolean, will likely fall into TypeError for datetime conversion | |
| # This else block might not be strictly necessary if TypeError is caught below | |
| # but added for logical completeness before attempting datetime conversion. | |
| pass # Let it fall through to TypeError if not numeric here | |
| # If we reached here, it means it wasn't numeric or bool, try direct comparison which will likely raise TypeError | |
| # and be caught by the except block for datetime conversion if applicable. | |
| # This line is kept to ensure non-numeric non-datetime-like strings also trigger the except. | |
| return (df[time_var] >= treatment_period_start).astype(int) | |
| except TypeError: | |
| # If direct comparison fails (e.g., comparing datetime with int/str, or non-numeric string with number), | |
| # attempt to convert both to datetime objects for comparison. | |
| logger.info(f"Direct comparison/numeric check failed for time_var '{time_var}'. Attempting datetime conversion.") | |
| try: | |
| time_series_dt = pd.to_datetime(df[time_var], errors='coerce') | |
| # Try to convert treatment_period_start to datetime if it's not already | |
| # This handles cases where treatment_period_start might be a date string | |
| try: | |
| treatment_start_dt = pd.to_datetime(treatment_period_start) | |
| except Exception as e_conv: | |
| logger.error(f"Could not convert treatment_period_start '{treatment_period_start}' to datetime: {e_conv}") | |
| raise TypeError(f"treatment_period_start '{treatment_period_start}' could not be converted to a comparable datetime format.") | |
| if time_series_dt.isna().all(): # if all values are NaT after conversion | |
| raise ValueError(f"Time variable '{time_var}' could not be converted to datetime (all values NaT).") | |
| if pd.isna(treatment_start_dt): | |
| raise ValueError(f"Treatment start period '{treatment_period_start}' converted to NaT.") | |
| logger.info(f"Comparing time_var '{time_var}' (as datetime) with treatment_start_dt '{treatment_start_dt}' (as datetime).") | |
| return (time_series_dt >= treatment_start_dt).astype(int) | |
| except Exception as e: | |
| logger.error(f"Failed to compare time variable '{time_var}' with treatment start '{treatment_period_start}' using datetime logic: {e}", exc_info=True) | |
| raise TypeError(f"Could not compare time variable '{time_var}' with treatment start '{treatment_period_start}'. Ensure they are comparable or convertible to datetime. Error: {e}") | |
| except Exception as ex: | |
| # Catch any other unexpected errors during the initial numeric processing | |
| logger.error(f"Unexpected error processing time_var '{time_var}' for post indicator: {ex}", exc_info=True) | |
| raise TypeError(f"Unexpected error processing time_var '{time_var}': {ex}") |