DML 双重/去偏机器学习 Double Machine Learning

# ============================================================
# 06.1 DoubleML 入门：PLR 模型 + 随机森林第一阶段
# 安装：pip install doubleml scikit-learn numpy pandas
# ============================================================
import numpy as np
from doubleml import DoubleMLData, DoubleMLPLR
from doubleml.datasets import make_plr_CCDDHNR2018
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.linear_model import LassoCV

np.random.seed(3141)

# ---- 1. 生成模拟数据（真实 theta0 = 0.5）----
# n=500, p=20 协变量，混淆函数为非线性
df = make_plr_CCDDHNR2018(alpha=0.5, return_type="dataframe")
# df 含列：y（结果）, d（处理）, X1...Xp（协变量）
print("数据形状:", df.shape, "| 真实 theta0 = 0.5")

# ---- 2. 包装成 DoubleMLData ----
obj_dml_data = DoubleMLData(df, y_col="y", d_cols="d")

# ---- 3. 定义第一阶段学习器 ----
# ml_l: 拟合 E[Y|X]（连续，回归树）
ml_l = RandomForestRegressor(n_estimators=500, max_depth=5, random_state=3141)
# ml_m: 拟合 E[D|X]（这里 D 是连续；若是二值处理用 RandomForestClassifier）
ml_m = RandomForestRegressor(n_estimators=500, max_depth=5, random_state=3141)

# ---- 4. 构造 PLR 对象并拟合（默认 5-fold cross-fitting）----
dml_plr = DoubleMLPLR(obj_dml_data,
                      ml_l=ml_l, ml_m=ml_m,
                      n_folds=5,          # 5 折 cross-fitting
                      score="partialling out",
                      dml_procedure="dml2")
dml_plr.fit()
print(dml_plr.summary)
# 输出 theta、std err、t、p-value、2.5%/97.5% 置信区间
# 与真实 0.5 对比：DML 的 theta 应接近 0.5，CI 覆盖率接近 95%

# ---- 5. 换 Lasso 当第一阶段 ----
ml_l_lasso = LassoCV(cv=5)
ml_m_lasso = LassoCV(cv=5)
dml_plr_lasso = DoubleMLPLR(obj_dml_data,
                            ml_l=ml_l_lasso, ml_m=ml_m_lasso,
                            n_folds=5)
dml_plr_lasso.fit()
print(dml_plr_lasso.summary)

# ============================================================
# 06.2 蒙特卡洛对比：OLS vs Lasso vs DML
# 重复 R 次实验，看偏误（bias）和 95% CI 覆盖率
# ============================================================
import numpy as np
import statsmodels.api as sm
from sklearn.linear_model import LassoCV
from doubleml import DoubleMLData, DoubleMLPLR
from sklearn.ensemble import RandomForestRegressor

np.random.seed(0)
R = 50                      # 重复次数（正式研究可调到 1000）
n, p = 1000, 20
theta0 = 0.5
res = {"OLS": [], "Lasso": [], "DML": []}

for r in range(R):
    # 模拟：X ~ N(0,1)；m(X)=非线性；g(X)=非线性
    X = np.random.normal(size=(n, p))
    m_x = np.tanh(X[:, 0] + X[:, 1])           # E[D|X]
    g_x = np.tanh(X[:, 2] + X[:, 3])           # E[Y|X]
    D = m_x + np.random.normal(size=n)
    Y = theta0 * D + g_x + np.random.normal(size=n)

    # --- (a) OLS：把 X 全塞进去 ---
    X_design = sm.add_constant(np.column_stack([D, X]))
    ols = sm.OLS(Y, X_design).fit()
    res["OLS"].append(ols.params[1])

    # --- (b) Lasso 选变量后 OLS ---
    lasso = LassoCV(cv=5).fit(X, Y)
    sel = np.where(np.abs(lasso.coef_) > 1e-8)[0]
    X_sel = sm.add_constant(np.column_stack([D, X[:, sel]]))
    l = sm.OLS(Y, X_sel).fit()
    res["Lasso"].append(l.params[1])

    # --- (c) DML（随机森林第一阶段）---
    import pandas as pd
    df_r = pd.DataFrame(np.column_stack([Y, D, X]),
                        columns=["y", "d"] + [f"X{i}" for i in range(p)])
    dd = DoubleMLData(df_r, y_col="y", d_cols="d",
                      x_cols=[f"X{i}" for i in range(p)])
    rf = RandomForestRegressor(n_estimators=200, max_depth=4, random_state=r)
    dml = DoubleMLPLR(dd, ml_l=rf, ml_m=rf, n_folds=5)
    dml.fit()
    res["DML"].append(float(dml.coef[0]))

