[数据分析]双重差分法：政策效果评估的利器

import pandas as pd
import numpy as np
import statsmodels.api as sm
import statsmodels.formula.api as smf
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import stats

# 设置随机种子确保结果可重现
np.random.seed(1234)

# 生成模拟数据
def generate_did_data(n=2000, pre_periods=3, post_periods=2):
    """
    生成DID分析所需的模拟数据
    
    参数:
    n: 总样本量
    pre_periods: 政策前期数
    post_periods: 政策后期数
    """
    # 创建基本数据结构
    individuals = range(n)
    periods = range(pre_periods + post_periods)
    
    # 生成个体特征
    data = []
    for i in individuals:
        # 随机分配处理组（40%的处理组）
        treated = np.random.binomial(1, 0.4)
        
        # 生成个体固定特征
        ability = np.random.normal(0, 1)  # 个人能力
        urban = np.random.binomial(1, 0.6)  # 是否城市居民
        age = np.random.randint(18, 60)  # 年龄
        
        for t in periods:
            # 政策实施时间点（第3期后）
            post = 1 if t >= pre_periods else 0
            
            # 基础收入生成过程（两组共享的时间趋势）
            time_trend = 0.1 * t  # 随时间增长的趋势
            base_income = (50 + 2*ability + 5*urban + 0.1*age + 
                          time_trend + np.random.normal(0, 2))
            
            # 处理效应（仅在处理组且政策后）
            if treated == 1 and post == 1:
                # 政策效应：提高收入，但效应因人而异
                treatment_effect = 8 + 1.5*ability + np.random.normal(0, 1)
            else:
                treatment_effect = 0
                
            # 最终收入
            income = base_income + treatment_effect
            
            # 就业状态（受收入和能力影响）
            employment_prob = 1 / (1 + np.exp(-(0.1*income + 0.5*ability - 5)))
            employed = np.random.binomial(1, employment_prob)
            
            data.append({
                'id': i,
                'period': t,
                'year': 2017 + t,  # 假设从2017年开始
                'treated': treated,
                'post': post,
                'ability': ability,
                'urban': urban,
                'age': age,
                'income': max(income, 0),  # 确保非负
                'employed': employed,
                'time_trend': time_trend
            })
    
    return pd.DataFrame(data)

# 生成数据
df = generate_did_data(n=2000, pre_periods=3, post_periods=2)
print(f"生成的数据集形状: {df.shape}")
print(df.head())

# 政策前数据子集
pre_policy = df[df['post'] == 0]

# 平衡性检验表格
balance_table = pre_policy.groupby('treated').agg({
    'ability': ['mean', 'std'],
    'urban': ['mean', 'std'],
    'age': ['mean', 'std'],
    'income': ['mean', 'std'],
    'employed': ['mean', 'std']
}).round(3)

print("政策前处理组与对照组特征比较:")
print(balance_table)

# 平衡性检验的统计检验
print("\n平衡性检验的t检验结果:")
variables = ['ability', 'urban', 'age', 'income', 'employed']
for var in variables:
    treated_var = pre_policy[pre_policy['treated'] == 1][var]
    control_var = pre_policy[pre_policy['treated'] == 0][var]
    t_stat, p_value = stats.ttest_ind(treated_var, control_var)
    print(f"{var}: t={t_stat:.3f}, p={p_value:.3f}")

# 设置图形风格
plt.style.use('seaborn-v0_8')
fig, axes = plt.subplots(2, 2, figsize=(15, 12))

# 收入随时间变化的趋势
treated_pre = df[(df['treated'] == 1) & (df['post'] == 0)]['income'].mean()
treated_post = df[(df['treated'] == 1) & (df['post'] == 1)]['income'].mean()
control_pre = df[(df['treated'] == 0) & (df['post'] == 0)]['income'].mean()
control_post = df[(df['treated'] == 0) & (df['post'] == 1)]['income'].mean()

periods = ['政策前', '政策后']
treated_means = [treated_pre, treated_post]
control_means = [control_pre, control_post]

axes[0, 0].plot(periods, treated_means, 'o-', label='处理组', linewidth=2, markersize=8)
axes[0, 0].plot(periods, control_means, 's-', label='对照组', linewidth=2, markersize=8)
axes[0, 0].set_title('收入随时间变化趋势', fontsize=14, fontweight='bold')
axes[0, 0].set_ylabel('平均收入', fontsize=12)
axes[0, 0].legend()
axes[0, 0].grid(True, alpha=0.3)

