Matplotlib xy line graph legend

Question

How to do add a legend to a xy line graph in Matplotlib, inside of IPython notebook? My current attempt:

x = np.linspace(0, 3*np.pi, 500)
a = plt.plot(x, np.sin(x**2))
b = plt.plot(x, np.sin(x**0.5))
plt.legend([a,b], ['square', 'square root'])

Doing this, I get the following error:

/Users/mc/.virtualenvs/kaggle/lib/python2.7/site-packages/matplotlib/legend.py:613: UserWarning: Legend does not support [] Use proxy artist instead.

http://matplotlib.sourceforge.net/users/legend_guide.html#using-proxy-artist

(str(orig_handle),)) /Users/mc/.virtualenvs/kaggle/lib/python2.7/site-packages/matplotlib/legend.py:613: UserWarning: Legend does not support [] Use proxy artist instead.

http://matplotlib.sourceforge.net/users/legend_guide.html#using-proxy-artist

(str(orig_handle),))

This command works if I do plt.scatter instead of plt.plot, but I want a line graph instead of x,y points.

Answer 1

How about doing this?

x = np.linspace(0, 3*np.pi, 500)
fig, ax = plt.subplots()
a = ax.plot(x, np.sin(x**2), label = 'square')
b = ax.plot(x, np.sin(x**0.5), label = 'square root')
handles, labels = ax.get_legend_handles_labels()
ax.legend(handles, labels)

To get this: