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 = LinearDML(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_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_stratify_orthoiv(self):
"""Test that we can properly stratify by treatment/instrument pair"""
T = [1, 0, 1, 1, 0, 0, 1, 0]
Z = [1, 0, 0, 1, 0, 1, 0, 1]
Y = [1, 2, 3, 4, 5, 6, 7, 8]
X = np.array([1, 1, 2, 2, 1, 2, 1, 2]).reshape(-1, 1)
est = LinearIntentToTreatDRIV(model_Y_X=LinearRegression(),
model_T_XZ=LogisticRegression(),
flexible_model_effect=LinearRegression(),
n_splits=2)
inference = BootstrapInference(n_bootstrap_samples=20)
est.fit(Y, T, Z=Z, X=X, inference=inference)
est.const_marginal_effect_interval(X)
def test_can_summarize(self):
LinearDMLCateEstimator().fit(TestInference.Y,
TestInference.T,
TestInference.X,
TestInference.W,
inference='statsmodels').summary()
LinearDRLearner(fit_cate_intercept=False).fit(
TestInference.Y,
TestInference.T > 0,
TestInference.X,
TestInference.W,
inference=BootstrapInference(5)).summary(1)
def test_can_summarize(self):
LinearDML(model_t=LinearRegression(),
model_y=LinearRegression()).fit(TestInference.Y,
TestInference.T,
TestInference.X,
TestInference.W).summary()
LinearDRLearner(model_regression=LinearRegression(),
model_propensity=LogisticRegression(),
fit_cate_intercept=False).fit(
TestInference.Y,
TestInference.T > 0,
TestInference.X,
TestInference.W,
inference=BootstrapInference(5)).summary(1)
def test_refit_final_inference(self):
"""Test that we can perform inference during refit_final"""
est = LinearDML(linear_first_stages=False, featurizer=PolynomialFeatures(1, include_bias=False))
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, 2))
W = np.random.normal(size=(500, 2))
est.fit(y, T, X=X, W=W, cache_values=True, inference='statsmodels')
assert isinstance(est.effect_inference(X), NormalInferenceResults)
with pytest.raises(ValueError):
est.refit_final(inference=BootstrapInference(2))
def test_translte(self):
Y, T, X, W = TestInference.Y, TestInference.T, TestInference.X, TestInference.W
for offset in [10, pd.Series(np.arange(TestInference.X.shape[0]))]:
for inf in ['auto', BootstrapInference(n_bootstrap_samples=5)]:
est = LinearDML().fit(Y, T, X=X, W=W, inference=inf)
inf = est.const_marginal_effect_inference(X)
pred, bounds, summary = inf.point_estimate, inf.conf_int(
), inf.summary_frame()
inf.translate(offset)
pred2, bounds2, summary2 = inf.point_estimate, inf.conf_int(
), inf.summary_frame()
np.testing.assert_array_equal(pred + offset, pred2)
np.testing.assert_array_almost_equal(bounds[0] + offset,
bounds2[0])
np.testing.assert_array_almost_equal(bounds[1] + offset,
bounds2[1])
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 = LinearDML(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()
def test_inference_results(self):
"""Tests the inference results summary."""
# Test inference results when `cate_feature_names` doesn not exist
for inference in [
BootstrapInference(n_bootstrap_samples=5), 'statsmodels'
]:
cate_est = LinearDMLCateEstimator(
featurizer=PolynomialFeatures(degree=1, include_bias=False))
wrapped_est = self._NoFeatNamesEst(cate_est)
wrapped_est.fit(TestInference.Y,
TestInference.T,
TestInference.X,
TestInference.W,
inference=inference)
summary_results = wrapped_est.summary()
coef_rows = np.asarray(summary_results.tables[0].data)[1:, 0]
np.testing.assert_array_equal(
coef_rows, ['X' + str(i) for i in range(TestInference.d_x)])
def test_internal_options(self):
"""Test that the internal use of bootstrap within an estimator using custom options works."""
x = np.random.normal(size=(1000, 2))
z = np.random.normal(size=(1000, 1))
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))
opts = BootstrapInference(50, 2)
est = NonparametricTwoStageLeastSquares(PolynomialFeatures(2),
PolynomialFeatures(2),
PolynomialFeatures(2), None)
est.fit(y, t, x, None, z, inference=opts)
# 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, t, 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.7
# test that we can do the same thing once we provide percentile bounds
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.65
def test_internal_options(self):
"""Test that the internal use of bootstrap within an estimator using custom options works."""
x = np.random.normal(size=(1000, 2))
z = np.random.normal(size=(1000, 1))
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))
opts = BootstrapInference(50, 2)
est = SieveTSLS(t_featurizer=PolynomialFeatures(2),
x_featurizer=PolynomialFeatures(2),
z_featurizer=PolynomialFeatures(2),
dt_featurizer=None)
est.fit(y, t, X=x, W=None, Z=z, inference=opts)
# 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()
# TODO: test that the estimated effect is usually within the bounds
# and that the true effect is also usually within the bounds
# test that we can do the same thing once we provide percentile bounds
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()
def test_cate_api(self):
"""Test that we correctly implement the CATE API."""
n = 20
def make_random(is_discrete, d):
if d is None:
return None
sz = (n, d) if d >= 0 else (n, )
if is_discrete:
while True:
arr = np.random.choice(['a', 'b', 'c'], size=sz)
# ensure that we've got at least two of every element
_, counts = np.unique(arr, return_counts=True)
if len(counts) == 3 and counts.min() > 1:
return arr
else:
return np.random.normal(size=sz)
for d_t in [2, 1, -1]:
for is_discrete in [True, False] if d_t <= 1 else [False]:
for d_y in [3, 1, -1]:
for d_x in [2, None]:
for d_w in [2, None]:
W, X, Y, T = [
make_random(is_discrete, d)
for is_discrete, d in [(
False,
d_w), (False,
d_x), (False,
d_y), (is_discrete, d_t)]
]
d_t_final = 2 if is_discrete else d_t
effect_shape = (n, ) + ((d_y, ) if d_y > 0 else ())
marginal_effect_shape = ((n, ) + (
(d_y, ) if d_y > 0 else
()) + ((d_t_final, ) if d_t_final > 0 else ()))
# since T isn't passed to const_marginal_effect, defaults to one row if X is None
const_marginal_effect_shape = (
(n if d_x else 1, ) + ((d_y, ) if d_y > 0 else
()) +
((d_t_final, ) if d_t_final > 0 else ()))
model_t = LogisticRegression(
) if is_discrete else Lasso()
# TODO: add stratification to bootstrap so that we can use it even with discrete treatments
all_infs = [None, 'statsmodels']
if not is_discrete:
all_infs.append(BootstrapInference(1))
for est, multi, infs in [
(LinearDMLCateEstimator(
model_y=Lasso(),
model_t='auto',
discrete_treatment=is_discrete), False,
all_infs),
(SparseLinearDMLCateEstimator(
model_y=LinearRegression(),
model_t=model_t,
discrete_treatment=is_discrete), True,
[None]),
(KernelDMLCateEstimator(
model_y=LinearRegression(),
model_t=model_t,
discrete_treatment=is_discrete), False,
[None])
]:
if not (multi) and d_y > 1:
continue
for inf in infs:
with self.subTest(d_w=d_w,
d_x=d_x,
d_y=d_y,
d_t=d_t,
is_discrete=is_discrete,
est=est,
inf=inf):
est.fit(Y, T, X, W, inference=inf)
# make sure we can call the marginal_effect and effect methods
const_marg_eff = est.const_marginal_effect(
X)
marg_eff = est.marginal_effect(T, X)
self.assertEqual(
shape(marg_eff),
marginal_effect_shape)
self.assertEqual(
shape(const_marg_eff),
const_marginal_effect_shape)
np.testing.assert_array_equal(
marg_eff if d_x else marg_eff[0:1],
const_marg_eff)
T0 = np.full_like(
T, 'a'
) if is_discrete else np.zeros_like(T)
eff = est.effect(X, T0=T0, T1=T)
self.assertEqual(
shape(eff), effect_shape)
if inf is not None:
const_marg_eff_int = est.const_marginal_effect_interval(
X)
marg_eff_int = est.marginal_effect_interval(
T, X)
self.assertEqual(
shape(marg_eff_int),
(2, ) + marginal_effect_shape)
self.assertEqual(
shape(const_marg_eff_int),
(2, ) +
const_marginal_effect_shape)
self.assertEqual(
shape(
est.effect_interval(X,
T0=T0,
T1=T)),
(2, ) + effect_shape)
est.score(Y, T, X, W)
# make sure we can call effect with implied scalar treatments, no matter the
# dimensions of T, and also that we warn when there are multiple treatments
if d_t > 1:
cm = self.assertWarns(Warning)
else:
cm = ExitStack(
) # ExitStack can be used as a "do nothing" ContextManager
with cm:
effect_shape2 = (
n if d_x else 1, ) + (
(d_y, ) if d_y > 0 else ())
eff = est.effect(
X
) if not is_discrete else est.effect(
X, T0='a', T1='b')
self.assertEqual(
shape(eff), effect_shape2)
def test_cate_api(self):
"""Test that we correctly implement the CATE API."""
n_panels = 100 # number of panels
n_periods = 3 # number of time periods per panel
n = n_panels * n_periods
groups = np.repeat(a=np.arange(n_panels), repeats=n_periods, axis=0)
def make_random(n, is_discrete, d):
if d is None:
return None
sz = (n, d) if d >= 0 else (n,)
if is_discrete:
return np.random.choice(['a', 'b', 'c'], size=sz)
else:
return np.random.normal(size=sz)
for d_t in [2, 1, -1]:
for is_discrete in [True, False] if d_t <= 1 else [False]:
# for is_discrete in [False]:
for d_y in [3, 1, -1]:
for d_x in [2, None]:
for d_w in [2, None]:
W, X, Y, T = [make_random(n, is_discrete, d)
for is_discrete, d in [(False, d_w),
(False, d_x),
(False, d_y),
(is_discrete, d_t)]]
T_test = np.hstack([(T.reshape(-1, 1) if d_t == -1 else T) for i in range(n_periods)])
for featurizer, fit_cate_intercept in\
[(None, True),
(PolynomialFeatures(degree=2, include_bias=False), True),
(PolynomialFeatures(degree=2, include_bias=True), False)]:
d_t_final = (2 if is_discrete else max(d_t, 1)) * n_periods
effect_shape = (n,) + ((d_y,) if d_y > 0 else ())
effect_summaryframe_shape = (n * (d_y if d_y > 0 else 1), 6)
marginal_effect_shape = ((n,) +
((d_y,) if d_y > 0 else ()) +
((d_t_final,) if d_t_final > 0 else ()))
marginal_effect_summaryframe_shape = (n * (d_y if d_y > 0 else 1) *
(d_t_final if d_t_final > 0 else 1), 6)
# since T isn't passed to const_marginal_effect, defaults to one row if X is None
const_marginal_effect_shape = ((n if d_x else 1,) +
((d_y,) if d_y > 0 else ()) +
((d_t_final,) if d_t_final > 0 else()))
const_marginal_effect_summaryframe_shape = (
(n if d_x else 1) * (d_y if d_y > 0 else 1) *
(d_t_final if d_t_final > 0 else 1), 6)
fd_x = featurizer.fit_transform(X).shape[1:] if featurizer and d_x\
else ((d_x,) if d_x else (0,))
coef_shape = Y.shape[1:] + (d_t_final, ) + fd_x
coef_summaryframe_shape = (
(d_y if d_y > 0 else 1) * (fd_x[0] if fd_x[0] >
0 else 1) * (d_t_final), 6)
intercept_shape = Y.shape[1:] + (d_t_final, )
intercept_summaryframe_shape = (
(d_y if d_y > 0 else 1) * (d_t_final if d_t_final > 0 else 1), 6)
all_infs = [None, 'auto', BootstrapInference(2)]
est = DynamicDML(model_y=Lasso() if d_y < 1 else MultiTaskLasso(),
model_t=LogisticRegression() if is_discrete else
(Lasso() if d_t < 1 else MultiTaskLasso()),
featurizer=featurizer,
fit_cate_intercept=fit_cate_intercept,
discrete_treatment=is_discrete)
# ensure we can serialize the unfit estimator
pickle.dumps(est)
for inf in all_infs:
with self.subTest(d_w=d_w, d_x=d_x, d_y=d_y, d_t=d_t,
is_discrete=is_discrete, est=est, inf=inf):
if X is None and (not fit_cate_intercept):
with pytest.raises(AttributeError):
est.fit(Y, T, X=X, W=W, groups=groups, inference=inf)
continue
est.fit(Y, T, X=X, W=W, groups=groups, inference=inf)
# ensure we can pickle the fit estimator
pickle.dumps(est)
# make sure we can call the marginal_effect and effect methods
const_marg_eff = est.const_marginal_effect(X)
marg_eff = est.marginal_effect(T_test, X)
self.assertEqual(shape(marg_eff), marginal_effect_shape)
self.assertEqual(shape(const_marg_eff), const_marginal_effect_shape)
np.testing.assert_allclose(
marg_eff if d_x else marg_eff[0:1], const_marg_eff)
assert len(est.score_) == n_periods
for score in est.nuisance_scores_y[0]:
assert score.shape == (n_periods, )
for score in est.nuisance_scores_t[0]:
assert score.shape == (n_periods, n_periods)
T0 = np.full_like(T_test, 'a') if is_discrete else np.zeros_like(T_test)
eff = est.effect(X, T0=T0, T1=T_test)
self.assertEqual(shape(eff), effect_shape)
self.assertEqual(shape(est.coef_), coef_shape)
if fit_cate_intercept:
self.assertEqual(shape(est.intercept_), intercept_shape)
else:
with pytest.raises(AttributeError):
self.assertEqual(shape(est.intercept_), intercept_shape)
if inf is not None:
const_marg_eff_int = est.const_marginal_effect_interval(X)
marg_eff_int = est.marginal_effect_interval(T_test, X)
self.assertEqual(shape(marg_eff_int),
(2,) + marginal_effect_shape)
self.assertEqual(shape(const_marg_eff_int),
(2,) + const_marginal_effect_shape)
self.assertEqual(shape(est.effect_interval(X, T0=T0, T1=T_test)),
(2,) + effect_shape)
self.assertEqual(shape(est.coef__interval()),
(2,) + coef_shape)
if fit_cate_intercept:
self.assertEqual(shape(est.intercept__interval()),
(2,) + intercept_shape)
else:
with pytest.raises(AttributeError):
self.assertEqual(shape(est.intercept__interval()),
(2,) + intercept_shape)
const_marg_effect_inf = est.const_marginal_effect_inference(X)
T1 = np.full_like(T_test, 'b') if is_discrete else T_test
effect_inf = est.effect_inference(X, T0=T0, T1=T1)
marg_effect_inf = est.marginal_effect_inference(T_test, X)
# test const marginal inference
self.assertEqual(shape(const_marg_effect_inf.summary_frame()),
const_marginal_effect_summaryframe_shape)
self.assertEqual(shape(const_marg_effect_inf.point_estimate),
const_marginal_effect_shape)
self.assertEqual(shape(const_marg_effect_inf.stderr),
const_marginal_effect_shape)
self.assertEqual(shape(const_marg_effect_inf.var),
const_marginal_effect_shape)
self.assertEqual(shape(const_marg_effect_inf.pvalue()),
const_marginal_effect_shape)
self.assertEqual(shape(const_marg_effect_inf.zstat()),
const_marginal_effect_shape)
self.assertEqual(shape(const_marg_effect_inf.conf_int()),
(2,) + const_marginal_effect_shape)
np.testing.assert_array_almost_equal(
const_marg_effect_inf.conf_int()[0],
const_marg_eff_int[0], decimal=5)
const_marg_effect_inf.population_summary()._repr_html_()
# test effect inference
self.assertEqual(shape(effect_inf.summary_frame()),
effect_summaryframe_shape)
self.assertEqual(shape(effect_inf.point_estimate),
effect_shape)
self.assertEqual(shape(effect_inf.stderr),
effect_shape)
self.assertEqual(shape(effect_inf.var),
effect_shape)
self.assertEqual(shape(effect_inf.pvalue()),
effect_shape)
self.assertEqual(shape(effect_inf.zstat()),
effect_shape)
self.assertEqual(shape(effect_inf.conf_int()),
(2,) + effect_shape)
np.testing.assert_array_almost_equal(
effect_inf.conf_int()[0],
est.effect_interval(X, T0=T0, T1=T1)[0], decimal=5)
effect_inf.population_summary()._repr_html_()
# test marginal effect inference
self.assertEqual(shape(marg_effect_inf.summary_frame()),
marginal_effect_summaryframe_shape)
self.assertEqual(shape(marg_effect_inf.point_estimate),
marginal_effect_shape)
self.assertEqual(shape(marg_effect_inf.stderr),
marginal_effect_shape)
self.assertEqual(shape(marg_effect_inf.var),
marginal_effect_shape)
self.assertEqual(shape(marg_effect_inf.pvalue()),
marginal_effect_shape)
self.assertEqual(shape(marg_effect_inf.zstat()),
marginal_effect_shape)
self.assertEqual(shape(marg_effect_inf.conf_int()),
(2,) + marginal_effect_shape)
np.testing.assert_array_almost_equal(
marg_effect_inf.conf_int()[0], marg_eff_int[0], decimal=5)
marg_effect_inf.population_summary()._repr_html_()
# test coef__inference and intercept__inference
if X is not None:
self.assertEqual(
shape(est.coef__inference().summary_frame()),
coef_summaryframe_shape)
np.testing.assert_array_almost_equal(
est.coef__inference().conf_int()
[0], est.coef__interval()[0], decimal=5)
if fit_cate_intercept:
cm = ExitStack()
# ExitStack can be used as a "do nothing" ContextManager
else:
cm = pytest.raises(AttributeError)
with cm:
self.assertEqual(shape(est.intercept__inference().
summary_frame()),
intercept_summaryframe_shape)
np.testing.assert_array_almost_equal(
est.intercept__inference().conf_int()
[0], est.intercept__interval()[0], decimal=5)
est.summary()
est.score(Y, T, X, W, groups=groups)
# make sure we can call effect with implied scalar treatments,
# no matter the dimensions of T, and also that we warn when there
# are multiple treatments
if d_t > 1:
cm = self.assertWarns(Warning)
else:
# ExitStack can be used as a "do nothing" ContextManager
cm = ExitStack()
with cm:
effect_shape2 = (n if d_x else 1,) + ((d_y,) if d_y > 0 else())
eff = est.effect(X) if not is_discrete else est.effect(
X, T0='a', T1='b')
self.assertEqual(shape(eff), effect_shape2)
def test_summary_discrete(self):
"""Tests the inference results summary for discrete treatment estimators."""
# Test inference results when `cate_feature_names` doesn not exist
for inference in [
BootstrapInference(n_bootstrap_samples=5), 'statsmodels'
]:
cate_est = LinearDRLearner(model_regression=LinearRegression(),
model_propensity=LogisticRegression(),
featurizer=PolynomialFeatures(
degree=2, include_bias=False))
cate_est.fit(TestInference.Y,
TestInference.T,
TestInference.X,
TestInference.W,
inference=inference)
summary_results = cate_est.summary(T=1)
coef_rows = np.asarray(summary_results.tables[0].data)[1:, 0]
fnames = PolynomialFeatures(degree=2, include_bias=False).fit(
TestInference.X).get_feature_names()
np.testing.assert_array_equal(coef_rows, fnames)
intercept_rows = np.asarray(summary_results.tables[1].data)[1:, 0]
np.testing.assert_array_equal(intercept_rows, ['intercept'])
cate_est = LinearDRLearner(model_regression=LinearRegression(),
model_propensity=LogisticRegression(),
featurizer=PolynomialFeatures(
degree=2, include_bias=False))
cate_est.fit(TestInference.Y,
TestInference.T,
TestInference.X,
TestInference.W,
inference=inference)
fnames = ['Q' + str(i) for i in range(TestInference.d_x)]
summary_results = cate_est.summary(T=1, feat_name=fnames)
coef_rows = np.asarray(summary_results.tables[0].data)[1:, 0]
fnames = PolynomialFeatures(degree=2, include_bias=False).fit(
TestInference.X).get_feature_names(input_features=fnames)
np.testing.assert_array_equal(coef_rows, fnames)
cate_est = LinearDRLearner(model_regression=LinearRegression(),
model_propensity=LogisticRegression(),
featurizer=None)
cate_est.fit(TestInference.Y,
TestInference.T,
TestInference.X,
TestInference.W,
inference=inference)
summary_results = cate_est.summary(T=1)
coef_rows = np.asarray(summary_results.tables[0].data)[1:, 0]
np.testing.assert_array_equal(
coef_rows, ['X' + str(i) for i in range(TestInference.d_x)])
cate_est = LinearDRLearner(model_regression=LinearRegression(),
model_propensity=LogisticRegression(),
featurizer=None)
cate_est.fit(TestInference.Y,
TestInference.T,
TestInference.X,
TestInference.W,
inference=inference)
fnames = ['Q' + str(i) for i in range(TestInference.d_x)]
summary_results = cate_est.summary(T=1, feat_name=fnames)
coef_rows = np.asarray(summary_results.tables[0].data)[1:, 0]
np.testing.assert_array_equal(coef_rows, fnames)
cate_est = LinearDRLearner(model_regression=LinearRegression(),
model_propensity=LogisticRegression(),
featurizer=None)
wrapped_est = self._NoFeatNamesEst(cate_est)
wrapped_est.fit(TestInference.Y,
TestInference.T,
TestInference.X,
TestInference.W,
inference=inference)
summary_results = wrapped_est.summary(T=1)
coef_rows = np.asarray(summary_results.tables[0].data)[1:, 0]
np.testing.assert_array_equal(
coef_rows, ['X' + str(i) for i in range(TestInference.d_x)])
cate_est = LinearDRLearner(model_regression=LinearRegression(),
model_propensity=LogisticRegression(),
featurizer=None)
wrapped_est = self._NoFeatNamesEst(cate_est)
wrapped_est.fit(TestInference.Y,
TestInference.T,
TestInference.X,
TestInference.W,
inference=inference)
fnames = ['Q' + str(i) for i in range(TestInference.d_x)]
summary_results = wrapped_est.summary(T=1, feat_name=fnames)
coef_rows = np.asarray(summary_results.tables[0].data)[1:, 0]
np.testing.assert_array_equal(coef_rows, fnames)
def test_cate_api(self):
"""Test that we correctly implement the CATE API."""
n = 30
def size(n, d):
return (n, d) if d >= 0 else (n, )
def make_random(is_discrete, d):
if d is None:
return None
sz = size(n, d)
if is_discrete:
while True:
arr = np.random.choice(['a', 'b', 'c'], size=sz)
# ensure that we've got at least two of every row
_, counts = np.unique(arr, return_counts=True, axis=0)
if len(counts) == 3**(d if d > 0 else
1) and counts.min() > 1:
return arr
else:
return np.random.normal(size=sz)
def eff_shape(n, d_y):
return (n, ) + ((d_y, ) if d_y > 0 else ())
def marg_eff_shape(n, d_y, d_t_final):
return ((n, ) + ((d_y, ) if d_y > 0 else
()) + ((d_t_final, ) if d_t_final > 0 else ()))
# since T isn't passed to const_marginal_effect, defaults to one row if X is None
def const_marg_eff_shape(n, d_x, d_y, d_t_final):
return ((n if d_x else 1, ) + ((d_y, ) if d_y > 0 else ()) +
((d_t_final, ) if d_t_final > 0 else ()))
for d_t in [2, 1, -1]:
n_t = d_t if d_t > 0 else 1
for discrete_t in [True, False] if n_t == 1 else [False]:
for d_y in [3, 1, -1]:
for d_q in [2, None]:
for d_z in [2, 1]:
if d_z < n_t:
continue
for discrete_z in [True, False
] if d_z == 1 else [False]:
Z1, Q, Y, T1 = [
make_random(is_discrete, d)
for is_discrete, d in [(
discrete_z,
d_z), (False,
d_q), (False,
d_y), (discrete_t, d_t)]
]
if discrete_t and discrete_z:
# need to make sure we get all *joint* combinations
arr = make_random(True, 2)
Z1 = arr[:, 0].reshape(size(n, d_z))
T1 = arr[:, 0].reshape(size(n, d_t))
d_t_final1 = 2 if discrete_t else d_t
if discrete_t:
# IntentToTreat only supports binary treatments/instruments
T2 = T1.copy()
T2[T1 == 'c'] = np.random.choice(
['a', 'b'],
size=np.count_nonzero(T1 == 'c'))
d_t_final2 = 1
if discrete_z:
# IntentToTreat only supports binary treatments/instruments
Z2 = Z1.copy()
Z2[Z1 == 'c'] = np.random.choice(
['a', 'b'],
size=np.count_nonzero(Z1 == 'c'))
effect_shape = eff_shape(n, d_y)
model_t = LogisticRegression(
) if discrete_t else Lasso()
model_z = LogisticRegression(
) if discrete_z else Lasso()
all_infs = [None, BootstrapInference(1)]
estimators = [
(DMLATEIV(model_Y_W=Lasso(),
model_T_W=model_t,
model_Z_W=model_z,
discrete_treatment=discrete_t,
discrete_instrument=discrete_z),
True, all_infs),
(ProjectedDMLATEIV(
model_Y_W=Lasso(),
model_T_W=model_t,
model_T_WZ=model_t,
discrete_treatment=discrete_t,
discrete_instrument=discrete_z), False,
all_infs),
(DMLIV(model_Y_X=Lasso(),
model_T_X=model_t,
model_T_XZ=model_t,
model_final=Lasso(),
discrete_treatment=discrete_t,
discrete_instrument=discrete_z),
False, all_infs)
]
if d_q and discrete_t and discrete_z:
# IntentToTreat requires X
estimators.append((LinearIntentToTreatDRIV(
model_Y_X=Lasso(),
model_T_XZ=model_t,
flexible_model_effect=WeightedLasso(),
cv=2), False, all_infs + ['auto']))
for est, multi, infs in estimators:
if not (
multi
) and d_y > 1 or d_t > 1 or d_z > 1:
continue
# ensure we can serialize unfit estimator
pickle.dumps(est)
d_ws = [None]
if isinstance(est,
LinearIntentToTreatDRIV):
d_ws.append(2)
for d_w in d_ws:
W = make_random(False, d_w)
for inf in infs:
with self.subTest(
d_z=d_z,
d_x=d_q,
d_y=d_y,
d_t=d_t,
discrete_t=discrete_t,
discrete_z=discrete_z,
est=est,
inf=inf):
Z = Z1
T = T1
d_t_final = d_t_final1
X = Q
d_x = d_q
if isinstance(
est,
(DMLATEIV,
ProjectedDMLATEIV)):
# these support only W but not X
W = Q
X = None
d_x = None
def fit():
return est.fit(
Y,
T,
Z=Z,
W=W,
inference=inf)
def score():
return est.score(Y,
T,
Z=Z,
W=W)
else:
# these support only binary, not general discrete T and Z
if discrete_t:
T = T2
d_t_final = d_t_final2
if discrete_z:
Z = Z2
if isinstance(
est,
LinearIntentToTreatDRIV
):
def fit():
return est.fit(
Y,
T,
Z=Z,
X=X,
W=W,
inference=inf)
def score():
return est.score(
Y,
T,
Z=Z,
X=X,
W=W)
else:
def fit():
return est.fit(
Y,
T,
Z=Z,
X=X,
inference=inf)
def score():
return est.score(
Y, T, Z=Z, X=X)
marginal_effect_shape = marg_eff_shape(
n, d_y, d_t_final)
const_marginal_effect_shape = const_marg_eff_shape(
n, d_x, d_y, d_t_final)
fit()
# ensure we can serialize fit estimator
pickle.dumps(est)
# make sure we can call the marginal_effect and effect methods
const_marg_eff = est.const_marginal_effect(
X)
marg_eff = est.marginal_effect(
T, X)
self.assertEqual(
shape(marg_eff),
marginal_effect_shape)
self.assertEqual(
shape(const_marg_eff),
const_marginal_effect_shape
)
np.testing.assert_array_equal(
marg_eff
if d_x else marg_eff[0:1],
const_marg_eff)
T0 = np.full_like(
T, 'a'
) if discrete_t else np.zeros_like(
T)
eff = est.effect(X,
T0=T0,
T1=T)
self.assertEqual(
shape(eff), effect_shape)
# TODO: add tests for extra properties like coef_ where they exist
if inf is not None:
const_marg_eff_int = est.const_marginal_effect_interval(
X)
marg_eff_int = est.marginal_effect_interval(
T, X)
self.assertEqual(
shape(marg_eff_int),
(2, ) +
marginal_effect_shape)
self.assertEqual(
shape(
const_marg_eff_int
), (2, ) +
const_marginal_effect_shape
)
self.assertEqual(
shape(
est.
effect_interval(
X, T0=T0,
T1=T)),
(2, ) + effect_shape)
# TODO: add tests for extra properties like coef_ where they exist
score()
# make sure we can call effect with implied scalar treatments,
# no matter the dimensions of T, and also that we warn when there
# are multiple treatments
if d_t > 1:
cm = self.assertWarns(
Warning)
else:
# ExitStack can be used as a "do nothing" ContextManager
cm = ExitStack()
with cm:
effect_shape2 = (
n if d_x else 1, ) + (
(d_y, )
if d_y > 0 else ())
eff = est.effect(
X
) if not discrete_t else est.effect(
X, T0='a', T1='b')
self.assertEqual(
shape(eff),
effect_shape2)
def test_summary(self):
"""Tests the inference results summary for continuous treatment estimators."""
# Test inference results when `cate_feature_names` doesn not exist
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(
TestInference.Y,
TestInference.T,
TestInference.X,
TestInference.W,
inference=inference
)
summary_results = cate_est.summary()
coef_rows = np.asarray(summary_results.tables[0].data)[1:, 0]
default_names = get_input_columns(TestInference.X)
fnames = PolynomialFeatures(degree=2, include_bias=False).fit(
TestInference.X).get_feature_names(default_names)
np.testing.assert_array_equal(coef_rows, fnames)
intercept_rows = np.asarray(summary_results.tables[1].data)[1:, 0]
np.testing.assert_array_equal(intercept_rows, ['cate_intercept'])
cate_est = LinearDML(model_t=LinearRegression(), model_y=LinearRegression(),
featurizer=PolynomialFeatures(degree=2,
include_bias=False)
)
cate_est.fit(
TestInference.Y,
TestInference.T,
TestInference.X,
TestInference.W,
inference=inference
)
fnames = ['Q' + str(i) for i in range(TestInference.d_x)]
summary_results = cate_est.summary(feature_names=fnames)
coef_rows = np.asarray(summary_results.tables[0].data)[1:, 0]
fnames = PolynomialFeatures(degree=2, include_bias=False).fit(
TestInference.X).get_feature_names(input_features=fnames)
np.testing.assert_array_equal(coef_rows, fnames)
cate_est = LinearDML(model_t=LinearRegression(), model_y=LinearRegression(), featurizer=None)
cate_est.fit(
TestInference.Y,
TestInference.T,
TestInference.X,
TestInference.W,
inference=inference
)
summary_results = cate_est.summary()
coef_rows = np.asarray(summary_results.tables[0].data)[1:, 0]
np.testing.assert_array_equal(coef_rows, ['X' + str(i) for i in range(TestInference.d_x)])
cate_est = LinearDML(model_t=LinearRegression(), model_y=LinearRegression(), featurizer=None)
cate_est.fit(
TestInference.Y,
TestInference.T,
TestInference.X,
TestInference.W,
inference=inference
)
fnames = ['Q' + str(i) for i in range(TestInference.d_x)]
summary_results = cate_est.summary(feature_names=fnames)
coef_rows = np.asarray(summary_results.tables[0].data)[1:, 0]
np.testing.assert_array_equal(coef_rows, fnames)
cate_est = LinearDML(model_t=LinearRegression(), model_y=LinearRegression(), featurizer=None)
wrapped_est = self._NoFeatNamesEst(cate_est)
wrapped_est.fit(
TestInference.Y,
TestInference.T,
TestInference.X,
TestInference.W,
inference=inference
)
summary_results = wrapped_est.summary()
coef_rows = np.asarray(summary_results.tables[0].data)[1:, 0]
np.testing.assert_array_equal(coef_rows, ['X' + str(i) for i in range(TestInference.d_x)])
cate_est = LinearDML(model_t=LinearRegression(), model_y=LinearRegression(), featurizer=None)
wrapped_est = self._NoFeatNamesEst(cate_est)
wrapped_est.fit(
TestInference.Y,
TestInference.T,
TestInference.X,
TestInference.W,
inference=inference
)
fnames = ['Q' + str(i) for i in range(TestInference.d_x)]
summary_results = wrapped_est.summary(feature_names=fnames)
coef_rows = np.asarray(summary_results.tables[0].data)[1:, 0]
np.testing.assert_array_equal(coef_rows, fnames)
def test_cate_api_nonparam(self):
"""Test that we correctly implement the CATE API."""
n = 20
def make_random(is_discrete, d):
if d is None:
return None
sz = (n, d) if d >= 0 else (n, )
if is_discrete:
while True:
arr = np.random.choice(['a', 'b'], size=sz)
# ensure that we've got at least two of every element
_, counts = np.unique(arr, return_counts=True)
if len(counts) == 2 and counts.min() > 2:
return arr
else:
return np.random.normal(size=sz)
for d_t in [1, -1]:
for is_discrete in [True, False] if d_t <= 1 else [False]:
for d_y in [3, 1, -1]:
for d_x in [2, None]:
for d_w in [2, None]:
W, X, Y, T = [
make_random(is_discrete, d)
for is_discrete, d in [(
False,
d_w), (False,
d_x), (False,
d_y), (is_discrete, d_t)]
]
d_t_final = 1 if is_discrete else d_t
effect_shape = (n, ) + ((d_y, ) if d_y > 0 else ())
marginal_effect_shape = ((n, ) + (
(d_y, ) if d_y > 0 else
()) + ((d_t_final, ) if d_t_final > 0 else ()))
# since T isn't passed to const_marginal_effect, defaults to one row if X is None
const_marginal_effect_shape = (
(n if d_x else 1, ) + ((d_y, ) if d_y > 0 else
()) +
((d_t_final, ) if d_t_final > 0 else ()))
model_t = LogisticRegression(
) if is_discrete else WeightedLasso()
# TODO Add bootstrap inference, once discrete treatment issue is fixed
base_infs = [None]
if not is_discrete:
base_infs += [BootstrapInference(2)]
for est, multi, infs in [
(NonParamDMLCateEstimator(
model_y=WeightedLasso(),
model_t=model_t,
model_final=WeightedLasso(),
featurizer=None,
discrete_treatment=is_discrete), True,
base_infs),
(NonParamDMLCateEstimator(
model_y=WeightedLasso(),
model_t=model_t,
model_final=WeightedLasso(),
featurizer=FunctionTransformer(),
discrete_treatment=is_discrete), True,
base_infs),
(ForestDMLCateEstimator(
model_y=WeightedLasso(),
model_t=model_t,
discrete_treatment=is_discrete), True,
base_infs + ['blb'])
]:
if not (multi) and d_y > 1:
continue
for inf in infs:
with self.subTest(d_w=d_w,
d_x=d_x,
d_y=d_y,
d_t=d_t,
is_discrete=is_discrete,
est=est,
inf=inf):
if X is None:
with pytest.raises(AttributeError):
est.fit(Y,
T,
X,
W,
inference=inf)
continue
est.fit(Y, T, X, W, inference=inf)
# make sure we can call the marginal_effect and effect methods
const_marg_eff = est.const_marginal_effect(
X)
marg_eff = est.marginal_effect(T, X)
self.assertEqual(
shape(marg_eff),
marginal_effect_shape)
self.assertEqual(
shape(const_marg_eff),
const_marginal_effect_shape)
np.testing.assert_array_equal(
marg_eff if d_x else marg_eff[0:1],
const_marg_eff)
T0 = np.full_like(
T, 'a'
) if is_discrete else np.zeros_like(T)
eff = est.effect(X, T0=T0, T1=T)
self.assertEqual(
shape(eff), effect_shape)
if inf is not None:
const_marg_eff_int = est.const_marginal_effect_interval(
X)
marg_eff_int = est.marginal_effect_interval(
T, X)
self.assertEqual(
shape(marg_eff_int),
(2, ) + marginal_effect_shape)
self.assertEqual(
shape(const_marg_eff_int),
(2, ) +
const_marginal_effect_shape)
self.assertEqual(
shape(
est.effect_interval(X,
T0=T0,
T1=T)),
(2, ) + effect_shape)
est.score(Y, T, X, W)
# make sure we can call effect with implied scalar treatments, no matter the
# dimensions of T, and also that we warn when there are multiple treatments
if d_t > 1:
cm = self.assertWarns(Warning)
else:
cm = ExitStack(
) # ExitStack can be used as a "do nothing" ContextManager
with cm:
effect_shape2 = (
n if d_x else 1, ) + (
(d_y, ) if d_y > 0 else ())
eff = est.effect(
X
) if not is_discrete else est.effect(
X, T0='a', T1='b')
self.assertEqual(
shape(eff), effect_shape2)