# 就业率随时间变化
employed_treated_pre = df[(df['treated'] == 1) & (df['post'] == 0)]['employed'].mean()
employed_treated_post = df[(df['treated'] == 1) & (df['post'] == 1)]['employed'].mean()
employed_control_pre = df[(df['treated'] == 0) & (df['post'] == 0)]['employed'].mean()
employed_control_post = df[(df['treated'] == 0) & (df['post'] == 1)]['employed'].mean()

employed_treated = [employed_treated_pre, employed_treated_post]
employed_control = [employed_control_pre, employed_control_post]

axes[0, 1].plot(periods, employed_treated, 'o-', label='处理组', linewidth=2, markersize=8)
axes[0, 1].plot(periods, employed_control, 's-', label='对照组', linewidth=2, markersize=8)
axes[0, 1].set_title('就业率随时间变化趋势', fontsize=14, fontweight='bold')
axes[0, 1].set_ylabel('就业率', fontsize=12)
axes[0, 1].legend()
axes[0, 1].grid(True, alpha=0.3)

# 处理组和对照组特征分布
variables = ['ability', 'age']
for i, var in enumerate(variables):
    axes[1, i].hist([df[df['treated'] == 1][var], df[df['treated'] == 0][var]], 
                   bins=20, alpha=0.7, label=['处理组', '对照组'])
    axes[1, i].set_title(f'{var}分布', fontsize=14, fontweight='bold')
    axes[1, i].legend()
    axes[1, i].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# 基础DID模型 - 收入效应
def basic_did_income(df):
    """估计政策对收入的影响"""
    # 方法1: 使用回归公式
    model_income = smf.ols('income ~ treated + post + treated:post', data=df).fit()
    
    # 方法2: 手动计算（验证结果）
    treated_pre = df[(df['treated'] == 1) & (df['post'] == 0)]['income'].mean()
    treated_post = df[(df['treated'] == 1) & (df['post'] == 1)]['income'].mean()
    control_pre = df[(df['treated'] == 0) & (df['post'] == 0)]['income'].mean()
    control_post = df[(df['treated'] == 0) & (df['post'] == 1)]['income'].mean()
    
    did_manual = (treated_post - treated_pre) - (control_post - control_pre)
    
    return model_income, did_manual

# 执行DID估计
model_income, did_manual = basic_did_income(df)

print("="*60)
print("基础DID模型结果 - 收入效应")
print("="*60)
print(model_income.summary())
print(f"\n手动计算的DID估计量: {did_manual:.4f}")

# 就业效应的DID模型（使用线性概率模型）
def basic_did_employment(df):
    """估计政策对就业的影响（线性概率模型）"""
    model_employment = smf.ols('employed ~ treated + post + treated:post', data=df).fit()
    
    # 手动计算
    treated_pre = df[(df['treated'] == 1) & (df['post'] == 0)]['employed'].mean()
    treated_post = df[(df['treated'] == 1) & (df['post'] == 1)]['employed'].mean()
    control_pre = df[(df['treated'] == 0) & (df['post'] == 0)]['employed'].mean()
    control_post = df[(df['treated'] == 0) & (df['post'] == 1)]['employed'].mean()
    
    did_manual_emp = (treated_post - treated_pre) - (control_post - control_pre)
    
    return model_employment, did_manual_emp

model_employment, did_manual_emp = basic_did_employment(df)

print("\n" + "="*60)
print("基础DID模型结果 - 就业效应（线性概率模型）")
print("="*60)
print(model_employment.summary())
print(f"\n手动计算的DID估计量: {did_manual_emp:.4f}")

# 双向固定效应模型
def twoway_fe_model(df):
    """双向固定效应模型"""
    # 为每个个体和时间创建虚拟变量
    df['id_factor'] = pd.Categorical(df['id'])
    df['period_factor'] = pd.Categorical(df['period'])
    
    # 使用固定效应模型
    fe_formula = 'income ~ treated:post + C(id_factor) + C(period_factor) - 1'
    fe_model = smf.ols(fe_formula, data=df).fit()
    
    return fe_model

# 执行双向固定效应估计
fe_model = twoway_fe_model(df)

print("="*60)
print("双向固定效应模型结果")
print("="*60)
# 提取我们关心的系数（交互项系数）
interaction_coef = fe_model.params[fe_model.params.index.str.contains('treated:post')]
print(f"政策效应估计: {interaction_coef.values[0]:.4f}")

# 简化版的固定效应模型（使用groupby去均值）
def simple_fe_did(df):
    """简化的固定效应DID"""
    # 计算个体均值和时间均值
    individual_means = df.groupby('id')['income'].mean()
    time_means = df.groupby('period')['income'].mean()
    overall_mean = df['income'].mean()
    
    # 对数据进行去均值处理
    df_fe = df.copy()
    df_fe['income_dm'] = (df_fe['income'] - 
                         df_fe['id'].map(individual_means) - 
                         df_fe['period'].map(time_means) + 
                         overall_mean)
    
    # 估计模型
    fe_simple = smf.ols('income_dm ~ treated:post', data=df_fe).fit()
    
    return fe_simple

