Python：如何使用示例创建图例

Question

这是来自Machine Learning In Action一书中的第2章，我试图在这里绘制情节：

作者已经发布了剧情的代码here，我认为这可能有点hacky（他还提到这段代码很草率，因为它超出了本书的范围）。

这是我尝试重新创建剧情：

首先，保存数据的.txt文件如下（来源：第2章here中的“datingTestSet2.txt”）：

40920   8.326976    0.953952    largeDoses
14488   7.153469    1.673904    smallDoses
26052   1.441871    0.805124    didntLike
75136   13.147394   0.428964    didntLike
38344   1.669788    0.134296    didntLike
...

假设datingDataMat是numpy.ndarray形状`（1000L，2L），其中第0列是“每年频繁飞行里数”，第1列是“播放视频游戏的时间百分比”，第2列是“每周消耗的冰淇淋”，​​如上面的样本所示。

假设datingLabels是整数1,2或3的list，意思是“不喜欢”，“喜欢小剂量”和“大剂量喜欢” - 与列相关联3以上。

以下是我创建图表的代码（file2matrix的完整详细信息在最后）：

datingDataMat,datingLabels = file2matrix("datingTestSet2.txt")
import matplotlib.pyplot as plt

fig = plt.figure()
ax = fig.add_subplot (111)
plt.xlabel("Freq flier miles")
plt.ylabel("% time video games")
# Not sure how to finish this: plt.legend([1, 2, 3], ["did not like", "small doses", "large doses"])
plt.scatter(datingDataMat[:,0], datingDataMat[:,1], 15.0*np.array(datingLabels), 15.0*np.array(datingLabels)) # Change marker color and size 
plt.show()

输出在这里：

我主要关心的是如何创建这个图例。有没有办法做到这一点，而无需直接处理积分？

接下来，我很好奇我是否能找到一种方法来切换颜色以匹配情节。有没有办法在没有对各个点进行某种“处理”的情况下做到这一点？

此外，如果有兴趣，这里是file2matrix实施：

def file2matrix(filename):
    fr = open(filename)
    numberOfLines = len(fr.readlines())
    returnMat = np.zeros((numberOfLines,3)) #numpy.zeros(shape, dtype=float, order='C') 
    classLabelVector = []
    fr = open(filename)
    index = 0
    for line in fr.readlines():
        line = line.strip()
        listFromLine = line.split('\t')
        returnMat[index,:] = listFromLine[0:3] # FFmiles/yr, % time gaming, L ice cream/wk
        classLabelVector.append(int(listFromLine[-1]))
        index += 1
    return returnMat,classLabelVector

Answer 1

要创建图例，您必须：

• 为每条曲线添加标签

• 从当前legend()对象调用AxesSubplot方法，例如，可以使用plt.gca()获取该方法。

请参阅以下示例：

plt.scatter(datingDataMat[:,0], datingDataMat[:,1],
            15.0*np.array(datingLabels), 15.0*np.array(datingLabels),
            label='Label for this data')
plt.gca().legend(loc='upper left')

Answer 2

这是一个模仿你已经拥有的代码的例子，它展示了Saullo Castro的例子中描述的方法。 它还显示了如何在示例中设置颜色。 如果您想了解有关可用颜色的更多信息，请参阅http://matplotlib.org/api/colors_api.html

上的文档

还值得查看http://matplotlib.org/1.3.1/api/pyplot_api.html#matplotlib.pyplot.scatter

处的散点图文档

from numpy.random import rand, randint
from matplotlib import pyplot as plt
n = 1000
# Generate random data
data = rand(n, 2)
# Make a random array to mimic datingLabels
labels = randint(1, 4, n)
# Separate the data according to the labels
data_1 = data[labels==1]
data_2 = data[labels==2]
data_3 = data[labels==3]
# Plot each set of points separately
# 's' is the size parameter.
# 'c' is the color parameter.
# I have chosen the colors so that they match the plot shown.
# With each set of points, input the desired label for the legend.
plt.scatter(data_1[:,0], data_1[:,1], s=15, c='r', label="label 1")
plt.scatter(data_2[:,0], data_2[:,1], s=30, c='g', label="label 2")
plt.scatter(data_3[:,0], data_3[:,1], s=45, c='b', label="label 3")
# Put labels on the axes
plt.ylabel("ylabel")
plt.xlabel("xlabel")
# Place the Legend in the plot.
plt.gca().legend(loc="upper left")
# Display it.
plt.show()

如果使用plt.savefig将图形保存到文件而不是显示，则灰色边框应变为白色。 请务必在保存到文件后运行plt.clf()或plt.cla()以清除轴，这样您就不会一遍又一遍地重新绘制相同的数据。