Matplotlib绘制3D散点图时predictions分类图例不显示问题

直接将分类数值列传给scatter的c参数着色时，matplotlib默认不会自动拆分不同取值生成对应图例项，因此直接调用ax.legend()无法正常渲染分组图例。

修复后完整代码 #

import random
import pandas as pd
from matplotlib import pyplot as plt

random.seed(0)
D = [[random.random() for x in range(3)] for y in range(1000)]
df = pd.DataFrame(D,columns=['Feature_1','Feature_2','Feature_3'])
predictions = [random.randint(0,4) for x in range(1000)]
df['predictions'] = predictions

plt.rcParams["figure.figsize"]=(10,10)
plt.rcParams['legend.fontsize'] = 10
from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
ax = Axes3D(fig)

# 保存scatter返回对象
scatter = ax.scatter(df['Feature_1'],df['Feature_2'],df['Feature_3'], c=df['predictions'], s =150,cmap='rainbow')
# 自动生成分组图例句柄与标签
handles, labels = scatter.legend_elements()
# 显式传入句柄和标签渲染图例，可自定义title
ax.legend(handles=handles, labels=labels, loc = 'upper left', title='predictions')

ax.set_xlabel('Feature_1',fontsize=20,labelpad=10)
ax.set_ylabel('Feature_2', fontsize=20, rotation=150,labelpad=10)
ax.set_zlabel('Feature_3', fontsize=20, rotation=60,labelpad=15)

plt.show()

核心修改说明 #

• 补全缺失的pandas导入语句，避免运行时报错
• 接收ax.scatter返回的PathCollection对象
• 调用scatter.legend_elements()自动生成每个颜色分组对应的图例句柄与标签
• 调用ax.legend()时显式传入句柄与标签，可通过title参数设置图例标题，也可手动修改labels列表自定义图例文本

原异常渲染效果图：

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