图例按组着色的散点图,无需多次调用plt.scatter

dan*_*n_g 5 python scatter matplotlib legend

pyplot.scatter允许传递到c=与组相对应的数组,然后将根据这些组为点着色。但是,这似乎不支持不单独绘制每个组而生成图例。

因此,例如,可以通过遍历各组并分别绘制每个图来生成带有彩色组的散点图:

import matplotlib.pyplot as plt
from sklearn.datasets import load_iris
feats = load_iris()['data']
target = load_iris()['target']

f, ax = plt.subplots(1)
for i in np.unique(target):
    mask = target == i
    plt.scatter(feats[mask, 0], feats[mask, 1], label=i)
ax.legend()
Run Code Online (Sandbox Code Playgroud)

会产生:

在此处输入图片说明

我可以实现类似外观的情节而无需遍历每个组:

f, ax = plt.subplots(1)
ax.scatter(feats[:, 0], feats[:, 1], c=np.array(['C0', 'C1', 'C2'])[target])
Run Code Online (Sandbox Code Playgroud)

但是我无法找到第二种策略来生成相应图例的方法。我遇到的所有示例都在各个组之间进行迭代,这似乎不理想。我知道我可以手动生成图例,但这又显得太麻烦了。

Bob*_*ley 4

解决此问题的 matplotlib 散点示例也使用循环,因此这可能是预期用法:https://matplotlib.org/examples/lines_bars_and_markers/scatter_with_legend.html

如果您的更大目标只是使分类数据的绘制和标记更加简单,那么您应该考虑Seaborn。这是与 Pandas/Pyplot 中的散点图类似的问题:如何按类别绘制

实现目标的一种方法是使用带有标签列的 pandas。一旦 Pandas Dataframe 中有数据,您就可以使用Seabornpairplot来绘制这种图。(Seaborn 还提供 iris 数据集作为标记的 DataFrame)

import seaborn as sns
iris = sns.load_dataset("iris")
sns.pairplot(iris, hue="species")
Run Code Online (Sandbox Code Playgroud)

在此输入图像描述

如果您只想要前两个功能,您可以使用

sns.pairplot(x_vars=['sepal_length'], y_vars=['sepal_width'], data=iris, hue="species", size=5)
Run Code Online (Sandbox Code Playgroud)

在此输入图像描述

如果你真的想使用 sklearn 数据字典,你可以将其放入数据框中,如下所示:

import pandas as pd
from sklearn.datasets import load_iris
import numpy as np

feats = load_iris()['data'].astype('O')
target = load_iris()['target']
feat_names = load_iris()['feature_names']
target_names = load_iris()['target_names'].astype('O')

sk_df = pd.DataFrame(
    np.hstack([feats,target_names[target][:,np.newaxis]]),
    columns=feat_names+['target',])
sns.pairplot(sk_df, vars=feat_names, hue="target")
Run Code Online (Sandbox Code Playgroud)