1"""
2===========
3Legend Demo
4===========
5
6Plotting legends in Matplotlib.
7
8There are many ways to create and customize legends in Matplotlib. Below
9we'll show a few examples for how to do so.
10
11First we'll show off how to make a legend for specific lines.
12"""
13from __future__ import (absolute_import, division,
14                        print_function, unicode_literals)
15
16import matplotlib.pyplot as plt
17import numpy as np
18from matplotlib.legend_handler import (HandlerLineCollection,
19                                       HandlerTuple)
20import matplotlib.collections as mcol
21from matplotlib.lines import Line2D
22
23t1 = np.arange(0.0, 2.0, 0.1)
24t2 = np.arange(0.0, 2.0, 0.01)
25
26fig, ax = plt.subplots()
27
28# note that plot returns a list of lines.  The "l1, = plot" usage
29# extracts the first element of the list into l1 using tuple
30# unpacking.  So l1 is a Line2D instance, not a sequence of lines
31l1, = ax.plot(t2, np.exp(-t2))
32l2, l3 = ax.plot(t2, np.sin(2 * np.pi * t2), '--o', t1, np.log(1 + t1), '.')
33l4, = ax.plot(t2, np.exp(-t2) * np.sin(2 * np.pi * t2), 's-.')
34
35ax.legend((l2, l4), ('oscillatory', 'damped'), loc='upper right', shadow=True)
36ax.set_xlabel('time')
37ax.set_ylabel('volts')
38ax.set_title('Damped oscillation')
39plt.show()
40
41
42###############################################################################
43# Next we'll demonstrate plotting more complex labels.
44
45x = np.linspace(0, 1)
46
47fig, (ax0, ax1) = plt.subplots(2, 1)
48
49# Plot the lines y=x**n for n=1..4.
50for n in range(1, 5):
51    ax0.plot(x, x**n, label="n={0}".format(n))
52leg = ax0.legend(loc="upper left", bbox_to_anchor=[0, 1],
53                 ncol=2, shadow=True, title="Legend", fancybox=True)
54leg.get_title().set_color("red")
55
56# Demonstrate some more complex labels.
57ax1.plot(x, x**2, label="multi\nline")
58half_pi = np.linspace(0, np.pi / 2)
59ax1.plot(np.sin(half_pi), np.cos(half_pi), label=r"$\frac{1}{2}\pi$")
60ax1.plot(x, 2**(x**2), label="$2^{x^2}$")
61ax1.legend(shadow=True, fancybox=True)
62
63plt.show()
64
65
66###############################################################################
67# Here we attach legends to more complex plots.
68
69fig, axes = plt.subplots(3, 1)
70top_ax, middle_ax, bottom_ax = axes
71
72top_ax.bar([0, 1, 2], [0.2, 0.3, 0.1], width=0.4, label="Bar 1",
73           align="center")
74top_ax.bar([0.5, 1.5, 2.5], [0.3, 0.2, 0.2], color="red", width=0.4,
75           label="Bar 2", align="center")
76top_ax.legend()
77
78middle_ax.errorbar([0, 1, 2], [2, 3, 1], xerr=0.4, fmt="s", label="test 1")
79middle_ax.errorbar([0, 1, 2], [3, 2, 4], yerr=0.3, fmt="o", label="test 2")
80middle_ax.errorbar([0, 1, 2], [1, 1, 3], xerr=0.4, yerr=0.3, fmt="^",
81                   label="test 3")
82middle_ax.legend()
83
84bottom_ax.stem([0.3, 1.5, 2.7], [1, 3.6, 2.7], label="stem test")
85bottom_ax.legend()
86
87plt.subplots_adjust(hspace=0.7)
88plt.show()
89
90###############################################################################
91# Now we'll showcase legend entries with more than one legend key.
92
93fig, (ax1, ax2) = plt.subplots(2, 1)
94
95# First plot: two legend keys for a single entry
96p1 = ax1.scatter([1], [5], c='r', marker='s', s=100)
97p2 = ax1.scatter([3], [2], c='b', marker='o', s=100)
98# `plot` returns a list, but we want the handle - thus the comma on the left
99p3, = ax1.plot([1, 5], [4, 4], 'm-d')
100
101# Assign two of the handles to the same legend entry by putting them in a tuple
102# and using a generic handler map (which would be used for any additional
103# tuples of handles like (p1, p3)).
104l = ax1.legend([(p1, p3), p2], ['two keys', 'one key'], scatterpoints=1,
105               numpoints=1, handler_map={tuple: HandlerTuple(ndivide=None)})
106
107# Second plot: plot two bar charts on top of each other and change the padding
108# between the legend keys
109x_left = [1, 2, 3]
110y_pos = [1, 3, 2]
111y_neg = [2, 1, 4]
112
113rneg = ax2.bar(x_left, y_neg, width=0.5, color='w', hatch='///', label='-1')
114rpos = ax2.bar(x_left, y_pos, width=0.5, color='k', label='+1')
115
116# Treat each legend entry differently by using specific `HandlerTuple`s
117l = ax2.legend([(rpos, rneg), (rneg, rpos)], ['pad!=0', 'pad=0'],
118               handler_map={(rpos, rneg): HandlerTuple(ndivide=None),
119                            (rneg, rpos): HandlerTuple(ndivide=None, pad=0.)})
120
121plt.show()
122
123###############################################################################
124# Finally, it is also possible to write custom objects that define
125# how to stylize legends.
126
127
128class HandlerDashedLines(HandlerLineCollection):
129    """
130    Custom Handler for LineCollection instances.
131    """
132    def create_artists(self, legend, orig_handle,
133                       xdescent, ydescent, width, height, fontsize, trans):
134        # figure out how many lines there are
135        numlines = len(orig_handle.get_segments())
136        xdata, xdata_marker = self.get_xdata(legend, xdescent, ydescent,
137                                             width, height, fontsize)
138        leglines = []
139        # divide the vertical space where the lines will go
140        # into equal parts based on the number of lines
141        ydata = ((height) / (numlines + 1)) * np.ones(xdata.shape, float)
142        # for each line, create the line at the proper location
143        # and set the dash pattern
144        for i in range(numlines):
145            legline = Line2D(xdata, ydata * (numlines - i) - ydescent)
146            self.update_prop(legline, orig_handle, legend)
147            # set color, dash pattern, and linewidth to that
148            # of the lines in linecollection
149            try:
150                color = orig_handle.get_colors()[i]
151            except IndexError:
152                color = orig_handle.get_colors()[0]
153            try:
154                dashes = orig_handle.get_dashes()[i]
155            except IndexError:
156                dashes = orig_handle.get_dashes()[0]
157            try:
158                lw = orig_handle.get_linewidths()[i]
159            except IndexError:
160                lw = orig_handle.get_linewidths()[0]
161            if dashes[0] is not None:
162                legline.set_dashes(dashes[1])
163            legline.set_color(color)
164            legline.set_transform(trans)
165            legline.set_linewidth(lw)
166            leglines.append(legline)
167        return leglines
168
169x = np.linspace(0, 5, 100)
170
171fig, ax = plt.subplots()
172colors = plt.rcParams['axes.prop_cycle'].by_key()['color'][:5]
173styles = ['solid', 'dashed', 'dashed', 'dashed', 'solid']
174lines = []
175for i, color, style in zip(range(5), colors, styles):
176    ax.plot(x, np.sin(x) - .1 * i, c=color, ls=style)
177
178
179# make proxy artists
180# make list of one line -- doesn't matter what the coordinates are
181line = [[(0, 0)]]
182# set up the proxy artist
183lc = mcol.LineCollection(5 * line, linestyles=styles, colors=colors)
184# create the legend
185ax.legend([lc], ['multi-line'], handler_map={type(lc): HandlerDashedLines()},
186          handlelength=2.5, handleheight=3)
187
188plt.show()
189