Altcademy - a Forbes magazine logo Best Coding Bootcamp 2023

How to make a legend in matplotlib

Understanding Legends in Matplotlib

When you're just starting out with programming, especially in the world of data visualization, you might feel like you're navigating a dense forest with all sorts of unfamiliar terms and concepts. One such concept is the "legend" in a plot or graph. Think of a legend as a guide or a map that helps you understand the symbols, colors, or line types in your chart, much like a legend on a geographical map that explains the symbols and colors used.

Why You Need a Legend

Imagine you're looking at a beautiful painting with various colors and shapes. Without an explanation, you might appreciate its beauty, but you might not understand what each color or shape represents. Similarly, in a plot, a legend tells the story of your data, making it clear what each plot element corresponds to in your dataset.

Creating Your First Legend with Matplotlib

Let's dive right in with an example. Assume you have a simple line plot. Here's how you can add a legend to it using Matplotlib:

import matplotlib.pyplot as plt

# Sample data
x = [1, 2, 3, 4]
y1 = [10, 20, 25, 30]
y2 = [15, 25, 30, 35]

# Plotting the first line
plt.plot(x, y1, label='Product A Sales')

# Plotting the second line
plt.plot(x, y2, label='Product B Sales')

# Adding a legend
plt.legend()

# Display the plot
plt.show()

In this code, we plot two lines representing sales data for products A and B. By using the label argument when plotting each line, we give a name to each dataset. The plt.legend() function then uses these labels to create a legend.

Customizing the Legend

Matplotlib offers various ways to customize the appearance of legends. You can change their location, font size, and even their frame. Let's explore some of these options:

Changing the Location

By default, Matplotlib tries to place the legend where it fits best, but sometimes you might want to place it elsewhere. You can do this by passing a location argument to plt.legend(). Here's how:

# ...
plt.legend(loc='upper left')
# ...

The loc parameter accepts various string arguments like 'upper right', 'lower left', 'center', etc., or you can use numbers between 0 and 10 to specify the location.

Adjusting the Font Size

If the text in your legend is too small or too big, you can adjust its size with the fontsize parameter:

# ...
plt.legend(loc='best', fontsize='large')
# ...

Adding a Frame

Legends come with a frame by default, but you can turn it off or customize it:

# ...
plt.legend(loc='best', frameon=False) # This turns off the frame
# ...

Or to customize the frame:

# ...
legend = plt.legend(loc='best')
legend.get_frame().set_color('lightblue') # This changes the frame color
# ...

Handling Multiple Legends

Sometimes you might have a complex plot with multiple groups of data, each needing its own legend. Here's a trick to handle such situations:

# First legend
first_legend = plt.legend(handles=[line1], loc='upper left')

# Add the legend manually to the current Axes.
ax = plt.gca().add_artist(first_legend)

# Create another legend for the second line.
plt.legend(handles=[line2], loc='lower right')

In this example, we create the first legend and add it to the current Axes, which is the area of the figure that contains the data space. Then, we create the second legend, which Matplotlib places without affecting the first one.

Using Auto-legend with plt.scatter()

For scatter plots, legends are just as important. Here's how you can automatically generate a legend for different groups of data:

# Sample scatter plot data
categories = ['Category A', 'Category B', 'Category C']
colors = ['red', 'green', 'blue']

for i, category in enumerate(categories):
    x = range(i, i+10)
    y = range(i, i+10)
    plt.scatter(x, y, color=colors[i], label=category)

# Auto-generate the legend
plt.legend()
plt.show()

In this code, we loop through our categories, creating a scatter plot for each. By setting a label for each plot, we let Matplotlib create a legend that matches colors to categories.

Intuitions and Analogies

To help you better understand the concept of legends, consider a legend as a key to a treasure map. It doesn't matter how many X's, dotted lines, or landmarks are on the map; without the key to decode what each symbol means, finding the treasure would be challenging. Similarly, in a plot, the legend acts as the key that decodes the meaning of each color, line type, or marker, allowing the viewer to understand the data's story.

Conclusion

Legends are the storytellers of your plots. They guide your audience through the visual elements, ensuring that the message you want to convey is understood. As you continue your journey into the world of programming and data visualization, remember that the legend is your ally. It helps transform a simple chart into a clear and informative visualization, just like a good map makes an adventure less daunting and more exciting. So next time you plot your data, don't forget to give your audience the legend they need to navigate the insights you're presenting. With these tools and examples at your disposal, you're now ready to create legends in Matplotlib that will make your data speak volumes. Happy plotting!