93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
|
|
from . import colors
|
|
|
|
|
|
def group_difference(
|
|
shap_values,
|
|
group_mask,
|
|
feature_names=None,
|
|
xlabel=None,
|
|
xmin=None,
|
|
xmax=None,
|
|
max_display=None,
|
|
sort=True,
|
|
show=True,
|
|
ax=None,
|
|
):
|
|
"""This plots the difference in mean SHAP values between two groups.
|
|
|
|
It is useful to decompose many group level metrics about the model output among the
|
|
input features. Quantitative fairness metrics for machine learning models are
|
|
a common example of such group level metrics.
|
|
|
|
Parameters
|
|
----------
|
|
shap_values : numpy.array
|
|
Matrix of SHAP values (# samples x # features) or a vector of model outputs (# samples).
|
|
|
|
group_mask : numpy.array
|
|
A boolean mask where True represents the first group of samples and False the second.
|
|
|
|
feature_names : list
|
|
A list of feature names.
|
|
|
|
"""
|
|
# Compute confidence bounds for the group difference value
|
|
vs = []
|
|
gmean = group_mask.mean()
|
|
for _ in range(200):
|
|
r = np.random.rand(shap_values.shape[0]) > gmean
|
|
vs.append(shap_values[r].mean(0) - shap_values[~r].mean(0))
|
|
vs_ = np.array(vs)
|
|
xerr = np.vstack([np.percentile(vs_, 95, axis=0), np.percentile(vs_, 5, axis=0)])
|
|
|
|
# See if we were passed a single model output vector and not a matrix of SHAP values
|
|
if len(shap_values.shape) == 1:
|
|
shap_values = shap_values.reshape(1, -1).T
|
|
if feature_names is None:
|
|
feature_names = [""]
|
|
|
|
# Fill in any missing feature names
|
|
if feature_names is None:
|
|
feature_names = [f"Feature {i}" for i in range(shap_values.shape[1])]
|
|
|
|
diff = shap_values[group_mask].mean(0) - shap_values[~group_mask].mean(0)
|
|
|
|
if sort is True:
|
|
inds = np.argsort(-np.abs(diff)).astype(int)
|
|
else:
|
|
inds = np.arange(len(diff))
|
|
|
|
if max_display is not None:
|
|
inds = inds[:max_display]
|
|
if ax:
|
|
# Disable plotting out if an ax has been provided
|
|
show = False
|
|
else:
|
|
# Draw the figure if no ax has been provided
|
|
figsize = (6.4, 0.2 + 0.9 * len(inds))
|
|
_, ax = plt.subplots(figsize=figsize)
|
|
ticks = range(len(inds) - 1, -1, -1)
|
|
ax.axvline(0, color="#999999", linewidth=0.5)
|
|
ax.barh(ticks, diff[inds], color=colors.blue_rgb, capsize=3, xerr=np.abs(xerr[:, inds]))
|
|
|
|
for i in range(len(inds)):
|
|
ax.axhline(y=i, color="#cccccc", lw=0.5, dashes=(1, 5), zorder=-1)
|
|
|
|
ax.xaxis.set_ticks_position("bottom")
|
|
ax.yaxis.set_ticks_position("none")
|
|
ax.set_yticks(ticks)
|
|
ax.set_yticklabels([feature_names[i] for i in inds], fontsize=13)
|
|
ax.spines["right"].set_visible(False)
|
|
ax.spines["top"].set_visible(False)
|
|
ax.spines["left"].set_visible(False)
|
|
ax.tick_params(labelsize=11)
|
|
if xlabel is None:
|
|
xlabel = "Group SHAP value difference"
|
|
ax.set_xlabel(xlabel, fontsize=13)
|
|
ax.set_xlim(xmin, xmax)
|
|
if show:
|
|
plt.show()
|