因果推断笔记——DML ：Double Machine Learning案例学习（十六）

# 连续数值的T
array([ 0.41083324,  0.57893171,  1.08130798, ..., -0.43799495,
        1.61941775,  1.64209826])

# 二分类01的T
array([0, 0, 0, 0, 1, 1, 0, 1, 1, 1, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 0, 0])

# Multi-T
array([[1.35325451, 1.15373159, 0.46373402],
       [1.35325451, 1.15373159, 0.98954119],
       [1.35325451, 0.87129337, 0.73716407],
       ...,
       [1.13462273, 1.03673688, 0.46373402],
       [1.0260416 , 0.95165788, 0.39877612],
       [1.05779029, 0.78390154, 0.55961579]])
       
# Multi-Y
array([[ 9.01869549,  8.40737833,  9.26482856],
       [ 8.72323127,  8.44934252,  8.98719682],
       [ 8.25322765,  9.91145572,  8.83171192],
       ...,
       [ 9.93653541,  9.51546936,  9.50599061],
       [10.85591733,  9.31793838, 10.9273763 ],
       [10.23652533, 11.20553036,  8.85936345]])

import econml
## Ignore warnings
import warnings
warnings.filterwarnings("ignore")

# Main imports
from econml.dml import DML, LinearDML, SparseLinearDML, CausalForestDML

# Helper imports
import numpy as np
from itertools import product
from sklearn.linear_model import (Lasso, LassoCV, LogisticRegression,
                                  LogisticRegressionCV,LinearRegression,
                                  MultiTaskElasticNet,MultiTaskElasticNetCV)
from sklearn.ensemble import RandomForestRegressor,RandomForestClassifier
from sklearn.preprocessing import PolynomialFeatures
import matplotlib.pyplot as plt
import matplotlib
from sklearn.model_selection import train_test_split

%matplotlib inline
# Treatment effect function
def exp_te(x):
    return np.exp(2*x[0])

# DGP constants
np.random.seed(123)
n = 2000
n_w = 30
support_size = 5
n_x = 1
# Outcome support
support_Y = np.random.choice(np.arange(n_w), size=support_size, replace=False)
coefs_Y = np.random.uniform(0, 1, size=support_size)
epsilon_sample = lambda n: np.random.uniform(-1, 1, size=n)
# Treatment support
support_T = support_Y
coefs_T = np.random.uniform(0, 1, size=support_size)
eta_sample = lambda n: np.random.uniform(-1, 1, size=n)

# Generate controls, covariates, treatments and outcomes
W = np.random.normal(0, 1, size=(n, n_w))
X = np.random.uniform(0, 1, size=(n, n_x))
# Heterogeneous treatment effects
TE = np.array([exp_te(x_i) for x_i in X])
T = np.dot(W[:, support_T], coefs_T) + eta_sample(n)
Y = TE * T + np.dot(W[:, support_Y], coefs_Y) + epsilon_sample(n)

Y_train, Y_val, T_train, T_val, X_train, X_val, W_train, W_val = train_test_split(Y, T, X, W, test_size=.2)
# Generate test data
X_test = np.array(list(product(np.arange(0, 1, 0.01), repeat=n_x)))

# 模型
est = LinearDML(model_y=RandomForestRegressor(),
                model_t=RandomForestRegressor(),
                random_state=123)
est.fit(Y_train, T_train, X=X_train, W=W_train)
te_pred = est.effect(X_test)
# 模型
est1 = SparseLinearDML(model_y=RandomForestRegressor(),
                       model_t=RandomForestRegressor(),
                       featurizer=PolynomialFeatures(degree=3),
                       random_state=123)
est1.fit(Y_train, T_train, X=X_train, W=W_train)
te_pred1 = est1.effect(X_test)

# 模型
est2 = DML(model_y=RandomForestRegressor(),
           model_t=RandomForestRegressor(),
           model_final=Lasso(alpha=0.1, fit_intercept=False),
           featurizer=PolynomialFeatures(degree=10),
           random_state=123)
est2.fit(Y_train, T_train, X=X_train, W=W_train)
te_pred2 = est2.effect(X_test)
# 模型
est3 = CausalForestDML(model_y=RandomForestRegressor(),
                       model_t=RandomForestRegressor(),
                       criterion='mse', n_estimators=1000,
                       min_impurity_decrease=0.001,
                       random_state=123)
est3.tune(Y_train, T_train, X=X_train, W=W_train)
est3.fit(Y_train, T_train, X=X_train, W=W_train)
te_pred3 = est3.effect(X_test)

# 画图
plt.figure(figsize=(10,6))
plt.plot(X_test, te_pred, label='DML default')
plt.plot(X_test, te_pred1, label='DML polynomial degree=3')
plt.plot(X_test, te_pred2, label='DML polynomial degree=10 with Lasso')
plt.plot(X_test, te_pred3, label='ForestDML')
expected_te = np.array([exp_te(x_i) for x_i in X_test])
plt.plot(X_test, expected_te, 'b--', label='True effect')
plt.ylabel('Treatment Effect')
plt.xlabel('x')
plt.legend()
plt.show()

score={}
score["DML default"] = est.score(Y_val, T_val, X_val, W_val)
score["DML polynomial degree=3"] = est1.score(Y_val, T_val, X_val, W_val)
score["DML polynomial degree=10 with Lasso"] = est2.score(Y_val, T_val, X_val, W_val)
score["ForestDML"] = est3.score(Y_val, T_val, X_val, W_val)

mse_te={}
mse_te["DML default"] = ((expected_te - te_pred)**2).mean()
mse_te["DML polynomial degree=3"] = ((expected_te - te_pred1)**2).mean()
mse_te["DML polynomial degree=10 with Lasso"] = ((expected_te - te_pred2)**2).mean()
mse_te["ForestDML"] = ((expected_te - te_pred3)**2).mean()

