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(), cv=2) inference = BootstrapInference(n_bootstrap_samples=20, n_jobs=-1, verbose=3) 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