CODEWITHSUNDEEP
pythonmatplotlibheatmapcolorbar
ahajib
ahajib

Reputation: 13500

plot heatmap colorbar

I am working on a plot which looks like this so far:

enter image description here

And I am using the following code:

import matplotlib.pyplot as plt
import numpy as np
column_labels = ['A','B']
row_labels = ['A','B']
data = np.random.rand(2,2)
print(type(data))
data = np.array([[1, 0.232], [0.119, 1]])
fig, ax = plt.subplots()
heatmap = ax.pcolor(data, cmap=plt.cm.Blues)

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

# want a more natural, table-like display
ax.invert_yaxis()
ax.xaxis.tick_top()

ax.set_xticklabels(row_labels, minor=False)
ax.set_yticklabels(column_labels, minor=False)
plt.show()

What I am trying to get is something like this one:

enter image description here

Where I can see the color bar as well as value of the cells on them. What would I need to add to my code to get these?

Thanks

Upvotes: 3

Views: 32325

Answers (2)

Sam
Sam

Reputation: 4090

Have you looked into seaborn?

import seaborn as sns
data = np.random.rand(2,2)
sns.heatmap(data, annot=True,  linewidths=.5)

enter image description here

Upvotes: 6

jojonas
jojonas

Reputation: 1675

For the color bar, see this example. You guessed correctly, the right function call is just ax.colorbar().

Concerning the values on the pcolor plot, I'm afraid you have to add them "manually" using ax.text(). Calculate and loop over the cell center values and call the text() function. You can use the keyword arguments horizontalalignment='center' and verticalalignment='center' to keep the text value centered around x, y.

Upvotes: 5

Related Questions