# ---- 输出对比 ----
for k, v in res.items():
    v = np.array(v)
    bias = v.mean() - theta0
    print(f"{k:6s}: 均值={v.mean():.4f}, 偏误={bias:+.4f}, 标准差={v.std():.4f}")
# 预期：OLS 有偏（未正确建模非线性）；
#       Lasso 偏误较大且 CI 不准；
#       DML 偏误最小、最接近 theta0=0.5

# ============================================================
# 06.3 econml 中的 DML（LinearDML）
# pip install econml
# ============================================================
from econml.dml import LinearDML
from sklearn.ensemble import GradientBoostingRegressor, GradientBoostingClassifier
import numpy as np, pandas as pd

# 复用上面的数据
X = np.random.normal(size=(n, p))
D = np.tanh(X[:, 0] + X[:, 1]) + np.random.normal(size=n)
Y = theta0 * D + np.tanh(X[:, 2] + X[:, 3]) + np.random.normal(size=n)

est = LinearDML(
    model_y=GradientBoostingRegressor(n_estimators=200, max_depth=3),
    model_t=GradientBoostingRegressor(n_estimators=200, max_depth=3),
    discrete_treatment=False,        # 连续处理
    cv=5,                            # cross-fitting
    random_state=0,
)
est.fit(Y, D, X=X)
# 平均处理效应
print("ATE =", est.ate(X), " 95% CI =", est.ate_interval(X, alpha=0.05))
# 线性异质性系数（CATE(x) = intercept + coefs @ x）
print("CATE 线性模型系数:")
print(est.intercept_, est.coef_)

# ============================================================
# 07.1 R 版 DML：install.packages(c("DoubleML","mlr3","mlr3learners"))
# ============================================================
library(DoubleML)
library(mlr3)
library(mlr3learners)
library(data.table)

set.seed(3141)

# ---- 1. 用内置 DGP（与 Python 版相同）----
data = make_plr_CCDDHNR2018(alpha = 0.5, n_obs = 5000,
                             return_type = "data.table")
# 列：y, d, X1..X20

# ---- 2. 包装数据 ----
obj_dml_data = DoubleMLData$new(data, y_col = "y", d_cols = "d")

# ---- 3. 定义第一阶段学习器（mlr3 语法）----
# 连续结果用 regr.lm / regr.ranger；二值处理用 classif.ranger
ml_l = lrn("regr.ranger", num.trees = 500, max.depth = 5)
ml_m = lrn("regr.ranger", num.trees = 500, max.depth = 5)

# ---- 4. 构造 PLR 对象并拟合 ----
dml_plr = DoubleMLPLR$new(obj_dml_data,
                          ml_l = ml_l,
                          ml_m = ml_m,
                          n_folds = 5,
                          score = "partialling out")
dml_plr$fit()
print(dml_plr)
# 输出：theta = 0.5 (0.0X)，置信区间与 Python 版一致

# ---- 5. 总结表 ----
dml_plr$summary()

# ---- 6. 换 Lasso（glmnet）----
library(mlr3learners)
ml_l_lasso = lrn("regr.glmnet", alpha = 1, lambda = 0.05)
ml_m_lasso = lrn("regr.glmnet", alpha = 1, lambda = 0.05)
dml_plr_lasso = DoubleMLPLR$new(obj_dml_data,
                                 ml_l = ml_l_lasso,
                                 ml_m = ml_m_lasso,
                                 n_folds = 5)
dml_plr_lasso$fit()
print(dml_plr_lasso$summary())

* ============================================================
* 08.1 Stata: ddml 做部分线性模型 DML
* ssc install ddml
* ssc install rforest    // 随机森林后端
* ============================================================
clear all
set seed 3141
set obs 5000

* --- 模拟数据 ---
drawnorm X1-X20, means(0) sds(1)
gen D = tanh(X1 + X2) + rnormal()
gen Y = 0.5*D + tanh(X3 + X4) + rnormal()

* --- 初始化 ddml 环境 ---
ddml init partial, kfolds(5)

* --- 第一阶段：用随机森林拟合 E[Y|X] 和 E[D|X] ---
ddml E[Y|X]: rforest Y X1-X20, type(reg)
ddml E[D|X]: rforest D X1-X20, type(reg)

* --- 第二阶段：在残差上估计 theta ---
ddml estimate, robust

* --- 查看结果 ---
ddml summary
* 输出 theta、标准误、p 值、置信区间
* 预期 theta 接近 0.5

* --- 换 Lasso ---
ddml init partial, kfolds(5)
ddml E[Y|X]: lasso Y X1-X20
ddml E[D|X]: lasso D X1-X20
ddml estimate, robust
ddml summary

# ============================================================
# 12.3 手工 K 折 cross-fitted AIPW（理解内部机制用）
# 安装：pip install numpy pandas scikit-learn
# ============================================================
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.model_selection import KFold

