Matplotlib not displaying legend and how to do it within a for loop

June 14, 2023

I am trying to scatter some datapoints in matplotlib but am having some issues with displaying the legend properly.
After importing and initializing

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.cm as cm

rng = np.random.default_rng(0)

x = rng.uniform(0, 3, 5)
y = rng.uniform(0, 5, 5)
labels = ["a", "b", "c", "d", "e"]
colors = [plt.cm.plasma(i / float(len(labels))) for i in range(len(labels))]

I run a (less than I thought) simpleplt.scatter

# scatter strategy 1
scatter1 = plt.figure()
scatter1 = plt.scatter(
    x=x,
    y=y,
    c=colors,
    cmap=cm.plasma,
    label=labels,
)
handles, labels = scatter1.legend_elements(num=list(np.unique(labels)))
plt.legend(handles, labels, loc="lower right")
plt.show()

just to discover that handles and labels are empty lists. This results in the legend not being displayed.
How can I fix this, and why?

On a sidenote, I also tried plotting each datapoint one by one with a for loop, as follows:

# scatter strategy 2
scatter2 = plt.figure()
unique = np.unique(labels)

for i, u, color in zip(range(len(unique)), unique, colors):
    xi = x[i]
    yi = y[i]
    plt.scatter(xi, yi, color=color, cmap=cm.plasma, label=str(unique))
plt.show()

to no good, as it seems like the figure is not being "filled" with the data. Why is it so, and how could I fix it? How to properly compile a legend in this latter example?

Thank you!

>Solution :

As you are looking for the answer with a for loop, the second (strategy 2) is the one I modified. See if this is what you are looking for…

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.cm as cm

rng = np.random.default_rng(0)

x = rng.uniform(0, 3, 5)
y = rng.uniform(0, 5, 5)
labels = ["a", "b", "c", "d", "e"]
colors = [plt.cm.plasma(i / float(len(labels))) for i in range(len(labels))]

# scatter strategy 2
scatter2 = plt.figure()
unique = np.unique(labels)

for i in range(len(unique)): ##Simplified a bit
    plt.scatter(x[i], y[i], color=colors[i], cmap=cm.plasma, label=labels[i])  ## Updates here... should be colors[i] and labels[i]

plt.legend()  ## You will need this
plt.show()