Home > OS >  plt.legend only adds first element to scatter plot
plt.legend only adds first element to scatter plot

Time:07-12

I am trying to add a legend to my scatter plot with 13 classes, however, with my code below, I am only able to get the first label. Can you assist me in generating the full list to show up in the legend of the scatter plot?

Here is my example code:

from sklearn.datasets import make_blobs
from matplotlib import pyplot
from pandas import DataFrame
# generate 2d classification dataset
X, y = make_blobs(n_samples=1000, centers=13, n_features=2)

classes = [f"class {i}" for i in range(13)]

#fig = plt.figure()
plt.figure(figsize=(15, 12))
scatter = plt.scatter(
    x=X[:,0],
    y=X[:,1],
    s = 20, 
    c = y, 
    cmap='Spectral'
    #c=[sns.color_palette()[x] for x in y_train_new]
    )
plt.gca().set_aspect('equal', 'datalim')
plt.legend(classes)
plt.title('Dataset', fontsize=24)

CodePudding user response:

You can do that by replacing the plt.legend(classes) in your code by this line... I hope this is what you are looking for. I am using matplotlib 3.3.4.

plt.legend(handles=scatter.legend_elements()[0], labels=classes)

Output plot

enter image description here

  • Related