Brian Kim
Brian Kim

Reputation: 33

Matplotlib scatter plot legend display problem

I'm struggling to make a scatterplot code which includes legend related to 'pitch' like below. The variable 'u' stands for the unique pitch list. And the number is 19. So the number of legend's label should be 19. But it turns out to be like below. There are only 10 labels on legend. And the color range assign looks weird. (label '15' should be the dark-blue but it's not though) What seems to be the problem?

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




df = pd.DataFrame({
"X" : [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37], 
"Y" : [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37], 
"pitch":[10,10,20,20,30,40,50,50,60,60,60,70,70,70,70,80,80,80,100,150,1,2,3,4,5,6,7,8,9,10,3,4,5,8,9,10,3],
})


color = cm.jet
u, div = np.unique(df.pitch.values, return_inverse=True)
colorlist = [colors.rgb2hex(color(i)) for i in np.linspace(1, 0, len(u))]
cmap = ListedColormap(colorlist)
fig,ax = plt.subplots()
scatter = plt.scatter(df['X'],df['Y'], c=div, cmap=cmap)
plt.legend(scatter.legend_elements()[0], u, loc=2)
plt.show()

enter image description here

Upvotes: 1

Views: 1073

Answers (2)

otavios
otavios

Reputation: 124

fig,ax = plt.subplots(figsize=(12,8))

for i, pitch in enumerate(u):
    df_p = df[df['pitch'] == pitch]
    scatter = ax.scatter(df_p['X'],df_p['Y'], c=colorlist[i], cmap=cmap, 
    label=pitch)

ax.legend(loc=2)
plt.show()

enter image description here

Upvotes: 2

Redox
Redox

Reputation: 10017

You need to replace the plt.legend(scatteer..) line by this.... Documentation on legend_element is available here.

plt.legend(scatter.legend_elements(prop='colors', num=len(colorlist))[0], u, loc=2)

Output Plot

enter image description here

Upvotes: 1

Related Questions