Matplotlib: colorbars и его текстовые метки

Я хотел бы создать colorbar легенду для heatmap, таким образом, что метки находятся в центре каждого дискретного цвета. Пример заимствован отсюда:

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

#discrete color scheme
cMap = ListedColormap(['white', 'green', 'blue','red'])

#data
np.random.seed(42)
data = np.random.rand(4, 4)
fig, ax = plt.subplots()
heatmap = ax.pcolor(data, cmap=cMap)

#legend
cbar = plt.colorbar(heatmap)
cbar.ax.set_yticklabels(['0','1','2','>3'])
cbar.set_label('# of contacts', rotation=270)

# put the major ticks at the middle of each cell
ax.set_xticks(np.arange(data.shape[1]) + 0.5, minor=False)
ax.set_yticks(np.arange(data.shape[0]) + 0.5, minor=False)
ax.invert_yaxis()

#labels
column_labels = list('ABCD')
row_labels = list('WXYZ')
ax.set_xticklabels(column_labels, minor=False)
ax.set_yticklabels(row_labels, minor=False)

plt.show()

Это создает следующий сюжет:

В идеале я хотел бы создать панель легенды, которая имеет четыре цвета и для каждого цвета метку в центре: 0,1,2,>3. Как этого достичь?

Ответ 1

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

#discrete color scheme
cMap = ListedColormap(['white', 'green', 'blue','red'])

#data
np.random.seed(42)
data = np.random.rand(4, 4)
fig, ax = plt.subplots()
heatmap = ax.pcolor(data, cmap=cMap)

#legend
cbar = plt.colorbar(heatmap)

cbar.ax.get_yaxis().set_ticks([])
for j, lab in enumerate(['$0$','$1$','$2$','$>3$']):
    cbar.ax.text(.5, (2 * j + 1) / 8.0, lab, ha='center', va='center')
cbar.ax.get_yaxis().labelpad = 15
cbar.ax.set_ylabel('# of contacts', rotation=270)


# put the major ticks at the middle of each cell
ax.set_xticks(np.arange(data.shape[1]) + 0.5, minor=False)
ax.set_yticks(np.arange(data.shape[0]) + 0.5, minor=False)
ax.invert_yaxis()

#labels
column_labels = list('ABCD')
row_labels = list('WXYZ')
ax.set_xticklabels(column_labels, minor=False)
ax.set_yticklabels(row_labels, minor=False)

plt.show()

Вы были очень близки. Когда у вас есть ссылка на ось цветных баров, вы можете делать все, что захотите, включая размещение текстовых меток посередине. Возможно, вы захотите сыграть с форматированием, чтобы сделать его более заметным.

Ответ 2

Чтобы добавить ответ colorbar() функция colorbar() имеет дополнительный ввод cax который можно использовать для прохождения оси, на которой должен быть нарисован colorbar. Если вы используете этот вход, вы можете напрямую установить метку, используя эту ось.

import matplotlib.pyplot as plt
from mpl_toolkits.axes.grid1 import make_axes_locatable

fig, ax = plt.subplots()
heatmap = ax.imshow(data)
divider = make_axes_locatable(ax)
cax = divider.append_axes('bottom', size='10%', pad=0.6)
cb = fig.colorbar(heatmap, cax=cax, orientation='horizontal')

cax.set_xlabel('data label')  # cax == cb.ax