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=WeightedModelWrapper(LassoCV(),
sample_type="weighted"),
model_Y_final=WeightedModelWrapper(LassoCV(),
sample_type="weighted"))
# 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=WeightedModelWrapper(LassoCV(),
sample_type="weighted"),
model_Y_final=WeightedModelWrapper(LassoCV(),
sample_type="weighted"))
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)
class OrthoForest(Model):
def __init__(self, *args, **kwargs):
self.reg = None
super(OrthoForest, self).__init__(*args, **kwargs)
def fit(self, x, t, y, nfolds=5, seed=282):
# splits = super().get_splits(x, nfolds, seed)
#### CLASSIFICATION ####
if self.binary:
self.reg = DiscreteTreatmentOrthoForest(n_trees=1,
max_depth=2,
n_jobs=100,
subsample_ratio=0.25,
random_state=282)
self.reg.fit(y, t, x)
#### REGRESSION ####
else:
self.reg = ContinuousTreatmentOrthoForest(n_trees=1,
max_depth=2,
n_jobs=100,
subsample_ratio=0.25,
random_state=282)
self.reg.fit(y, t, x)
def predict(self, x, t):
if self.reg is None:
raise Exception('OrthoForest not Initialized')
# print("x", x.shape, x)
# print("t", t.shape, t)
effect = self.reg.const_marginal_effect(x).reshape(-1)
# print("effect", effect.shape, effect)
return effect * t
def get_predictors(self, x, t):
return np.hstack([x, (t - 0.5).reshape(-1, 1) * x])
