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