fe_simple = simple_fe_did(df)
print("\n简化固定效应模型结果:")
print(fe_simple.summary())

# 平行趋势检验
def parallel_trend_test(df, outcome_var='income'):
    """平行趋势检验"""
    # 为每个时期创建虚拟变量
    period_dummies = pd.get_dummies(df['period'], prefix='period')
    df_with_dummies = pd.concat([df, period_dummies], axis=1)
    
    # 将政策前第一期作为基准期
    base_period = df['period'].min()
    period_cols = [col for col in df_with_dummies.columns if 'period_' in col]
    period_cols.remove(f'period_{base_period}')  # 删除基准期
    
    # 创建交互项
    interaction_terms = []
    for col in period_cols:
        interac_col = f'treated_{col}'
        df_with_dummies[interac_col] = df_with_dummies['treated'] * df_with_dummies[col]
        interaction_terms.append(interac_col)
    
    # 构建回归模型（排除政策后时期）
    df_pre = df_with_dummies[df_with_dummies['post'] == 0].copy()
    
    # 回归公式
    formula = f"{outcome_var} ~ treated + " + " + ".join(period_cols) + " + " + " + ".join(interaction_terms)
    pt_model = smf.ols(formula, data=df_pre).fit()
    
    return pt_model, interaction_terms

# 执行平行趋势检验
pt_model, interaction_terms = parallel_trend_test(df, 'income')

print("="*60)
print("平行趋势检验结果 - 收入")
print("="*60)
print(pt_model.summary())

# 提取交互项系数（政策前各期）
pre_periods = [term for term in interaction_terms if int(term.split('_')[-1]) < 3]  # 政策前期数
pt_coefficients = pt_model.params[pre_periods]
pt_pvalues = pt_model.pvalues[pre_periods]

print("\n政策前各期处理组与对照组差异:")
for period, coef, pval in zip(pre_periods, pt_coefficients, pt_pvalues):
    period_num = period.split('_')[-1]
    print(f"时期 {period_num}: 系数 = {coef:.4f}, p值 = {pval:.4f}")

# 可视化平行趋势
def plot_parallel_trends(df, outcome_var='income'):
    """绘制平行趋势图"""
    # 计算各期均值
    period_means = df.groupby(['period', 'treated'])[outcome_var].mean().reset_index()
    
    # 分离处理组和对照组
    treated_means = period_means[period_means['treated'] == 1]
    control_means = period_means[period_means['treated'] == 0]
    
    # 创建图形
    plt.figure(figsize=(10, 6))
    periods = sorted(df['period'].unique())
    
    # 政策实施时间线
    policy_period = 3  # 第3期后实施政策
    plt.axvline(x=policy_period - 0.5, color='red', linestyle='--', alpha=0.7, label='政策实施')
    
    # 绘制趋势线
    plt.plot(treated_means['period'], treated_means[outcome_var], 
             'o-', linewidth=2, markersize=8, label='处理组')
    plt.plot(control_means['period'], control_means[outcome_var], 
             's-', linewidth=2, markersize=8, label='对照组')
    
    plt.xlabel('时期', fontsize=12)
    plt.ylabel(f'平均{outcome_var}', fontsize=12)
    plt.title('平行趋势检验', fontsize=14, fontweight='bold')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    # 添加政策前后标注
    plt.text(policy_period/2, plt.ylim()[0] + 0.1*(plt.ylim()[1]-plt.ylim()[0]), 
             '政策前', ha='center', fontsize=11, style='italic')
    plt.text(policy_period + (max(periods)-policy_period)/2, 
             plt.ylim()[0] + 0.1*(plt.ylim()[1]-plt.ylim()[0]), 
             '政策后', ha='center', fontsize=11, style='italic')
    
    plt.tight_layout()
    plt.show()

# 绘制平行趋势图
plot_parallel_trends(df, 'income')
plot_parallel_trends(df, 'employed')

