def test_stratify_orthoiv(self):
"""Test that we can properly stratify by treatment/instrument pair"""
T = [1, 0, 1, 1, 0, 0, 1, 0]
Z = [1, 0, 0, 1, 0, 1, 0, 1]
Y = [1, 2, 3, 4, 5, 6, 7, 8]
X = np.array([1, 1, 2, 2, 1, 2, 1, 2]).reshape(-1, 1)
est = LinearIntentToTreatDRIV(model_Y_X=LinearRegression(),
model_T_XZ=LogisticRegression(),
flexible_model_effect=LinearRegression(),
n_splits=2)
inference = BootstrapInference(n_bootstrap_samples=20)
est.fit(Y, T, Z=Z, X=X, inference=inference)
est.const_marginal_effect_interval(X)
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