seaborn 热图获取颜色代码值数组

seaborn heatmap get array of color codes values

我正在尝试获取与热图的每个单元格关联的颜色代码:

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)

因为sns.heatmap return是一个matplotlib轴对象,我们不能直接使用hm。但是我们可以使用 cmap 对象本身来 return 数据的 rgba 值。 编辑代码已更新以包含数据规范化。

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))

所有颜色现在都包含在 rgba_values 中。所以要获得热图中左上角方块的颜色,你可以简单地做

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

有关更多信息,请查看 Getting individual colors from a color map in matplotlib


更新

要在对 sns.heatmap 的调用中使用 centerrobust 关键字重新调整颜色图,基本上只需重新定义 vminvmax.查看相关的 seaborn 源代码 (http://github.com/mwaskom/seaborn/blob/master/seaborn/matrix.py#L202),下面对 vminvmax 的更改应该可以解决问题。

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))

在不了解 heatmap 的输入数据和参数的情况下,您可以从底层 QuadMesh 获取颜色,知道热图应该是第一个也是唯一的 collection 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()))

另见 。与 AxesImage 相比,QuadMesh returns 颜色列表。因此,上面的代码将为您提供一个二维数组,其中列是 RGBA 颜色通道。如果您需要前两个维度与输入数据相同的 3D 输出,则需要重塑

rgba_values = rgba_values.reshape((im._meshHeight, im._meshWidth, 4))