Matplotlib热图指定行使用不同cmap配色的实现问题问询

原代码错误原因 #

• ax.imshow()返回的是Matplotlib的AxesImage绘图对象，并非数值数组，不能直接调用np.vstack拼接
• 两次调用imshow默认从坐标轴0行索引位置开始绘制，后绘制的内容会直接覆盖先绘制的内容，无法实现分区分行配色的效果

修正后可运行代码 #

import numpy as np
import matplotlib.pyplot as plt

y_ticks = ["f", "e", "d", "c",
           "b", "a", "baseline"]
x_ticks = ["A", "B", "C",
           "D", "E", "F", "G"]

data = np.array([[0.8, 2.4, 2.5, 3.9, 0.0, 4.0, 0.0],
                 [2.4, 0.0, 4.0, 1.0, 2.7, 0.0, 0.0],
                 [1.1, 2.4, 0.8, 4.3, 1.9, 4.4, 0.0],
                 [0.6, 0.0, 0.3, 0.0, 3.1, 0.0, 0.0],
                 [0.7, 1.7, 0.6, 2.6, 2.2, 6.2, 0.0],
                 [1.3, 1.2, 0.0, 0.0, 0.0, 3.2, 5.1],
                 [0.1, 2.0, 0.0, 1.4, 0.0, 1.9, 6.3]])

# 定位baseline所在行索引
baseline_idx = y_ticks.index("baseline")
# 统一配色数值映射范围，保证两组颜色深浅的对比逻辑一致
vmin, vmax = data.min(), data.max()

fig, ax = plt.subplots()

# 构造仅baseline行有值、其余位置为NaN的数组，用Blues渲染，NaN位置透明不覆盖
blues_data = np.full_like(data, np.nan)
blues_data[baseline_idx] = data[baseline_idx]
ax.imshow(blues_data, cmap='Blues', vmin=vmin, vmax=vmax)

# 构造除baseline外其余行有值、baseline位置为NaN的数组，用Reds渲染
reds_data = np.full_like(data, np.nan)
reds_data[:baseline_idx] = data[:baseline_idx]
ax.imshow(reds_data, cmap='Reds', vmin=vmin, vmax=vmax)

ax.set_xticks(np.arange(len(x_ticks)))
ax.set_yticks(np.arange(len(y_ticks)))

ax.set_xticklabels(x_ticks)
ax.set_yticklabels(y_ticks)

plt.setp(ax.get_xticklabels(), ha="right", rotation=45, rotation_mode="anchor")

# 循环添加数值标注
for i in range(len(y_ticks)):
    for j in range(len(x_ticks)):
        text = ax.text(j, i, data[i, j], ha="center", va="center", color="black")

fig.tight_layout()
plt.show()

核心逻辑说明 #

• 利用NaN值在imshow渲染时透明的特性，将不同配色规则的行拆分为两个独立的掩码数组，分别渲染互不覆盖
• 显式指定统一的vmin和vmax，避免两组配色各自独立映射数值导致的颜色对比逻辑混乱
• 动态获取baseline行索引，后续调整y轴标签顺序时无需手动修改硬编码的索引值

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