def test_dml(self):
#################################
# Single treatment and outcome #
#################################
X = TestPandasIntegration.df[TestPandasIntegration.features]
W = TestPandasIntegration.df[TestPandasIntegration.controls]
Y = TestPandasIntegration.df[TestPandasIntegration.outcome]
T = TestPandasIntegration.df[TestPandasIntegration.cont_treat]
# Test LinearDML
est = LinearDML(model_y=LassoCV(), model_t=LassoCV())
est.fit(Y, T, X=X, W=W, inference='statsmodels')
treatment_effects = est.effect(X)
lb, ub = est.effect_interval(X, alpha=0.05)
self._check_input_names(
est.summary()) # Check that names propagate as expected
# Test re-fit
X1 = X.rename(columns={c: "{}_1".format(c) for c in X.columns})
est.fit(Y, T, X=X1, W=W, inference='statsmodels')
self._check_input_names(est.summary(), feat_comp=X1.columns)
# Test SparseLinearDML
est = SparseLinearDML(model_y=LassoCV(), model_t=LassoCV())
est.fit(Y, T, X=X, W=W, inference='debiasedlasso')
treatment_effects = est.effect(X)
lb, ub = est.effect_interval(X, alpha=0.05)
self._check_input_names(
est.summary()) # Check that names propagate as expected
# ForestDML
est = ForestDML(model_y=GradientBoostingRegressor(),
model_t=GradientBoostingRegressor())
est.fit(Y, T, X=X, W=W, inference='blb')
treatment_effects = est.effect(X)
lb, ub = est.effect_interval(X, alpha=0.05)
####################################
# Mutiple treatments and outcomes #
####################################
Y = TestPandasIntegration.df[TestPandasIntegration.outcome_multi]
T = TestPandasIntegration.df[TestPandasIntegration.cont_treat_multi]
# Test LinearDML
est = LinearDML(model_y=MultiTaskLasso(), model_t=MultiTaskLasso())
est.fit(Y, T, X=X, W=W, inference='statsmodels')
self._check_input_names(est.summary(), True,
True) # Check that names propagate as expected
self._check_popsum_names(
est.effect_inference(X).population_summary(), True)
est.fit(Y, T, X=X, W=W,
inference='bootstrap') # Check bootstrap as well
self._check_input_names(est.summary(), True, True)
self._check_popsum_names(
est.effect_inference(X).population_summary(), True)
# Test SparseLinearDML
est = SparseLinearDML(model_y=MultiTaskLasso(),
model_t=MultiTaskLasso())
est.fit(Y, T, X=X, W=W, inference='debiasedlasso')
treatment_effects = est.effect(X)
lb, ub = est.effect_interval(X, alpha=0.05)
self._check_input_names(est.summary(), True,
True) # Check that names propagate as expected
self._check_popsum_names(
est.effect_inference(X).population_summary(), True)
def test_dml_random_state(self):
Y, T, X, W, X_test = TestRandomState._make_data(500, 2)
for est in [
NonParamDML(model_y=RandomForestRegressor(n_estimators=10,
max_depth=4,
random_state=123),
model_t=RandomForestClassifier(n_estimators=10,
max_depth=4,
random_state=123),
model_final=RandomForestRegressor(
max_depth=3,
n_estimators=10,
min_samples_leaf=100,
bootstrap=True,
random_state=123),
discrete_treatment=True,
n_splits=2,
random_state=123),
CausalForestDML(
model_y=RandomForestRegressor(n_estimators=10,
max_depth=4,
random_state=123),
model_t=RandomForestClassifier(n_estimators=10,
max_depth=4,
random_state=123),
n_estimators=8,
discrete_treatment=True,
cv=2,
random_state=123),
LinearDML(model_y=RandomForestRegressor(n_estimators=10,
max_depth=4,
random_state=123),
model_t=RandomForestClassifier(n_estimators=10,
max_depth=4,
random_state=123),
discrete_treatment=True,
n_splits=2,
random_state=123),
SparseLinearDML(discrete_treatment=True,
n_splits=2,
random_state=123),
KernelDML(discrete_treatment=True,
n_splits=2,
random_state=123)
]:
TestRandomState._test_random_state(est, X_test, Y, T, X=X, W=W)
def test_dml(self):
"""Test setting attributes and refitting"""
y, T, X, W = self._get_data()
dml = DML(model_y=LinearRegression(),
model_t=LinearRegression(),
model_final=StatsModelsLinearRegression(fit_intercept=False),
linear_first_stages=False,
random_state=123)
dml.fit(y, T, X=X, W=W)
with pytest.raises(Exception):
dml.refit_final()
dml.fit(y, T, X=X, W=W, cache_values=True)
dml.model_final = StatsModelsRLM(fit_intercept=False)
dml.refit_final()
assert isinstance(dml.model_cate, StatsModelsRLM)
np.testing.assert_array_equal(dml.model_cate.coef_[1:].flatten(), dml.coef_.flatten())
lb, ub = dml.model_cate.coef__interval(alpha=0.01)
lbt, ubt = dml.coef__interval(alpha=0.01)
np.testing.assert_array_equal(lb[1:].flatten(), lbt.flatten())
np.testing.assert_array_equal(ub[1:].flatten(), ubt.flatten())
intcpt = dml.intercept_
dml.fit_cate_intercept = False
np.testing.assert_equal(dml.intercept_, intcpt)
dml.refit_final()
np.testing.assert_array_equal(dml.model_cate.coef_.flatten(), dml.coef_.flatten())
lb, ub = dml.model_cate.coef__interval(alpha=0.01)
lbt, ubt = dml.coef__interval(alpha=0.01)
np.testing.assert_array_equal(lb.flatten(), lbt.flatten())
np.testing.assert_array_equal(ub.flatten(), ubt.flatten())
with pytest.raises(AttributeError):
dml.intercept_
with pytest.raises(AttributeError):
dml.intercept__interval()
dml.model_final = DebiasedLasso(fit_intercept=False)
dml.refit_final()
assert isinstance(dml.model_cate, DebiasedLasso)
dml.featurizer = PolynomialFeatures(degree=2, include_bias=False)
dml.model_final = StatsModelsLinearRegression(fit_intercept=False)
dml.refit_final()
assert isinstance(dml.featurizer_, PolynomialFeatures)
dml.fit_cate_intercept = True
dml.refit_final()
assert isinstance(dml.featurizer_, Pipeline)
np.testing.assert_array_equal(dml.coef_.shape, (X.shape[1]**2))
np.testing.assert_array_equal(dml.coef__interval()[0].shape, (X.shape[1]**2))
coefpre = dml.coef_
coefpreint = dml.coef__interval()
dml.fit(y, T, X=X, W=W)
np.testing.assert_array_equal(coefpre, dml.coef_)
np.testing.assert_array_equal(coefpreint[0], dml.coef__interval()[0])
dml.discrete_treatment = True
dml.featurizer = None
dml.linear_first_stages = True
dml.model_t = LogisticRegression()
dml.fit(y, T, X=X, W=W)
newdml = DML(model_y=LinearRegression(),
model_t=LogisticRegression(),
model_final=StatsModelsLinearRegression(fit_intercept=False),
discrete_treatment=True,
linear_first_stages=True,
random_state=123).fit(y, T, X=X, W=W)
np.testing.assert_array_equal(dml.coef_, newdml.coef_)
np.testing.assert_array_equal(dml.coef__interval()[0], newdml.coef__interval()[0])
ldml = LinearDML(model_y=LinearRegression(),
model_t=LinearRegression(),
linear_first_stages=False)
ldml.fit(y, T, X=X, W=W, cache_values=True)
# can set final model for plain DML, but can't for LinearDML (hardcoded to StatsModelsRegression)
with pytest.raises(ValueError):
ldml.model_final = StatsModelsRLM()
ldml = SparseLinearDML(model_y=LinearRegression(),
model_t=LinearRegression(),
linear_first_stages=False)
ldml.fit(y, T, X=X, W=W, cache_values=True)
# can set final model for plain DML, but can't for LinearDML (hardcoded to StatsModelsRegression)
with pytest.raises(ValueError):
ldml.model_final = StatsModelsRLM()
ldml.alpha = 0.01
ldml.max_iter = 10
ldml.tol = 0.01
ldml.refit_final()
np.testing.assert_equal(ldml.model_cate.estimators_[0].alpha, 0.01)
np.testing.assert_equal(ldml.model_cate.estimators_[0].max_iter, 10)
np.testing.assert_equal(ldml.model_cate.estimators_[0].tol, 0.01)
def test_comparison(self):
def reg():
return LinearRegression()
def clf():
return LogisticRegression()
y, T, X, true_eff = self._get_data()
(X_train, X_val, T_train, T_val, Y_train, Y_val, _,
true_eff_val) = train_test_split(X, T, y, true_eff, test_size=.4)
models = [
('ldml',
LinearDML(model_y=reg(),
model_t=clf(),
discrete_treatment=True,
linear_first_stages=False,
cv=3)),
('sldml',
SparseLinearDML(model_y=reg(),
model_t=clf(),
discrete_treatment=True,
featurizer=PolynomialFeatures(degree=2,
include_bias=False),
linear_first_stages=False,
cv=3)),
('xlearner',
XLearner(models=reg(), cate_models=reg(),
propensity_model=clf())),
('dalearner',
DomainAdaptationLearner(models=reg(),
final_models=reg(),
propensity_model=clf())),
('slearner', SLearner(overall_model=reg())),
('tlearner', TLearner(models=reg())),
('drlearner',
DRLearner(model_propensity=clf(),
model_regression=reg(),
model_final=reg(),
cv=3)),
('rlearner',
NonParamDML(model_y=reg(),
model_t=clf(),
model_final=reg(),
discrete_treatment=True,
cv=3)),
('dml3dlasso',
DML(model_y=reg(),
model_t=clf(),
model_final=reg(),
discrete_treatment=True,
featurizer=PolynomialFeatures(degree=3),
linear_first_stages=False,
cv=3))
]
models = Parallel(n_jobs=-1, verbose=1)(
delayed(_fit_model)(name, mdl, Y_train, T_train, X_train)
for name, mdl in models)
scorer = RScorer(model_y=reg(),
model_t=clf(),
discrete_treatment=True,
cv=3,
mc_iters=2,
mc_agg='median')
scorer.fit(Y_val, T_val, X=X_val)
rscore = [scorer.score(mdl) for _, mdl in models]
rootpehe_score = [
np.sqrt(
np.mean(
(true_eff_val.flatten() - mdl.effect(X_val).flatten())**2))
for _, mdl in models
]
assert LinearRegression().fit(
np.array(rscore).reshape(-1, 1),
np.array(rootpehe_score)).coef_ < 0.5
mdl, _ = scorer.best_model([mdl for _, mdl in models])
rootpehe_best = np.sqrt(
np.mean((true_eff_val.flatten() - mdl.effect(X_val).flatten())**2))
assert rootpehe_best < 1.2 * np.min(rootpehe_score)
mdl, _ = scorer.ensemble([mdl for _, mdl in models])
rootpehe_ensemble = np.sqrt(
np.mean((true_eff_val.flatten() - mdl.effect(X_val).flatten())**2))
assert rootpehe_ensemble < 1.2 * np.min(rootpehe_score)
def test_dml(self):
#################################
# Single treatment and outcome #
#################################
X = TestPandasIntegration.df[TestPandasIntegration.features]
W = TestPandasIntegration.df[TestPandasIntegration.controls]
Y = TestPandasIntegration.df[TestPandasIntegration.outcome]
T = TestPandasIntegration.df[TestPandasIntegration.cont_treat]
# Test LinearDML
est = LinearDML(model_y=LassoCV(), model_t=LassoCV())
est.fit(Y, T, X=X, W=W, inference='statsmodels')
treatment_effects = est.effect(X)
lb, ub = est.effect_interval(X, alpha=0.05)
self._check_input_names(
est.summary()) # Check that names propagate as expected
# |--> Test featurizers
est.featurizer = PolynomialFeatures(degree=2, include_bias=False)
est.fit(Y, T, X=X, W=W, inference='statsmodels')
self._check_input_names(
est.summary(),
feat_comp=est.original_featurizer.get_feature_names(X.columns))
est.featurizer = FunctionTransformer()
est.fit(Y, T, X=X, W=W, inference='statsmodels')
self._check_input_names(
est.summary(),
feat_comp=[
f"feat(X){i}" for i in range(TestPandasIntegration.n_features)
])
est.featurizer = ColumnTransformer([('passthrough', 'passthrough', [0])
])
est.fit(Y, T, X=X, W=W, inference='statsmodels')
# ColumnTransformer doesn't propagate column names
self._check_input_names(est.summary(), feat_comp=["x0"])
# |--> Test re-fit
est.featurizer = None
X1 = X.rename(columns={c: "{}_1".format(c) for c in X.columns})
est.fit(Y, T, X=X1, W=W, inference='statsmodels')
self._check_input_names(est.summary(), feat_comp=X1.columns)
# Test SparseLinearDML
est = SparseLinearDML(model_y=LassoCV(), model_t=LassoCV())
est.fit(Y, T, X=X, W=W, inference='debiasedlasso')
treatment_effects = est.effect(X)
lb, ub = est.effect_interval(X, alpha=0.05)
self._check_input_names(
est.summary()) # Check that names propagate as expected
# Test ForestDML
est = ForestDML(model_y=GradientBoostingRegressor(),
model_t=GradientBoostingRegressor())
est.fit(Y, T, X=X, W=W, inference='blb')
treatment_effects = est.effect(X)
lb, ub = est.effect_interval(X, alpha=0.05)
####################################
# Mutiple treatments and outcomes #
####################################
Y = TestPandasIntegration.df[TestPandasIntegration.outcome_multi]
T = TestPandasIntegration.df[TestPandasIntegration.cont_treat_multi]
# Test LinearDML
est = LinearDML(model_y=MultiTaskLasso(), model_t=MultiTaskLasso())
est.fit(Y, T, X=X, W=W, inference='statsmodels')
self._check_input_names(est.summary(), True,
True) # Check that names propagate as expected
self._check_popsum_names(
est.effect_inference(X).population_summary(), True)
est.fit(Y, T, X=X, W=W,
inference='bootstrap') # Check bootstrap as well
self._check_input_names(est.summary(), True, True)
self._check_popsum_names(
est.effect_inference(X).population_summary(), True)
# Test SparseLinearDML
est = SparseLinearDML(model_y=MultiTaskLasso(),
model_t=MultiTaskLasso())
est.fit(Y, T, X=X, W=W, inference='debiasedlasso')
treatment_effects = est.effect(X)
lb, ub = est.effect_interval(X, alpha=0.05)
self._check_input_names(est.summary(), True,
True) # Check that names propagate as expected
self._check_popsum_names(
est.effect_inference(X).population_summary(), True)