def test_continuous_treatments(self):
np.random.seed(123)
for global_residualization in [False, True]:
# Generate data with continuous treatments
T = np.dot(TestOrthoForest.W[:, TestOrthoForest.support], TestOrthoForest.coefs_T) + \
TestOrthoForest.eta_sample(TestOrthoForest.n)
TE = np.array([self._exp_te(x) for x in TestOrthoForest.X])
Y = np.dot(TestOrthoForest.W[:, TestOrthoForest.support], TestOrthoForest.coefs_Y) + \
T * TE + TestOrthoForest.epsilon_sample(TestOrthoForest.n)
# Instantiate model with most of the default parameters. Using n_jobs=1 since code coverage
# does not work well with parallelism.
est = DMLOrthoForest(n_jobs=1,
n_trees=10,
model_T=Lasso(),
model_Y=Lasso(),
model_T_final=WeightedLassoCVWrapper(),
model_Y_final=WeightedLassoCVWrapper(),
global_residualization=global_residualization)
# Test inputs for continuous treatments
# --> Check that one can pass in regular lists
est.fit(list(Y),
list(T),
X=list(TestOrthoForest.X),
W=list(TestOrthoForest.W))
# --> Check that it fails correctly if lists of different shape are passed in
self.assertRaises(ValueError,
est.fit,
Y[:TestOrthoForest.n // 2],
T[:TestOrthoForest.n // 2],
X=TestOrthoForest.X,
W=TestOrthoForest.W)
# Check that outputs have the correct shape
out_te = est.const_marginal_effect(TestOrthoForest.x_test)
self.assertEqual(TestOrthoForest.x_test.shape[0], out_te.shape[0])
# Test continuous treatments with controls
est = DMLOrthoForest(n_trees=100,
min_leaf_size=10,
max_depth=50,
subsample_ratio=0.50,
bootstrap=False,
n_jobs=1,
model_T=Lasso(alpha=0.024),
model_Y=Lasso(alpha=0.024),
model_T_final=WeightedLassoCVWrapper(cv=5),
model_Y_final=WeightedLassoCVWrapper(cv=5),
global_residualization=global_residualization,
global_res_cv=5)
est.fit(Y,
T,
X=TestOrthoForest.X,
W=TestOrthoForest.W,
inference="blb")
self._test_te(est, TestOrthoForest.expected_exp_te, tol=0.5)
self._test_ci(est, TestOrthoForest.expected_exp_te, tol=1.5)
# Test continuous treatments without controls
T = TestOrthoForest.eta_sample(TestOrthoForest.n)
Y = T * TE + TestOrthoForest.epsilon_sample(TestOrthoForest.n)
est.fit(Y, T, X=TestOrthoForest.X, inference="blb")
self._test_te(est, TestOrthoForest.expected_exp_te, tol=0.5)
self._test_ci(est, TestOrthoForest.expected_exp_te, tol=1.5)
def test_multiple_treatments(self):
np.random.seed(123)
# Only applicable to continuous treatments
# Generate data for 2 treatments
TE = np.array(
[[TestOrthoForest._exp_te(x),
TestOrthoForest._const_te(x)] for x in TestOrthoForest.X])
coefs_T = uniform(0, 1, size=(TestOrthoForest.support_size, 2))
T = np.matmul(TestOrthoForest.W[:, TestOrthoForest.support], coefs_T) + \
uniform(-1, 1, size=(TestOrthoForest.n, 2))
delta_Y = np.array(
[np.dot(TE[i], T[i]) for i in range(TestOrthoForest.n)])
Y = delta_Y + np.dot(TestOrthoForest.W[:, TestOrthoForest.support], TestOrthoForest.coefs_Y) + \
TestOrthoForest.epsilon_sample(TestOrthoForest.n)
# Test multiple treatments with controls
est = ContinuousTreatmentOrthoForest(
n_trees=50,
min_leaf_size=10,
max_depth=50,
subsample_ratio=0.30,
bootstrap=False,
n_jobs=4,
model_T=MultiOutputRegressor(Lasso(alpha=0.024)),
model_Y=Lasso(alpha=0.024),
model_T_final=WeightedLassoCVWrapper(),
model_Y_final=WeightedLassoCVWrapper())
est.fit(Y, T, TestOrthoForest.X, TestOrthoForest.W, inference="blb")
expected_te = np.array([
TestOrthoForest.expected_exp_te, TestOrthoForest.expected_const_te
]).T
self._test_te(est, expected_te, tol=0.5, treatment_type='multi')
self._test_ci(est, expected_te, tol=2.0, treatment_type='multi')
def test_binary_treatments(self):
np.random.seed(123)
# Generate data with binary treatments
log_odds = np.dot(TestOrthoForest.W[:, TestOrthoForest.support], TestOrthoForest.coefs_T) + \
TestOrthoForest.eta_sample(TestOrthoForest.n)
T_sigmoid = 1 / (1 + np.exp(-log_odds))
T = np.array([np.random.binomial(1, p) for p in T_sigmoid])
TE = np.array([self._exp_te(x) for x in TestOrthoForest.X])
Y = np.dot(TestOrthoForest.W[:, TestOrthoForest.support], TestOrthoForest.coefs_Y) + \
T * TE + TestOrthoForest.epsilon_sample(TestOrthoForest.n)
# Instantiate model with default params. Using n_jobs=1 since code coverage
# does not work well with parallelism.
est = DiscreteTreatmentOrthoForest(n_trees=10, n_jobs=1,
propensity_model=LogisticRegression(), model_Y=Lasso(),
propensity_model_final=LogisticRegressionCV(penalty='l1', solver='saga'),
model_Y_final=WeightedLassoCVWrapper())
# Test inputs for binary treatments
# --> Check that one can pass in regular lists
est.fit(list(Y), list(T), list(TestOrthoForest.X), list(TestOrthoForest.W))
# --> Check that it fails correctly if lists of different shape are passed in
self.assertRaises(ValueError, est.fit, Y[:TestOrthoForest.n // 2], T[:TestOrthoForest.n // 2],
TestOrthoForest.X, TestOrthoForest.W)
# --> Check that it works when T, Y have shape (n, 1)
est.fit(Y.reshape(-1, 1), T.reshape(-1, 1), TestOrthoForest.X, TestOrthoForest.W)
# --> Check that it fails correctly when T has shape (n, 2)
self.assertRaises(ValueError, est.fit, Y, np.ones((TestOrthoForest.n, 2)),
TestOrthoForest.X, TestOrthoForest.W)
# --> Check that it fails correctly when the treatments are not numeric
self.assertRaises(ValueError, est.fit, Y, np.array(["a"] * TestOrthoForest.n),
TestOrthoForest.X, TestOrthoForest.W)
# Check that outputs have the correct shape
out_te = est.const_marginal_effect(TestOrthoForest.x_test)
self.assertSequenceEqual((TestOrthoForest.x_test.shape[0], 1), out_te.shape)
# Test binary treatments with controls
est = DiscreteTreatmentOrthoForest(n_trees=100, min_leaf_size=10,
max_depth=30, subsample_ratio=0.30, bootstrap=False, n_jobs=4,
propensity_model=LogisticRegression(
C=1 / 0.024, penalty='l1', solver='saga'),
model_Y=Lasso(alpha=0.024),
propensity_model_final=LogisticRegressionCV(penalty='l1', solver='saga'),
model_Y_final=WeightedLassoCVWrapper())
est.fit(Y, T, TestOrthoForest.X, TestOrthoForest.W, inference="blb")
self._test_te(est, TestOrthoForest.expected_exp_te, tol=0.7, treatment_type='discrete')
self._test_ci(est, TestOrthoForest.expected_exp_te, tol=1.5, treatment_type='discrete')
# Test binary treatments without controls
log_odds = TestOrthoForest.eta_sample(TestOrthoForest.n)
T_sigmoid = 1 / (1 + np.exp(-log_odds))
T = np.array([np.random.binomial(1, p) for p in T_sigmoid])
Y = T * TE + TestOrthoForest.epsilon_sample(TestOrthoForest.n)
est.fit(Y, T, TestOrthoForest.X, inference="blb")
self._test_te(est, TestOrthoForest.expected_exp_te, tol=0.5, treatment_type='discrete')
self._test_ci(est, TestOrthoForest.expected_exp_te, tol=1.5, treatment_type='discrete')
def __init__(self,
model_y=WeightedLassoCVWrapper(), model_t='auto',
alpha='auto',
max_iter=1000,
tol=1e-4,
featurizer=None,
fit_cate_intercept=True,
linear_first_stages=True,
discrete_treatment=False,
n_splits=2,
random_state=None):
model_final = MultiOutputDebiasedLasso(
alpha=alpha,
fit_intercept=False,
max_iter=max_iter,
tol=tol)
super().__init__(model_y=model_y,
model_t=model_t,
model_final=model_final,
featurizer=featurizer,
fit_cate_intercept=fit_cate_intercept,
linear_first_stages=linear_first_stages,
discrete_treatment=discrete_treatment,
n_splits=n_splits,
random_state=random_state)
def __init__(self,
model_y, model_t, model_final,
featurizer=None,
fit_cate_intercept=True,
linear_first_stages=False,
discrete_treatment=False,
n_splits=2,
random_state=None):
# TODO: consider whether we need more care around stateful featurizers,
# since we clone it and fit separate copies
if model_t == 'auto':
if discrete_treatment:
model_t = LogisticRegressionCV(cv=WeightedStratifiedKFold())
else:
model_t = WeightedLassoCVWrapper()
self.bias_part_of_coef = fit_cate_intercept
self.fit_cate_intercept = fit_cate_intercept
super().__init__(model_y=_FirstStageWrapper(model_y, True,
featurizer, linear_first_stages, discrete_treatment),
model_t=_FirstStageWrapper(model_t, False,
featurizer, linear_first_stages, discrete_treatment),
model_final=_FinalWrapper(model_final, fit_cate_intercept, featurizer, False),
discrete_treatment=discrete_treatment,
n_splits=n_splits,
random_state=random_state)
def test_continuous_treatments(self):
np.random.seed(123)
# Generate data with continuous treatments
T = np.dot(TestOrthoForest.W[:, TestOrthoForest.support], TestOrthoForest.coefs_T) + \
TestOrthoForest.eta_sample(TestOrthoForest.n)
TE = np.array([self._exp_te(x) for x in TestOrthoForest.X])
Y = np.dot(TestOrthoForest.W[:, TestOrthoForest.support], TestOrthoForest.coefs_Y) + \
T * TE + TestOrthoForest.epsilon_sample(TestOrthoForest.n)
# Instantiate model with most of the default parameters
est = ContinuousTreatmentOrthoForest(
n_jobs=4,
n_trees=10,
model_T=Lasso(),
model_Y=Lasso(),
model_T_final=WeightedLassoCVWrapper(),
model_Y_final=WeightedLassoCVWrapper())
# Test inputs for continuous treatments
# --> Check that one can pass in regular lists
est.fit(list(Y), list(T), list(TestOrthoForest.X),
list(TestOrthoForest.W))
# --> Check that it fails correctly if lists of different shape are passed in
self.assertRaises(ValueError, est.fit, Y[:TestOrthoForest.n // 2],
T[:TestOrthoForest.n // 2], TestOrthoForest.X,
TestOrthoForest.W)
# Check that outputs have the correct shape
out_te = est.const_marginal_effect(TestOrthoForest.x_test)
self.assertSequenceEqual((TestOrthoForest.x_test.shape[0], 1),
out_te.shape)
# Test continuous treatments with controls
est = ContinuousTreatmentOrthoForest(
n_trees=50,
min_leaf_size=10,
max_depth=50,
subsample_ratio=0.30,
bootstrap=False,
n_jobs=4,
model_T=Lasso(alpha=0.024),
model_Y=Lasso(alpha=0.024),
model_T_final=WeightedLassoCVWrapper(),
model_Y_final=WeightedLassoCVWrapper())
est.fit(Y, T, TestOrthoForest.X, TestOrthoForest.W)
self._test_te(est, TestOrthoForest.expected_exp_te, tol=0.5)
# Test continuous treatments without controls
T = TestOrthoForest.eta_sample(TestOrthoForest.n)
Y = T * TE + TestOrthoForest.epsilon_sample(TestOrthoForest.n)
est.fit(Y, T, TestOrthoForest.X)
self._test_te(est, TestOrthoForest.expected_exp_te, tol=0.5)
def __init__(self,
model_y=WeightedLassoCVWrapper(), model_t='auto',
featurizer=None,
fit_cate_intercept=True,
linear_first_stages=True,
discrete_treatment=False,
n_splits=2,
random_state=None):
super().__init__(model_y=model_y,
model_t=model_t,
model_final=StatsModelsLinearRegression(fit_intercept=False),
featurizer=featurizer,
fit_cate_intercept=fit_cate_intercept,
linear_first_stages=linear_first_stages,
discrete_treatment=discrete_treatment,
n_splits=n_splits,
random_state=random_state)
def __init__(self, model_y=WeightedLassoCVWrapper(), model_t='auto', fit_cate_intercept=True,
dim=20, bw=1.0, discrete_treatment=False, n_splits=2, random_state=None):
class RandomFeatures(TransformerMixin):
def __init__(self, random_state):
self._random_state = check_random_state(random_state)
def fit(self, X):
self.omegas = self._random_state.normal(0, 1 / bw, size=(shape(X)[1], dim))
self.biases = self._random_state.uniform(0, 2 * np.pi, size=(1, dim))
return self
def transform(self, X):
return np.sqrt(2 / dim) * np.cos(np.matmul(X, self.omegas) + self.biases)
super().__init__(model_y=model_y, model_t=model_t,
model_final=ElasticNetCV(fit_intercept=False),
featurizer=RandomFeatures(random_state),
fit_cate_intercept=fit_cate_intercept,
discrete_treatment=discrete_treatment, n_splits=n_splits, random_state=random_state)
def test_wrapper_attributes(self):
"""Test that attributes are properly maintained across calls to fit that switch between 1- and 2-D"""
wrapper = WeightedLassoCVWrapper(alphas=[5, 10], max_iter=100)
wrapper.tol = 0.01 # set an attribute manually as well
assert wrapper.alphas == [5, 10]
assert wrapper.max_iter == 100
assert wrapper.tol == 0.01
# perform 1D fit
wrapper.fit(np.random.normal(size=(100, 3)),
np.random.normal(size=100))
assert wrapper.alphas == [5, 10]
assert wrapper.max_iter == 100
assert wrapper.tol == 0.01
# perform 2D fit
wrapper.fit(np.random.normal(size=(100, 3)),
np.random.normal(size=(100, 2)))
assert wrapper.alphas == [5, 10]
assert wrapper.max_iter == 100
assert wrapper.tol == 0.01
def __init__(self,
model_y, model_t, model_final,
featurizer=None,
fit_cate_intercept=True,
linear_first_stages=False,
discrete_treatment=False,
n_splits=2,
random_state=None):
# TODO: consider whether we need more care around stateful featurizers,
# since we clone it and fit separate copies
if model_t == 'auto':
if discrete_treatment:
model_t = LogisticRegressionCV(cv=WeightedStratifiedKFold())
else:
model_t = WeightedLassoCVWrapper()
class FirstStageWrapper:
def __init__(self, model, is_Y):
self._model = clone(model, safe=False)
self._featurizer = clone(featurizer, safe=False)
self._is_Y = is_Y
def _combine(self, X, W, n_samples, fitting=True):
if X is None:
# if both X and W are None, just return a column of ones
return (W if W is not None else np.ones((n_samples, 1)))
XW = hstack([X, W]) if W is not None else X
if self._is_Y and linear_first_stages:
if self._featurizer is None:
F = X
else:
F = self._featurizer.fit_transform(X) if fitting else self._featurizer.transform(X)
return cross_product(XW, hstack([np.ones((shape(XW)[0], 1)), F]))
else:
return XW
def fit(self, X, W, Target, sample_weight=None):
if (not self._is_Y) and discrete_treatment:
# In this case, the Target is the one-hot-encoding of the treatment variable
# We need to go back to the label representation of the one-hot so as to call
# the classifier.
if np.any(np.all(Target == 0, axis=0)) or (not np.any(np.all(Target == 0, axis=1))):
raise AttributeError("Provided crossfit folds contain training splits that " +
"don't contain all treatments")
Target = inverse_onehot(Target)
if sample_weight is not None:
self._model.fit(self._combine(X, W, Target.shape[0]), Target, sample_weight=sample_weight)
else:
self._model.fit(self._combine(X, W, Target.shape[0]), Target)
def predict(self, X, W):
n_samples = X.shape[0] if X is not None else (W.shape[0] if W is not None else 1)
if (not self._is_Y) and discrete_treatment:
return self._model.predict_proba(self._combine(X, W, n_samples, fitting=False))[:, 1:]
else:
return self._model.predict(self._combine(X, W, n_samples, fitting=False))
class FinalWrapper:
def __init__(self):
self._model = clone(model_final, safe=False)
self._original_featurizer = clone(featurizer, safe=False)
self._fit_cate_intercept = fit_cate_intercept
if self._fit_cate_intercept:
add_intercept = FunctionTransformer(lambda F:
hstack([np.ones((F.shape[0], 1)), F]))
if featurizer:
self._featurizer = Pipeline([('featurize', self._original_featurizer),
('add_intercept', add_intercept)])
else:
self._featurizer = add_intercept
else:
self._featurizer = self._original_featurizer
def _combine(self, X, T, fitting=True):
if X is not None:
if self._featurizer is not None:
F = self._featurizer.fit_transform(X) if fitting else self._featurizer.transform(X)
else:
F = X
else:
if not self._fit_cate_intercept:
raise AttributeError("Cannot have X=None and also not allow for a CATE intercept!")
F = np.ones((T.shape[0], 1))
return cross_product(F, T)
def fit(self, X, T_res, Y_res, sample_weight=None, sample_var=None):
# Track training dimensions to see if Y or T is a vector instead of a 2-dimensional array
self._d_t = shape(T_res)[1:]
self._d_y = shape(Y_res)[1:]
fts = self._combine(X, T_res)
if sample_weight is not None:
if sample_var is not None:
self._model.fit(fts,
Y_res, sample_weight=sample_weight, sample_var=sample_var)
else:
self._model.fit(fts,
Y_res, sample_weight=sample_weight)
else:
self._model.fit(fts, Y_res)
self._intercept = None
intercept = self._model.predict(np.zeros_like(fts[0:1]))
if (np.count_nonzero(intercept) > 0):
warn("The final model has a nonzero intercept for at least one outcome; "
"it will be subtracted, but consider fitting a model without an intercept if possible.",
UserWarning)
self._intercept = intercept
def predict(self, X):
X2, T = broadcast_unit_treatments(X if X is not None else np.empty((1, 0)),
self._d_t[0] if self._d_t else 1)
prediction = self._model.predict(self._combine(None if X is None else X2, T, fitting=False))
if self._intercept is not None:
prediction -= self._intercept
return reshape_treatmentwise_effects(prediction,
self._d_t, self._d_y)
self.bias_part_of_coef = fit_cate_intercept
self.fit_cate_intercept = fit_cate_intercept
super().__init__(model_y=FirstStageWrapper(model_y, is_Y=True),
model_t=FirstStageWrapper(model_t, is_Y=False),
model_final=FinalWrapper(),
discrete_treatment=discrete_treatment,
n_splits=n_splits,
random_state=random_state)
subsample_ratio = 0.04
#%%
# Definition of range of variable tested for heterogeneity
min_tfsum = 0.0
max_tfsum = 24.0
delta = (max_tfsum - min_tfsum) / 100
X_test = np.arange(min_tfsum, max_tfsum + delta - 0.001, delta).reshape(-1, 1)
#%%
# Estimation of causal tree
est = CausalForest(n_trees=n_trees,
min_leaf_size=min_leaf_size,
max_depth=max_depth,
subsample_ratio=subsample_ratio,
model_T=WeightedLassoCVWrapper(cv=3),
model_Y=WeightedLassoCVWrapper(cv=3),
random_state=123)
est.fit(Y, T, X=X, W=W)
treatment_effects = est.effect(X_test)
te_lower, te_upper = est.effect_interval(X_test)
#%%
# Plot results
plt.figure(figsize=(15, 5))
plt.plot(X_test.flatten(), treatment_effects)
plt.fill_between(X_test.flatten(),
te_lower,
te_upper,
label="90% CI",
alpha=0.3)