Matplotlib使用分类变量散布带有颜色的图例

时间:2018-05-08 17:36:01

标签: python pandas matplotlib

我使用matplotlib制作了一个简单的散点图,显示了来自2个数值变量(varA和varB)的数据,颜色是我用第3个分类字符串变量(col)定义的,包含10个唯一颜色(对应于另一个具有10个唯一名称的字符串变量) ),都在同一个100多行的Pandas DataFrame中。 有没有一种简单的方法可以为此散点图创建一个图例,该图例显示了唯一的彩色圆点及其相应的类别名称?或者我应该以某种方式对数据进行分组并在子图中绘制每个类别来执行此操作?这就是我到目前为止所做的:

import matplotlib.pyplot as plt
from matplotlib import colors as mcolors

varA = df['A']
varB = df['B'] 
col = df['Color']

plt.scatter(varA,varB, c=col, alpha=0.8)
plt.legend()

plt.show()

2 个答案:

答案 0 :(得分:0)

一种简单的方法是按颜色对数据进行分组,然后在一个图上绘制所有数据。 Pandas有一个内置的groupby功能。例如:

import matplotlib.pyplot as plt
from matplotlib import colors as mcolors

for color, group in df.groupby(['Color']):
    plt.scatter(group['A'], group['B'], c=color, alpha=0.8, label=color)

plt.legend()
plt.show()

请注意,我们会针对每组数据调用plt.scatter一次。然后,只要所有数据都在我们的情节中,我们只需要调用plt.legendplt.show

答案 1 :(得分:0)

考虑到,Color是包含所有颜色和标签的列,您只需执行以下操作即可。

colors = list(df['Color'].unique())
for i in range(0 , len(colors)):
    data = df.loc[df['Color'] == colors[i]]
    plt.scatter('A', 'B', data=data, color='Color', label=colors[i])
plt.legend()
plt.show()