使用調色板進行散點圖繪制時顯示正確的圖例

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'])

這段代碼很完美，我們只用一行就可以繪制所有類。 但是，圖例這次沒有正確顯示。

題

使用以下內容繪制圖形時如何維護正確的圖例？

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 plt
from sklearn.datasets import make_blobs
import numpy as np

X, 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 plt
import seaborn as sns
from sklearn.datasets import make_blobs
import numpy as np

X, y = make_blobs()
ngroups = 3

ax = 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()