如何在matplotlib的pcolormesh中为指定像素设置自定义颜色？

问题背景 #

用户使用matplotlib.pcolormesh绘制热力图，原始代码如下：

import numpy as np
import matplotlib.pyplot as plt

# 生成随机数组
array = np.random.rand(10,10)
x = np.arange(0,10)
y = np.arange(0,10)

fig, ax = plt.subplots(nrows=1,ncols=1)
colormesh = ax.pcolormesh(x,y,array)
ax.set_xticks(x)
ax.set_yticks(y)
cb_ax = fig.add_axes([0.93, 0.1, 0.02, 0.8])
cbar = fig.colorbar(colormesh,cax=cb_ax)
colormesh.set_clim(0,1)
plt.show()

随后将部分像素值设为-1：

# 选择要着色的像素
array[5,5] = -1
array[0,7] = -1

需求：将值为-1的像素整体染成红色，且不改变色条范围（保持[0,1]）；因原始数据x、y非等间距，需保留使用pcolormesh而非imshow，且需填充整个像素区域而非仅添加点。

解决方案 #

方法1：掩码数组标记坏值并指定颜色 #

利用numpy掩码数组将-1的元素标记为"坏值"，通过色图的set_bad()方法指定坏值显示为红色，同时维持色条的[0,1]范围。

完整代码：

import numpy as np
import matplotlib.pyplot as plt

# 生成数据并设置目标值
array = np.random.rand(10,10)
array[5,5] = -1
array[0,7] = -1

# 创建掩码数组：将-1的元素设为掩码（坏值）
masked_array = np.ma.masked_where(array == -1, array)

x = np.arange(0,10)
y = np.arange(0,10)

fig, ax = plt.subplots(nrows=1,ncols=1)
# 使用掩码数组绘制热力图
colormesh = ax.pcolormesh(x, y, masked_array)
# 设置坏值颜色为红色
colormesh.cmap.set_bad('red')

ax.set_xticks(x)
ax.set_yticks(y)

# 添加色条，保持范围0-1
cb_ax = fig.add_axes([0.93, 0.1, 0.02, 0.8])
cbar = fig.colorbar(colormesh, cax=cb_ax)
colormesh.set_clim(0,1)

plt.show()

方法2：手动绘制红色矩形覆盖目标像素 #

直接在目标像素对应的矩形区域绘制红色填充块，适合需要更灵活控制单个像素样式的场景。

完整代码：

import numpy as np
import matplotlib.pyplot as plt

# 生成数据并记录目标像素位置
array = np.random.rand(10,10)
target_pixels = [(5,5), (0,7)]
for i, j in target_pixels:
    array[i,j] = -1

x = np.arange(0,10)
y = np.arange(0,10)

fig, ax = plt.subplots(nrows=1,ncols=1)
# 正常绘制热力图，设置色条范围0-1（-1值会被截断）
colormesh = ax.pcolormesh(x, y, array)
colormesh.set_clim(0,1)

# 遍历目标像素，绘制红色矩形
for i, j in target_pixels:
    # 每个像素对应x[j]~x[j+1]、y[i]~y[i+1]的矩形区域
    rect = plt.Rectangle(
        (x[j], y[i]), 
        width=x[j+1]-x[j], 
        height=y[i+1]-y[i], 
        color='red'
    )
    ax.add_patch(rect)

ax.set_xticks(x)
ax.set_yticks(y)

# 添加色条
cb_ax = fig.add_axes([0.93, 0.1, 0.02, 0.8])
cbar = fig.colorbar(colormesh, cax=cb_ax)

plt.show()

方案对比 #

• 方法1代码更简洁，适合批量处理同类型特殊值的像素，颜色管理更统一；
• 方法2灵活性更高，可单独调整每个目标像素的颜色、透明度等样式；
• 两种方案均不会压缩[0,1]区间数据的颜色展示范围，符合需求。

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