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)