How to fix seaborn heatmap color mapping when values are in wide range

I have a dataframe with 6 unique values in range(0-9). I want to assign specif color to each value but mapping is not working for me.

This is how my dataframe looks like:

cmap_new = {0: '#faf5f5', 1: '#ff0303', 6: '#1f78b4', 7: '#b2df8a', 8: '#33a02c', 9: '#fb9a99'}
cmap = ListedColormap([cmap_new[i] for i in cmap_new.keys()])
ax = sns.heatmap(data=tmp_df, cmap=cmap, yticklabels=True, xticklabels=False,linewidths=1,square=True,annot=True)

My plot looks like this:

In my data, though I dont have values [2-5], they are assigned a color. I want to fix this problem and assign colors only to keys in the cmap_new dictionary.

Can anyone help me with this?

Solution

You can use a BoundaryNorm to assign a color to each of the values used. An extra value is needed, as 7 boundaries define 6 color regions. In order to get a nice colorbar, the ticks can be moved to the centers of each region.

import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap, BoundaryNorm
import pandas as pd

cmap_new = {0: '#faf5f5', 1: '#ff0303', 6: '#1f78b4', 7: '#b2df8a', 8: '#33a02c', 9: '#fb9a99'}
keys = sorted(cmap_new.keys())
keys += [max(keys) + 1]
tmp_df = pd.DataFrame(np.random.choice(keys, size=(8, 7)),
                      index=['Zoros', 'Lyra', 'Elara', 'Drakor', 'Astra', 'Mystalin', 'Terra', 'Vora'])
cmap = ListedColormap([cmap_new[k] for k in keys[:-1]])
norm = BoundaryNorm(keys, ncolors=len(keys) - 1)
ax = sns.heatmap(data=tmp_df, cmap=cmap, norm=norm,
                 yticklabels=True, xticklabels=False, linewidths=1, square=True, annot=True)
ax.collections[0].colorbar.set_ticks([(k1 + k2) / 2 for k1, k2 in zip(keys[:-1], keys[1:])], labels=keys[:-1])
plt.show()