How do I put legend on scatter plot from a pandas data frame based on categories in the data?

Question

I have this code

import pandas as pd
import matplotlib.pyplot as plt

dataset = pd.read_csv('iris.csv', header=None, names=['sepal length', 'sepal width', 'petal length', 'petal width', 'class'])

dataset.drop(index=dataset.index[dataset['class'] == 'Iris-setosa'], inplace=True)

dataset.drop('petal width', axis='columns', inplace=True)
dataset.drop('sepal width', axis='columns', inplace=True)


fig, ax = plt.subplots()

groups = dataset.groupby('class')

colors = {'Iris-versicolor': 'blue', 'Iris-virginica': 'red'}

ax.scatter(dataset['petal length'], dataset['sepal length'], c=dataset['class'].map(colors))

plt.legend(['aaa', 'bbb'])

plt.show()

The data set is from this link: https://archive.ics.uci.edu/ml/datasets/Iris

I get a colored scatterplot, color based on 'class'. But in the legend it only shows one color: 'aaa'.

How do I make the put a legend on scatter with pyplotlib showing multiple colors?

`ns.scatterplot(data=dataset, x='petal length', y='sepal length', hue='class')` — Trenton McKinney, Aug 20 '21 at 14:03

score 1 · Answer 1 · answered Aug 20 '21 at 13:50

Turning the plotting code into a for loop over the groups works for me:

import pandas as pd
import matplotlib.pyplot as plt

# Removed data wrangling code

fig, ax = plt.subplots()

for name, group in dataset.groupby('class'):
    ax.scatter(
        group['petal length'], group['sepal length'], label=name)

plt.legend()

Output

How do I put legend on scatter plot from a pandas data frame based on categories in the data?

1 Answers1