ENH: Label dots in scatterplot according to classes and add a legend

[sroener]

Problem Description

Hi,

I decided to interpret a few of my models using shap. To get a good grasp on the most important features, I want to plot them using shap.scatter and see how the shap values are based on the feature value distribution. To make these plots even more meaningful, I would like to add color to the dots according to their respective class labels.

Alternative Solutions

I tried producing an example based on one of the official examples for the scatterplot:

import sklearn.linear_model
import shap

X, y = shap.datasets.adult()
model = sklearn.linear_model.LogisticRegression(max_iter = 10000).fit(X,y)

explainer = shap.Explainer(lambda x: model.predict_proba(x)[:, 1], X)
shap_values = explainer(X[:1000])

shap.plots.scatter(shap_values[:, "Age"], color=y[:1000]) #limit y to 1000, equal to the shap_values 

The following plot shows the output of the code:

This solution kind of works, but abusing the color option feels a bit hacky and the colorbar representing the classes is also suboptimal.