def test_fetch_bonus_return_types():
res = fetch_bonus('DoubleMLData')
assert isinstance(res, DoubleMLData)
res = fetch_bonus('DataFrame')
assert isinstance(res, pd.DataFrame)
with pytest.raises(ValueError, match=msg_inv_return_type):
_ = fetch_bonus('matrix')
import numpy as np
import pytest
import math
from sklearn.base import clone
from sklearn.linear_model import Lasso, LogisticRegression
from sklearn.ensemble import RandomForestClassifier
import doubleml as dml
from doubleml.datasets import fetch_bonus
from ._utils import draw_smpls
from ._utils_plr_manual import fit_plr, boot_plr
bonus_data = fetch_bonus()
@pytest.fixture(scope='module',
params=[Lasso(),
RandomForestClassifier(max_depth=2, n_estimators=10),
LogisticRegression()])
def learner(request):
return request.param
@pytest.fixture(scope='module',
params=['IV-type', 'partialling out'])
def score(request):
return request.param
def test_fetch_bonus_poly():
data_bonus_wo_poly = fetch_bonus(polynomial_features=False)
n_x = len(data_bonus_wo_poly.x_cols)
data_bonus_w_poly = fetch_bonus(polynomial_features=True)
assert len(data_bonus_w_poly.x_cols) == ((n_x + 1) * n_x / 2 + n_x)