np.random.seed(20240101)
n, p = 2000, 6
X = np.random.normal(size=(n, p))

# --- DGP：真实 ATE ≈ 1 ---
# e(X) = logit^{-1}(X1 + 0.5*X2)
e_x = 1 / (1 + np.exp(-(X[:, 0] + 0.5 * X[:, 1])))
D = (np.random.uniform(size=n) < e_x).astype(float)
m1 = 1 + X[:, 0]                 # 处理组均值函数
m0 = 0.5 * X[:, 1]               # 对照组均值函数
Y = D * m1 + (1 - D) * m0 + np.random.normal(scale=0.5, size=n)

# --- K=5 折 cross-fitting ---
kf = KFold(n_splits=5, shuffle=True, random_state=20240101)
# 存放 out-of-fold 的 e_hat, m1_hat, m0_hat
e_hat = np.zeros(n); m1_hat = np.zeros(n); m0_hat = np.zeros(n)

for train_idx, test_idx in kf.split(X):
    Xtr, Xte = X[train_idx], X[test_idx]
    Dtr, Ytr = D[train_idx], Y[train_idx]
    # 倾向得分：分类器
    ps = RandomForestClassifier(n_estimators=300, max_depth=4, random_state=0)
    ps.fit(Xtr, Dtr)
    e_hat[test_idx] = ps.predict_proba(Xte)[:, 1]
    # 结果回归：分别在 D=1 / D=0 上训练，预测全折
    r1 = RandomForestRegressor(n_estimators=300, max_depth=4, random_state=0)
    r1.fit(Xtr[Dtr == 1], Ytr[Dtr == 1])
    r0 = RandomForestRegressor(n_estimators=300, max_depth=4, random_state=0)
    r0.fit(Xtr[Dtr == 0], Ytr[Dtr == 0])
    m1_hat[test_idx] = r1.predict(Xte)
    m0_hat[test_idx] = r0.predict(Xte)

# --- 数值稳定：把倾向得分截断到 [0.01, 0.99]（trim 极端权重）---
e_hat = np.clip(e_hat, 0.01, 0.99)

# --- AIPW score（逐个体），见 s11 的公式 ---
score = (D * (Y - m1_hat) / e_hat
         - (1 - D) * (Y - m0_hat) / (1 - e_hat)
         + m1_hat - m0_hat)
tau_aipw = score.mean()
print(f"cross-fitted AIPW ATE = {tau_aipw:.4f}  (真实 ATE = 1.0)")

# --- 对比：不做 cross-fitting（1 折，全样本训练）---
ps_all = RandomForestClassifier(n_estimators=300, max_depth=4, random_state=0).fit(X, D)
e_all = np.clip(ps_all.predict_proba(X)[:, 1], 0.01, 0.99)
r1a = RandomForestRegressor(n_estimators=300, max_depth=4, random_state=0).fit(X[D==1], Y[D==1])
r0a = RandomForestRegressor(n_estimators=300, max_depth=4, random_state=0).fit(X[D==0], Y[D==0])
m1a, m0a = r1a.predict(X), r0a.predict(X)
score_nocf = (D*(Y-m1a)/e_all - (1-D)*(Y-m0a)/(1-e_all) + m1a - m0a)
print(f"不做 cross-fitting 的 ATE = {score_nocf.mean():.4f}  (通常偏差更大)")

* ============================================================
* 13.1 Stata: teffects aipw 双重稳健 ATE
* 模拟数据：真实 ATE ≈ 1
* ============================================================
clear all
set seed 20240101
set obs 2000

drawnorm X1 X2 X3, means(0) sds(1)

* 倾向得分 e(X) = invlogit(X1 + 0.5*X2)
gen double linp = X1 + 0.5*X2
gen double e_x  = invlogit(linp)
gen byte D = runiform() < e_x

* 结果：m1(X)=1+X1, m0(X)=0.5*X2
gen double Y = D*(1 + X1) + (1-D)*(0.5*X2) + rnormal()

* --- AIPW：(结果回归变量) 在前，(处理变量) 在后 ---
* 语法：teffects aipw y (处理方程协变量) (处理变量)
teffects aipw Y (X1 X2 X3) (D), osample(oor)
* osample() 生成共同支撑标记 oor（=0 表示倾向得分极端、缺共同支撑）

* --- 诊断：检查共同支撑 ---
teffects overlap        // 画处理组/对照组倾向得分分布图
count if oor == 1        // 数出被剔除的极端权重样本
teffects stats          // 报告 ATE、POT0、POT1 及其 SE

* --- 对比：纯 IPW 与纯回归调整，看 AIPW 的稳健性 ---
teffects ipw Y (X1 X2 X3) (D)
teffects ra Y (X1 X2 X3) (D)
* 三列 ATE 应接近；若 IPW 偏离很大而 AIPW 接近，说明结果模型在"兜底"

