Automated legend creation in Matplotlib

Matplotlib can automatically create legends for scatter plots using the legend_elements() method. This is particularly useful when plotting data with multiple categories or varying sizes.

Basic Automated Legend

The legend_elements() method extracts legend information from scatter plots ?

import matplotlib.pyplot as plt
import numpy as np

# Set figure size
plt.figure(figsize=(8, 6))

# Generate sample data
N = 30
x = np.random.rand(N)
y = np.random.rand(N)
colors = np.random.randint(1, 4, size=N)  # 3 categories

# Create scatter plot
scatter = plt.scatter(x, y, c=colors, s=100, cmap='viridis')

# Automatically create legend for colors
legend = plt.legend(*scatter.legend_elements(), title="Categories")
plt.title("Automated Legend for Categories")
plt.show()

Legend for Both Colors and Sizes

You can create separate legends for colors and sizes using multiple calls to legend_elements() ?

import matplotlib.pyplot as plt
import numpy as np

plt.figure(figsize=(10, 6))

# Generate data with varying colors and sizes
N = 45
x = np.random.rand(N)
y = np.random.rand(N)
colors = np.random.randint(1, 5, size=N)    # 4 color categories
sizes = np.random.randint(10, 220, size=N)  # Variable sizes

# Create scatter plot
fig, ax = plt.subplots(figsize=(10, 6))
scatter = ax.scatter(x, y, c=colors, s=sizes, alpha=0.7, cmap='tab10')

# Create legend for colors
legend1 = ax.legend(*scatter.legend_elements(), 
                   loc="lower left", title="Classes")
ax.add_artist(legend1)  # Keep first legend when adding second

# Create legend for sizes
handles, labels = scatter.legend_elements(prop="sizes", alpha=0.6)
legend2 = ax.legend(handles, labels, loc="upper right", title="Sizes")

ax.set_title("Automated Legends for Colors and Sizes")
plt.show()

Customizing Legend Elements

The legend_elements() method accepts parameters to customize the legend appearance ?

import matplotlib.pyplot as plt
import numpy as np

plt.figure(figsize=(9, 6))

# Create sample data
N = 40
x = np.random.rand(N)
y = np.random.rand(N)
categories = np.random.choice(['A', 'B', 'C', 'D'], size=N)
sizes = np.random.randint(20, 200, size=N)

# Map categories to numbers for coloring
color_map = {'A': 1, 'B': 2, 'C': 3, 'D': 4}
colors = [color_map[cat] for cat in categories]

# Create scatter plot
fig, ax = plt.subplots(figsize=(9, 6))
scatter = ax.scatter(x, y, c=colors, s=sizes, alpha=0.6, cmap='Set1')

# Custom legend with specific number of entries
legend1 = ax.legend(*scatter.legend_elements(num=4), 
                   loc="lower left", title="Groups")
ax.add_artist(legend1)

# Size legend with custom formatting
handles, labels = scatter.legend_elements(prop="sizes", num=5, 
                                        func=lambda s: s/4)  # Scale down labels
legend2 = ax.legend(handles, labels, loc="upper right", 
                   title="Size Scale")

ax.set_title("Customized Automated Legends")
ax.set_xlabel("X Values")
ax.set_ylabel("Y Values")
plt.show()

Key Parameters

Conclusion

Use legend_elements() to automatically generate legends from scatter plot data. The add_artist() method allows multiple legends on the same plot, making it easy to show both color and size information.