如何在Pandas散点图中实现自定义分类颜色条？

我有一个Pandas DataFrame，想要绘制两个变量的散点图，按类别给散点上色，并且已经定义了颜色映射字典。当前代码生成的颜色条不符合需求，需要实现：

• 带有类别颜色和对应标签的离散颜色条
• 沿用自定义的颜色映射字典

当前代码如下：

import pandas as pd
import matplotlib.pyplot as plt

weights = [1.0, 1.1, 1.3, 2.6, 5.1] 
volumes = [2.1, 4.3, 2.6, 2.7, 9.6]
fruits = ['apple', 'banana', 'banana', 'apple', 'coconut']

data_dict = {'weight': weights, 'vol': volumes, 'class': fruits}
df = pd.DataFrame(data_dict)

color_dict = {'apple' : 'r', 'banana' : 'yellow', 'coconut' : 'lime'}

plt.scatter(df['weight'], df['vol'], c = [color_dict[i] for i in df['class'].iloc[:]])
plt.xlabel('Weight')
plt.ylabel('Volume')
plt.colorbar()
plt.show()

补充说明：

• 以下单行代码可在Pandas中实现与上述相同的基础上色效果，但同样无法生成符合需求的分类色条：

df.plot(kind="scatter", x="weight", y="vol", c=df['class'].map(color_dict), colorbar=True)

• 将类别转为分类类型后，Pandas可生成带标签的分类色条，但无法直接传入自定义颜色映射：

df['class'] = df['class'].astype('category')
df.plot(kind="scatter", x="weight", y="vol", c='class')

方法一：直接用Matplotlib实现（灵活可控） #

核心思路是把类别映射为数值，基于自定义颜色字典构建离散色图，再手动将色条的刻度替换为类别标签：

import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

weights = [1.0, 1.1, 1.3, 2.6, 5.1] 
volumes = [2.1, 4.3, 2.6, 2.7, 9.6]
fruits = ['apple', 'banana', 'banana', 'apple', 'coconut']

data_dict = {'weight': weights, 'vol': volumes, 'class': fruits}
df = pd.DataFrame(data_dict)

color_dict = {'apple' : 'r', 'banana' : 'yellow', 'coconut' : 'lime'}

# 获取唯一类别并匹配颜色字典顺序
unique_classes = list(color_dict.keys())
# 构建类别到数值的映射
class_to_num = {cls: i for i, cls in enumerate(unique_classes)}
# 将数据中的类别转为数值
df['class_num'] = df['class'].map(class_to_num)

# 基于颜色字典构建自定义离散色图
cmap = ListedColormap([color_dict[cls] for cls in unique_classes])

# 绘制散点图，用数值作为颜色参数
scatter = plt.scatter(df['weight'], df['vol'], c=df['class_num'], cmap=cmap)

# 设置色条，替换刻度为类别标签
cbar = plt.colorbar(scatter, ticks=range(len(unique_classes)), label='水果类别')
cbar.set_ticklabels(unique_classes)

plt.xlabel('Weight')
plt.ylabel('Volume')
plt.show()

方法二：Pandas结合Matplotlib实现 #

Pandas的plot接口可以配合Matplotlib的色图和归一化工具，实现自定义分类色条：

import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap, BoundaryNorm

weights = [1.0, 1.1, 1.3, 2.6, 5.1] 
volumes = [2.1, 4.3, 2.6, 2.7, 9.6]
fruits = ['apple', 'banana', 'banana', 'apple', 'coconut']

data_dict = {'weight': weights, 'vol': volumes, 'class': fruits}
df = pd.DataFrame(data_dict)

color_dict = {'apple' : 'r', 'banana' : 'yellow', 'coconut' : 'lime'}

# 处理类别和颜色映射
unique_classes = list(color_dict.keys())
class_to_num = {cls: i for i, cls in enumerate(unique_classes)}
df['class_num'] = df['class'].map(class_to_num)

cmap = ListedColormap([color_dict[cls] for cls in unique_classes])
# 定义边界，让每个类别对应一个离散区间
bounds = [i - 0.5 for i in range(len(unique_classes)+1)]
norm = BoundaryNorm(bounds, cmap.N)

# 用Pandas绘图，传入自定义cmap和norm
ax = df.plot(kind="scatter", x="weight", y="vol", c=df['class_num'], cmap=cmap, norm=norm, colorbar=False)
# 手动添加色条并设置标签
scatter = ax.collections[0]
cbar = plt.colorbar(scatter, ticks=range(len(unique_classes)), ax=ax)
cbar.set_ticklabels(unique_classes)
cbar.set_label('水果类别')

plt.xlabel('Weight')
plt.ylabel('Volume')
plt.show()

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