python - python中的两种颜色热图

Question

我有以下数据:

my_array = array([[0, 0, 1, 0, 0],
       [0, 1, 1, 1, 0],
       [0, 0, 0, 1, 1],
       [0, 0, 1, 1, 1],
       [0, 1, 1, 0, 0],
       [1, 1, 1, 1, 0],
       [0, 1, 1, 1, 1],
       [0, 0, 0, 0, 1],
       [0, 1, 0, 1, 0]])

和

df.values = array([246360,  76663,  29045,  11712,   5526,   3930,   3754,   1677,
         1328])

我正在制作这样的热图:

import seaborn as sns
import matplotlib.pyplot as plt
cmap = sns.cm.rocket_r
ax = sns.heatmap(my_array, xticklabels=["A", "B", "C", "D", "E"], yticklabels=df.values, cmap = cmap)
ax.set(xlabel='Test Type', ylabel='Number', title='patterns of missingness')
fig=plt.figure(figsize=(40,30), dpi= 20, facecolor='w', edgecolor='k')
fig

我得到以下信息:

我的问题是，如何摆脱连续的色标并只选择两种不同的颜色:0 为白色，1 为绿色？

Answer 1

您可以使用 LinearSegmentedColormap 而不是您正在使用的当前 cmap 来执行此操作。

import matplotlib as mpl
cmap = mpl.colors.LinearSegmentedColormap.from_list('my_cmap', ['white', 'green'], 2)

然后在调用 sns.heatmap 时将 ticks 参数传递给 cbar_kws 参数:

ax = sns.heatmap(my_array, xticklabels=["A", "B", "C", "D", "E"],
                 yticklabels=df.values, cmap=cmap, cbar_kws={"ticks":[0,1]})

给出:

如果很清楚绿色和白色代表什么，您也可以使用 cbar=False 完全关闭 cbar。

ax = sns.heatmap(my_array, xticklabels=["A", "B", "C", "D", "E"],
             yticklabels=df.values, cmap=cmap, cbar=False)