python seaborn heatmap replace colorbar with absolute colorboxes

Question

I have a heatmap that contains only absolute values of -1, 0 and 1

import random 
import numpy as np
import matplotlib
import seaborn as sb
import matplotlib.pyplot as plt

array = []
for x in range(10):
    array.append(random.choices([-1,0, 1], k = 5))
array = np.array(array)


heatmap = sb.heatmap(array, cbar_kws={'ticks': [-1, 0 , 1]}, cmap = ["red", "grey", "green"])
plt.show()

is it possible to remove the colorbar and replace it by three different collored boxes and a custom label, like you would expect in the legend of a barplot? I.E. a red box and the word "No" next to it, a grey box and the word "N/A" next to it and a green box with the word "yes" next to it

sunnytown sunnytown · Accepted Answer · 2021-02-04T11:30:34

You can do it manually using matplotlib. To create the boxes, we do a scatter plot with squares for each color but without data, so they won't show up. We save the return value of those scatter plots as handles to pass to the legend. We then grab the matplotlib figure object from the heatmap variable, which contains the axis on which the plot is located. There, we create a legend with the custom handles and labels. Calling subplots_adjust on that figure, we make room for the legend on the right.

import random 
import numpy as np
import matplotlib
import seaborn as sb
import matplotlib.pyplot as plt

array = []
for x in range(10):
    array.append(random.choices([-1,0, 1], k = 5))
array = np.array(array)

colors = ["red", "grey", "green"]
heatmap = sb.heatmap(array, cmap = ["red", "grey", "green"], cbar=False)

#Create dummy handles using scatter
handles = [plt.scatter([], [], marker='s', s=50, color=color) for color in colors]
labels = [-1, 0 , 1]
#Creating the legend using dummy handles
heatmap.figure.legend(handles=handles, labels=labels, loc='center right', frameon=False)
#Adjusting the plot space to the right to make room for the legend
heatmap.figure.subplots_adjust(right=0.8)
plt.show()

On a sidenote:

You can replace your code for the generation of the random array with a numpy function, that does exactly what you want but is way more conventient.

So replace this:

array = []
for x in range(10):
    array.append(random.choices([-1,0, 1], k = 5))
array = np.array(array)

With this:

array = np.random.choice((-1, 0, 1), (10, 5))

where the first argument is the choices and the second argument is the shape of the array, so in your case 10 by 5.

python seaborn heatmap replace colorbar with absolute colorboxes

1 Answers