[Python] Adding row colors to a heatmap

I'm trying to visualize different groups within my data, on the axis of a heatmap (using sns.heatmap). I want the ticks to be categorized according to a specific dictionary, with names and colors, and eventually to be presented in a legend.

I know this can be obtained by using sns.clustermap but this function also clusters the values in the heatmap, which I don't want to happen.

Any idea how can I make such visualization in a heatmap?

Example

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import random
import seaborn as sns


data = pd.DataFrame(np.random.randn(10, 3), columns=(list('ABC')))
data['group'] = pd.Series(random.choices(['group_1', 'group_2', 'group_3'], weights=[1,1,1], k=len(data)))

data
A B C group
0 0.366822 0.583965 1.629740 group_3
1 0.557286 0.450663 0.255852 group_3
2 -0.265515 -0.153028 0.670448 group_3
3 0.132278 -0.226668 1.365583 group_3
4 0.595304 -0.577290 0.395477 group_2
5 -0.805420 0.168376 0.748649 group_1
6 0.105664 -0.568047 -0.281488 group_2
7 -0.046202 0.173409 -0.250321 group_1
8 -0.132696 -0.877354 0.086954 group_3
9 -0.843666 0.655146 -1.629453 group_2

lut = {'group_1': 'red', 'group_2': 'blue', 'group_3': 'green'}

row_colors = data['group'].map(lut)

data.drop(['group'], axis=1, inplace=True)

fig = sns.heatmap(data)



I aim to use row_colors to visually indicate which indices correspond to which group, in the heatmap to be.

any help regarding this would be highly appreciated, either directly from seaborn, or somehow else

Continue reading...