Search code examples
pythonmatplotlibgraph

colors in legend do not match graph colors

I am trying to make customized legends showing only some of the lines plotted. However, the color in the legend does not always match the color of the plotted line. In the example legend 1 and legend 2 are fine, but legend 3 should show blue and green, not blue and orange. See attached image. How can I fix this?

legend 3 should be blue (0) and green (2)

Here is a mwe which I used to generate the image, but note that in the real program the plot function is not directly accessible and loaded from a different file. Also, the plot function should not be modified as it is used in other programs as well.

import matplotlib.pyplot as plt
import numpy as np

def my_plot_function(x,y,ax):
    #do complicated stuff here, manipulating y for example
    ax.plot(x,y)

x = np.linspace(0,1,100)
fig = plt.figure()
ax = fig.add_subplot(111)

my_plot_function(x, x, ax)
my_plot_function(x, x**2, ax)
my_plot_function(x, x**3, ax)
my_plot_function(x, x**4, ax)

lines = ax.get_lines()
print(lines[0])
print(lines[1])
print(lines[2])
print(lines[3])

fig.legend(lines, labels=range(4), loc=1, title="legend 1")
fig.legend([lines[0],lines[1]], labels=["0","1"], loc=2, title="legend 2")
fig.legend([lines[0],lines[2]], labels=["0","2"], loc=3, title="legend 3")
plt.show()

EDIT:
To clarify what I was asking for: Plotting needs to be done by my_plot_function which is defined in a separate file and may not be modified. Thus I cannot pass it additional keywords such as label.

Solution

  • When looking at the docs of figure.legend we can see that labels has to be used with handles (and the other way). You must put both handles and labels or none of them.

    This works

    fig.legend(lines, range(4), loc=1, title="legend 1")
fig.legend((lines[0],lines[1]), ("0","1"), loc=2, title="legend 2")
fig.legend((lines[0],lines[2]), ("0","2"), loc=3, title="legend 3")

    and this works

    fig.legend(handles=lines, labels=range(4), loc=1, title="legend 1")
fig.legend(handles=(lines[0],lines[1]), labels=("0","1"), loc=2, title="legend 2")
fig.legend(handles=(lines[0],lines[2]), labels=("0","2"), loc=3, title="legend 3")