matplotlib: кольорові смуги та його текстові мітки

Я хотів би створити 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. Як цього можна досягти?

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()

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

Щоб додати відповідь tacaswell , colorbar()функція має необов'язковий caxвхід, який можна використовувати для передачі осі, на якій слід намалювати кольорову панель. Якщо ви використовуєте цей вхід, ви можете безпосередньо встановити мітку за допомогою цієї осі.

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