CODEWITHSUNDEEP
pythonpandasmatplotlibscatter-plot
Alhpa Delta
Alhpa Delta

Reputation: 3660

Matplotlib Legend on Scatterplot from pandas dataframe

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. enter image description here

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

Upvotes: 3

Views: 10309

Answers (2)

ImportanceOfBeingErnest
ImportanceOfBeingErnest

Reputation: 339052

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

enter image description here

Upvotes: 6

Quang Hoang
Quang Hoang

Reputation: 150725

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:

enter code here

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:

enter image description here

Upvotes: 8

Related Questions