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