如何为Matplotlib绘制的1D热力图添加颜色条？

基于一段实现1D热力图+折线图的Matplotlib代码，需要为顶部的热力图添加对应颜色数值的参考颜色条，但原代码中ax.imshow没有直接选项，也不清楚ax的对应方法。使用环境为Matplotlib 3.8.0、Python 3.10.12。

解决方案 #

Matplotlib中imshow()函数会返回一个ScalarMappable对象，我们可以基于这个对象调用fig.colorbar()生成对应颜色条。为了让颜色条和热力图高度匹配，可使用make_axes_locator调整布局，避免颜色条与图表重叠。

修改后的完整代码 #

import matplotlib.pyplot as plt
from matplotlib import ticker
import numpy as np; np.random.seed(1)
plt.rcParams["figure.figsize"] = 5, 3

x = np.arange(0,99)
y = np.random.uniform(-10, 10, 99)

fig, (ax, ax2) = plt.subplots(nrows=2, sharex=True)

extent = [x[0]-(x[1]-x[0])/2., x[-1]+(x[1]-x[0])/2., 0, 1]
# 保存imshow返回的颜色映射对象
im = ax.imshow(y[np.newaxis,:], cmap="gist_rainbow", aspect="auto", extent=extent)

ax.set_yticks([])
ax.set_xlim(extent[0], extent[1])

# 创建与热力图同高度的颜色条轴
from matplotlib.pyplot import make_axes_locator
divider = make_axes_locator(ax)
cax = divider.append_axes("right", size="5%", pad=0.05)
# 绑定颜色条到热力图的颜色映射
fig.colorbar(im, cax=cax)

ax2.plot(x,y)

plt.tight_layout()
plt.show()

关键改动说明 #

• 将ax.imshow()的返回值赋值给变量im，该对象包含热力图的颜色映射规则
• 使用make_axes_locator生成与热力图高度一致的颜色条容器，避免颜色条拉伸变形
• 通过fig.colorbar(im, cax=cax)完成颜色条与热力图的绑定

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