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

时间:2014-01-03 17:09:37

标签: python matplotlib heatmap

我正在尝试使用指定的着色要求制作热图。我想设置数据的间隔并判断为好,并将其着色为绿色,其余结果应为红色。有没有人知道如何做到这一点? 我有一个简单的例子,使用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()

example heatmap

2 个答案:

答案 0 :(得分:5)

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

最简单的方法是将布尔数组传递给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()

enter image description here

或者,正如@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()

enter image description here

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

另一方面,使用pcolormesh而不是pcolor的速度要快得多。对于小型数组,它不会产生显着差异,但对于大型数组pcolor来说速度过慢。如果你不介意光栅输出,imshow甚至更快。使用imshow(data, interpolation='nearest', aspect='auto', origin='lower')匹配默认pcolorpcolormesh

答案 1 :(得分:0)

你可以制作2色colormap。 然后,您可以设置红色和绿色之间的截止值。