如何修复Python中Plotly 3D散点图的图例问题

使用Python的Plotly绘制3D散点图，数据集包含二分类人格特质（0/1）与5类聚类标签（0-4），需求如下：

• 用颜色区分人格特质
• 用符号区分聚类标签
• 可视化TSNE嵌入结果

现有代码运行后出现渐变图例（不符合分类数据需求），图例混乱，期望仅显示聚类标签对应的符号、以及人格特质对应的离散颜色分类，效果为颜色区分人格特质有无、符号区分聚类归属的清晰图例。

原有代码：

from plotly.subplots import make_subplots
import plotly.express as px

# 替换y/n为数字
df2 = df.replace({'y': 1, 'n': 0})

# 分配聚类标签
df2 = df2.assign(Cluster=labels)

# 定义各人格特质的图表标题
plot_titles = ['Extraversion', 'Neuroticism', 'Agreeableness', 'Conscientiousness', 'Openness']

# 绘制每个人格特质对应的TSNE嵌入图
for i, trait in enumerate(['cEXT', 'cNEU', 'cAGR', 'cCON', 'cOPN']):
    fig = px.scatter_3d(df2,
                        x=embeddings_3d[:, 0], y=embeddings_3d[:, 1], z=embeddings_3d[:, 2],
                        color = df2[trait], symbol=df2['Cluster'],
                        color_discrete_map={0: '#FF0000', 1: '#0000FF'},
                        size_max=1, symbol_map={0: 'circle', 1: 'square', 2: 'diamond', 3: 'cross', 4: 'x'},
                        opacity=0.3)
    fig.update_layout(title=plot_titles[i])
    fig.show()

问题核心是Plotly将数值型的人格特质列识别为连续变量，自动生成渐变图例。只需将人格特质列转换为字符串类型强制按分类变量处理，即可生成离散图例；同时可优化图例布局提升可读性：

修改后的代码：

import plotly.express as px

# 替换y/n为数字
df2 = df.replace({'y': 1, 'n': 0})

# 将人格特质列转换为字符串类型，强制Plotly按分类变量处理
trait_cols = ['cEXT', 'cNEU', 'cAGR', 'cCON', 'cOPN']
df2[trait_cols] = df2[trait_cols].astype(str)

# 分配聚类标签
df2 = df2.assign(Cluster=labels)

# 定义各人格特质的图表标题
plot_titles = ['Extraversion', 'Neuroticism', 'Agreeableness', 'Conscientiousness', 'Openness']

# 绘制每个人格特质对应的TSNE嵌入图
for i, trait in enumerate(trait_cols):
    fig = px.scatter_3d(df2,
                        x=embeddings_3d[:, 0], y=embeddings_3d[:, 1], z=embeddings_3d[:, 2],
                        color=df2[trait], 
                        symbol=df2['Cluster'],
                        # 匹配字符串类型的键值
                        color_discrete_map={'0': '#FF0000', '1': '#0000FF'},
                        size_max=1, 
                        symbol_map={0: 'circle', 1: 'square', 2: 'diamond', 3: 'cross', 4: 'x'},
                        opacity=0.3)
    
    # 优化图例布局，避免遮挡图表
    fig.update_layout(
        title=plot_titles[i],
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=-0.2,
            xanchor="center",
            x=0.5
        )
    )
    
    fig.show()

关键修改说明 #

1. 转换人格特质列为字符串：df2[trait_cols] = df2[trait_cols].astype(str)，强制Plotly将其识别为分类变量，自动使用离散色标，消除渐变图例。
2. 同步颜色映射键值：由于列类型变为字符串，color_discrete_map的键需改为'0'和'1'，确保颜色匹配正确。
3. 优化图例位置：通过update_layout将图例移至底部居中，避免遮挡图表内容。

如果需要仅保留聚类标签的符号图例（隐藏人格特质的颜色图例），可在update_layout后添加以下代码：

# 隐藏所有图例，再仅显示聚类符号的图例
fig.update_layout(showlegend=False)
fig.for_each_trace(lambda t: t.update(showlegend=True) if 'Cluster' in t.name else t.update(showlegend=False))

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