Matplotlib scatter plot legend display problem-CodePudding

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

CodePudding user response：

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

CodePudding user response：

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