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_with_econml(self):
"""Test that we can bootstrap econml estimators."""
x = np.random.normal(size=(1000, 2))
t = np.random.normal(size=(1000, 1))
t2 = np.random.normal(size=(1000, 1))
y = x[:, 0:1] * 0.5 + t + np.random.normal(size=(1000, 1))
est = LinearDML(model_y=LinearRegression(), model_t=LinearRegression())
est.fit(y, t, X=x)
bs = BootstrapEstimator(est, 50)
# test that we can fit with the same arguments as the base estimator
bs.fit(y, t, X=x)
# test that we can get the same attribute for the bootstrap as the original, with the same shape
self.assertEqual(np.shape(est.coef_), np.shape(bs.coef_))
# test that we can get an interval for the same attribute for the bootstrap as the original,
# with the same shape for the lower and upper bounds
lower, upper = bs.coef__interval()
for bound in [lower, upper]:
self.assertEqual(np.shape(est.coef_), np.shape(bound))
# test that the lower and upper bounds differ
assert (lower <= upper).all()
assert (lower < upper).any()
# test that we can do the same thing once we provide percentile bounds
lower, upper = bs.coef__interval(lower=10, upper=90)
for bound in [lower, upper]:
self.assertEqual(np.shape(est.coef_), np.shape(bound))
# test that we can do the same thing with the results of a method, rather than an attribute
self.assertEqual(np.shape(est.effect(x, T0=t, T1=t2)),
np.shape(bs.effect(x, T0=t, T1=t2)))
# test that we can get an interval for the same attribute for the bootstrap as the original,
# with the same shape for the lower and upper bounds
lower, upper = bs.effect_interval(x, T0=t, T1=t2)
for bound in [lower, upper]:
self.assertEqual(np.shape(est.effect(x, T0=t, T1=t2)),
np.shape(bound))
# test that the lower and upper bounds differ
assert (lower <= upper).all()
assert (lower < upper).any()
# test that we can do the same thing once we provide percentile bounds
lower, upper = bs.effect_interval(x, T0=t, T1=t2, lower=10, upper=90)
for bound in [lower, upper]:
self.assertEqual(np.shape(est.effect(x, T0=t, T1=t2)),
np.shape(bound))
# test that the lower and upper bounds differ
assert (lower <= upper).all()
assert (lower < upper).any()
def test_can_set_discrete_treatment(self):
X = np.random.choice(np.arange(5), size=(500, 3))
y = np.random.normal(size=(500,))
T = np.random.choice(np.arange(3), size=(500, 1))
W = np.random.normal(size=(500, 2))
est = LinearDML(model_y=RandomForestRegressor(),
model_t=RandomForestClassifier(min_samples_leaf=10),
discrete_treatment=True,
linear_first_stages=False,
cv=3)
est.fit(y, T, X=X, W=W)
est.effect(X)
est.discrete_treatment = False
est.fit(y, T, X=X, W=W)
est.effect(X)
def test_ate_inference(self):
"""Tests the ate inference results."""
Y, T, X, W = TestATEInference.Y, TestATEInference.T, TestATEInference.X, TestATEInference.W
for inference in [BootstrapInference(n_bootstrap_samples=5), 'auto']:
cate_est = LinearDML(model_t=LinearRegression(),
model_y=LinearRegression(),
featurizer=PolynomialFeatures(
degree=2, include_bias=False))
cate_est.fit(Y, T, X=X, W=W, inference=inference)
cate_est.ate(X)
cate_est.ate_inference(X)
cate_est.ate_interval(X, alpha=.01)
lb, _ = cate_est.ate_inference(X).conf_int_mean()
np.testing.assert_array_equal(lb.shape, Y.shape[1:])
cate_est.marginal_ate(T, X)
cate_est.marginal_ate_interval(T, X, alpha=.01)
cate_est.marginal_ate_inference(T, X)
lb, _ = cate_est.marginal_ate_inference(T, X).conf_int_mean()
np.testing.assert_array_equal(lb.shape, Y.shape[1:] + T.shape[1:])
cate_est.const_marginal_ate(X)
cate_est.const_marginal_ate_interval(X, alpha=.01)
cate_est.const_marginal_ate_inference(X)
lb, _ = cate_est.const_marginal_ate_inference(X).conf_int_mean()
np.testing.assert_array_equal(lb.shape, Y.shape[1:] + T.shape[1:])
summary = cate_est.ate_inference(X).summary(value=10)
for i in range(Y.shape[1]):
assert summary.tables[0].data[1 + i][4] < 1e-5
summary = cate_est.ate_inference(X).summary(
value=np.mean(cate_est.effect(X), axis=0))
for i in range(Y.shape[1]):
np.testing.assert_almost_equal(
summary.tables[0].data[1 + i][4], 1.0)
summary = cate_est.marginal_ate_inference(T, X).summary(value=10)
for i in range(Y.shape[1] * T.shape[1]):
assert summary.tables[0].data[1 + i][4] < 1e-5
summary = cate_est.marginal_ate_inference(T, X).summary(
value=np.mean(cate_est.marginal_effect(T, X), axis=0))
for i in range(Y.shape[1] * T.shape[1]):
np.testing.assert_almost_equal(
summary.tables[0].data[1 + i][4], 1.0)
summary = cate_est.const_marginal_ate_inference(X).summary(
value=10)
for i in range(Y.shape[1] * T.shape[1]):
assert summary.tables[0].data[1 + i][4] < 1e-5
summary = cate_est.const_marginal_ate_inference(X).summary(
value=np.mean(cate_est.const_marginal_effect(X), axis=0))
for i in range(Y.shape[1] * T.shape[1]):
np.testing.assert_almost_equal(
summary.tables[0].data[1 + i][4], 1.0)
def test_internal(self):
"""Test that the internal use of bootstrap within an estimator works."""
x = np.random.normal(size=(1000, 2))
t = np.random.normal(size=(1000, 1))
t2 = np.random.normal(size=(1000, 1))
y = x[:, 0:1] * 0.5 + t + np.random.normal(size=(1000, 1))
est = LinearDML(model_y=LinearRegression(), model_t=LinearRegression())
est.fit(y, t, X=x, inference='bootstrap')
# test that we can get an interval for the same attribute for the bootstrap as the original,
# with the same shape for the lower and upper bounds
eff = est.effect(x, T0=t, T1=t2)
lower, upper = est.effect_interval(x, T0=t, T1=t2)
for bound in [lower, upper]:
self.assertEqual(np.shape(eff), np.shape(bound))
# test that the lower and upper bounds differ
assert (lower <= upper).all()
assert (lower < upper).any()
# test that the estimated effect is usually within the bounds
assert np.mean(np.logical_and(lower <= eff, eff <= upper)) >= 0.9
# test that we can do the same thing once we provide alpha explicitly
lower, upper = est.effect_interval(x, T0=t, T1=t2, alpha=0.2)
for bound in [lower, upper]:
self.assertEqual(np.shape(eff), np.shape(bound))
# test that the lower and upper bounds differ
assert (lower <= upper).all()
assert (lower < upper).any()
# test that the estimated effect is usually within the bounds
assert np.mean(np.logical_and(lower <= eff, eff <= upper)) >= 0.8
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)