def test_shap_rfe_randomized_search_cols_to_keep(X, y, capsys): """ Test with ShapRFECV with column to keep param. """ clf = DecisionTreeClassifier(max_depth=1) param_grid = {"criterion": ["gini"], "min_samples_split": [1, 2]} search = RandomizedSearchCV(clf, param_grid, cv=2, n_iter=2) with pytest.warns(None) as record: shap_elimination = ShapRFECV(search, step=0.8, cv=2, scoring="roc_auc", n_jobs=4, random_state=1) report = shap_elimination.fit_compute( X, y, columns_to_keep=["col_2", "col_3"]) assert shap_elimination.fitted shap_elimination._check_if_fitted() assert report.shape[0] == 2 reduced_feature_set = set( shap_elimination.get_reduced_features_set(num_features=2)) assert reduced_feature_set == set(["col_2", "col_3"]) _ = shap_elimination.plot(show=False) # Ensure that number of warnings was at least 2 for the verbose (2 generated by probatus + possibly more by SHAP) assert len(record) >= 2 # Check if there is any prints out, _ = capsys.readouterr() assert len(out) == 0
def test_shap_rfe_randomized_search(X, y, capsys): clf = DecisionTreeClassifier(max_depth=1) param_grid = {'criterion': ['gini'], 'min_samples_split': [1, 2]} search = RandomizedSearchCV(clf, param_grid, cv=2, n_iter=2) with pytest.warns(None) as record: shap_elimination = ShapRFECV(search, step=0.8, cv=2, scoring='roc_auc', n_jobs=4, verbose=150) report = shap_elimination.fit_compute(X, y) assert shap_elimination.fitted == True shap_elimination._check_if_fitted() assert report.shape[0] == 2 assert shap_elimination.get_reduced_features_set(1) == ['col_3'] ax1 = shap_elimination.plot(show=False) # Ensure that number of warnings was at least 2 for the verbose (2 generated by probatus + possibly more by SHAP) assert len(record) >= 2 # Check if there is any prints out, _ = capsys.readouterr() assert len(out) > 0
def test_complex_dataset(complex_data, complex_lightgbm): """ Test on complex dataset. """ X, y = complex_data param_grid = { "n_estimators": [5, 7, 10], "num_leaves": [3, 5, 7, 10], } search = RandomizedSearchCV(complex_lightgbm, param_grid, n_iter=1) shap_elimination = ShapRFECV(clf=search, step=1, cv=10, scoring="roc_auc", n_jobs=3, verbose=50) with pytest.warns(None) as record: report = shap_elimination.fit_compute(X, y) assert report.shape[0] == X.shape[1] assert len(record) >= 2