如何在Matplotlib中为基于DataFrame的3D散点图添加自定义图例？

我来帮你搞定这个3D散点图的图例问题～你之前遇到的报错，核心原因是没有给每个类别生成独立的绘图句柄（handle），Matplotlib找不到能对应到图例标签的元素，自然就报错了。

为什么之前的方法行不通？ #

你一开始把所有数据点一次性用scatter画出来，这时候只会生成一个全局的handle，而且你没给它设置label参数，所以plt.legend()根本找不到可以关联的标签信息。第二次尝试直接传标签文字，但还是没有对应的handle，所以同样报错。

正确的解决方案：为每个类别单独绘制散点 #

要生成自定义图例，我们需要给每个数字标签对应的类别单独画散点，这样每个类别都会有自己的handle，再给每个handle绑定自定义的文字标签，最后就能正常生成图例了。

完整代码示例

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
# %matplotlib notebook

# 第一步：定义数字标签到自定义文字的映射关系
label_map = {0: 'G', 1: 'A', 2: 'B', 3: 'C', 4: 'Test'}

# 创建3D绘图对象
threedee = plt.figure(figsize=(10,10)).gca(projection='3d')

# 第二步：循环每个唯一的数字标签，单独绘制对应类别的散点
for label_num in y["y_num"].unique():
    # 筛选当前标签对应的数据行
    mask = y["y_num"] == label_num
    pc1_data = df_pca_test.loc[mask, "PC1"]
    pc2_data = df_pca_test.loc[mask, "PC2"]
    pc3_data = df_pca_test.loc[mask, "PC3"]
    
    # 绘制散点，并设置label为自定义文字
    threedee.scatter(pc1_data, pc2_data, pc3_data, label=label_map[label_num])

# 设置坐标轴标签
threedee.set_xlabel('PC1', fontsize=12)
threedee.set_ylabel('PC2', fontsize=12)
threedee.set_zlabel('PC3', fontsize=12)

# 生成图例
plt.legend(scatterpoints=1, loc='lower left', fontsize=8)
plt.show()

关键说明 #

• 句柄（handle）到底是什么？ 简单说就是每个绘图元素的“身份标识”，比如这里每个类别的散点组就是一个handle。图例需要把这些标识和对应的文字标签绑定，才能正确显示每个类别代表什么。
• 我们通过label_map把数字标签（0-4）映射成你想要的文字（G、A、B、C、Test），确保图例显示的是你需要的内容。
• 循环y["y_num"].unique()可以自动覆盖所有存在的标签，不用担心漏画某个类别。

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