seaborn熱圖獲取顏色代碼值的數組

Question

我正在嘗試獲取與熱圖的每個單元格關聯的顏色代碼：

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)

Answer 1

由於sns.heatmap返回一個matplotlib軸對象，因此我們不能真正直接使用hm 。 但是我們可以使用cmap對象本身來返回數據的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.        ])

有關更多信息，請查看從matplotlib的顏色圖中獲取單個顏色

---

更新

要從對sns.heatmap的調用中使用center和robust關鍵字重新調整顏色sns.heatmap ，基本上只需要重新定義vmin和vmax 。 查看相關的seaborn源代碼（ http://github.com/mwaskom/seaborn/blob/master/seaborn/matrix.py#L202 ），對vmin和vmax進行以下更改應該可以解決問題。

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

Answer 2

無需了解heatmap的輸入數據和參數，就可以從底層QuadMesh獲取顏色，因為知道熱圖應該是heatmap返回的軸內的第一個也是唯一的collection 。

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返回顏色列表。 因此，上面的代碼將為您提供一個2D數組，其中的列是RGBA顏色通道。 如果您需要3D輸出且前兩個維度與輸入數據相同，則需要重塑形狀

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