# 动态效应分析
def dynamic_effects(df, outcome_var='income'):
    """估计政策的动态效应"""
    # 创建时期虚拟变量与处理组的交互项
    period_dummies = pd.get_dummies(df['period'], prefix='period')
    df_dynamic = pd.concat([df, period_dummies], axis=1)
    
    # 生成交互项（排除基准期）
    base_period = df['period'].min()
    period_cols = [col for col in df_dynamic.columns if 'period_' in col]
    period_cols.remove(f'period_{base_period}')
    
    interaction_terms = []
    for col in period_cols:
        interac_col = f'treated_{col}'
        df_dynamic[interac_col] = df_dynamic['treated'] * df_dynamic[col]
        interaction_terms.append(interac_col)
    
    # 构建动态效应模型
    formula = f"{outcome_var} ~ treated + " + " + ".join(period_cols) + " + " + " + ".join(interaction_terms)
    dynamic_model = smf.ols(formula, data=df_dynamic).fit()
    
    return dynamic_model, interaction_terms

# 执行动态效应分析
dynamic_model, interaction_terms = dynamic_effects(df, 'income')

print("="*60)
print("动态效应分析结果 - 收入")
print("="*60)

# 提取各期政策效应
dynamic_coefs = dynamic_model.params[interaction_terms]
dynamic_ci = dynamic_model.conf_int().loc[interaction_terms]

print("各期政策效应估计:")
for period, coef in zip(interaction_terms, dynamic_coefs):
    period_num = period.split('_')[-1]
    print(f"时期 {period_num}: 效应 = {coef:.4f}")

# 可视化动态效应
def plot_dynamic_effects(dynamic_coefs, interaction_terms):
    """绘制动态效应图"""
    # 提取时期和系数
    periods = [int(term.split('_')[-1]) for term in interaction_terms]
    coefficients = [dynamic_coefs[term] for term in interaction_terms]
    
    # 计算置信区间
    conf_int = dynamic_model.conf_int().loc[interaction_terms]
    lower_bound = conf_int[0].values
    upper_bound = conf_int[1].values
    
    # 创建图形
    plt.figure(figsize=(12, 6))
    
    # 政策实施时间线
    policy_period = 3
    plt.axvline(x=policy_period - 0.5, color='red', linestyle='--', alpha=0.7, label='政策实施')
    plt.axhline(y=0, color='black', linestyle='-', alpha=0.3)
    
    # 绘制点估计和置信区间
    plt.plot(periods, coefficients, 'o-', linewidth=2, markersize=8, label='政策效应')
    plt.fill_between(periods, lower_bound, upper_bound, alpha=0.2, label='95%置信区间')
    
    plt.xlabel('时期', fontsize=12)
    plt.ylabel('政策效应', fontsize=12)
    plt.title('政策动态效应', fontsize=14, fontweight='bold')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    # 添加政策前后标注
    plt.text(1.5, plt.ylim()[0] + 0.9*(plt.ylim()[1]-plt.ylim()[0]), 
             '政策前', ha='center', fontsize=11, style='italic', color='blue')
    plt.text(4, plt.ylim()[0] + 0.9*(plt.ylim()[1]-plt.ylim()[0]), 
             '政策后', ha='center', fontsize=11, style='italic', color='blue')
    
    plt.tight_layout()
    plt.show()

# 绘制动态效应图
plot_dynamic_effects(dynamic_coefs, interaction_terms)

# 安慰剂检验
def placebo_test(df, n_simulations=500):
    """进行安慰剂检验"""
    placebo_effects = []
    
    for i in range(n_simulations):
        # 随机分配处理组（与真实处理组无关）
        df_placebo = df.copy()
        df_placebo['treated_placebo'] = np.random.binomial(1, 0.4, size=len(df))
        
        # 估计安慰剂DID
        model_placebo = smf.ols('income ~ treated_placebo + post + treated_placebo:post', 
                               data=df_placebo).fit()
        
        placebo_effect = model_placebo.params['treated_placebo:post']
        placebo_effects.append(placebo_effect)
    
    return placebo_effects

# 执行安慰剂检验
placebo_effects = placebo_test(df, n_simulations=500)

# 真实政策效应
real_effect = model_income.params['treated:post']

# 可视化安慰剂检验结果
plt.figure(figsize=(10, 6))
plt.hist(placebo_effects, bins=30, alpha=0.7, color='lightblue', edgecolor='black')
plt.axvline(x=real_effect, color='red', linestyle='--', linewidth=2, label=f'真实效应: {real_effect:.3f}')
plt.axvline(x=np.mean(placebo_effects), color='green', linestyle='--', linewidth=2, 
           label=f'安慰剂均值: {np.mean(placebo_effects):.3f}')

plt.xlabel('安慰剂效应估计值', fontsize=12)
plt.ylabel('频数', fontsize=12)
plt.title('安慰剂检验: 政策效应估计值的分布', fontsize=14, fontweight='bold')
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()

# 计算p值（真实效应在安慰剂分布中的位置）
p_value = np.mean(np.abs(placebo_effects) > np.abs(real_effect))
print(f"安慰剂检验p值: {p_value:.4f}")

