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)