Matplotlib：为散点图的标记与颜色添加对应图例

我尝试用下方DataFrame绘制散点图，要求用不同标记（圆形、方形等）区分不同Stage，不同颜色（红色、蓝色等）区分不同Product。目前已实现标记和颜色的设置，但无法生成对应的图例。

我的代码 #

import pandas as pd
from sklearn.preprocessing import LabelEncoder
import numpy as np
import matplotlib.pyplot as plt

df = pd.DataFrame([[1500,24,'open','drive'],
                   [2900, 30, 'open', 'walk'],
                   [1200, 50, 'closed', 'drive'],
                   [4000, 80, 'open', 'air'],
                   [8000, 70, 'ongoing', 'air'],
                   [6100, 40, 'ongoing', 'walk'],
                   [7200, 85, 'closed', 'drive'],
                   [3300, 25, 'closed', 'drive'],
                   [5400, 45, 'open', 'walk'],
                   [5900, 53, 'open', 'air']])
df.columns = ['Cost','Duration','Stage','Product']

label_encoder = LabelEncoder()
markers = {0: 'o', 1: 's', 2: '^'}
df['Product_encoded'] = label_encoder.fit_transform(df['Product'])
df['Stage_encoded'] = label_encoder.fit_transform(df['Stage'])
df['Stage_encoded']= df['Stage_encoded'].map(markers)
colors = ('red', 'blue', 'lightgreen', 'gray', 'cyan')

X= np.array(df)
for idx,cl in enumerate(np.unique(df['Stage_encoded'])):
    plt.scatter(x=X[df['Stage_encoded']== cl,0],y=X[df['Stage_encoded']== cl,1],marker=cl,c=[colors[i] for i in X[df['Stage_encoded'] == cl, 4]])
    plt.legend()

问题原因 #

• 循环内重复调用plt.legend()，每次都会覆盖之前的图例，最终只显示最后一组的内容。
• 没有为Stage（标记）和Product（颜色）单独创建图例条目，matplotlib无法自动识别这种多维度的图例规则。

解决代码及说明 #

直接用原始标签映射标记和颜色，同时创建虚拟散点作为图例手柄，分两次生成图例：

import pandas as pd
import matplotlib.pyplot as plt

df = pd.DataFrame([[1500,24,'open','drive'],
                   [2900, 30, 'open', 'walk'],
                   [1200, 50, 'closed', 'drive'],
                   [4000, 80, 'open', 'air'],
                   [8000, 70, 'ongoing', 'air'],
                   [6100, 40, 'ongoing', 'walk'],
                   [7200, 85, 'closed', 'drive'],
                   [3300, 25, 'closed', 'drive'],
                   [5400, 45, 'open', 'walk'],
                   [5900, 53, 'open', 'air']])
df.columns = ['Cost','Duration','Stage','Product']

# 直接用原始标签定义标记和颜色映射，更直观
stage_markers = {'open': 'o', 'closed': 's', 'ongoing': '^'}
product_colors = {'drive': 'red', 'walk': 'blue', 'air': 'lightgreen'}

# 保存图例需要的手柄对象
stage_handles = []
product_handles = []

# 遍历每个Stage绘制数据
for stage in stage_markers:
    stage_data = df[df['Stage'] == stage]
    marker = stage_markers[stage]
    # 创建虚拟散点作为Stage的图例条目（空数据，只显示标记）
    stage_handle = plt.scatter([], [], marker=marker, color='black', label=stage)
    stage_handles.append(stage_handle)
    
    # 遍历当前Stage下的每个Product
    for product in product_colors:
        product_data = stage_data[stage_data['Product'] == product]
        if not product_data.empty:
            plt.scatter(product_data['Cost'], product_data['Duration'], 
                       marker=marker, color=product_colors[product])
            
# 创建Product的图例条目（虚拟散点，只显示颜色）
for product in product_colors:
    product_handle = plt.scatter([], [], marker='o', color=product_colors[product], label=product)
    product_handles.append(product_handle)

# 生成两个独立图例，分别对应Stage和Product
legend_stage = plt.legend(handles=stage_handles, title='Stage', loc='upper left')
plt.gca().add_artist(legend_stage)  # 保留第一个图例
plt.legend(handles=product_handles, title='Product', loc='lower right')

plt.xlabel('Cost')
plt.ylabel('Duration')
plt.title('Cost vs Duration by Stage and Product')
plt.show()

关键说明 #

• 虚拟散点：通过绘制空数据的散点，单独定义图例的标记或颜色，避免和实际数据点混在一起。
• 双图例设置：用add_artist保留第一个图例，让两个图例分别显示Stage标记和Product颜色，布局更清晰。
• 直接映射标签：放弃LabelEncoder，用原始类别名直接对应标记和颜色，减少编码转换的错误概率。

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