# 倾向得分加权DID
def propensity_score_weighting(df):
    """倾向得分加权DID"""
    from sklearn.linear_model import LogisticRegression
    
    # 政策前数据用于估计倾向得分
    df_pre = df[df['post'] == 0].copy()
    
    # 估计倾向得分模型
    X_ps = df_pre[['ability', 'urban', 'age']]
    y_ps = df_pre['treated']
    
    ps_model = LogisticRegression()
    ps_model.fit(X_ps, y_ps)
    
    # 为所有观测值计算倾向得分
    X_all = df[['ability', 'urban', 'age']]
    df['propensity_score'] = ps_model.predict_proba(X_all)[:, 1]
    
    # 计算逆概率权重
    df['weight'] = np.where(df['treated'] == 1, 
                           1/df['propensity_score'], 
                           1/(1-df['propensity_score']))
    
    # 加权DID回归
    weighted_model = smf.wls('income ~ treated + post + treated:post', 
                            data=df, weights=df['weight']).fit()
    
    return weighted_model, df

# 执行倾向得分加权DID
weighted_model, df_weighted = propensity_score_weighting(df)

print("="*60)
print("倾向得分加权DID结果")
print("="*60)
print(weighted_model.summary())

# 加权后的平衡性检验
def check_balance_weighted(df):
    """检查加权后的平衡性"""
    from statsmodels.stats.weightstats import DescrStatsW
    
    df_pre = df[df['post'] == 0]
    
    # 加权描述性统计
    variables = ['ability', 'urban', 'age']
    balance_results = []
    
    for var in variables:
        treated_data = df_pre[df_pre['treated'] == 1][var]
        treated_weights = df_pre[df_pre['treated'] == 1]['weight']
        control_data = df_pre[df_pre['treated'] == 0][var]
        control_weights = df_pre[df_pre['treated'] == 0]['weight']
        
        # 加权均值
        treated_mean = np.average(treated_data, weights=treated_weights)
        control_mean = np.average(control_data, weights=control_weights)
        
        # 加权t检验
        weighted_stats = DescrStatsW(np.concatenate([treated_data, control_data]), 
                                   weights=np.concatenate([treated_weights, control_weights]), 
                                   group=np.concatenate([np.ones(len(treated_data)), 
                                                       np.zeros(len(control_data))]))
        t_stat, p_value, df = weighted_stats.ttest_mean()
        
        balance_results.append({
            'variable': var,
            'treated_mean': treated_mean,
            'control_mean': control_mean,
            'difference': treated_mean - control_mean,
            'p_value': p_value
        })
    
    return pd.DataFrame(balance_results)

balance_weighted = check_balance_weighted(df_weighted)
print("\n加权后平衡性检验:")
print(balance_weighted)

# 异质性分析
def heterogeneity_analysis(df):
    """异质性分析"""
    # 按能力分组分析
    df['high_ability'] = (df['ability'] > df['ability'].median()).astype(int)
    
    # 三向交互模型
    hetero_model = smf.ols('income ~ treated*post*high_ability + ability + urban + age', 
                          data=df).fit()
    
    # 按城乡分组
    urban_model = smf.ols('income ~ treated*post*urban + ability + age', 
                         data=df).fit()
    
    return hetero_model, urban_model

# 执行异质性分析
hetero_model, urban_model = heterogeneity_analysis(df)

print("="*60)
print("异质性分析结果 - 能力分组")
print("="*60)
print(hetero_model.summary())

print("\n" + "="*60)
print("异质性分析结果 - 城乡分组")
print("="*60)
print(urban_model.summary())

# 边际效应分析
def plot_marginal_effects(model, df, moderator_var='high_ability'):
    """绘制边际效应图"""
    # 创建预测数据
    prediction_data = pd.DataFrame({
        'treated': [0, 0, 1, 1],
        'post': [0, 1, 0, 1],
        moderator_var: [0, 0, 0, 0]  # 基准组
    })
    
    # 添加控制变量均值
    control_vars = ['ability', 'urban', 'age']
    for var in control_vars:
        if var in df.columns and var not in prediction_data.columns:
            prediction_data[var] = df[var].mean()
    
    # 预测
    predictions = model.predict(prediction_data)
    
    # 可视化
    plt.figure(figsize=(10, 6))
    scenarios = ['对照组-政策前', '对照组-政策后', '处理组-政策前', '处理组-政策后']
    plt.bar(scenarios, predictions, alpha=0.7, color=['lightblue', 'lightgreen', 'lightcoral', 'lightsalmon'])
    
    plt.ylabel('预测收入', fontsize=12)
    plt.title('政策效应的异质性分析', fontsize=14, fontweight='bold')
    plt.xticks(rotation=45)
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()

