149 lines
4.6 KiB
Python
149 lines
4.6 KiB
Python
"""This file contains tests for the `shap.datasets` module."""
|
|
|
|
import pytest
|
|
|
|
import shap
|
|
|
|
|
|
@pytest.mark.parametrize("n_points", [None, 12])
|
|
def test_imagenet50(n_points):
|
|
# test that fetch/download works fine
|
|
X, y = shap.datasets.imagenet50(n_points=n_points)
|
|
|
|
# check the shape of the result
|
|
# check that the n_points parameter samples the dataset
|
|
n_points = 50 if n_points is None else n_points
|
|
assert X.shape == (n_points, 224, 224, 3)
|
|
assert y.shape == (n_points,)
|
|
|
|
|
|
@pytest.mark.parametrize("n_points", [None, 12])
|
|
def test_california(n_points):
|
|
# test that fetch/download works fine
|
|
X, y = shap.datasets.california(n_points=n_points)
|
|
|
|
# check the shape of the result
|
|
# check that the n_points parameter samples the dataset
|
|
n_points = 20_640 if n_points is None else n_points
|
|
assert X.shape == (n_points, 8)
|
|
assert y.shape == (n_points,)
|
|
|
|
|
|
@pytest.mark.parametrize("n_points", [None, 12])
|
|
def test_linnerud(n_points):
|
|
# test that fetch/download works fine
|
|
X, y = shap.datasets.linnerud(n_points=n_points)
|
|
|
|
# check the shape of the result
|
|
# check that the n_points parameter samples the dataset
|
|
n_points = 20 if n_points is None else n_points
|
|
assert X.shape == (n_points, 3)
|
|
assert y.shape == (n_points, 3)
|
|
|
|
|
|
@pytest.mark.parametrize("n_points", [None, 12])
|
|
def test_imdb(n_points):
|
|
# test that fetch/download works fine
|
|
X, y = shap.datasets.imdb(n_points=n_points)
|
|
|
|
# check the shape of the result
|
|
# check that the n_points parameter samples the dataset
|
|
n_points = 25_000 if n_points is None else n_points
|
|
assert len(X) == n_points
|
|
assert len(y) == n_points
|
|
|
|
|
|
@pytest.mark.parametrize("n_points", [None, 12])
|
|
def test_diabetes(n_points):
|
|
# test that fetch/download works fine
|
|
X, y = shap.datasets.diabetes(n_points=n_points)
|
|
|
|
# check the shape of the result
|
|
# check that the n_points parameter samples the dataset
|
|
n_points = 442 if n_points is None else n_points
|
|
assert X.shape == (n_points, 10)
|
|
assert y.shape == (n_points,)
|
|
|
|
|
|
@pytest.mark.parametrize("n_points", [None, 12])
|
|
def test_iris(n_points):
|
|
# test that fetch/download works fine
|
|
X, y = shap.datasets.iris(n_points=n_points)
|
|
|
|
# check the shape of the result
|
|
# check that the n_points parameter samples the dataset
|
|
n_points = 150 if n_points is None else n_points
|
|
assert X.shape == (n_points, 4)
|
|
assert y.shape == (n_points,)
|
|
|
|
|
|
@pytest.mark.parametrize("n_points", [None, 12])
|
|
def test_adult(n_points):
|
|
# test that fetch/download works fine
|
|
X, y = shap.datasets.adult(n_points=n_points)
|
|
|
|
# check the shape of the result
|
|
# check that the n_points parameter samples the dataset
|
|
n_points = 32_561 if n_points is None else n_points
|
|
assert X.shape == (n_points, 12)
|
|
assert y.shape == (n_points,)
|
|
|
|
|
|
@pytest.mark.parametrize("n_points", [None, 12])
|
|
def test_nhanesi(n_points):
|
|
# test that fetch/download works fine
|
|
X, y = shap.datasets.nhanesi(n_points=n_points)
|
|
|
|
# check the shape of the result
|
|
# check that the n_points parameter samples the dataset
|
|
n_points = 14_264 if n_points is None else n_points
|
|
assert X.shape == (n_points, 79)
|
|
assert y.shape == (n_points,)
|
|
|
|
|
|
@pytest.mark.parametrize("n_points", [100, 2_000])
|
|
def test_corrgroups60(n_points):
|
|
# test that fetch/download works fine
|
|
X, y = shap.datasets.corrgroups60(n_points=n_points)
|
|
|
|
# check the shape of the result
|
|
# check that the n_points parameter samples the dataset
|
|
assert X.shape == (n_points, 60)
|
|
assert y.shape == (n_points,)
|
|
|
|
|
|
@pytest.mark.parametrize("n_points", [100, 2_000])
|
|
def test_independentlinear60(n_points):
|
|
# test that fetch/download works fine
|
|
X, y = shap.datasets.independentlinear60(n_points=n_points)
|
|
|
|
# check the shape of the result
|
|
# check that the n_points parameter samples the dataset
|
|
assert X.shape == (n_points, 60)
|
|
assert y.shape == (n_points,)
|
|
|
|
|
|
@pytest.mark.parametrize("n_points", [None, 12])
|
|
def test_a1a(n_points):
|
|
# test that fetch/download works fine
|
|
X, y = shap.datasets.a1a(n_points=n_points)
|
|
|
|
# check the shape of the result
|
|
# check that the n_points parameter samples the dataset
|
|
n_points = 1_605 if n_points is None else n_points
|
|
assert X.shape == (n_points, 119)
|
|
assert y.shape == (n_points,)
|
|
|
|
|
|
def test_rank():
|
|
# test that fetch/download works fine
|
|
X1, y1, X2, y2, q1, q2 = shap.datasets.rank()
|
|
|
|
# check the shape of the result
|
|
assert X1.shape == (3_005, 300)
|
|
assert y1.shape == (3_005,)
|
|
assert X2.shape == (768, 300)
|
|
assert y2.shape == (768,)
|
|
assert q1.shape == (201,)
|
|
assert q2.shape == (50,)
|