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

Question

这来自 Machine Learning In Action 一书的第 2 章，我正在尝试制作此处所示的情节:

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

这是我尝试重新创建情节的尝试:

首先，保存数据的 .txt 文件如下(来源:“datingTestSet2.txt” in Ch.2 here ):

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 是形状为 `(1000L, 2L) 的 numpy.ndarray，其中第 0 列是“每年的飞行里程数”，第 1 列是“% Time玩电子游戏”，第 2 列是“每周消耗的冰淇淋升数”，如上例所示。

假设 datingLabels 是整数 1、2 或 3 的 列表，表示“不喜欢”、“小剂量喜欢”和“大剂量喜欢” "分别 - 与上面的第 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

这是一个模仿您已有代码的示例，它展示了 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() 以清除轴，这样您就不会最终在上面重新绘制相同的数据一遍又一遍。