python - 根据要求在 Python/Matplotlib 中为热图着色

Question

我正在尝试制作具有特定着色要求的热图。我想给数据设置一个区间，判断ok就给它涂上绿色，其余的结果应该涂成红色。有谁知道如何做到这一点？为了更好地理解，我附上了一个使用 pandas 和 matplotlib 的简单示例。

import numpy as np 
from pandas import *
import matplotlib.pyplot as plt

Index= ['aaa', 'bbb', 'ccc', 'ddd', 'eee']
Cols = ['A', 'B', 'C', 'D']
data= abs(np.random.randn(5, 4))
df = DataFrame(data, index=Index, columns=Cols)

plt.pcolor(df)
plt.yticks(np.arange(0.5, len(df.index), 1), df.index)
plt.xticks(np.arange(0.5, len(df.columns), 1), df.columns)
plt.show()

Answer 1

有不止一种方法可以做到这一点。

最简单的方法是将 bool 数组传递给 pcolor然后选择一个颜色图，其中绿色高，红色低。

例如:

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

Index= ['aaa', 'bbb', 'ccc', 'ddd', 'eee']
Cols = ['A', 'B', 'C', 'D']
data= np.random.random((5, 4))
df = pd.DataFrame(data, index=Index, columns=Cols)

plt.pcolor(df > 0.5, cmap='RdYlGn')
plt.yticks(np.arange(0.5, len(df.index), 1), df.index)
plt.xticks(np.arange(0.5, len(df.columns), 1), df.columns)
plt.show()

或者，正如@Cyber​​ 提到的，您可以根据您的值制作一个双色颜色图并使用它:

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors

Index= ['aaa', 'bbb', 'ccc', 'ddd', 'eee']
Cols = ['A', 'B', 'C', 'D']
data= np.random.random((5, 4))
df = pd.DataFrame(data, index=Index, columns=Cols)

# Values from 0-0.5 will be red and 0.5-1 will be green
cmap, norm = mcolors.from_levels_and_colors([0, 0.5, 1], ['red', 'green'])

plt.pcolor(df, cmap=cmap, norm=norm)
plt.yticks(np.arange(0.5, len(df.index), 1), df.index)
plt.xticks(np.arange(0.5, len(df.columns), 1), df.columns)
plt.show()

(颜色差异只是因为“RdYlGn”颜色图使用较深的绿色和红色作为其端点。)

附带说明一下，使用 pcolormesh 也快得多为此，而不是 pcolor .对于小型阵列，它不会产生显着差异，但对于大型阵列 pcolor太慢了。 imshow甚至更快，如果你不介意光栅输出。使用 imshow(data, interpolation='nearest', aspect='auto', origin='lower')匹配 pcolor 的默认值和 pcolormesh .