|
| 1 | +from __future__ import (absolute_import, division, |
| 2 | + print_function, unicode_literals) |
| 3 | +import six |
| 4 | +from matplotlib import pyplot as plt |
| 5 | +import numpy as np |
| 6 | +from matplotlib.legend_handler import HandlerLineCollection |
| 7 | +import matplotlib.collections as mcol |
| 8 | +from matplotlib.lines import Line2D |
| 9 | + |
| 10 | + |
| 11 | +class HandlerDashedLines(HandlerLineCollection): |
| 12 | + """ |
| 13 | + Custom Handler for LineCollection instances. |
| 14 | + """ |
| 15 | + def create_artists(self, legend, orig_handle, |
| 16 | + xdescent, ydescent, width, height, fontsize, trans): |
| 17 | + # figure out how many lines there are |
| 18 | + numlines = len(orig_handle.get_segments()) |
| 19 | + xdata, xdata_marker = self.get_xdata(legend, xdescent, ydescent, |
| 20 | + width, height, fontsize) |
| 21 | + leglines = [] |
| 22 | + # divide the vertical space where the lines will go |
| 23 | + # into equal parts based on the number of lines |
| 24 | + ydata = ((height) / (numlines + 1)) * np.ones(xdata.shape, float) |
| 25 | + # for each line, create the line at the proper location |
| 26 | + # and set the dash pattern |
| 27 | + for i in range(numlines): |
| 28 | + legline = Line2D(xdata, ydata * (numlines - i) - ydescent) |
| 29 | + self.update_prop(legline, orig_handle, legend) |
| 30 | + # set color, dash pattern, and linewidth to that |
| 31 | + # of the lines in linecollection |
| 32 | + try: |
| 33 | + color = orig_handle.get_colors()[i] |
| 34 | + except IndexError: |
| 35 | + color = orig_handle.get_colors()[0] |
| 36 | + try: |
| 37 | + dashes = orig_handle.get_dashes()[i] |
| 38 | + except IndexError: |
| 39 | + dashes = orig_handle.get_dashes()[0] |
| 40 | + try: |
| 41 | + lw = orig_handle.get_linewidths()[i] |
| 42 | + except IndexError: |
| 43 | + lw = orig_handle.get_linewidths()[0] |
| 44 | + if dashes[0] != None: |
| 45 | + legline.set_dashes(dashes[1]) |
| 46 | + legline.set_color(color) |
| 47 | + legline.set_transform(trans) |
| 48 | + legline.set_linewidth(lw) |
| 49 | + leglines.append(legline) |
| 50 | + return leglines |
| 51 | + |
| 52 | +x = np.linspace(0, 5, 100) |
| 53 | + |
| 54 | +plt.figure() |
| 55 | +colors = ['red', 'orange', 'yellow', 'green', 'blue'] |
| 56 | +styles = ['solid', 'dashed', 'dashed', 'dashed', 'solid'] |
| 57 | +lines = [] |
| 58 | +for i, color, style in zip(range(5), colors, styles): |
| 59 | + plt.plot(x, np.sin(x) - .1 * i, c=color, ls=style) |
| 60 | + |
| 61 | + |
| 62 | +# make proxy artists |
| 63 | +# make list of one line -- doesn't matter what the coordinates are |
| 64 | +line = [[(0, 0)]] |
| 65 | +# set up the proxy artist |
| 66 | +lc = mcol.LineCollection(5 * line, linestyles=styles, colors=colors) |
| 67 | +# create the legend |
| 68 | +plt.legend([lc], ['multi-line'], handler_map={type(lc): HandlerDashedLines()}, |
| 69 | + handlelength=2.5, handleheight=3) |
| 70 | + |
| 71 | +plt.show() |
0 commit comments