def test_drlearners(self): X = TestPandasIntegration.df[TestPandasIntegration.features] W = TestPandasIntegration.df[TestPandasIntegration.controls] Y = TestPandasIntegration.df[TestPandasIntegration.outcome] T = TestPandasIntegration.df[TestPandasIntegration.bin_treat] # Test LinearDRLearner est = LinearDRLearner(model_propensity=GradientBoostingClassifier(), model_regression=GradientBoostingRegressor()) est.fit(Y, T, X=X, W=W, inference='statsmodels') treatment_effects = est.effect(X) lb, ub = est.effect_interval(X, alpha=0.05) self._check_input_names(est.summary(T=1)) self._check_popsum_names(est.effect_inference(X).population_summary()) # Test SparseLinearDRLearner est = SparseLinearDRLearner( model_propensity=GradientBoostingClassifier(), model_regression=GradientBoostingRegressor()) est.fit(Y, T, X=X, W=W, inference='debiasedlasso') treatment_effects = est.effect(X) lb, ub = est.effect_interval(X, alpha=0.05) self._check_input_names(est.summary(T=1)) self._check_popsum_names(est.effect_inference(X).population_summary()) # Test ForestDRLearner est = ForestDRLearner(model_propensity=GradientBoostingClassifier(), model_regression=GradientBoostingRegressor()) est.fit(Y, T, X=X, W=W, inference='blb') treatment_effects = est.effect(X) lb, ub = est.effect_interval(X, alpha=0.05) self._check_popsum_names(est.effect_inference(X).population_summary())
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)