plot_marginal_effects(hetero_model, df, 'high_ability')

# 多期DID与事件研究法
def event_study_did(df, outcome_var='income'):
    """事件研究法DID"""
    # 生成相对时期变量
    df_es = df.copy()
    policy_year = 2020  # 政策实施年份
    df_es['relative_year'] = df_es['year'] - policy_year
    
    # 创建相对时期虚拟变量（排除基准期-1）
    relative_years = sorted(df_es['relative_year'].unique())
    base_year = -1  # 政策前一年作为基准
    
    # 生成时期虚拟变量
    for ry in relative_years:
        if ry != base_year:
            df_es[f'rel_year_{ry}'] = (df_es['relative_year'] == ry).astype(int)
    
    # 生成交互项
    rel_year_cols = [col for col in df_es.columns if 'rel_year_' in col]
    interaction_terms = []
    for col in rel_year_cols:
        interac_col = f'treated_{col}'
        df_es[interac_col] = df_es['treated'] * df_es[col]
        interaction_terms.append(interac_col)
    
    # 事件研究回归
    es_formula = f"{outcome_var} ~ treated + " + " + ".join(rel_year_cols) + " + " + " + ".join(interaction_terms)
    es_model = smf.ols(es_formula, data=df_es).fit()
    
    return es_model, interaction_terms, relative_years

# 执行事件研究分析
es_model, interaction_terms, relative_years = event_study_did(df)

print("="*60)
print("事件研究法结果")
print("="*60)

# 提取事件研究系数
es_coefs = es_model.params[interaction_terms]
es_ci = es_model.conf_int().loc[interaction_terms]

# 可视化事件研究结果
def plot_event_study(es_coefs, interaction_terms, relative_years, base_year=-1):
    """绘制事件研究图"""
    # 提取相对时期和系数
    periods = []
    coefficients = []
    lower_bounds = []
    upper_bounds = []
    
    for term in interaction_terms:
        period = int(term.split('_')[-1])
        periods.append(period)
        coefficients.append(es_coefs[term])
        lower_bounds.append(es_ci.loc[term, 0])
        upper_bounds.append(es_ci.loc[term, 1])
    
    # 添加基准期（系数为0）
    periods.append(base_year)
    coefficients.append(0)
    lower_bounds.append(0)
    upper_bounds.append(0)
    
    # 排序
    sorted_indices = np.argsort(periods)
    periods = np.array(periods)[sorted_indices]
    coefficients = np.array(coefficients)[sorted_indices]
    lower_bounds = np.array(lower_bounds)[sorted_indices]
    upper_bounds = np.array(upper_bounds)[sorted_indices]
    
    # 创建图形
    plt.figure(figsize=(12, 6))
    
    # 政策实施时间线
    plt.axvline(x=-0.5, color='red', linestyle='--', alpha=0.7, label='政策实施')
    plt.axhline(y=0, color='black', linestyle='-', alpha=0.3)
    
    # 绘制点估计和置信区间
    plt.plot(periods, coefficients, 'o-', linewidth=2, markersize=8, label='政策效应')
    plt.fill_between(periods, lower_bounds, upper_bounds, alpha=0.2, label='95%置信区间')
    
    plt.xlabel('相对于政策实施的年份', fontsize=12)
    plt.ylabel('政策效应', fontsize=12)
    plt.title('事件研究: 政策效应的动态路径', fontsize=14, fontweight='bold')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.show()

plot_event_study(es_coefs, interaction_terms, relative_years)

# 堆叠DID模拟
def stacked_did_simulation():
    """堆叠DID模拟"""
    # 生成多时点政策数据
    np.random.seed(1234)
    n_units = 1000
    n_periods = 10
    
    stacked_data = []
    unit_ids = range(n_units)
    
    for unit in unit_ids:
        # 随机分配政策实施时间（有些单位从未被处理）
        treatment_time = np.random.choice([0, 4, 6, 8, 0], p=[0.3, 0.2, 0.2, 0.2, 0.1])
        # 0表示从未被处理
        
        for period in range(n_periods):
            # 基础结果
            base_outcome = 50 + 0.5 * period + np.random.normal(0, 2)
            
            # 处理效应（如果被处理且在当前时期之后）
            if treatment_time > 0 and period >= treatment_time:
                treatment_effect = 5 + np.random.normal(0, 1)
            else:
                treatment_effect = 0
                
            outcome = base_outcome + treatment_effect
            
            stacked_data.append({
                'id': unit,
                'period': period,
                'treatment_time': treatment_time,
                'treated': 1 if treatment_time > 0 else 0,
                'post_treatment': 1 if treatment_time > 0 and period >= treatment_time else 0,
                'outcome': outcome,
                'relative_time': period - treatment_time if treatment_time > 0 else -999
            })
    
    return pd.DataFrame(stacked_data)

