Matplotlib使用彩色数字作为散点符号绘图并添加对应图例的问题

适配现有annotate实现的方案 #

核心逻辑 #

自定义Matplotlib的图例渲染处理器，让图例符号位置直接渲染对应颜色的数字，完全匹配图中标记样式。

步骤代码 #

1. 导入依赖模块

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.legend_handler import HandlerBase
from matplotlib.text import Text

2. 定义自定义图例处理器

class TextLegendHandler(HandlerBase):
    def create_artists(self, legend, orig_handle, xdescent, ydescent, width, height, fontsize, trans):
        num, color = orig_handle
        return [Text(x=width/2, y=height/2, text=str(num), color=color, fontsize=fontsize, ha='center', va='center')]

3. 在你现有绘图代码后添加图例生成逻辑

# 替换为你自己的标签列表
legend_labels = [f"分类 {n[i]}" for i in range(len(n))]
# 构造图例句柄
legend_handles = [(n[i], x_y_colours[i]) for i in range(len(n))]
# 生成图例
axs[0].legend(
    legend_handles, 
    legend_labels,
    handler_map={tuple: TextLegendHandler()},
    handlelength=1.5, # 调整符号宽度，数字位数多可以适当调大
    fontsize=10
)

完整可运行测试代码 #

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.legend_handler import HandlerBase
from matplotlib.text import Text

class TextLegendHandler(HandlerBase):
    def create_artists(self, legend, orig_handle, xdescent, ydescent, width, height, fontsize, trans):
        num, color = orig_handle
        return [Text(x=width/2, y=height/2, text=str(num), color=color, fontsize=fontsize, ha='center', va='center')]

# 模拟测试数据，可替换为你自己的真实数据
np.random.seed(42)
x = np.random.rand(13)
y = np.random.rand(13)
x_y_colours = plt.cm.tab20(np.arange(13))
n = np.arange(1,14,1)

# 你的原有绘图逻辑
fig, axs = plt.subplots(1, 2, figsize=(12,5))
axs[0].scatter(x, y, linestyle='None', color="white")
for i, txt in enumerate(n):
     axs[0].annotate(txt, (x[i], y[i]), color=x_y_colours[i], ha="center", va="center")

# 生成图例
legend_labels = [f"分类 {n[i]}" for i in range(len(n))]
legend_handles = [(n[i], x_y_colours[i]) for i in range(len(n))]
axs[0].legend(
    legend_handles, 
    legend_labels,
    handler_map={tuple: TextLegendHandler()},
    handlelength=1.5,
    fontsize=10
)

plt.tight_layout()
plt.show()

更简洁的替代实现方案 #

无需循环添加annotate，直接通过scatter的marker参数指定数字作为标记，原生支持生成对应图例：

import numpy as np
import matplotlib.pyplot as plt

# 模拟数据
np.random.seed(42)
x = np.random.rand(13)
y = np.random.rand(13)
x_y_colours = plt.cm.tab20(np.arange(13))
n = np.arange(1,14,1)

fig, axs = plt.subplots(1, 2, figsize=(12,5))
scatter_handles = []
for i in range(len(n)):
    # 直接用数字作为marker，$包裹是为了正确渲染字符
    h = axs[0].scatter(x[i], y[i], marker=f'${n[i]}$', color=x_y_colours[i], s=100)
    scatter_handles.append(h)

# 直接生成图例
legend_labels = [f"分类 {n[i]}" for i in range(len(n))]
axs[0].legend(scatter_handles, legend_labels, fontsize=10)

plt.tight_layout()
plt.show()

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