使用调色板绘制散点图时显示正确的图例

Question

绘制散点图的愚蠢方法

假设我有一个包含 3 个类的数据，以下代码可以为我提供一个具有正确图例的完美图形，其中我逐类绘制数据。

import matplotlib.pyplot as plt

import seaborn as sns

from sklearn.datasets import make_blobs

import numpy as np

X, y = make_blobs()

X0 = X[y==0]

X1 = X[y==1]

X2 = X[y==2]

ax = plt.subplot(1,1,1)

ax.scatter(X0[:,0],X0[:,1], lw=0, s=40)

ax.scatter(X1[:,0],X1[:,1], lw=0, s=40)

ax.scatter(X2[:,0],X2[:,1], lw=0, s=40)

ax.legend(['0','1','2'])

绘制散点图的更好方法

但是，如果我有一个包含 3000 个类的数据集，则上述方法不再有效。（你不会指望我写3000行对应每个类吧？）于是我想出了下面的绘图代码。

num_classes = len(set(y))

palette = np.array(sns.color_palette("hls", num_classes))

ax = plt.subplot(1,1,1)

ax.scatter(X[:,0], X[:,1], lw=0, s=40, c=palette[y.astype(np.int)])

ax.legend(['0','1','2'])

这段代码很完美，我们可以用 1 行绘制出所有的类。但是，图例这次没有正确显示。

题

当我们使用以下方法绘制图形时，如何保持正确的图例？

ax.scatter(X[:,0], X[:,1], lw=0, s=40, c=palette[y.astype(np.int)])

Answer 1

plt.legend()当情节中有多个“艺术家”时效果最佳。在您的第一个示例中就是这种情况，这就是调用plt.legend(labels)可以毫不费力地工作的原因。如果您担心编写大量代码行，那么您可以利用for循环。正如我们在这个例子中看到的，使用了 5 个类：import matplotlib.pyplot as pltfrom sklearn.datasets import make_blobsimport numpy as npX, y = make_blobs(centers=5)ax = plt.subplot(1,1,1)for c in np.unique(y):    ax.scatter(X[y==c,0],X[y==c,1],label=c)ax.legend()np.unique()返回 y 的唯一元素的排序数组，通过循环遍历这些元素并使用自己的艺术家绘制每个类plt.legend()可以轻松提供图例。编辑：您还可以在制作时为图分配标签，这可能更安全。plt.scatter(..., label=c) 其次是 plt.legend()

Answer 2

为什么不简单地执行以下操作？import matplotlib.pyplot as pltimport seaborn as snsfrom sklearn.datasets import make_blobsimport numpy as npX, y = make_blobs()ngroups = 3ax = plt.subplot(1, 1, 1)for i in range(ngroups):    ax.scatter(X[y==i][:,0], X[y==i][:,1], lw=0, s=40, label=i)ax.legend()