276 lines
9.5 KiB
Python
276 lines
9.5 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
|
|
from .. import Explanation
|
|
from ..utils import OpChain
|
|
from . import colors
|
|
|
|
if TYPE_CHECKING:
|
|
from matplotlib.colors import Colormap
|
|
|
|
|
|
def convert_color(color: str | np.ndarray) -> np.ndarray | Colormap | str:
|
|
"""Converts a color specification alias into its actual representation"""
|
|
if isinstance(color, np.ndarray):
|
|
return color
|
|
elif isinstance(color, str) and color == "shap_red":
|
|
return colors.red_rgb
|
|
elif isinstance(color, str) and color == "shap_blue":
|
|
return colors.blue_rgb
|
|
else:
|
|
try:
|
|
return plt.get_cmap(color)
|
|
except ValueError:
|
|
return color
|
|
|
|
|
|
def convert_ordering(ordering, shap_values):
|
|
if issubclass(type(ordering), OpChain):
|
|
ordering = ordering.apply(Explanation(shap_values))
|
|
if issubclass(type(ordering), Explanation):
|
|
if any(op.name == "argsort" for op in ordering.op_history):
|
|
ordering = ordering.values
|
|
else:
|
|
ordering = ordering.argsort.flip.values
|
|
return ordering
|
|
|
|
|
|
def get_sort_order(dist, clust_order, cluster_threshold, feature_order):
|
|
"""Returns a sorted order of the values where we respect the clustering order when dist[i,j] < cluster_threshold"""
|
|
# feature_imp = np.abs(values)
|
|
|
|
# if partition_tree is not None:
|
|
# new_tree = fill_internal_max_values(partition_tree, shap_values)
|
|
# clust_order = sort_inds(new_tree, np.abs(shap_values))
|
|
clust_inds = np.argsort(clust_order)
|
|
|
|
feature_order = feature_order.copy() # order.apply(Explanation(shap_values))
|
|
# print("feature_order", feature_order)
|
|
for i in range(len(feature_order) - 1):
|
|
ind1 = feature_order[i]
|
|
next_ind = feature_order[i + 1]
|
|
next_ind_pos = i + 1
|
|
for j in range(i + 1, len(feature_order)):
|
|
ind2 = feature_order[j]
|
|
|
|
# if feature_imp[ind] >
|
|
# if ind1 == 2:
|
|
# print(ind1, ind2, dist[ind1,ind2])
|
|
if dist[ind1, ind2] <= cluster_threshold:
|
|
# if ind1 == 2:
|
|
# print(clust_inds)
|
|
# print(ind1, ind2, next_ind, dist[ind1,ind2], clust_inds[ind2], clust_inds[next_ind])
|
|
if dist[ind1, next_ind] > cluster_threshold or clust_inds[ind2] < clust_inds[next_ind]:
|
|
next_ind = ind2
|
|
next_ind_pos = j
|
|
# print("next_ind", next_ind)
|
|
# print("next_ind_pos", next_ind_pos)
|
|
|
|
# insert the next_ind next
|
|
for j in range(next_ind_pos, i + 1, -1):
|
|
# print("j", j)
|
|
feature_order[j] = feature_order[j - 1]
|
|
feature_order[i + 1] = next_ind
|
|
# print(feature_order)
|
|
|
|
return feature_order
|
|
|
|
|
|
def merge_nodes(values, partition_tree):
|
|
"""This merges the two clustered leaf nodes with the smallest total value."""
|
|
M = partition_tree.shape[0] + 1
|
|
|
|
ptind = 0
|
|
min_val = np.inf
|
|
for i in range(partition_tree.shape[0]):
|
|
ind1 = int(partition_tree[i, 0])
|
|
ind2 = int(partition_tree[i, 1])
|
|
if ind1 < M and ind2 < M:
|
|
val = np.abs(values[ind1]) + np.abs(values[ind2])
|
|
if val < min_val:
|
|
min_val = val
|
|
ptind = i
|
|
# print("ptind", ptind, min_val)
|
|
|
|
ind1 = int(partition_tree[ptind, 0])
|
|
ind2 = int(partition_tree[ptind, 1])
|
|
if ind1 > ind2:
|
|
tmp = ind1
|
|
ind1 = ind2
|
|
ind2 = tmp
|
|
|
|
partition_tree_new = partition_tree.copy()
|
|
for i in range(partition_tree_new.shape[0]):
|
|
i0 = int(partition_tree_new[i, 0])
|
|
i1 = int(partition_tree_new[i, 1])
|
|
if i0 == ind2:
|
|
partition_tree_new[i, 0] = ind1
|
|
elif i0 > ind2:
|
|
partition_tree_new[i, 0] -= 1
|
|
if i0 == ptind + M:
|
|
partition_tree_new[i, 0] = ind1
|
|
elif i0 > ptind + M:
|
|
partition_tree_new[i, 0] -= 1
|
|
|
|
if i1 == ind2:
|
|
partition_tree_new[i, 1] = ind1
|
|
elif i1 > ind2:
|
|
partition_tree_new[i, 1] -= 1
|
|
if i1 == ptind + M:
|
|
partition_tree_new[i, 1] = ind1
|
|
elif i1 > ptind + M:
|
|
partition_tree_new[i, 1] -= 1
|
|
partition_tree_new = np.delete(partition_tree_new, ptind, axis=0)
|
|
|
|
# update the counts to be correct
|
|
fill_counts(partition_tree_new)
|
|
|
|
return partition_tree_new, ind1, ind2
|
|
|
|
|
|
def dendrogram_coords(leaf_positions, partition_tree):
|
|
"""Returns the x and y coords of the lines of a dendrogram where the leaf order is given.
|
|
|
|
Note that scipy can compute these coords as well, but it does not allow you to easily specify
|
|
a specific leaf order, hence this reimplementation.
|
|
"""
|
|
xout: list[list[float]] = []
|
|
yout: list[list[float]] = []
|
|
_dendrogram_coords_rec(partition_tree.shape[0] - 1, leaf_positions, partition_tree, xout, yout)
|
|
|
|
return np.array(xout), np.array(yout)
|
|
|
|
|
|
def _dendrogram_coords_rec(pos, leaf_positions, partition_tree, xout, yout):
|
|
M = partition_tree.shape[0] + 1
|
|
|
|
if pos < 0:
|
|
return leaf_positions[pos + M], 0
|
|
|
|
left = int(partition_tree[pos, 0]) - M
|
|
right = int(partition_tree[pos, 1]) - M
|
|
|
|
x_left, y_left = _dendrogram_coords_rec(left, leaf_positions, partition_tree, xout, yout)
|
|
x_right, y_right = _dendrogram_coords_rec(right, leaf_positions, partition_tree, xout, yout)
|
|
|
|
y_curr = partition_tree[pos, 2]
|
|
|
|
xout.append([x_left, x_left, x_right, x_right])
|
|
yout.append([y_left, y_curr, y_curr, y_right])
|
|
|
|
return (x_left + x_right) / 2, y_curr
|
|
|
|
|
|
def fill_internal_max_values(partition_tree, leaf_values):
|
|
"""This fills the forth column of the partition tree matrix with the max leaf value in that cluster."""
|
|
M = partition_tree.shape[0] + 1
|
|
new_tree = partition_tree.copy()
|
|
for i in range(new_tree.shape[0]):
|
|
val = 0
|
|
if new_tree[i, 0] < M:
|
|
ind = int(new_tree[i, 0])
|
|
val = max(val, np.abs(leaf_values[ind]))
|
|
else:
|
|
ind = int(new_tree[i, 0]) - M
|
|
val = max(val, np.abs(new_tree[ind, 3])) # / partition_tree[ind,2])
|
|
if new_tree[i, 1] < M:
|
|
ind = int(new_tree[i, 1])
|
|
val = max(val, np.abs(leaf_values[ind]))
|
|
else:
|
|
ind = int(new_tree[i, 1]) - M
|
|
val = max(val, np.abs(new_tree[ind, 3])) # / partition_tree[ind,2])
|
|
new_tree[i, 3] = val
|
|
return new_tree
|
|
|
|
|
|
def fill_counts(partition_tree):
|
|
"""This updates the"""
|
|
M = partition_tree.shape[0] + 1
|
|
for i in range(partition_tree.shape[0]):
|
|
val = 0
|
|
if partition_tree[i, 0] < M:
|
|
ind = int(partition_tree[i, 0])
|
|
val += 1
|
|
else:
|
|
ind = int(partition_tree[i, 0]) - M
|
|
val += partition_tree[ind, 3]
|
|
if partition_tree[i, 1] < M:
|
|
ind = int(partition_tree[i, 1])
|
|
val += 1
|
|
else:
|
|
ind = int(partition_tree[i, 1]) - M
|
|
val += partition_tree[ind, 3]
|
|
partition_tree[i, 3] = val
|
|
|
|
|
|
def sort_inds(partition_tree, leaf_values, pos=None, inds=None):
|
|
if inds is None:
|
|
inds = []
|
|
|
|
if pos is None:
|
|
partition_tree = fill_internal_max_values(partition_tree, leaf_values)
|
|
pos = partition_tree.shape[0] - 1
|
|
|
|
M = partition_tree.shape[0] + 1
|
|
|
|
if pos < 0:
|
|
inds.append(pos + M)
|
|
return
|
|
|
|
left = int(partition_tree[pos, 0]) - M
|
|
right = int(partition_tree[pos, 1]) - M
|
|
|
|
left_val = partition_tree[left, 3] if left >= 0 else leaf_values[left + M]
|
|
right_val = partition_tree[right, 3] if right >= 0 else leaf_values[right + M]
|
|
|
|
if left_val < right_val:
|
|
tmp = right
|
|
right = left
|
|
left = tmp
|
|
|
|
sort_inds(partition_tree, leaf_values, left, inds)
|
|
sort_inds(partition_tree, leaf_values, right, inds)
|
|
|
|
return inds
|
|
|
|
|
|
# Various ways to specify a desired axis limit in plots
|
|
type AxisLimitSpec = Explanation | str | float | None
|
|
|
|
|
|
def parse_axis_limit(ax_limit: AxisLimitSpec, ax_values: np.ndarray, *, is_shap_axis: bool) -> float | None:
|
|
"""Handle axis limits in "percentile(float)" format or from Explanation objects.
|
|
|
|
Parameters
|
|
----------
|
|
ax_limit : LimitSpec
|
|
Represents one of the lower or upper bounds of an xlim / ylim of a plot.
|
|
Can be in "percentile(float)" format, a float or an :class:`.Explanation` object that has been aggregated
|
|
into a single value, e.g. ``explanation[:, "feature_name"].percentile(20)``.
|
|
|
|
ax_values : np.ndarray
|
|
The values represented by the axis in question. Usually the SHAP values or the feature values. Only used if
|
|
``ax_limit`` is a string of the "percentile(float)" form.
|
|
|
|
is_shap_axis : bool
|
|
Whether the ``ax_limit`` is describing the axis representing SHAP values, or not. Only relevant when
|
|
``ax_limit`` is an :class:`.Explanation` object. It is assumed that when False, the axis is representing
|
|
the feature values instead of the SHAP values.
|
|
|
|
"""
|
|
if isinstance(ax_limit, str):
|
|
try:
|
|
percentage = float(ax_limit.removeprefix("percentile(").removesuffix(")"))
|
|
except ValueError as e:
|
|
raise ValueError("Only strings of the format `percentile(x)` are supported.") from e
|
|
return np.nanpercentile(ax_values, percentage)
|
|
if isinstance(ax_limit, Explanation):
|
|
# Extract relevant attributes, depending on whether the axis is a SHAP-value axis or not
|
|
return float(ax_limit.values) if is_shap_axis else float(ax_limit.data)
|
|
# Else, should be float or None
|
|
return ax_limit
|