Amilovsky
Amilovsky

Reputation: 499

Display legend of seaborn clustermap corresponding to the row colors

''' For sake of simplicity, let's use the iris dataset. I would like to add a legend matching each species to its color code (blue,green, red in this example). By the way, I found similar problems at the following links but were a bit more complicated. How to express classes on the axis of a heatmap in Seaborn

The solution proposed at Seaborn clustermap row color with legend would have worked but for df[['tissue type','label']] when definig legend_TN, I'm not sure how to define the label similarly, like iris['species','xxxx'] Thank you in advance for helping me out. '''

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
%matplotlib inline

iris = sns.load_dataset('iris')
species = iris.pop('species')


lut = dict(zip(species.unique(), "rbg"))
row_colors = species.map(lut)
g = sns.clustermap(iris, row_colors=row_colors)
plt.show()

Upvotes: 6

Views: 7880

Answers (1)

JohanC
JohanC

Reputation: 80459

Following the examples at the docs, one could create a custom legend as:

from matplotlib.patches import Patch

handles = [Patch(facecolor=lut[name]) for name in lut]
plt.legend(handles, lut, title='Species',
           bbox_to_anchor=(1, 1), bbox_transform=plt.gcf().transFigure, loc='upper right')

example plot

Upvotes: 10

Related Questions