def test_nonparam_dml(self): y, T, X, W = self._get_data() dml = NonParamDML(model_y=LinearRegression(), model_t=LinearRegression(), model_final=WeightedLasso(), random_state=123) dml.fit(y, T, X=X, W=W) with pytest.raises(Exception): dml.refit_final() dml.fit(y, T, X=X, W=W, cache_values=True) dml.model_final = DebiasedLasso(fit_intercept=False) dml.refit_final() assert isinstance(dml.model_cate, DebiasedLasso) dml.effect_interval(X[:1]) dml.featurizer = PolynomialFeatures(degree=2, include_bias=False) dml.refit_final() assert isinstance(dml.featurizer_, PolynomialFeatures) dml.effect_interval(X[:1]) dml.discrete_treatment = True dml.featurizer = None dml.linear_first_stages = True dml.model_t = LogisticRegression() dml.model_final = DebiasedLasso() dml.fit(y, T, X=X, W=W) newdml = NonParamDML(model_y=LinearRegression(), model_t=LogisticRegression(), model_final=DebiasedLasso(), discrete_treatment=True, random_state=123).fit(y, T, X=X, W=W) np.testing.assert_array_equal(dml.effect(X[:1]), newdml.effect(X[:1])) np.testing.assert_array_equal(dml.effect_interval(X[:1])[0], newdml.effect_interval(X[:1])[0])
def test_access_to_internal_models(self): """ Test that API related to accessing the nuisance models, cate_model and featurizer is working. """ from econml.dml import DMLCateEstimator Y = np.array([2, 3, 1, 3, 2, 1, 1, 1]) T = np.array([3, 2, 1, 2, 1, 2, 1, 3]) X = np.ones((8, 1)) est = DMLCateEstimator(model_y=WeightedLasso(), model_t=LogisticRegression(), model_final=WeightedLasso(), featurizer=PolynomialFeatures( degree=2, include_bias=False), fit_cate_intercept=True, discrete_treatment=True) est.fit(Y, T, X) assert isinstance(est.original_featurizer, PolynomialFeatures) assert isinstance(est.featurizer, Pipeline) assert isinstance(est.model_cate, WeightedLasso) for mdl in est.models_y: assert isinstance(mdl, WeightedLasso) for mdl in est.models_t: assert isinstance(mdl, LogisticRegression) np.testing.assert_array_equal(est.cate_feature_names(['A']), ['A', 'A^2']) np.testing.assert_array_equal(est.cate_feature_names(), ['x0', 'x0^2']) est = DMLCateEstimator(model_y=WeightedLasso(), model_t=LogisticRegression(), model_final=WeightedLasso(), featurizer=None, fit_cate_intercept=True, discrete_treatment=True) est.fit(Y, T, X) assert est.original_featurizer is None assert isinstance(est.featurizer, FunctionTransformer) assert isinstance(est.model_cate, WeightedLasso) for mdl in est.models_y: assert isinstance(mdl, WeightedLasso) for mdl in est.models_t: assert isinstance(mdl, LogisticRegression) np.testing.assert_array_equal(est.cate_feature_names(['A']), ['A'])
def test_multidim_arrays_fail(self): Y = np.array([2, 3, 1, 3, 2, 1, 1, 1]) three_class = np.array([1, 2, 3, 1, 2, 3, 1, 2]) two_class = np.array([1, 2, 1, 1, 2, 1, 1, 2]) est = NonParamDMLIV(model_Y_X=Lasso(), model_T_X=LogisticRegression(), model_T_XZ=LogisticRegression(), model_final=WeightedLasso(), discrete_treatment=True) with pytest.raises(AttributeError): est.fit(Y, T=three_class, Z=two_class) est = IntentToTreatDRIV(model_Y_X=Lasso(), model_T_XZ=LogisticRegression(), flexible_model_effect=WeightedLasso()) with pytest.raises(AttributeError): est.fit(Y, T=three_class, Z=two_class) with pytest.raises(AttributeError): est.fit(Y, T=two_class, Z=three_class)
def test_bad_treatment_nonparam(self): """ Test that the non-parametric dml raises errors when treatment is not binary or single dimensional """ Y = np.array([2, 3, 1, 3, 2, 1, 1, 1]) T = np.array([3, 2, 1, 2, 1, 2, 1, 3]) X = np.ones((8, 1)) est = NonParamDMLCateEstimator(model_y=WeightedLasso(), model_t=LogisticRegression(), model_final=WeightedLasso(), discrete_treatment=True) with pytest.raises(AttributeError): est.fit(Y, T, X) T = np.ones((8, 2)) est = NonParamDMLCateEstimator(model_y=WeightedLasso(), model_t=LinearRegression(), model_final=WeightedLasso(), discrete_treatment=False) with pytest.raises(AttributeError): est.fit(Y, T, X)
def _compare_with_lasso(self, lasso_X, lasso_y, wlasso_X, wlasso_y, sample_weight, alpha_range=[0.01], params={}): for alpha in alpha_range: lasso = Lasso(alpha=alpha) lasso.set_params(**params) lasso.fit(lasso_X, lasso_y) wlasso = WeightedLasso(alpha=alpha) wlasso.set_params(**params) wlasso.fit(wlasso_X, wlasso_y, sample_weight=sample_weight) # Check results are similar with tolerance 1e-6 if np.ndim(lasso_y) > 1: for i in range(lasso_y.shape[1]): np.testing.assert_allclose(lasso.coef_[i], wlasso.coef_[i]) if lasso.get_params()["fit_intercept"]: self.assertAlmostEqual(lasso.intercept_[i], wlasso.intercept_[i]) else: np.testing.assert_allclose(lasso.coef_, wlasso.coef_) self.assertAlmostEqual(lasso.intercept_, wlasso.intercept_)
def test_orf(self): # Single outcome only, ORF does not support multiple outcomes X = TestPandasIntegration.df[TestPandasIntegration.features] W = TestPandasIntegration.df[TestPandasIntegration.controls] Y = TestPandasIntegration.df[TestPandasIntegration.outcome] T = TestPandasIntegration.df[TestPandasIntegration.cont_treat] # Test DMLOrthoForest est = DMLOrthoForest(n_trees=100, max_depth=2, model_T=WeightedLasso(), model_Y=WeightedLasso()) 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()) # Test DROrthoForest est = DROrthoForest(n_trees=100, max_depth=2) T = TestPandasIntegration.df[TestPandasIntegration.bin_treat] 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_access_to_internal_models(self): """ Test that API related to accessing the nuisance models, cate_model and featurizer is working. """ est = LinearIntentToTreatDRIV(model_Y_X=LinearRegression(), model_T_XZ=LogisticRegression(C=1000), flexible_model_effect=WeightedLasso(), featurizer=PolynomialFeatures( degree=2, include_bias=False)) Y = np.array([1, 1, 2, 2, 1, 1, 2, 2, 1, 1, 2, 2]) T = np.array([1, 1, 2, 2, 1, 1, 2, 2, 1, 1, 2, 2]) Z = np.array([1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2]) X = np.array([1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 6]).reshape(-1, 1) est.fit(Y, T, Z=Z, X=X) assert isinstance(est.original_featurizer, PolynomialFeatures) assert isinstance(est.featurizer_, Pipeline) assert isinstance(est.model_final_, StatsModelsLinearRegression) for mdl in est.models_Y_X: assert isinstance(mdl, LinearRegression) for mdl in est.models_T_XZ: assert isinstance(mdl, LogisticRegression) np.testing.assert_array_equal(est.cate_feature_names(['A']), ['A', 'A^2']) np.testing.assert_array_equal(est.cate_feature_names(), ['x0', 'x0^2']) est = LinearIntentToTreatDRIV(model_Y_X=LinearRegression(), model_T_XZ=LogisticRegression(C=1000), flexible_model_effect=WeightedLasso(), featurizer=None) est.fit(Y, T, Z=Z, X=X) assert est.original_featurizer is None assert isinstance(est.featurizer_, FunctionTransformer) assert isinstance(est.model_final_, StatsModelsLinearRegression) for mdl in est.models_Y_X: assert isinstance(mdl, LinearRegression) for mdl in est.models_T_XZ: assert isinstance(mdl, LogisticRegression) np.testing.assert_array_equal(est.cate_feature_names(['A']), ['A'])
def monte_carlo_lasso(first_stage=lambda: WeightedLasso( alpha=0.01, fit_intercept=True, tol=1e-6, random_state=123), folder='lasso'): n_exp = 1000 n_list = [500] hetero_coef_list = [1] d_list = [20] d_x_list = [5] p_list = [1, 2] t_list = [1, 3] cov_type_list = ['HC1'] alpha_list = [.01, .05, .2] run_all_mc(first_stage, folder, n_list, n_exp, hetero_coef_list, d_list, d_x_list, p_list, t_list, cov_type_list, alpha_list)
def test_can_use_statsmodel_inference(self): """Test that we can use statsmodels to generate confidence intervals""" est = LinearIntentToTreatDRIV(model_Y_X=LinearRegression(), model_T_XZ=LogisticRegression(C=1000), flexible_model_effect=WeightedLasso()) est.fit(np.array([1, 1, 2, 2, 1, 1, 2, 2, 1, 1, 2, 2]), np.array([1, 1, 2, 2, 1, 1, 2, 2, 1, 1, 2, 2]), Z=np.array([1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2]), X=np.array([1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 6]).reshape(-1, 1)) interval = est.effect_interval(np.ones((9, 1)), T0=np.array([1, 1, 1, 2, 2, 2, 1, 1, 1]), T1=np.array([1, 2, 1, 1, 2, 2, 2, 2, 1]), alpha=0.05) point = est.effect(np.ones((9, 1)), T0=np.array([1, 1, 1, 2, 2, 2, 1, 1, 1]), T1=np.array([1, 2, 1, 1, 2, 2, 2, 2, 1])) assert len(interval) == 2 lo, hi = interval assert lo.shape == hi.shape == point.shape assert np.all(lo <= point) assert np.all(point <= hi) assert np.any(lo < hi) # for at least some of the examples, the CI should have nonzero width interval = est.const_marginal_effect_interval(np.ones((9, 1)), alpha=0.05) point = est.const_marginal_effect(np.ones((9, 1))) assert len(interval) == 2 lo, hi = interval assert lo.shape == hi.shape == point.shape assert np.all(lo <= point) assert np.all(point <= hi) assert np.any(lo < hi) # for at least some of the examples, the CI should have nonzero width interval = est.coef__interval(alpha=0.05) point = est.coef_ assert len(interval) == 2 lo, hi = interval assert lo.shape == hi.shape == point.shape assert np.all(lo <= point) assert np.all(point <= hi) assert np.any(lo < hi) # for at least some of the examples, the CI should have nonzero width interval = est.intercept__interval(alpha=0.05) point = est.intercept_ assert len(interval) == 2 lo, hi = interval assert np.all(lo <= point) assert np.all(point <= hi) assert np.any(lo < hi) # for at least some of the examples, the CI should have nonzero width
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_orthoiv(self): y, T, X, W = self._get_data() Z = T.copy() est = DMLATEIV(model_Y_W=LinearRegression(), model_T_W=LinearRegression(), model_Z_W=LinearRegression(), mc_iters=2) est.fit(y, T, W=W, Z=Z, cache_values=True) est.refit_final() est.model_Y_W = Lasso() est.model_T_W = ElasticNet() est.model_Z_W = WeightedLasso() est.fit(y, T, W=W, Z=Z, cache_values=True) assert isinstance(est.models_nuisance_[0]._model_Y_W._model, Lasso) assert isinstance(est.models_nuisance_[0]._model_T_W._model, ElasticNet) assert isinstance(est.models_nuisance_[0]._model_Z_W._model, WeightedLasso) est = ProjectedDMLATEIV(model_Y_W=LinearRegression(), model_T_W=LinearRegression(), model_T_WZ=LinearRegression(), mc_iters=2) est.fit(y, T, W=W, Z=Z, cache_values=True) est.refit_final() est.model_Y_W = Lasso() est.model_T_W = ElasticNet() est.model_T_WZ = WeightedLasso() est.fit(y, T, W=W, Z=Z, cache_values=True) assert isinstance(est.models_nuisance_[0]._model_Y_W._model, Lasso) assert isinstance(est.models_nuisance_[0]._model_T_W._model, ElasticNet) assert isinstance(est.models_nuisance_[0]._model_T_WZ._model, WeightedLasso) est = DMLIV(model_Y_X=LinearRegression(), model_T_X=LinearRegression(), model_T_XZ=LinearRegression(), model_final=LinearRegression(fit_intercept=False), mc_iters=2) est.fit(y, T, X=X, Z=Z, cache_values=True) np.testing.assert_equal(len(est.coef_), X.shape[1]) est.featurizer = PolynomialFeatures(degree=2, include_bias=False) est.refit_final() np.testing.assert_equal(len(est.coef_), X.shape[1]**2) est.intercept_ est.fit_cate_intercept = False est.intercept_ est.refit_final() with pytest.raises(AttributeError): est.intercept_ est.model_Y_X = Lasso() est.model_T_X = ElasticNet() est.model_T_XZ = WeightedLasso() est.fit(y, T, X=X, Z=Z, cache_values=True) assert isinstance(est.models_Y_X[0], Lasso) assert isinstance(est.models_T_X[0], ElasticNet) assert isinstance(est.models_T_XZ[0], WeightedLasso) est = DMLIV(model_Y_X=LinearRegression(), model_T_X=LinearRegression(), model_T_XZ=LinearRegression(), model_final=LinearRegression(fit_intercept=False), mc_iters=2) est.fit(y, T, X=X, Z=Z, cache_values=True) np.testing.assert_equal(len(est.coef_), X.shape[1]) est.featurizer = PolynomialFeatures(degree=2, include_bias=False) est.refit_final() np.testing.assert_equal(len(est.coef_), X.shape[1]**2) est.intercept_ est.fit_cate_intercept = False est.intercept_ est.refit_final() with pytest.raises(AttributeError): est.intercept_ est.model_Y_X = Lasso() est.model_T_X = ElasticNet() est.model_T_XZ = WeightedLasso() est.fit(y, T, X=X, Z=Z, cache_values=True) assert isinstance(est.models_nuisance_[0]._model_Y_X._model, Lasso) assert isinstance(est.models_nuisance_[0]._model_T_X._model, ElasticNet) assert isinstance(est.models_nuisance_[0]._model_T_XZ._model, WeightedLasso) est = NonParamDMLIV(model_Y_X=LinearRegression(), model_T_X=LinearRegression(), model_T_XZ=LinearRegression(), model_final=LinearRegression(fit_intercept=True), mc_iters=2) est.fit(y, T, X=X, Z=Z, cache_values=True) est.featurizer = PolynomialFeatures(degree=2, include_bias=False) est.model_final = WeightedLasso() est.refit_final() assert isinstance(est.model_cate, WeightedLasso) assert isinstance(est.featurizer_, PolynomialFeatures) est = IntentToTreatDRIV(model_Y_X=LinearRegression(), model_T_XZ=LogisticRegression(), flexible_model_effect=LinearRegression()) est.fit(y, T, X=X, W=W, Z=Z, cache_values=True) assert est.model_final is None assert isinstance(est.model_final_, LinearRegression) est.flexible_model_effect = Lasso() est.refit_final() assert est.model_final is None assert isinstance(est.model_final_, Lasso) est.model_final = Lasso() est.refit_final() assert isinstance(est.model_final, Lasso) assert isinstance(est.model_final_, Lasso) assert isinstance(est.models_nuisance_[0]._prel_model_effect.model_final_, LinearRegression) est.fit(y, T, X=X, W=W, Z=Z, cache_values=True) assert isinstance(est.models_nuisance_[0]._prel_model_effect.model_final_, Lasso) est = LinearIntentToTreatDRIV(model_Y_X=LinearRegression(), model_T_XZ=LogisticRegression(), flexible_model_effect=LinearRegression()) est.fit(y, T, X=X, W=W, Z=Z, cache_values=True) est.fit_cate_intercept = False est.intercept_ est.intercept__interval() est.refit_final() with pytest.raises(AttributeError): est.intercept_ with pytest.raises(AttributeError): est.intercept__interval() with pytest.raises(ValueError): est.model_final = LinearRegression() est.flexible_model_effect = Lasso() est.fit(y, T, X=X, W=W, Z=Z, cache_values=True) assert isinstance(est.models_nuisance_[0]._prel_model_effect.model_final_, Lasso)
model_Y=sklearn.linear_model.LinearRegression()) est.fit(Y, T, W, W) print(est.effect(W[:2])) # advanced example with many confounders X = np.random.uniform(-1, 1, size=(4000, 1)) W = np.random.normal(size=(4000, 50)) support = np.random.choice(50, 4, replace=False) T = np.dot(W[:, support], np.random.normal(size=4)) + np.random.normal(size=4000) Y = np.exp(2 * X[:, 0]) * T + np.dot(W[:, support], np.random.normal(size=4)) + .5 est = ContinuousTreatmentOrthoForest(n_trees=100, max_depth=5, model_Y=WeightedLasso(alpha=0.01), model_T=WeightedLasso(alpha=0.01)) est.fit(Y, T, X, W) X_test = np.linspace(-1, 1, 30).reshape(-1, 1) treatment_effects = est.effect(X_test) plt.plot(X_test[:, 0], treatment_effects, label='ORF estimate') plt.plot(X_test[:, 0], np.exp(2 * X_test[:, 0]), 'b--', label='True effect') plt.legend() plt.show(block=False) ########## get my data in the framework ####################################### # generate controls/confounders ## W ## region_fe = pd.get_dummies(dta['ags']) time_fe = pd.get_dummies(data=dta[['dow_num', 'month', 'year']],
def first_stage_model(): return WeightedLasso(alpha=0.01, fit_intercept=True, tol=1e-12, random_state=123)
def test_orthoiv(self): y, T, X, W = self._get_data() Z = T.copy() est = OrthoIV(model_y_xw=LinearRegression(), model_t_xw=LinearRegression(), model_z_xw=LinearRegression(), mc_iters=2) est.fit(y, T, Z=Z, W=W, cache_values=True) est.refit_final() est.model_y_xw = Lasso() est.model_t_xw = ElasticNet() est.model_z_xw = WeightedLasso() est.fit(y, T, Z=Z, W=W, cache_values=True) assert isinstance(est.models_nuisance_[0][0]._model_y_xw._model, Lasso) assert isinstance(est.models_nuisance_[0][0]._model_t_xw._model, ElasticNet) assert isinstance(est.models_nuisance_[0][0]._model_z_xw._model, WeightedLasso) est = DMLIV(model_y_xw=LinearRegression(), model_t_xw=LinearRegression(), model_t_xwz=LinearRegression(), model_final=LinearRegression(fit_intercept=False), mc_iters=2) est.fit(y, T, Z=Z, X=X, W=W, cache_values=True) est.model_y_xw = Lasso() est.model_t_xw = ElasticNet() est.model_t_xwz = WeightedLasso() est.fit(y, T, Z=Z, X=X, W=W, cache_values=True) assert isinstance(est.models_nuisance_[0][0]._model_y_xw._model, Lasso) assert isinstance(est.models_nuisance_[0][0]._model_t_xw._model, ElasticNet) assert isinstance(est.models_nuisance_[0][0]._model_t_xwz._model, WeightedLasso) est = NonParamDMLIV(model_y_xw=LinearRegression(), model_t_xw=LinearRegression(), model_t_xwz=LinearRegression(), model_final=LinearRegression(fit_intercept=True), mc_iters=2) est.fit(y, T, Z=Z, X=X, W=W, cache_values=True) est.featurizer = PolynomialFeatures(degree=2, include_bias=False) est.model_final = WeightedLasso() est.refit_final() assert isinstance(est.model_cate, WeightedLasso) assert isinstance(est.featurizer_, PolynomialFeatures) est = IntentToTreatDRIV(model_y_xw=LinearRegression(), model_t_xwz=LogisticRegression(), flexible_model_effect=LinearRegression()) est.fit(y, T, Z=Z, X=X, W=W, cache_values=True) assert est.model_final is None assert isinstance(est.model_final_, LinearRegression) est.flexible_model_effect = Lasso() est.refit_final() assert est.model_final is None assert isinstance(est.model_final_, Lasso) est.model_final = Lasso() est.refit_final() assert isinstance(est.model_final, Lasso) assert isinstance(est.model_final_, Lasso) assert isinstance( est.models_nuisance_[0][0]._prel_model_effect.model_final_, LinearRegression) est.fit(y, T, Z=Z, X=X, W=W, cache_values=True) assert isinstance( est.models_nuisance_[0][0]._prel_model_effect.model_final_, Lasso) est = LinearIntentToTreatDRIV(model_y_xw=LinearRegression(), model_t_xwz=LogisticRegression(), flexible_model_effect=LinearRegression()) est.fit(y, T, Z=Z, X=X, W=W, cache_values=True) est.fit_cate_intercept = False est.intercept_ est.intercept__interval() est.refit_final() with pytest.raises(AttributeError): est.intercept_ with pytest.raises(AttributeError): est.intercept__interval() with pytest.raises(ValueError): est.model_final = LinearRegression() est.flexible_model_effect = Lasso() est.fit(y, T, Z=Z, X=X, W=W, cache_values=True) assert isinstance( est.models_nuisance_[0][0]._prel_model_effect.model_final_, Lasso)
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)] ] 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 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 ())) 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:] + ( T.shape[1:] if not is_discrete else (2, )) + fd_x intercept_shape = Y.shape[1:] + ( T.shape[1:] if not is_discrete else (2, )) 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', featurizer=featurizer, fit_cate_intercept=fit_cate_intercept, discrete_treatment=is_discrete), True, all_infs), (SparseLinearDMLCateEstimator(model_y=WeightedLasso(), model_t=model_t, featurizer=featurizer, fit_cate_intercept=fit_cate_intercept, discrete_treatment=is_discrete), True, [None, 'debiasedlasso']), (KernelDMLCateEstimator(model_y=WeightedLasso(), model_t=model_t, fit_cate_intercept=fit_cate_intercept, 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): if X is None and ( not fit_cate_intercept): 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 isinstance(est, LinearDMLCateEstimator) or\ isinstance(est, SparseLinearDMLCateEstimator): 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, 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) if (isinstance( est, LinearDMLCateEstimator ) or isinstance( est, SparseLinearDMLCateEstimator )): 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 ) 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: # 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_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)