Matplotlib Legend on Scatterplot from pandas dataframe

Question

I am trying to get the legend right on the figure below. It should be just 'green', 'blue' and 'red' with the corresponding color. But it is all over the place.

the code is below:

import pandas as pd
import matplotlib.pyplot as plt

df = pd.DataFrame({
        'category':['blue','green','red','blue','green','red','blue','green','red'],
        'attempts':[8955,7881,6723,100,200,300,4567,876,54],
        'success':[3000,7500,2000, 256,4567,4567,7665,543,43]
})
fig,ax = plt.subplots()
plt.scatter(df['attempts'],df['success'],c=df['category'],label=df['category'])
plt.legend(loc=2)
plt.savefig('scatter.png')
plt.show()

How do I get this right? (There is a similar one here: https://pythonspot.com/matplotlib-scatterplot/ in the second part "Scatter plot with groups", but this is not based on pandas dataframe).

Quang Hoang · Accepted Answer · 2019-09-18 19:31:55Z

8

You can use seaborn's scatterplot:

fig,ax = plt.subplots()
sns.scatterplot(data=df, hue='category', x='attempts', y='success')
plt.legend(loc=2)
plt.savefig('scatter.png')
plt.show()

Output:

Or pure matplotlib:

fig,ax = plt.subplots()
for k,d in df.groupby('category'):
    ax.scatter(d['attempts'], d['success'], label=k)
plt.legend(loc=2)
plt.savefig('scatter.png')
plt.show()

output:

edited Sep 18, 2019 at 19:31

answered Sep 18, 2019 at 19:27

Quang Hoang

151k11 gold badges64 silver badges86 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Alhpa Delta Over a year ago

I never knew about being able to use groupby like this. Is there a weblink or some document you can refer to where I can read about this? Not just in the context of plotting but in any other context in which it might be useful

Quang Hoang Over a year ago

This is from pandas doc, the groupby() returns a GroupBy object that has an __iter__() method, which allows for ... in.

ImportanceOfBeingErnest · Accepted Answer · 2019-09-18 20:30:15Z

6

If you want to use a single scatter with matplotlib, it would look like this:

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

df = pd.DataFrame({
        'category':['blue','green','red','blue','green','red','blue','green','red'],
        'attempts':[8955,7881,6723,100,200,300,4567,876,54],
        'success':[3000,7500,2000, 256,4567,4567,7665,543,43]
})

u, inv = np.unique(df.category.values, return_inverse=True)
cmap = ListedColormap(u)

fig,ax = plt.subplots()
scatter = plt.scatter(df['attempts'],df['success'],c=inv, cmap=cmap)
plt.legend(scatter.legend_elements()[0], u, loc=2)
plt.savefig('scatter.png')
plt.show()

answered Sep 18, 2019 at 20:30

ImportanceOfBeingErnest

342k61 gold badges737 silver badges771 bronze badges

Collectives™ on Stack Overflow

Matplotlib Legend on Scatterplot from pandas dataframe

2 Answers 2

2 Comments

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

2 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related