# 生成堆叠DID数据
stacked_df = stacked_did_simulation()

# 堆叠DID估计
def stacked_did_estimation(df):
    """堆叠DID估计"""
    # 仅使用被处理单位和从未被处理单位
    analysis_df = df[df['treatment_time'].isin([0, 4, 6, 8])].copy()
    
    # 为每个政策队列创建堆叠数据集
    cohorts = [4, 6, 8]  # 政策实施时间点
    stacked_datasets = []
    
    for cohort in cohorts:
        cohort_df = analysis_df.copy()
        # 定义处理组和对照组
        cohort_df['stacked_treated'] = (cohort_df['treatment_time'] == cohort).astype(int)
        cohort_df['stacked_control'] = (cohort_df['treatment_time'] == 0).astype(int)
        
        # 选择相关时间窗口（政策前后各3期）
        time_window = 3
        cohort_df = cohort_df[((cohort_df['period'] >= cohort - time_window) & 
                              (cohort_df['period'] <= cohort + time_window)) |
                             (cohort_df['treatment_time'] == 0)]
        
        # 定义堆叠样本的post变量
        cohort_df['stacked_post'] = (cohort_df['period'] >= cohort).astype(int)
        
        # 添加队列标识
        cohort_df['cohort'] = cohort
        stacked_datasets.append(cohort_df)
    
    # 合并堆叠数据集
    stacked_data = pd.concat(stacked_datasets, ignore_index=True)
    
    # 堆叠DID回归
    stacked_model = smf.ols('outcome ~ stacked_treated + stacked_post + stacked_treated:stacked_post + C(cohort)', 
                           data=stacked_data).fit()
    
    return stacked_model, stacked_data

# 执行堆叠DID估计
stacked_model, stacked_data = stacked_did_estimation(stacked_df)

print("="*60)
print("堆叠DID估计结果")
print("="*60)
print(stacked_model.summary())

# 连续DID
def continuous_did(df):
    """连续DID分析"""
    # 生成连续处理强度
    np.random.seed(1234)
    df_cont = df.copy()
    
    # 模拟处理强度（如培训时长）
    df_cont['treatment_intensity'] = np.where(
        df_cont['treated'] == 1,
        np.random.exponential(10, size=len(df)),  # 处理组有不同强度
        0  # 对照组强度为0
    )
    
    # 连续DID模型
    cont_model = smf.ols('income ~ treatment_intensity + post + treatment_intensity:post + ability + urban + age', 
                        data=df_cont).fit()
    
    return cont_model, df_cont

# 执行连续DID分析
cont_model, df_cont = continuous_did(df)

print("="*60)
print("连续DID估计结果")
print("="*60)
print(cont_model.summary())

# 剂量响应曲线
def plot_dose_response(model, df):
    """绘制剂量响应曲线"""
    # 创建预测数据
    intensity_levels = np.linspace(0, 30, 100)
    prediction_data = pd.DataFrame({
        'treatment_intensity': intensity_levels,
        'post': 1,  # 政策后
        'ability': df['ability'].mean(),
        'urban': 1,
        'age': df['age'].mean()
    })
    
    # 预测
    predictions = model.predict(prediction_data)
    
    # 计算置信区间
    design_matrix = sm.add_constant(prediction_data)
    cov_matrix = model.cov_params()
    prediction_var = np.diag(design_matrix @ cov_matrix @ design_matrix.T)
    prediction_se = np.sqrt(prediction_var)
    ci_lower = predictions - 1.96 * prediction_se
    ci_upper = predictions + 1.96 * prediction_se
    
    # 可视化
    plt.figure(figsize=(10, 6))
    plt.plot(intensity_levels, predictions, linewidth=2, label='预测效应')
    plt.fill_between(intensity_levels, ci_lower, ci_upper, alpha=0.2, label='95%置信区间')
    
    plt.xlabel('处理强度', fontsize=12)
    plt.ylabel('预测收入', fontsize=12)
    plt.title('剂量响应关系', fontsize=14, fontweight='bold')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()

plot_dose_response(cont_model, df_cont)

