shap--shap/shap/utils/_clustering.py

from __future__ import annotations

import itertools as it
import warnings
from typing import TYPE_CHECKING, Any, Literal

import numpy as np
import numpy.typing as npt
import pandas as pd
import scipy.cluster
import scipy.spatial
import sklearn
from numba import njit

from ..utils._exceptions import DimensionError
from ._show_progress import show_progress

if TYPE_CHECKING:
    from ._types import _ArrayLike


def partition_tree(X: pd.DataFrame, metric: str = "correlation") -> npt.NDArray[Any]:
    X_full_rank = X + np.random.randn(*X.shape) * 1e-8
    D = scipy.spatial.distance.pdist(X_full_rank.fillna(X_full_rank.mean()).T, metric=metric)
    return scipy.cluster.hierarchy.complete(D)


def partition_tree_shuffle(
    indexes: npt.NDArray[Any],
    index_mask: npt.NDArray[np.bool_],
    partition_tree: npt.NDArray[Any],
) -> None:
    """Randomly shuffle the indexes in a way that is consistent with the given partition tree.

    Parameters
    ----------
    indexes: np.array
        The output location of the indexes we want shuffled. Note that len(indexes) should equal index_mask.sum().
    index_mask: np.array
        A bool mask of which indexes we want to include in the shuffled list.
    partition_tree: np.array
        The partition tree we should follow.

    """
    M = len(index_mask)
    # switch = np.random.randn(M) < 0
    _pt_shuffle_rec(partition_tree.shape[0] - 1, indexes, index_mask, partition_tree, M, 0)


@njit
def _pt_shuffle_rec(
    i: int,
    indexes: npt.NDArray[Any],
    index_mask: npt.NDArray[np.bool_],
    partition_tree: npt.NDArray[Any],
    M: int,
    pos: int,
) -> int:  # type: ignore[misc]
    if i < 0:
        # see if we should include this index in the ordering
        if index_mask[i + M]:
            indexes[pos] = i + M
            return pos + 1
        else:
            return pos
    left = int(partition_tree[i, 0] - M)
    right = int(partition_tree[i, 1] - M)
    if np.random.randn() < 0:
        pos = _pt_shuffle_rec(left, indexes, index_mask, partition_tree, M, pos)
        pos = _pt_shuffle_rec(right, indexes, index_mask, partition_tree, M, pos)
    else:
        pos = _pt_shuffle_rec(right, indexes, index_mask, partition_tree, M, pos)
        pos = _pt_shuffle_rec(left, indexes, index_mask, partition_tree, M, pos)
    return pos


@njit
def delta_minimization_order(
    all_masks: npt.NDArray[Any],
    max_swap_size: int = 100,
    num_passes: int = 2,
) -> npt.NDArray[Any]:  # type: ignore[misc]
    order = np.arange(len(all_masks))
    for _ in range(num_passes):
        for length in list(range(2, max_swap_size)):
            for i in range(1, len(order) - length):
                if _reverse_window_score_gain(all_masks, order, i, length) > 0:
                    _reverse_window(order, i, length)
    return order


