matplotlib put legend side-by-side

mysticSnake 0
Viewed 42

I am doing a series of plots inside a for loop:

import numpy as np
import matplotlib.pyplot as plt

x = np.linspace(0,1, 100)
y = np.sin(x)

for i in range(1,6):
    plt.plot(x, i*y, label=f't= {i}')
    
    plt.plot(x[::2], i*y[::2], marker='o', linestyle='None', markersize=2,label=f'a= {i}')
    
plt.legend(loc='best', ncol=2)

The output is:

enter image description here

I would like the legend to be:

enter image description here

How can I access the legend and make it like in the image above?

legend matplotlib python
1 Answers

If you don't give extra arguments to legend, it will put the artists in the legend in the order they were created and also sort them by container (but here you don't care since you only have artists belonging to ax.lines). You should sort the handles manually to get the result you desire then give it to legend.

Here it's pretty simple:

import numpy as np
import matplotlib.pyplot as plt

fig, ax = plt.subplots()
x = np.linspace(0,1, 100)
y = np.sin(x)

for i in range(1,6):
    ax.plot(x, i*y, label=f't= {i}')    
    ax.plot(x[::2], i*y[::2], marker='o', linestyle='None', markersize=2,label=f'a= {i}')

handles = [line for x in ("t", "a") for line in ax.lines  if line.get_label().startswith(x)]
# handles = [line for m in ("None", "o") for line in ax.lines if line.get_marker() == m]      
ax.legend(handles=handles, loc='best', ncol=2)

Notice that we have several solutions to order the artists. Here we can follow the pattern in the line labels, or simply their marker type.

enter image description here

This seems like a clean and object-orientated solution to me. For a solution which modifies your code much less, you could do:

for i in range(1,6):
    plt.plot(x, i*y, label=f't= {i}')
for i in range(1,6):    
    plt.plot(x[::2], i*y[::2], marker='o', linestyle='None', markersize=2,label=f'a= {i}')
Guimoute0
Related