# DID分析检查清单
def did_checklist():
    """DID分析检查清单"""
    checklist = {
        '研究设计': [
            '明确政策冲击时点',
            '确定处理组和对照组',
            '确保政策外生性',
            '选择合适的结果变量'
        ],
        '数据准备': [
            '平衡面板数据结构',
            '处理缺失值问题',
            '检查异常值',
            '变量标准化处理'
        ],
        '模型设定': [
            '平行趋势检验',
            '选择合适的固定效应',
            '考虑聚类标准误',
            '函数形式检验'
        ],
        '稳健性检验': [
            '安慰剂检验',
            '不同对照组选择',
            '模型设定变化',
            '异质性分析'
        ],
        '结果解释': [
            '经济显著性评估',
            '机制分析',
            '政策含义讨论',
            '局限性说明'
        ]
    }
    
    return checklist

# 显示检查清单
checklist = did_checklist()
print("="*60)
print("DID分析检查清单")
print("="*60)

for category, items in checklist.items():
    print(f"\n{category}:")
    for i, item in enumerate(items, 1):
        print(f"  {i}. {item}")

# 常见陷阱检测函数
def common_pitfalls_detection(df, model):
    """检测常见DID陷阱"""
    pitfalls = {}
    
    # 1. 样本选择偏差检测
    pre_treated = df[(df['treated'] == 1) & (df['post'] == 0)]
    pre_control = df[(df['treated'] == 0) & (df['post'] == 0)]
    
    # 均值差异检验
    for var in ['ability', 'age', 'urban']:
        t_stat, p_value = stats.ttest_ind(pre_treated[var], pre_control[var])
        pitfalls[f'前测平衡性_{var}'] = {
            'p_value': p_value,
            '结论': '通过' if p_value > 0.05 else '未通过'
        }
    
    # 2. 政策预期效应检测
    policy_year = 2020
    df['pre_policy'] = ((df['year'] == policy_year - 1) & (df['treated'] == 1)).astype(int)
    anticipation_test = smf.ols('income ~ treated + pre_policy + post + treated:post', 
                               data=df).fit()
    anticipation_p = anticipation_test.pvalues['pre_policy']
    pitfalls['政策预期效应'] = {
        'p_value': anticipation_p,
        '结论': '无预期效应' if anticipation_p > 0.05 else '存在预期效应'
    }
    
    # 3. 处理组稳定性检测
    treatment_stability = df.groupby('id')['treated'].std()
    pitfalls['处理组稳定性'] = {
        '处理组变动比例': (treatment_stability > 0).mean(),
        '结论': '稳定' if (treatment_stability > 0).mean() < 0.05 else '不稳定'
    }
    
    return pitfalls

# 检测常见陷阱
pitfalls = common_pitfalls_detection(df, model_income)

print("\n" + "="*60)
print("常见陷阱检测结果")
print("="*60)
for pitfall, results in pitfalls.items():
    print(f"{pitfall}:")
    for key, value in results.items():
        print(f"  {key}: {value}")

# 标准误比较
def compare_standard_errors(model, df):
    """比较不同标准误调整方法"""
    import linearmodels as lm
    from linearmodels.panel import PanelOLS
    
    # 准备面板数据
    df_panel = df.set_index(['id', 'period'])
    
    # 使用linearmodels进行面板数据估计
    # 注意：这里需要将数据转换为适合面板分析的形式
    try:
        # 双向固定效应模型
        panel_model = PanelOLS.from_formula(
            'income ~ treated*post + EntityEffects + TimeEffects', 
            data=df_panel
        ).fit(cov_type='clustered', cluster_entity=True)
        
        # 比较标准误
        se_comparison = {
            '传统标准误': model.bse['treated:post'],
            '聚类标准误': '需要完整面板数据模型'  # 实际应用中这里应该是具体数值
        }
    except:
        se_comparison = {'提示': '需要安装linearmodels库进行完整的面板数据分析'}
    
    return se_comparison

# 标准误比较结果
se_comparison = compare_standard_errors(model_income, df)
print("\n标准误比较:")
for method, value in se_comparison.items():
    print(f"{method}: {value}")

# 小样本偏差校正
def small_sample_correction(model, df):
    """小样本偏差校正"""
    n = len(df['id'].unique())  # 个体数量
    k = len(model.params)       # 参数数量
    
    # 自由度校正
    df_correction = n - k - 1
    
    # 使用t分布而非正态分布
    t_critical = stats.t.ppf(0.975, df_correction)
    normal_critical = stats.norm.ppf(0.975)
    
    correction_info = {
        '样本量': n,
        '参数个数': k,
        '校正后自由度': df_correction,
        't分布临界值': t_critical,
        '正态分布临界值': normal_critical,
        '建议': '使用t分布' if n < 100 else '正态分布近似可用'
    }
    
    return correction_info

# 小样本校正信息
correction_info = small_sample_correction(model_income, df)
print("\n小样本偏差校正:")
for key, value in correction_info.items():
    print(f"{key}: {value}")