Tuyen
Tuyen

Reputation: 1111

Manually set legend colors

I try to plot the following line plot, but I have difficulties to manually set legend colors. Currently the legend colors did not match the line colors. Any help would be very helpful. Thank you.

import random
import matplotlib.pyplot as plt
random.seed(10)

data=[(i, i+random.randint(1,20), random.choice(list("ABC"))) for i in range(2000,2025)]

plt.figure(figsize=(14,8))
for x, y,z in data:
    a=(x,x+y)
    b=(y+random.random(),y+random.random())
    if z=="A":
        a=(x,x)
        plt.plot(a,b,"bo",linestyle="-", linewidth=0.4, color="blue")
    elif z=="B":
        plt.plot(a,b,"bo",linestyle="-", linewidth=0.4, color="green")
    else:
        plt.plot(a,b,"bo",linestyle="-", linewidth=0.4, color="red")
ax = plt.gca()
plt.legend(['A', 'B',"C"])

Upvotes: 2

Views: 6114

Answers (3)

JohanC
JohanC

Reputation: 80339

A simple approach would be to save handles to each type of element. In the code below, handleA, = plt.plot(..., label='A') stored the line element created by plt.plot into a variable named handleA. The handle will keep its label to automatically use in the legend. (A comma is needed because plt.plot always returns a tuple, even if only one line element is created.)

import random
import matplotlib.pyplot as plt

random.seed(10)
data = [(i, i + random.randint(1, 20), random.choice(list("ABC"))) for i in range(2000, 2025)]

plt.figure(figsize=(14, 8))
for x, y, z in data:
    a = (x, x + y)
    b = (y + random.random(), y + random.random())
    if z == "A":
        a = (x, x)
        handleA, = plt.plot(a, b, '-o', linewidth=0.4, color="blue", label='A')
    elif z == "B":
        handleB, = plt.plot(a, b, '-o', linewidth=0.4, color="green", label='B')
    else:
        handleC, = plt.plot(a, b, '-o', linewidth=0.4, color="red", label='C')

plt.legend(handles=[handleA, handleB, handleC], bbox_to_anchor=(1.01, 1.01), loc='upper left')
plt.tight_layout()
plt.show()

custom legend created from handles

Upvotes: 3

warped
warped

Reputation: 9481

You can create symbols as you add data to the plot, and then vet the legend to unique entries, like so:

import random
import matplotlib.pyplot as plt
random.seed(10)

data=[(i, i+random.randint(1,20), random.choice(list("ABC"))) for i in range(2000,2025)]

plt.figure(figsize=(14,8))
for x, y,z in data:
    a=(x,x+y)
    b=(y+random.random(),y+random.random())
    if z=="A":
        a=(x,x)
        plt.plot(a,b, '-o', linewidth=0.4, color="blue", label='A')
    elif z=="B":
        plt.plot(a,b, '-o', linewidth=0.4, color="green", label='B')
    else:
        plt.plot(a,b, '-o', linewidth=0.4, color="red", label='C')

        

symbols, names = plt.gca().get_legend_handles_labels()
new_symbols, new_names = [], []
for name in sorted(list(set(names))):
    index = [i for i,n in enumerate(names) if n==name][0]
    new_symbols.append(symbols[index])
    new_names.append(name)
    
plt.legend(new_symbols, new_names)

Upvotes: 1

NsaNinja
NsaNinja

Reputation: 314

For custom generation of legends you can use this link.Composing Custom Legends


from matplotlib.patches import Patch
from matplotlib.lines import Line2D

legend_elements = [Line2D([0], [0], color='b', lw=4, label='Line'),
                   Line2D([0], [0], marker='o', color='w', label='Scatter',
                          markerfacecolor='g', markersize=15),
                   Patch(facecolor='orange', edgecolor='r',
                         label='Color Patch')]

# Create the figure
fig, ax = plt.subplots()
ax.legend(handles=legend_elements, loc='center')

plt.show()

Output will be like this: Output

Upvotes: 3

Related Questions