Matplotlib legend, add items across columns instead of down

Question

For a simple plot below, is there a way to make matplotlib populate the legend so that it fills the rows left to right, instead of first column then second column?

>>> from pylab import *
>>> x = arange(-2*pi, 2*pi, 0.1)
>>> plot(x, sin(x), label='Sine')
>>> plot(x, cos(x), label='Cosine')
>>> plot(x, arctan(x), label='Inverse tan')
>>> legend(loc=9,ncol=2)
>>> grid('on')

Answer 1

I can think of one possible way. You can order your legend items as you like. All you need to do is to switch the order so that it will give you the result you want.

import matplotlib.pyplot as plt
import numpy as np
import itertools

def flip(items, ncol):
    return itertools.chain(*[items[i::ncol] for i in range(ncol)])

x = np.arange(-2*np.pi, 2*np.pi, 0.1)
ax = plt.subplot(111)
ax.plot(x, np.sin(x), label='Sine')
ax.plot(x, np.cos(x), label='Cosine')
ax.plot(x, np.arctan(x), label='Inverse tan')

handles, labels = ax.get_legend_handles_labels()
plt.legend(flip(handles, 2), flip(labels, 2), loc=9, ncol=2)

plt.grid('on')
plt.show()

Answer 2

By default, the legend will fill all allocated columns before adding a new row. You can therefore re-order the handles and labels together to take advantage of this:

handles, labels = ax1.get_legend_handles_labels()
handles = np.concatenate((handles[::2],handles[1::2]),axis=0)
labels = np.concatenate((labels[::2],labels[1::2]),axis=0)