def test_TMLELearner(generate_regression_data):
y, X, treatment, tau, b, e = generate_regression_data()
learner = TMLELearner(learner=XGBRegressor())
# check the accuracy of the ATE estimation
ate_p, lb, ub = learner.estimate_ate(X=X, p=e, treatment=treatment, y=y)
assert (ate_p >= lb) and (ate_p <= ub)
assert ape(tau.mean(), ate_p) < ERROR_THRESHOLD
def test_pandas_input(generate_regression_data):
y, X, treatment, tau, b, e = generate_regression_data()
# convert to pandas types
y = pd.Series(y)
X = pd.DataFrame(X)
treatment = pd.Series(treatment)
try:
learner = BaseSLearner(learner=LinearRegression())
ate_p, lb, ub = learner.estimate_ate(X=X, treatment=treatment, y=y, return_ci=True)
except AttributeError:
assert False
try:
learner = BaseTLearner(learner=LinearRegression())
ate_p, lb, ub = learner.estimate_ate(X=X, treatment=treatment, y=y)
except AttributeError:
assert False
try:
learner = BaseXLearner(learner=LinearRegression())
ate_p, lb, ub = learner.estimate_ate(X=X, treatment=treatment, y=y, p=e)
except AttributeError:
assert False
try:
learner = BaseRLearner(learner=LinearRegression())
ate_p, lb, ub = learner.estimate_ate(X=X, treatment=treatment, y=y, p=e)
except AttributeError:
assert False
try:
learner = TMLELearner(learner=LinearRegression())
ate_p, lb, ub = learner.estimate_ate(X=X, treatment=treatment, y=y, p=e)
except AttributeError:
assert False