How to add legend within a for loop

Question

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?

You want each point with the possibility of a unique color / label? So you got x, y, colors and labels properly defined? — J. M. Arnold, Jun 14 '23 at 10:34
Yes! Each point has its own unique label in the actual dataset. Labels are strings formatted in LaTeX. — Antonio Carnevali, Jun 14 '23 at 10:35

score 2 · Accepted Answer · answered Jun 14 '23 at 10:41

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

Stef · Answer 2 · 2023-06-14T11:13:16.517

1

You don't need to calculate the colors yourself, let matplotlib do it and just use your labels as is:

import matplotlib.pyplot as plt
import numpy as np

rng = np.random.default_rng(0)

x = rng.uniform(0, 3, 5)
y = rng.uniform(0, 5, 5)
labels = ["a", "b", "c", "d", "e"]

scatter1 = plt.scatter(
    x=x,
    y=y,
    c=range(len(labels)),
    cmap="plasma",
    vmax=len(labels)
)
handles, _ = scatter1.legend_elements()
plt.legend(handles, labels, loc="lower right")

edited Jun 14 '23 at 11:13

answered Jun 14 '23 at 11:06

Stef

28,728
2
24
52

this solution seems neat, but I get `AttributeError: 'Figure' object has no attribute 'legend_elements'`. Any idea what's going on? – Antonio Carnevali Jun 14 '23 at 12:26
Do you get the error when you just copy and paste this very example? Please verify that you call `legend_elements` on the return value of `plt.scatter` without any re-assignments to this value in-between. – Stef Jun 14 '23 at 12:53

How to add legend within a for loop

2 Answers2