print("best model selected by MSE of TE: ", min(mse_te, key=lambda x: mse_te.get(x)))

import econml
# Main imports
from econml.dynamic.dml import DynamicDML
from econml.tests.dgp import DynamicPanelDGP, add_vlines

# Helper imports
import numpy as np
from sklearn.linear_model import Lasso, LassoCV, LogisticRegression, LogisticRegressionCV, MultiTaskLassoCV
import matplotlib.pyplot as plt

%matplotlib inline

# Define DGP parameters
np.random.seed(123)
n_panels = 5000 # number of panels
n_periods = 3 # number of time periods in each panel
n_treatments = 2 # number of treatments in each period
n_x = 100 # number of features + controls
s_x = 10 # number of controls (endogeneous variables)
s_t = 10 # treatment support size

# Generate data
dgp = DynamicPanelDGP(n_periods, n_treatments, n_x).create_instance(
            s_x, random_seed=12345)
Y, T, X, W, groups = dgp.observational_data(n_panels, s_t=s_t, random_seed=12345)
true_effect = dgp.true_effect

# Train Estimator
est = DynamicDML(
    model_y=LassoCV(cv=3, max_iter=1000), 
    model_t=MultiTaskLassoCV(cv=3, max_iter=1000), 
    cv=3)
est.fit(Y, T, X=None, W=W, groups=groups)

# Effect of target policy over baseline policy
# Must specify a treatment for each period
baseline_policy = np.zeros((1, n_periods * n_treatments))
target_policy = np.ones((1, n_periods * n_treatments))
eff = est.effect(T0=baseline_policy, T1=target_policy)
print(f"Effect of target policy over baseline policy: {eff[0]:0.2f}")

# Period treatment effects + interpretation
for i, theta in enumerate(est.intercept_.reshape(-1, n_treatments)):
    print(f"Marginal effect of a treatments in period {i+1} on period {n_periods} outcome: {theta}")

# Period treatment effects with confidence intervals
est.summary()

conf_ints = est.intercept__interval(alpha=0.05)

# Some plotting boilerplate code
plt.figure(figsize=(15, 5))
plt.errorbar(np.arange(n_periods*n_treatments)-.04, est.intercept_, yerr=(conf_ints[1] - est.intercept_,
                                                    est.intercept_ - conf_ints[0]), fmt='o', label='DynamicDML')
plt.errorbar(np.arange(n_periods*n_treatments), true_effect.flatten(), fmt='o', alpha=.6, label='Ground truth')
for t in np.arange(1, n_periods):
    plt.axvline(x=t * n_treatments - .5, linestyle='--', alpha=.4)
plt.xticks([t * n_treatments - .5 + n_treatments/2 for t in range(n_periods)],
           ["$\\theta_{}$".format(t) for t in range(n_periods)])
plt.gca().set_xlim([-.5, n_periods*n_treatments - .5])
plt.ylabel("Effect")
plt.legend()
plt.show()

# imports
from econml.data.dynamic_panel_dgp import SemiSynthetic
from sklearn.linear_model import LassoCV, MultiTaskLassoCV
import numpy as np
import matplotlib.pyplot as plt

# generate historical dataset (training purpose)
np.random.seed(43)
dgp = SemiSynthetic()
dgp.create_instance()
n_periods = 4    # Y包含了几个时期，默认为4个时间段
n_units = 5000   # 样本量
n_treatments = dgp.n_treatments  # 3 - 多干预
random_seed = 43
thetas = np.random.uniform(0, 2, size=(dgp.n_proxies, n_treatments))  # 误差项

panelX, panelT, panelY, panelGroups, true_effect = dgp.gen_data(
    n_units, n_periods, thetas, random_seed
)

# on historical data construct adjusted outcomes
from econml.dynamic.dml import DynamicDML

panelYadj = panelY.copy()

est = DynamicDML(
    model_y=LassoCV(max_iter=2000), model_t=MultiTaskLassoCV(max_iter=2000), cv=2
)
for t in range(1, n_periods):  # for each target period 1...m
    # learn period effect for each period treatment on target period t
    est.fit(
        long(panelY[:, 1 : t + 1]),
        long(panelT[:, 1 : t + 1, :]),  # reshape data to long format
        X=None,
        W=long(panelX[:, 1 : t + 1, :]),
        groups=long(panelGroups[:, 1 : t + 1]),
    )
    # print(f'y的shape {long(panelY[:, 1 : t + 1]).shape} , t 为 {t}， x为{long(panelX[:, 1 : t + 1, :]).shape}')
    # remove effect of observed treatments
    T1 = wide(panelT[:, 1 : t + 1, :])
    panelYadj[:, t] = panelY[:, t] - est.effect(
        T0=np.zeros_like(T1), T1=T1
    )  # reshape data to wide format

# 单独训练一个Y ~ X的回归模型进行预测
# train surrogate index on historical dataset
XS = np.hstack(
    [panelX[:, 1], panelYadj[:, :1]]
)  # concatenate controls and surrogates from historical dataset
# 有四期，第2个时期的X / Y 
# (5000, 71) + (5000, 1) => (5000, 72)

TotalYadj = np.sum(panelYadj, axis=1)  # total revenue from historical dataset  
adjusted_proxy_model = LassoCV().fit(
    XS, TotalYadj
)  # train proxy model from historical dataset


# predict new long term revenue
XSnew = np.hstack(
    [panelXnew[:, 1], panelYnew[:, :1]]
)  # concatenate controls and surrogates from new dataset
sindex_adj = adjusted_proxy_model.predict(XSnew)