3D聚类散点图自动生成图例 解决无标签句柄报错并展示集群信息

问题根因 #

• 原函数未定义clusterSizes入参，内部无法获取聚类规模数据，标签拼接逻辑失效
• 代码中重复调用了两次legend()方法，第二次全局调用会覆盖分支内的图例逻辑，若分支未正确设置标签就会触发No handles with labels found to put in legend.报错
• 使用type() is np.ndarray的类型判断方式不严谨，无法匹配子类继承场景
• 手动循环绘制不同聚类散点的方式容易出现颜色、标签不匹配的问题，也无法适配legend_elements()的自动生成逻辑
• 对DataFrame使用[:, 0]的切片方式会触发索引报错，需要调整为pandas支持的索引写法

修复后完整代码 #

import numpy as np
import matplotlib.pyplot as plt

color_names = ["red", "blue", "yellow", "black", "pink", "purple", "orange"]

def plot_3d_transformed_data(df, title, colors="red", clusterSizes=None):
    ax = plt.figure(figsize=(12,10)).gca(projection='3d')
    
    if isinstance(colors, np.ndarray) and clusterSizes is not None:
        # 一次性绘制所有散点，适配legend_elements自动生成句柄逻辑
        scatter = ax.scatter(df.iloc[:, 0], df.iloc[:, 1], df.iloc[:, 2], 
                            marker="x", c=colors, 
                            cmap=plt.cm.get_cmap('tab10', len(clusterSizes)))
        # 自动获取图例句柄和原始标签
        handles, labels = scatter.legend_elements()
        # 自定义标签，拼接聚类ID和对应样本规模
        custom_labels = []
        for lab in labels:
            # 提取聚类ID数字
            cluster_id = int(lab.split('{')[1].split('}')[0])
            custom_labels.append(f"Cluster {cluster_id} | 样本数: {clusterSizes[cluster_id]}")
        ax.legend(handles, custom_labels, bbox_to_anchor=(1.04,1), loc="upper left")
    else:
        # 单颜色绘制分支逻辑
        ax.scatter(df.Type, df.Length, df.Freq, alpha=0.6, c=colors, marker="x")
        if clusterSizes is not None:
            ax.legend([str(clusterSizes)], bbox_to_anchor=(1.04,1), loc="upper left")

    ax.set_xlabel("PC1: Type")
    ax.set_ylabel("PC2: Length")
    ax.set_zlabel("PC3: Frequency")
    ax.set_title(title)
    
    plt.show()

调用示例 #

plot_3d_transformed_data(pdf_km_pred,
                        f'Clustering rare URL parameters for data of date: {DATE_FROM}  \nMethod: KMeans over PCA \nn_clusters={n_clusters} , Distance_Measure={DistanceMeasure}',
                        colors=pdf_km_pred.prediction_km.values,
                        clusterSizes=clusterSizes)

效果说明 #

• 图例自动生成，不会出现无句柄报错
• 单个图例同时展示聚类标签和对应样本数量，符合需求
• 图例固定在画布右侧，不会遮挡3D散点图内容
• 如果需要使用自定义的color_names配色，只需要将cmap参数替换为colors=color_names即可

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