Skip to content

You're a Legend! Solution


Solution 1

import matplotlib.pyplot as plt

# Make figure with one Axes
fig, ax = plt.subplots()

# Initialize an empty list to store each artist
artists = []

# Put y arrays in a list
Ys = [y_A, y_B, y_C, y_D]

# Iterate over each y array and color at the same time
for c, y in zip(colors.values(), Ys):

    # Plot the (x,y) scatter points with color c. Fetch the resulting artist
    artist, = ax.plot(x, y, linestyle='None', marker='o', color=c)

    # Append the artist to list of artists
    artists.append(artist)

# Add the legend
ax.legend(handles=artists, labels=colors.keys(), loc='upper right', framealpha=1)

Explanation

The outline of this method is to iterate over each category (A, B, C, and D), draw its scatter plot, append its artist to a list, and then add the legend.

Keep in mind, there are multiple ways to add a legend.

legend()
legend(handles, labels)
legend(handles=handles)
legend(labels)

The call signature we use here is legend(handles, labels) where handles is a list of artists and labels is a list of corresponding category names.

  1. Create a figure with a single subplot Axes.

    import matplotlib.pyplot as plt
    fig, ax = plt.subplots()
    
  2. Initialize an empty list to store each artist.

    artists = []
    

    You can think of an artist as a painter with specific instructions on how to paint something. For example, plt.plot(...) returns an artist who knows how to plot your data. Even Axes and Figure are artists as they inherit from the Artist class.

  3. Wrap the y_* arrays into a list (so we can iterate over them).

    Ys = [y_A, y_B, y_C, y_D]
    
  4. Iterate over each y_* array and color at the same time.

    for c, y in zip(colors.values(), Ys):
        # ...
    
    • colors.values() returns dict_values(['#9D44B5', '#B5446E', '#525252', '#BADEFC']) which is very similar to a list of strings.
    • zip(colors.values(), Ys) lets us iterate over each (color, y_* array) pair (which we unpack as c and y).
  5. Draw the scatter plot for the current y_* array. Then grab and append the artist to artists.

    for c, y in zip(colors.values(), Ys):
        artist, = ax.plot(x, y, linestyle='None', marker='o', color=c)
        artists.append(artist)
    
    1. ax.plot(...) returns a list with one artist.

      ax.plot(x, y, linestyle='None', marker='o', color=c)
      # [<matplotlib.lines.Line2D object at 0x13552a140>] (1)
      
      1. This is a list with one matplotlib.lines.Line2D object inside it. Line2D is a primitive Artist.

      To unpack it into the artist variable, we do artist, = ax.plot(...). (Alternatively, we could say artist = ax.plot(...)[0]).

    2. We append artist to artists via artists.append(artist).

  6. Add the legend.

    For this we use Axes.legend().

    ax.legend(
        handles=artists,       # (1)!
        labels=colors.keys(),  # (2)!
        loc='upper right',     # (3)!
        framealpha=1           # (4)!
    )
    
    1. handles is our list of artists.
    2. labels is a sequence of colors corresponding to each handle. In this case, the keys of the colors dict.
    3. loc='upper right' positions the legend in the top right corner of the Axes.
    4. framealpha=1 removes the transparency of the legend.

      framealpha should range from 0 (invisible) to 1 (non-transparent).

Solution 2

import matplotlib.pyplot as plt

# Put y arrays in a list
Ys = [y_A, y_B, y_C, y_D]

# Make figure with one Axes
fig, ax = plt.subplots()

# Iterate over each y array and color key at the same time
for c, y in zip(colors.keys(), Ys):

    # Plot the (x,y) scatter points, labelled by the current color key, c
    ax.plot(x, y, linestyle='None', marker='o', color=colors[c], label=c)

# Add the legend
ax.legend(loc='upper right', framealpha=1)

Explanation

This solution is nearly identical to Solution 1, but here we specify the label of each artist when we instantiate it.

for c, y in zip(colors.keys(), Ys):
    ax.plot(x, y, linestyle='None', marker='o', color=colors[c], label=c)

Notice the label=c parameter.

Since each artist has a label, we can use the Axes.legend() call signature without parameters to build the legend.

ax.legend()

Although, as in Solution 1, we include keywords loc='upper right' and framealpha=1 to position the legend in the top right and remove the legend's transparency.

ax.legend(loc='upper right', framealpha=1)

See the problem