Custom features in beeswarm plot of shap

Question

I have a causal inference model with

featurizer=PolynomialFeatures(degree=3)

which includes a degree 3 polynomial in X variable. I get the plot for interpretability too with:

shap.plots.beeswarm(shap_values['Y0']['T0'])

but the plot shows me the shap values for X0, X0^2 and X0^3 as:

However, is it possible that the beeswarm plot show me only the values for X0?

Best regards

Answer 1

X0 and X0^3 are the column 0 and 2 of the array named main_effects inside the object T0, the subplot with these features could be:

ind = [0, 2]

raw_shap_values = np.take(shap_values['Y0']['T0'].main_effects, ind, axis=1)
raw_data_values = np.take(shap_values['Y0']['T0'].data, ind, axis=1)
shap.summary_plot(raw_shap_values, features=raw_data_values, plot_type="violin")