Manually set legend colors using matplotlib Python-CodePudding

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"])

CodePudding user response：

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

CodePudding user response：

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

CodePudding user response：

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)