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_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()