优化PCA相关圈函数：自动生成数字标签并添加对应图例

直接上修改完成的display_circles函数，满足自动生成数字标签、添加带列名的图例，同时解决NameError问题：

import matplotlib.pyplot as plt
import numpy as np

def display_circles(pcs, n_comp, pca, labels=None, label_rotation=0, figsize=(10,10)):
    # 初始化图例标签列表
    legend_labels = []
    
    # 处理labels参数：如果传入DataFrame，自动生成数字标签和图例文本
    if labels is not None:
        if hasattr(labels, 'columns'):  # 判断是否为DataFrame
            # 生成从1开始的数字标签
            num_labels = [str(i+1) for i in range(len(labels.columns))]
            # 构建"数字, 列名"格式的图例文本
            legend_labels = [f"{num}, {col}" for num, col in zip(num_labels, labels.columns)]
            # 替换labels为数字标签，用于绘图标注
            labels = num_labels
        else:
            # 若传入的是列表/数组，直接使用原标签作为图例
            legend_labels = labels
    
    fig, ax = plt.subplots(figsize=figsize)
    
    # 绘制单位参考圆
    circle = plt.Circle((0,0), 1, color='gray', fill=False)
    ax.add_artist(circle)
    
    # 绘制主成分轴
    for i in range(n_comp):
        ax.arrow(0, 0, pcs[i,0], pcs[i,1], head_width=0.03, head_length=0.03, color='red', linewidth=2)
        ax.text(pcs[i,0] + 0.05, pcs[i,1] + 0.05, f'PC{i+1}', fontsize=12, color='red')
    
    # 绘制变量相关箭头
    for i, (x, y) in enumerate(pca.components_[:n_comp].T):
        ax.arrow(0, 0, x, y, head_width=0.02, head_length=0.02, color='blue', linewidth=1)
        if labels is not None:
            ax.text(x + 0.03, y + 0.03, labels[i], fontsize=10, ha='center', va='center')
    
    # 添加图例：用虚拟线条匹配变量箭头样式
    if legend_labels:
        dummy_line = plt.Line2D([0], [0], color='blue', linewidth=1, marker='>', markersize=8)
        ax.legend([dummy_line], legend_labels, loc='upper left', bbox_to_anchor=(1,1))
    
    # 设置坐标轴样式
    ax.set_xlim(-1.2, 1.2)
    ax.set_ylim(-1.2, 1.2)
    ax.set_xlabel(f'PC1 ({round(pca.explained_variance_ratio_[0]*100, 2)}%)')
    ax.set_ylabel(f'PC2 ({round(pca.explained_variance_ratio_[1]*100, 2)}%)')
    ax.grid(True, linestyle='--', alpha=0.7)
    plt.title('PCA Correlation Circle')
    plt.xticks(rotation=label_rotation)
    plt.show()

关键改动说明 #

• 自动生成数字标签：

  • 函数内部判断labels是否为DataFrame（通过检查columns属性），如果是，自动生成从1开始的数字序列作为标注文本，同时构建对应列名的图例内容。
  • 若传入的是普通列表/数组，直接沿用原标签内容。

• 添加带列名的图例：

  • 创建与变量箭头样式一致的虚拟线条，将"数字, 列名"格式的文本绑定到图例，放置在图外右上角避免遮挡绘图区域。

• 解决NameError问题：

  • 所有绘图相关变量（如fig、ax、circle）均在函数内部定义，避免引用外部未声明的变量。
  • 确保依赖模块（matplotlib.pyplot、numpy）在函数开头导入，防止调用时出现模块未找到的错误。

使用示例 #

假设你有一个名为df的DataFrame，包含待分析的变量，调用方式如下：

from sklearn.decomposition import PCA

# 实例化PCA并拟合数据
pca = PCA(n_components=2)
pca.fit(df)
pcs = pca.components_

# 直接传入df作为labels参数，自动生成标签和图例
display_circles(pcs, n_comp=2, pca=pca, labels=df)

调用后会生成带数字标注的相关圈，右侧图例清晰展示每个数字对应的变量列名，无需手动在函数外预处理标签。

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