* ============================================================
* 13.2 Stata: drdid 双重稳健 DID（交错处理 / staggered adoption）
* ssc install drdid
* ssc install csdid
* ============================================================

* --- 截面 / 重复截面数据：drdid ---
* drdid y [treat] time, ivar(panel id) [repeated]
* 例：drdid Y D post, ivar(id)
*   - Y     结果变量
*   - D     处理组虚拟（=1 表示最终会被处理）
*   - post  政策实施后虚拟
*   drdid 内部用 AIPW/DR 估计 ATT，并报告双重稳健标准误

* --- 面板数据：csdid（推荐，按首次处理年份 gvar 分组）---
* webuse Hanson_2021, clear
* csdid Y year gvar, notyet ivar(id)
*   - gvar = 个体首次接受处理的年份（从未处理 = 0）
* 输出：ATT(g,t) 矩阵 + 加权聚合的 overall ATT
* 事后画动态效应图：
* estat event

* --- 下面用模拟面板演示语法结构 ---
clear all
set seed 20240101
set obs 1000
gen long id = _n
expand 10
bysort id: gen year = 2010 + _n - 1
bysort id: gen double x = rnormal()
* 组别 gvar：id 尾号 1=2015 处理，2=2017 处理，其余=0（对照）
gen gvar = cond(mod(id,3)==1, 2015, cond(mod(id,3)==2, 2017, 0))
gen post = (year >= gvar) & (gvar > 0)
gen Y = 1.2*post + x + rnormal()

* csdid Y gvar, timevar(year) ivar(id) notyet
* estat event
* 若 csdid 不可用，退回 drdid：
* drdid Y post, ivar(id)

# ============================================================
# 13.3 econml DRLearner（AIPW + 元学习器）
# pip install econml scikit-learn
# ============================================================
import numpy as np
from econml.dr import DRLearner
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier

np.random.seed(20240101)
n, p = 2000, 6
X = np.random.normal(size=(n, p))
e_x = 1/(1+np.exp(-(X[:,0] + 0.5*X[:,1])))
D = (np.random.uniform(size=n) < e_x).astype(int)
Y = D*(1 + X[:,0]) + (1-D)*(0.5*X[:,1]) + np.random.normal(scale=0.5, size=n)

est = DRLearner(
    # 倾向得分模型（分类器）
    model_propensity=RandomForestClassifier(n_estimators=300, max_depth=4),
    # 结果模型（回归器，分别拟合处理/对照）
    model_regression=RandomForestRegressor(n_estimators=300, max_depth=4),
    # 第二阶段元学习器：在 AIPW 个别效应上拟合 CATE(X)
    model_final=RandomForestRegressor(n_estimators=300, max_depth=3),
    cv=5,                      # 5 折 cross-fitting
    random_state=20240101,
)
est.fit(Y, D, X=X)
print("ATE =", est.ate(X), " 95% CI =", est.ate_interval(X, alpha=0.05))
# 真实 ATE ≈ 1.0
# 异质性：对每个个体给出 CATE(x)
cate = est.effect(X)
print("CATE 均值 ≈", cate.mean(), "| CATE 标准差（异质性大小）=", cate.std())

# ============================================================
# 13.4 DoubleMLIRM = cross-fitted AIPW
# pip install doubleml scikit-learn
# ============================================================
import numpy as np, pandas as pd
from doubleml import DoubleMLData, DoubleMLIRM
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor

np.random.seed(20240101)
n, p = 2000, 6
X = np.random.normal(size=(n, p))
e_x = 1/(1+np.exp(-(X[:,0] + 0.5*X[:,1])))
D = (np.random.uniform(size=n) < e_x).astype(int)
Y = D*(1 + X[:,0]) + (1-D)*(0.5*X[:,1]) + np.random.normal(scale=0.5, size=n)

df = pd.DataFrame(np.column_stack([Y, D, X]),
                  columns=["y", "d"] + [f"X{i}" for i in range(p)])
dml_data = DoubleMLData(df, y_col="y", d_cols="d",
                        x_cols=[f"X{i}" for i in range(p)])

# ml_propensity: 倾向得分 e(X)；ml_outcome: 结果 m1/m0
irm = DoubleMLIRM(
    dml_data,
    ml_propensity=RandomForestClassifier(n_estimators=300, max_depth=4),
    ml_outcome=RandomForestRegressor(n_estimators=300, max_depth=4),
    n_folds=5,
    score="AIPW",              # 双重稳健 score
    dml_procedure="dml2",
)
irm.fit()
print(irm.summary)
# 输出 theta(≈1.0)、SE、p 值、95% CI，与手工 12.3 / econml 13.3 一致