def test_bad_splits_discrete(self):
"""
Tests that when some training splits in a crossfit fold don't contain all treatments then an error
is raised.
"""
Y = np.array([2, 3, 1, 3, 2, 1, 1, 1])
bad = np.array([2, 2, 1, 2, 1, 1, 1, 1])
W = np.ones((8, 1))
ok = np.array([1, 2, 3, 1, 2, 3, 1, 2])
models = [Lasso(), Lasso(), Lasso()]
est = DMLATEIV(*models, n_splits=[(np.arange(4, 8), np.arange(4))])
est.fit(Y, T=bad, Z=bad,
W=W) # imbalance ok with continuous instrument/treatment
models = [Lasso(), LogisticRegression(), Lasso()]
est = DMLATEIV(*models,
n_splits=[(np.arange(4, 8), np.arange(4))],
discrete_treatment=True)
with pytest.raises(AttributeError):
est.fit(Y, T=bad, Z=ok, W=W)
models = [Lasso(), Lasso(), LogisticRegression()]
est = DMLATEIV(*models,
n_splits=[(np.arange(4, 8), np.arange(4))],
discrete_instrument=True)
with pytest.raises(AttributeError):
est.fit(Y, T=ok, Z=bad, W=W)
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)