带有图例的matplotlib散点图

时间:2017-02-05 20:08:13

标签: python matplotlib plot

我有兴趣在我的散点图中绘制一个图例。我目前的代码看起来像这样

x=[1,2,3,4]
y=[5,6,7,8]
classes = [2,4,4,2]
plt.scatter(x, y, c=classes, label=classes)
plt.legend()

问题在于创建绘图时,图例显示为数组,而不是显示唯一标签及其类。

This is how the plot looks

我知道这是先前在诸如此one之类的线程中讨论的问题,但是我觉得我的问题更简单,并且那里的解决方案不适合它。此外,在该示例中,该人指定颜色,但在我的情况下,我事先知道我需要多少颜色。此外,在this示例中,用户正在创建多个散点图,每个散点图都具有唯一的颜色。再说一遍,这不是我想要的。我的目标是使用x,y数组和标签简单地创建绘图。这可能吗?

感谢。

2 个答案:

答案 0 :(得分:8)

实际上,这两个相关问题都提供了一种如何达到预期效果的方法。

最简单的方法是创建与唯一类一样多的散点图,并为每个散列图提供单个颜色和图例条目。

import matplotlib.pyplot as plt

x=[1,2,3,4]
y=[5,6,7,8]
classes = [2,4,4,2]
unique = list(set(classes))
colors = [plt.cm.jet(float(i)/max(unique)) for i in unique]
for i, u in enumerate(unique):
    xi = [x[j] for j  in range(len(x)) if classes[j] == u]
    yi = [y[j] for j  in range(len(x)) if classes[j] == u]
    plt.scatter(xi, yi, c=colors[i], label=str(u))
plt.legend()

plt.show()

enter image description here

如果类是字符串标签,解决方案看起来会略有不同,因为您需要从索引中获取颜色而不是使用类本身。

import numpy as np
import matplotlib.pyplot as plt

x=[1,2,3,4]
y=[5,6,7,8]
classes = ['X','Y','Z','X']
unique = np.unique(classes)
colors = [plt.cm.jet(i/float(len(unique)-1)) for i in range(len(unique))]
for i, u in enumerate(unique):
    xi = [x[j] for j  in range(len(x)) if classes[j] == u]
    yi = [y[j] for j  in range(len(x)) if classes[j] == u]
    plt.scatter(xi, yi, c=colors[i], label=str(u))
plt.legend()

plt.show()

enter image description here

答案 1 :(得分:5)

也许手动填写table在这里很有用。如果您的类是连续的数字,则另一个想法是使用colorbar。我在一个方面展示了这两种方法。

import matplotlib.pyplot as plt
import numpy as np

x=[1,2,3,4,5,6,7]
y=[1,2,3,4,5,6,7]
classes = [2,4,4,2,1,3,5]
cmap = plt.cm.get_cmap("viridis",5)
plt.scatter(x, y, c=classes, label=classes,cmap=cmap,vmin=0.5,vmax=5.5)
plt.colorbar()
unique_classes = list(set(classes))
plt.table(cellText=[[x] for x in unique_classes], loc='lower right',
          colWidths=[0.2],rowColours=cmap(np.array(unique_classes)-1),
         rowLabels=['label%d'%x for x in unique_classes],
          colLabels=['classes'])

enter image description here