如何在Matplotlib中同一画布的不同子轴并排绘制两个SHAP蜂群图

问题背景 #

SHAP 0.39.0版本的shap.plots.beeswarm函数不支持ax参数，无法直接指定子轴绘制，导致报错TypeError: beeswarm() got an unexpected keyword argument 'ax'。

可行方案 #

通过手动切换matplotlib的当前活动轴，让SHAP绘图函数在指定子轴上渲染：

from matplotlib import pyplot as plt
import shap

# 创建1行2列的画布与子轴
fig, axes = plt.subplots(1, 2, figsize=(35, 7), gridspec_kw={"wspace": 1.0}) 

# 绑定第一个子轴并绘制蜂群图
plt.sca(axes[0])
shap.plots.beeswarm(explainer_a(X_test_a), max_display=10)

# 绑定第二个子轴并绘制蜂群图
plt.sca(axes[1])
shap.plots.beeswarm(explainer_b(X_test_b), max_display=10)

plt.show()

核心原理 #

• plt.sca(ax)会将传入的子轴设置为matplotlib的当前活动轴，所有后续绘图操作都会定向到该轴。
• SHAP的beeswarm函数默认使用当前活动轴，因此通过切换轴的方式即可实现多子轴布局。

内容的提问来源于stack exchange，提问作者Paul