def test_all_kinds(self):
T = [1, 0, 1, 2, 0, 2] * 5
Y = [1, 2, 3, 4, 5, 6] * 5
X = np.array([1, 1, 2, 2, 1, 2] * 5).reshape(-1, 1)
est = LinearDMLCateEstimator(n_splits=2)
for kind in ['percentile', 'pivot', 'normal']:
with self.subTest(kind=kind):
inference = BootstrapInference(n_bootstrap_samples=5,
bootstrap_type=kind)
est.fit(Y, T, inference=inference)
i = est.const_marginal_effect_interval()
inf = est.const_marginal_effect_inference()
assert i[0].shape == i[1].shape == inf.point_estimate.shape
assert np.allclose(i[0], inf.conf_int()[0])
assert np.allclose(i[1], inf.conf_int()[1])
est.fit(Y, T, X=X, inference=inference)
i = est.const_marginal_effect_interval(X)
inf = est.const_marginal_effect_inference(X)
assert i[0].shape == i[1].shape == inf.point_estimate.shape
assert np.allclose(i[0], inf.conf_int()[0])
assert np.allclose(i[1], inf.conf_int()[1])
i = est.coef__interval()
inf = est.coef__inference()
assert i[0].shape == i[1].shape == inf.point_estimate.shape
assert np.allclose(i[0], inf.conf_int()[0])
assert np.allclose(i[1], inf.conf_int()[1])
i = est.effect_interval(X)
inf = est.effect_inference(X)
assert i[0].shape == i[1].shape == inf.point_estimate.shape
assert np.allclose(i[0], inf.conf_int()[0])
assert np.allclose(i[1], inf.conf_int()[1])
def test_can_use_statsmodel_inference(self):
"""Test that we can use statsmodels to generate confidence intervals"""
dml = LinearDMLCateEstimator(LinearRegression(),
LogisticRegression(C=1000),
discrete_treatment=True)
dml.fit(np.array([2, 3, 1, 3, 2, 1, 1, 1]),
np.array([3, 2, 1, 2, 3, 1, 1, 1]),
np.ones((8, 1)),
inference='statsmodels')
interval = dml.effect_interval(np.ones((9, 1)),
T0=np.array([1, 1, 1, 2, 2, 2, 3, 3,
3]),
T1=np.array([1, 2, 3, 1, 2, 3, 1, 2,
3]),
alpha=0.05)
point = dml.effect(np.ones((9, 1)),
T0=np.array([1, 1, 1, 2, 2, 2, 3, 3, 3]),
T1=np.array([1, 2, 3, 1, 2, 3, 1, 2, 3]))
assert len(interval) == 2
lo, hi = interval
assert lo.shape == hi.shape == point.shape
assert (lo <= point).all()
assert (point <= hi).all()
assert (lo < hi).any(
) # for at least some of the examples, the CI should have nonzero width
interval = dml.const_marginal_effect_interval(np.ones((9, 1)),
alpha=0.05)
point = dml.const_marginal_effect(np.ones((9, 1)))
assert len(interval) == 2
lo, hi = interval
assert lo.shape == hi.shape == point.shape
assert (lo <= point).all()
assert (point <= hi).all()
assert (lo < hi).any(
) # for at least some of the examples, the CI should have nonzero width
interval = dml.coef__interval(alpha=0.05)
point = dml.coef_
assert len(interval) == 2
lo, hi = interval
assert lo.shape == hi.shape == point.shape
assert (lo <= point).all()
assert (point <= hi).all()
assert (lo < hi).any(
) # for at least some of the examples, the CI should have nonzero width
interval = dml.intercept__interval(alpha=0.05)
point = dml.intercept_
assert len(interval) == 2
lo, hi = interval
assert (lo <= point).all()
assert (point <= hi).all()
assert (lo < hi).any(
) # for at least some of the examples, the CI should have nonzero width
def test_stratify(self):
"""Test that we can properly stratify by treatment"""
T = [1, 0, 1, 2, 0, 2]
Y = [1, 2, 3, 4, 5, 6]
X = np.array([1, 1, 2, 2, 1, 2]).reshape(-1, 1)
est = LinearDMLCateEstimator(model_y=LinearRegression(),
model_t=LogisticRegression(),
discrete_treatment=True)
inference = BootstrapInference(n_bootstrap_samples=5)
est.fit(Y, T, inference=inference)
est.const_marginal_effect_interval()
est.fit(Y, T, X=X, inference=inference)
est.const_marginal_effect_interval(X)
est.fit(Y, np.asarray(T).reshape(-1, 1),
inference=inference) # test stratifying 2D treatment
est.const_marginal_effect_interval()
