75 lines
2.7 KiB
Python
75 lines
2.7 KiB
Python
from typing import Any
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
import shap
|
|
|
|
|
|
def test_dependence_one_string_feature():
|
|
"""Test the dependence plot with a string feature."""
|
|
X = _create_sample_dataset(string_features={"Sex"})
|
|
|
|
shap.dependence_plot("Sex", np.random.randn(*X.values.shape), X, interaction_index="Age", show=False)
|
|
|
|
|
|
def test_dependence_two_string_features():
|
|
"""Test the dependence plot with two string features."""
|
|
X = _create_sample_dataset(string_features={"Sex", "Blood group"})
|
|
|
|
shap.dependence_plot("Sex", np.random.randn(*X.values.shape), X, interaction_index="Blood group", show=False)
|
|
|
|
|
|
def test_dependence_one_string_feature_no_interaction():
|
|
"""Test the dependence plot with no interactions."""
|
|
X = _create_sample_dataset(string_features={"Sex"})
|
|
|
|
shap.dependence_plot("Sex", np.random.randn(*X.values.shape), X, interaction_index=None, show=False)
|
|
|
|
|
|
def test_dependence_one_string_feature_auto_interaction():
|
|
"""Test the dependence plot with auto interaction detection."""
|
|
X = _create_sample_dataset(string_features={"Sex"})
|
|
|
|
shap.dependence_plot("Sex", np.random.randn(*X.values.shape), X, interaction_index="auto", show=False)
|
|
|
|
|
|
def test_approximate_interactions():
|
|
"""Test the approximate interaction detector."""
|
|
X_no_string_features = _create_sample_dataset(string_features={})
|
|
X_one_string_feature = _create_sample_dataset(string_features={"Sex"})
|
|
X_two_string_features = _create_sample_dataset(string_features={"Sex", "Blood group"})
|
|
|
|
shap_values = np.random.randn(*X_one_string_feature.values.shape)
|
|
|
|
interactions_no_features = shap.approximate_interactions(0, shap_values, X_no_string_features)
|
|
interactions_one_string_feature = shap.approximate_interactions(0, shap_values, X_one_string_feature)
|
|
interactions_two_string_feature = shap.approximate_interactions(0, shap_values, X_two_string_features)
|
|
|
|
assert (interactions_no_features == interactions_one_string_feature).all()
|
|
assert (interactions_no_features == interactions_two_string_feature).all()
|
|
|
|
|
|
def _create_sample_dataset(string_features):
|
|
sex_values: list[Any]
|
|
if "Sex" in string_features:
|
|
sex_values = ["Male", "Female", "Male", "Male", "Female", "Female"]
|
|
else:
|
|
sex_values = [0, 1, 0, 0, 1, 1]
|
|
|
|
blood_values: list[Any]
|
|
if "Blood group" in string_features:
|
|
blood_values = ["A", "B", "A", "O", "O", "O"]
|
|
else:
|
|
blood_values = [1, 2, 1, 3, 3, 3]
|
|
|
|
X = pd.DataFrame(
|
|
{
|
|
"Sex": sex_values,
|
|
"Blood group": blood_values,
|
|
"Age": [10, 15, 28, 3, 84, 56],
|
|
"Height": [130, 170, 185, 40, 150, 164],
|
|
}
|
|
)
|
|
return X
|