散点图,其中图例按组着色,没有多次调用plt.scatter

时间:2017-05-14 18:52:14

标签: python matplotlib legend scatter

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

生成:

enter image description here

我可以实现类似的情节而不会遍历每个组:

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

但是我无法想出用这个第二个策略生成相应图例的方法。我遇到的所有例子都在群体中进行迭代,这看起来......不太理想。我知道我可以手动生成一个图例,但再次看起来过于繁琐。

1 个答案:

答案 0 :(得分:2)

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

如果您的更大目标是让绘图和标记分类数据更直接,则应考虑Seaborn。这是与Scatter plots in Pandas/Pyplot: How to plot by category

类似的问题

实现目标的一种方法是使用带有标记列的pandas。在Pandas Dataframe中获得数据后,您可以使用Seaborn pairplot制作此类图表。 (Seaborn还有虹膜数据集作为标记的DataFrame提供)

import seaborn as sns
iris = sns.load_dataset("iris")
sns.pairplot(iris, hue="species")

enter image description here

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

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

enter image description here

如果您真的想使用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")