@njit
def _reverse_window(order: npt.NDArray[Any], start: int, length: int) -> None:  # type: ignore[misc]
    for i in range(length // 2):
        tmp = order[start + i]
        order[start + i] = order[start + length - i - 1]
        order[start + length - i - 1] = tmp


@njit
def _reverse_window_score_gain(
    masks: npt.NDArray[Any],
    order: npt.NDArray[Any],
    start: int,
    length: int,
) -> int:  # type: ignore[misc]
    forward_score = _mask_delta_score(masks[order[start - 1]], masks[order[start]]) + _mask_delta_score(
        masks[order[start + length - 1]], masks[order[start + length]]
    )
    reverse_score = _mask_delta_score(masks[order[start - 1]], masks[order[start + length - 1]]) + _mask_delta_score(
        masks[order[start]], masks[order[start + length]]
    )

    return forward_score - reverse_score


@njit
def _mask_delta_score(m1: npt.NDArray[Any], m2: npt.NDArray[Any]) -> int:  # type: ignore[misc]
    return (m1 ^ m2).sum()


def hclust_ordering(
    X: npt.NDArray[Any],
    metric: str = "sqeuclidean",
    anchor_first: bool = False,
) -> npt.NDArray[Any]:
    """A leaf ordering is under-defined, this picks the ordering that keeps nearby samples similar."""
    # compute a hierarchical clustering and return the optimal leaf ordering
    D = scipy.spatial.distance.pdist(X, metric)
    cluster_matrix = scipy.cluster.hierarchy.complete(D)
    return scipy.cluster.hierarchy.leaves_list(scipy.cluster.hierarchy.optimal_leaf_ordering(cluster_matrix, D))


def xgboost_distances_r2(
    X: npt.NDArray[Any],
    y: npt.NDArray[Any],
    learning_rate: float = 0.6,
    early_stopping_rounds: int | None = 2,
    subsample: float | None = 1.0,
    max_estimators: int | None = 10_000,
    random_state: int | np.random.RandomState = 0,
) -> npt.NDArray[Any]:
    """Compute redundancy distances scaled from 0-1 among all the features in X relative to the label y.

    Distances are measured by training univariate XGBoost models of y for all the features, and then
    predicting the output of these models using univariate XGBoost models of other features. If one
    feature can effectively predict the output of another feature's univariate XGBoost model of y,
    then the second feature is redundant with the first with respect to y. A distance of 1 corresponds
    to no redundancy while a distance of 0 corresponds to perfect redundancy (measured using the
    proportion of variance explained). Note these distances are not symmetric.

    Returns
    -------
    np.ndarray
        A square matrix of shape (n_features, n_features) containing the pairwise
        redundancy distances between features. Each element [i, j] represents the
        redundancy distance from feature i to feature j with respect to y.

    """
    import xgboost

    # pick our train/text split
    X_train, X_test, y_train, y_test = sklearn.model_selection.train_test_split(X, y, random_state=random_state)

    # fit an XGBoost model on each of the features
    num_features = X.shape[1]
    train_preds_list: list[npt.NDArray[Any]] = []
    test_preds_list: list[npt.NDArray[Any]] = []
    for i in range(num_features):
        model = xgboost.XGBRegressor(
            subsample=subsample,
            n_estimators=max_estimators,
            learning_rate=learning_rate,
            max_depth=1,
            early_stopping_rounds=early_stopping_rounds,
        )
        model.fit(X_train[:, i : i + 1], y_train, eval_set=[(X_test[:, i : i + 1], y_test)], verbose=False)
        train_preds_list.append(model.predict(X_train[:, i : i + 1]))
        test_preds_list.append(model.predict(X_test[:, i : i + 1]))
    train_preds: npt.NDArray[Any] = np.vstack(train_preds_list).T
    test_preds: npt.NDArray[Any] = np.vstack(test_preds_list).T

    # fit XGBoost models to predict the outputs of other XGBoost models to see how redundant features are
    dist: npt.NDArray[Any] = np.zeros((num_features, num_features))
    for i, j in show_progress(
        it.product(range(num_features), range(num_features)),
        total=num_features * num_features,
    ):
        if i == j:
            continue

        # skip features that have no variance in their predictions (likely because the feature is a constant)
        preds_var: float = np.var(test_preds[:, i])
        if preds_var < 1e-4:
            warnings.warn(
                f"No/low signal found from feature {i} (this is typically caused by constant or "
                "near-constant features)! Cluster distances can't be computed for it (so setting "
                "all redundancy distances to 1)."
            )
            r2: float = 0

        # fit the model
        else:
            model = xgboost.XGBRegressor(
                subsample=subsample,
                n_estimators=max_estimators,
                learning_rate=learning_rate,
                max_depth=1,
                early_stopping_rounds=early_stopping_rounds,
            )
            model.fit(
                X_train[:, j : j + 1],
                train_preds[:, i],
                eval_set=[(X_test[:, j : j + 1], test_preds[:, i])],
                verbose=False,
            )
            r2 = max(0, 1 - np.mean((test_preds[:, i] - model.predict(X_test[:, j : j + 1])) ** 2) / preds_var)
        dist[i, j] = 1 - r2

    return dist


def hclust(
    X: _ArrayLike,
    y: _ArrayLike | None = None,
    linkage: Literal["single", "complete", "average"] = "single",
    metric: str = "auto",
    random_state: int | np.random.RandomState = 0,
) -> npt.NDArray[Any]:
    """Fit a hierarchical clustering model for features X relative to target variable y.

    For more information on clustering methods, see :external+scipy:func:`scipy.cluster.hierarchy.linkage`.

    For more information on scipy distance metrics, see :external+scipy:func:`scipy.spatial.distance.pdist`.

    Parameters
    ----------
    X: 2d-array-like
        Features to cluster
    y: array-like or None
        Target variable
    linkage: str
        Defines the method to calculate the distance between clusters. Must be
        one of "single", "complete" or "average".
    metric: str
        Scipy distance metric or "xgboost_distances_r2".

        * If ``xgboost_distances_r2``, estimate redundancy distances between
          features X with respect to target variable y using
          :func:`shap.utils.xgboost_distances_r2`.
        * Otherwise, calculate distances between features using the given
          distance metric.
        * If ``auto`` (default), use ``xgboost_distances_r2`` if target variable
          is provided, or else ``cosine`` distance metric.
    random_state: int or np.random.RandomState
        Numpy random state, defaults to 0.

    Returns
    -------
    clustering: np.array
        The hierarchical clustering encoded as a linkage matrix.

    """
    if isinstance(X, pd.DataFrame):
        X_arr: npt.NDArray[Any] = X.values
    else:
        X_arr = np.array(X)
    if len(X_arr.shape) != 2:
        raise DimensionError("X needs to be a 2-dimensional array-like object")

    known_linkages: tuple[str, str, str] = ("single", "complete", "average")
    if linkage not in known_linkages:
        raise ValueError(f"Unknown linkage type: {linkage}")

    if metric == "auto":
        if y is not None:
            metric = "xgboost_distances_r2"
        else:
            metric = "cosine"

    # build the distance matrix
    if metric == "xgboost_distances_r2":
        dist_full: npt.NDArray[Any] = xgboost_distances_r2(X_arr, y, random_state=random_state)  # type: ignore[arg-type]

        # build a condensed upper triangular version by taking the max distance from either direction
        dist_list: list[float] = []
        for i, j in it.combinations(range(len(dist_full)), 2):
            if linkage == "single":
                dist_list.append(min(dist_full[i, j], dist_full[j, i]))
            elif linkage == "complete":
                dist_list.append(max(dist_full[i, j], dist_full[j, i]))
            elif linkage == "average":
                dist_list.append((dist_full[i, j] + dist_full[j, i]) / 2)
        dist: npt.NDArray[Any] = np.array(dist_list)

    else:
        if y is not None:
            warnings.warn(
                "Ignoring the y argument passed to shap.utils.hclust since the given clustering metric is "
                "not based on label fitting!"
            )
        bg_no_nan: npt.NDArray[Any] = X_arr.copy()
        for i in range(bg_no_nan.shape[1]):
            np.nan_to_num(bg_no_nan[:, i], nan=np.nanmean(bg_no_nan[:, i]), copy=False)
        dist = scipy.spatial.distance.pdist(bg_no_nan.T + np.random.randn(*bg_no_nan.T.shape) * 1e-8, metric=metric)

    # build linkage
    if linkage == "single":
        return scipy.cluster.hierarchy.single(dist)
    elif linkage == "complete":
        return scipy.cluster.hierarchy.complete(dist)
    elif linkage == "average":
        return scipy.cluster.hierarchy.average(dist)