如何在散点图中显示图例以区分类别

问题描述

我正在研究来自 sklearn 的 iris 数据集。您可能知道 iris 数据集有 3 个类 ['setosa'、'versicolor'、'virginica']。我为这个数据集做了一个散点图。详情如下

from sklearn.datasets import load_iris
iris=load_iris()
Y_train=iris.target
X_train=iris.data
class_labels=iris.target_names
plt.scatter(X_train[:,0],X_train[:,1],c=Y_train)
plt.xlabel('attr1')
plt.ylabel('attr2')
plt.show()

我有散点图，您可以在其中看到黄色、绿色和紫色的点。我想知道哪个颜色点属于哪个类（'setosa'、'versicolor'、'virginica'）。我想显示图例，以便我知道哪种颜色代表哪个类

解决方法

在这种情况下，您可以通过循环标签并使用与散点图相同的 custom legend 和 colormap 来创建 norm。默认情况下，使用 'viridis' 颜色图，以及将最小颜色值映射为 0，将最大值映射为 1 的范数。

import matplotlib.pyplot as plt
from sklearn.datasets import load_iris

iris = load_iris()
Y_train = iris.target
X_train = iris.data
class_labels = iris.target_names
cmap = plt.get_cmap('viridis')
norm = plt.Normalize(Y_train.min(),Y_train.max())
plt.scatter(X_train[:,0],X_train[:,1],c=Y_train,cmap='viridis',norm=norm)
handles = [plt.Line2D([0,[0,color=cmap(norm(i)),marker='o',linestyle='',label=label)
           for i,label in enumerate(class_labels)]
plt.legend(handles=handles,title='Species')
plt.show()

您也可以使用 seaborn，但目前设置图例标签并不简单。

import seaborn as sns

sns.set()
ax = sns.scatterplot(x=X_train[:,y=X_train[:,hue=Y_train,palette='viridis')
ax.legend(ax.legend_.legendHandles,class_labels,title='Species')