Matplotlib.pyplot如何巧用参数c为不同子矩阵点集匹配不同色图

我当前使用的第三方库内置了一个基于matplotlib.pyplot的自动绘图函数，该函数仅需传入单个numpy.ndarray或pandas.DataFrame即可生成自定义图表。我希望通过灵活调整（即“滥用”）matplotlib.pyplot.plot(..., c=colors, ...)的c参数，实现以下效果：将多个矩阵通过np.vstack拼接后传入该函数时，每个原始子矩阵对应的点集可以使用不同的色表着色。

相关示例代码如下：

import numpy as np

a = np.zeros((10, 2))
b = np.zeros((20, 2))
c = np.zeros((5, 2))

plot_mat = np.vstack([a, b, c])

colors = list(range(a.shape[0])) + list(range(b.shape[0])) + list(range(c.shape[0]))
external_fn(plot_mat, ..., c=colors)

实现方案 #

你可以通过给不同子矩阵分配互不重叠的颜色数值区间，配合自定义分段配色映射来实现需求，具体操作如下：

• 给每个子矩阵设置唯一的数值偏移量，确保不同子矩阵的颜色取值落在完全独立的区间，避免配色混淆
• 用ListedColormap和BoundaryNorm定义你需要的多色表映射规则，每个区间对应一个子矩阵的目标色表
• 同一个子矩阵内部可以保留渐变效果，也可以设置为固定值让整个子矩阵的点颜色统一

可直接运行的参考代码如下：

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap, BoundaryNorm

# 1. 准备测试子矩阵，每个子矩阵加偏移方便区分显示效果
a = np.random.randn(10, 2)
b = np.random.randn(20, 2) + 3
c = np.random.randn(5, 2) + 6
sub_mat_list = [a, b, c]
# 按需设置每个子矩阵对应的色表
sub_cmap_list = [plt.cm.Blues, plt.cm.Oranges, plt.cm.Greens]

# 2. 生成拼接矩阵和颜色数组
max_point_count = max(mat.shape[0] for mat in sub_mat_list)
plot_mat = np.vstack(sub_mat_list)
colors = []
for mat_idx, mat in enumerate(sub_mat_list):
    # 每个子矩阵的颜色值落在 [mat_idx*max_point_count, (mat_idx+1)*max_point_count) 区间
    colors.extend([mat_idx * max_point_count + point_idx for point_idx in range(mat.shape[0])])

# 3. 自定义多色表映射规则
merged_color_list = []
for cmap in sub_cmap_list:
    # 每个色表取max_point_count个色值，0.3到0.9的区间避免过浅或过深看不清
    merged_color_list.extend(cmap(np.linspace(0.3, 0.9, max_point_count)))
custom_cmap = ListedColormap(merged_color_list)
# 定义区间边界，确保数值和色表一一对应
bounds = [i * max_point_count for i in range(len(sub_mat_list) + 1)]
custom_norm = BoundaryNorm(bounds, custom_cmap.N)

# 4. 调用第三方绘图函数，传入自定义的cmap和norm即可生效
# 只要你用的第三方函数支持透传cmap、norm参数给底层的matplotlib接口，就能实现目标效果
external_fn(plot_mat, c=colors, cmap=custom_cmap, norm=custom_norm)

# 本地测试效果可以直接用scatter验证
# plt.scatter(plot_mat[:, 0], plot_mat[:, 1], c=colors, cmap=custom_cmap, norm=custom_norm)
# plt.show()

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