因果推断机器学习¶

相关: A 和 B 同时升高（可能有共同原因 C）
因果: 改变 A 会导致 B 改变（do-演算，干预）

import numpy as np
import pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler

# ===== 1. 倾向性得分匹配（PSM）=====
# 场景：比较两组患者（治疗组 vs 对照组），消除混杂因子

np.random.seed(42)
n = 500  # 样本量

# 模拟观察数据：年龄、BMI、吸烟状态 → 影响是否接受治疗
age = np.random.normal(50, 10, n)           # 年龄
bmi = np.random.normal(25, 5, n)            # BMI
smoking = np.random.binomial(1, 0.3, n)    # 吸烟（0/1）

# 治疗概率受年龄和 BMI 影响（混杂！）
treatment_prob = 1 / (1 + np.exp(-(0.05*age + 0.1*bmi - 6 + 0.5*smoking)))
treatment = np.random.binomial(1, treatment_prob, n)  # 治疗分配（1=治疗，0=对照）

# 真实结局（治疗效果 = 0.3，另有年龄影响）
outcome = 0.3 * treatment + 0.02 * age + np.random.normal(0, 0.5, n)

df = pd.DataFrame({
    'age': age, 'bmi': bmi, 'smoking': smoking,
    'treatment': treatment, 'outcome': outcome
})

# 第1步：用逻辑回归估计倾向性得分
X_confounders = df[['age', 'bmi', 'smoking']]
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X_confounders)

ps_model = LogisticRegression()
ps_model.fit(X_scaled, df['treatment'])
# 倾向性得分 = 每个样本接受治疗的预测概率
df['propensity_score'] = ps_model.predict_proba(X_scaled)[:, 1]

# 第2步：最近邻匹配（每个治疗样本找最相近倾向得分的对照样本）
treated = df[df['treatment'] == 1].copy()
control = df[df['treatment'] == 0].copy()

matched_pairs = []
used_control_idx = set()

for i, t_row in treated.iterrows():
    # 找倾向得分最接近的未被匹配的对照样本
    available_control = control[~control.index.isin(used_control_idx)]
    if len(available_control) == 0:
        break
    ps_diff = abs(available_control['propensity_score'] - t_row['propensity_score'])
    best_match_idx = ps_diff.idxmin()          # 最小差异的对照样本

    matched_pairs.append({
        'treated_outcome': t_row['outcome'],
        'control_outcome': available_control.loc[best_match_idx, 'outcome']
    })
    used_control_idx.add(best_match_idx)

matched_df = pd.DataFrame(matched_pairs)

# 第3步：计算匹配后的平均治疗效应（ATT）
att = (matched_df['treated_outcome'] - matched_df['control_outcome']).mean()
print(f'倾向得分匹配估计的平均治疗效应 ATT: {att:.3f}')
print(f'真实治疗效应: 0.300')

# ===== 2. 因果森林（异质处理效应）=====
# pip install econml
try:
    from econml.grf import CausalForest

    cf = CausalForest(
        n_estimators=100,     # 树的数量
        min_samples_leaf=10,  # 叶子节点最少样本
        random_state=42
    )
    X = df[['age', 'bmi', 'smoking']].values
    T = df['treatment'].values
    Y = df['outcome'].values

    cf.fit(X, T, Y)  # 拟合因果森林

    # 为每个个体估计处理效应（个性化治疗效应）
    te_pred = cf.effect(X)
    print(f'\n个体处理效应估计（前5个）: {te_pred[:5].round(3)}')
    print(f'平均处理效应: {te_pred.mean():.3f}')
except ImportError:
    print("需要安装 econml: pip install econml")

# ===== 3. 孟德尔随机化（简化演示）=====
# 工具变量分析：SNP → 暴露（BMI）→ 结局（糖尿病风险）
def mendelian_randomization_wald(beta_exposure, beta_outcome, se_exposure, se_outcome):
    """
    Wald 估计量：最简单的 MR 方法
    单一 SNP 工具变量
    """
    # 因果效应估计 = SNP对结局的效应 / SNP对暴露的效应
    causal_estimate = beta_outcome / beta_exposure

    # Delta 法估计标准误
    se_ratio = abs(causal_estimate) * np.sqrt(
        (se_outcome/beta_outcome)**2 + (se_exposure/beta_exposure)**2
    )
    return causal_estimate, se_ratio

# 假设数据：某 SNP 与 BMI 的 GWAS 效应估计
beta_bmi_snp = 0.3       # SNP 对 BMI 的效应
se_bmi_snp = 0.05
beta_diabetes_snp = 0.12  # SNP 对糖尿病风险的效应
se_diabetes_snp = 0.03

causal_effect, se = mendelian_randomization_wald(
    beta_bmi_snp, beta_diabetes_snp, se_bmi_snp, se_diabetes_snp
)
print(f'\nMR Wald 估计：BMI → 糖尿病因果效应 = {causal_effect:.3f} ± {se:.3f}')
print(f'95% CI: [{causal_effect - 1.96*se:.3f}, {causal_effect + 1.96*se:.3f}]')