2

I am trying to get the color codes associated with each cell of a heatmap:

import seaborn as sns
import numpy as np
import matplotlib.cm as cm

hm = sns.heatmap(
np.random.randn(10,10),
cmap = cm.coolwarm)

# hm.<some function>[0][0] would return the color code of the cell indexed (0,0)
lanery
  • 5,222
  • 3
  • 29
  • 43
spfraib
  • 109
  • 3
  • 8

2 Answers2

4

Because sns.heatmap returns a matplotlib axis object, we can't really use hm directly. But we can use the cmap object itself to return the rgba values of the data. Edit Code has been updated to include normalization of data.

from matplotlib.colors import Normalize

data = np.random.randn(10, 10)
cmap = cm.get_cmap('Greens')
hm = sns.heatmap(data, cmap=cmap)

# Normalize data
norm = Normalize(vmin=data.min(), vmax=data.max())
rgba_values = cmap(norm(data))

All of the colors are now contained in rgba_values. So to get the color of the upper left square in the heatmap you could simply do

In [13]: rgba_values[0,0]
Out[13]: array([ 0.        ,  0.26666668,  0.10588235,  1.        ])

For more, check out Getting individual colors from a color map in matplotlib

Update

To readjust the colormap from using the center and robust keywords in the call to sns.heatmap, you basically just have to redefine vmin and vmax. Looking at the relevant seaborn source code (http://github.com/mwaskom/seaborn/blob/master/seaborn/matrix.py#L202), the below changes to vmin and vmax should do the trick.

data = np.random.randn(10, 10)
center = 2
robust = False
cmap = cm.coolwarm
hm = sns.heatmap(data, cmap=cmap, center=center, robust=robust)

vmin = np.percentile(data, 2) if robust else data.min()
vmax = np.percentile(data, 98) if robust else data.max()
vmin += center
vmax += center

norm = Normalize(vmin=vmin, vmax=vmax)
rgba_values = cmap(norm(data))
Community
  • 1
  • 1
lanery
  • 5,222
  • 3
  • 29
  • 43
  • Thank you lanery. What if I were to use the seaborn options "center" and "robust" to adjust the colormap, is there a way to then use the modified colormap? say for instance: heatmap = sns.heatmap( np.random.randn(10, 10), cmap = cm.coolwarm, center=1, robust=True) – spfraib Feb 09 '17 at 02:12
  • You need to add a norm object — as called you are only mapping values in [0–1]. – mwaskom Feb 09 '17 at 14:33
  • Also cm.get_cmap('Greens') just returns the cm.Greens object. – mwaskom Feb 09 '17 at 15:34
  • Would you mind giving an example of how to do that ? – spfraib Feb 09 '17 at 18:25
  • @spfraib Check out my updated answer regarding the normalization and adjusting for the use of center and robust. – lanery Feb 09 '17 at 20:00
  • @mwaskom 1st comment: Thanks, updated answer to show that. 2nd comment: True. I didn't mean to make it sound as though there was something special about `cm.get_cmap`, but I guess my wording may have implied that, so I changed the wording. – lanery Feb 09 '17 at 20:21
2

Without any knowledge on the input data and arguments of heatmap you can get the colors from the underlying QuadMesh, knowing that the heatmap should be the first and only collection inside the axes that is returned by heatmap.

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns

data = np.array([[0,-2],[10,5]])
ax = sns.heatmap(data, center=0, cmap="bwr", robust=False)
im = ax.collections[0]
rgba_values = im.cmap(im.norm(im.get_array()))

Also see this answer. In contrast to AxesImage, QuadMesh returns a list of colors. Hence the above code will give you a 2D array where the columns are the RGBA color channels. If you need a 3D output with the first two dimensions being the same as the input data you would need to reshape

rgba_values = rgba_values.reshape((im._meshHeight, im._meshWidth, 4))
ImportanceOfBeingErnest
  • 321,279
  • 53
  • 665
  • 712