wehub-resource-sync
207 lines
7.4 KiB
Python
207 lines
7.4 KiB
Python
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
|
|
from .. import Explanation
|
|
from ..utils import OpChain
|
|
from . import colors
|
|
from ._labels import labels
|
|
from ._utils import convert_ordering
|
|
|
|
|
|
def heatmap(
|
|
shap_values: Explanation,
|
|
instance_order=Explanation.hclust(), # type: ignore
|
|
feature_values=Explanation.abs.mean(0), # type: ignore
|
|
feature_order=None,
|
|
max_display=10,
|
|
cmap=colors.red_white_blue,
|
|
show=True,
|
|
plot_width=8,
|
|
ax=None,
|
|
):
|
|
"""Create a heatmap plot of a set of SHAP values.
|
|
|
|
This plot is designed to show the population substructure of a dataset using supervised
|
|
clustering and a heatmap.
|
|
Supervised clustering involves clustering data points not by their original
|
|
feature values but by their explanations.
|
|
By default, we cluster using :func:`shap.utils.hclust_ordering`,
|
|
but any clustering can be used to order the samples.
|
|
|
|
Parameters
|
|
----------
|
|
shap_values : shap.Explanation
|
|
A multi-row :class:`.Explanation` object that we want to visualize in a
|
|
cluster ordering.
|
|
|
|
instance_order : OpChain or numpy.ndarray
|
|
A function that returns a sort ordering given a matrix of SHAP values and an axis, or
|
|
a direct sample ordering given as an ``numpy.ndarray``.
|
|
|
|
feature_values : OpChain or numpy.ndarray
|
|
A function that returns a global summary value for each input feature, or an array of such values.
|
|
|
|
feature_order : None, OpChain, or numpy.ndarray
|
|
A function that returns a sort ordering given a matrix of SHAP values and an axis, or
|
|
a direct input feature ordering given as an ``numpy.ndarray``.
|
|
If ``None``, then we use ``feature_values.argsort``.
|
|
|
|
max_display : int
|
|
The maximum number of features to display (default is 10).
|
|
|
|
show : bool
|
|
Whether :external+mpl:func:`matplotlib.pyplot.show()` is called before returning.
|
|
Setting this to ``False`` allows the plot
|
|
to be customized further after it has been created.
|
|
|
|
plot_width : int, default 8
|
|
The width of the heatmap plot.
|
|
|
|
ax : matplotlib Axes
|
|
Axes object to draw the plot onto, otherwise uses the current Axes.
|
|
|
|
Returns
|
|
-------
|
|
ax: matplotlib Axes
|
|
Returns the :external+mpl:class:`~matplotlib.axes.Axes` object with the plot drawn onto it.
|
|
|
|
Examples
|
|
--------
|
|
See `heatmap plot examples <https://shap.readthedocs.io/en/latest/example_notebooks/api_examples/plots/heatmap.html>`_.
|
|
|
|
"""
|
|
# sort the SHAP values matrix by rows and columns
|
|
values = shap_values.values
|
|
if issubclass(type(feature_values), OpChain):
|
|
feature_values = feature_values.apply(Explanation(values))
|
|
if issubclass(type(feature_values), Explanation):
|
|
feature_values = feature_values.values
|
|
if feature_order is None:
|
|
feature_order = np.argsort(-feature_values)
|
|
elif issubclass(type(feature_order), OpChain):
|
|
feature_order = feature_order.apply(Explanation(values))
|
|
elif not hasattr(feature_order, "__len__"):
|
|
raise Exception(f"Unsupported feature_order: {str(feature_order)}!")
|
|
xlabel = "Instances"
|
|
instance_order = convert_ordering(instance_order, shap_values)
|
|
# if issubclass(type(instance_order), OpChain):
|
|
# #xlabel += " " + instance_order.summary_string("SHAP values")
|
|
# instance_order = instance_order.apply(Explanation(values))
|
|
# elif not hasattr(instance_order, "__len__"):
|
|
# raise Exception("Unsupported instance_order: %s!" % str(instance_order))
|
|
# else:
|
|
# instance_order_ops = None
|
|
|
|
feature_names = np.array(shap_values.feature_names)[feature_order]
|
|
values = shap_values.values[instance_order][:, feature_order]
|
|
feature_values = feature_values[feature_order]
|
|
|
|
# if we have more features than `max_display`, then group all the excess features
|
|
# into a single feature
|
|
if values.shape[1] > max_display:
|
|
new_values = np.zeros((values.shape[0], max_display))
|
|
new_values[:, :-1] = values[:, : max_display - 1]
|
|
new_values[:, -1] = values[:, max_display - 1 :].sum(1)
|
|
new_feature_values = np.zeros(max_display)
|
|
new_feature_values[:-1] = feature_values[: max_display - 1]
|
|
new_feature_values[-1] = feature_values[max_display - 1 :].sum()
|
|
feature_names = [
|
|
*feature_names[: max_display - 1],
|
|
f"Sum of {values.shape[1] - max_display + 1} other features",
|
|
]
|
|
values = new_values
|
|
feature_values = new_feature_values
|
|
|
|
# define the plot size based on how many features we are plotting
|
|
row_height = 0.5
|
|
if ax is None:
|
|
plt.gcf().set_size_inches(plot_width, values.shape[1] * row_height + 2.5)
|
|
ax = plt.gca()
|
|
|
|
# plot the matrix of SHAP values as a heat map
|
|
vmin, vmax = np.nanpercentile(values.flatten(), [1, 99])
|
|
ax.imshow(
|
|
values.T,
|
|
aspect=0.7 * values.shape[0] / values.shape[1],
|
|
interpolation="nearest",
|
|
vmin=min(vmin, -vmax),
|
|
vmax=max(-vmin, vmax),
|
|
cmap=cmap,
|
|
)
|
|
|
|
# adjust the axes ticks and spines for the heat map + f(x) line chart
|
|
ax.xaxis.set_ticks_position("bottom")
|
|
ax.yaxis.set_ticks_position("left")
|
|
ax.spines[["left", "right"]].set_visible(True)
|
|
ax.spines[["left", "right"]].set_bounds(values.shape[1] - row_height, -row_height)
|
|
ax.spines[["top", "bottom"]].set_visible(False)
|
|
ax.tick_params(axis="both", direction="out")
|
|
|
|
ax.set_ylim(values.shape[1] - row_height, -3)
|
|
heatmap_yticks_pos = np.arange(values.shape[1])
|
|
heatmap_yticks_labels = feature_names
|
|
ax.yaxis.set_ticks(
|
|
[-1.5, *heatmap_yticks_pos],
|
|
[r"$f(x)$", *heatmap_yticks_labels],
|
|
fontsize=13,
|
|
)
|
|
# remove the y-tick line for the f(x) label
|
|
ax.yaxis.get_ticklines()[0].set_visible(False)
|
|
|
|
ax.set_xlim(-0.5, values.shape[0] - 0.5)
|
|
ax.set_xlabel(xlabel)
|
|
|
|
# plot the f(x) line chart above the heat map
|
|
ax.axhline(-1.5, color="#aaaaaa", linestyle="--", linewidth=0.5)
|
|
fx = values.T.sum(0)
|
|
fx_max = np.abs(fx).max()
|
|
if fx_max > 0:
|
|
fx_normalized = fx / fx_max
|
|
else:
|
|
fx_normalized = fx # all zeros, flat line — no division needed
|
|
ax.plot(
|
|
-fx_normalized - 1.5,
|
|
color="#000000",
|
|
linewidth=1,
|
|
)
|
|
|
|
# plot the bar plot on the right spine of the heat map
|
|
fv_max = np.abs(feature_values).max()
|
|
bar_widths = (feature_values / fv_max) * values.shape[0] / 20 if fv_max > 0 else np.zeros_like(feature_values)
|
|
bar_container = ax.barh(
|
|
heatmap_yticks_pos,
|
|
bar_widths,
|
|
height=0.7,
|
|
align="center",
|
|
color="#000000",
|
|
left=values.shape[0] * 1.0 - 0.5,
|
|
)
|
|
for b in bar_container:
|
|
b.set_clip_on(False)
|
|
|
|
# draw the color bar
|
|
import matplotlib.cm as cm
|
|
|
|
m = cm.ScalarMappable(cmap=cmap)
|
|
m.set_array([min(vmin, -vmax), max(-vmin, vmax)])
|
|
cb = plt.colorbar(
|
|
m,
|
|
ticks=[min(vmin, -vmax), max(-vmin, vmax)],
|
|
ax=ax,
|
|
aspect=80,
|
|
fraction=0.01,
|
|
pad=0.10, # padding between the cb and the main axes
|
|
)
|
|
cb.set_label(labels["VALUE"], size=12, labelpad=-10)
|
|
cb.ax.tick_params(labelsize=11, length=0)
|
|
cb.set_alpha(1)
|
|
cb.outline.set_visible(False) # type: ignore
|
|
# bbox = cb.ax.get_window_extent().transformed(plt.gcf().dpi_scale_trans.inverted())
|
|
# cb.ax.set_aspect((bbox.height - 0.9) * 15)
|
|
# cb.draw_all()
|
|
|
|
if show:
|
|
plt.show()
|
|
|
|
return ax
|