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