因果森林与政策学习 Causal Forest & Policy Learning

# ============================================================
# 04.1 R: 用 grf 估计 CATE、变量重要性、置信区间
# install.packages("grf")
# ============================================================
library(grf)
set.seed(123)

# ---- 1. 模拟数据：n=2000, p=6；CATE 与 X1, X3 有关 ----
n <- 2000
p <- 6
X <- matrix(rnorm(n * p), n, p)
W <- rbinom(n, 1, 0.5)                                   # 随机分配处理
# 真实 CATE(x) = X[,1] + X[,3]（其他变量不影响效应）
tau_true <- X[, 1] + X[, 3]
Y <- tau_true * W + X[, 2] + pmin(X[, 4], 0) + rnorm(n)

# ---- 2. 拟合因果森林（默认 honesty=TRUE）----
cf <- causal_forest(
  X = X,
  Y = Y,
  W = W,
  num.trees = 2000,          # 树数
  min.node.size = 5,         # 叶子最小样本
  honesty = TRUE,            # 诚实树
  honesty.fraction = 0.5,    # 一半选分裂、一半估效应
  ci.group.size = 2          # 用于方差估计
)

# ---- 3. 用 out-of-bag 预测训练样本的 CATE ----
tau.hat.oob <- predict(cf)
hist(tau.hat.oob$predictions,
     main = "OOB CATE 估计分布", xlab = "tau_hat(x)")

# ---- 4. 置信区间（在测试样本上）----
X.test <- matrix(rnorm(100 * p), 100, p)
tau.test <- predict(cf, X.test, estimate.variance = TRUE)
# tau.test$predictions       : CATE 点估计
# tau.test$variance.estimates: 方差估计
# 95% CI = predictions ± 1.96 * sqrt(variance.estimates)
ci_lower <- tau.test$predictions - 1.96 * sqrt(tau.test$variance.estimates)
ci_upper <- tau.test$predictions + 1.96 * sqrt(tau.test$variance.estimates)

# ---- 5. 模型校准诊断（calibration test）----
test_calibration(cf)
# 输出：估计 ATE、最优线性投影 R^2、是否校准良好

# ---- 6. 变量重要性（哪些 X 决定了 CATE）----
varimp <- variable_importance(cf)
print(data.frame(variable = paste0("X", 1:p),
                 importance = round(varimp, 3)))
# X1, X3 的 importance 应显著高于 X2, X4, X5, X6

# ---- 7. 平均处理效应 ATE ----
ate_cf <- average_treatment_effect(cf)
print(ate_cf)   # 估计 ATE 与标准误

# ============================================================
# 04.2 R: 用 policy_tree 从 CATE 到最优政策
# ============================================================
library(grf)

# 用上一步的 cf
# 政策树：直接在观察数据上学习"哪些 x 应被处理"
pt <- policy_tree(X, Y, W,
                  depth = 2,             # 树深
                  policy = "overall")    # 最大化整体福利

# 在测试样本上预测政策：1=建议处理, 0=建议不处理
policy.predict <- predict(pt, X.test)
table(policy.predict$predictions)

# 画政策树
plot(pt)
# 树的每个叶子给出"按哪个变量切分 -> 该叶子里是否建议处理"
# 这就是可解释的"政策规则"

# ============================================================
# 05.1 Python: econml CausalForestDML
# pip install econml scikit-learn
# ============================================================
import numpy as np
import matplotlib.pyplot as plt
from sklearn.ensemble import GradientBoostingRegressor, GradientBoostingClassifier
from econml.dml import CausalForestDML

np.random.seed(42)

# ---- 1. 模拟数据（与 R 版相同）----
n, p = 2000, 6
X = np.random.normal(size=(n, p))
W = np.random.binomial(1, 0.5, size=n)
tau_true = X[:, 0] + X[:, 2]            # 真实 CATE(x)
Y = tau_true * W + X[:, 1] + np.minimum(X[:, 3], 0) + np.random.normal(size=n)

# ---- 2. 拟合 CausalForestDML ----
est = CausalForestDML(
    model_y=GradientBoostingRegressor(n_estimators=200, max_depth=3,
                                      random_state=42),
    model_t=GradientBoostingClassifier(n_estimators=200, max_depth=3,
                                       random_state=42),
    discrete_treatment=True,              # W 是二值
    n_estimators=1000,                    # 树数
    min_samples_leaf=20,
    max_samples=0.5,
    honest=True,                          # 诚实分裂
    inference=True,                       # 启用置信区间
    random_state=42,
)
est.fit(Y, W, X=X)

# ---- 3. CATE 点估计与置信区间 ----
cate_hat = est.effect(X)                  # 每个个体的 CATE
lb, ub = est.effect_interval(X, alpha=0.05)

print("CATE 均值 =", cate_hat.mean())
print("CATE 标准差 =", cate_hat.std())
print("95% CI 半宽 均值 =", ((ub - lb) / 2).mean())

# ---- 4. 与真实 CATE 对比 ----
plt.figure(figsize=(6, 6))
plt.scatter(tau_true, cate_hat, alpha=0.5, s=10)
plt.plot([-4, 4], [-4, 4], "r--")
plt.xlabel("true CATE"); plt.ylabel("estimated CATE")
plt.title("Causal Forest: true vs estimated CATE")
plt.tight_layout()
plt.savefig("cate_true_vs_hat.png", dpi=120)

# ---- 5. 变量重要性（基于分裂频率）----
imp = est.feature_importances_
for i in range(p):
    print(f"X{i+1}: importance = {imp[i]:.3f}")
# X1, X3 应最高

# ---- 6. 政策解释：用单棵决策树逼近 CATE ----
from econml.cate_interpreter import SingleTreeCateInterpreter
interp = SingleTreeCateInterpreter(include_model_uncertainty=True,
                                   max_depth=2)
interp.interpret(est, X)
# interp.plot()  # 在 Jupyter 中画可解释的 CATE 树