如何使用matplotlib.pyplot.imshow实现基于列的colormap色阶映射

修改思路 #

默认的imshow会对整个矩阵做全局色阶归一，要实现每列单独归一，只需：

• 对原始计数矩阵按列做min-max归一，得到归一化后的矩阵用于颜色映射
• 单元格文本标注仍保留原始的真实计数值，不影响数值读取

完整可运行代码 #

import numpy as np
import matplotlib.pyplot as plt

x = np.round(np.array([-1, 0,     1,  2, 3, 2, 0, 0, 1, 2, 3]))
y = np.round(np.array([-1, 1.2, -3.1, 2, 3, 2, 1, 1, 1, 3, 3]))

xrange = np.arange(np.min(x),np.max(x) + 1)
yrange = np.arange(np.min(y),np.max(y) + 1)
a = np.zeros((len(yrange), len(xrange)))
for i in range(len(x)):
    a[yrange == y[i],xrange == x[i]] = a[yrange == y[i],xrange == x[i]] + 1

# 新增按列归一逻辑
a_norm = np.zeros_like(a, dtype=np.float32)
col_max = a.max(axis=0, keepdims=True)
# 过滤全0列避免除以0报错
valid_col_mask = col_max[0] > 0
a_norm[:, valid_col_mask] = a[:, valid_col_mask] / col_max[:, valid_col_mask]

fig, ax = plt.subplots()
# 传入归一化后的矩阵做颜色映射
im = ax.imshow(a_norm, cmap='Wistia', origin="lower")
# 文本标注仍使用原始计数值
for i in range(len(yrange)):
    for j in range(len(xrange)):
        text = ax.text(j, i, str(int(a[i, j])),
                       ha="center", va="center", color="k")
# 可选优化：设置坐标轴刻度与实际x、y值对应
ax.set_xticks(np.arange(len(xrange)))
ax.set_xticklabels(xrange)
ax.set_yticks(np.arange(len(yrange)))
ax.set_yticklabels(yrange)
ax.set(xlabel = 'Input', ylabel = 'Output')
fig.tight_layout()
plt.show()

效果说明 #

修改后每一列的最大值都会映射为colormap的最深色，最小值映射为最浅色，不同列之间的色阶独立，完全符合要求的映射逻辑。

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