Matplotlib: How to add legend to scatter plot colors?

Question

My data frame has three columns of SKU, Saving and label (categorical var). When I call plt.legend(), it adds legend of "Saving", but I want to add legend to my colors (a,b,c,d)?

from numpy import *
import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
import matplotlib.pyplot as plt

df = pd.DataFrame(np.random.rand(100,1), columns=['Saving'])
df['SKU'] = np.arange(100)
df['label'] = np.random.choice(['a', 'b', 'c','d'], 100)

fig, ax = plt.subplots()
colors = {'a':'red', 'b':'blue', 'c':'green', 'd':'white'}
figSaving = ax.scatter(df['SKU'], df['Saving'], c=df['label'].apply(lambda x: colors[x]))

plt.show()

score 1 · Answer 1 · answered Sep 15 '17 at 19:15

1

plt.legend is a callable. By writing plt.legend={'a', 'b', 'c', 'd'}, you are replacing that callable by a set, which in itself does nothing (except making it impossible to call legend afterwards. What you want to do is to call plt.legend(). See https://matplotlib.org/users/legend_guide.html.

answered Sep 15 '17 at 19:15

fuglede

17,388
2
54
99

Thank you, but when I call plt.legend, it shows Saving, not a,b,c,d groups. – P.J Sep 15 '17 at 19:17
Knowing what you want to plot is unfortunately almost impossible from the question itself, but given your comment, it sounds like you are looking for something like https://stackoverflow.com/questions/21654635/scatter-plots-in-pandas-pyplot-how-to-plot-by-category#21655256. – fuglede Sep 15 '17 at 19:25

score 1 · Accepted Answer · answered Sep 15 '17 at 20:02

import pandas as pd
import numpy as np
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt

df = pd.DataFrame(np.random.rand(100,1), columns=['Saving'])
df['SKU'] = np.arange(100)
df['label'] = np.random.choice(['a', 'b', 'c','d'], 100)

fig, ax = plt.subplots()
colors = {'a':'red', 'b':'blue', 'c':'green', 'd':'white'}
figSaving = ax.scatter(df['SKU'], df['Saving'], c=df['label'].apply(lambda x: colors[x]))


# build the legend
red_patch = mpatches.Patch(color='red', label='a')
blue_patch = mpatches.Patch(color='blue', label='b')
green_patch = mpatches.Patch(color='green', label='c')
white_patch = mpatches.Patch(color='white', label='d')

# set up for handles declaration
patches = [red_patch, blue_patch, green_patch, white_patch]

# define and place the legend
#legend = ax.legend(handles=patches,loc='upper right')

# alternative declaration for placing legend outside of plot
legend = ax.legend(handles=patches,bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)

plt.show()

Matplotlib: How to add legend to scatter plot colors?

2 Answers2