Beeswarm plot colorbar is too narrow on jupyter notebooks

[jagmoreira]

Hello there, I noticed that on my jupyter notebook I can barely see the colorbars of beeswarm plots. I believe the issue was fixed in this PR in the original shap package:

Shap beeswarm

FastTreeSHAP beeswarm

Code to reproduce in notebook

import xgboost
import shap
import fasttreeshap
from sklearn.model_selection import train_test_split

X,y = shap.datasets.adult()

# create a train/test split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=7)
d_train = xgboost.DMatrix(X_train, label=y_train)
d_test = xgboost.DMatrix(X_test, label=y_test)

# train model
params = {
    "objective": "binary:logistic",
    "eval_metric": "logloss"
}
model = xgboost.train(params, d_train, 200, evals = [(d_test, "test")], verbose_eval=100)

# Calculate shap values
shap_explainer = shap.TreeExplainer(model)
fastshap_explainer = fasttreeshap.TreeExplainer(model)
shap_values = shap_explainer(X_test)
fastshap_values = fastshap_explainer(X_test)

# Create beeswarm plots
shap.plots.beeswarm(shap_values)
fasttreeshap.plots.beeswarm(fastshap_values)

Environment

• python version: 3.9.13
• relevant pip packages:

  fasttreeshap==0.1.3
  matplotlib==3.5.3
  notebook==6.4.12
  numpy==1.21.6
  pandas==1.3.5
  scikit-learn==1.0.2
  shap==0.41.0
  xgboost==1.6.2

[jlyang1990]

Thanks @jagmoreira for pointing out this issue! I have pushed a new commit https://github.com/linkedin/FastTreeSHAP/commit/b08edb76b1a06c4b2a2ad3c9042fedc2ee373473 to fix it. Let me know if this issue still exists.