def test_can_use_vectors(self): """Test that we can pass vectors for T and Y (not only 2-dimensional arrays).""" dml = DMLCateEstimator(LinearRegression(), LinearRegression(), featurizer=FunctionTransformer()) dml.fit(np.array([1, 2, 3, 1, 2, 3]), np.array([1, 2, 3, 1, 2, 3]), np.ones((6, 1))) self.assertAlmostEqual(dml.coef_.reshape(())[()], 1)
def test_cate_api(self): """Test that we correctly implement the CATE API.""" for d_t in [2, -1]: for d_y in [2, -1]: for d_x in [2, None]: for d_w in [2, None]: for est in [ DMLCateEstimator(model_y=LinearRegression(), model_t=LinearRegression()), SparseLinearDMLCateEstimator( linear_model_y=LinearRegression(), linear_model_t=LinearRegression()), KernelDMLCateEstimator( model_y=LinearRegression(), model_t=LinearRegression()) ]: n = 20 with self.subTest(d_w=d_w, d_x=d_x, d_y=d_y, d_t=d_t): W, X, Y, T = [ np.random.normal(size=(n, d)) if (d and d >= 0) else np.random.normal( size=(n, )) if (d and d < 0) else None for d in [d_w, d_x, d_y, d_t] ] est.fit(Y, T, X, W) # just make sure we can call the marginal_effect and effect methods est.marginal_effect(None, X) est.effect(0, T, X)
def test_discrete_treatments(self): """Test that we can use discrete treatments""" dml = DMLCateEstimator(LinearRegression(), LogisticRegression(C=1000), featurizer=FunctionTransformer(), discrete_treatment=True) # create a simple artificial setup where effect of moving from treatment # 1 -> 2 is 2, # 1 -> 3 is 1, and # 2 -> 3 is -1 (necessarily, by composing the previous two effects) # Using an uneven number of examples from different classes, # and having the treatments in non-lexicographic order, # Should rule out some basic issues. dml.fit(np.array([2, 3, 1, 3, 2, 1, 1, 1]), np.array([3, 2, 1, 2, 3, 1, 1, 1]), np.ones((8, 1))) np.testing.assert_almost_equal( dml.effect(np.array([1, 1, 1, 2, 2, 2, 3, 3, 3]), np.array([1, 2, 3, 1, 2, 3, 1, 2, 3]), np.ones((9, 1))), [0, 2, 1, -2, 0, -1, -1, 1, 0])
def test_ignores_final_intercept(self): """Test that final model intercepts are ignored (with a warning)""" class InterceptModel: def fit(Y, X): pass def predict(X): return X + 1 # (incorrectly) use a final model with an intercept dml = DMLCateEstimator(LinearRegression(), LinearRegression(), model_final=InterceptModel) # Because final model is fixed, actual values of T and Y don't matter t = np.random.normal(size=100) y = np.random.normal(size=100) with self.assertWarns( Warning): # we should warn whenever there's an intercept dml.fit(y, t) assert dml.const_marginal_effect( ) == 1 # coefficient on X in InterceptModel is 1
def test_complex_features(self): # recover simple features by initializing complex features appropriately for _ in range(10): d_w = np.random.randint(0, 4) d_x = np.random.randint(1, 3) d_y = np.random.randint(1, 3) d_t = np.random.randint(1, 3) n = 20 with self.subTest(d_w=d_w, d_x=d_x, d_y=d_y, d_t=d_t): W, X, Y, T = [np.random.normal(size=(n, d)) for d in [d_w, d_x, d_y, d_t]] # using full set of matrix features should be equivalent to using non-matrix featurizer dml = DMLCateEstimator(model_y=LinearRegression(), model_t=LinearRegression(), featurizer=MatrixFeatures(d_y, d_t)) dml.fit(Y, T, X, W) coef1 = dml.coef_ dml = DMLCateEstimator(model_y=LinearRegression(), model_t=LinearRegression()) dml.fit(Y, T, X, W) coef2 = dml.coef_ np.testing.assert_allclose(coef1, reshape(coef2, -1))
def test_with_econml(self): """ Test that we can bootstrap econml estimators """ x = np.random.normal(size=(1000, 2)) t = np.random.normal(size=(1000, 1)) t2 = np.random.normal(size=(1000, 1)) y = x[:, 0] * 0.5 + t + np.random.normal(size=(1000, 1)) est = DMLCateEstimator(LinearRegression(), LinearRegression()) est.fit(y, t, x) bs = BootstrapEstimator(est, 50) # test that we can fit with the same arguments as the base estimator bs.fit(y, t, x) # test that we can get the same attribute for the bootstrap as the original, with the same shape self.assertEqual(np.shape(est.coef_), np.shape(bs.coef_)) # 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 lower, upper = bs.coef__interval() for bound in [lower, upper]: self.assertEqual(np.shape(est.coef_), np.shape(bound)) # test that we can do the same thing once we provide percentile bounds lower, upper = bs.coef__interval(lower=10, upper=90) for bound in [lower, upper]: self.assertEqual(np.shape(est.coef_), np.shape(bound)) # test that we can do the same thing with the results of a method, rather than an attribute self.assertEqual(np.shape(est.effect(t, t2, x)), np.shape(bs.effect(t, t2, x))) # 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 lower, upper = bs.effect_interval(t, t2, x) for bound in [lower, upper]: self.assertEqual(np.shape(est.effect(t, t2, x)), np.shape(bound)) # test that we can do the same thing once we provide percentile bounds lower, upper = bs.effect_interval(t, t2, x, lower=10, upper=90) for bound in [lower, upper]: self.assertEqual(np.shape(est.effect(t, t2, x)), np.shape(bound))
def test_internal(self): """Test that the internal use of bootstrap within an estimator works.""" x = np.random.normal(size=(1000, 2)) t = np.random.normal(size=(1000, 1)) t2 = np.random.normal(size=(1000, 1)) y = x[:, 0] * 0.5 + t + np.random.normal(size=(1000, 1)) est = DMLCateEstimator(LinearRegression(), LinearRegression(), inference='bootstrap') est.fit(y, t, x) # 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 lower, upper = est.effect_interval(x, t, t2) for bound in [lower, upper]: self.assertEqual(np.shape(est.effect(x, t, t2)), np.shape(bound)) self.assertFalse(np.allclose(lower, upper)) # test that we can do the same thing once we provide percentile bounds lower, upper = est.effect_interval(x, t, t2, lower=10, upper=90) for bound in [lower, upper]: self.assertEqual(np.shape(est.effect(x, t, t2)), np.shape(bound)) self.assertFalse(np.allclose(lower, upper))
''' EconML from Microsoft python econometrics library with ML automation for causal inference ''' # Double ML from econml.dml import DMLCateEstimator from sklearn.linear_model import LassoCV est = DMLCateEstimator(model_y=LassoCV(), model_t=LassoCV()) est.fit(Y, T, X, W) # W -> high-dimensional confounders, X -> features treatment_effects = est.effect(X_test) # Orthogonal Random Forest from econml.ortho_forest import ContinuousTreatmentOrthoForest # Use defaults est = ContinuousTreatmentOrthoForest() # Or specify hyperparameters est = ContinuousTreatmentOrthoForest(n_trees=500, min_leaf_size=10, max_depth=10, subsample_ratio=0.7, lambda_reg=0.01, model_T=LassoCV(cv=3), model_Y=LassoCV(cv=3) ) est.fit(Y, T, X, W) treatment_effects = est.effect(X_test)
def test_cate_api(self): """Test that we correctly implement the CATE API.""" d_w = 2 d_x = 2 d_y = 2 d_t = 2 n = 20 with self.subTest(d_w=d_w, d_x=d_x, d_y=d_y, d_t=d_t): W, X, Y, T = [np.random.normal(size=(n, d)) for d in [d_w, d_x, d_y, d_t]] dml = DMLCateEstimator(model_y=LinearRegression(), model_t=LinearRegression()) dml.fit(Y, T, X, W) # just make sure we can call the marginal_effect and effect methods dml.marginal_effect(None, X) dml.effect(0, T, X) dml = DMLCateEstimator(model_y=LinearRegression(), model_t=LinearRegression(), featurizer=MatrixFeatures(d_y, d_t)) dml.fit(Y, T, X, W) # just make sure we can call the marginal_effect and effect methods dml.marginal_effect(None, X) dml.effect(0, T, X)
def test_dominicks(): file_name = "oj_large.csv" if not os.path.isfile(file_name): print("Downloading file (this might take a few seconds)...") urllib.request.urlretrieve( "https://msalicedatapublic.blob.core.windows.net/datasets/OrangeJuice/oj_large.csv", file_name) oj_data = pd.read_csv(file_name) brands = sorted(set(oj_data["brand"])) stores = sorted(set(oj_data["store"])) featnames = ["week", "feat"] + list(oj_data.columns[6:]) # Preprocess data import datetime import numpy as np # Convert 'week' to a date # week_zero = datetime.datetime.strptime("09/07/89", "%m/%d/%y") # oj_data["week"] = pd.to_timedelta(oj_data["week"], unit='w') + week_zero # Take log of price oj_data["logprice"] = np.log(oj_data["price"]) oj_data.drop("price", axis=1, inplace=True) # Make brand numeric oj_data["brand"] = [brands.index(b) for b in oj_data["brand"]] class PriceFeaturizer(TransformerMixin): def __init__(self, n_prods, own_price=True, cross_price_groups=False, cross_price_indiv=True, per_product_effects=True): base_arrays = [] effect_names = [] one_hots = [(0, ) * p + (1, ) + (0, ) * (n_prods - p - 1) for p in range(n_prods)] if own_price: base_arrays.append(np.eye(n_prods)) effect_names.append("own price") if cross_price_groups: base_arrays.append((np.ones( (n_prods, n_prods)) - np.eye(n_prods)) / (n_prods - 1)) effect_names.append("group cross price") if cross_price_indiv: for p in range(n_prods): base_arrays.append(one_hots[p] * np.ones((n_prods, 1)) - np.diag(one_hots[p])) effect_names.append("cross price effect {} ->".format(p)) if per_product_effects: all = [(np.diag(one_hots[p]) @ arr, nm + " {}".format(p)) for arr, nm in zip(base_arrays, effect_names) for p in range(n_prods)] # remove meaningless features (e.g. cross-price effects of products on themselves), # which have all zero coeffs nonempty = [(arr, nm) for arr, nm in all if np.count_nonzero(arr) > 0] self._features = [arr for arr, _ in nonempty] self._names = [nm for _, nm in nonempty] else: self._features = base_arrays self._names = effect_names def fit(self, X): self._is_fitted = True assert shape(X)[1] == 0 return self def transform(self, X): assert self._is_fitted assert shape(X)[1] == 0 return np.tile(self._features, (shape(X)[0], 1, 1, 1)) @property def names(self): return self._names for name, op, xp_g, xp_i, pp in [ ("Homogeneous treatment effect", True, False, False, False), ("Heterogeneous treatment effects", True, False, False, True), (("Heterogeneous treatment effects" " with group effects"), True, True, False, True), (("Heterogeneous treatment effects" " with cross price effects"), True, False, True, True) ]: print(name) np.random.seed(42) ft = PriceFeaturizer(n_prods=3, own_price=op, cross_price_groups=xp_g, cross_price_indiv=xp_i, per_product_effects=pp) names = ft.names dml = DMLCateEstimator(model_y=RandomForestRegressor(), model_t=RandomForestRegressor(), featurizer=ft, n_splits=2) effects = [] for store in stores: data = oj_data[oj_data['store'] == store].sort_values( by=['week', 'brand']) dml.fit(T=reshape(data.as_matrix(["logprice"]), (-1, 3)), Y=reshape(data.as_matrix(["logmove"]), (-1, 3)), W=reshape(data.as_matrix(featnames), (-1, 3 * len(featnames)))) effects.append(dml.coef_) effects = np.array(effects) for nm, eff in zip(names, effects.T): print(" Effect: {}".format(nm)) print(" Mean: {}".format(np.mean(eff))) print(" Std.: {}".format(np.std(eff))) class ConstFt(TransformerMixin): def fit(self, X): return self def transform(self, X): return np.ones((shape(X)[0], 1)) print("Vanilla HTE+XP") np.random.seed(42) dml = DMLCateEstimator(model_y=RandomForestRegressor(), model_t=RandomForestRegressor(), featurizer=ConstFt(), n_splits=2) effects = [] for store in stores: data = oj_data[oj_data['store'] == store].sort_values( by=['week', 'brand']) dml.fit(T=reshape(data.as_matrix(["logprice"]), (-1, 3)), Y=reshape(data.as_matrix(["logmove"]), (-1, 3)), W=reshape(data.as_matrix(featnames), (-1, 3 * len(featnames)))) effects.append(dml.coef_) effects = np.array(effects) names = ["{} on {}".format(i, j) for j in range(3) for i in range(3)] for nm, eff in zip(names, reshape(effects, (-1, 9)).T): print(" Effect: {}".format(nm)) print(" Mean: {}".format(np.mean(eff))) print(" Std.: {}".format(np.std(eff)))
def test_access_to_internal_models(self): """ Test that API related to accessing the nuisance models, cate_model and featurizer is working. """ 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'])