I'd like to create a colorbar legend for a heatmap, such that the labels are in the center of each discrete color. Example borrowed from here:
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()
This generates the following plot:
Ideally I'd like to generate a legend bar which has the four colors and for each color, a label in its center: 0,1,2,>3. How can this be achieved?
This question is related to
python
matplotlib
label
legend
colorbar
The answer is
To add to tacaswell's answer, the colorbar() function has an optional cax input you can use to pass an axis on which the colorbar should be drawn. If you are using that input, you can directly set a label using that axis.
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
Source: Stackoverflow.com