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_alpha(self):
Y, T, X, W = TestInference.Y, TestInference.T, TestInference.X, TestInference.W
est = LinearDML(model_y=LinearRegression(), model_t=LinearRegression())
est.fit(Y, T, X=X, W=W)
# ensure alpha is passed
lb, ub = est.const_marginal_ate_interval(X, alpha=1)
assert (lb == ub).all()
lb, ub = est.const_marginal_ate_interval(X)
assert (lb != ub).all()