项目文件夹

文件
2026-07-13 13:22:52 +08:00

121 行
3.7 KiB
Python

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pytest
from sklearn.neighbors import KNeighborsClassifier
from sklearn.tree import DecisionTreeRegressor
import shap
from shap.plots import _style
def test_waterfall_input_is_explanation():
"""Checks an error is raised if a non-Explanation object is passed as input."""
with pytest.raises(
TypeError,
match="waterfall plot requires an `Explanation` object",
):
_ = shap.plots.waterfall(np.random.randn(20, 5), show=False)
def test_waterfall_wrong_explanation_shape(explainer):
explanation = explainer(explainer.data)
emsg = "waterfall plot can currently only plot a single explanation"
with pytest.raises(ValueError, match=emsg):
shap.plots.waterfall(explanation, show=False)
@pytest.mark.mpl_image_compare(tolerance=3)
def test_waterfall(explainer):
"""Test the new waterfall plot."""
fig = plt.figure()
explanation = explainer(explainer.data)
shap.plots.waterfall(explanation[0], show=False)
plt.tight_layout()
return fig
@pytest.mark.mpl_image_compare(tolerance=3)
def test_waterfall_legacy(explainer):
"""Test the old waterfall plot."""
shap_values = explainer.shap_values(explainer.data)
fig = plt.figure()
shap.plots._waterfall.waterfall_legacy(explainer.expected_value, shap_values[0], show=False)
plt.tight_layout()
return fig
@pytest.mark.mpl_image_compare(tolerance=3)
def test_waterfall_bounds(explainer):
"""Test the waterfall plot with upper and lower error bounds plotted."""
fig = plt.figure()
explanation = explainer(explainer.data)
explanation._s.lower_bounds = explanation.values - 0.1
explanation._s.upper_bounds = explanation.values + 0.1
shap.plots.waterfall(explanation[0])
plt.tight_layout()
return fig
@pytest.mark.mpl_image_compare(tolerance=5)
def test_waterfall_custom_style(explainer):
"""Test the waterfall plot in the context of custom styles"""
# Note: the tolerance is set to 5 because matplotlib 3.10 changed the way negative values are displayed
# There is now an increased space before the negative sign, which leads to a RMS diff of ~4.4
# See: GH #3946
# TODO: reset tolerance to 3 when python 3.9 is dropped, and all tests use matplotlib 3.10+
with _style.style_context(
primary_color_positive="#9ACD32",
primary_color_negative="#FFA500",
text_color="black",
hlines_color="red",
vlines_color="red",
tick_labels_color="red",
):
fig = plt.figure()
explanation = explainer(explainer.data)
shap.plots.waterfall(explanation[0])
plt.tight_layout()
return fig
def test_waterfall_plot_for_decision_tree_explanation():
# Regression tests for GH #3129
X = pd.DataFrame({"A": [1, 2, 3], "B": [2, 1, 3]})
y = pd.Series([1, 2, 3])
model = DecisionTreeRegressor()
model.fit(X, y)
explainer = shap.TreeExplainer(model)
explanation = explainer(X)
shap.plots.waterfall(explanation[0], show=False)
def test_waterfall_plot_for_data_with_number_columns():
# GH 4150
model = KNeighborsClassifier()
def f(x):
return model.predict_proba(x)[:, 1]
X = pd.DataFrame(
[
[1, 0, 0, 0, 0],
[0, 1, 0, 0, 0],
[0, 0, 1, 0, 0],
[0, 0, 0, 1, 0],
[0, 0, 0, 0, 1],
[0, 0, 0, 1, 0],
[0, 0, 1, 0, 0],
[0, 1, 0, 0, 0],
]
)
y = pd.Series([0, 1, 2, 3, 4, 3, 2, 1])
model.fit(X, y)
med = X.median().values.reshape((1, X.shape[1]))
explainer = shap.Explainer(f, med)
shap_values = explainer(X)
shap.plots.waterfall(shap_values[0], show=False)