web-dev-qa-db-fra.com

matplotlib: les barres de couleur et ses étiquettes de texte

J'aimerais créer une légende colorbar pour un heatmap, de sorte que les étiquettes se trouvent au centre de chaque couleur distincte. Exemple emprunté à partir d'ici :

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

Cela génère l'intrigue suivante:

pmesh plot

Idéalement, j'aimerais créer une barre de légende comportant les quatre couleurs et, pour chaque couleur, une étiquette en son centre: 0,1,2,>3. Comment cela peut il etre accompli?

pythonmatplotliblabellegendcolorbar
76
dimka
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()

Tu étais très proche. Une fois que vous avez une référence à l'axe de la barre de couleur, vous pouvez faire ce que vous voulez, y compris mettre des étiquettes de texte au milieu. Vous voudrez peut-être jouer avec la mise en forme pour la rendre plus visible.

demo

84
tacaswell

Pour ajouter à réponse de tacaswell , la fonction colorbar() a une entrée facultative cax que vous pouvez utiliser pour transmettre un axe sur lequel la barre de couleur doit être dessinée. Si vous utilisez cette entrée, vous pouvez directement définir une étiquette à l'aide de cet axe.

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
1
Steven C. Howell