import
numpy
as
np
from
matplotlib
import
pyplot
as
plt
from
sklearn.linear_model
import
LinearRegression
from
sklearn.model_selection
import
train_test_split
import
statsmodels.api
as
sm
from
xgboost
import
XGBRegressor
import
warnings
from
causalml.inference.meta
import
LRSRegressor
from
causalml.inference.meta
import
XGBTRegressor, MLPTRegressor
from
causalml.inference.meta
import
BaseXRegressor, BaseRRegressor, BaseSRegressor, BaseTRegressor
from
causalml.
match
import
NearestNeighborMatch, MatchOptimizer, create_table_one
from
causalml.propensity
import
ElasticNetPropensityModel
from
causalml.dataset
import
*
from
causalml.metrics
import
*
# Generate synthetic data using mode 1
y, X, treatment, tau, b, e = synthetic_data(mode=
1
, n=
10000
, p=
8
, sigma=
1.0
)
s-learner、t-learner、x-learner、r-learner 预估
# Ready-to-use S-Learner using LinearRegressionlearner_s = LRSRegressor()
# Calling the Base Learner class and feeding in XGBlearner_t = BaseTRegressor(learner=XGBRegressor())
# Calling the Base Learner class and feeding in XGBlearner_x = BaseXRegressor(XGBRegressor())
# Calling the Base Learner class and feeding in XGBlearner_r = BaseRRegressor(learner=XGBRegressor())
# 估计ATE(以t-learner为例,其他均相同) ate_t = learner_t.estimate_ate(X, treatment, y)
# 估计ITE(以t-learner为例,其他均相同)# fit_predict先后执行fit -> predict两个过程cate_t = learner_t.fit_predict(X=X, treatment=treatment, y=y)