def test2(ax): # bbox=round has two optional argument. pad and rounding_size. # They can be set during the initialization. p_fancy = FancyBboxPatch((bb.xmin, bb.ymin), abs(bb.width), abs(bb.height), boxstyle="round,pad=0.1", fc=(1., .8, 1.), ec=(1., 0.5, 1.)) ax.add_patch(p_fancy) # boxstyle and its argument can be later modified with # set_boxstyle method. Note that the old attributes are simply # forgotten even if the boxstyle name is same. p_fancy.set_boxstyle("round,pad=0.1, rounding_size=0.2") # or # p_fancy.set_boxstyle("round", pad=0.1, rounding_size=0.2) ax.text(0.1, 0.8, ' boxstyle="round,pad=0.1\n rounding_size=0.2"', size=10, transform=ax.transAxes) # draws control points for the fancy box. # l = p_fancy.get_path().vertices # ax.plot(l[:,0], l[:,1], ".") draw_bbox(ax, bb)
def test2(ax): p_fancy = FancyBboxPatch((bb.xmin, bb.ymin), abs(bb.width), abs(bb.height), boxstyle="round,pad=0.1", fc=(1., .8, 1.), ec=(1., 0.5, 1.)) ax.add_patch(p_fancy) p_fancy.set_boxstyle("round,pad=0.1, rounding_size=0.2") ax.text(0.1, 0.8, ' boxstyle="round,pad=0.1\n rounding\\_size=0.2"', size=10, transform=ax.transAxes) draw_bbox(ax, bb)
def __init__(self, child, pad=None, draw_frame=False, patch_attrs=None): """ *pad* : boundary pad .. note:: *pad* need to given in points and will be scale with the renderer dpi, while *width* and *hight* need to be in pixels. """ super(PaddedBox, self).__init__() self.pad = pad self._children = [child] self.patch = FancyBboxPatch( xy=(0.0, 0.0), width=1., height=1., facecolor='w', edgecolor='k', mutation_scale=1, #self.prop.get_size_in_points(), snap=True ) self.patch.set_boxstyle("square",pad=0) if patch_attrs is not None: self.patch.update(patch_attrs) self._drawFrame = draw_frame
def _basic_outline(self, colors): if self.need != 0: #center = Circle( # self.coords, # self.r / float(20), # color=color['lc'], zorder=4, facecolor=color['lc'], fill=True, alpha=self.alpha #) if self.layout_mode: border = Circle( self.coords, 0.99 * self.r, color=colors['lc'], zorder=3, fill=False, edgecolor=colors['lc'], alpha=self.alpha, linewidth=70 * self.scale ) role_fill = Circle( self.coords, self.r, color=self.role_color, zorder=4, facecolor=self.role_color, fill=True, alpha=0.5 * self.alpha ) self.basic_structure = [role_fill, border] else: border = Circle( self.coords, self.r, color=colors['lc'], zorder=3, fill=False, edgecolor=colors['lc'], alpha=self.alpha, linewidth=0.5 ) self.basic_structure = [border, ]#center] else: center = Circle( self.coords, self.r / self._transit_fraction, color=colors['lc'], zorder=4, facecolor=colors['lc'], fill=True, alpha=self.alpha ) self.basic_structure = [center] if not self.functional: size = self.r if self.need != 0. else self.r / self._transit_fraction cross_bar = FancyBboxPatch( (0, 0), 2.4 * size, 0.2 * size, boxstyle="square,pad={}".format(0.05 * self.scale * self.label_scale), color=colors['ac'], alpha=min(1, 2 * self.alpha), zorder=6 ) rotator = Affine2D().rotate_deg(-40) corner_coords = ( self.coords[0] - size, self.coords[1] + 0.7 * size ) translator = Affine2D().translate(*corner_coords) transform = rotator + translator cross_bar.set_transform(transform) self.super_patch_collection.append(cross_bar)
def __init__(self, loc, pad=0.4, borderpad=0.5, child=None, prop=None, frameon=True, bbox_to_anchor=None, bbox_transform=None, **kwargs): """ loc is a string or an integer specifying the legend location. The valid location codes are:: 'upper right' : 1, 'upper left' : 2, 'lower left' : 3, 'lower right' : 4, 'right' : 5, 'center left' : 6, 'center right' : 7, 'lower center' : 8, 'upper center' : 9, 'center' : 10, pad : pad around the child for drawing a frame. given in fraction of fontsize. borderpad : pad between offsetbox frame and the bbox_to_anchor, child : OffsetBox instance that will be anchored. prop : font property. This is only used as a reference for paddings. frameon : draw a frame box if True. bbox_to_anchor : bbox to anchor. Use self.axes.bbox if None. bbox_transform : with which the bbox_to_anchor will be transformed. """ super(AnchoredOffsetbox, self).__init__(**kwargs) self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform) self.set_child(child) self.loc = loc self.borderpad=borderpad self.pad = pad if prop is None: self.prop=FontProperties(size=rcParams["legend.fontsize"]) elif isinstance(prop, dict): self.prop=FontProperties(**prop) if "size" not in prop: self.prop.set_size(rcParams["legend.fontsize"]) else: self.prop = prop self.patch = FancyBboxPatch( xy=(0.0, 0.0), width=1., height=1., facecolor='w', edgecolor='k', mutation_scale=self.prop.get_size_in_points(), snap=True ) self.patch.set_boxstyle("square",pad=0) self._drawFrame = frameon
def __init__(self, offsetbox, xy, xybox=None, xycoords='data', boxcoords=None, frameon=True, pad=0.4, # BboxPatch annotation_clip=None, box_alignment=(0.5, 0.5), bboxprops=None, arrowprops=None, fontsize=None, **kwargs): """ *offsetbox* : OffsetBox instance *xycoords* : same as Annotation but can be a tuple of two strings which are interpreted as x and y coordinates. *boxcoords* : similar to textcoords as Annotation but can be a tuple of two strings which are interpreted as x and y coordinates. *box_alignment* : a tuple of two floats for a vertical and horizontal alignment of the offset box w.r.t. the *boxcoords*. The lower-left corner is (0.0) and upper-right corner is (1.1). other parameters are identical to that of Annotation. """ self.offsetbox = offsetbox self.arrowprops = arrowprops self.set_fontsize(fontsize) if arrowprops is not None: self._arrow_relpos = self.arrowprops.pop("relpos", (0.5, 0.5)) self.arrow_patch = FancyArrowPatch((0, 0), (1,1), **self.arrowprops) else: self._arrow_relpos = None self.arrow_patch = None _AnnotationBase.__init__(self, xy, xytext=xybox, xycoords=xycoords, textcoords=boxcoords, annotation_clip=annotation_clip) martist.Artist.__init__(self, **kwargs) self._box_alignment = box_alignment self.patch = FancyBboxPatch( xy=(0.0, 0.0), width=1., height=1., facecolor='w', edgecolor='k', mutation_scale=self.prop.get_size_in_points(), snap=True ) self.patch.set_boxstyle("square",pad=pad) if bboxprops: self.patch.set(**bboxprops) self._drawFrame = frameon
class Legend(Artist): """ Place a legend on the axes at location loc. Labels are a sequence of strings and loc can be a string or an integer specifying the legend location The location codes are:: 'best' : 0, (only implemented for axis legends) 'upper right' : 1, 'upper left' : 2, 'lower left' : 3, 'lower right' : 4, 'right' : 5, 'center left' : 6, 'center right' : 7, 'lower center' : 8, 'upper center' : 9, 'center' : 10, loc can be a tuple of the noramilzed coordinate values with respect its parent. Return value is a sequence of text, line instances that make up the legend """ codes = {'best' : 0, # only implemented for axis legends 'upper right' : 1, 'upper left' : 2, 'lower left' : 3, 'lower right' : 4, 'right' : 5, 'center left' : 6, 'center right' : 7, 'lower center' : 8, 'upper center' : 9, 'center' : 10, } zorder = 5 def __str__(self): return "Legend" def __init__(self, parent, handles, labels, loc = None, numpoints = None, # the number of points in the legend line markerscale = None, # the relative size of legend markers vs. original scatterpoints = 3, # TODO: may be an rcParam scatteryoffsets=None, prop = None, # properties for the legend texts # the following dimensions are in axes coords pad = None, # deprecated; use borderpad labelsep = None, # deprecated; use labelspacing handlelen = None, # deprecated; use handlelength handletextsep = None, # deprecated; use handletextpad axespad = None, # deprecated; use borderaxespad # spacing & pad defined as a fractionof the font-size borderpad = None, # the whitespace inside the legend border labelspacing=None, #the vertical space between the legend entries handlelength=None, # the length of the legend handles handletextpad=None, # the pad between the legend handle and text borderaxespad=None, # the pad between the axes and legend border columnspacing=None, # spacing between columns ncol=1, # number of columns mode=None, # mode for horizontal distribution of columns. None, "expand" fancybox=None, # True use a fancy box, false use a rounded box, none use rc shadow = None, ): """ - *parent* : the artist that contains the legend - *handles* : a list of artists (lines, patches) to add to the legend - *labels* : a list of strings to label the legend Optional keyword arguments: ================ ================================================================== Keyword Description ================ ================================================================== loc a location code or a tuple of coordinates numpoints the number of points in the legend line prop the font property markerscale the relative size of legend markers vs. original fancybox if True, draw a frame with a round fancybox. If None, use rc shadow if True, draw a shadow behind legend scatteryoffsets a list of yoffsets for scatter symbols in legend borderpad the fractional whitespace inside the legend border labelspacing the vertical space between the legend entries handlelength the length of the legend handles handletextpad the pad between the legend handle and text borderaxespad the pad between the axes and legend border columnspacing the spacing between columns ================ ================================================================== The dimensions of pad and spacing are given as a fraction of the fontsize. Values from rcParams will be used if None. """ from matplotlib.axes import Axes # local import only to avoid circularity from matplotlib.figure import Figure # local import only to avoid circularity Artist.__init__(self) if prop is None: self.prop=FontProperties(size=rcParams["legend.fontsize"]) else: self.prop=prop self.fontsize = self.prop.get_size_in_points() propnames=['numpoints', 'markerscale', 'shadow', "columnspacing", "scatterpoints"] localdict = locals() for name in propnames: if localdict[name] is None: value = rcParams["legend."+name] else: value = localdict[name] setattr(self, name, value) # Take care the deprecated keywords deprecated_kwds = {"pad":"borderpad", "labelsep":"labelspacing", "handlelen":"handlelength", "handletextsep":"handletextpad", "axespad":"borderaxespad"} # convert values of deprecated keywords (ginve in axes coords) # to new vaules in a fraction of the font size # conversion factor bbox = parent.bbox axessize_fontsize = min(bbox.width, bbox.height)/self.fontsize for k, v in deprecated_kwds.items(): # use deprecated value if not None and if their newer # counter part is None. if localdict[k] is not None and localdict[v] is None: warnings.warn("Use '%s' instead of '%s'." % (v, k), DeprecationWarning) setattr(self, v, localdict[k]*axessize_fontsize) continue # Otherwise, use new keywords if localdict[v] is None: setattr(self, v, rcParams["legend."+v]) else: setattr(self, v, localdict[v]) del localdict self._ncol = ncol if self.numpoints <= 0: raise ValueError("numpoints must be >= 0; it was %d"% numpoints) # introduce y-offset for handles of the scatter plot if scatteryoffsets is None: self._scatteryoffsets = np.array([3./8., 4./8., 2.5/8.]) else: self._scatteryoffsets = np.asarray(scatteryoffsets) reps = int(self.numpoints / len(self._scatteryoffsets)) + 1 self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[:self.scatterpoints] # _legend_box is an OffsetBox instance that contains all # legend items and will be initialized from _init_legend_box() # method. self._legend_box = None if isinstance(parent,Axes): self.isaxes = True self.set_figure(parent.figure) elif isinstance(parent,Figure): self.isaxes = False self.set_figure(parent) else: raise TypeError("Legend needs either Axes or Figure as parent") self.parent = parent if loc is None: loc = rcParams["legend.loc"] if not self.isaxes and loc in [0,'best']: loc = 'upper right' if is_string_like(loc): if loc not in self.codes: if self.isaxes: warnings.warn('Unrecognized location "%s". Falling back on "best"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes.keys()))) loc = 0 else: warnings.warn('Unrecognized location "%s". Falling back on "upper right"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes.keys()))) loc = 1 else: loc = self.codes[loc] if not self.isaxes and loc == 0: warnings.warn('Automatic legend placement (loc="best") not implemented for figure legend. ' 'Falling back on "upper right".') loc = 1 self._loc = loc self._mode = mode # We use FancyBboxPatch to draw a legend frame. The location # and size of the box will be updated during the drawing time. self.legendPatch = FancyBboxPatch( xy=(0.0, 0.0), width=1., height=1., facecolor='w', edgecolor='k', mutation_scale=self.fontsize, snap=True ) # The width and height of the legendPatch will be set (in the # draw()) to the length that includes the padding. Thus we set # pad=0 here. if fancybox is None: fancybox = rcParams["legend.fancybox"] if fancybox == True: self.legendPatch.set_boxstyle("round",pad=0, rounding_size=0.2) else: self.legendPatch.set_boxstyle("square",pad=0) self._set_artist_props(self.legendPatch) self._drawFrame = True # init with null renderer self._init_legend_box(handles, labels) self._last_fontsize_points = self.fontsize def _set_artist_props(self, a): """ set the boilerplate props for artists added to axes """ a.set_figure(self.figure) for c in self.get_children(): c.set_figure(self.figure) a.set_transform(self.get_transform()) def _findoffset_best(self, width, height, xdescent, ydescent, renderer): "Heper function to locate the legend at its best position" ox, oy = self._find_best_position(width, height, renderer) return ox+xdescent, oy+ydescent def _findoffset_loc(self, width, height, xdescent, ydescent, renderer): "Heper function to locate the legend using the location code" if iterable(self._loc) and len(self._loc)==2: # when loc is a tuple of axes(or figure) coordinates. fx, fy = self._loc bbox = self.parent.bbox x, y = bbox.x0 + bbox.width * fx, bbox.y0 + bbox.height * fy else: bbox = Bbox.from_bounds(0, 0, width, height) x, y = self._get_anchored_bbox(self._loc, bbox, self.parent.bbox, renderer) return x+xdescent, y+ydescent def draw(self, renderer): "Draw everything that belongs to the legend" if not self.get_visible(): return self._update_legend_box(renderer) renderer.open_group('legend') # find_offset function will be provided to _legend_box and # _legend_box will draw itself at the location of the return # value of the find_offset. if self._loc == 0: _findoffset = self._findoffset_best else: _findoffset = self._findoffset_loc def findoffset(width, height, xdescent, ydescent): return _findoffset(width, height, xdescent, ydescent, renderer) self._legend_box.set_offset(findoffset) fontsize = renderer.points_to_pixels(self.fontsize) # if mode == fill, set the width of the legend_box to the # width of the paret (minus pads) if self._mode in ["expand"]: pad = 2*(self.borderaxespad+self.borderpad)*fontsize self._legend_box.set_width(self.parent.bbox.width-pad) if self._drawFrame: # update the location and size of the legend bbox = self._legend_box.get_window_extent(renderer) self.legendPatch.set_bounds(bbox.x0, bbox.y0, bbox.width, bbox.height) self.legendPatch.set_mutation_scale(fontsize) if self.shadow: shadow = Shadow(self.legendPatch, 2, -2) shadow.draw(renderer) self.legendPatch.draw(renderer) self._legend_box.draw(renderer) renderer.close_group('legend') def _approx_text_height(self, renderer=None): """ Return the approximate height of the text. This is used to place the legend handle. """ if renderer is None: return self.fontsize else: return renderer.points_to_pixels(self.fontsize) def _init_legend_box(self, handles, labels): """ Initiallize the legend_box. The legend_box is an instance of the OffsetBox, which is packed with legend handles and texts. Once packed, their location is calculated during the drawing time. """ fontsize = self.fontsize # legend_box is a HPacker, horizontally packed with # columns. Each column is a VPacker, vertically packed with # legend items. Each legend item is HPacker packed with # legend handleBox and labelBox. handleBox is an instance of # offsetbox.DrawingArea which contains legend handle. labelBox # is an instance of offsetbox.TextArea which contains legend # text. text_list = [] # the list of text instances handle_list = [] # the list of text instances label_prop = dict(verticalalignment='baseline', horizontalalignment='left', fontproperties=self.prop, ) labelboxes = [] for l in labels: textbox = TextArea(l, textprops=label_prop, multilinebaseline=True, minimumdescent=True) text_list.append(textbox._text) labelboxes.append(textbox) handleboxes = [] # The approximate height and descent of text. These values are # only used for plotting the legend handle. height = self._approx_text_height() * 0.7 descent = 0. # each handle needs to be drawn inside a box of (x, y, w, h) = # (0, -descent, width, height). And their corrdinates should # be given in the display coordinates. # NOTE : the coordinates will be updated again in # _update_legend_box() method. # The transformation of each handle will be automatically set # to self.get_trasnform(). If the artist does not uses its # default trasnform (eg, Collections), you need to # manually set their transform to the self.get_transform(). for handle in handles: if isinstance(handle, RegularPolyCollection): npoints = self.scatterpoints else: npoints = self.numpoints if npoints > 1: # we put some pad here to compensate the size of the # marker xdata = np.linspace(0.3*fontsize, (self.handlelength-0.3)*fontsize, npoints) xdata_marker = xdata elif npoints == 1: xdata = np.linspace(0, self.handlelength*fontsize, 2) xdata_marker = [0.5*self.handlelength*fontsize] if isinstance(handle, Line2D): ydata = ((height-descent)/2.)*np.ones(xdata.shape, float) legline = Line2D(xdata, ydata) legline.update_from(handle) self._set_artist_props(legline) # after update legline.set_clip_box(None) legline.set_clip_path(None) legline.set_drawstyle('default') legline.set_marker('None') handle_list.append(legline) legline_marker = Line2D(xdata_marker, ydata[:len(xdata_marker)]) legline_marker.update_from(handle) self._set_artist_props(legline_marker) legline_marker.set_clip_box(None) legline_marker.set_clip_path(None) legline_marker.set_linestyle('None') # we don't want to add this to the return list because # the texts and handles are assumed to be in one-to-one # correpondence. legline._legmarker = legline_marker elif isinstance(handle, Patch): p = Rectangle(xy=(0., 0.), width = self.handlelength*fontsize, height=(height-descent), ) p.update_from(handle) self._set_artist_props(p) p.set_clip_box(None) p.set_clip_path(None) handle_list.append(p) elif isinstance(handle, LineCollection): ydata = ((height-descent)/2.)*np.ones(xdata.shape, float) legline = Line2D(xdata, ydata) self._set_artist_props(legline) legline.set_clip_box(None) legline.set_clip_path(None) lw = handle.get_linewidth()[0] dashes = handle.get_dashes()[0] color = handle.get_colors()[0] legline.set_color(color) legline.set_linewidth(lw) legline.set_dashes(dashes) handle_list.append(legline) elif isinstance(handle, RegularPolyCollection): #ydata = self._scatteryoffsets ydata = height*self._scatteryoffsets size_max, size_min = max(handle.get_sizes()),\ min(handle.get_sizes()) # we may need to scale these sizes by "markerscale" # attribute. But other handle types does not seem # to care about this attribute and it is currently ignored. if self.scatterpoints < 4: sizes = [.5*(size_max+size_min), size_max, size_min] else: sizes = (size_max-size_min)*np.linspace(0,1,self.scatterpoints)+size_min p = type(handle)(handle.get_numsides(), rotation=handle.get_rotation(), sizes=sizes, offsets=zip(xdata_marker,ydata), transOffset=self.get_transform(), ) p.update_from(handle) p.set_figure(self.figure) p.set_clip_box(None) p.set_clip_path(None) handle_list.append(p) else: handle_list.append(None) handlebox = DrawingArea(width=self.handlelength*fontsize, height=height, xdescent=0., ydescent=descent) handle = handle_list[-1] handlebox.add_artist(handle) if hasattr(handle, "_legmarker"): handlebox.add_artist(handle._legmarker) handleboxes.append(handlebox) # We calculate number of lows in each column. The first # (num_largecol) columns will have (nrows+1) rows, and remaing # (num_smallcol) columns will have (nrows) rows. nrows, num_largecol = divmod(len(handleboxes), self._ncol) num_smallcol = self._ncol-num_largecol # starting index of each column and number of rows in it. largecol = safezip(range(0, num_largecol*(nrows+1), (nrows+1)), [nrows+1] * num_largecol) smallcol = safezip(range(num_largecol*(nrows+1), len(handleboxes), nrows), [nrows] * num_smallcol) handle_label = safezip(handleboxes, labelboxes) columnbox = [] for i0, di in largecol+smallcol: # pack handleBox and labelBox into itemBox itemBoxes = [HPacker(pad=0, sep=self.handletextpad*fontsize, children=[h, t], align="baseline") for h, t in handle_label[i0:i0+di]] # minimumdescent=False for the text of the last row of the column itemBoxes[-1].get_children()[1].set_minimumdescent(False) # pack columnBox columnbox.append(VPacker(pad=0, sep=self.labelspacing*fontsize, align="baseline", children=itemBoxes)) if self._mode == "expand": mode = "expand" else: mode = "fixed" sep = self.columnspacing*fontsize self._legend_box = HPacker(pad=self.borderpad*fontsize, sep=sep, align="baseline", mode=mode, children=columnbox) self._legend_box.set_figure(self.figure) self.texts = text_list self.legendHandles = handle_list def _update_legend_box(self, renderer): """ Update the dimension of the legend_box. This is required becuase the paddings, the hadle size etc. depends on the dpi of the renderer. """ # fontsize in points. fontsize = renderer.points_to_pixels(self.fontsize) if self._last_fontsize_points == fontsize: # no update is needed return # each handle needs to be drawn inside a box of # (x, y, w, h) = (0, -descent, width, height). # And their corrdinates should be given in the display coordinates. # The approximate height and descent of text. These values are # only used for plotting the legend handle. height = self._approx_text_height(renderer) * 0.7 descent = 0. for handle in self.legendHandles: if isinstance(handle, RegularPolyCollection): npoints = self.scatterpoints else: npoints = self.numpoints if npoints > 1: # we put some pad here to compensate the size of the # marker xdata = np.linspace(0.3*fontsize, (self.handlelength-0.3)*fontsize, npoints) xdata_marker = xdata elif npoints == 1: xdata = np.linspace(0, self.handlelength*fontsize, 2) xdata_marker = [0.5*self.handlelength*fontsize] if isinstance(handle, Line2D): legline = handle ydata = ((height-descent)/2.)*np.ones(xdata.shape, float) legline.set_data(xdata, ydata) legline_marker = legline._legmarker legline_marker.set_data(xdata_marker, ydata[:len(xdata_marker)]) elif isinstance(handle, Patch): p = handle p.set_bounds(0., 0., self.handlelength*fontsize, (height-descent), ) elif isinstance(handle, RegularPolyCollection): p = handle ydata = height*self._scatteryoffsets p.set_offsets(zip(xdata_marker,ydata)) # correction factor cor = fontsize / self._last_fontsize_points # helper function to iterate over all children def all_children(parent): yield parent for c in parent.get_children(): for cc in all_children(c): yield cc #now update paddings for box in all_children(self._legend_box): if isinstance(box, PackerBase): box.pad = box.pad * cor box.sep = box.sep * cor elif isinstance(box, DrawingArea): box.width = self.handlelength*fontsize box.height = height box.xdescent = 0. box.ydescent=descent self._last_fontsize_points = fontsize def _auto_legend_data(self): """ Returns list of vertices and extents covered by the plot. Returns a two long list. First element is a list of (x, y) vertices (in display-coordinates) covered by all the lines and line collections, in the legend's handles. Second element is a list of bounding boxes for all the patches in the legend's handles. """ assert self.isaxes # should always hold because function is only called internally ax = self.parent vertices = [] bboxes = [] lines = [] for handle in ax.lines: assert isinstance(handle, Line2D) path = handle.get_path() trans = handle.get_transform() tpath = trans.transform_path(path) lines.append(tpath) for handle in ax.patches: assert isinstance(handle, Patch) if isinstance(handle, Rectangle): transform = handle.get_data_transform() bboxes.append(handle.get_bbox().transformed(transform)) else: transform = handle.get_transform() bboxes.append(handle.get_path().get_extents(transform)) return [vertices, bboxes, lines] def draw_frame(self, b): 'b is a boolean. Set draw frame to b' self._drawFrame = b def get_children(self): 'return a list of child artists' children = [] if self._legend_box: children.append(self._legend_box) return children def get_frame(self): 'return the Rectangle instance used to frame the legend' return self.legendPatch def get_lines(self): 'return a list of lines.Line2D instances in the legend' return [h for h in self.legendHandles if isinstance(h, Line2D)] def get_patches(self): 'return a list of patch instances in the legend' return silent_list('Patch', [h for h in self.legendHandles if isinstance(h, Patch)]) def get_texts(self): 'return a list of text.Text instance in the legend' return silent_list('Text', self.texts) def get_window_extent(self): 'return a extent of the the legend' return self.legendPatch.get_window_extent() def _get_anchored_bbox(self, loc, bbox, parentbbox, renderer): """ Place the *bbox* inside the *parentbbox* according to a given location code. Return the (x,y) coordinate of the bbox. - loc: a location code in range(1, 11). This corresponds to the possible values for self._loc, excluding "best". - bbox: bbox to be placed, display coodinate units. - parentbbox: a parent box which will contain the bbox. In display coordinates. """ assert loc in range(1,11) # called only internally BEST, UR, UL, LL, LR, R, CL, CR, LC, UC, C = range(11) anchor_coefs={UR:"NE", UL:"NW", LL:"SW", LR:"SE", R:"E", CL:"W", CR:"E", LC:"S", UC:"N", C:"C"} c = anchor_coefs[loc] fontsize = renderer.points_to_pixels(self.fontsize) container = parentbbox.padded(-(self.borderaxespad) * fontsize) anchored_box = bbox.anchored(c, container=container) return anchored_box.x0, anchored_box.y0 def _find_best_position(self, width, height, renderer, consider=None): """ Determine the best location to place the legend. `consider` is a list of (x, y) pairs to consider as a potential lower-left corner of the legend. All are display coords. """ assert self.isaxes # should always hold because function is only called internally verts, bboxes, lines = self._auto_legend_data() bbox = Bbox.from_bounds(0, 0, width, height) consider = [self._get_anchored_bbox(x, bbox, self.parent.bbox, renderer) for x in range(1, len(self.codes))] #tx, ty = self.legendPatch.get_x(), self.legendPatch.get_y() candidates = [] for l, b in consider: legendBox = Bbox.from_bounds(l, b, width, height) badness = 0 badness = legendBox.count_contains(verts) badness += legendBox.count_overlaps(bboxes) for line in lines: if line.intersects_bbox(legendBox): badness += 1 ox, oy = l, b if badness == 0: return ox, oy candidates.append((badness, (l, b))) # rather than use min() or list.sort(), do this so that we are assured # that in the case of two equal badnesses, the one first considered is # returned. # NOTE: list.sort() is stable.But leave as it is for now. -JJL minCandidate = candidates[0] for candidate in candidates: if candidate[0] < minCandidate[0]: minCandidate = candidate ox, oy = minCandidate[1] return ox, oy
for coll in list(ax.collections): coll.remove() bounds = np.array( [[mne_path.vertices[:, ii].min(), mne_path.vertices[:, ii].max()] for ii in range(2)]) bounds *= (plot_dims / dims) xy = np.mean(bounds, axis=1) - [100, 0] r = np.diff(bounds, axis=1).max() * 1.2 w, h = r, r * (2 / 3) box_xy = [xy[0] - w * 0.5, xy[1] - h * (2 / 5)] ax.set_ylim(box_xy[1] + h * 1.001, box_xy[1] - h * 0.001) patch = FancyBboxPatch(box_xy, w, h, clip_on=False, zorder=-1, fc='k', ec='none', alpha=0.75, boxstyle="round,rounding_size=200.0", mutation_aspect=1) ax.add_patch(patch) fig.set_size_inches((512 / dpi, 512 * (h / w) / dpi)) plt.savefig(op.join(data_dir, 'mne_splash.png'), transparent=True) patch.remove() # modify to make an icon ax.patches.pop(-1) # no tag line for our icon patch = Ellipse(xy, r, r, clip_on=False, zorder=-1, fc='k') ax.add_patch(patch) ax.set_ylim(xy[1] + r / 1.9, xy[1] - r / 1.9) fig.set_size_inches((256 / dpi, 256 / dpi))
class Legend(Artist): """ Place a legend on the axes at location loc. """ codes = {'best': 0, # only implemented for axes legends 'upper right': 1, 'upper left': 2, 'lower left': 3, 'lower right': 4, 'right': 5, 'center left': 6, 'center right': 7, 'lower center': 8, 'upper center': 9, 'center': 10, } zorder = 5 def __str__(self): return "Legend" @docstring.dedent_interpd def __init__(self, parent, handles, labels, loc=None, numpoints=None, # the number of points in the legend line markerscale=None, # the relative size of legend markers # vs. original markerfirst=True, # controls ordering (left-to-right) of # legend marker and label scatterpoints=None, # number of scatter points scatteryoffsets=None, prop=None, # properties for the legend texts fontsize=None, # keyword to set font size directly # spacing & pad defined as a fraction of the font-size borderpad=None, # the whitespace inside the legend border labelspacing=None, # the vertical space between the legend # entries handlelength=None, # the length of the legend handles handleheight=None, # the height of the legend handles handletextpad=None, # the pad between the legend handle # and text borderaxespad=None, # the pad between the axes and legend # border columnspacing=None, # spacing between columns ncol=1, # number of columns mode=None, # mode for horizontal distribution of columns. # None, "expand" fancybox=None, # True use a fancy box, false use a rounded # box, none use rc shadow=None, title=None, # set a title for the legend title_fontsize=None, # set to ax.fontsize if None framealpha=None, # set frame alpha edgecolor=None, # frame patch edgecolor facecolor=None, # frame patch facecolor bbox_to_anchor=None, # bbox that the legend will be anchored. bbox_transform=None, # transform for the bbox frameon=None, # draw frame handler_map=None, ): """ Parameters ---------- parent : `~matplotlib.axes.Axes` or `.Figure` The artist that contains the legend. handles : sequence of `.Artist` A list of Artists (lines, patches) to be added to the legend. labels : sequence of strings A list of labels to show next to the artists. The length of handles and labels should be the same. If they are not, they are truncated to the smaller of both lengths. Other Parameters ---------------- %(_legend_kw_doc)s Notes ----- Users can specify any arbitrary location for the legend using the *bbox_to_anchor* keyword argument. bbox_to_anchor can be an instance of BboxBase(or its derivatives) or a tuple of 2 or 4 floats. See :meth:`set_bbox_to_anchor` for more detail. The legend location can be specified by setting *loc* with a tuple of 2 floats, which is interpreted as the lower-left corner of the legend in the normalized axes coordinate. """ # local import only to avoid circularity from matplotlib.axes import Axes from matplotlib.figure import Figure Artist.__init__(self) if prop is None: if fontsize is not None: self.prop = FontProperties(size=fontsize) else: self.prop = FontProperties(size=rcParams["legend.fontsize"]) elif isinstance(prop, dict): self.prop = FontProperties(**prop) if "size" not in prop: self.prop.set_size(rcParams["legend.fontsize"]) else: self.prop = prop self._fontsize = self.prop.get_size_in_points() self.texts = [] self.legendHandles = [] self._legend_title_box = None #: A dictionary with the extra handler mappings for this Legend #: instance. self._custom_handler_map = handler_map locals_view = locals() for name in ["numpoints", "markerscale", "shadow", "columnspacing", "scatterpoints", "handleheight", 'borderpad', 'labelspacing', 'handlelength', 'handletextpad', 'borderaxespad']: if locals_view[name] is None: value = rcParams["legend." + name] else: value = locals_view[name] setattr(self, name, value) del locals_view # trim handles and labels if illegal label... _lab, _hand = [], [] for label, handle in zip(labels, handles): if isinstance(label, str) and label.startswith('_'): cbook._warn_external('The handle {!r} has a label of {!r} ' 'which cannot be automatically added to' ' the legend.'.format(handle, label)) else: _lab.append(label) _hand.append(handle) labels, handles = _lab, _hand handles = list(handles) if len(handles) < 2: ncol = 1 self._ncol = ncol if self.numpoints <= 0: raise ValueError("numpoints must be > 0; it was %d" % numpoints) # introduce y-offset for handles of the scatter plot if scatteryoffsets is None: self._scatteryoffsets = np.array([3. / 8., 4. / 8., 2.5 / 8.]) else: self._scatteryoffsets = np.asarray(scatteryoffsets) reps = self.scatterpoints // len(self._scatteryoffsets) + 1 self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[:self.scatterpoints] # _legend_box is an OffsetBox instance that contains all # legend items and will be initialized from _init_legend_box() # method. self._legend_box = None if isinstance(parent, Axes): self.isaxes = True self.axes = parent self.set_figure(parent.figure) elif isinstance(parent, Figure): self.isaxes = False self.set_figure(parent) else: raise TypeError("Legend needs either Axes or Figure as parent") self.parent = parent self._loc_used_default = loc is None if loc is None: loc = rcParams["legend.loc"] if not self.isaxes and loc in [0, 'best']: loc = 'upper right' if isinstance(loc, str): if loc not in self.codes: if self.isaxes: cbook.warn_deprecated( "3.1", message="Unrecognized location {!r}. Falling " "back on 'best'; valid locations are\n\t{}\n" "This will raise an exception %(removal)s." .format(loc, '\n\t'.join(self.codes))) loc = 0 else: cbook.warn_deprecated( "3.1", message="Unrecognized location {!r}. Falling " "back on 'upper right'; valid locations are\n\t{}\n'" "This will raise an exception %(removal)s." .format(loc, '\n\t'.join(self.codes))) loc = 1 else: loc = self.codes[loc] if not self.isaxes and loc == 0: cbook.warn_deprecated( "3.1", message="Automatic legend placement (loc='best') not " "implemented for figure legend. Falling back on 'upper " "right'. This will raise an exception %(removal)s.") loc = 1 self._mode = mode self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform) # We use FancyBboxPatch to draw a legend frame. The location # and size of the box will be updated during the drawing time. if facecolor is None: facecolor = rcParams["legend.facecolor"] if facecolor == 'inherit': facecolor = rcParams["axes.facecolor"] if edgecolor is None: edgecolor = rcParams["legend.edgecolor"] if edgecolor == 'inherit': edgecolor = rcParams["axes.edgecolor"] self.legendPatch = FancyBboxPatch( xy=(0.0, 0.0), width=1., height=1., facecolor=facecolor, edgecolor=edgecolor, mutation_scale=self._fontsize, snap=True ) # The width and height of the legendPatch will be set (in the # draw()) to the length that includes the padding. Thus we set # pad=0 here. if fancybox is None: fancybox = rcParams["legend.fancybox"] if fancybox: self.legendPatch.set_boxstyle("round", pad=0, rounding_size=0.2) else: self.legendPatch.set_boxstyle("square", pad=0) self._set_artist_props(self.legendPatch) self._drawFrame = frameon if frameon is None: self._drawFrame = rcParams["legend.frameon"] # init with null renderer self._init_legend_box(handles, labels, markerfirst) # If shadow is activated use framealpha if not # explicitly passed. See Issue 8943 if framealpha is None: if shadow: self.get_frame().set_alpha(1) else: self.get_frame().set_alpha(rcParams["legend.framealpha"]) else: self.get_frame().set_alpha(framealpha) tmp = self._loc_used_default self._set_loc(loc) self._loc_used_default = tmp # ignore changes done by _set_loc # figure out title fontsize: if title_fontsize is None: title_fontsize = rcParams['legend.title_fontsize'] tprop = FontProperties(size=title_fontsize) self.set_title(title, prop=tprop) self._last_fontsize_points = self._fontsize self._draggable = None def _set_artist_props(self, a): """ Set the boilerplate props for artists added to axes. """ a.set_figure(self.figure) if self.isaxes: # a.set_axes(self.axes) a.axes = self.axes a.set_transform(self.get_transform()) def _set_loc(self, loc): # find_offset function will be provided to _legend_box and # _legend_box will draw itself at the location of the return # value of the find_offset. self._loc_used_default = False self._loc_real = loc self.stale = True self._legend_box.set_offset(self._findoffset) def _get_loc(self): return self._loc_real _loc = property(_get_loc, _set_loc) def _findoffset(self, width, height, xdescent, ydescent, renderer): "Helper function to locate the legend." if self._loc == 0: # "best". x, y = self._find_best_position(width, height, renderer) elif self._loc in Legend.codes.values(): # Fixed location. bbox = Bbox.from_bounds(0, 0, width, height) x, y = self._get_anchored_bbox(self._loc, bbox, self.get_bbox_to_anchor(), renderer) else: # Axes or figure coordinates. fx, fy = self._loc bbox = self.get_bbox_to_anchor() x, y = bbox.x0 + bbox.width * fx, bbox.y0 + bbox.height * fy return x + xdescent, y + ydescent @allow_rasterization def draw(self, renderer): "Draw everything that belongs to the legend." if not self.get_visible(): return renderer.open_group('legend') fontsize = renderer.points_to_pixels(self._fontsize) # if mode == fill, set the width of the legend_box to the # width of the parent (minus pads) if self._mode in ["expand"]: pad = 2 * (self.borderaxespad + self.borderpad) * fontsize self._legend_box.set_width(self.get_bbox_to_anchor().width - pad) # update the location and size of the legend. This needs to # be done in any case to clip the figure right. bbox = self._legend_box.get_window_extent(renderer) self.legendPatch.set_bounds(bbox.x0, bbox.y0, bbox.width, bbox.height) self.legendPatch.set_mutation_scale(fontsize) if self._drawFrame: if self.shadow: shadow = Shadow(self.legendPatch, 2, -2) shadow.draw(renderer) self.legendPatch.draw(renderer) self._legend_box.draw(renderer) renderer.close_group('legend') self.stale = False def _approx_text_height(self, renderer=None): """ Return the approximate height of the text. This is used to place the legend handle. """ if renderer is None: return self._fontsize else: return renderer.points_to_pixels(self._fontsize) # _default_handler_map defines the default mapping between plot # elements and the legend handlers. _default_handler_map = { StemContainer: legend_handler.HandlerStem(), ErrorbarContainer: legend_handler.HandlerErrorbar(), Line2D: legend_handler.HandlerLine2D(), Patch: legend_handler.HandlerPatch(), LineCollection: legend_handler.HandlerLineCollection(), RegularPolyCollection: legend_handler.HandlerRegularPolyCollection(), CircleCollection: legend_handler.HandlerCircleCollection(), BarContainer: legend_handler.HandlerPatch( update_func=legend_handler.update_from_first_child), tuple: legend_handler.HandlerTuple(), PathCollection: legend_handler.HandlerPathCollection(), PolyCollection: legend_handler.HandlerPolyCollection() } # (get|set|update)_default_handler_maps are public interfaces to # modify the default handler map. @classmethod def get_default_handler_map(cls): """ A class method that returns the default handler map. """ return cls._default_handler_map @classmethod def set_default_handler_map(cls, handler_map): """ A class method to set the default handler map. """ cls._default_handler_map = handler_map @classmethod def update_default_handler_map(cls, handler_map): """ A class method to update the default handler map. """ cls._default_handler_map.update(handler_map) def get_legend_handler_map(self): """ Return the handler map. """ default_handler_map = self.get_default_handler_map() if self._custom_handler_map: hm = default_handler_map.copy() hm.update(self._custom_handler_map) return hm else: return default_handler_map @staticmethod def get_legend_handler(legend_handler_map, orig_handle): """ Return a legend handler from *legend_handler_map* that corresponds to *orig_handler*. *legend_handler_map* should be a dictionary object (that is returned by the get_legend_handler_map method). It first checks if the *orig_handle* itself is a key in the *legend_handler_map* and return the associated value. Otherwise, it checks for each of the classes in its method-resolution-order. If no matching key is found, it returns ``None``. """ if is_hashable(orig_handle): try: return legend_handler_map[orig_handle] except KeyError: pass for handle_type in type(orig_handle).mro(): try: return legend_handler_map[handle_type] except KeyError: pass return None def _init_legend_box(self, handles, labels, markerfirst=True): """ Initialize the legend_box. The legend_box is an instance of the OffsetBox, which is packed with legend handles and texts. Once packed, their location is calculated during the drawing time. """ fontsize = self._fontsize # legend_box is a HPacker, horizontally packed with # columns. Each column is a VPacker, vertically packed with # legend items. Each legend item is HPacker packed with # legend handleBox and labelBox. handleBox is an instance of # offsetbox.DrawingArea which contains legend handle. labelBox # is an instance of offsetbox.TextArea which contains legend # text. text_list = [] # the list of text instances handle_list = [] # the list of text instances handles_and_labels = [] label_prop = dict(verticalalignment='baseline', horizontalalignment='left', fontproperties=self.prop, ) # The approximate height and descent of text. These values are # only used for plotting the legend handle. descent = 0.35 * self._approx_text_height() * (self.handleheight - 0.7) # 0.35 and 0.7 are just heuristic numbers and may need to be improved. height = self._approx_text_height() * self.handleheight - descent # each handle needs to be drawn inside a box of (x, y, w, h) = # (0, -descent, width, height). And their coordinates should # be given in the display coordinates. # The transformation of each handle will be automatically set # to self.get_transform(). If the artist does not use its # default transform (e.g., Collections), you need to # manually set their transform to the self.get_transform(). legend_handler_map = self.get_legend_handler_map() for orig_handle, lab in zip(handles, labels): handler = self.get_legend_handler(legend_handler_map, orig_handle) if handler is None: cbook._warn_external( "Legend does not support {!r} instances.\nA proxy artist " "may be used instead.\nSee: " "http://matplotlib.org/users/legend_guide.html" "#creating-artists-specifically-for-adding-to-the-legend-" "aka-proxy-artists".format(orig_handle)) # We don't have a handle for this artist, so we just defer # to None. handle_list.append(None) else: textbox = TextArea(lab, textprops=label_prop, multilinebaseline=True, minimumdescent=True) handlebox = DrawingArea(width=self.handlelength * fontsize, height=height, xdescent=0., ydescent=descent) text_list.append(textbox._text) # Create the artist for the legend which represents the # original artist/handle. handle_list.append(handler.legend_artist(self, orig_handle, fontsize, handlebox)) handles_and_labels.append((handlebox, textbox)) if handles_and_labels: # We calculate number of rows in each column. The first # (num_largecol) columns will have (nrows+1) rows, and remaining # (num_smallcol) columns will have (nrows) rows. ncol = min(self._ncol, len(handles_and_labels)) nrows, num_largecol = divmod(len(handles_and_labels), ncol) num_smallcol = ncol - num_largecol # starting index of each column and number of rows in it. rows_per_col = [nrows + 1] * num_largecol + [nrows] * num_smallcol start_idxs = np.concatenate([[0], np.cumsum(rows_per_col)[:-1]]) cols = zip(start_idxs, rows_per_col) else: cols = [] columnbox = [] for i0, di in cols: # pack handleBox and labelBox into itemBox itemBoxes = [HPacker(pad=0, sep=self.handletextpad * fontsize, children=[h, t] if markerfirst else [t, h], align="baseline") for h, t in handles_and_labels[i0:i0 + di]] # minimumdescent=False for the text of the last row of the column if markerfirst: itemBoxes[-1].get_children()[1].set_minimumdescent(False) else: itemBoxes[-1].get_children()[0].set_minimumdescent(False) # pack columnBox alignment = "baseline" if markerfirst else "right" columnbox.append(VPacker(pad=0, sep=self.labelspacing * fontsize, align=alignment, children=itemBoxes)) mode = "expand" if self._mode == "expand" else "fixed" sep = self.columnspacing * fontsize self._legend_handle_box = HPacker(pad=0, sep=sep, align="baseline", mode=mode, children=columnbox) self._legend_title_box = TextArea("") self._legend_box = VPacker(pad=self.borderpad * fontsize, sep=self.labelspacing * fontsize, align="center", children=[self._legend_title_box, self._legend_handle_box]) self._legend_box.set_figure(self.figure) self.texts = text_list self.legendHandles = handle_list def _auto_legend_data(self): """ Returns list of vertices and extents covered by the plot. Returns a two long list. First element is a list of (x, y) vertices (in display-coordinates) covered by all the lines and line collections, in the legend's handles. Second element is a list of bounding boxes for all the patches in the legend's handles. """ # should always hold because function is only called internally assert self.isaxes ax = self.parent bboxes = [] lines = [] offsets = [] for handle in ax.lines: assert isinstance(handle, Line2D) path = handle.get_path() trans = handle.get_transform() tpath = trans.transform_path(path) lines.append(tpath) for handle in ax.patches: assert isinstance(handle, Patch) if isinstance(handle, Rectangle): transform = handle.get_data_transform() bboxes.append(handle.get_bbox().transformed(transform)) else: transform = handle.get_transform() bboxes.append(handle.get_path().get_extents(transform)) for handle in ax.collections: transform, transOffset, hoffsets, paths = handle._prepare_points() if len(hoffsets): for offset in transOffset.transform(hoffsets): offsets.append(offset) try: vertices = np.concatenate([l.vertices for l in lines]) except ValueError: vertices = np.array([]) return [vertices, bboxes, lines, offsets] def draw_frame(self, b): ''' Set draw frame to b. Parameters ---------- b : bool ''' self.set_frame_on(b) def get_children(self): 'Return a list of child artists.' children = [] if self._legend_box: children.append(self._legend_box) children.append(self.get_frame()) return children def get_frame(self): ''' Return the `~.patches.Rectangle` instances used to frame the legend. ''' return self.legendPatch def get_lines(self): 'Return a list of `~.lines.Line2D` instances in the legend.' return [h for h in self.legendHandles if isinstance(h, Line2D)] def get_patches(self): 'Return a list of `~.patches.Patch` instances in the legend.' return silent_list('Patch', [h for h in self.legendHandles if isinstance(h, Patch)]) def get_texts(self): 'Return a list of `~.text.Text` instances in the legend.' return silent_list('Text', self.texts) def set_title(self, title, prop=None): """ Set the legend title. Fontproperties can be optionally set with *prop* parameter. """ self._legend_title_box._text.set_text(title) if title: self._legend_title_box._text.set_visible(True) self._legend_title_box.set_visible(True) else: self._legend_title_box._text.set_visible(False) self._legend_title_box.set_visible(False) if prop is not None: if isinstance(prop, dict): prop = FontProperties(**prop) self._legend_title_box._text.set_fontproperties(prop) self.stale = True def get_title(self): 'Return the `.Text` instance for the legend title.' return self._legend_title_box._text def get_window_extent(self, renderer=None): 'Return extent of the legend.' if renderer is None: renderer = self.figure._cachedRenderer return self._legend_box.get_window_extent(renderer=renderer) def get_tightbbox(self, renderer): """ Like `.Legend.get_window_extent`, but uses the box for the legend. Parameters ---------- renderer : `.RendererBase` instance renderer that will be used to draw the figures (i.e. ``fig.canvas.get_renderer()``) Returns ------- `.BboxBase` : containing the bounding box in figure pixel co-ordinates. """ return self._legend_box.get_window_extent(renderer) def get_frame_on(self): """Get whether the legend box patch is drawn.""" return self._drawFrame def set_frame_on(self, b): """ Set whether the legend box patch is drawn. Parameters ---------- b : bool """ self._drawFrame = b self.stale = True def get_bbox_to_anchor(self): """Return the bbox that the legend will be anchored to.""" if self._bbox_to_anchor is None: return self.parent.bbox else: return self._bbox_to_anchor def set_bbox_to_anchor(self, bbox, transform=None): """ Set the bbox that the legend will be anchored to. *bbox* can be - A `.BboxBase` instance - A tuple of ``(left, bottom, width, height)`` in the given transform (normalized axes coordinate if None) - A tuple of ``(left, bottom)`` where the width and height will be assumed to be zero. """ if bbox is None: self._bbox_to_anchor = None return elif isinstance(bbox, BboxBase): self._bbox_to_anchor = bbox else: try: l = len(bbox) except TypeError: raise ValueError("Invalid argument for bbox : %s" % str(bbox)) if l == 2: bbox = [bbox[0], bbox[1], 0, 0] self._bbox_to_anchor = Bbox.from_bounds(*bbox) if transform is None: transform = BboxTransformTo(self.parent.bbox) self._bbox_to_anchor = TransformedBbox(self._bbox_to_anchor, transform) self.stale = True def _get_anchored_bbox(self, loc, bbox, parentbbox, renderer): """ Place the *bbox* inside the *parentbbox* according to a given location code. Return the (x,y) coordinate of the bbox. - loc: a location code in range(1, 11). This corresponds to the possible values for self._loc, excluding "best". - bbox: bbox to be placed, display coordinate units. - parentbbox: a parent box which will contain the bbox. In display coordinates. """ assert loc in range(1, 11) # called only internally BEST, UR, UL, LL, LR, R, CL, CR, LC, UC, C = range(11) anchor_coefs = {UR: "NE", UL: "NW", LL: "SW", LR: "SE", R: "E", CL: "W", CR: "E", LC: "S", UC: "N", C: "C"} c = anchor_coefs[loc] fontsize = renderer.points_to_pixels(self._fontsize) container = parentbbox.padded(-(self.borderaxespad) * fontsize) anchored_box = bbox.anchored(c, container=container) return anchored_box.x0, anchored_box.y0 def _find_best_position(self, width, height, renderer, consider=None): """ Determine the best location to place the legend. *consider* is a list of ``(x, y)`` pairs to consider as a potential lower-left corner of the legend. All are display coords. """ # should always hold because function is only called internally assert self.isaxes verts, bboxes, lines, offsets = self._auto_legend_data() if self._loc_used_default and verts.shape[0] > 200000: # this size results in a 3+ second render time on a good machine cbook._warn_external( 'Creating legend with loc="best" can be slow with large ' 'amounts of data.' ) bbox = Bbox.from_bounds(0, 0, width, height) if consider is None: consider = [self._get_anchored_bbox(x, bbox, self.get_bbox_to_anchor(), renderer) for x in range(1, len(self.codes))] candidates = [] for idx, (l, b) in enumerate(consider): legendBox = Bbox.from_bounds(l, b, width, height) badness = 0 # XXX TODO: If markers are present, it would be good to # take them into account when checking vertex overlaps in # the next line. badness = (legendBox.count_contains(verts) + legendBox.count_contains(offsets) + legendBox.count_overlaps(bboxes) + sum(line.intersects_bbox(legendBox, filled=False) for line in lines)) if badness == 0: return l, b # Include the index to favor lower codes in case of a tie. candidates.append((badness, idx, (l, b))) _, _, (l, b) = min(candidates) return l, b def contains(self, event): return self.legendPatch.contains(event) def set_draggable(self, state, use_blit=False, update='loc'): """ Enable or disable mouse dragging support of the legend. Parameters ---------- state : bool Whether mouse dragging is enabled. use_blit : bool, optional Use blitting for faster image composition. For details see :ref:`func-animation`. update : {'loc', 'bbox'}, optional The legend parameter to be changed when dragged: - 'loc': update the *loc* parameter of the legend - 'bbox': update the *bbox_to_anchor* parameter of the legend Returns ------- If *state* is ``True`` this returns the `~.DraggableLegend` helper instance. Otherwise this returns ``None``. """ if state: if self._draggable is None: self._draggable = DraggableLegend(self, use_blit, update=update) else: if self._draggable is not None: self._draggable.disconnect() self._draggable = None return self._draggable def get_draggable(self): """Return ``True`` if the legend is draggable, ``False`` otherwise.""" return self._draggable is not None
class AnnotationBbox(martist.Artist, _AnnotationBase): """ Annotation-like class, but with offsetbox instead of Text. """ zorder = 3 def __str__(self): return "AnnotationBbox(%g,%g)"%(self.xy[0],self.xy[1]) @docstring.dedent_interpd def __init__(self, offsetbox, xy, xybox=None, xycoords='data', boxcoords=None, frameon=True, pad=0.4, # BboxPatch annotation_clip=None, box_alignment=(0.5, 0.5), bboxprops=None, arrowprops=None, fontsize=None, **kwargs): """ *offsetbox* : OffsetBox instance *xycoords* : same as Annotation but can be a tuple of two strings which are interpreted as x and y coordinates. *boxcoords* : similar to textcoords as Annotation but can be a tuple of two strings which are interpreted as x and y coordinates. *box_alignment* : a tuple of two floats for a vertical and horizontal alignment of the offset box w.r.t. the *boxcoords*. The lower-left corner is (0.0) and upper-right corner is (1.1). other parameters are identical to that of Annotation. """ self.offsetbox = offsetbox self.arrowprops = arrowprops self.set_fontsize(fontsize) if arrowprops is not None: self._arrow_relpos = self.arrowprops.pop("relpos", (0.5, 0.5)) self.arrow_patch = FancyArrowPatch((0, 0), (1,1), **self.arrowprops) else: self._arrow_relpos = None self.arrow_patch = None _AnnotationBase.__init__(self, xy, xytext=xybox, xycoords=xycoords, textcoords=boxcoords, annotation_clip=annotation_clip) martist.Artist.__init__(self, **kwargs) #self._fw, self._fh = 0., 0. # for alignment self._box_alignment = box_alignment # frame self.patch = FancyBboxPatch( xy=(0.0, 0.0), width=1., height=1., facecolor='w', edgecolor='k', mutation_scale=self.prop.get_size_in_points(), snap=True ) self.patch.set_boxstyle("square",pad=pad) if bboxprops: self.patch.set(**bboxprops) self._drawFrame = frameon def contains(self,event): t,tinfo = self.offsetbox.contains(event) #if self.arrow_patch is not None: # a,ainfo=self.arrow_patch.contains(event) # t = t or a # self.arrow_patch is currently not checked as this can be a line - JJ return t,tinfo def get_children(self): children = [self.offsetbox, self.patch] if self.arrow_patch: children.append(self.arrow_patch) return children def set_figure(self, fig): if self.arrow_patch is not None: self.arrow_patch.set_figure(fig) self.offsetbox.set_figure(fig) martist.Artist.set_figure(self, fig) def set_fontsize(self, s=None): """ set fontsize in points """ if s is None: s = rcParams["legend.fontsize"] self.prop=FontProperties(size=s) def get_fontsize(self, s=None): """ return fontsize in points """ return self.prop.get_size_in_points() def update_positions(self, renderer): "Update the pixel positions of the annotated point and the text." xy_pixel = self._get_position_xy(renderer) self._update_position_xybox(renderer, xy_pixel) mutation_scale = renderer.points_to_pixels(self.get_fontsize()) self.patch.set_mutation_scale(mutation_scale) if self.arrow_patch: self.arrow_patch.set_mutation_scale(mutation_scale) def _update_position_xybox(self, renderer, xy_pixel): "Update the pixel positions of the annotation text and the arrow patch." x, y = self.xytext if isinstance(self.textcoords, tuple): xcoord, ycoord = self.textcoords x1, y1 = self._get_xy(renderer, x, y, xcoord) x2, y2 = self._get_xy(renderer, x, y, ycoord) ox0, oy0 = x1, y2 else: ox0, oy0 = self._get_xy(renderer, x, y, self.textcoords) w, h, xd, yd = self.offsetbox.get_extent(renderer) _fw, _fh = self._box_alignment self.offsetbox.set_offset((ox0-_fw*w+xd, oy0-_fh*h+yd)) # update patch position bbox = self.offsetbox.get_window_extent(renderer) #self.offsetbox.set_offset((ox0-_fw*w, oy0-_fh*h)) self.patch.set_bounds(bbox.x0, bbox.y0, bbox.width, bbox.height) x, y = xy_pixel ox1, oy1 = x, y if self.arrowprops: x0, y0 = x, y d = self.arrowprops.copy() # Use FancyArrowPatch if self.arrowprops has "arrowstyle" key. # adjust the starting point of the arrow relative to # the textbox. # TODO : Rotation needs to be accounted. relpos = self._arrow_relpos ox0 = bbox.x0 + bbox.width * relpos[0] oy0 = bbox.y0 + bbox.height * relpos[1] # The arrow will be drawn from (ox0, oy0) to (ox1, # oy1). It will be first clipped by patchA and patchB. # Then it will be shrinked by shirnkA and shrinkB # (in points). If patch A is not set, self.bbox_patch # is used. self.arrow_patch.set_positions((ox0, oy0), (ox1,oy1)) fs = self.prop.get_size_in_points() mutation_scale = d.pop("mutation_scale", fs) mutation_scale = renderer.points_to_pixels(mutation_scale) self.arrow_patch.set_mutation_scale(mutation_scale) patchA = d.pop("patchA", self.patch) self.arrow_patch.set_patchA(patchA) def draw(self, renderer): """ Draw the :class:`Annotation` object to the given *renderer*. """ if renderer is not None: self._renderer = renderer if not self.get_visible(): return xy_pixel = self._get_position_xy(renderer) if not self._check_xy(renderer, xy_pixel): return self.update_positions(renderer) if self.arrow_patch is not None: if self.arrow_patch.figure is None and self.figure is not None: self.arrow_patch.figure = self.figure self.arrow_patch.draw(renderer) if self._drawFrame: self.patch.draw(renderer) self.offsetbox.draw(renderer)
class AnchoredOffsetbox(OffsetBox): """ An offset box placed according to the legend location loc. AnchoredOffsetbox has a single child. When multiple children is needed, use other OffsetBox class to enlose them. By default, the offset box is anchored against its parent axes. You may explicitly specify the bbox_to_anchor. """ zorder = 5 # zorder of the legend def __init__(self, loc, pad=0.4, borderpad=0.5, child=None, prop=None, frameon=True, bbox_to_anchor=None, bbox_transform=None, **kwargs): """ loc is a string or an integer specifying the legend location. The valid location codes are:: 'upper right' : 1, 'upper left' : 2, 'lower left' : 3, 'lower right' : 4, 'right' : 5, 'center left' : 6, 'center right' : 7, 'lower center' : 8, 'upper center' : 9, 'center' : 10, pad : pad around the child for drawing a frame. given in fraction of fontsize. borderpad : pad between offsetbox frame and the bbox_to_anchor, child : OffsetBox instance that will be anchored. prop : font property. This is only used as a reference for paddings. frameon : draw a frame box if True. bbox_to_anchor : bbox to anchor. Use self.axes.bbox if None. bbox_transform : with which the bbox_to_anchor will be transformed. """ super(AnchoredOffsetbox, self).__init__(**kwargs) self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform) self.set_child(child) self.loc = loc self.borderpad=borderpad self.pad = pad if prop is None: self.prop=FontProperties(size=rcParams["legend.fontsize"]) elif isinstance(prop, dict): self.prop=FontProperties(**prop) if "size" not in prop: self.prop.set_size(rcParams["legend.fontsize"]) else: self.prop = prop self.patch = FancyBboxPatch( xy=(0.0, 0.0), width=1., height=1., facecolor='w', edgecolor='k', mutation_scale=self.prop.get_size_in_points(), snap=True ) self.patch.set_boxstyle("square",pad=0) self._drawFrame = frameon def set_child(self, child): "set the child to be anchored" self._child = child def get_child(self): "return the child" return self._child def get_children(self): "return the list of children" return [self._child] def get_extent(self, renderer): """ return the extent of the artist. The extent of the child added with the pad is returned """ w, h, xd, yd = self.get_child().get_extent(renderer) fontsize = renderer.points_to_pixels(self.prop.get_size_in_points()) pad = self.pad * fontsize return w+2*pad, h+2*pad, xd+pad, yd+pad def get_bbox_to_anchor(self): """ return the bbox that the legend will be anchored """ if self._bbox_to_anchor is None: return self.axes.bbox else: transform = self._bbox_to_anchor_transform if transform is None: return self._bbox_to_anchor else: return TransformedBbox(self._bbox_to_anchor, transform) def set_bbox_to_anchor(self, bbox, transform=None): """ set the bbox that the child will be anchored. *bbox* can be a Bbox instance, a list of [left, bottom, width, height], or a list of [left, bottom] where the width and height will be assumed to be zero. The bbox will be transformed to display coordinate by the given transform. """ if bbox is None or isinstance(bbox, BboxBase): self._bbox_to_anchor = bbox else: try: l = len(bbox) except TypeError: raise ValueError("Invalid argument for bbox : %s" % str(bbox)) if l == 2: bbox = [bbox[0], bbox[1], 0, 0] self._bbox_to_anchor = Bbox.from_bounds(*bbox) self._bbox_to_anchor_transform = transform def get_window_extent(self, renderer): ''' get the bounding box in display space. ''' self._update_offset_func(renderer) w, h, xd, yd = self.get_extent(renderer) ox, oy = self.get_offset(w, h, xd, yd, renderer) return Bbox.from_bounds(ox-xd, oy-yd, w, h) def _update_offset_func(self, renderer, fontsize=None): """ Update the offset func which depends on the dpi of the renderer (because of the padding). """ if fontsize is None: fontsize = renderer.points_to_pixels(self.prop.get_size_in_points()) def _offset(w, h, xd, yd, renderer, fontsize=fontsize, self=self): bbox = Bbox.from_bounds(0, 0, w, h) borderpad = self.borderpad*fontsize bbox_to_anchor = self.get_bbox_to_anchor() x0, y0 = self._get_anchored_bbox(self.loc, bbox, bbox_to_anchor, borderpad) return x0+xd, y0+yd self.set_offset(_offset) def update_frame(self, bbox, fontsize=None): self.patch.set_bounds(bbox.x0, bbox.y0, bbox.width, bbox.height) if fontsize: self.patch.set_mutation_scale(fontsize) def draw(self, renderer): "draw the artist" if not self.get_visible(): return fontsize = renderer.points_to_pixels(self.prop.get_size_in_points()) self._update_offset_func(renderer, fontsize) if self._drawFrame: # update the location and size of the legend bbox = self.get_window_extent(renderer) self.update_frame(bbox, fontsize) self.patch.draw(renderer) width, height, xdescent, ydescent = self.get_extent(renderer) px, py = self.get_offset(width, height, xdescent, ydescent, renderer) self.get_child().set_offset((px, py)) self.get_child().draw(renderer) def _get_anchored_bbox(self, loc, bbox, parentbbox, borderpad): """ return the position of the bbox anchored at the parentbbox with the loc code, with the borderpad. """ assert loc in range(1,11) # called only internally BEST, UR, UL, LL, LR, R, CL, CR, LC, UC, C = range(11) anchor_coefs={UR:"NE", UL:"NW", LL:"SW", LR:"SE", R:"E", CL:"W", CR:"E", LC:"S", UC:"N", C:"C"} c = anchor_coefs[loc] container = parentbbox.padded(-borderpad) anchored_box = bbox.anchored(c, container=container) return anchored_box.x0, anchored_box.y0
def __init__(self, parent, handles, labels, loc = None, numpoints = None, # the number of points in the legend line markerscale = None, # the relative size of legend markers vs. original scatterpoints = 3, # TODO: may be an rcParam scatteryoffsets=None, prop = None, # properties for the legend texts # the following dimensions are in axes coords pad = None, # deprecated; use borderpad labelsep = None, # deprecated; use labelspacing handlelen = None, # deprecated; use handlelength handletextsep = None, # deprecated; use handletextpad axespad = None, # deprecated; use borderaxespad # spacing & pad defined as a fraction of the font-size borderpad = None, # the whitespace inside the legend border labelspacing=None, #the vertical space between the legend entries handlelength=None, # the length of the legend handles handleheight=None, # the height of the legend handles handletextpad=None, # the pad between the legend handle and text borderaxespad=None, # the pad between the axes and legend border columnspacing=None, # spacing between columns ncol=1, # number of columns mode=None, # mode for horizontal distribution of columns. None, "expand" fancybox=None, # True use a fancy box, false use a rounded box, none use rc shadow = None, title = None, # set a title for the legend bbox_to_anchor = None, # bbox that the legend will be anchored. bbox_transform = None, # transform for the bbox frameon = None, # draw frame handler_map = None, ): """ - *parent* : the artist that contains the legend - *handles* : a list of artists (lines, patches) to be added to the legend - *labels* : a list of strings to label the legend Optional keyword arguments: ================ ================================================================== Keyword Description ================ ================================================================== loc a location code prop the font property markerscale the relative size of legend markers vs. original numpoints the number of points in the legend for line scatterpoints the number of points in the legend for scatter plot scatteryoffsets a list of yoffsets for scatter symbols in legend frameon if True, draw a frame around the legend. If None, use rc fancybox if True, draw a frame with a round fancybox. If None, use rc shadow if True, draw a shadow behind legend ncol number of columns borderpad the fractional whitespace inside the legend border labelspacing the vertical space between the legend entries handlelength the length of the legend handles handleheight the length of the legend handles handletextpad the pad between the legend handle and text borderaxespad the pad between the axes and legend border columnspacing the spacing between columns title the legend title bbox_to_anchor the bbox that the legend will be anchored. bbox_transform the transform for the bbox. transAxes if None. ================ ================================================================== The pad and spacing parameters are measured in font-size units. E.g., a fontsize of 10 points and a handlelength=5 implies a handlelength of 50 points. Values from rcParams will be used if None. Users can specify any arbitrary location for the legend using the *bbox_to_anchor* keyword argument. bbox_to_anchor can be an instance of BboxBase(or its derivatives) or a tuple of 2 or 4 floats. See :meth:`set_bbox_to_anchor` for more detail. The legend location can be specified by setting *loc* with a tuple of 2 floats, which is interpreted as the lower-left corner of the legend in the normalized axes coordinate. """ from matplotlib.axes import Axes # local import only to avoid circularity from matplotlib.figure import Figure # local import only to avoid circularity Artist.__init__(self) if prop is None: self.prop=FontProperties(size=rcParams["legend.fontsize"]) elif isinstance(prop, dict): self.prop=FontProperties(**prop) if "size" not in prop: self.prop.set_size(rcParams["legend.fontsize"]) else: self.prop=prop self._fontsize = self.prop.get_size_in_points() propnames=['numpoints', 'markerscale', 'shadow', "columnspacing", "scatterpoints", "handleheight"] self.texts = [] self.legendHandles = [] self._legend_title_box = None self._handler_map = handler_map localdict = locals() for name in propnames: if localdict[name] is None: value = rcParams["legend."+name] else: value = localdict[name] setattr(self, name, value) # Take care the deprecated keywords deprecated_kwds = {"pad":"borderpad", "labelsep":"labelspacing", "handlelen":"handlelength", "handletextsep":"handletextpad", "axespad":"borderaxespad"} # convert values of deprecated keywords (ginve in axes coords) # to new vaules in a fraction of the font size # conversion factor bbox = parent.bbox axessize_fontsize = min(bbox.width, bbox.height)/self._fontsize for k, v in deprecated_kwds.iteritems(): # use deprecated value if not None and if their newer # counter part is None. if localdict[k] is not None and localdict[v] is None: warnings.warn("Use '%s' instead of '%s'." % (v, k), DeprecationWarning) setattr(self, v, localdict[k]*axessize_fontsize) continue # Otherwise, use new keywords if localdict[v] is None: setattr(self, v, rcParams["legend."+v]) else: setattr(self, v, localdict[v]) del localdict handles = list(handles) if len(handles)<2: ncol = 1 self._ncol = ncol if self.numpoints <= 0: raise ValueError("numpoints must be > 0; it was %d"% numpoints) # introduce y-offset for handles of the scatter plot if scatteryoffsets is None: self._scatteryoffsets = np.array([3./8., 4./8., 2.5/8.]) else: self._scatteryoffsets = np.asarray(scatteryoffsets) reps = int(self.scatterpoints / len(self._scatteryoffsets)) + 1 self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[:self.scatterpoints] # _legend_box is an OffsetBox instance that contains all # legend items and will be initialized from _init_legend_box() # method. self._legend_box = None if isinstance(parent,Axes): self.isaxes = True self.set_axes(parent) self.set_figure(parent.figure) elif isinstance(parent,Figure): self.isaxes = False self.set_figure(parent) else: raise TypeError("Legend needs either Axes or Figure as parent") self.parent = parent if loc is None: loc = rcParams["legend.loc"] if not self.isaxes and loc in [0,'best']: loc = 'upper right' if is_string_like(loc): if loc not in self.codes: if self.isaxes: warnings.warn('Unrecognized location "%s". Falling back on "best"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes.iterkeys()))) loc = 0 else: warnings.warn('Unrecognized location "%s". Falling back on "upper right"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes.iterkeys()))) loc = 1 else: loc = self.codes[loc] if not self.isaxes and loc == 0: warnings.warn('Automatic legend placement (loc="best") not implemented for figure legend. ' 'Falling back on "upper right".') loc = 1 self._mode = mode self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform) # We use FancyBboxPatch to draw a legend frame. The location # and size of the box will be updated during the drawing time. self.legendPatch = FancyBboxPatch( xy=(0.0, 0.0), width=1., height=1., facecolor=rcParams["axes.facecolor"], edgecolor=rcParams["axes.edgecolor"], mutation_scale=self._fontsize, snap=True ) # The width and height of the legendPatch will be set (in the # draw()) to the length that includes the padding. Thus we set # pad=0 here. if fancybox is None: fancybox = rcParams["legend.fancybox"] if fancybox == True: self.legendPatch.set_boxstyle("round",pad=0, rounding_size=0.2) else: self.legendPatch.set_boxstyle("square",pad=0) self._set_artist_props(self.legendPatch) self._drawFrame = frameon if frameon is None: self._drawFrame = rcParams["legend.frameon"] # init with null renderer self._init_legend_box(handles, labels) self._loc = loc self.set_title(title) self._last_fontsize_points = self._fontsize self._draggable = None
def __init__( self, parent, handles, labels, loc=None, numpoints=None, # the number of points in the legend line markerscale=None, # the relative size of legend markers vs. original scatterpoints=3, # TODO: may be an rcParam scatteryoffsets=None, prop=None, # properties for the legend texts # the following dimensions are in axes coords pad=None, # deprecated; use borderpad labelsep=None, # deprecated; use labelspacing handlelen=None, # deprecated; use handlelength handletextsep=None, # deprecated; use handletextpad axespad=None, # deprecated; use borderaxespad # spacing & pad defined as a fraction of the font-size borderpad=None, # the whitespace inside the legend border labelspacing=None, #the vertical space between the legend entries handlelength=None, # the length of the legend handles handletextpad=None, # the pad between the legend handle and text borderaxespad=None, # the pad between the axes and legend border columnspacing=None, # spacing between columns ncol=1, # number of columns mode=None, # mode for horizontal distribution of columns. None, "expand" fancybox=None, # True use a fancy box, false use a rounded box, none use rc shadow=None, title=None, # set a title for the legend bbox_to_anchor=None, # bbox that the legend will be anchored. bbox_transform=None, # transform for the bbox frameon=True, # draw frame handler_map=None, ): """ - *parent* : the artist that contains the legend - *handles* : a list of artists (lines, patches) to be added to the legend - *labels* : a list of strings to label the legend Optional keyword arguments: ================ ================================================================== Keyword Description ================ ================================================================== loc a location code prop the font property markerscale the relative size of legend markers vs. original numpoints the number of points in the legend for line scatterpoints the number of points in the legend for scatter plot scatteryoffsets a list of yoffsets for scatter symbols in legend frameon if True, draw a frame (default is True) fancybox if True, draw a frame with a round fancybox. If None, use rc shadow if True, draw a shadow behind legend ncol number of columns borderpad the fractional whitespace inside the legend border labelspacing the vertical space between the legend entries handlelength the length of the legend handles handletextpad the pad between the legend handle and text borderaxespad the pad between the axes and legend border columnspacing the spacing between columns title the legend title bbox_to_anchor the bbox that the legend will be anchored. bbox_transform the transform for the bbox. transAxes if None. ================ ================================================================== The pad and spacing parameters are measured in font-size units. E.g., a fontsize of 10 points and a handlelength=5 implies a handlelength of 50 points. Values from rcParams will be used if None. Users can specify any arbitrary location for the legend using the *bbox_to_anchor* keyword argument. bbox_to_anchor can be an instance of BboxBase(or its derivatives) or a tuple of 2 or 4 floats. See :meth:`set_bbox_to_anchor` for more detail. The legend location can be specified by setting *loc* with a tuple of 2 floats, which is interpreted as the lower-left corner of the legend in the normalized axes coordinate. """ from matplotlib.axes import Axes # local import only to avoid circularity from matplotlib.figure import Figure # local import only to avoid circularity Artist.__init__(self) if prop is None: self.prop = FontProperties(size=rcParams["legend.fontsize"]) elif isinstance(prop, dict): self.prop = FontProperties(**prop) if "size" not in prop: self.prop.set_size(rcParams["legend.fontsize"]) else: self.prop = prop self._fontsize = self.prop.get_size_in_points() propnames = [ 'numpoints', 'markerscale', 'shadow', "columnspacing", "scatterpoints" ] self.texts = [] self.legendHandles = [] self._legend_title_box = None self._handler_map = handler_map localdict = locals() for name in propnames: if localdict[name] is None: value = rcParams["legend." + name] else: value = localdict[name] setattr(self, name, value) # Take care the deprecated keywords deprecated_kwds = { "pad": "borderpad", "labelsep": "labelspacing", "handlelen": "handlelength", "handletextsep": "handletextpad", "axespad": "borderaxespad" } # convert values of deprecated keywords (ginve in axes coords) # to new vaules in a fraction of the font size # conversion factor bbox = parent.bbox axessize_fontsize = min(bbox.width, bbox.height) / self._fontsize for k, v in deprecated_kwds.items(): # use deprecated value if not None and if their newer # counter part is None. if localdict[k] is not None and localdict[v] is None: warnings.warn("Use '%s' instead of '%s'." % (v, k), DeprecationWarning) setattr(self, v, localdict[k] * axessize_fontsize) continue # Otherwise, use new keywords if localdict[v] is None: setattr(self, v, rcParams["legend." + v]) else: setattr(self, v, localdict[v]) del localdict handles = list(handles) if len(handles) < 2: ncol = 1 self._ncol = ncol if self.numpoints <= 0: raise ValueError("numpoints must be >= 0; it was %d" % numpoints) # introduce y-offset for handles of the scatter plot if scatteryoffsets is None: self._scatteryoffsets = np.array([3. / 8., 4. / 8., 2.5 / 8.]) else: self._scatteryoffsets = np.asarray(scatteryoffsets) reps = int(self.scatterpoints / len(self._scatteryoffsets)) + 1 self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[:self.scatterpoints] # _legend_box is an OffsetBox instance that contains all # legend items and will be initialized from _init_legend_box() # method. self._legend_box = None if isinstance(parent, Axes): self.isaxes = True self.set_axes(parent) self.set_figure(parent.figure) elif isinstance(parent, Figure): self.isaxes = False self.set_figure(parent) else: raise TypeError("Legend needs either Axes or Figure as parent") self.parent = parent if loc is None: loc = rcParams["legend.loc"] if not self.isaxes and loc in [0, 'best']: loc = 'upper right' if is_string_like(loc): if loc not in self.codes: if self.isaxes: warnings.warn( 'Unrecognized location "%s". Falling back on "best"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes.keys()))) loc = 0 else: warnings.warn( 'Unrecognized location "%s". Falling back on "upper right"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes.keys()))) loc = 1 else: loc = self.codes[loc] if not self.isaxes and loc == 0: warnings.warn( 'Automatic legend placement (loc="best") not implemented for figure legend. ' 'Falling back on "upper right".') loc = 1 self._mode = mode self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform) # We use FancyBboxPatch to draw a legend frame. The location # and size of the box will be updated during the drawing time. self.legendPatch = FancyBboxPatch(xy=(0.0, 0.0), width=1., height=1., facecolor=rcParams["axes.facecolor"], edgecolor=rcParams["axes.edgecolor"], mutation_scale=self._fontsize, snap=True) # The width and height of the legendPatch will be set (in the # draw()) to the length that includes the padding. Thus we set # pad=0 here. if fancybox is None: fancybox = rcParams["legend.fancybox"] if fancybox == True: self.legendPatch.set_boxstyle("round", pad=0, rounding_size=0.2) else: self.legendPatch.set_boxstyle("square", pad=0) self._set_artist_props(self.legendPatch) self._drawFrame = frameon # init with null renderer self._init_legend_box(handles, labels) self._loc = loc self.set_title(title) self._last_fontsize_points = self._fontsize self._draggable = None
def draw_nodes(nodes, fig, scaling_factor, nw_height_inches): """Create a list of FancyBbox Patches, one for each node. Args: nodes (iterable collection of _Node): collection of nodes fig (matplotlib.figure.Figure): the figure of the network scaling_factor (float): scaling_factor to decrease or increase the figure size nw_height_inches (float): height of the network in inches Returns: list of matplotlib.patches - FancyBboxPatch, Ellipse, or Polygon """ node_patches = [] for node in nodes: if node.shape == "round_box": print(node.lower_left_point) print(node.lower_left_point[0] * INCHES_PER_POINT + WIDTH_SHIFT, node.lower_left_point[1] * INCHES_PER_POINT + HEIGHT_SHIFT) node_patch = FancyBboxPatch( xy=(scaling_factor * node.lower_left_point[0] * INCHES_PER_POINT + WIDTH_SHIFT, scaling_factor * (nw_height_inches - node.lower_left_point[1] * INCHES_PER_POINT - node.height * INCHES_PER_POINT) + HEIGHT_SHIFT), width=scaling_factor * node.width * INCHES_PER_POINT, height=scaling_factor * node.height * INCHES_PER_POINT, edgecolor=node.edge_color, facecolor=node.fill_color, linewidth=node.edge_width, boxstyle=BoxStyle("round", pad=0, rounding_size=node.rectangle_rounding), transform=fig.dpi_scale_trans) elif node.shape == "ellipse": node_patch = Ellipse( (scaling_factor * node.center.x * INCHES_PER_POINT + WIDTH_SHIFT, scaling_factor * (nw_height_inches - node.center.y * INCHES_PER_POINT) + HEIGHT_SHIFT), scaling_factor * node.width * INCHES_PER_POINT, scaling_factor * node.height * INCHES_PER_POINT, edgecolor=node.edge_color, facecolor=node.fill_color, linewidth=node.edge_width, transform=fig.dpi_scale_trans) elif node.shape == "polygon": node_points = _adjust_x_and_y_values(node.polygon_points, scaling_factor, nw_height_inches, node) node_points = np.array(node_points) # need to adjust the y's path = Path(node_points, node.polygon_codes) node_patch = PathPatch(path, facecolor=node.fill_color, edgecolor=node.edge_color, linewidth=node.edge_width, transform=fig.dpi_scale_trans) else: pass node_patches.append(node_patch) return node_patches
class PaddedBox(OffsetBox): def __init__(self, child, pad=None, draw_frame=False, patch_attrs=None): """ *pad* : boundary pad .. note:: *pad* need to given in points and will be scale with the renderer dpi, while *width* and *hight* need to be in pixels. """ super(PaddedBox, self).__init__() self.pad = pad self._children = [child] self.patch = FancyBboxPatch( xy=(0.0, 0.0), width=1., height=1., facecolor='w', edgecolor='k', mutation_scale=1, #self.prop.get_size_in_points(), snap=True ) self.patch.set_boxstyle("square",pad=0) if patch_attrs is not None: self.patch.update(patch_attrs) self._drawFrame = draw_frame def get_extent_offsets(self, renderer): """ update offset of childrens and return the extents of the box """ dpicor = renderer.points_to_pixels(1.) pad = self.pad * dpicor w, h, xd, yd = self._children[0].get_extent(renderer) return w + 2*pad, h + 2*pad, \ xd+pad, yd+pad, \ [(0, 0)] def draw(self, renderer): """ Update the location of children if necessary and draw them to the given *renderer*. """ width, height, xdescent, ydescent, offsets = self.get_extent_offsets(renderer) px, py = self.get_offset(width, height, xdescent, ydescent, renderer) for c, (ox, oy) in zip(self.get_visible_children(), offsets): c.set_offset((px+ox, py+oy)) self.draw_frame(renderer) for c in self.get_visible_children(): c.draw(renderer) #bbox_artist(self, renderer, fill=False, props=dict(pad=0.)) def update_frame(self, bbox, fontsize=None): self.patch.set_bounds(bbox.x0, bbox.y0, bbox.width, bbox.height) if fontsize: self.patch.set_mutation_scale(fontsize) def draw_frame(self, renderer): # update the location and size of the legend bbox = self.get_window_extent(renderer) self.update_frame(bbox) if self._drawFrame: self.patch.draw(renderer)
def __init__(self, # Variables controlling flow of the program save_history=save_history_globals, # Variables used for physical simulation dt=dt_main_simulation_globals, m=m_globals, # mass of pend, kg M=M_globals, # mass of cart, kg L=L_globals, # half length of pend, m u_max=u_max_globals, # max cart force, N M_fric=M_fric_globals, # 1.0, # cart friction, N/m/s J_fric=J_fric_globals, # 10.0, # friction coefficient on angular velocity, Nm/rad/s v_max=v_max_globals, # max DC motor speed, m/s, in absense of friction, used for motor back EMF model controlDisturbance=controlDisturbance_globals, # 0.01, # disturbance, as factor of u_max sensorNoise=sensorNoise_globals, # 0.01, # noise, as factor of max values # Variables for random trace generation N=N_globals, # Complexity of the random trace, number of random points used for interpolation random_length=random_length_globals, # Number of points in the random length trece mode_init = mode_globals # In which mode the Cart should be initialized ): # State of the cart self.s = SimpleNamespace() # s like state self.s.position = 0.0 self.s.positionD = 0.0 self.s.positionDD = 0.0 self.s.angle = 0.0 self.s.angleD = 0.0 self.s.angleDD = 0.0 self.u = 0.0 self.Q = 0.0 self.target_position = 0.0 # Other variables controlling flow or initial state of the program self.Q_thread_enabled = False # If True, control input is computed asynchronously to the simulation in a separate thread self.Q_max = 1.0 self.slider_value = 0.0 self.dt = dt self.time = 0.0 self.dict_history = {} self.reset_dict_history() self.save_history = save_history self.random_trace_generated = False self.use_pregenerated_target_position = False # Physical parameters of the cart self.p = SimpleNamespace() # p like parameters self.p.m = m # mass of pend, kg self.p.M = M # mass of cart, kg self.p.M_fric = M_fric # cart friction, N/m/s self.p.J_fric = J_fric # friction coefficient on angular velocity, Nm/rad/s g = g_globals self.p.g = g_globals self.p.L = L # half length of pend, m k = k_globals self.p.k = k_globals self.p.u_max = u_max # max cart force, N self.p.v_max = v_max # max DC motor speed, m/s, in absence of friction, used for motor back EMF model self.p.controlDisturbance = controlDisturbance # disturbance, as factor of u_max self.p.sensorNoise = sensorNoise # noise, as factor of max values self.p.force_damping = 1.0 # Jacobian of the system linearized around upper equilibrium position # x' = f(x) # x = [x, v, theta, omega] # TODO if parameters change in runtime this Jacobian wont be updated self.Jacobian_UP = array([ [0.0, 1.0, 0.0, 0.0], [0.0, (-(1 + k) * M_fric) / (-m + (1 + k) * (m + M)), (g * m) / (-m + (1 + k) * (m + M)), (-J_fric) / (L * (-m + (1 + k) * (m + M)))], [0.0, 0.0, 0.0, 1.0], [0.0, (-M_fric) / (L * (-m + (1 + k) * (m + M))), (g * (M + m)) / (L * (-m + (1 + k) * (m + M))), -m * (M + m) * J_fric / (L * L * (-m + (1 + k) * (m + M)))], ]) # Array gathering control around equilibrium self.B = u_max * array([ [0.0], [(1 + k) / (-m + (1 + k) * (m + M))], [0.0], [1.0 / (L * (-m + (1 + k) * (m + M)))], ]) # Cost matrices for LQR controller self.Q_matrix = diag([10.0, 1.0, 1.0, 1.0]) # How much to punish x, v, theta, omega self.R_matrix = 1.0e9 # How much to punish Q self.controller_lqr = controller_lqr(self.Jacobian_UP, self.B, self.Q_matrix, self.R_matrix) self.controller_do_mpc = controller_do_mpc() self.controller_do_mpc_discrete = controller_do_mpc_discrete() self.controller = None self.controller_name = '' # Variables for pre-generated random trace self.N = int(N) self.random_length = random_length self.random_track_f = None self.new_track_generated = False self.t_max_pre = None # THE REMAINING PART OF __init__ METHOD CONCERNS DRAWING SETTINGS ONLY # DIMENSIONS OF THE DRAWING ONLY!!! # NOTHING TO DO WITH THE SIMULATION AND NOT INTENDED TO BE MANIPULATED BY USER !!! self.CartLength = 10.0 self.WheelRadius = 0.5 self.WheelToMiddle = 4.0 self.y_plane = 0.0 self.y_wheel = self.y_plane + self.WheelRadius self.MastHight = 10.0 # For drowing only. For calculation see L self.MastThickness = 0.05 self.HalfLength = 50.0 # Length of the track # Elements of the drawing self.Mast = FancyBboxPatch(xy=(self.s.position - (self.MastThickness / 2.0), 1.25 * self.WheelRadius), width=self.MastThickness, height=self.MastHight, fc='g') self.Chassis = FancyBboxPatch((self.s.position - (self.CartLength / 2.0), self.WheelRadius), self.CartLength, 1 * self.WheelRadius, fc='r') self.WheelLeft = Circle((self.s.position - self.WheelToMiddle, self.y_wheel), radius=self.WheelRadius, fc='y', ec='k', lw=5) self.WheelRight = Circle((self.s.position + self.WheelToMiddle, self.y_wheel), radius=self.WheelRadius, fc='y', ec='k', lw=5) self.Slider = Rectangle((0.0, 0.0), self.slider_value, 1.0) self.t2 = transforms.Affine2D().rotate(0.0) # An abstract container for the transform rotating the mast # Set starting mode of operation self.set_mode(new_mode=mode_init)
def draw_compartments(compartments, fig, scaling_factor, nw_height_inches): """Create a list of FancyBbox Patches, one for each compartment. Args: compartments (iterable collection of Compartment): collection of compartments fig (matplotlib.figure.Figure): the figure of the network scaling_factor (float): scaling_factor to decrease or increase the figure size nw_height_inches (float): height of the network in inches Returns: list of matplotlib.patches - FancyBboxPatch, Ellipse, or Polygon """ compartment_patch = None compartment_patches = [] for compartment in compartments: if compartment.shape == "round_box": # todo change the following to reduce padding around species boxes compartment_patch = FancyBboxPatch( xy=(scaling_factor * compartment.lower_left_point[0] * INCHES_PER_POINT + WIDTH_SHIFT, scaling_factor * (nw_height_inches - compartment.lower_left_point[1] * INCHES_PER_POINT - compartment.height * INCHES_PER_POINT) + HEIGHT_SHIFT), width=scaling_factor * compartment.width * INCHES_PER_POINT, height=scaling_factor * compartment.height * INCHES_PER_POINT, edgecolor=compartment.edge_color, facecolor=compartment.fill_color, linewidth=compartment.line_width, boxstyle=BoxStyle("round", pad=0, rounding_size=0.6), transform=fig.dpi_scale_trans) elif compartment.shape == "ellipse": compartment_patch = Ellipse( (scaling_factor * compartment.center_x * INCHES_PER_POINT + WIDTH_SHIFT, scaling_factor * (nw_height_inches - compartment.center_y * INCHES_PER_POINT) + HEIGHT_SHIFT), scaling_factor * compartment.width * INCHES_PER_POINT, scaling_factor * compartment.height * INCHES_PER_POINT, edgecolor=compartment.edge_color, facecolor=compartment.fill_color, linewidth=compartment.line_width, transform=fig.dpi_scale_trans) elif compartment.shape == "polygon": compartment_points = _adjust_x_and_y_values( compartment.polygon_points, scaling_factor, nw_height_inches, compartment) compartment_points = np.array(compartment_points) # need to adjust the y's path = Path(compartment_points, compartment.polygon_codes) compartment_patch = PathPatch(path, facecolor=compartment.fill_color, edgecolor=compartment.edge_color, linewidth=compartment.line_width, transform=fig.dpi_scale_trans) if compartment_patch is None: raise ValueError compartment_patches.append(compartment_patch) return compartment_patches
class Cart: def __init__(self, # Variables controlling flow of the program save_history=save_history_globals, # Variables used for physical simulation dt=dt_main_simulation_globals, m=m_globals, # mass of pend, kg M=M_globals, # mass of cart, kg L=L_globals, # half length of pend, m u_max=u_max_globals, # max cart force, N M_fric=M_fric_globals, # 1.0, # cart friction, N/m/s J_fric=J_fric_globals, # 10.0, # friction coefficient on angular velocity, Nm/rad/s v_max=v_max_globals, # max DC motor speed, m/s, in absense of friction, used for motor back EMF model controlDisturbance=controlDisturbance_globals, # 0.01, # disturbance, as factor of u_max sensorNoise=sensorNoise_globals, # 0.01, # noise, as factor of max values # Variables for random trace generation N=N_globals, # Complexity of the random trace, number of random points used for interpolation random_length=random_length_globals, # Number of points in the random length trece mode_init = mode_globals # In which mode the Cart should be initialized ): # State of the cart self.s = SimpleNamespace() # s like state self.s.position = 0.0 self.s.positionD = 0.0 self.s.positionDD = 0.0 self.s.angle = 0.0 self.s.angleD = 0.0 self.s.angleDD = 0.0 self.u = 0.0 self.Q = 0.0 self.target_position = 0.0 # Other variables controlling flow or initial state of the program self.Q_thread_enabled = False # If True, control input is computed asynchronously to the simulation in a separate thread self.Q_max = 1.0 self.slider_value = 0.0 self.dt = dt self.time = 0.0 self.dict_history = {} self.reset_dict_history() self.save_history = save_history self.random_trace_generated = False self.use_pregenerated_target_position = False # Physical parameters of the cart self.p = SimpleNamespace() # p like parameters self.p.m = m # mass of pend, kg self.p.M = M # mass of cart, kg self.p.M_fric = M_fric # cart friction, N/m/s self.p.J_fric = J_fric # friction coefficient on angular velocity, Nm/rad/s g = g_globals self.p.g = g_globals self.p.L = L # half length of pend, m k = k_globals self.p.k = k_globals self.p.u_max = u_max # max cart force, N self.p.v_max = v_max # max DC motor speed, m/s, in absence of friction, used for motor back EMF model self.p.controlDisturbance = controlDisturbance # disturbance, as factor of u_max self.p.sensorNoise = sensorNoise # noise, as factor of max values self.p.force_damping = 1.0 # Jacobian of the system linearized around upper equilibrium position # x' = f(x) # x = [x, v, theta, omega] # TODO if parameters change in runtime this Jacobian wont be updated self.Jacobian_UP = array([ [0.0, 1.0, 0.0, 0.0], [0.0, (-(1 + k) * M_fric) / (-m + (1 + k) * (m + M)), (g * m) / (-m + (1 + k) * (m + M)), (-J_fric) / (L * (-m + (1 + k) * (m + M)))], [0.0, 0.0, 0.0, 1.0], [0.0, (-M_fric) / (L * (-m + (1 + k) * (m + M))), (g * (M + m)) / (L * (-m + (1 + k) * (m + M))), -m * (M + m) * J_fric / (L * L * (-m + (1 + k) * (m + M)))], ]) # Array gathering control around equilibrium self.B = u_max * array([ [0.0], [(1 + k) / (-m + (1 + k) * (m + M))], [0.0], [1.0 / (L * (-m + (1 + k) * (m + M)))], ]) # Cost matrices for LQR controller self.Q_matrix = diag([10.0, 1.0, 1.0, 1.0]) # How much to punish x, v, theta, omega self.R_matrix = 1.0e9 # How much to punish Q self.controller_lqr = controller_lqr(self.Jacobian_UP, self.B, self.Q_matrix, self.R_matrix) self.controller_do_mpc = controller_do_mpc() self.controller_do_mpc_discrete = controller_do_mpc_discrete() self.controller = None self.controller_name = '' # Variables for pre-generated random trace self.N = int(N) self.random_length = random_length self.random_track_f = None self.new_track_generated = False self.t_max_pre = None # THE REMAINING PART OF __init__ METHOD CONCERNS DRAWING SETTINGS ONLY # DIMENSIONS OF THE DRAWING ONLY!!! # NOTHING TO DO WITH THE SIMULATION AND NOT INTENDED TO BE MANIPULATED BY USER !!! self.CartLength = 10.0 self.WheelRadius = 0.5 self.WheelToMiddle = 4.0 self.y_plane = 0.0 self.y_wheel = self.y_plane + self.WheelRadius self.MastHight = 10.0 # For drowing only. For calculation see L self.MastThickness = 0.05 self.HalfLength = 50.0 # Length of the track # Elements of the drawing self.Mast = FancyBboxPatch(xy=(self.s.position - (self.MastThickness / 2.0), 1.25 * self.WheelRadius), width=self.MastThickness, height=self.MastHight, fc='g') self.Chassis = FancyBboxPatch((self.s.position - (self.CartLength / 2.0), self.WheelRadius), self.CartLength, 1 * self.WheelRadius, fc='r') self.WheelLeft = Circle((self.s.position - self.WheelToMiddle, self.y_wheel), radius=self.WheelRadius, fc='y', ec='k', lw=5) self.WheelRight = Circle((self.s.position + self.WheelToMiddle, self.y_wheel), radius=self.WheelRadius, fc='y', ec='k', lw=5) self.Slider = Rectangle((0.0, 0.0), self.slider_value, 1.0) self.t2 = transforms.Affine2D().rotate(0.0) # An abstract container for the transform rotating the mast # Set starting mode of operation self.set_mode(new_mode=mode_init) def Generate_Random_Trace_Function(self): # t_pre = arange(0, self.random_length)*self.dt self.t_max_pre = self.random_length * self.dt # t_init = linspace(0, self.t_max_pre, num=self.N, endpoint=True) t_init = np.sort(random.uniform(self.dt, self.t_max_pre, self.N)) t_init = np.insert(t_init, 0, 0.0) t_init = np.append(t_init, self.t_max_pre) y = 2.0 * (random.random(self.N) - 0.5) y = y * 0.8 * self.HalfLength / max(abs(y)) y = np.insert(y, 0, 0.0) y = np.append(y, 0.0) # t1 = timeit.default_timer() # self.random_track_f = interp1d(t_init, y, kind='cubic') # t2 = timeit.default_timer() # Try algorithm setting derivative to 0 a each point yder = [[y[i], 0] for i in range(len(y))] self.random_track_f = BPoly.from_derivatives(t_init, yder) # t3 = timeit.default_timer() # print('Time for simple method: {:.3f}ms'.format((t2-t1)*1000)) # print('Time for complex method: {:.3f}ms'.format((t3 - t2) * 1000)) self.new_track_generated = True # Function gathering equations to update the CartPole state: # x' = f(x) # x = [x, v, theta, omega] def Equations_of_motion(self): self.s.position, self.s.positionD, self.s.angle, self.s.angleD = \ cartpole_integration(self.s, self.dt) self.s.angleDD, self.s.positionDD = cartpole_ode(self.p, self.s, self.u) # Determine the dimensionales [-1,1] value of the motor power Q def Update_Q(self): if self.mode == 0: # in this case slider corresponds already to the power of the motor self.Q = self.slider_value else: # in this case slider gives a target position, lqr regulator # mode == 1 -> lqr controller # mode == 2 -> mpc controller # tic = timeit.default_timer() self.Q = self.controller.step(self.s, self.target_position) # toc = timeit.default_timer() # print("Time to find control input = {} ms".format((toc - tic) * 1000.0)) # if self.Q > 1.0: # print('Q to big! ' + str(self.Q)) # elif self.Q < -1.0: # print('Q to small! ' + str(self.Q)) # This method changes the internal state of the CartPole # from a state at time t to a state at t+dt def update_state(self, slider=None, dt=None, save_history=True): # Optionally update slider, mode and dt values if slider: self.slider_value = slider if dt: self.dt = dt self.save_history = save_history if self.use_pregenerated_target_position == True: self.target_position = self.random_track_f(self.time) if self.target_position > 0.8 * self.HalfLength: self.target_position = 0.8 * self.HalfLength elif self.target_position < -0.8 * self.HalfLength: self.target_position = -0.8 * self.HalfLength self.slider_value = self.target_position else: if self.mode == 0: self.target_position = 0.0 else: self.target_position = self.slider_value # Calculate the next state self.Equations_of_motion() # Normalize angle self.s.angle = normalize_angle_rad(self.s.angle) # In case in the next step the wheel of the cart # went beyond the track # Bump elastically into an (invisible) boarder if (abs(self.s.position) + self.WheelToMiddle) > self.HalfLength: self.s.positionD = -self.s.positionD # Determine the dimensionales [-1,1] value of the motor power Q if not self.Q_thread_enabled: self.Update_Q() self.u = Q2u(self.Q, self.p) # Update the total time of the simulation self.time = self.time + self.dt # If user chose to save history of the simulation it is saved now # It is saved first internally to a dictionary in the Cart instance if self.save_history: # Saving simulation data self.dict_history['time'].append(around(self.time, 5)) self.dict_history['dt'].append(around(self.dt * 1000.0, 3)) self.dict_history['s.position'].append(around(self.s.position, 3)) self.dict_history['s.positionD'].append(around(self.s.positionD, 4)) self.dict_history['s.positionDD'].append(around(self.s.positionDD, 4)) self.dict_history['s.angle'].append(around(self.s.angle, 4)) self.dict_history['s.angleD'].append(around(self.s.angleD, 4)) self.dict_history['s.angleDD'].append(around(self.s.angleDD, 4)) self.dict_history['u'].append(around(self.u, 4)) self.dict_history['Q'].append(around(self.Q, 4)) # The target_position is not always meaningful # If it is not meaningful all values in this column are set to 0 self.dict_history['target_position'].append(around(self.target_position, 4)) # Return the state of the CartPole return self.s.position, self.s.positionD, self.s.positionDD, \ self.s.angle, self.s.angleD, self.s.angleDD, \ self.u # This method only returns the state of the CartPole instance def get_state(self): return self.s.position, self.s.positionD, self.s.positionDD, \ self.s.angle, self.s.angleD, self.s.angleDD, \ self.u # This method resets the internal state of the CartPole instance def reset_state(self): self.s.position = 0.0 self.s.positionD = 0.0 self.s.positionDD = 0.0 self.s.angle = (2.0 * random.normal() - 1.0) * pi / 180.0 self.s.angleD = 0.0 self.s.angleDD = 0.0 self.u = 0.0 self.slider_value = 0.0 self.time = 0.0 # This method draws elements and set properties of the CartPole figure # which do not change at every frame of the animation def draw_constant_elements(self, fig, AxCart, AxSlider): # Delete all elements of the Figure AxCart.clear() AxSlider.clear() ## Upper chart with Cart Picture # Set x and y limits AxCart.set_xlim((-self.HalfLength * 1.1, self.HalfLength * 1.1)) AxCart.set_ylim((-1.0, 15.0)) # Remove ticks on the y-axes AxCart.yaxis.set_major_locator(NullLocator()) # Draw track Floor = Rectangle((-self.HalfLength, -1.0), 2 * self.HalfLength, 1.0, fc='brown') AxCart.add_patch(Floor) # Draw an invisible point at constant position # Thanks to it the axes is drawn high enough for the mast InvisiblePointUp = Rectangle((0, self.MastHight + 2.0), self.MastThickness, 0.0001, fc='w', ec='w') AxCart.add_patch(InvisiblePointUp) # Apply scaling AxCart.axis('scaled') ## Lower Chart with Slider # Set y limits AxSlider.set(xlim=(-1.1 * self.slider_max, self.slider_max * 1.1)) # Remove ticks on the y-axes AxSlider.yaxis.set_major_locator(NullLocator()) # Apply scaling AxSlider.set_aspect("auto") return fig, AxCart, AxSlider # This method accepts the mouse position and updated the slider value accordingly # The mouse position has to be captured by a function not included in this class def update_slider(self, mouse_position): # The if statement formulates a saturation condition if mouse_position > self.slider_max: self.slider_value = self.slider_max elif mouse_position < -self.slider_max: self.slider_value = -self.slider_max else: self.slider_value = mouse_position # This method updates the elements of the Cart Figure which change at every frame. # Not that these elements are not ploted directly by this method # but rather returned as objects which can be used by another function # e.g. animation function from matplotlib package def update_drawing(self): # Draw mast mast_position = (self.s.position - (self.MastThickness / 2.0)) self.Mast.set_x(mast_position) # Draw rotated mast t21 = transforms.Affine2D().translate(-mast_position, -1.25 * self.WheelRadius) t22 = transforms.Affine2D().rotate(self.s.angle) t23 = transforms.Affine2D().translate(mast_position, 1.25 * self.WheelRadius) self.t2 = t21 + t22 + t23 # Draw Chassis self.Chassis.set_x(self.s.position - (self.CartLength / 2.0)) # Draw Wheels self.WheelLeft.center = (self.s.position - self.WheelToMiddle, self.y_wheel) self.WheelRight.center = (self.s.position + self.WheelToMiddle, self.y_wheel) # Draw SLider self.Slider.set_width(self.slider_value) return self.Mast, self.t2, self.Chassis, self.WheelRight, self.WheelLeft, self.Slider # This method resets the dictionary keeping the history of simulation def reset_dict_history(self): self.dict_history = {'time': [0.0], 'dt': [0.0], 's.position': [self.s.position], 's.positionD': [self.s.positionD], 's.positionDD': [self.s.positionDD], 's.angle': [self.s.angle], 's.angleD': [self.s.angleD], 's.angleDD': [self.s.angleDD], 'u': [self.u], 'Q': [self.Q], 'target_position': [self.target_position]} def augment_dict_history(self): # Augment dict self.dict_history['s.angle.sin'] = [around(np.sin(x), 4) for x in self.dict_history['s.angle']] self.dict_history['s.angle.cos'] = [around(np.cos(x), 4) for x in self.dict_history['s.angle']] # This method saves the dictionary keeping the history of simulation to a .csv file def save_history_csv(self, csv_name=None): # Make folder to save data (if not yet existing) try: os.makedirs('./data') except FileExistsError: pass # Set path where to save the data if csv_name is None or csv_name == '': logpath = './data/' + 'CP_' + self.controller_name + str(datetime.now().strftime('_%Y-%m-%d_%H-%M-%S')) + '.csv' else: logpath = './data/' + csv_name if csv_name[-4:] != '.csv': logpath += '.csv' # If such file exists, add index to the end (do not overwrite) net_index = 1 logpath_new = logpath while True: if os.path.isfile(logpath_new): logpath_new = logpath[:-4] else: logpath = logpath_new break logpath_new = logpath_new + '-' + str(net_index) + '.csv' net_index += 1 # Write the .csv file with open(logpath, "w") as outfile: writer = csv.writer(outfile) writer.writerow(['# ' + 'This is CartPole experiment from {} at time {}' .format(datetime.now().strftime('%d.%m.%Y'), datetime.now().strftime('%H:%M:%S'))]) writer.writerow(['# Number of data points: {}'.format(len(self.dict_history['time']))]) writer.writerow(['# Controller: {}'.format(self.controller_name)]) writer.writerow(['#']) writer.writerow(['# Parameters:']) for k in self.p.__dict__: writer.writerow(['# ' + k + ': ' + str(getattr(self.p, k))]) writer.writerow(['#']) writer.writerow(['# Data:']) writer.writerow(self.dict_history.keys()) writer.writerows(zip(*self.dict_history.values())) def get_history(self): # Augment dict self.dict_history['s.angle.sin'] = [around(np.sin(x), 4) for x in self.dict_history['s.angle']] self.dict_history['s.angle.cos'] = [around(np.cos(x), 4) for x in self.dict_history['s.angle']] return pd.DataFrame(self.dict_history) def load_history_csv(self, csv_name=None, visualisation_only=True): # Set path where to save the data if csv_name is None or csv_name == '': # get the latest file try: list_of_files = glob.glob('./data/' + '/*.csv') file_path = max(list_of_files, key=os.path.getctime) except FileNotFoundError: print('Cannot load: No experiment recording found in data folder ' + './data/') return False else: filename = csv_name if csv_name[-4:] != '.csv': filename += '.csv' # check if file found in DATA_FOLDER_NAME or at local starting point if not os.path.isfile(filename): file_path = os.path.join('data', filename) if not os.path.isfile(file_path): print( 'Cannot load: There is no experiment recording file with name {} at local folder or in {}'.format( filename, './data/')) return False # Get race recording print('Loading file {}'.format(file_path)) try: data: pd.DataFrame = pd.read_csv(file_path, comment='#') # skip comment lines starting with # except Exception as e: print('Cannot load: Caught {} trying to read CSV file {}'.format(e, file_path)) return False if visualisation_only: data = data[['time', 'dt', 's.position', 's.positionD', 's.angle', 'u', 'target_position']] return data def set_mode(self, new_mode=0): if new_mode == 0: self.mode = 0 self.controller = None self.controller_name = 'free' self.slider_max = self.Q_max elif new_mode == 1: self.mode = 1 self.controller = self.controller_lqr self.controller_name = 'lqr' self.slider_max = self.HalfLength # Set the maximal allowed value of the slider elif new_mode == 2: self.mode = 2 self.controller = self.controller_do_mpc self.controller_name = 'do-mpc' self.slider_max = self.HalfLength # Set the maximal allowed value of the slider elif new_mode == 3: self.mode = 3 self.controller = self.controller_do_mpc_discrete self.controller_name = 'do-mpc-discrete' self.slider_max = self.HalfLength # Set the maximal allowed value of the slider
def init_graphical_elements(self): self.CartLength = 10.0 self.WheelRadius = 0.5 self.WheelToMiddle = 4.0 self.y_plane = 0.0 self.y_wheel = self.y_plane + self.WheelRadius self.MastHight = 10.0 # For drawing only. For calculation see L self.MastThickness = 0.05 self.HalfLength = 50.0 # Length of the track self.y_acceleration_arrow = 1.5 * self.WheelRadius self.scaling_dx_acceleration_arrow = 20.0 self.x_acceleration_arrow = ( self.s.position + # np.sign(self.Q) * (self.CartLength / 2.0) + self.scaling_dx_acceleration_arrow * self.Q) # Initialize elements of the drawing self.Mast = FancyBboxPatch( xy=(self.s.position - (self.MastThickness / 2.0), 1.25 * self.WheelRadius), width=self.MastThickness, height=self.MastHight, fc='g') self.Chassis = FancyBboxPatch( (self.s.position - (self.CartLength / 2.0), self.WheelRadius), self.CartLength, 1 * self.WheelRadius, fc='r') self.WheelLeft = Circle( (self.s.position - self.WheelToMiddle, self.y_wheel), radius=self.WheelRadius, fc='y', ec='k', lw=5) self.WheelRight = Circle( (self.s.position + self.WheelToMiddle, self.y_wheel), radius=self.WheelRadius, fc='y', ec='k', lw=5) self.Acceleration_Arrow = FancyArrowPatch( (self.s.position, self.y_acceleration_arrow), (self.x_acceleration_arrow, self.y_acceleration_arrow), arrowstyle='simple', mutation_scale=10, facecolor='gold', edgecolor='orange') self.Slider_Arrow = FancyArrowPatch((self.slider_value, 0), (self.slider_value, 0), arrowstyle='fancy', mutation_scale=50) self.Slider_Bar = Rectangle((0.0, 0.0), self.slider_value, 1.0) self.t2 = transforms.Affine2D().rotate( 0.0) # An abstract container for the transform rotating the mast
class CartPole: def __init__(self): # Global time of the simulation self.time = 0.0 # CartPole is initialized with state, control input, target position all zero # This is however usually changed before running the simulation. Treat it just as placeholders. # Container for the augmented state (angle, position and their first and second derivatives)of the cart self.s = s0 # (s like state) # Variables for control input and target position. self.u = 0.0 # Physical force acting on the cart self.Q = 0.0 # Dimensionless motor power in the range [-1,1] from which force is calculated with Q2u() method self.target_position = 0.0 # Physical parameters of the CartPole self.p = p_globals # region Time scales for simulation step, controller update and saving data # See last paragraph of "Time scales" section for explanations # ∆t in number of steps (related to simulation time step) # is set while setting corresponding dt through @property self.dt_controller_number_of_steps = 0 self.dt_save_number_of_steps = 0 # Counts time steps from last controller update or saving # is set while setting corresponding dt through @property self.dt_controller_steps_counter = 0 self.dt_save_steps_counter = 0 # Helper variables to set timescales self._dt_simulation = None self._dt_controller = None self._dt_save = None self.dt_simulation = None # s, Update CartPole dynamical state every dt_simulation seconds self.dt_controller = None # s, Recalculate control input every dt_controller_default seconds self.dt_save = None # s, Save CartPole state every dt_save_default seconds # endregion # region Variables controlling operation of the program - can be modified directly from CartPole environment self.rounding_decimals = 5 # Sets number of digits after coma to save in experiment history for each feature self.save_data_in_cart = True # Decides whether to store whole data of the experiment in dict_history or not self.stop_at_90 = False # If true pole is blocked after reaching the horizontal position # endregion # region Variables controlling operation of the program - should not be modified directly self.save_flag = False # Signalizes that the current time step should be saved self.csv_filepath = None # Where to save the experiment history. self.controller = None # Placeholder for the currently used controller function self.controller_name = '' # Placeholder for the currently used controller name self.controller_idx = None # Placeholder for the currently used controller index self.controller_names = self.get_available_controller_names( ) # list of controllers available in controllers folder # endregion # region Variables for generating experiments with random target trace # Parameters for random trace generation # These need to be set, before CartPole can generate random trace and random experiment self.track_relative_complexity = None # randomly placed target points/s self.length_of_experiment = None # seconds, length of the random length trace self.interpolation_type = None # Sets how to interpolate between turning points of random trace # Possible choices: '0-derivative-smooth', 'linear', 'previous' # '0-derivative-smooth' # -> turning points are connected with smooth interpolation curve having derivative = 0 at each turning p. # 'linear' -> turning points are connected with line segments # 'previous' -> between two turning points the value of the preceding point is kept constant. # In this last setting endpoint if set has no visible effect # (may however appear in the last line of the recording - TODO: not checked) self.turning_points_period = None # How turning points should be distributed # Possible choices: 'regular', 'random' # Regular means that they are equidistant from each other # Random means we pick randomly points at time axis at which we place turning points # Where the target position of the random experiment starts and end: self.start_random_target_position_at = None self.end_random_target_position_at = None # Alternatively you can provide a list of target positions. # e.g. self.turning_points = [10.0, 0.0, 0.0] # If not None this variable has precedence - # track_relative_complexity, start/end_random_target_position_at_globals have no effect. self.turning_points = None self.random_track_f = None # Function interpolataing the random target position between turning points self.new_track_generated = False # Flag informing that a new target position track is generated self.t_max_pre = None # Placeholder for the end time of the generated random experiment self.number_of_timesteps_in_random_experiment = None self.use_pregenerated_target_position = False # Informs method performing experiment # not to take target position from environment # endregion and self.dict_history = {} # Dictionary holding the experiment history # region Variables initialization for drawing/animating a CartPole # DIMENSIONS OF THE DRAWING ONLY!!! # NOTHING TO DO WITH THE SIMULATION AND NOT INTENDED TO BE MANIPULATED BY USER !!! # Variable relevant for interactive use of slider self.slider_max = 0.0 self.slider_value = 0.0 # Parameters needed to display CartPole in GUI # They are assigned with values in self.init_elements() self.CartLength = None self.WheelRadius = None self.WheelToMiddle = None self.y_plane = None self.y_wheel = None self.MastHight = None # For drawing only. For calculation see L self.MastThickness = None self.HalfLength = None # Length of the track # Elements of the drawing self.Mast = None self.Chassis = None self.WheelLeft = None self.WheelRight = None # Arrow indicating acceleration (=motor power) self.Acceleration_Arrow = None self.y_acceleration_arrow = None self.scaling_dx_acceleration_arrow = None self.x_acceleration_arrow = None # Depending on mode, slider may be displayed either as bar or as an arrow self.Slider_Bar = None self.Slider_Arrow = None self.t2 = None # An abstract container for the transform rotating the mast self.init_graphical_elements( ) # Assign proper object to the above variables # endregion # region Initialize CartPole in manual-stabilization mode self.set_controller('manual-stabilization') # endregion # region 1. Methods related to dynamic evolution of CartPole system # This method changes the internal state of the CartPole # from a state at time t to a state at t+dt # We assume this function is called for the first time to calculate first time step # @profile(precision=4) def update_state(self): # Update the total time of the simulation self.time = self.time + self.dt_simulation # Update target position depending on the mode of operation if self.use_pregenerated_target_position: # If time exceeds the max time for which target position was defined if self.time >= self.t_max_pre: return self.target_position = self.random_track_f(self.time) self.slider_value = self.target_position # Assign target position to slider to display it else: if self.controller_name == 'manual-stabilization': self.target_position = 0.0 # In this case target position is not used. # This just fill the corresponding column in history with zeros else: self.target_position = self.slider_value # Get target position from slider # Calculate the next state self.cartpole_integration() # Snippet to stop pole at +/- 90 deg if enabled zero_DD = None if self.stop_at_90: if self.s.angle >= np.pi / 2: self.s.angle = np.pi / 2 self.s.angleD = 0.0 zero_DD = True # Make also second derivatives 0 after they are calculated elif self.s.angle <= -np.pi / 2: self.s.angle = -np.pi / 2 self.s.angleD = 0.0 zero_DD = True # Make also second derivatives 0 after they are calculated else: zero_DD = False # Wrap angle to +/-π self.s.angle = wrap_angle_rad(self.s.angle) # In case in the next step the wheel of the cart # went beyond the track # Bump elastically into an (invisible) boarder if (abs(self.s.position) + self.WheelToMiddle) > self.HalfLength: self.s.positionD = -self.s.positionD # Determine the dimensionless [-1,1] value of the motor power Q self.Update_Q() # Convert dimensionless motor power to a physical force acting on the Cart self.u = Q2u(self.Q, self.p) # Update second derivatives self.s.angleDD, self.s.positionDD = cartpole_ode( self.p, self.s, self.u) if zero_DD: self.s.angleDD = 0.0 # Calculate time steps from last saving # The counter should be initialized at max-1 to start with a control input update self.dt_save_steps_counter += 1 # If update time interval elapsed save current state and zero the counter if self.dt_save_steps_counter == self.dt_save_number_of_steps: # If user chose to save history of the simulation it is saved now # It is saved first internally to a dictionary in the Cart instance if self.save_data_in_cart: # Saving simulation data self.dict_history['time'].append(self.time) self.dict_history['s.position'].append(self.s.position) self.dict_history['s.positionD'].append(self.s.positionD) self.dict_history['s.positionDD'].append(self.s.positionDD) self.dict_history['s.angle'].append(self.s.angle) self.dict_history['s.angleD'].append(self.s.angleD) self.dict_history['s.angleDD'].append(self.s.angleDD) self.dict_history['u'].append(self.u) self.dict_history['Q'].append(self.Q) # The target_position is not always meaningful # If it is not meaningful all values in this column are set to 0 self.dict_history['target_position'].append( self.target_position) self.dict_history['s.angle.sin'].append(np.sin(self.s.angle)) self.dict_history['s.angle.cos'].append(np.cos(self.s.angle)) else: self.dict_history = { 'time': [self.time], 's.position': [self.s.position], 's.positionD': [self.s.positionD], 's.positionDD': [self.s.positionDD], 's.angle': [self.s.angle], 's.angleD': [self.s.angleD], 's.angleDD': [self.s.angleDD], 'u': [self.u], 'Q': [self.Q], 'target_position': [self.target_position], 's.angle.sin': [np.sin(self.s.angle)], 's.angle.cos': [np.cos(self.s.angle)] } self.save_flag = True self.dt_save_steps_counter = 0 # A method integrating the cartpole ode over time step dt # Currently we use a simple single step Euler stepping def cartpole_integration(self): """Simple single step integration of CartPole state by dt Takes state as SimpleNamespace, but returns as separate variables :param s: state of the CartPole (contains: s.position, s.positionD, s.angle and s.angleD) :param dt: time step by which the CartPole state should be integrated """ self.s.position = self.s.position + self.s.positionD * self.dt_simulation self.s.positionD = self.s.positionD + self.s.positionDD * self.dt_simulation self.s.angle = self.s.angle + self.s.angleD * self.dt_simulation self.s.angleD = self.s.angleD + self.s.angleDD * self.dt_simulation # Determine the dimensionless [-1,1] value of the motor power Q # The function loads an external controller from PATH_TO_CONTROLLERS # This function should be called for the first time to calculate 0th time step # Otherwise it goes out of sync with saving def Update_Q(self): # Calculate time steps from last update # The counter should be initialized at max-1 to start with a control input update self.dt_controller_steps_counter += 1 # If update time interval elapsed update control input and zero the counter if self.dt_controller_steps_counter == self.dt_controller_number_of_steps: if self.controller_name == 'manual-stabilization': # in this case slider corresponds already to the power of the motor self.Q = self.slider_value else: # in this case slider gives a target position, lqr regulator self.Q = self.controller.step(self.s, self.target_position, self.time) self.dt_controller_steps_counter = 0 # endregion # region 2. Methods related to experiment history as a whole: saving, loading, plotting, resetting # This method saves the dictionary keeping the history of simulation to a .csv file def save_history_csv(self, csv_name=None, mode='init', length_of_experiment='unknown'): if mode == 'init': # Make folder to save data (if not yet existing) try: os.makedirs(PATH_TO_EXPERIMENT_RECORDINGS[:-1]) except FileExistsError: pass # Set path where to save the data if csv_name is None or csv_name == '': self.csv_filepath = PATH_TO_EXPERIMENT_RECORDINGS + 'CP_' + self.controller_name + str( datetime.now().strftime('_%Y-%m-%d_%H-%M-%S')) + '.csv' else: self.csv_filepath = PATH_TO_EXPERIMENT_RECORDINGS + csv_name if csv_name[-4:] != '.csv': self.csv_filepath += '.csv' # If such file exists, append index to the end (do not overwrite) net_index = 1 logpath_new = self.csv_filepath while True: if os.path.isfile(logpath_new): logpath_new = self.csv_filepath[:-4] else: self.csv_filepath = logpath_new break logpath_new = logpath_new + '-' + str(net_index) + '.csv' net_index += 1 # Write the .csv file with open(self.csv_filepath, "a") as outfile: writer = csv.writer(outfile) writer.writerow([ '# ' + 'This is CartPole experiment from {} at time {}'.format( datetime.now().strftime('%d.%m.%Y'), datetime.now().strftime('%H:%M:%S')) ]) try: repo = Repo() git_revision = repo.head.object.hexsha except: git_revision = 'unknown' writer.writerow( ['# ' + 'Done with git-revision: {}'.format(git_revision)]) writer.writerow(['#']) writer.writerow([ '# Length of experiment: {} s'.format( str(length_of_experiment)) ]) writer.writerow(['#']) writer.writerow(['# Time intervals dt:']) writer.writerow( ['# Simulation: {} s'.format(str(self.dt_simulation))]) writer.writerow([ '# Controller update: {} s'.format(str(self.dt_controller)) ]) writer.writerow(['# Saving: {} s'.format(str(self.dt_save))]) writer.writerow(['#']) writer.writerow( ['# Controller: {}'.format(self.controller_name)]) writer.writerow(['#']) writer.writerow(['# Parameters:']) for k in self.p.__dict__: writer.writerow( ['# ' + k + ': ' + str(getattr(self.p, k))]) writer.writerow(['#']) writer.writerow(['# Data:']) writer.writerow(self.dict_history.keys()) elif mode == 'save online': # Save this dict with open(self.csv_filepath, "a") as outfile: writer = csv.writer(outfile) self.dict_history = { key: np.around(value, self.rounding_decimals) for key, value in self.dict_history.items() } writer.writerows(zip(*self.dict_history.values())) self.save_now = False elif mode == 'save offline': # Round data to a set precision with open(self.csv_filepath, "a") as outfile: writer = csv.writer(outfile) self.dict_history = { key: np.around(value, self.rounding_decimals) for key, value in self.dict_history.items() } writer.writerows(zip(*self.dict_history.values())) self.save_now = False # Another possibility to save data. # DF_history = pd.DataFrame.from_dict(self.dict_history).round(self.rounding_decimals) # DF_history.to_csv(self.csv_filepath, index=False, header=False, mode='a') # Mode (a)ppend # load csv file with experiment recording (e.g. for replay) def load_history_csv(self, csv_name=None): # Set path where to save the data if csv_name is None or csv_name == '': # get the latest file try: list_of_files = glob.glob(PATH_TO_EXPERIMENT_RECORDINGS + '/*.csv') file_path = max(list_of_files, key=os.path.getctime) except FileNotFoundError: print( 'Cannot load: No experiment recording found in data folder ' + './data/') return False else: filename = csv_name if csv_name[-4:] != '.csv': filename += '.csv' # check if file found in DATA_FOLDER_NAME or at local starting point if not os.path.isfile(filename): file_path = os.path.join(PATH_TO_EXPERIMENT_RECORDINGS, filename) if not os.path.isfile(file_path): print( 'Cannot load: There is no experiment recording file with name {} at local folder or in {}' .format(filename, PATH_TO_EXPERIMENT_RECORDINGS)) return False # Get race recording print('Loading file {}'.format(file_path)) try: data: pd.DataFrame = pd.read_csv( file_path, comment='#') # skip comment lines starting with # except Exception as e: print('Cannot load: Caught {} trying to read CSV file {}'.format( e, file_path)) return False return data # Method plotting the dynamic evolution over time of the CartPole # It should be called after an experiment and only if experiment data was saved def summary_plots(self): fontsize_labels = 14 fontsize_ticks = 12 fig, axs = plt.subplots( 4, 1, figsize=(16, 9), sharex=True) # share x axis so zoom zooms all plots # Plot angle error axs[0].set_ylabel("Angle (deg)", fontsize=fontsize_labels) axs[0].plot(np.array(self.dict_history['time']), np.array(self.dict_history['s.angle']) * 180.0 / np.pi, 'b', markersize=12, label='Ground Truth') axs[0].tick_params(axis='both', which='major', labelsize=fontsize_ticks) # Plot position axs[1].set_ylabel("position (m)", fontsize=fontsize_labels) axs[1].plot(self.dict_history['time'], self.dict_history['s.position'], 'g', markersize=12, label='Ground Truth') axs[1].tick_params(axis='both', which='major', labelsize=fontsize_ticks) # Plot motor input command axs[2].set_ylabel("motor (N)", fontsize=fontsize_labels) axs[2].plot(self.dict_history['time'], self.dict_history['u'], 'r', markersize=12, label='motor') axs[2].tick_params(axis='both', which='major', labelsize=fontsize_ticks) # Plot target position axs[3].set_ylabel("position target (m)", fontsize=fontsize_labels) axs[3].plot(self.dict_history['time'], self.dict_history['target_position'], 'k') axs[3].tick_params(axis='both', which='major', labelsize=fontsize_ticks) axs[3].set_xlabel('Time (s)', fontsize=fontsize_labels) fig.align_ylabels() plt.show() return fig, axs # endregion # region 3. Methods for generating random target position for generation of random experiment # Generates a random target position # in a form of a function interpolating between turning points def Generate_Random_Trace_Function(self): if (self.turning_points is None) or (self.turning_points == []): number_of_turning_points = int( np.floor(self.length_of_experiment * self.track_relative_complexity)) y = 2.0 * (np.random.random(number_of_turning_points) - 0.5) y = y * 0.5 * self.HalfLength if number_of_turning_points == 0: y = np.append(y, 0.0) y = np.append(y, 0.0) elif number_of_turning_points == 1: if self.start_random_target_position_at is not None: y[0] = self.start_random_target_position_at elif self.end_random_target_position_at is not None: y[0] = self.end_random_target_position_at else: pass y = np.append(y, y[0]) else: if self.start_random_target_position_at is not None: y[0] = self.start_random_target_position_at if self.end_random_target_position_at is not None: y[-1] = self.end_random_target_position_at else: number_of_turning_points = len(self.turning_points) if number_of_turning_points == 0: raise ValueError('You should not be here!') elif number_of_turning_points == 1: y = np.array([self.turning_points[0], self.turning_points[0]]) else: y = np.array(self.turning_points) number_of_timesteps = np.ceil(self.length_of_experiment / self.dt_simulation) self.t_max_pre = number_of_timesteps * self.dt_simulation random_samples = number_of_turning_points - 2 if number_of_turning_points - 2 >= 0 else 0 # t_init = linspace(0, self.t_max_pre, num=self.track_relative_complexity, endpoint=True) if self.turning_points_period == 'random': t_init = np.sort( np.random.uniform(self.dt_simulation, self.t_max_pre - self.dt_simulation, random_samples)) t_init = np.insert(t_init, 0, 0.0) t_init = np.append(t_init, self.t_max_pre) elif self.turning_points_period == 'regular': t_init = np.linspace(0, self.t_max_pre, num=random_samples + 2, endpoint=True) else: raise NotImplementedError( 'There is no mode corresponding to this value of turning_points_period variable' ) # Try algorithm setting derivative to 0 a each point if self.interpolation_type == '0-derivative-smooth': yder = [[y[i], 0] for i in range(len(y))] random_track_f = BPoly.from_derivatives(t_init, yder) elif self.interpolation_type == 'linear': random_track_f = interp1d(t_init, y, kind='linear') elif self.interpolation_type == 'previous': random_track_f = interp1d(t_init, y, kind='previous') else: raise ValueError('Unknown interpolation type.') # Truncate the target position to be not grater than 80% of track length def random_track_f_truncated(time): target_position = random_track_f(time) if target_position > 0.8 * self.HalfLength: target_position = 0.8 * self.HalfLength elif target_position < -0.8 * self.HalfLength: target_position = -0.8 * self.HalfLength return target_position self.random_track_f = random_track_f_truncated self.new_track_generated = True # Prepare CartPole Instance to perform an experiment with random target position trace def setup_cartpole_random_experiment( self, # Initial state s0=None, controller=None, dt_simulation=None, dt_controller=None, dt_save=None, # Settings related to random trace generation track_relative_complexity=None, length_of_experiment=None, interpolation_type=None, turning_points_period=None, start_random_target_position_at=None, end_random_target_position_at=None, turning_points=None): # Set time scales: self.dt_simulation = dt_simulation self.dt_controller = dt_controller self.dt_save = dt_save # Set CartPole in the right (automatic control) mode # You may want to provide it before this function not to reload it every time if controller is not None: self.set_controller(controller) # Set initial state self.s = s0 self.track_relative_complexity = track_relative_complexity self.length_of_experiment = length_of_experiment self.interpolation_type = interpolation_type self.turning_points_period = turning_points_period self.start_random_target_position_at = start_random_target_position_at self.end_random_target_position_at = end_random_target_position_at self.turning_points = turning_points self.Generate_Random_Trace_Function() self.use_pregenerated_target_position = 1 self.number_of_timesteps_in_random_experiment = int( np.ceil(self.length_of_experiment / self.dt_simulation)) # Target position at time 0 self.target_position = self.random_track_f(self.time) # Make already in the first timestep Q appropriate to the initial state, target position and controller self.Update_Q() self.set_cartpole_state_at_t0(reset_mode=2, s=self.s, Q=self.Q, target_position=self.target_position) # Runs a random experiment with parameters set with setup_cartpole_random_experiment # And saves the experiment recording to csv file # @profile(precision=4) def run_cartpole_random_experiment(self, csv=None, save_mode='offline'): """ This function runs a random CartPole experiment and returns the history of CartPole states, control inputs and desired cart position """ if save_mode == 'offline': self.save_data_in_cart = True elif save_mode == 'online': self.save_data_in_cart = False else: raise ValueError('Unknown save mode value') # Create csv file for savign self.save_history_csv(csv_name=csv, mode='init', length_of_experiment=self.length_of_experiment) # Save 0th timestep if save_mode == 'online': self.save_history_csv(csv_name=csv, mode='save online') # Run the CartPole experiment for number of time for _ in trange(self.number_of_timesteps_in_random_experiment): # Print an error message if it runs already to long (should stop before) if self.time > self.t_max_pre: raise Exception( 'ERROR: It seems the experiment is running too long...') self.update_state() # Additional option to stop the experiment if abs(self.s.position) > 45.0: break print('Cart went out of safety boundaries') # if abs(CartPoleInstance.s.angle) > 0.8*np.pi: # raise ValueError('Cart went unstable') if save_mode == 'online' and self.save_flag: self.save_history_csv(csv_name=csv, mode='save online') self.save_flag = False data = pd.DataFrame(self.dict_history) if save_mode == 'offline': self.save_history_csv(csv_name=csv, mode='save offline') self.summary_plots() # Set CartPole state - the only use is to make sure that experiment history is discared # Maybe you can delete this line self.set_cartpole_state_at_t0(reset_mode=0) return data # endregion # region 4. Methods "Get, set, reset" # Method returns the list of controllers available in the PATH_TO_CONTROLLERS folder def get_available_controller_names(self): """ Method returns the list of controllers available in the PATH_TO_CONTROLLERS folder """ controller_files = glob.glob(PATH_TO_CONTROLLERS + 'controller_' + '*.py') controller_names = ['manual-stabilization'] controller_names.extend( np.sort([ os.path.basename(item)[len('controller_'):-len('.py')].replace( '_', '-') for item in controller_files ])) return controller_names # Set the controller of CartPole def set_controller(self, controller_name=None, controller_idx=None): """ The method sets a new controller as the current controller of the CartPole instance. The controller may be indicated either by its name or by the index on the controller list (see get_available_controller_names method). """ # Check if the proper information was provided: either controller_name or controller_idx if (controller_name is None) and (controller_idx is None): raise ValueError( 'You have to specify either controller_name or controller_idx to set a new controller.' 'You have specified none of the two.') elif (controller_name is not None) and (controller_idx is not None): raise ValueError( 'You have to specify either controller_name or controller_idx to set a new controller.' 'You have specified both.') else: pass # If controller name provided get controller index and vice versa if (controller_name is not None): try: controller_idx = self.controller_names.index(controller_name) except ValueError: raise ValueError( '{} is not in list. \n In list are: {}'.format( controller_name, self.controller_names)) else: controller_name = self.controller_names[controller_idx] # save controller name and index to variables in the CartPole namespace self.controller_name = controller_name self.controller_idx = controller_idx # Load controller if self.controller_name == 'manual-stabilization': self.controller = None else: controller_full_name = 'controller_' + self.controller_name.replace( '-', '_') path_import = PATH_TO_CONTROLLERS[2:].replace('/', '.').replace( r'\\', '.') import_str = 'from ' + path_import + controller_full_name + ' import ' + controller_full_name exec(import_str) self.controller = eval(controller_full_name + '()') # Set the maximal allowed value of the slider - relevant only for GUI if self.controller_name == 'manual-stabilization': self.slider_max = 1.0 self.Slider_Arrow.set_positions((0, 0), (0, 0)) else: self.slider_max = self.p.TrackHalfLength self.Slider_Bar.set_width(0.0) # This method resets the internal state of the CartPole instance # The starting state (for t = 0) may be # all zeros (reset_mode = 0) # set in this function (reset_mode = 1) # provide by user (reset_mode = 1), by giving s, Q and target_position def set_cartpole_state_at_t0(self, reset_mode=1, s=None, Q=None, target_position=None): self.time = 0.0 if reset_mode == 0: # Don't change it self.s.position = self.s.positionD = self.s.positionDD = 0.0 self.s.angle = self.s.angleD = self.s.angleDD = 0.0 self.Q = self.u = 0.0 self.slider = self.target_position = 0.0 elif reset_mode == 1: # You may change this but be carefull with other user. Better use 3 # You can change here with which initial parameters you wish to start the simulation self.s.position = 0.0 self.s.positionD = 0.0 self.s.angle = (2.0 * np.random.normal() - 1.0) * np.pi / 180.0 # np.pi/2.0 # self.s.angleD = 0.0 # 1.0 self.target_position = self.slider_value self.Q = 0.0 self.u = Q2u(self.Q, self.p) self.s.angleDD, self.s.positionDD = cartpole_ode( self.p, self.s, self.u) elif reset_mode == 2: # Don't change it if (s is not None) and (Q is not None) and (target_position is not None): self.s = s self.Q = Q self.slider = self.target_position = target_position self.u = Q2u(self.Q, self.p) # Calculate CURRENT control input self.s.angleDD, self.s.positionDD = cartpole_ode( self.p, self.s, self.u) # Calculate CURRENT second derivatives else: raise ValueError( 's, Q or target position not provided for initial state') # Reset the dict keeping the experiment history and save the state for t = 0 self.dict_history = { 'time': [self.time], 's.position': [self.s.position], 's.positionD': [self.s.positionD], 's.positionDD': [self.s.positionDD], 's.angle': [self.s.angle], 's.angleD': [self.s.angleD], 's.angleDD': [self.s.angleDD], 'u': [self.u], 'Q': [self.Q], 'target_position': [self.target_position], 's.angle.sin': [np.sin(self.s.angle)], 's.angle.cos': [np.cos(self.s.angle)] } # region Get and set timescales # Makes sure that when dt is updated also related variables are updated @property def dt_simulation(self): return self._dt_simulation @dt_simulation.setter def dt_simulation(self, value): self._dt_simulation = value if self._dt_controller is not None: self.dt_controller_number_of_steps = np.rint( self._dt_controller / value).astype(np.int32) if self.dt_controller_number_of_steps == 0: self.dt_controller_number_of_steps = 1 # Initialize counter at max value to start with update self.dt_controller_steps_counter = self.dt_controller_number_of_steps - 1 if self._dt_save is not None: self.dt_save_number_of_steps = np.rint(self._dt_save / value).astype(np.int32) if self.dt_save_number_of_steps == 0: self.dt_save_number_of_steps = 1 self.dt_save_steps_counter = 0 @property def dt_controller(self): return self._dt_controller @dt_controller.setter def dt_controller(self, value): self._dt_controller = value if self._dt_simulation is not None: self.dt_controller_number_of_steps = np.rint( value / self._dt_simulation).astype(np.int32) if self.dt_controller_number_of_steps == 0: self.dt_controller_number_of_steps = 1 # Initialize counter at max value to start with update self.dt_controller_steps_counter = self.dt_controller_number_of_steps - 1 @property def dt_save(self): return self._dt_save @dt_save.setter def dt_save(self, value): self._dt_save = value if self._dt_simulation is not None: self.dt_save_number_of_steps = np.rint( value / self._dt_simulation).astype(np.int32) if self.dt_save_number_of_steps == 0: self.dt_save_number_of_steps = 1 # This counter is initialized at 0 - 0th step is saved manually self.dt_save_steps_counter = 0 # endregion # endregion # region 5. Methods needed to display CartPole in GUI """ This section contains methods related to displaying CartPole in GUI of the simulator. One could think of moving these function outside of CartPole class and connecting them rather more tightly with GUI of the simulator. We leave them however as a part of CartPole class as they rely on variables of the CartPole. """ # This method initializes CartPole elements to be plotted in CartPole GUI def init_graphical_elements(self): self.CartLength = 10.0 self.WheelRadius = 0.5 self.WheelToMiddle = 4.0 self.y_plane = 0.0 self.y_wheel = self.y_plane + self.WheelRadius self.MastHight = 10.0 # For drawing only. For calculation see L self.MastThickness = 0.05 self.HalfLength = 50.0 # Length of the track self.y_acceleration_arrow = 1.5 * self.WheelRadius self.scaling_dx_acceleration_arrow = 20.0 self.x_acceleration_arrow = ( self.s.position + # np.sign(self.Q) * (self.CartLength / 2.0) + self.scaling_dx_acceleration_arrow * self.Q) # Initialize elements of the drawing self.Mast = FancyBboxPatch( xy=(self.s.position - (self.MastThickness / 2.0), 1.25 * self.WheelRadius), width=self.MastThickness, height=self.MastHight, fc='g') self.Chassis = FancyBboxPatch( (self.s.position - (self.CartLength / 2.0), self.WheelRadius), self.CartLength, 1 * self.WheelRadius, fc='r') self.WheelLeft = Circle( (self.s.position - self.WheelToMiddle, self.y_wheel), radius=self.WheelRadius, fc='y', ec='k', lw=5) self.WheelRight = Circle( (self.s.position + self.WheelToMiddle, self.y_wheel), radius=self.WheelRadius, fc='y', ec='k', lw=5) self.Acceleration_Arrow = FancyArrowPatch( (self.s.position, self.y_acceleration_arrow), (self.x_acceleration_arrow, self.y_acceleration_arrow), arrowstyle='simple', mutation_scale=10, facecolor='gold', edgecolor='orange') self.Slider_Arrow = FancyArrowPatch((self.slider_value, 0), (self.slider_value, 0), arrowstyle='fancy', mutation_scale=50) self.Slider_Bar = Rectangle((0.0, 0.0), self.slider_value, 1.0) self.t2 = transforms.Affine2D().rotate( 0.0) # An abstract container for the transform rotating the mast # This method accepts the mouse position and updated the slider value accordingly # The mouse position has to be captured by a function not included in this class def update_slider(self, mouse_position): # The if statement formulates a saturation condition if mouse_position > self.slider_max: self.slider_value = self.slider_max elif mouse_position < -self.slider_max: self.slider_value = -self.slider_max else: self.slider_value = mouse_position # This method draws elements and set properties of the CartPole figure # which do not change at every frame of the animation def draw_constant_elements(self, fig, AxCart, AxSlider): # Delete all elements of the Figure AxCart.clear() AxSlider.clear() ## Upper chart with Cart Picture # Set x and y limits AxCart.set_xlim((-self.HalfLength * 1.1, self.HalfLength * 1.1)) AxCart.set_ylim((-1.0, 15.0)) # Remove ticks on the y-axes AxCart.yaxis.set_major_locator(plt.NullLocator( )) # NullLocator is used to disable ticks on the Figures # Draw track Floor = Rectangle((-self.HalfLength, -1.0), 2 * self.HalfLength, 1.0, fc='brown') AxCart.add_patch(Floor) # Draw an invisible point at constant position # Thanks to it the axes is drawn high enough for the mast InvisiblePointUp = Rectangle((0, self.MastHight + 2.0), self.MastThickness, 0.0001, fc='w', ec='w') AxCart.add_patch(InvisiblePointUp) # Apply scaling AxCart.axis('scaled') ## Lower Chart with Slider # Set y limits AxSlider.set(xlim=(-1.1 * self.slider_max, self.slider_max * 1.1)) # Remove ticks on the y-axes AxSlider.yaxis.set_major_locator(plt.NullLocator( )) # NullLocator is used to disable ticks on the Figures # Apply scaling AxSlider.set_aspect("auto") return fig, AxCart, AxSlider # This method updates the elements of the Cart Figure which change at every frame. # Not that these elements are not ploted directly by this method # but rather returned as objects which can be used by another function # e.g. animation function from matplotlib package def update_drawing(self): self.x_acceleration_arrow = ( self.s.position + # np.sign(self.Q) * (self.CartLength / 2.0) + self.scaling_dx_acceleration_arrow * self.Q) self.Acceleration_Arrow.set_positions( (self.s.position, self.y_acceleration_arrow), (self.x_acceleration_arrow, self.y_acceleration_arrow)) # Draw mast mast_position = (self.s.position - (self.MastThickness / 2.0)) self.Mast.set_x(mast_position) # Draw rotated mast t21 = transforms.Affine2D().translate(-mast_position, -1.25 * self.WheelRadius) if ANGLE_CONVENTION == 'CLOCK-NEG': t22 = transforms.Affine2D().rotate(self.s.angle) elif ANGLE_CONVENTION == 'CLOCK-POS': t22 = transforms.Affine2D().rotate(-self.s.angle) else: raise ValueError('Unknown angle convention') t23 = transforms.Affine2D().translate(mast_position, 1.25 * self.WheelRadius) self.t2 = t21 + t22 + t23 # Draw Chassis self.Chassis.set_x(self.s.position - (self.CartLength / 2.0)) # Draw Wheels self.WheelLeft.center = (self.s.position - self.WheelToMiddle, self.y_wheel) self.WheelRight.center = (self.s.position + self.WheelToMiddle, self.y_wheel) # Draw SLider if self.controller_name == 'manual-stabilization': self.Slider_Bar.set_width(self.slider_value) else: self.Slider_Arrow.set_positions((self.slider_value, 0), (self.slider_value, 1.0)) return self.Mast, self.t2, self.Chassis, self.WheelRight, self.WheelLeft,\ self.Slider_Bar, self.Slider_Arrow, self.Acceleration_Arrow # A function redrawing the changing elements of the Figure def run_animation(self, fig): def init(): # Adding variable elements to the Figure fig.AxCart.add_patch(self.Mast) fig.AxCart.add_patch(self.Chassis) fig.AxCart.add_patch(self.WheelLeft) fig.AxCart.add_patch(self.WheelRight) fig.AxCart.add_patch(self.Acceleration_Arrow) fig.AxSlider.add_patch(self.Slider_Bar) fig.AxSlider.add_patch(self.Slider_Arrow) return self.Mast, self.Chassis, self.WheelLeft, self.WheelRight,\ self.Slider_Bar, self.Slider_Arrow, self.Acceleration_Arrow def animationManage(i): # Updating variable elements self.update_drawing() # Special care has to be taken of the mast rotation self.t2 = self.t2 + fig.AxCart.transData self.Mast.set_transform(self.t2) return self.Mast, self.Chassis, self.WheelLeft, self.WheelRight,\ self.Slider_Bar, self.Slider_Arrow, self.Acceleration_Arrow # Initialize animation object anim = animation.FuncAnimation( fig, animationManage, init_func=init, frames=300, # fargs=(CartPoleInstance,), # It was used when this function was a part of GUI class. Now left as an example how to add arguments to FuncAnimation interval=10, blit=True, repeat=True) return anim
def __init__(self, parent, handles, labels, loc = None, numpoints = None, # the number of points in the legend line markerscale = None, # the relative size of legend markers vs. original scatterpoints = 3, # TODO: may be an rcParam scatteryoffsets=None, prop = None, # properties for the legend texts # the following dimensions are in axes coords pad = None, # deprecated; use borderpad labelsep = None, # deprecated; use labelspacing handlelen = None, # deprecated; use handlelength handletextsep = None, # deprecated; use handletextpad axespad = None, # deprecated; use borderaxespad # spacing & pad defined as a fractionof the font-size borderpad = None, # the whitespace inside the legend border labelspacing=None, #the vertical space between the legend entries handlelength=None, # the length of the legend handles handletextpad=None, # the pad between the legend handle and text borderaxespad=None, # the pad between the axes and legend border columnspacing=None, # spacing between columns ncol=1, # number of columns mode=None, # mode for horizontal distribution of columns. None, "expand" shadow = None, ): """ - *parent* : the artist that contains the legend - *handles* : a list of artists (lines, patches) to add to the legend - *labels* : a list of strings to label the legend Optional keyword arguments: ================ ================================================= Keyword Description ================ ================================================= loc a location code or a tuple of coordinates numpoints the number of points in the legend line prop the font property markerscale the relative size of legend markers vs. original shadow if True, draw a shadow behind legend scatteryoffsets a list of yoffsets for scatter symbols in legend borderpad the fractional whitespace inside the legend border labelspacing the vertical space between the legend entries handlelength the length of the legend handles handletextpad the pad between the legend handle and text borderaxespad the pad between the axes and legend border columnspacing the spacing between columns The dimensions of pad and spacing are given as a fraction of the fontsize. Values from rcParams will be used if None. """ from matplotlib.axes import Axes # local import only to avoid circularity from matplotlib.figure import Figure # local import only to avoid circularity Artist.__init__(self) if prop is None: self.prop=FontProperties(size=rcParams["legend.fontsize"]) else: self.prop=prop self.fontsize = self.prop.get_size_in_points() propnames=['numpoints', 'markerscale', 'shadow', "columnspacing", "scatterpoints"] localdict = locals() for name in propnames: if localdict[name] is None: value = rcParams["legend."+name] else: value = localdict[name] setattr(self, name, value) # Take care the deprecated keywords deprecated_kwds = {"pad":"borderpad", "labelsep":"labelspacing", "handlelen":"handlelength", "handletextsep":"handletextpad", "axespad":"borderaxespad"} # convert values of deprecated keywords (ginve in axes coords) # to new vaules in a fraction of the font size # conversion factor bbox = parent.bbox axessize_fontsize = min(bbox.width, bbox.height)/self.fontsize for k, v in deprecated_kwds.items(): # use deprecated value if not None and if their newer # counter part is None. if localdict[k] is not None and localdict[v] is None: warnings.warn("Use '%s' instead of '%s'." % (v, k), DeprecationWarning) setattr(self, v, localdict[k]*axessize_fontsize) continue # Otherwise, use new keywords if localdict[v] is None: setattr(self, v, rcParams["legend."+v]) else: setattr(self, v, localdict[v]) del localdict self._ncol = ncol if self.numpoints <= 0: raise ValueError("numpoints must be >= 0; it was %d"% numpoints) # introduce y-offset for handles of the scatter plot if scatteryoffsets is None: self._scatteryoffsets = np.array([3./8., 4./8., 2.5/8.]) else: self._scatteryoffsets = np.asarray(scatteryoffsets) reps = int(self.numpoints / len(self._scatteryoffsets)) + 1 self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[:self.scatterpoints] # _legend_box is an OffsetBox instance that contains all # legend items and will be initialized from _init_legend_box() # method. self._legend_box = None if isinstance(parent,Axes): self.isaxes = True self.set_figure(parent.figure) elif isinstance(parent,Figure): self.isaxes = False self.set_figure(parent) else: raise TypeError("Legend needs either Axes or Figure as parent") self.parent = parent if loc is None: loc = rcParams["legend.loc"] if not self.isaxes and loc in [0,'best']: loc = 'upper right' if is_string_like(loc): if loc not in self.codes: if self.isaxes: warnings.warn('Unrecognized location "%s". Falling back on "best"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes.keys()))) loc = 0 else: warnings.warn('Unrecognized location "%s". Falling back on "upper right"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes.keys()))) loc = 1 else: loc = self.codes[loc] if not self.isaxes and loc == 0: warnings.warn('Automatic legend placement (loc="best") not implemented for figure legend. ' 'Falling back on "upper right".') loc = 1 self._loc = loc self._mode = mode # We use FancyBboxPatch to draw a legend frame. The location # and size of the box will be updated during the drawing time. self.legendPatch = FancyBboxPatch( xy=(0.0, 0.0), width=1., height=1., facecolor='w', edgecolor='k', mutation_scale=self.fontsize, ) # The width and height of the legendPatch will be set (in the # draw()) to the length that includes the padding. Thus we set # pad=0 here. self.legendPatch.set_boxstyle("round",pad=0, #self.borderpad, rounding_size=0.2) self._set_artist_props(self.legendPatch) self._drawFrame = True # populate the legend_box with legend items. self._init_legend_box(handles, labels) self._legend_box.set_figure(self.figure)
def montage(nifti, anat, roi_dict, thr=2, fig=None, out_file=None, feature_dict=None, target_stat=None, target_value=None): if isinstance(anat, str): anat = load_image(anat) assert nifti is not None assert anat is not None assert roi_dict is not None texcol = 1 bgcol = 0 iscale = 2 weights = nifti.get_data() #weights = weights / weights.std(axis=3) features = weights.shape[-1] indices = [0] y = 8 x = int(ceil(1.0 * features / y)) font = {"size": 8} rc("font", **font) if fig is None: fig = plt.figure(figsize=[iscale * y, iscale * x / 2.5]) plt.subplots_adjust(left=0.01, right=0.99, bottom=0.01, top=0.99, wspace=0.1, hspace=0) for f in xrange(features): roi = roi_dict.get(f, None) if roi is None: continue coords = roi["top_clust"]["coords"] assert coords is not None feat = weights[:, :, :, f] feat = feat / feat.std() imax = np.max(np.absolute(feat)) imin = -imax imshow_args = {"vmax": imax, "vmin": imin} coords = ([-coords[0], -coords[1], coords[2]]) ax = fig.add_subplot(x, y, f + 1) plt.axis("off") try: plot_map(feat, xyz_affine(nifti), anat=anat.get_data(), anat_affine=xyz_affine(anat), threshold=thr, figure=fig, axes=ax, cut_coords=coords, annotate=False, cmap=cmap, draw_cross=False, **imshow_args) except Exception as e: logger.exception(e) plt.text(0.05, 0.8, str(f), transform=ax.transAxes, horizontalalignment="center", color=(texcol, texcol, texcol)) pos = [(0.05, 0.05), (0.4, 0.05), (0.8, 0.05)] colors = ["purple", "yellow", "green"] if feature_dict is not None and feature_dict.get(f, None) is not None: d = feature_dict[f] for i, key in enumerate([k for k in d if k != "real_id"]): plt.text(pos[i][0], pos[i][1], "%s=%.2f" % (key, d[key]), transform=ax.transAxes, horizontalalignment="left", color=colors[i]) if key == target_stat: assert target_value is not None if d[key] >= target_value: p_fancy = FancyBboxPatch((0.1, 0.1), 2.5 - .1, 1 - .1, boxstyle="round,pad=0.1", ec=(1., 0.5, 1.), fc="none") ax.add_patch(p_fancy) elif d[key] <= -target_value: p_fancy = FancyBboxPatch((0.1, 0.1), iscale * 2.5 - .1, iscale - .1, boxstyle="round,pad=0.1", ec=(0., 0.5, 0.), fc="none") ax.add_patch(p_fancy) # stdout.write("\rSaving montage: DONE\n") if out_file is not None: plt.savefig(out_file, transparent=True, facecolor=(bgcol, bgcol, bgcol)) else: plt.draw()
def montage(weights, fig=None, out_file=None, feature_dict=None, target_stat=None, target_value=None): features = weights.shape[0] iscale = 1 y = 8 x = int(ceil(1.0 * features / y)) font = {'size': 8} rc('font', **font) if fig is None: fig = plt.figure(figsize=[iscale * y, iscale * x]) plt.subplots_adjust(left=0.01, right=0.99, bottom=0.01, top=0.99, wspace=0.1, hspace=0) for f in xrange(features): logger.debug("Saving simtb montage %d" % f) feat = weights[f] assert feat.shape[2] == 1, feat.shape feat = feat.reshape(feat.shape[0], feat.shape[1]) feat = feat / feat.std() imax = np.max(np.absolute(feat)) imin = -imax imshow_args = {'vmax': imax, 'vmin': imin} ax = fig.add_subplot(x, y, f + 1) plt.axis("off") ax.imshow(feat, cmap=cm.RdBu_r, **imshow_args) plt.text(0.05, 0.8, str(f), transform=ax.transAxes, horizontalalignment='center', color="white") pos = [(0.05, 0.05), (0.4, 0.05), (0.8, 0.05)] colors = ["purple", "yellow", "green"] if (feature_dict is not None and feature_dict.get(f, None) is not None): d = feature_dict[f] for i, key in enumerate([k for k in d if k != "real_id"]): plt.text(pos[i][0], pos[i][1], "%s=%.2f" % (key, d[key]), transform=ax.transAxes, horizontalalignment="left", color=colors[i]) if key == target_stat: assert target_value is not None if d[key] >= target_value: p_fancy = FancyBboxPatch((0.1, 0.1), 2.5 - .1, 1 - .1, boxstyle="round,pad=0.1", ec=(1., 0.5, 1.), fc="none") ax.add_patch(p_fancy) elif d[key] <= -target_value: p_fancy = FancyBboxPatch((0.1, 0.1), iscale * 2.5 - .1, iscale - .1, boxstyle="round,pad=0.1", ec=(0., 0.5, 0.), fc="none") ax.add_patch(p_fancy) logger.info("Finished processing simtb montage") if out_file is not None: logger.info("Saving montage to %s" % out_file) plt.savefig(out_file) else: plt.draw()
axes.grid(which='major', axis='x', linewidth=0.75, linestyle='-', color='0.75') axes.grid(which='minor', axis='x', linewidth=0.25, linestyle='-', color='0.75') axes.grid(which='major', axis='y', linewidth=0.75, linestyle='-', color='0.75') axes.grid(which='minor', axis='y', linewidth=0.25, linestyle='-', color='0.75') axes.set_xticklabels([]) axes.set_yticklabels([]) # Add a title and a box around it from matplotlib.patches import FancyBboxPatch ax = plt.gca() ax.add_patch( FancyBboxPatch((-0.05, .87), width=.66, height=.165, clip_on=False, boxstyle="square,pad=0", zorder=3, facecolor='white', alpha=1.0, transform=plt.gca().transAxes)) plt.text(-0.05, 1.02, " Grid: plt.grid(...)\n", horizontalalignment='left', verticalalignment='top', size='xx-large', transform=axes.transAxes) plt.text(-0.05, 1.01,
def plot_mauve_compare(refgb, assembly_list, backbones_list, bufferlen=10000, breakwidth=40, aspect=.6, names=["Position", "Entropy"], title="Shannon Entropy by Position", output_prefix="entropy_plot.png"): assert len(assembly_list) == len(backbones_list), \ "must have same amount of assemblies as backbones" with open(refgb, "r") as rg: ref_recs = list(SeqIO.parse(rg, "genbank")) assembly_lens = [[sum([len(x) for x in ref_recs])]] for seq in assembly_list: with open(seq, "r") as inseq: assembly_lens.append( [len(x) for x in list(SeqIO.parse(inseq, "fasta"))]) backbones = parseBackbones(backbones_list) npanels = len(assembly_list) + 1 max_combined_len = max([sum(x) for x in assembly_lens]) + bufferlen print(max_combined_len) fig, ax = plt.subplots(1, 1) ax.set_title(title, y=1.08) relheight = max_combined_len * aspect coding_height = .05 * relheight # set the centers as starting relative to relheight - (2* codingdepth) relinner = relheight - (coding_height * 3) centers = [] for i in range(npanels): if i == 0: centers.append(relheight - (coding_height * 1.5)) elif i == npanels - 1: centers.append(0 + (coding_height * 1.5)) else: centers.append(relheight - ((coding_height * 1.5) + (relinner / float(npanels - 1)) * i)) xmin, xmax = 0, max_combined_len ymin, ymax = 0, relheight ax.set_xlim([xmin, xmax]) ax.set_ylim([ymin, ymax]) # plot the color shadings unused_cols = ["red", "green", "yellow", "purple", "red", "blue"] nudge = coding_height / 2 patch_list = [] for i, bblist in enumerate(backbones): for As, Ae, Bs, Be in bblist: if (Bs == 0 and Be == 0) or \ (As == 0 and Ae == 0): continue verts = [ (Bs, centers[i + 1] + nudge), # left, bottom (As, centers[0] - nudge), # left, top (Ae, centers[0] - nudge), # right, top (Be, centers[i + 1] + nudge), # right, bottom (Bs, centers[i + 1] + nudge), # ignored ] codes = [ mpl.path.Path.MOVETO, mpl.path.Path.LINETO, mpl.path.Path.LINETO, mpl.path.Path.LINETO, mpl.path.Path.CLOSEPOLY ] path = mpl.path.Path(verts, codes) patch = patches.PathPatch(path, facecolor=bgcols.get(unused_cols[0]), edgecolor=mycolors.get("clear"), lw=2) patch_list.append(patch) unused_cols.pop(0) # we want the first annotation on top [ax.add_patch(p) for p in list(reversed(patch_list))] # add annotations last_chrom_end = 0 for record in ref_recs: # coding sequence print(centers[0] * .005) coding_box = FancyBboxPatch( (last_chrom_end, centers[0] - coding_height / 2), len(record), coding_height, boxstyle="round,pad=0,rounding_size=" + str(centers[0] / 50), mutation_aspect=.5, # mutation_scale=.5, fc=mycolors['greyish'], ec=mycolors['clear']) # buffer_box = FancyBboxPatch( # (last_chrom_end + len(record), centers[0] - coding_height / 2), # last_chrom_end + len(record) + bufferlen, coding_height, # boxstyle="round,pad=0,rounding_size=0", # mutation_aspect=.5, # # mutation_scale=.5, # fc=mycolors['clear'], # ec=mycolors['clear'] # ) last_chrom_end = last_chrom_end + len(record) ax.add_patch(coding_box) # ax.add_patch(buffer_box) for i, feature in enumerate(record.features): if feature.type != "rRNA" and i == 0: #Exclude this feature continue feat_len = \ feature.location.end.position - feature.location.start.position anno_box = FancyBboxPatch( (feature.location.start.position, centers[0] - coding_height), feat_len, coding_height * 2, boxstyle="round,pad=0,rounding_size=" + str(feat_len / 2), mutation_aspect=.5, # mutation_scale=.5, fc=mycolors['redish'], ec=mycolors['redish']) ax.add_patch(anno_box) for i in range(npanels): # for each assembly if i == 0: continue with open(assembly_list[i - 1], "r") as infile: contigs = list(SeqIO.parse(infile, "fasta")) last_contig_end = 0 for record in contigs: coding_box = FancyBboxPatch( (last_contig_end, centers[i] - coding_height / 2), len(record), coding_height, boxstyle="round,pad=0,rounding_size=" + str(centers[i] / 50), mutation_aspect=.5, # mutation_scale=.5, fc=mycolors['greyish'], ec=mycolors['clear']) buffer_box = FancyBboxPatch( (last_contig_end + len(record) - breakwidth, centers[i] - coding_height), breakwidth, coding_height * 2, boxstyle="round,pad=0,rounding_size=0", mutation_aspect=.5, # mutation_scale=.5, fc="black", ec=mycolors['clear']) last_contig_end = last_contig_end + len(record) ax.add_patch(coding_box) ax.add_patch(buffer_box) ax.set_yticks(np.array(centers)) ax.set_yticklabels(names) ax.get_yaxis().set_label_coords(-.05, .1) ax.yaxis.set_ticks_position('left') ax.xaxis.set_ticks_position('top') # ax.tick_params(axis='y', colors='dimgrey') ax.tick_params(axis='x', colors='dimgrey') ax.yaxis.label.set_color('black') ax.xaxis.label.set_color('black') ax.spines['top'].set_visible(True) ax.spines["left"].set_visible(False) ax.spines["right"].set_visible(False) ax.spines["bottom"].set_visible(False) plt.tight_layout() fig.subplots_adjust(hspace=0) fig.set_size_inches(12, 12 * aspect) fig.savefig(str(output_prefix + '.png'), dpi=(200)) fig.savefig(str(output_prefix + '.pdf'), dpi=(200)) return 0
class Legend(Artist): """ Place a legend on the axes at location loc. """ # 'best' is only implemented for axes legends codes = {'best': 0, **AnchoredOffsetbox.codes} zorder = 5 def __str__(self): return "Legend" @docstring.dedent_interpd def __init__( self, parent, handles, labels, loc=None, numpoints=None, # number of points in the legend line markerscale=None, # relative size of legend markers vs. original markerfirst=True, # left/right ordering of legend marker and label scatterpoints=None, # number of scatter points scatteryoffsets=None, prop=None, # properties for the legend texts fontsize=None, # keyword to set font size directly labelcolor=None, # keyword to set the text color # spacing & pad defined as a fraction of the font-size borderpad=None, # whitespace inside the legend border labelspacing=None, # vertical space between the legend entries handlelength=None, # length of the legend handles handleheight=None, # height of the legend handles handletextpad=None, # pad between the legend handle and text borderaxespad=None, # pad between the axes and legend border columnspacing=None, # spacing between columns ncol=1, # number of columns mode=None, # horizontal distribution of columns: None or "expand" fancybox=None, # True: fancy box, False: rounded box, None: rcParam shadow=None, title=None, # legend title title_fontsize=None, # legend title font size framealpha=None, # set frame alpha edgecolor=None, # frame patch edgecolor facecolor=None, # frame patch facecolor bbox_to_anchor=None, # bbox to which the legend will be anchored bbox_transform=None, # transform for the bbox frameon=None, # draw frame handler_map=None, title_fontproperties=None, # properties for the legend title ): """ Parameters ---------- parent : `~matplotlib.axes.Axes` or `.Figure` The artist that contains the legend. handles : list of `.Artist` A list of Artists (lines, patches) to be added to the legend. labels : list of str A list of labels to show next to the artists. The length of handles and labels should be the same. If they are not, they are truncated to the smaller of both lengths. Other Parameters ---------------- %(_legend_kw_doc)s Notes ----- Users can specify any arbitrary location for the legend using the *bbox_to_anchor* keyword argument. *bbox_to_anchor* can be a `.BboxBase` (or derived there from) or a tuple of 2 or 4 floats. See `set_bbox_to_anchor` for more detail. The legend location can be specified by setting *loc* with a tuple of 2 floats, which is interpreted as the lower-left corner of the legend in the normalized axes coordinate. """ # local import only to avoid circularity from matplotlib.axes import Axes from matplotlib.figure import FigureBase super().__init__() if prop is None: if fontsize is not None: self.prop = FontProperties(size=fontsize) else: self.prop = FontProperties( size=mpl.rcParams["legend.fontsize"]) else: self.prop = FontProperties._from_any(prop) if isinstance(prop, dict) and "size" not in prop: self.prop.set_size(mpl.rcParams["legend.fontsize"]) self._fontsize = self.prop.get_size_in_points() self.texts = [] self.legendHandles = [] self._legend_title_box = None #: A dictionary with the extra handler mappings for this Legend #: instance. self._custom_handler_map = handler_map def val_or_rc(val, rc_name): return val if val is not None else mpl.rcParams[rc_name] self.numpoints = val_or_rc(numpoints, 'legend.numpoints') self.markerscale = val_or_rc(markerscale, 'legend.markerscale') self.scatterpoints = val_or_rc(scatterpoints, 'legend.scatterpoints') self.borderpad = val_or_rc(borderpad, 'legend.borderpad') self.labelspacing = val_or_rc(labelspacing, 'legend.labelspacing') self.handlelength = val_or_rc(handlelength, 'legend.handlelength') self.handleheight = val_or_rc(handleheight, 'legend.handleheight') self.handletextpad = val_or_rc(handletextpad, 'legend.handletextpad') self.borderaxespad = val_or_rc(borderaxespad, 'legend.borderaxespad') self.columnspacing = val_or_rc(columnspacing, 'legend.columnspacing') self.shadow = val_or_rc(shadow, 'legend.shadow') # trim handles and labels if illegal label... _lab, _hand = [], [] for label, handle in zip(labels, handles): if isinstance(label, str) and label.startswith('_'): _api.warn_external(f"The label {label!r} of {handle!r} starts " "with '_'. It is thus excluded from the " "legend.") else: _lab.append(label) _hand.append(handle) labels, handles = _lab, _hand handles = list(handles) if len(handles) < 2: ncol = 1 self._ncol = ncol if self.numpoints <= 0: raise ValueError("numpoints must be > 0; it was %d" % numpoints) # introduce y-offset for handles of the scatter plot if scatteryoffsets is None: self._scatteryoffsets = np.array([3. / 8., 4. / 8., 2.5 / 8.]) else: self._scatteryoffsets = np.asarray(scatteryoffsets) reps = self.scatterpoints // len(self._scatteryoffsets) + 1 self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[:self.scatterpoints] # _legend_box is a VPacker instance that contains all # legend items and will be initialized from _init_legend_box() # method. self._legend_box = None if isinstance(parent, Axes): self.isaxes = True self.axes = parent self.set_figure(parent.figure) elif isinstance(parent, FigureBase): self.isaxes = False self.set_figure(parent) else: raise TypeError("Legend needs either Axes or FigureBase as parent") self.parent = parent self._loc_used_default = loc is None if loc is None: loc = mpl.rcParams["legend.loc"] if not self.isaxes and loc in [0, 'best']: loc = 'upper right' if isinstance(loc, str): loc = _api.check_getitem(self.codes, loc=loc) if not self.isaxes and loc == 0: raise ValueError( "Automatic legend placement (loc='best') not implemented for " "figure legend") self._mode = mode self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform) # We use FancyBboxPatch to draw a legend frame. The location # and size of the box will be updated during the drawing time. if facecolor is None: facecolor = mpl.rcParams["legend.facecolor"] if facecolor == 'inherit': facecolor = mpl.rcParams["axes.facecolor"] if edgecolor is None: edgecolor = mpl.rcParams["legend.edgecolor"] if edgecolor == 'inherit': edgecolor = mpl.rcParams["axes.edgecolor"] if fancybox is None: fancybox = mpl.rcParams["legend.fancybox"] self.legendPatch = FancyBboxPatch( xy=(0, 0), width=1, height=1, facecolor=facecolor, edgecolor=edgecolor, # If shadow is used, default to alpha=1 (#8943). alpha=(framealpha if framealpha is not None else 1 if shadow else mpl.rcParams["legend.framealpha"]), # The width and height of the legendPatch will be set (in draw()) # to the length that includes the padding. Thus we set pad=0 here. boxstyle=("round,pad=0,rounding_size=0.2" if fancybox else "square,pad=0"), mutation_scale=self._fontsize, snap=True, visible=(frameon if frameon is not None else mpl.rcParams["legend.frameon"])) self._set_artist_props(self.legendPatch) # init with null renderer self._init_legend_box(handles, labels, markerfirst) tmp = self._loc_used_default self._set_loc(loc) self._loc_used_default = tmp # ignore changes done by _set_loc # figure out title font properties: if title_fontsize is not None and title_fontproperties is not None: raise ValueError( "title_fontsize and title_fontproperties can't be specified " "at the same time. Only use one of them. ") title_prop_fp = FontProperties._from_any(title_fontproperties) if isinstance(title_fontproperties, dict): if "size" not in title_fontproperties: title_fontsize = mpl.rcParams["legend.title_fontsize"] title_prop_fp.set_size(title_fontsize) elif title_fontsize is not None: title_prop_fp.set_size(title_fontsize) elif not isinstance(title_fontproperties, FontProperties): title_fontsize = mpl.rcParams["legend.title_fontsize"] title_prop_fp.set_size(title_fontsize) self.set_title(title, prop=title_prop_fp) self._draggable = None # set the text color color_getters = { # getter function depends on line or patch 'linecolor': ['get_color', 'get_facecolor'], 'markerfacecolor': ['get_markerfacecolor', 'get_facecolor'], 'mfc': ['get_markerfacecolor', 'get_facecolor'], 'markeredgecolor': ['get_markeredgecolor', 'get_edgecolor'], 'mec': ['get_markeredgecolor', 'get_edgecolor'], } if labelcolor is None: if mpl.rcParams['legend.labelcolor'] is not None: labelcolor = mpl.rcParams['legend.labelcolor'] else: labelcolor = mpl.rcParams['text.color'] if isinstance(labelcolor, str) and labelcolor in color_getters: getter_names = color_getters[labelcolor] for handle, text in zip(self.legendHandles, self.texts): for getter_name in getter_names: try: color = getattr(handle, getter_name)() text.set_color(color) break except AttributeError: pass elif isinstance(labelcolor, str) and labelcolor == 'none': for text in self.texts: text.set_color(labelcolor) elif np.iterable(labelcolor): for text, color in zip( self.texts, itertools.cycle(colors.to_rgba_array(labelcolor))): text.set_color(color) else: raise ValueError("Invalid argument for labelcolor : %s" % str(labelcolor)) def _set_artist_props(self, a): """ Set the boilerplate props for artists added to axes. """ a.set_figure(self.figure) if self.isaxes: # a.set_axes(self.axes) a.axes = self.axes a.set_transform(self.get_transform()) def _set_loc(self, loc): # find_offset function will be provided to _legend_box and # _legend_box will draw itself at the location of the return # value of the find_offset. self._loc_used_default = False self._loc_real = loc self.stale = True self._legend_box.set_offset(self._findoffset) def _get_loc(self): return self._loc_real _loc = property(_get_loc, _set_loc) def _findoffset(self, width, height, xdescent, ydescent, renderer): """Helper function to locate the legend.""" if self._loc == 0: # "best". x, y = self._find_best_position(width, height, renderer) elif self._loc in Legend.codes.values(): # Fixed location. bbox = Bbox.from_bounds(0, 0, width, height) x, y = self._get_anchored_bbox(self._loc, bbox, self.get_bbox_to_anchor(), renderer) else: # Axes or figure coordinates. fx, fy = self._loc bbox = self.get_bbox_to_anchor() x, y = bbox.x0 + bbox.width * fx, bbox.y0 + bbox.height * fy return x + xdescent, y + ydescent @allow_rasterization def draw(self, renderer): # docstring inherited if not self.get_visible(): return renderer.open_group('legend', gid=self.get_gid()) fontsize = renderer.points_to_pixels(self._fontsize) # if mode == fill, set the width of the legend_box to the # width of the parent (minus pads) if self._mode in ["expand"]: pad = 2 * (self.borderaxespad + self.borderpad) * fontsize self._legend_box.set_width(self.get_bbox_to_anchor().width - pad) # update the location and size of the legend. This needs to # be done in any case to clip the figure right. bbox = self._legend_box.get_window_extent(renderer) self.legendPatch.set_bounds(bbox.bounds) self.legendPatch.set_mutation_scale(fontsize) if self.shadow: Shadow(self.legendPatch, 2, -2).draw(renderer) self.legendPatch.draw(renderer) self._legend_box.draw(renderer) renderer.close_group('legend') self.stale = False # _default_handler_map defines the default mapping between plot # elements and the legend handlers. _default_handler_map = { StemContainer: legend_handler.HandlerStem(), ErrorbarContainer: legend_handler.HandlerErrorbar(), Line2D: legend_handler.HandlerLine2D(), Patch: legend_handler.HandlerPatch(), StepPatch: legend_handler.HandlerStepPatch(), LineCollection: legend_handler.HandlerLineCollection(), RegularPolyCollection: legend_handler.HandlerRegularPolyCollection(), CircleCollection: legend_handler.HandlerCircleCollection(), BarContainer: legend_handler.HandlerPatch( update_func=legend_handler.update_from_first_child), tuple: legend_handler.HandlerTuple(), PathCollection: legend_handler.HandlerPathCollection(), PolyCollection: legend_handler.HandlerPolyCollection() } # (get|set|update)_default_handler_maps are public interfaces to # modify the default handler map. @classmethod def get_default_handler_map(cls): """Return the global default handler map, shared by all legends.""" return cls._default_handler_map @classmethod def set_default_handler_map(cls, handler_map): """Set the global default handler map, shared by all legends.""" cls._default_handler_map = handler_map @classmethod def update_default_handler_map(cls, handler_map): """Update the global default handler map, shared by all legends.""" cls._default_handler_map.update(handler_map) def get_legend_handler_map(self): """Return this legend instance's handler map.""" default_handler_map = self.get_default_handler_map() return ({ **default_handler_map, **self._custom_handler_map } if self._custom_handler_map else default_handler_map) @staticmethod def get_legend_handler(legend_handler_map, orig_handle): """ Return a legend handler from *legend_handler_map* that corresponds to *orig_handler*. *legend_handler_map* should be a dictionary object (that is returned by the get_legend_handler_map method). It first checks if the *orig_handle* itself is a key in the *legend_handler_map* and return the associated value. Otherwise, it checks for each of the classes in its method-resolution-order. If no matching key is found, it returns ``None``. """ try: return legend_handler_map[orig_handle] except (TypeError, KeyError): # TypeError if unhashable. pass for handle_type in type(orig_handle).mro(): try: return legend_handler_map[handle_type] except KeyError: pass return None def _init_legend_box(self, handles, labels, markerfirst=True): """ Initialize the legend_box. The legend_box is an instance of the OffsetBox, which is packed with legend handles and texts. Once packed, their location is calculated during the drawing time. """ fontsize = self._fontsize # legend_box is a HPacker, horizontally packed with columns. # Each column is a VPacker, vertically packed with legend items. # Each legend item is a HPacker packed with: # - handlebox: a DrawingArea which contains the legend handle. # - labelbox: a TextArea which contains the legend text. text_list = [] # the list of text instances handle_list = [] # the list of handle instances handles_and_labels = [] # The approximate height and descent of text. These values are # only used for plotting the legend handle. descent = 0.35 * fontsize * (self.handleheight - 0.7) # heuristic. height = fontsize * self.handleheight - descent # each handle needs to be drawn inside a box of (x, y, w, h) = # (0, -descent, width, height). And their coordinates should # be given in the display coordinates. # The transformation of each handle will be automatically set # to self.get_transform(). If the artist does not use its # default transform (e.g., Collections), you need to # manually set their transform to the self.get_transform(). legend_handler_map = self.get_legend_handler_map() for orig_handle, label in zip(handles, labels): handler = self.get_legend_handler(legend_handler_map, orig_handle) if handler is None: _api.warn_external( "Legend does not support handles for {0} " "instances.\nA proxy artist may be used " "instead.\nSee: https://matplotlib.org/" "stable/tutorials/intermediate/legend_guide.html" "#controlling-the-legend-entries".format( type(orig_handle).__name__)) # No handle for this artist, so we just defer to None. handle_list.append(None) else: textbox = TextArea(label, multilinebaseline=True, textprops=dict(verticalalignment='baseline', horizontalalignment='left', fontproperties=self.prop)) handlebox = DrawingArea(width=self.handlelength * fontsize, height=height, xdescent=0., ydescent=descent) text_list.append(textbox._text) # Create the artist for the legend which represents the # original artist/handle. handle_list.append( handler.legend_artist(self, orig_handle, fontsize, handlebox)) handles_and_labels.append((handlebox, textbox)) columnbox = [] # array_split splits n handles_and_labels into ncol columns, with the # first n%ncol columns having an extra entry. filter(len, ...) handles # the case where n < ncol: the last ncol-n columns are empty and get # filtered out. for handles_and_labels_column \ in filter(len, np.array_split(handles_and_labels, self._ncol)): # pack handlebox and labelbox into itembox itemboxes = [ HPacker(pad=0, sep=self.handletextpad * fontsize, children=[h, t] if markerfirst else [t, h], align="baseline") for h, t in handles_and_labels_column ] # pack columnbox alignment = "baseline" if markerfirst else "right" columnbox.append( VPacker(pad=0, sep=self.labelspacing * fontsize, align=alignment, children=itemboxes)) mode = "expand" if self._mode == "expand" else "fixed" sep = self.columnspacing * fontsize self._legend_handle_box = HPacker(pad=0, sep=sep, align="baseline", mode=mode, children=columnbox) self._legend_title_box = TextArea("") self._legend_box = VPacker( pad=self.borderpad * fontsize, sep=self.labelspacing * fontsize, align="center", children=[self._legend_title_box, self._legend_handle_box]) self._legend_box.set_figure(self.figure) self._legend_box.axes = self.axes self.texts = text_list self.legendHandles = handle_list def _auto_legend_data(self): """ Return display coordinates for hit testing for "best" positioning. Returns ------- bboxes List of bounding boxes of all patches. lines List of `.Path` corresponding to each line. offsets List of (x, y) offsets of all collection. """ assert self.isaxes # always holds, as this is only called internally bboxes = [] lines = [] offsets = [] for artist in self.parent._children: if isinstance(artist, Line2D): lines.append(artist.get_transform().transform_path( artist.get_path())) elif isinstance(artist, Rectangle): bboxes.append(artist.get_bbox().transformed( artist.get_data_transform())) elif isinstance(artist, Patch): bboxes.append(artist.get_path().get_extents( artist.get_transform())) elif isinstance(artist, Collection): _, offset_trf, hoffsets, _ = artist._prepare_points() for offset in offset_trf.transform(hoffsets): offsets.append(offset) return bboxes, lines, offsets def get_children(self): # docstring inherited return [self._legend_box, self.get_frame()] def get_frame(self): """Return the `~.patches.Rectangle` used to frame the legend.""" return self.legendPatch def get_lines(self): r"""Return the list of `~.lines.Line2D`\s in the legend.""" return [h for h in self.legendHandles if isinstance(h, Line2D)] def get_patches(self): r"""Return the list of `~.patches.Patch`\s in the legend.""" return silent_list( 'Patch', [h for h in self.legendHandles if isinstance(h, Patch)]) def get_texts(self): r"""Return the list of `~.text.Text`\s in the legend.""" return silent_list('Text', self.texts) def set_title(self, title, prop=None): """ Set the legend title. Fontproperties can be optionally set with *prop* parameter. """ self._legend_title_box._text.set_text(title) if title: self._legend_title_box._text.set_visible(True) self._legend_title_box.set_visible(True) else: self._legend_title_box._text.set_visible(False) self._legend_title_box.set_visible(False) if prop is not None: self._legend_title_box._text.set_fontproperties(prop) self.stale = True def get_title(self): """Return the `.Text` instance for the legend title.""" return self._legend_title_box._text def get_window_extent(self, renderer=None): # docstring inherited if renderer is None: renderer = self.figure._cachedRenderer return self._legend_box.get_window_extent(renderer=renderer) def get_tightbbox(self, renderer): # docstring inherited return self._legend_box.get_window_extent(renderer) def get_frame_on(self): """Get whether the legend box patch is drawn.""" return self.legendPatch.get_visible() def set_frame_on(self, b): """ Set whether the legend box patch is drawn. Parameters ---------- b : bool """ self.legendPatch.set_visible(b) self.stale = True draw_frame = set_frame_on # Backcompat alias. def get_bbox_to_anchor(self): """Return the bbox that the legend will be anchored to.""" if self._bbox_to_anchor is None: return self.parent.bbox else: return self._bbox_to_anchor def set_bbox_to_anchor(self, bbox, transform=None): """ Set the bbox that the legend will be anchored to. Parameters ---------- bbox : `~matplotlib.transforms.BboxBase` or tuple The bounding box can be specified in the following ways: - A `.BboxBase` instance - A tuple of ``(left, bottom, width, height)`` in the given transform (normalized axes coordinate if None) - A tuple of ``(left, bottom)`` where the width and height will be assumed to be zero. - *None*, to remove the bbox anchoring, and use the parent bbox. transform : `~matplotlib.transforms.Transform`, optional A transform to apply to the bounding box. If not specified, this will use a transform to the bounding box of the parent. """ if bbox is None: self._bbox_to_anchor = None return elif isinstance(bbox, BboxBase): self._bbox_to_anchor = bbox else: try: l = len(bbox) except TypeError as err: raise ValueError("Invalid argument for bbox : %s" % str(bbox)) from err if l == 2: bbox = [bbox[0], bbox[1], 0, 0] self._bbox_to_anchor = Bbox.from_bounds(*bbox) if transform is None: transform = BboxTransformTo(self.parent.bbox) self._bbox_to_anchor = TransformedBbox(self._bbox_to_anchor, transform) self.stale = True def _get_anchored_bbox(self, loc, bbox, parentbbox, renderer): """ Place the *bbox* inside the *parentbbox* according to a given location code. Return the (x, y) coordinate of the bbox. Parameters ---------- loc : int A location code in range(1, 11). This corresponds to the possible values for ``self._loc``, excluding "best". bbox : `~matplotlib.transforms.Bbox` bbox to be placed, in display coordinates. parentbbox : `~matplotlib.transforms.Bbox` A parent box which will contain the bbox, in display coordinates. """ return offsetbox._get_anchored_bbox( loc, bbox, parentbbox, self.borderaxespad * renderer.points_to_pixels(self._fontsize)) def _find_best_position(self, width, height, renderer, consider=None): """ Determine the best location to place the legend. *consider* is a list of ``(x, y)`` pairs to consider as a potential lower-left corner of the legend. All are display coords. """ assert self.isaxes # always holds, as this is only called internally start_time = time.perf_counter() bboxes, lines, offsets = self._auto_legend_data() bbox = Bbox.from_bounds(0, 0, width, height) if consider is None: consider = [ self._get_anchored_bbox(x, bbox, self.get_bbox_to_anchor(), renderer) for x in range(1, len(self.codes)) ] candidates = [] for idx, (l, b) in enumerate(consider): legendBox = Bbox.from_bounds(l, b, width, height) badness = 0 # XXX TODO: If markers are present, it would be good to take them # into account when checking vertex overlaps in the next line. badness = ( sum(legendBox.count_contains(line.vertices) for line in lines) + legendBox.count_contains(offsets) + legendBox.count_overlaps(bboxes) + sum( line.intersects_bbox(legendBox, filled=False) for line in lines)) if badness == 0: return l, b # Include the index to favor lower codes in case of a tie. candidates.append((badness, idx, (l, b))) _, _, (l, b) = min(candidates) if self._loc_used_default and time.perf_counter() - start_time > 1: _api.warn_external( 'Creating legend with loc="best" can be slow with large ' 'amounts of data.') return l, b def contains(self, event): inside, info = self._default_contains(event) if inside is not None: return inside, info return self.legendPatch.contains(event) def set_draggable(self, state, use_blit=False, update='loc'): """ Enable or disable mouse dragging support of the legend. Parameters ---------- state : bool Whether mouse dragging is enabled. use_blit : bool, optional Use blitting for faster image composition. For details see :ref:`func-animation`. update : {'loc', 'bbox'}, optional The legend parameter to be changed when dragged: - 'loc': update the *loc* parameter of the legend - 'bbox': update the *bbox_to_anchor* parameter of the legend Returns ------- `.DraggableLegend` or *None* If *state* is ``True`` this returns the `.DraggableLegend` helper instance. Otherwise this returns *None*. """ if state: if self._draggable is None: self._draggable = DraggableLegend(self, use_blit, update=update) else: if self._draggable is not None: self._draggable.disconnect() self._draggable = None return self._draggable def get_draggable(self): """Return ``True`` if the legend is draggable, ``False`` otherwise.""" return self._draggable is not None
class Legend(Artist): """ Place a legend on the axes at location loc. Labels are a sequence of strings and loc can be a string or an integer specifying the legend location The location codes are:: 'best' : 0, (only implemented for axis legends) 'upper right' : 1, 'upper left' : 2, 'lower left' : 3, 'lower right' : 4, 'right' : 5, 'center left' : 6, 'center right' : 7, 'lower center' : 8, 'upper center' : 9, 'center' : 10, loc can be a tuple of the noramilzed coordinate values with respect its parent. """ codes = { 'best': 0, # only implemented for axis legends 'upper right': 1, 'upper left': 2, 'lower left': 3, 'lower right': 4, 'right': 5, 'center left': 6, 'center right': 7, 'lower center': 8, 'upper center': 9, 'center': 10, } zorder = 5 def __str__(self): return "Legend" def __init__( self, parent, handles, labels, loc=None, numpoints=None, # the number of points in the legend line markerscale=None, # the relative size of legend markers vs. original scatterpoints=3, # TODO: may be an rcParam scatteryoffsets=None, prop=None, # properties for the legend texts # the following dimensions are in axes coords pad=None, # deprecated; use borderpad labelsep=None, # deprecated; use labelspacing handlelen=None, # deprecated; use handlelength handletextsep=None, # deprecated; use handletextpad axespad=None, # deprecated; use borderaxespad # spacing & pad defined as a fraction of the font-size borderpad=None, # the whitespace inside the legend border labelspacing=None, #the vertical space between the legend entries handlelength=None, # the length of the legend handles handletextpad=None, # the pad between the legend handle and text borderaxespad=None, # the pad between the axes and legend border columnspacing=None, # spacing between columns ncol=1, # number of columns mode=None, # mode for horizontal distribution of columns. None, "expand" fancybox=None, # True use a fancy box, false use a rounded box, none use rc shadow=None, title=None, # set a title for the legend bbox_to_anchor=None, # bbox that the legend will be anchored. bbox_transform=None, # transform for the bbox frameon=True, # draw frame handler_map=None, ): """ - *parent* : the artist that contains the legend - *handles* : a list of artists (lines, patches) to be added to the legend - *labels* : a list of strings to label the legend Optional keyword arguments: ================ ================================================================== Keyword Description ================ ================================================================== loc a location code prop the font property markerscale the relative size of legend markers vs. original numpoints the number of points in the legend for line scatterpoints the number of points in the legend for scatter plot scatteryoffsets a list of yoffsets for scatter symbols in legend frameon if True, draw a frame (default is True) fancybox if True, draw a frame with a round fancybox. If None, use rc shadow if True, draw a shadow behind legend ncol number of columns borderpad the fractional whitespace inside the legend border labelspacing the vertical space between the legend entries handlelength the length of the legend handles handletextpad the pad between the legend handle and text borderaxespad the pad between the axes and legend border columnspacing the spacing between columns title the legend title bbox_to_anchor the bbox that the legend will be anchored. bbox_transform the transform for the bbox. transAxes if None. ================ ================================================================== The pad and spacing parameters are measured in font-size units. E.g., a fontsize of 10 points and a handlelength=5 implies a handlelength of 50 points. Values from rcParams will be used if None. Users can specify any arbitrary location for the legend using the *bbox_to_anchor* keyword argument. bbox_to_anchor can be an instance of BboxBase(or its derivatives) or a tuple of 2 or 4 floats. See :meth:`set_bbox_to_anchor` for more detail. The legend location can be specified by setting *loc* with a tuple of 2 floats, which is interpreted as the lower-left corner of the legend in the normalized axes coordinate. """ from matplotlib.axes import Axes # local import only to avoid circularity from matplotlib.figure import Figure # local import only to avoid circularity Artist.__init__(self) if prop is None: self.prop = FontProperties(size=rcParams["legend.fontsize"]) elif isinstance(prop, dict): self.prop = FontProperties(**prop) if "size" not in prop: self.prop.set_size(rcParams["legend.fontsize"]) else: self.prop = prop self._fontsize = self.prop.get_size_in_points() propnames = [ 'numpoints', 'markerscale', 'shadow', "columnspacing", "scatterpoints" ] self.texts = [] self.legendHandles = [] self._legend_title_box = None self._handler_map = handler_map localdict = locals() for name in propnames: if localdict[name] is None: value = rcParams["legend." + name] else: value = localdict[name] setattr(self, name, value) # Take care the deprecated keywords deprecated_kwds = { "pad": "borderpad", "labelsep": "labelspacing", "handlelen": "handlelength", "handletextsep": "handletextpad", "axespad": "borderaxespad" } # convert values of deprecated keywords (ginve in axes coords) # to new vaules in a fraction of the font size # conversion factor bbox = parent.bbox axessize_fontsize = min(bbox.width, bbox.height) / self._fontsize for k, v in deprecated_kwds.items(): # use deprecated value if not None and if their newer # counter part is None. if localdict[k] is not None and localdict[v] is None: warnings.warn("Use '%s' instead of '%s'." % (v, k), DeprecationWarning) setattr(self, v, localdict[k] * axessize_fontsize) continue # Otherwise, use new keywords if localdict[v] is None: setattr(self, v, rcParams["legend." + v]) else: setattr(self, v, localdict[v]) del localdict handles = list(handles) if len(handles) < 2: ncol = 1 self._ncol = ncol if self.numpoints <= 0: raise ValueError("numpoints must be >= 0; it was %d" % numpoints) # introduce y-offset for handles of the scatter plot if scatteryoffsets is None: self._scatteryoffsets = np.array([3. / 8., 4. / 8., 2.5 / 8.]) else: self._scatteryoffsets = np.asarray(scatteryoffsets) reps = int(self.scatterpoints / len(self._scatteryoffsets)) + 1 self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[:self.scatterpoints] # _legend_box is an OffsetBox instance that contains all # legend items and will be initialized from _init_legend_box() # method. self._legend_box = None if isinstance(parent, Axes): self.isaxes = True self.set_axes(parent) self.set_figure(parent.figure) elif isinstance(parent, Figure): self.isaxes = False self.set_figure(parent) else: raise TypeError("Legend needs either Axes or Figure as parent") self.parent = parent if loc is None: loc = rcParams["legend.loc"] if not self.isaxes and loc in [0, 'best']: loc = 'upper right' if is_string_like(loc): if loc not in self.codes: if self.isaxes: warnings.warn( 'Unrecognized location "%s". Falling back on "best"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes.keys()))) loc = 0 else: warnings.warn( 'Unrecognized location "%s". Falling back on "upper right"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes.keys()))) loc = 1 else: loc = self.codes[loc] if not self.isaxes and loc == 0: warnings.warn( 'Automatic legend placement (loc="best") not implemented for figure legend. ' 'Falling back on "upper right".') loc = 1 self._mode = mode self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform) # We use FancyBboxPatch to draw a legend frame. The location # and size of the box will be updated during the drawing time. self.legendPatch = FancyBboxPatch(xy=(0.0, 0.0), width=1., height=1., facecolor=rcParams["axes.facecolor"], edgecolor=rcParams["axes.edgecolor"], mutation_scale=self._fontsize, snap=True) # The width and height of the legendPatch will be set (in the # draw()) to the length that includes the padding. Thus we set # pad=0 here. if fancybox is None: fancybox = rcParams["legend.fancybox"] if fancybox == True: self.legendPatch.set_boxstyle("round", pad=0, rounding_size=0.2) else: self.legendPatch.set_boxstyle("square", pad=0) self._set_artist_props(self.legendPatch) self._drawFrame = frameon # init with null renderer self._init_legend_box(handles, labels) self._loc = loc self.set_title(title) self._last_fontsize_points = self._fontsize self._draggable = None def _set_artist_props(self, a): """ set the boilerplate props for artists added to axes """ a.set_figure(self.figure) if self.isaxes: a.set_axes(self.axes) a.set_transform(self.get_transform()) def _set_loc(self, loc): # find_offset function will be provided to _legend_box and # _legend_box will draw itself at the location of the return # value of the find_offset. self._loc_real = loc if loc == 0: _findoffset = self._findoffset_best else: _findoffset = self._findoffset_loc #def findoffset(width, height, xdescent, ydescent): # return _findoffset(width, height, xdescent, ydescent, renderer) self._legend_box.set_offset(_findoffset) self._loc_real = loc def _get_loc(self): return self._loc_real _loc = property(_get_loc, _set_loc) def _findoffset_best(self, width, height, xdescent, ydescent, renderer): "Helper function to locate the legend at its best position" ox, oy = self._find_best_position(width, height, renderer) return ox + xdescent, oy + ydescent def _findoffset_loc(self, width, height, xdescent, ydescent, renderer): "Heper function to locate the legend using the location code" if iterable(self._loc) and len(self._loc) == 2: # when loc is a tuple of axes(or figure) coordinates. fx, fy = self._loc bbox = self.get_bbox_to_anchor() x, y = bbox.x0 + bbox.width * fx, bbox.y0 + bbox.height * fy else: bbox = Bbox.from_bounds(0, 0, width, height) x, y = self._get_anchored_bbox(self._loc, bbox, self.get_bbox_to_anchor(), renderer) return x + xdescent, y + ydescent @allow_rasterization def draw(self, renderer): "Draw everything that belongs to the legend" if not self.get_visible(): return renderer.open_group('legend') fontsize = renderer.points_to_pixels(self._fontsize) # if mode == fill, set the width of the legend_box to the # width of the paret (minus pads) if self._mode in ["expand"]: pad = 2 * (self.borderaxespad + self.borderpad) * fontsize self._legend_box.set_width(self.get_bbox_to_anchor().width - pad) if self._drawFrame: # update the location and size of the legend bbox = self._legend_box.get_window_extent(renderer) self.legendPatch.set_bounds(bbox.x0, bbox.y0, bbox.width, bbox.height) self.legendPatch.set_mutation_scale(fontsize) if self.shadow: shadow = Shadow(self.legendPatch, 2, -2) shadow.draw(renderer) self.legendPatch.draw(renderer) self._legend_box.draw(renderer) renderer.close_group('legend') def _approx_text_height(self, renderer=None): """ Return the approximate height of the text. This is used to place the legend handle. """ if renderer is None: return self._fontsize else: return renderer.points_to_pixels(self._fontsize) # _default_handler_map defines the default mapping between plot # elements and the legend handlers. _default_handler_map = { ErrorbarContainer: legend_handler.HandlerErrorbar(), Line2D: legend_handler.HandlerLine2D(), Patch: legend_handler.HandlerPatch(), LineCollection: legend_handler.HandlerLineCollection(), RegularPolyCollection: legend_handler.HandlerRegularPolyCollection(), CircleCollection: legend_handler.HandlerCircleCollection(), BarContainer: legend_handler.HandlerPatch( update_func=legend_handler.update_from_first_child), tuple: legend_handler.HandlerTuple(), } # (get|set|update)_default_handler_maps are public interfaces to # modify the defalut handler map. @classmethod def get_default_handler_map(cls): """ A class method that returns the default handler map. """ return cls._default_handler_map @classmethod def set_default_handler_map(cls, handler_map): """ A class method to set the default handler map. """ cls._default_handler_map = handler_map @classmethod def update_default_handler_map(cls, handler_map): """ A class method to update the default handler map. """ cls._default_handler_map.update(handler_map) def get_legend_handler_map(self): """ return the handler map. """ default_handler_map = self.get_default_handler_map() if self._handler_map: hm = default_handler_map.copy() hm.update(self._handler_map) return hm else: return default_handler_map @staticmethod def get_legend_handler(legend_handler_map, orig_handle): """ return a legend handler from *legend_handler_map* that corresponds to *orig_handler*. *legend_handler_map* should be a dictionary object (that is returned by the get_legend_handler_map method). It first checks if the *orig_handle* itself is a key in the *legend_hanler_map* and return the associated value. Otherwised, it checks for each of the classes in its method-resolution-order. If no matching key is found, it returns None. """ legend_handler_keys = legend_handler_map.keys() if orig_handle in legend_handler_keys: handler = legend_handler_map[orig_handle] else: for handle_type in type(orig_handle).mro(): if handle_type in legend_handler_map: handler = legend_handler_map[handle_type] break else: handler = None return handler def _init_legend_box(self, handles, labels): """ Initiallize the legend_box. The legend_box is an instance of the OffsetBox, which is packed with legend handles and texts. Once packed, their location is calculated during the drawing time. """ fontsize = self._fontsize # legend_box is a HPacker, horizontally packed with # columns. Each column is a VPacker, vertically packed with # legend items. Each legend item is HPacker packed with # legend handleBox and labelBox. handleBox is an instance of # offsetbox.DrawingArea which contains legend handle. labelBox # is an instance of offsetbox.TextArea which contains legend # text. text_list = [] # the list of text instances handle_list = [] # the list of text instances label_prop = dict( verticalalignment='baseline', horizontalalignment='left', fontproperties=self.prop, ) labelboxes = [] handleboxes = [] # The approximate height and descent of text. These values are # only used for plotting the legend handle. height = self._approx_text_height() * 0.7 descent = 0. # each handle needs to be drawn inside a box of (x, y, w, h) = # (0, -descent, width, height). And their corrdinates should # be given in the display coordinates. # The transformation of each handle will be automatically set # to self.get_trasnform(). If the artist does not uses its # default trasnform (eg, Collections), you need to # manually set their transform to the self.get_transform(). legend_handler_map = self.get_legend_handler_map() for orig_handle, lab in zip(handles, labels): handler = self.get_legend_handler(legend_handler_map, orig_handle) if handler is None: warnings.warn( "Legend does not support %s\nUse proxy artist instead.\n\nhttp://matplotlib.sourceforge.net/users/legend_guide.html#using-proxy-artist\n" % (str(orig_handle), )) handle_list.append(None) continue textbox = TextArea(lab, textprops=label_prop, multilinebaseline=True, minimumdescent=True) text_list.append(textbox._text) labelboxes.append(textbox) handlebox = DrawingArea(width=self.handlelength * fontsize, height=height, xdescent=0., ydescent=descent) handle = handler(self, orig_handle, \ #xdescent, ydescent, width, height, fontsize, handlebox) handle_list.append(handle) handleboxes.append(handlebox) if len(handleboxes) > 0: # We calculate number of lows in each column. The first # (num_largecol) columns will have (nrows+1) rows, and remaing # (num_smallcol) columns will have (nrows) rows. ncol = min(self._ncol, len(handleboxes)) nrows, num_largecol = divmod(len(handleboxes), ncol) num_smallcol = ncol - num_largecol # starting index of each column and number of rows in it. largecol = safezip( range(0, num_largecol * (nrows + 1), (nrows + 1)), [nrows + 1] * num_largecol) smallcol = safezip( range(num_largecol * (nrows + 1), len(handleboxes), nrows), [nrows] * num_smallcol) else: largecol, smallcol = [], [] handle_label = safezip(handleboxes, labelboxes) columnbox = [] for i0, di in largecol + smallcol: # pack handleBox and labelBox into itemBox itemBoxes = [ HPacker(pad=0, sep=self.handletextpad * fontsize, children=[h, t], align="baseline") for h, t in handle_label[i0:i0 + di] ] # minimumdescent=False for the text of the last row of the column itemBoxes[-1].get_children()[1].set_minimumdescent(False) # pack columnBox columnbox.append( VPacker(pad=0, sep=self.labelspacing * fontsize, align="baseline", children=itemBoxes)) if self._mode == "expand": mode = "expand" else: mode = "fixed" sep = self.columnspacing * fontsize self._legend_handle_box = HPacker(pad=0, sep=sep, align="baseline", mode=mode, children=columnbox) self._legend_title_box = TextArea("") self._legend_box = VPacker( pad=self.borderpad * fontsize, sep=self.labelspacing * fontsize, align="center", children=[self._legend_title_box, self._legend_handle_box]) self._legend_box.set_figure(self.figure) self.texts = text_list self.legendHandles = handle_list def _auto_legend_data(self): """ Returns list of vertices and extents covered by the plot. Returns a two long list. First element is a list of (x, y) vertices (in display-coordinates) covered by all the lines and line collections, in the legend's handles. Second element is a list of bounding boxes for all the patches in the legend's handles. """ assert self.isaxes # should always hold because function is only called internally ax = self.parent vertices = [] bboxes = [] lines = [] for handle in ax.lines: assert isinstance(handle, Line2D) path = handle.get_path() trans = handle.get_transform() tpath = trans.transform_path(path) lines.append(tpath) for handle in ax.patches: assert isinstance(handle, Patch) if isinstance(handle, Rectangle): transform = handle.get_data_transform() bboxes.append(handle.get_bbox().transformed(transform)) else: transform = handle.get_transform() bboxes.append(handle.get_path().get_extents(transform)) return [vertices, bboxes, lines] def draw_frame(self, b): 'b is a boolean. Set draw frame to b' self.set_frame_on(b) def get_children(self): 'return a list of child artists' children = [] if self._legend_box: children.append(self._legend_box) children.extend(self.get_lines()) children.extend(self.get_patches()) children.extend(self.get_texts()) children.append(self.get_frame()) if self._legend_title_box: children.append(self.get_title()) return children def get_frame(self): 'return the Rectangle instance used to frame the legend' return self.legendPatch def get_lines(self): 'return a list of lines.Line2D instances in the legend' return [h for h in self.legendHandles if isinstance(h, Line2D)] def get_patches(self): 'return a list of patch instances in the legend' return silent_list( 'Patch', [h for h in self.legendHandles if isinstance(h, Patch)]) def get_texts(self): 'return a list of text.Text instance in the legend' return silent_list('Text', self.texts) def set_title(self, title): 'set the legend title' self._legend_title_box._text.set_text(title) if title: self._legend_title_box.set_visible(True) else: self._legend_title_box.set_visible(False) def get_title(self): 'return Text instance for the legend title' return self._legend_title_box._text def get_window_extent(self, *args, **kwargs): 'return a extent of the the legend' return self.legendPatch.get_window_extent(*args, **kwargs) def get_frame_on(self): """ Get whether the legend box patch is drawn """ return self._drawFrame def set_frame_on(self, b): """ Set whether the legend box patch is drawn ACCEPTS: [ *True* | *False* ] """ self._drawFrame = b def get_bbox_to_anchor(self): """ return the bbox that the legend will be anchored """ if self._bbox_to_anchor is None: return self.parent.bbox else: return self._bbox_to_anchor def set_bbox_to_anchor(self, bbox, transform=None): """ set the bbox that the legend will be anchored. *bbox* can be a BboxBase instance, a tuple of [left, bottom, width, height] in the given transform (normalized axes coordinate if None), or a tuple of [left, bottom] where the width and height will be assumed to be zero. """ if bbox is None: self._bbox_to_anchor = None return elif isinstance(bbox, BboxBase): self._bbox_to_anchor = bbox else: try: l = len(bbox) except TypeError: raise ValueError("Invalid argument for bbox : %s" % str(bbox)) if l == 2: bbox = [bbox[0], bbox[1], 0, 0] self._bbox_to_anchor = Bbox.from_bounds(*bbox) if transform is None: transform = BboxTransformTo(self.parent.bbox) self._bbox_to_anchor = TransformedBbox(self._bbox_to_anchor, transform) def _get_anchored_bbox(self, loc, bbox, parentbbox, renderer): """ Place the *bbox* inside the *parentbbox* according to a given location code. Return the (x,y) coordinate of the bbox. - loc: a location code in range(1, 11). This corresponds to the possible values for self._loc, excluding "best". - bbox: bbox to be placed, display coodinate units. - parentbbox: a parent box which will contain the bbox. In display coordinates. """ assert loc in range(1, 11) # called only internally BEST, UR, UL, LL, LR, R, CL, CR, LC, UC, C = range(11) anchor_coefs = { UR: "NE", UL: "NW", LL: "SW", LR: "SE", R: "E", CL: "W", CR: "E", LC: "S", UC: "N", C: "C" } c = anchor_coefs[loc] fontsize = renderer.points_to_pixels(self._fontsize) container = parentbbox.padded(-(self.borderaxespad) * fontsize) anchored_box = bbox.anchored(c, container=container) return anchored_box.x0, anchored_box.y0 def _find_best_position(self, width, height, renderer, consider=None): """ Determine the best location to place the legend. `consider` is a list of (x, y) pairs to consider as a potential lower-left corner of the legend. All are display coords. """ assert self.isaxes # should always hold because function is only called internally verts, bboxes, lines = self._auto_legend_data() bbox = Bbox.from_bounds(0, 0, width, height) consider = [ self._get_anchored_bbox(x, bbox, self.get_bbox_to_anchor(), renderer) for x in range(1, len(self.codes)) ] #tx, ty = self.legendPatch.get_x(), self.legendPatch.get_y() candidates = [] for l, b in consider: legendBox = Bbox.from_bounds(l, b, width, height) badness = 0 badness = legendBox.count_contains(verts) badness += legendBox.count_overlaps(bboxes) for line in lines: if line.intersects_bbox(legendBox): badness += 1 ox, oy = l, b if badness == 0: return ox, oy candidates.append((badness, (l, b))) # rather than use min() or list.sort(), do this so that we are assured # that in the case of two equal badnesses, the one first considered is # returned. # NOTE: list.sort() is stable.But leave as it is for now. -JJL minCandidate = candidates[0] for candidate in candidates: if candidate[0] < minCandidate[0]: minCandidate = candidate ox, oy = minCandidate[1] return ox, oy def draggable(self, state=None, use_blit=False, update="loc"): """ Set the draggable state -- if state is * None : toggle the current state * True : turn draggable on * False : turn draggable off If draggable is on, you can drag the legend on the canvas with the mouse. The DraggableLegend helper instance is returned if draggable is on. The update parameter control which parameter of the legend changes when dragged. If update is "loc", the *loc* paramter of the legend is changed. If "bbox", the *bbox_to_anchor* parameter is changed. """ is_draggable = self._draggable is not None # if state is None we'll toggle if state is None: state = not is_draggable if state: if self._draggable is None: self._draggable = DraggableLegend(self, use_blit, update=update) else: if self._draggable is not None: self._draggable.disconnect() self._draggable = None return self._draggable
detector = contour.MultiScaleStructuredForest() detector.load(os.path.join(dir_path, '..', 'data', 'sf.dat')) for im in argv[1:]: t0 = time() s = segmentation.geodesicKMeans(imgproc.imread(im), detector, 1000) t1 = time() b = prop.propose(s) t2 = time() # If you just want the boxes use boxes = s.maskToBox(b) print("Generated %d proposals in %0.2fs (OverSeg: %0.2fs, Prop: %0.2fs)" % (b.shape[0], t2 - t0, t1 - t0, t2 - t1)) figure() for i in range(min(20, b.shape[0])): im = np.array(s.image) im[b[i, s.s]] = (255, 0, 0) ax = subplot(4, 5, i + 1) ax.imshow(im) # Draw the bounding box from matplotlib.patches import FancyBboxPatch ax.add_patch( FancyBboxPatch((boxes[i, 0], boxes[i, 1]), boxes[i, 2] - boxes[i, 0], boxes[i, 3] - boxes[i, 1], boxstyle="square,pad=0.", ec="b", fc="none", lw=2)) show()
def __init__(self, offsetbox, xy, xybox=None, xycoords='data', boxcoords=None, frameon=True, pad=0.4, # BboxPatch annotation_clip=None, box_alignment=(0.5, 0.5), bboxprops=None, arrowprops=None, fontsize=None, **kwargs): """ *offsetbox* : OffsetBox instance *xycoords* : same as Annotation but can be a tuple of two strings which are interpreted as x and y coordinates. *boxcoords* : similar to textcoords as Annotation but can be a tuple of two strings which are interpreted as x and y coordinates. *box_alignment* : a tuple of two floats for a vertical and horizontal alignment of the offset box w.r.t. the *boxcoords*. The lower-left corner is (0.0) and upper-right corner is (1.1). other parameters are identical to that of Annotation. """ self.offsetbox = offsetbox self.arrowprops = arrowprops self.set_fontsize(fontsize) if arrowprops is not None: self._arrow_relpos = self.arrowprops.pop("relpos", (0.5, 0.5)) self.arrow_patch = FancyArrowPatch((0, 0), (1,1), **self.arrowprops) else: self._arrow_relpos = None self.arrow_patch = None _AnnotationBase.__init__(self, xy, xytext=xybox, xycoords=xycoords, textcoords=boxcoords, annotation_clip=annotation_clip) martist.Artist.__init__(self, **kwargs) #self._fw, self._fh = 0., 0. # for alignment self._box_alignment = box_alignment # frame self.patch = FancyBboxPatch( xy=(0.0, 0.0), width=1., height=1., facecolor='w', edgecolor='k', mutation_scale=self.prop.get_size_in_points(), snap=True ) self.patch.set_boxstyle("square",pad=pad) if bboxprops: self.patch.set(**bboxprops) self._drawFrame = frameon
def plot(self, ax, key, args, output): xval = output['time'] yval = output['value'] if self.cur_state is None: for line in self.lines: if line.startswith('graph'): items = line.split() width, height = map(float, items[2:4]) offs = 0.1 ax.set_xlim(0 - offs, width + offs) ax.set_ylim(0 - offs, height + offs) elif line.startswith('node'): items = line.split() name = items[1] xpos, ypos, width, height = map(float, items[2:6]) pad = 0.15 pos = xpos - width / 2 + pad, ypos - height / 2 + pad box = FancyBboxPatch( pos, width - 2 * pad, height - 2 * pad, boxstyle="round,pad={}".format(pad), fc='skyblue', #~ ec='royalblue' ec='navy') ax.add_patch(box) label = line.partition('<')[2].rpartition('>')[0] label = parser.feed(label) label = ax.text( xpos, ypos, label, ha='center', va='center', clip_on=True, color='navy', ) self.states[name] = box elif line.startswith('edge'): items = line.split() n1, n2 = items[1:3] num = int(items[3]) xdatas = list(map(float, items[4:4 + 2 * num:2])) ydatas = list(map(float, items[5:5 + 2 * num:2])) path = [(Path.MOVETO, (xdatas[0], ydatas[0]))] for x, y in zip(xdatas[1:], ydatas[1:]): point = Path.CURVE4, (x, y) path.append(point) cmds, verts = zip(*path) ppath = Path(np.array(verts), cmds) patch = PathPatch(ppath, fc='none', ec='navy') ax.add_patch(patch) dx = (xdatas[-1] - xdatas[-2]) * 0.05 dy = (ydatas[-1] - ydatas[-2]) * 0.05 arrow = ax.arrow(xdatas[-1], ydatas[-1], dx, dy, head_width=0.025, head_length=0.05, overhang=0.2, fc='navy', ec='navy') xypos = items[-4:-2] label = line.partition('<')[2].rpartition('>')[0] label = parser.feed(label) ax.annotate( label, xypos, ha='center', va='center', color='navy', ) event = label.partition('\n')[0] self.transitions[(n1, event)] = patch, arrow ax.set_title(self.name) ax.grid(False) ax.set_xticks([]) ax.set_yticks([]) for line in ax.spines.values(): line.set_visible(False) ax.set_aspect(1) self.cur_state = yval self.times[xval] = self.cur_state, self.hyfsm._last_transition self.keys[key] = xval, yval self.set_cursor(ax, xval, yval)
def _label(self, colors): """ Method to create an edges label :return: """ its_label = '{:g}'.format(abs(self.edge.capacity) if self.edge.capacity is not None else '') status_color = colors['bc'] if not self.layout_mode: in_percent = int(round(100 * self.edge.flux / float(self.edge.capacity), 0)) if in_percent < 100: status_color = colors['wc'] its_label += '; {:g}%'.format(in_percent) # get the label coordinates based on self.segments (use the end of the first segment if several segments are # present, if it is a single, place the label half the way s_coords, e_coords = self.segments[0] edge_vec = Vec(s_coords, e_coords) start_vec = Vec(s_coords) # maybe always use the last part... if len(self.segments) == 1: # single segment: place it half way part_way_vec = start_vec + self.label_position * edge_vec coords_label = part_way_vec.coords[:2] else: # several segments: place label at the end of the first #half_way_vec = start_vec + 0.5 * edge_vec #coords_label = half_way_vec.coords[:2] coords_label = e_coords if self.with_label: self.labels = [ # so far there is just a single label for an edge ( its_label, coords_label, { 'horizontalalignment': 'center', 'verticalalignment': 'center', 'size': self.scale * 200 * self.label_scale, 'color': colors['tc'], 'box_fc': status_color } ) ] else: self.labels = None # cross through the label if it is not functional if not self.functional: cross_bar = FancyBboxPatch( (0, 0), self.scale * self.label_scale, 0.07 * self.scale * self.label_scale, boxstyle="square,pad={}".format(0.05 * self.scale * self.label_scale), color=colors['ac'], alpha=min(1, 2 * self.alpha), zorder=6 ) rotator = Affine2D().rotate_deg(-45) corner_coords = ( coords_label[0] - 0.4 * self.scale * self.label_scale, coords_label[1] + 0.3 * self.scale * self.label_scale ) translator = Affine2D().translate(*corner_coords) transform = rotator + translator cross_bar.set_transform(transform) self.super_patch_collection.append(cross_bar)
def __init__(self, parent, handles, labels, loc=None, numpoints=None, # the number of points in the legend line markerscale=None, # the relative size of legend markers # vs. original markerfirst=True, # controls ordering (left-to-right) of # legend marker and label scatterpoints=None, # number of scatter points scatteryoffsets=None, prop=None, # properties for the legend texts fontsize=None, # keyword to set font size directly # spacing & pad defined as a fraction of the font-size borderpad=None, # the whitespace inside the legend border labelspacing=None, # the vertical space between the legend # entries handlelength=None, # the length of the legend handles handleheight=None, # the height of the legend handles handletextpad=None, # the pad between the legend handle # and text borderaxespad=None, # the pad between the axes and legend # border columnspacing=None, # spacing between columns ncol=1, # number of columns mode=None, # mode for horizontal distribution of columns. # None, "expand" fancybox=None, # True use a fancy box, false use a rounded # box, none use rc shadow=None, title=None, # set a title for the legend framealpha=None, # set frame alpha bbox_to_anchor=None, # bbox that the legend will be anchored. bbox_transform=None, # transform for the bbox frameon=None, # draw frame handler_map=None, displace=None, ): """ - *parent*: the artist that contains the legend - *handles*: a list of artists (lines, patches) to be added to the legend - *labels*: a list of strings to label the legend Optional keyword arguments: ================ ==================================================== Keyword Description ================ ==================================================== loc a location code prop the font property fontsize the font size (used only if prop is not specified) markerscale the relative size of legend markers vs. original markerfirst If true, place legend marker to left of label If false, place legend marker to right of label numpoints the number of points in the legend for line scatterpoints the number of points in the legend for scatter plot scatteryoffsets a list of yoffsets for scatter symbols in legend frameon if True, draw a frame around the legend. If None, use rc fancybox if True, draw a frame with a round fancybox. If None, use rc shadow if True, draw a shadow behind legend framealpha If not None, alpha channel for the frame. ncol number of columns borderpad the fractional whitespace inside the legend border labelspacing the vertical space between the legend entries handlelength the length of the legend handles handleheight the height of the legend handles handletextpad the pad between the legend handle and text borderaxespad the pad between the axes and legend border columnspacing the spacing between columns title the legend title bbox_to_anchor the bbox that the legend will be anchored. bbox_transform the transform for the bbox. transAxes if None. displace the legend posistion will be anchored ================ ==================================================== The pad and spacing parameters are measured in font-size units. e.g., a fontsize of 10 points and a handlelength=5 implies a handlelength of 50 points. Values from rcParams will be used if None. Users can specify any arbitrary location for the legend using the *bbox_to_anchor* keyword argument. bbox_to_anchor can be an instance of BboxBase(or its derivatives) or a tuple of 2 or 4 floats. See :meth:`set_bbox_to_anchor` for more detail. The legend location can be specified by setting *loc* with a tuple of 2 floats, which is interpreted as the lower-left corner of the legend in the normalized axes coordinate. """ # local import only to avoid circularity from matplotlib.axes import Axes from matplotlib.figure import Figure Artist.__init__(self) if prop is None: if fontsize is not None: self.prop = FontProperties(size=fontsize) else: self.prop = FontProperties(size=rcParams["legend.fontsize"]) elif isinstance(prop, dict): self.prop = FontProperties(**prop) if "size" not in prop: self.prop.set_size(rcParams["legend.fontsize"]) else: self.prop = prop self._fontsize = self.prop.get_size_in_points() self.texts = [] self.legendHandles = [] self._legend_title_box = None #: A dictionary with the extra handler mappings for this Legend #: instance. self._custom_handler_map = handler_map locals_view = locals() for name in ["numpoints", "markerscale", "shadow", "columnspacing", "scatterpoints", "handleheight", 'borderpad', 'labelspacing', 'handlelength', 'handletextpad', 'borderaxespad']: if locals_view[name] is None: value = rcParams["legend." + name] else: value = locals_view[name] setattr(self, name, value) del locals_view handles = list(handles) if len(handles) < 2: ncol = 1 self._ncol = ncol if self.numpoints <= 0: raise ValueError("numpoints must be > 0; it was %d" % numpoints) # introduce y-offset for handles of the scatter plot if scatteryoffsets is None: self._scatteryoffsets = np.array([3. / 8., 4. / 8., 2.5 / 8.]) else: self._scatteryoffsets = np.asarray(scatteryoffsets) reps = int(self.scatterpoints / len(self._scatteryoffsets)) + 1 self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[:self.scatterpoints] # _legend_box is an OffsetBox instance that contains all # legend items and will be initialized from _init_legend_box() # method. self._legend_box = None if isinstance(parent, Axes): self.isaxes = True self.axes = parent self.set_figure(parent.figure) elif isinstance(parent, Figure): self.isaxes = False self.set_figure(parent) else: raise TypeError("Legend needs either Axes or Figure as parent") self.parent = parent if loc is None: loc = rcParams["legend.loc"] if not self.isaxes and loc in [0, 'best']: loc = 'upper right' if is_string_like(loc): if loc not in self.codes: if self.isaxes: warnings.warn('Unrecognized location "%s". Falling back ' 'on "best"; valid locations are\n\t%s\n' % (loc, '\n\t'.join( six.iterkeys(self.codes)))) loc = 0 else: warnings.warn('Unrecognized location "%s". Falling back ' 'on "upper right"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join( six.iterkeys(self.codes)))) loc = 1 else: loc = self.codes[loc] if not self.isaxes and loc == 0: warnings.warn('Automatic legend placement (loc="best") not ' 'implemented for figure legend. ' 'Falling back on "upper right".') loc = 1 self._mode = mode self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform) # We use FancyBboxPatch to draw a legend frame. The location # and size of the box will be updated during the drawing time. if rcParams["legend.facecolor"] == 'inherit': facecolor = rcParams["axes.facecolor"] else: facecolor = rcParams["legend.facecolor"] if rcParams["legend.edgecolor"] == 'inherit': edgecolor = rcParams["axes.edgecolor"] else: edgecolor = rcParams["legend.edgecolor"] self.legendPatch = FancyBboxPatch( xy=(0.0, 0.0), width=1., height=1., facecolor=facecolor, edgecolor=edgecolor, mutation_scale=self._fontsize, snap=True ) # The width and height of the legendPatch will be set (in the # draw()) to the length that includes the padding. Thus we set # pad=0 here. if fancybox is None: fancybox = rcParams["legend.fancybox"] if fancybox: self.legendPatch.set_boxstyle("round", pad=0, rounding_size=0.2) else: self.legendPatch.set_boxstyle("square", pad=0) self._set_artist_props(self.legendPatch) self._drawFrame = frameon if frameon is None: self._drawFrame = rcParams["legend.frameon"] # init with null renderer self._init_legend_box(handles, labels, markerfirst) if framealpha is None: self.get_frame().set_alpha(rcParams["legend.framealpha"]) else: self.get_frame().set_alpha(framealpha) self._loc = loc self.set_title(title) self._last_fontsize_points = self._fontsize self._draggable = None if displace is not None: self.set_displace(displace)
def __init__( self, parent, handles, labels, loc=None, numpoints=None, # the number of points in the legend line markerscale=None, # the relative size of legend markers vs. original scatterpoints=3, # TODO: may be an rcParam scatteryoffsets=None, prop=None, # properties for the legend texts pad=None, # deprecated; use borderpad labelsep=None, # deprecated; use labelspacing handlelen=None, # deprecated; use handlelength handletextsep=None, # deprecated; use handletextpad axespad=None, # deprecated; use borderaxespad borderpad=None, # the whitespace inside the legend border labelspacing=None, # the vertical space between the legend entries handlelength=None, # the length of the legend handles handletextpad=None, # the pad between the legend handle and text borderaxespad=None, # the pad between the axes and legend border columnspacing=None, # spacing between columns ncol=1, # number of columns mode=None, # mode for horizontal distribution of columns. None, "expand" fancybox=None, # True use a fancy box, false use a rounded box, none use rc shadow=None, title=None, # set a title for the legend bbox_to_anchor=None, # bbox that the legend will be anchored. bbox_transform=None, # transform for the bbox ): """ - *parent* : the artist that contains the legend - *handles* : a list of artists (lines, patches) to add to the legend - *labels* : a list of strings to label the legend Optional keyword arguments: ================ ================================================================== Keyword Description ================ ================================================================== loc a location code prop the font property markerscale the relative size of legend markers vs. original numpoints the number of points in the legend for line scatterpoints the number of points in the legend for scatter plot scatteryoffsets a list of yoffsets for scatter symbols in legend fancybox if True, draw a frame with a round fancybox. If None, use rc shadow if True, draw a shadow behind legend ncol number of columns borderpad the fractional whitespace inside the legend border labelspacing the vertical space between the legend entries handlelength the length of the legend handles handletextpad the pad between the legend handle and text borderaxespad the pad between the axes and legend border columnspacing the spacing between columns title the legend title bbox_to_anchor the bbox that the legend will be anchored. bbox_transform the transform for the bbox. transAxes if None. ================ ================================================================== The dimensions of pad and spacing are given as a fraction of the _fontsize. Values from rcParams will be used if None. Users can specify any arbitrary location for the legend using the *bbox_to_anchor* keyword argument. bbox_to_anchor can be an instance of BboxBase(or its derivatives) or a tuple of 2 or 4 floats. See :meth:`set_bbox_to_anchor` for more detail. The legend location can be specified by setting *loc* with a tuple of 2 floats, which is interpreted as the lower-left corner of the legend in the normalized axes coordinate. """ from matplotlib.axes import Axes # local import only to avoid circularity from matplotlib.figure import Figure # local import only to avoid circularity Artist.__init__(self) if prop is None: self.prop = FontProperties(size=rcParams["legend.fontsize"]) elif isinstance(prop, dict): self.prop = FontProperties(**prop) if "size" not in prop: self.prop.set_size(rcParams["legend.fontsize"]) else: self.prop = prop self._fontsize = self.prop.get_size_in_points() propnames = ["numpoints", "markerscale", "shadow", "columnspacing", "scatterpoints"] self.texts = [] self.legendHandles = [] self._legend_title_box = None localdict = locals() for name in propnames: if localdict[name] is None: value = rcParams["legend." + name] else: value = localdict[name] setattr(self, name, value) deprecated_kwds = { "pad": "borderpad", "labelsep": "labelspacing", "handlelen": "handlelength", "handletextsep": "handletextpad", "axespad": "borderaxespad", } bbox = parent.bbox axessize_fontsize = min(bbox.width, bbox.height) / self._fontsize for k, v in deprecated_kwds.items(): if localdict[k] is not None and localdict[v] is None: warnings.warn("Use '%s' instead of '%s'." % (v, k), DeprecationWarning) setattr(self, v, localdict[k] * axessize_fontsize) continue if localdict[v] is None: setattr(self, v, rcParams["legend." + v]) else: setattr(self, v, localdict[v]) del localdict handles = list(handles) if len(handles) < 2: ncol = 1 self._ncol = ncol if self.numpoints <= 0: raise ValueError("numpoints must be >= 0; it was %d" % numpoints) if scatteryoffsets is None: self._scatteryoffsets = np.array([3.0 / 8.0, 4.0 / 8.0, 2.5 / 8.0]) else: self._scatteryoffsets = np.asarray(scatteryoffsets) reps = int(self.numpoints / len(self._scatteryoffsets)) + 1 self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[: self.scatterpoints] self._legend_box = None if isinstance(parent, Axes): self.isaxes = True self.set_axes(parent) self.set_figure(parent.figure) elif isinstance(parent, Figure): self.isaxes = False self.set_figure(parent) else: raise TypeError("Legend needs either Axes or Figure as parent") self.parent = parent if loc is None: loc = rcParams["legend.loc"] if not self.isaxes and loc in [0, "best"]: loc = "upper right" if is_string_like(loc): if loc not in self.codes: if self.isaxes: warnings.warn( 'Unrecognized location "%s". Falling back on "best"; ' "valid locations are\n\t%s\n" % (loc, "\n\t".join(self.codes.keys())) ) loc = 0 else: warnings.warn( 'Unrecognized location "%s". Falling back on "upper right"; ' "valid locations are\n\t%s\n" % (loc, "\n\t".join(self.codes.keys())) ) loc = 1 else: loc = self.codes[loc] if not self.isaxes and loc == 0: warnings.warn( 'Automatic legend placement (loc="best") not implemented for figure legend. ' 'Falling back on "upper right".' ) loc = 1 self._loc = loc self._mode = mode self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform) self.legendPatch = FancyBboxPatch( xy=(0.0, 0.0), width=1.0, height=1.0, facecolor="w", edgecolor="k", mutation_scale=self._fontsize, snap=True ) if fancybox is None: fancybox = rcParams["legend.fancybox"] if fancybox == True: self.legendPatch.set_boxstyle("round", pad=0, rounding_size=0.2) else: self.legendPatch.set_boxstyle("square", pad=0) self._set_artist_props(self.legendPatch) self._drawFrame = True self._init_legend_box(handles, labels) self.set_title(title) self._last_fontsize_points = self._fontsize
class Legend(Artist): """ Place a legend on the axes at location loc. Labels are a sequence of strings and loc can be a string or an integer specifying the legend location The location codes are:: 'best' : 0, (only implemented for axis legends) 'upper right' : 1, 'upper left' : 2, 'lower left' : 3, 'lower right' : 4, 'right' : 5, 'center left' : 6, 'center right' : 7, 'lower center' : 8, 'upper center' : 9, 'center' : 10, loc can be a tuple of the noramilzed coordinate values with respect its parent. Return value is a sequence of text, line instances that make up the legend """ codes = { 'best': 0, # only implemented for axis legends 'upper right': 1, 'upper left': 2, 'lower left': 3, 'lower right': 4, 'right': 5, 'center left': 6, 'center right': 7, 'lower center': 8, 'upper center': 9, 'center': 10, } zorder = 5 def __str__(self): return "Legend" def __init__( self, parent, handles, labels, loc=None, numpoints=None, # the number of points in the legend line markerscale=None, # the relative size of legend markers vs. original scatterpoints=3, # TODO: may be an rcParam scatteryoffsets=None, prop=None, # properties for the legend texts # the following dimensions are in axes coords pad=None, # deprecated; use borderpad labelsep=None, # deprecated; use labelspacing handlelen=None, # deprecated; use handlelength handletextsep=None, # deprecated; use handletextpad axespad=None, # deprecated; use borderaxespad # spacing & pad defined as a fractionof the font-size borderpad=None, # the whitespace inside the legend border labelspacing=None, #the vertical space between the legend entries handlelength=None, # the length of the legend handles handletextpad=None, # the pad between the legend handle and text borderaxespad=None, # the pad between the axes and legend border columnspacing=None, # spacing between columns ncol=1, # number of columns mode=None, # mode for horizontal distribution of columns. None, "expand" fancybox=None, # True use a fancy box, false use a rounded box, none use rc shadow=None, ): """ - *parent* : the artist that contains the legend - *handles* : a list of artists (lines, patches) to add to the legend - *labels* : a list of strings to label the legend Optional keyword arguments: ================ ================================================================== Keyword Description ================ ================================================================== loc a location code or a tuple of coordinates numpoints the number of points in the legend line prop the font property markerscale the relative size of legend markers vs. original fancybox if True, draw a frame with a round fancybox. If None, use rc shadow if True, draw a shadow behind legend scatteryoffsets a list of yoffsets for scatter symbols in legend borderpad the fractional whitespace inside the legend border labelspacing the vertical space between the legend entries handlelength the length of the legend handles handletextpad the pad between the legend handle and text borderaxespad the pad between the axes and legend border columnspacing the spacing between columns ================ ================================================================== The dimensions of pad and spacing are given as a fraction of the fontsize. Values from rcParams will be used if None. """ from matplotlib.axes import Axes # local import only to avoid circularity from matplotlib.figure import Figure # local import only to avoid circularity Artist.__init__(self) if prop is None: self.prop = FontProperties(size=rcParams["legend.fontsize"]) else: self.prop = prop self.fontsize = self.prop.get_size_in_points() propnames = [ 'numpoints', 'markerscale', 'shadow', "columnspacing", "scatterpoints" ] localdict = locals() for name in propnames: if localdict[name] is None: value = rcParams["legend." + name] else: value = localdict[name] setattr(self, name, value) # Take care the deprecated keywords deprecated_kwds = { "pad": "borderpad", "labelsep": "labelspacing", "handlelen": "handlelength", "handletextsep": "handletextpad", "axespad": "borderaxespad" } # convert values of deprecated keywords (ginve in axes coords) # to new vaules in a fraction of the font size # conversion factor bbox = parent.bbox axessize_fontsize = min(bbox.width, bbox.height) / self.fontsize for k, v in deprecated_kwds.items(): # use deprecated value if not None and if their newer # counter part is None. if localdict[k] is not None and localdict[v] is None: warnings.warn("Use '%s' instead of '%s'." % (v, k), DeprecationWarning) setattr(self, v, localdict[k] * axessize_fontsize) continue # Otherwise, use new keywords if localdict[v] is None: setattr(self, v, rcParams["legend." + v]) else: setattr(self, v, localdict[v]) del localdict self._ncol = ncol if self.numpoints <= 0: raise ValueError("numpoints must be >= 0; it was %d" % numpoints) # introduce y-offset for handles of the scatter plot if scatteryoffsets is None: self._scatteryoffsets = np.array([3. / 8., 4. / 8., 2.5 / 8.]) else: self._scatteryoffsets = np.asarray(scatteryoffsets) reps = int(self.numpoints / len(self._scatteryoffsets)) + 1 self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[:self.scatterpoints] # _legend_box is an OffsetBox instance that contains all # legend items and will be initialized from _init_legend_box() # method. self._legend_box = None if isinstance(parent, Axes): self.isaxes = True self.set_figure(parent.figure) elif isinstance(parent, Figure): self.isaxes = False self.set_figure(parent) else: raise TypeError("Legend needs either Axes or Figure as parent") self.parent = parent if loc is None: loc = rcParams["legend.loc"] if not self.isaxes and loc in [0, 'best']: loc = 'upper right' if is_string_like(loc): if loc not in self.codes: if self.isaxes: warnings.warn( 'Unrecognized location "%s". Falling back on "best"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes.keys()))) loc = 0 else: warnings.warn( 'Unrecognized location "%s". Falling back on "upper right"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes.keys()))) loc = 1 else: loc = self.codes[loc] if not self.isaxes and loc == 0: warnings.warn( 'Automatic legend placement (loc="best") not implemented for figure legend. ' 'Falling back on "upper right".') loc = 1 self._loc = loc self._mode = mode # We use FancyBboxPatch to draw a legend frame. The location # and size of the box will be updated during the drawing time. self.legendPatch = FancyBboxPatch(xy=(0.0, 0.0), width=1., height=1., facecolor='w', edgecolor='k', mutation_scale=self.fontsize, snap=True) # The width and height of the legendPatch will be set (in the # draw()) to the length that includes the padding. Thus we set # pad=0 here. if fancybox is None: fancybox = rcParams["legend.fancybox"] if fancybox == True: self.legendPatch.set_boxstyle("round", pad=0, rounding_size=0.2) else: self.legendPatch.set_boxstyle("square", pad=0) self._set_artist_props(self.legendPatch) self._drawFrame = True # populate the legend_box with legend items. self._init_legend_box(handles, labels) self._legend_box.set_figure(self.figure) def _set_artist_props(self, a): """ set the boilerplate props for artists added to axes """ a.set_figure(self.figure) for c in self.get_children(): c.set_figure(self.figure) a.set_transform(self.get_transform()) def _findoffset_best(self, width, height, xdescent, ydescent): "Heper function to locate the legend at its best position" ox, oy = self._find_best_position(width, height) return ox + xdescent, oy + ydescent def _findoffset_loc(self, width, height, xdescent, ydescent): "Heper function to locate the legend using the location code" if iterable(self._loc) and len(self._loc) == 2: # when loc is a tuple of axes(or figure) coordinates. fx, fy = self._loc bbox = self.parent.bbox x, y = bbox.x0 + bbox.width * fx, bbox.y0 + bbox.height * fy else: bbox = Bbox.from_bounds(0, 0, width, height) x, y = self._get_anchored_bbox(self._loc, bbox, self.parent.bbox) return x + xdescent, y + ydescent def draw(self, renderer): "Draw everything that belongs to the legend" if not self.get_visible(): return renderer.open_group('legend') # find_offset function will be provided to _legend_box and # _legend_box will draw itself at the location of the return # value of the find_offset. if self._loc == 0: self._legend_box.set_offset(self._findoffset_best) else: self._legend_box.set_offset(self._findoffset_loc) # if mode == fill, set the width of the legend_box to the # width of the paret (minus pads) if self._mode in ["expand"]: pad = 2 * (self.borderaxespad + self.borderpad) * self.fontsize self._legend_box.set_width(self.parent.bbox.width - pad) if self._drawFrame: # update the location and size of the legend bbox = self._legend_box.get_window_extent(renderer) self.legendPatch.set_bounds(bbox.x0, bbox.y0, bbox.width, bbox.height) if self.shadow: shadow = Shadow(self.legendPatch, 2, -2) shadow.draw(renderer) self.legendPatch.draw(renderer) self._legend_box.draw(renderer) renderer.close_group('legend') def _approx_text_height(self): """ Return the approximate height of the text. This is used to place the legend handle. """ return self.fontsize / 72.0 * self.figure.dpi def _init_legend_box(self, handles, labels): """ Initiallize the legend_box. The legend_box is an instance of the OffsetBox, which is packed with legend handles and texts. Once packed, their location is calculated during the drawing time. """ # legend_box is a HPacker, horizontally packed with # columns. Each column is a VPacker, vertically packed with # legend items. Each legend item is HPacker packed with # legend handleBox and labelBox. handleBox is an instance of # offsetbox.DrawingArea which contains legend handle. labelBox # is an instance of offsetbox.TextArea which contains legend # text. text_list = [] # the list of text instances handle_list = [] # the list of text instances label_prop = dict( verticalalignment='baseline', horizontalalignment='left', fontproperties=self.prop, ) labelboxes = [] for l in labels: textbox = TextArea(l, textprops=label_prop, multilinebaseline=True, minimumdescent=True) text_list.append(textbox._text) labelboxes.append(textbox) handleboxes = [] # The approximate height and descent of text. These values are # only used for plotting the legend handle. height = self._approx_text_height() * 0.7 descent = 0. # each handle needs to be drawn inside a box of # (x, y, w, h) = (0, -descent, width, height). # And their corrdinates should be given in the display coordinates. # The transformation of each handle will be automatically set # to self.get_trasnform(). If the artist does not uses its # default trasnform (eg, Collections), you need to # manually set their transform to the self.get_transform(). for handle in handles: if isinstance(handle, RegularPolyCollection): npoints = self.scatterpoints else: npoints = self.numpoints if npoints > 1: # we put some pad here to compensate the size of the # marker xdata = np.linspace(0.3 * self.fontsize, (self.handlelength - 0.3) * self.fontsize, npoints) xdata_marker = xdata elif npoints == 1: xdata = np.linspace(0, self.handlelength * self.fontsize, 2) xdata_marker = [0.5 * self.handlelength * self.fontsize] if isinstance(handle, Line2D): ydata = ((height - descent) / 2.) * np.ones(xdata.shape, float) legline = Line2D(xdata, ydata) legline.update_from(handle) self._set_artist_props(legline) # after update legline.set_clip_box(None) legline.set_clip_path(None) legline.set_drawstyle('default') legline.set_marker('None') handle_list.append(legline) legline_marker = Line2D(xdata_marker, ydata[:len(xdata_marker)]) legline_marker.update_from(handle) self._set_artist_props(legline_marker) legline_marker.set_clip_box(None) legline_marker.set_clip_path(None) legline_marker.set_linestyle('None') # we don't want to add this to the return list because # the texts and handles are assumed to be in one-to-one # correpondence. legline._legmarker = legline_marker elif isinstance(handle, Patch): p = Rectangle( xy=(0., 0.), width=self.handlelength * self.fontsize, height=(height - descent), ) p.update_from(handle) self._set_artist_props(p) p.set_clip_box(None) p.set_clip_path(None) handle_list.append(p) elif isinstance(handle, LineCollection): ydata = ((height - descent) / 2.) * np.ones(xdata.shape, float) legline = Line2D(xdata, ydata) self._set_artist_props(legline) legline.set_clip_box(None) legline.set_clip_path(None) lw = handle.get_linewidth()[0] dashes = handle.get_dashes()[0] color = handle.get_colors()[0] legline.set_color(color) legline.set_linewidth(lw) legline.set_dashes(dashes) handle_list.append(legline) elif isinstance(handle, RegularPolyCollection): #ydata = self._scatteryoffsets ydata = height * self._scatteryoffsets size_max, size_min = max(handle.get_sizes()),\ min(handle.get_sizes()) # we may need to scale these sizes by "markerscale" # attribute. But other handle types does not seem # to care about this attribute and it is currently ignored. if self.scatterpoints < 4: sizes = [.5 * (size_max + size_min), size_max, size_min] else: sizes = (size_max - size_min) * np.linspace( 0, 1, self.scatterpoints) + size_min p = type(handle)( handle.get_numsides(), rotation=handle.get_rotation(), sizes=sizes, offsets=zip(xdata_marker, ydata), transOffset=self.get_transform(), ) p.update_from(handle) p.set_figure(self.figure) p.set_clip_box(None) p.set_clip_path(None) handle_list.append(p) else: handle_list.append(None) handlebox = DrawingArea(width=self.handlelength * self.fontsize, height=height, xdescent=0., ydescent=descent) handle = handle_list[-1] handlebox.add_artist(handle) if hasattr(handle, "_legmarker"): handlebox.add_artist(handle._legmarker) handleboxes.append(handlebox) # We calculate number of lows in each column. The first # (num_largecol) columns will have (nrows+1) rows, and remaing # (num_smallcol) columns will have (nrows) rows. nrows, num_largecol = divmod(len(handleboxes), self._ncol) num_smallcol = self._ncol - num_largecol # starting index of each column and number of rows in it. largecol = safezip(range(0, num_largecol * (nrows + 1), (nrows + 1)), [nrows + 1] * num_largecol) smallcol = safezip( range(num_largecol * (nrows + 1), len(handleboxes), nrows), [nrows] * num_smallcol) handle_label = safezip(handleboxes, labelboxes) columnbox = [] for i0, di in largecol + smallcol: # pack handleBox and labelBox into itemBox itemBoxes = [ HPacker(pad=0, sep=self.handletextpad * self.fontsize, children=[h, t], align="baseline") for h, t in handle_label[i0:i0 + di] ] # minimumdescent=False for the text of the last row of the column itemBoxes[-1].get_children()[1].set_minimumdescent(False) # pack columnBox columnbox.append( VPacker(pad=0, sep=self.labelspacing * self.fontsize, align="baseline", children=itemBoxes)) if self._mode == "expand": mode = "expand" else: mode = "fixed" sep = self.columnspacing * self.fontsize self._legend_box = HPacker(pad=self.borderpad * self.fontsize, sep=sep, align="baseline", mode=mode, children=columnbox) self.texts = text_list self.legendHandles = handle_list def _auto_legend_data(self): """ Returns list of vertices and extents covered by the plot. Returns a two long list. First element is a list of (x, y) vertices (in display-coordinates) covered by all the lines and line collections, in the legend's handles. Second element is a list of bounding boxes for all the patches in the legend's handles. """ assert self.isaxes # should always hold because function is only called internally ax = self.parent vertices = [] bboxes = [] lines = [] for handle in ax.lines: assert isinstance(handle, Line2D) path = handle.get_path() trans = handle.get_transform() tpath = trans.transform_path(path) lines.append(tpath) for handle in ax.patches: assert isinstance(handle, Patch) if isinstance(handle, Rectangle): transform = handle.get_data_transform() bboxes.append(handle.get_bbox().transformed(transform)) else: transform = handle.get_transform() bboxes.append(handle.get_path().get_extents(transform)) return [vertices, bboxes, lines] def draw_frame(self, b): 'b is a boolean. Set draw frame to b' self._drawFrame = b def get_children(self): 'return a list of child artists' children = [] if self._legend_box: children.append(self._legend_box) return children def get_frame(self): 'return the Rectangle instance used to frame the legend' return self.legendPatch def get_lines(self): 'return a list of lines.Line2D instances in the legend' return [h for h in self.legendHandles if isinstance(h, Line2D)] def get_patches(self): 'return a list of patch instances in the legend' return silent_list( 'Patch', [h for h in self.legendHandles if isinstance(h, Patch)]) def get_texts(self): 'return a list of text.Text instance in the legend' return silent_list('Text', self.texts) def get_window_extent(self): 'return a extent of the the legend' return self.legendPatch.get_window_extent() def _get_anchored_bbox(self, loc, bbox, parentbbox): """ Place the *bbox* inside the *parentbbox* according to a given location code. Return the (x,y) coordinate of the bbox. - loc: a location code in range(1, 11). This corresponds to the possible values for self._loc, excluding "best". - bbox: bbox to be placed, display coodinate units. - parentbbox: a parent box which will contain the bbox. In display coordinates. """ assert loc in range(1, 11) # called only internally BEST, UR, UL, LL, LR, R, CL, CR, LC, UC, C = range(11) anchor_coefs = { UR: "NE", UL: "NW", LL: "SW", LR: "SE", R: "E", CL: "W", CR: "E", LC: "S", UC: "N", C: "C" } c = anchor_coefs[loc] container = parentbbox.padded(-(self.borderaxespad) * self.fontsize) anchored_box = bbox.anchored(c, container=container) return anchored_box.x0, anchored_box.y0 def _find_best_position(self, width, height, consider=None): """ Determine the best location to place the legend. `consider` is a list of (x, y) pairs to consider as a potential lower-left corner of the legend. All are display coords. """ assert self.isaxes # should always hold because function is only called internally verts, bboxes, lines = self._auto_legend_data() bbox = Bbox.from_bounds(0, 0, width, height) consider = [ self._get_anchored_bbox(x, bbox, self.parent.bbox) for x in range(1, len(self.codes)) ] #tx, ty = self.legendPatch.get_x(), self.legendPatch.get_y() candidates = [] for l, b in consider: legendBox = Bbox.from_bounds(l, b, width, height) badness = 0 badness = legendBox.count_contains(verts) badness += legendBox.count_overlaps(bboxes) for line in lines: if line.intersects_bbox(legendBox): badness += 1 ox, oy = l, b if badness == 0: return ox, oy candidates.append((badness, (l, b))) # rather than use min() or list.sort(), do this so that we are assured # that in the case of two equal badnesses, the one first considered is # returned. # NOTE: list.sort() is stable.But leave as it is for now. -JJL minCandidate = candidates[0] for candidate in candidates: if candidate[0] < minCandidate[0]: minCandidate = candidate ox, oy = minCandidate[1] return ox, oy
def __init__(self, parent, handles, labels, loc=None, numpoints=None, # the number of points in the legend line markerscale=None, # the relative size of legend markers # vs. original markerfirst=True, # controls ordering (left-to-right) of # legend marker and label scatterpoints=None, # number of scatter points scatteryoffsets=None, prop=None, # properties for the legend texts fontsize=None, # keyword to set font size directly # spacing & pad defined as a fraction of the font-size borderpad=None, # the whitespace inside the legend border labelspacing=None, # the vertical space between the legend # entries handlelength=None, # the length of the legend handles handleheight=None, # the height of the legend handles handletextpad=None, # the pad between the legend handle # and text borderaxespad=None, # the pad between the axes and legend # border columnspacing=None, # spacing between columns ncol=1, # number of columns mode=None, # mode for horizontal distribution of columns. # None, "expand" fancybox=None, # True use a fancy box, false use a rounded # box, none use rc shadow=None, title=None, # set a title for the legend title_fontsize=None, # set to ax.fontsize if None framealpha=None, # set frame alpha edgecolor=None, # frame patch edgecolor facecolor=None, # frame patch facecolor bbox_to_anchor=None, # bbox that the legend will be anchored. bbox_transform=None, # transform for the bbox frameon=None, # draw frame handler_map=None, ): """ Parameters ---------- parent : `~matplotlib.axes.Axes` or `.Figure` The artist that contains the legend. handles : sequence of `.Artist` A list of Artists (lines, patches) to be added to the legend. labels : sequence of strings A list of labels to show next to the artists. The length of handles and labels should be the same. If they are not, they are truncated to the smaller of both lengths. Other Parameters ---------------- %(_legend_kw_doc)s Notes ----- Users can specify any arbitrary location for the legend using the *bbox_to_anchor* keyword argument. bbox_to_anchor can be an instance of BboxBase(or its derivatives) or a tuple of 2 or 4 floats. See :meth:`set_bbox_to_anchor` for more detail. The legend location can be specified by setting *loc* with a tuple of 2 floats, which is interpreted as the lower-left corner of the legend in the normalized axes coordinate. """ # local import only to avoid circularity from matplotlib.axes import Axes from matplotlib.figure import Figure Artist.__init__(self) if prop is None: if fontsize is not None: self.prop = FontProperties(size=fontsize) else: self.prop = FontProperties(size=rcParams["legend.fontsize"]) elif isinstance(prop, dict): self.prop = FontProperties(**prop) if "size" not in prop: self.prop.set_size(rcParams["legend.fontsize"]) else: self.prop = prop self._fontsize = self.prop.get_size_in_points() self.texts = [] self.legendHandles = [] self._legend_title_box = None #: A dictionary with the extra handler mappings for this Legend #: instance. self._custom_handler_map = handler_map locals_view = locals() for name in ["numpoints", "markerscale", "shadow", "columnspacing", "scatterpoints", "handleheight", 'borderpad', 'labelspacing', 'handlelength', 'handletextpad', 'borderaxespad']: if locals_view[name] is None: value = rcParams["legend." + name] else: value = locals_view[name] setattr(self, name, value) del locals_view # trim handles and labels if illegal label... _lab, _hand = [], [] for label, handle in zip(labels, handles): if isinstance(label, str) and label.startswith('_'): cbook._warn_external('The handle {!r} has a label of {!r} ' 'which cannot be automatically added to' ' the legend.'.format(handle, label)) else: _lab.append(label) _hand.append(handle) labels, handles = _lab, _hand handles = list(handles) if len(handles) < 2: ncol = 1 self._ncol = ncol if self.numpoints <= 0: raise ValueError("numpoints must be > 0; it was %d" % numpoints) # introduce y-offset for handles of the scatter plot if scatteryoffsets is None: self._scatteryoffsets = np.array([3. / 8., 4. / 8., 2.5 / 8.]) else: self._scatteryoffsets = np.asarray(scatteryoffsets) reps = self.scatterpoints // len(self._scatteryoffsets) + 1 self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[:self.scatterpoints] # _legend_box is an OffsetBox instance that contains all # legend items and will be initialized from _init_legend_box() # method. self._legend_box = None if isinstance(parent, Axes): self.isaxes = True self.axes = parent self.set_figure(parent.figure) elif isinstance(parent, Figure): self.isaxes = False self.set_figure(parent) else: raise TypeError("Legend needs either Axes or Figure as parent") self.parent = parent self._loc_used_default = loc is None if loc is None: loc = rcParams["legend.loc"] if not self.isaxes and loc in [0, 'best']: loc = 'upper right' if isinstance(loc, str): if loc not in self.codes: if self.isaxes: cbook.warn_deprecated( "3.1", message="Unrecognized location {!r}. Falling " "back on 'best'; valid locations are\n\t{}\n" "This will raise an exception %(removal)s." .format(loc, '\n\t'.join(self.codes))) loc = 0 else: cbook.warn_deprecated( "3.1", message="Unrecognized location {!r}. Falling " "back on 'upper right'; valid locations are\n\t{}\n'" "This will raise an exception %(removal)s." .format(loc, '\n\t'.join(self.codes))) loc = 1 else: loc = self.codes[loc] if not self.isaxes and loc == 0: cbook.warn_deprecated( "3.1", message="Automatic legend placement (loc='best') not " "implemented for figure legend. Falling back on 'upper " "right'. This will raise an exception %(removal)s.") loc = 1 self._mode = mode self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform) # We use FancyBboxPatch to draw a legend frame. The location # and size of the box will be updated during the drawing time. if facecolor is None: facecolor = rcParams["legend.facecolor"] if facecolor == 'inherit': facecolor = rcParams["axes.facecolor"] if edgecolor is None: edgecolor = rcParams["legend.edgecolor"] if edgecolor == 'inherit': edgecolor = rcParams["axes.edgecolor"] self.legendPatch = FancyBboxPatch( xy=(0.0, 0.0), width=1., height=1., facecolor=facecolor, edgecolor=edgecolor, mutation_scale=self._fontsize, snap=True ) # The width and height of the legendPatch will be set (in the # draw()) to the length that includes the padding. Thus we set # pad=0 here. if fancybox is None: fancybox = rcParams["legend.fancybox"] if fancybox: self.legendPatch.set_boxstyle("round", pad=0, rounding_size=0.2) else: self.legendPatch.set_boxstyle("square", pad=0) self._set_artist_props(self.legendPatch) self._drawFrame = frameon if frameon is None: self._drawFrame = rcParams["legend.frameon"] # init with null renderer self._init_legend_box(handles, labels, markerfirst) # If shadow is activated use framealpha if not # explicitly passed. See Issue 8943 if framealpha is None: if shadow: self.get_frame().set_alpha(1) else: self.get_frame().set_alpha(rcParams["legend.framealpha"]) else: self.get_frame().set_alpha(framealpha) tmp = self._loc_used_default self._set_loc(loc) self._loc_used_default = tmp # ignore changes done by _set_loc # figure out title fontsize: if title_fontsize is None: title_fontsize = rcParams['legend.title_fontsize'] tprop = FontProperties(size=title_fontsize) self.set_title(title, prop=tprop) self._last_fontsize_points = self._fontsize self._draggable = None
class Legend(Artist): """ Place a legend on the axes at location loc. Labels are a sequence of strings and loc can be a string or an integer specifying the legend location The location codes are:: 'best' : 0, (only implemented for axes legends) 'upper right' : 1, 'upper left' : 2, 'lower left' : 3, 'lower right' : 4, 'right' : 5, (same as 'center right', for back-compatibility) 'center left' : 6, 'center right' : 7, 'lower center' : 8, 'upper center' : 9, 'center' : 10, loc can be a tuple of the normalized coordinate values with respect its parent. """ codes = { 'best': 0, # only implemented for axes legends 'upper right': 1, 'upper left': 2, 'lower left': 3, 'lower right': 4, 'right': 5, 'center left': 6, 'center right': 7, 'lower center': 8, 'upper center': 9, 'center': 10, } zorder = 5 def __str__(self): return "Legend" def __init__( self, parent, handles, labels, loc=None, numpoints=None, # the number of points in the legend line markerscale=None, # the relative size of legend markers # vs. original markerfirst=True, # controls ordering (left-to-right) of # legend marker and label scatterpoints=None, # number of scatter points scatteryoffsets=None, prop=None, # properties for the legend texts fontsize=None, # keyword to set font size directly # spacing & pad defined as a fraction of the font-size borderpad=None, # the whitespace inside the legend border labelspacing=None, # the vertical space between the legend # entries handlelength=None, # the length of the legend handles handleheight=None, # the height of the legend handles handletextpad=None, # the pad between the legend handle # and text borderaxespad=None, # the pad between the axes and legend # border columnspacing=None, # spacing between columns ncol=1, # number of columns mode=None, # mode for horizontal distribution of columns. # None, "expand" fancybox=None, # True use a fancy box, false use a rounded # box, none use rc shadow=None, title=None, # set a title for the legend framealpha=None, # set frame alpha edgecolor=None, # frame patch edgecolor facecolor=None, # frame patch facecolor bbox_to_anchor=None, # bbox that the legend will be anchored. bbox_transform=None, # transform for the bbox frameon=None, # draw frame handler_map=None, ): """ - *parent*: the artist that contains the legend - *handles*: a list of artists (lines, patches) to be added to the legend - *labels*: a list of strings to label the legend Optional keyword arguments: ================ ==================================================== Keyword Description ================ ==================================================== loc Location code string, or tuple (see below). prop the font property fontsize the font size (used only if prop is not specified) markerscale the relative size of legend markers vs. original markerfirst If True (default), marker is to left of the label. numpoints the number of points in the legend for line scatterpoints the number of points in the legend for scatter plot scatteryoffsets a list of yoffsets for scatter symbols in legend frameon If True, draw the legend on a patch (frame). fancybox If True, draw the frame with a round fancybox. shadow If True, draw a shadow behind legend. framealpha Transparency of the frame. edgecolor Frame edgecolor. facecolor Frame facecolor. ncol number of columns borderpad the fractional whitespace inside the legend border labelspacing the vertical space between the legend entries handlelength the length of the legend handles handleheight the height of the legend handles handletextpad the pad between the legend handle and text borderaxespad the pad between the axes and legend border columnspacing the spacing between columns title the legend title bbox_to_anchor the bbox that the legend will be anchored. bbox_transform the transform for the bbox. transAxes if None. ================ ==================================================== The pad and spacing parameters are measured in font-size units. e.g., a fontsize of 10 points and a handlelength=5 implies a handlelength of 50 points. Values from rcParams will be used if None. Users can specify any arbitrary location for the legend using the *bbox_to_anchor* keyword argument. bbox_to_anchor can be an instance of BboxBase(or its derivatives) or a tuple of 2 or 4 floats. See :meth:`set_bbox_to_anchor` for more detail. The legend location can be specified by setting *loc* with a tuple of 2 floats, which is interpreted as the lower-left corner of the legend in the normalized axes coordinate. """ # local import only to avoid circularity from matplotlib.axes import Axes from matplotlib.figure import Figure Artist.__init__(self) if prop is None: if fontsize is not None: self.prop = FontProperties(size=fontsize) else: self.prop = FontProperties(size=rcParams["legend.fontsize"]) elif isinstance(prop, dict): self.prop = FontProperties(**prop) if "size" not in prop: self.prop.set_size(rcParams["legend.fontsize"]) else: self.prop = prop self._fontsize = self.prop.get_size_in_points() self.texts = [] self.legendHandles = [] self._legend_title_box = None #: A dictionary with the extra handler mappings for this Legend #: instance. self._custom_handler_map = handler_map locals_view = locals() for name in [ "numpoints", "markerscale", "shadow", "columnspacing", "scatterpoints", "handleheight", 'borderpad', 'labelspacing', 'handlelength', 'handletextpad', 'borderaxespad' ]: if locals_view[name] is None: value = rcParams["legend." + name] else: value = locals_view[name] setattr(self, name, value) del locals_view handles = list(handles) if len(handles) < 2: ncol = 1 self._ncol = ncol if self.numpoints <= 0: raise ValueError("numpoints must be > 0; it was %d" % numpoints) # introduce y-offset for handles of the scatter plot if scatteryoffsets is None: self._scatteryoffsets = np.array([3. / 8., 4. / 8., 2.5 / 8.]) else: self._scatteryoffsets = np.asarray(scatteryoffsets) reps = self.scatterpoints // len(self._scatteryoffsets) + 1 self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[:self.scatterpoints] # _legend_box is an OffsetBox instance that contains all # legend items and will be initialized from _init_legend_box() # method. self._legend_box = None if isinstance(parent, Axes): self.isaxes = True self.axes = parent self.set_figure(parent.figure) elif isinstance(parent, Figure): self.isaxes = False self.set_figure(parent) else: raise TypeError("Legend needs either Axes or Figure as parent") self.parent = parent if loc is None: loc = rcParams["legend.loc"] if not self.isaxes and loc in [0, 'best']: loc = 'upper right' if isinstance(loc, six.string_types): if loc not in self.codes: if self.isaxes: warnings.warn('Unrecognized location "%s". Falling back ' 'on "best"; valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes))) loc = 0 else: warnings.warn('Unrecognized location "%s". Falling back ' 'on "upper right"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes))) loc = 1 else: loc = self.codes[loc] if not self.isaxes and loc == 0: warnings.warn('Automatic legend placement (loc="best") not ' 'implemented for figure legend. ' 'Falling back on "upper right".') loc = 1 self._mode = mode self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform) # We use FancyBboxPatch to draw a legend frame. The location # and size of the box will be updated during the drawing time. if facecolor is None: facecolor = rcParams["legend.facecolor"] if facecolor == 'inherit': facecolor = rcParams["axes.facecolor"] if edgecolor is None: edgecolor = rcParams["legend.edgecolor"] if edgecolor == 'inherit': edgecolor = rcParams["axes.edgecolor"] self.legendPatch = FancyBboxPatch(xy=(0.0, 0.0), width=1., height=1., facecolor=facecolor, edgecolor=edgecolor, mutation_scale=self._fontsize, snap=True) # The width and height of the legendPatch will be set (in the # draw()) to the length that includes the padding. Thus we set # pad=0 here. if fancybox is None: fancybox = rcParams["legend.fancybox"] if fancybox: self.legendPatch.set_boxstyle("round", pad=0, rounding_size=0.2) else: self.legendPatch.set_boxstyle("square", pad=0) self._set_artist_props(self.legendPatch) self._drawFrame = frameon if frameon is None: self._drawFrame = rcParams["legend.frameon"] # init with null renderer self._init_legend_box(handles, labels, markerfirst) # If shadow is activated use framealpha if not # explicitly passed. See Issue 8943 if framealpha is None: if shadow: self.get_frame().set_alpha(1) else: self.get_frame().set_alpha(rcParams["legend.framealpha"]) else: self.get_frame().set_alpha(framealpha) self._loc = loc self.set_title(title) self._last_fontsize_points = self._fontsize self._draggable = None def _set_artist_props(self, a): """ set the boilerplate props for artists added to axes """ a.set_figure(self.figure) if self.isaxes: # a.set_axes(self.axes) a.axes = self.axes a.set_transform(self.get_transform()) def _set_loc(self, loc): # find_offset function will be provided to _legend_box and # _legend_box will draw itself at the location of the return # value of the find_offset. self._loc_real = loc self.stale = True def _get_loc(self): return self._loc_real _loc = property(_get_loc, _set_loc) def _findoffset(self, width, height, xdescent, ydescent, renderer): "Helper function to locate the legend" if self._loc == 0: # "best". x, y = self._find_best_position(width, height, renderer) elif self._loc in Legend.codes.values(): # Fixed location. bbox = Bbox.from_bounds(0, 0, width, height) x, y = self._get_anchored_bbox(self._loc, bbox, self.get_bbox_to_anchor(), renderer) else: # Axes or figure coordinates. fx, fy = self._loc bbox = self.get_bbox_to_anchor() x, y = bbox.x0 + bbox.width * fx, bbox.y0 + bbox.height * fy return x + xdescent, y + ydescent @allow_rasterization def draw(self, renderer): "Draw everything that belongs to the legend" if not self.get_visible(): return renderer.open_group('legend') fontsize = renderer.points_to_pixels(self._fontsize) # if mode == fill, set the width of the legend_box to the # width of the paret (minus pads) if self._mode in ["expand"]: pad = 2 * (self.borderaxespad + self.borderpad) * fontsize self._legend_box.set_width(self.get_bbox_to_anchor().width - pad) # update the location and size of the legend. This needs to # be done in any case to clip the figure right. bbox = self._legend_box.get_window_extent(renderer) self.legendPatch.set_bounds(bbox.x0, bbox.y0, bbox.width, bbox.height) self.legendPatch.set_mutation_scale(fontsize) if self._drawFrame: if self.shadow: shadow = Shadow(self.legendPatch, 2, -2) shadow.draw(renderer) self.legendPatch.draw(renderer) self._legend_box.draw(renderer) renderer.close_group('legend') self.stale = False def _approx_text_height(self, renderer=None): """ Return the approximate height of the text. This is used to place the legend handle. """ if renderer is None: return self._fontsize else: return renderer.points_to_pixels(self._fontsize) # _default_handler_map defines the default mapping between plot # elements and the legend handlers. _default_handler_map = { StemContainer: legend_handler.HandlerStem(), ErrorbarContainer: legend_handler.HandlerErrorbar(), Line2D: legend_handler.HandlerLine2D(), Patch: legend_handler.HandlerPatch(), LineCollection: legend_handler.HandlerLineCollection(), RegularPolyCollection: legend_handler.HandlerRegularPolyCollection(), CircleCollection: legend_handler.HandlerCircleCollection(), BarContainer: legend_handler.HandlerPatch( update_func=legend_handler.update_from_first_child), tuple: legend_handler.HandlerTuple(), PathCollection: legend_handler.HandlerPathCollection(), PolyCollection: legend_handler.HandlerPolyCollection() } # (get|set|update)_default_handler_maps are public interfaces to # modify the default handler map. @classmethod def get_default_handler_map(cls): """ A class method that returns the default handler map. """ return cls._default_handler_map @classmethod def set_default_handler_map(cls, handler_map): """ A class method to set the default handler map. """ cls._default_handler_map = handler_map @classmethod def update_default_handler_map(cls, handler_map): """ A class method to update the default handler map. """ cls._default_handler_map.update(handler_map) def get_legend_handler_map(self): """ return the handler map. """ default_handler_map = self.get_default_handler_map() if self._custom_handler_map: hm = default_handler_map.copy() hm.update(self._custom_handler_map) return hm else: return default_handler_map @staticmethod def get_legend_handler(legend_handler_map, orig_handle): """ return a legend handler from *legend_handler_map* that corresponds to *orig_handler*. *legend_handler_map* should be a dictionary object (that is returned by the get_legend_handler_map method). It first checks if the *orig_handle* itself is a key in the *legend_hanler_map* and return the associated value. Otherwise, it checks for each of the classes in its method-resolution-order. If no matching key is found, it returns None. """ if is_hashable(orig_handle): try: return legend_handler_map[orig_handle] except KeyError: pass for handle_type in type(orig_handle).mro(): try: return legend_handler_map[handle_type] except KeyError: pass return None def _init_legend_box(self, handles, labels, markerfirst=True): """ Initialize the legend_box. The legend_box is an instance of the OffsetBox, which is packed with legend handles and texts. Once packed, their location is calculated during the drawing time. """ fontsize = self._fontsize # legend_box is a HPacker, horizontally packed with # columns. Each column is a VPacker, vertically packed with # legend items. Each legend item is HPacker packed with # legend handleBox and labelBox. handleBox is an instance of # offsetbox.DrawingArea which contains legend handle. labelBox # is an instance of offsetbox.TextArea which contains legend # text. text_list = [] # the list of text instances handle_list = [] # the list of text instances label_prop = dict( verticalalignment='baseline', horizontalalignment='left', fontproperties=self.prop, ) labelboxes = [] handleboxes = [] # The approximate height and descent of text. These values are # only used for plotting the legend handle. descent = 0.35 * self._approx_text_height() * (self.handleheight - 0.7) # 0.35 and 0.7 are just heuristic numbers and may need to be improved. height = self._approx_text_height() * self.handleheight - descent # each handle needs to be drawn inside a box of (x, y, w, h) = # (0, -descent, width, height). And their coordinates should # be given in the display coordinates. # The transformation of each handle will be automatically set # to self.get_trasnform(). If the artist does not use its # default transform (e.g., Collections), you need to # manually set their transform to the self.get_transform(). legend_handler_map = self.get_legend_handler_map() for orig_handle, lab in zip(handles, labels): handler = self.get_legend_handler(legend_handler_map, orig_handle) if handler is None: warnings.warn( "Legend does not support {!r} instances.\nA proxy artist " "may be used instead.\nSee: " "http://matplotlib.org/users/legend_guide.html" "#creating-artists-specifically-for-adding-to-the-legend-" "aka-proxy-artists".format(orig_handle)) # We don't have a handle for this artist, so we just defer # to None. handle_list.append(None) else: textbox = TextArea(lab, textprops=label_prop, multilinebaseline=True, minimumdescent=True) text_list.append(textbox._text) labelboxes.append(textbox) handlebox = DrawingArea(width=self.handlelength * fontsize, height=height, xdescent=0., ydescent=descent) handleboxes.append(handlebox) # Create the artist for the legend which represents the # original artist/handle. handle_list.append( handler.legend_artist(self, orig_handle, fontsize, handlebox)) if handleboxes: # We calculate number of rows in each column. The first # (num_largecol) columns will have (nrows+1) rows, and remaining # (num_smallcol) columns will have (nrows) rows. ncol = min(self._ncol, len(handleboxes)) nrows, num_largecol = divmod(len(handleboxes), ncol) num_smallcol = ncol - num_largecol # starting index of each column and number of rows in it. rows_per_col = [nrows + 1] * num_largecol + [nrows] * num_smallcol start_idxs = np.concatenate([[0], np.cumsum(rows_per_col)[:-1]]) cols = zip(start_idxs, rows_per_col) else: cols = [] handle_label = list(zip(handleboxes, labelboxes)) columnbox = [] for i0, di in cols: # pack handleBox and labelBox into itemBox itemBoxes = [ HPacker(pad=0, sep=self.handletextpad * fontsize, children=[h, t] if markerfirst else [t, h], align="baseline") for h, t in handle_label[i0:i0 + di] ] # minimumdescent=False for the text of the last row of the column if markerfirst: itemBoxes[-1].get_children()[1].set_minimumdescent(False) else: itemBoxes[-1].get_children()[0].set_minimumdescent(False) # pack columnBox alignment = "baseline" if markerfirst else "right" columnbox.append( VPacker(pad=0, sep=self.labelspacing * fontsize, align=alignment, children=itemBoxes)) mode = "expand" if self._mode == "expand" else "fixed" sep = self.columnspacing * fontsize self._legend_handle_box = HPacker(pad=0, sep=sep, align="baseline", mode=mode, children=columnbox) self._legend_title_box = TextArea("") self._legend_box = VPacker( pad=self.borderpad * fontsize, sep=self.labelspacing * fontsize, align="center", children=[self._legend_title_box, self._legend_handle_box]) self._legend_box.set_figure(self.figure) self._legend_box.set_offset(self._findoffset) self.texts = text_list self.legendHandles = handle_list def _auto_legend_data(self): """ Returns list of vertices and extents covered by the plot. Returns a two long list. First element is a list of (x, y) vertices (in display-coordinates) covered by all the lines and line collections, in the legend's handles. Second element is a list of bounding boxes for all the patches in the legend's handles. """ # should always hold because function is only called internally assert self.isaxes ax = self.parent bboxes = [] lines = [] offsets = [] for handle in ax.lines: assert isinstance(handle, Line2D) path = handle.get_path() trans = handle.get_transform() tpath = trans.transform_path(path) lines.append(tpath) for handle in ax.patches: assert isinstance(handle, Patch) if isinstance(handle, Rectangle): transform = handle.get_data_transform() bboxes.append(handle.get_bbox().transformed(transform)) else: transform = handle.get_transform() bboxes.append(handle.get_path().get_extents(transform)) for handle in ax.collections: transform, transOffset, hoffsets, paths = handle._prepare_points() if len(hoffsets): for offset in transOffset.transform(hoffsets): offsets.append(offset) try: vertices = np.concatenate([l.vertices for l in lines]) except ValueError: vertices = np.array([]) return [vertices, bboxes, lines, offsets] def draw_frame(self, b): 'b is a boolean. Set draw frame to b' self.set_frame_on(b) def get_children(self): 'return a list of child artists' children = [] if self._legend_box: children.append(self._legend_box) children.append(self.get_frame()) return children def get_frame(self): 'return the Rectangle instance used to frame the legend' return self.legendPatch def get_lines(self): 'return a list of lines.Line2D instances in the legend' return [h for h in self.legendHandles if isinstance(h, Line2D)] def get_patches(self): 'return a list of patch instances in the legend' return silent_list( 'Patch', [h for h in self.legendHandles if isinstance(h, Patch)]) def get_texts(self): 'return a list of text.Text instance in the legend' return silent_list('Text', self.texts) def set_title(self, title, prop=None): """ set the legend title. Fontproperties can be optionally set with *prop* parameter. """ self._legend_title_box._text.set_text(title) if prop is not None: if isinstance(prop, dict): prop = FontProperties(**prop) self._legend_title_box._text.set_fontproperties(prop) if title: self._legend_title_box.set_visible(True) else: self._legend_title_box.set_visible(False) self.stale = True def get_title(self): 'return Text instance for the legend title' return self._legend_title_box._text def get_window_extent(self, *args, **kwargs): 'return a extent of the legend' return self.legendPatch.get_window_extent(*args, **kwargs) def get_frame_on(self): """ Get whether the legend box patch is drawn """ return self._drawFrame def set_frame_on(self, b): """ Set whether the legend box patch is drawn ACCEPTS: [ *True* | *False* ] """ self._drawFrame = b self.stale = True def get_bbox_to_anchor(self): """ return the bbox that the legend will be anchored """ if self._bbox_to_anchor is None: return self.parent.bbox else: return self._bbox_to_anchor def set_bbox_to_anchor(self, bbox, transform=None): """ set the bbox that the legend will be anchored. *bbox* can be a BboxBase instance, a tuple of [left, bottom, width, height] in the given transform (normalized axes coordinate if None), or a tuple of [left, bottom] where the width and height will be assumed to be zero. """ if bbox is None: self._bbox_to_anchor = None return elif isinstance(bbox, BboxBase): self._bbox_to_anchor = bbox else: try: l = len(bbox) except TypeError: raise ValueError("Invalid argument for bbox : %s" % str(bbox)) if l == 2: bbox = [bbox[0], bbox[1], 0, 0] self._bbox_to_anchor = Bbox.from_bounds(*bbox) if transform is None: transform = BboxTransformTo(self.parent.bbox) self._bbox_to_anchor = TransformedBbox(self._bbox_to_anchor, transform) self.stale = True def _get_anchored_bbox(self, loc, bbox, parentbbox, renderer): """ Place the *bbox* inside the *parentbbox* according to a given location code. Return the (x,y) coordinate of the bbox. - loc: a location code in range(1, 11). This corresponds to the possible values for self._loc, excluding "best". - bbox: bbox to be placed, display coodinate units. - parentbbox: a parent box which will contain the bbox. In display coordinates. """ assert loc in range(1, 11) # called only internally BEST, UR, UL, LL, LR, R, CL, CR, LC, UC, C = list(xrange(11)) anchor_coefs = { UR: "NE", UL: "NW", LL: "SW", LR: "SE", R: "E", CL: "W", CR: "E", LC: "S", UC: "N", C: "C" } c = anchor_coefs[loc] fontsize = renderer.points_to_pixels(self._fontsize) container = parentbbox.padded(-(self.borderaxespad) * fontsize) anchored_box = bbox.anchored(c, container=container) return anchored_box.x0, anchored_box.y0 def _find_best_position(self, width, height, renderer, consider=None): """ Determine the best location to place the legend. `consider` is a list of (x, y) pairs to consider as a potential lower-left corner of the legend. All are display coords. """ # should always hold because function is only called internally assert self.isaxes verts, bboxes, lines, offsets = self._auto_legend_data() bbox = Bbox.from_bounds(0, 0, width, height) if consider is None: consider = [ self._get_anchored_bbox(x, bbox, self.get_bbox_to_anchor(), renderer) for x in range(1, len(self.codes)) ] candidates = [] for idx, (l, b) in enumerate(consider): legendBox = Bbox.from_bounds(l, b, width, height) badness = 0 # XXX TODO: If markers are present, it would be good to # take them into account when checking vertex overlaps in # the next line. badness = (legendBox.count_contains(verts) + legendBox.count_contains(offsets) + legendBox.count_overlaps(bboxes) + sum( line.intersects_bbox(legendBox, filled=False) for line in lines)) if badness == 0: return l, b # Include the index to favor lower codes in case of a tie. candidates.append((badness, idx, (l, b))) _, _, (l, b) = min(candidates) return l, b def contains(self, event): return self.legendPatch.contains(event) def draggable(self, state=None, use_blit=False, update="loc"): """ Set the draggable state -- if state is * None : toggle the current state * True : turn draggable on * False : turn draggable off If draggable is on, you can drag the legend on the canvas with the mouse. The DraggableLegend helper instance is returned if draggable is on. The update parameter control which parameter of the legend changes when dragged. If update is "loc", the *loc* parameter of the legend is changed. If "bbox", the *bbox_to_anchor* parameter is changed. """ is_draggable = self._draggable is not None # if state is None we'll toggle if state is None: state = not is_draggable if state: if self._draggable is None: self._draggable = DraggableLegend(self, use_blit, update=update) else: if self._draggable is not None: self._draggable.disconnect() self._draggable = None return self._draggable
def draw(self): try: if self._mplot_objects: self.clear() self._mplot_objects['text'] = [] UIM = UIManager() controller = UIM.get(self._controller_uid) toc = self.get_parent_controller() index_type = toc.get_well_plot_index_type() # ydata = toc.get_last_dimension_index_data() if ydata is None: return equivalent_ydata = toc.get_equivalent_index_data( datatype=index_type) if equivalent_ydata is None: # There is no equivalent data to be plotted. For example, # we have a TWT index_type and do not have TWT indexed data. # Then, no plot! return # ''' y_min, y_max = np.nanmin(ydata), np.nanmax(ydata) if y_min%controller.step: y_min = (y_min//controller.step + 1) * controller.step y_values = np.arange(y_min, y_max, controller.step) ''' y_min, y_max = np.nanmin(equivalent_ydata), np.nanmax( equivalent_ydata) if y_min % controller.step: y_min = (y_min // controller.step + 1) * controller.step y_values = np.arange(y_min, y_max, controller.step) # # track_controller_uid = UIM._getparentuid(toc_uid) track_controller = self.get_track_controller() # canvas = self.get_canvas() transformated_pos_x = canvas.transform(controller.pos_x, 0.0, 1.0) # # print('y_values:', y_values) for y_value in y_values: # print('\ny_value:', y_value) y_pos_index = (np.abs(equivalent_ydata - y_value)).argmin() y_pos = ydata[y_pos_index] text = track_controller.append_artist( 'Text', transformated_pos_x, y_pos, "%g" % y_value, color=controller.color, horizontalalignment=controller.ha, verticalalignment=controller.va, fontsize=controller.fontsize) if controller.bbox: pad = 0.2 boxstyle = controller.bbox_style boxstyle += ",pad=%0.2f" % pad text._bbox_patch = FancyBboxPatch( (0., 0.), 1., 1., boxstyle=boxstyle, color=controller.bbox_color, alpha=controller.bbox_alpha) #text.zorder = controller.zorder self._mplot_objects['text'].append(text) self.draw_canvas() except Exception as e: print('ERROR IndexRepresentationView.draw', e) raise
class Legend(Artist): """ Place a legend on the axes at location loc. Labels are a sequence of strings and loc can be a string or an integer specifying the legend location The location codes are:: 'best' : 0, (only implemented for axes legends) 'upper right' : 1, 'upper left' : 2, 'lower left' : 3, 'lower right' : 4, 'right' : 5, (same as 'center right', for back-compatibility) 'center left' : 6, 'center right' : 7, 'lower center' : 8, 'upper center' : 9, 'center' : 10, loc can be a tuple of the normalized coordinate values with respect its parent. """ codes = {'best': 0, # only implemented for axes legends 'upper right': 1, 'upper left': 2, 'lower left': 3, 'lower right': 4, 'right': 5, 'center left': 6, 'center right': 7, 'lower center': 8, 'upper center': 9, 'center': 10, } zorder = 5 def __str__(self): return "Legend" def __init__(self, parent, handles, labels, loc=None, numpoints=None, # the number of points in the legend line markerscale=None, # the relative size of legend markers # vs. original markerfirst=True, # controls ordering (left-to-right) of # legend marker and label scatterpoints=None, # number of scatter points scatteryoffsets=None, prop=None, # properties for the legend texts fontsize=None, # keyword to set font size directly # spacing & pad defined as a fraction of the font-size borderpad=None, # the whitespace inside the legend border labelspacing=None, # the vertical space between the legend # entries handlelength=None, # the length of the legend handles handleheight=None, # the height of the legend handles handletextpad=None, # the pad between the legend handle # and text borderaxespad=None, # the pad between the axes and legend # border columnspacing=None, # spacing between columns ncol=1, # number of columns mode=None, # mode for horizontal distribution of columns. # None, "expand" fancybox=None, # True use a fancy box, false use a rounded # box, none use rc shadow=None, title=None, # set a title for the legend framealpha=None, # set frame alpha edgecolor=None, # frame patch edgecolor facecolor=None, # frame patch facecolor bbox_to_anchor=None, # bbox that the legend will be anchored. bbox_transform=None, # transform for the bbox frameon=None, # draw frame handler_map=None, ): """ - *parent*: the artist that contains the legend - *handles*: a list of artists (lines, patches) to be added to the legend - *labels*: a list of strings to label the legend Optional keyword arguments: ================ ==================================================== Keyword Description ================ ==================================================== loc Location code string, or tuple (see below). prop the font property fontsize the font size (used only if prop is not specified) markerscale the relative size of legend markers vs. original markerfirst If True (default), marker is to left of the label. numpoints the number of points in the legend for line scatterpoints the number of points in the legend for scatter plot scatteryoffsets a list of yoffsets for scatter symbols in legend frameon If True, draw the legend on a patch (frame). fancybox If True, draw the frame with a round fancybox. shadow If True, draw a shadow behind legend. framealpha Transparency of the frame. edgecolor Frame edgecolor. facecolor Frame facecolor. ncol number of columns borderpad the fractional whitespace inside the legend border labelspacing the vertical space between the legend entries handlelength the length of the legend handles handleheight the height of the legend handles handletextpad the pad between the legend handle and text borderaxespad the pad between the axes and legend border columnspacing the spacing between columns title the legend title bbox_to_anchor the bbox that the legend will be anchored. bbox_transform the transform for the bbox. transAxes if None. ================ ==================================================== The pad and spacing parameters are measured in font-size units. e.g., a fontsize of 10 points and a handlelength=5 implies a handlelength of 50 points. Values from rcParams will be used if None. Users can specify any arbitrary location for the legend using the *bbox_to_anchor* keyword argument. bbox_to_anchor can be an instance of BboxBase(or its derivatives) or a tuple of 2 or 4 floats. See :meth:`set_bbox_to_anchor` for more detail. The legend location can be specified by setting *loc* with a tuple of 2 floats, which is interpreted as the lower-left corner of the legend in the normalized axes coordinate. """ # local import only to avoid circularity from matplotlib.axes import Axes from matplotlib.figure import Figure Artist.__init__(self) if prop is None: if fontsize is not None: self.prop = FontProperties(size=fontsize) else: self.prop = FontProperties(size=rcParams["legend.fontsize"]) elif isinstance(prop, dict): self.prop = FontProperties(**prop) if "size" not in prop: self.prop.set_size(rcParams["legend.fontsize"]) else: self.prop = prop self._fontsize = self.prop.get_size_in_points() self.texts = [] self.legendHandles = [] self._legend_title_box = None #: A dictionary with the extra handler mappings for this Legend #: instance. self._custom_handler_map = handler_map locals_view = locals() for name in ["numpoints", "markerscale", "shadow", "columnspacing", "scatterpoints", "handleheight", 'borderpad', 'labelspacing', 'handlelength', 'handletextpad', 'borderaxespad']: if locals_view[name] is None: value = rcParams["legend." + name] else: value = locals_view[name] setattr(self, name, value) del locals_view handles = list(handles) if len(handles) < 2: ncol = 1 self._ncol = ncol if self.numpoints <= 0: raise ValueError("numpoints must be > 0; it was %d" % numpoints) # introduce y-offset for handles of the scatter plot if scatteryoffsets is None: self._scatteryoffsets = np.array([3. / 8., 4. / 8., 2.5 / 8.]) else: self._scatteryoffsets = np.asarray(scatteryoffsets) reps = int(self.scatterpoints / len(self._scatteryoffsets)) + 1 self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[:self.scatterpoints] # _legend_box is an OffsetBox instance that contains all # legend items and will be initialized from _init_legend_box() # method. self._legend_box = None if isinstance(parent, Axes): self.isaxes = True self.axes = parent self.set_figure(parent.figure) elif isinstance(parent, Figure): self.isaxes = False self.set_figure(parent) else: raise TypeError("Legend needs either Axes or Figure as parent") self.parent = parent if loc is None: loc = rcParams["legend.loc"] if not self.isaxes and loc in [0, 'best']: loc = 'upper right' if isinstance(loc, six.string_types): if loc not in self.codes: if self.isaxes: warnings.warn('Unrecognized location "%s". Falling back ' 'on "best"; valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes))) loc = 0 else: warnings.warn('Unrecognized location "%s". Falling back ' 'on "upper right"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes))) loc = 1 else: loc = self.codes[loc] if not self.isaxes and loc == 0: warnings.warn('Automatic legend placement (loc="best") not ' 'implemented for figure legend. ' 'Falling back on "upper right".') loc = 1 self._mode = mode self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform) # We use FancyBboxPatch to draw a legend frame. The location # and size of the box will be updated during the drawing time. if facecolor is None: facecolor = rcParams["legend.facecolor"] if facecolor == 'inherit': facecolor = rcParams["axes.facecolor"] if edgecolor is None: edgecolor = rcParams["legend.edgecolor"] if edgecolor == 'inherit': edgecolor = rcParams["axes.edgecolor"] self.legendPatch = FancyBboxPatch( xy=(0.0, 0.0), width=1., height=1., facecolor=facecolor, edgecolor=edgecolor, mutation_scale=self._fontsize, snap=True ) # The width and height of the legendPatch will be set (in the # draw()) to the length that includes the padding. Thus we set # pad=0 here. if fancybox is None: fancybox = rcParams["legend.fancybox"] if fancybox: self.legendPatch.set_boxstyle("round", pad=0, rounding_size=0.2) else: self.legendPatch.set_boxstyle("square", pad=0) self._set_artist_props(self.legendPatch) self._drawFrame = frameon if frameon is None: self._drawFrame = rcParams["legend.frameon"] # init with null renderer self._init_legend_box(handles, labels, markerfirst) if framealpha is None: self.get_frame().set_alpha(rcParams["legend.framealpha"]) else: self.get_frame().set_alpha(framealpha) self._loc = loc self.set_title(title) self._last_fontsize_points = self._fontsize self._draggable = None def _set_artist_props(self, a): """ set the boilerplate props for artists added to axes """ a.set_figure(self.figure) if self.isaxes: # a.set_axes(self.axes) a.axes = self.axes a.set_transform(self.get_transform()) def _set_loc(self, loc): # find_offset function will be provided to _legend_box and # _legend_box will draw itself at the location of the return # value of the find_offset. self._loc_real = loc self.stale = True def _get_loc(self): return self._loc_real _loc = property(_get_loc, _set_loc) def _findoffset(self, width, height, xdescent, ydescent, renderer): "Helper function to locate the legend" if self._loc == 0: # "best". x, y = self._find_best_position(width, height, renderer) elif self._loc in Legend.codes.values(): # Fixed location. bbox = Bbox.from_bounds(0, 0, width, height) x, y = self._get_anchored_bbox(self._loc, bbox, self.get_bbox_to_anchor(), renderer) else: # Axes or figure coordinates. fx, fy = self._loc bbox = self.get_bbox_to_anchor() x, y = bbox.x0 + bbox.width * fx, bbox.y0 + bbox.height * fy return x + xdescent, y + ydescent @allow_rasterization def draw(self, renderer): "Draw everything that belongs to the legend" if not self.get_visible(): return renderer.open_group('legend') fontsize = renderer.points_to_pixels(self._fontsize) # if mode == fill, set the width of the legend_box to the # width of the paret (minus pads) if self._mode in ["expand"]: pad = 2 * (self.borderaxespad + self.borderpad) * fontsize self._legend_box.set_width(self.get_bbox_to_anchor().width - pad) # update the location and size of the legend. This needs to # be done in any case to clip the figure right. bbox = self._legend_box.get_window_extent(renderer) self.legendPatch.set_bounds(bbox.x0, bbox.y0, bbox.width, bbox.height) self.legendPatch.set_mutation_scale(fontsize) if self._drawFrame: if self.shadow: shadow = Shadow(self.legendPatch, 2, -2) shadow.draw(renderer) self.legendPatch.draw(renderer) self._legend_box.draw(renderer) renderer.close_group('legend') self.stale = False def _approx_text_height(self, renderer=None): """ Return the approximate height of the text. This is used to place the legend handle. """ if renderer is None: return self._fontsize else: return renderer.points_to_pixels(self._fontsize) # _default_handler_map defines the default mapping between plot # elements and the legend handlers. _default_handler_map = { StemContainer: legend_handler.HandlerStem(), ErrorbarContainer: legend_handler.HandlerErrorbar(), Line2D: legend_handler.HandlerLine2D(), Patch: legend_handler.HandlerPatch(), LineCollection: legend_handler.HandlerLineCollection(), RegularPolyCollection: legend_handler.HandlerRegularPolyCollection(), CircleCollection: legend_handler.HandlerCircleCollection(), BarContainer: legend_handler.HandlerPatch( update_func=legend_handler.update_from_first_child), tuple: legend_handler.HandlerTuple(), PathCollection: legend_handler.HandlerPathCollection(), PolyCollection: legend_handler.HandlerPolyCollection() } # (get|set|update)_default_handler_maps are public interfaces to # modify the default handler map. @classmethod def get_default_handler_map(cls): """ A class method that returns the default handler map. """ return cls._default_handler_map @classmethod def set_default_handler_map(cls, handler_map): """ A class method to set the default handler map. """ cls._default_handler_map = handler_map @classmethod def update_default_handler_map(cls, handler_map): """ A class method to update the default handler map. """ cls._default_handler_map.update(handler_map) def get_legend_handler_map(self): """ return the handler map. """ default_handler_map = self.get_default_handler_map() if self._custom_handler_map: hm = default_handler_map.copy() hm.update(self._custom_handler_map) return hm else: return default_handler_map @staticmethod def get_legend_handler(legend_handler_map, orig_handle): """ return a legend handler from *legend_handler_map* that corresponds to *orig_handler*. *legend_handler_map* should be a dictionary object (that is returned by the get_legend_handler_map method). It first checks if the *orig_handle* itself is a key in the *legend_hanler_map* and return the associated value. Otherwise, it checks for each of the classes in its method-resolution-order. If no matching key is found, it returns None. """ if is_hashable(orig_handle): try: return legend_handler_map[orig_handle] except KeyError: pass for handle_type in type(orig_handle).mro(): try: return legend_handler_map[handle_type] except KeyError: pass return None def _init_legend_box(self, handles, labels, markerfirst=True): """ Initialize the legend_box. The legend_box is an instance of the OffsetBox, which is packed with legend handles and texts. Once packed, their location is calculated during the drawing time. """ fontsize = self._fontsize # legend_box is a HPacker, horizontally packed with # columns. Each column is a VPacker, vertically packed with # legend items. Each legend item is HPacker packed with # legend handleBox and labelBox. handleBox is an instance of # offsetbox.DrawingArea which contains legend handle. labelBox # is an instance of offsetbox.TextArea which contains legend # text. text_list = [] # the list of text instances handle_list = [] # the list of text instances label_prop = dict(verticalalignment='baseline', horizontalalignment='left', fontproperties=self.prop, ) labelboxes = [] handleboxes = [] # The approximate height and descent of text. These values are # only used for plotting the legend handle. descent = 0.35 * self._approx_text_height() * (self.handleheight - 0.7) # 0.35 and 0.7 are just heuristic numbers and may need to be improved. height = self._approx_text_height() * self.handleheight - descent # each handle needs to be drawn inside a box of (x, y, w, h) = # (0, -descent, width, height). And their coordinates should # be given in the display coordinates. # The transformation of each handle will be automatically set # to self.get_trasnform(). If the artist does not use its # default transform (e.g., Collections), you need to # manually set their transform to the self.get_transform(). legend_handler_map = self.get_legend_handler_map() for orig_handle, lab in zip(handles, labels): handler = self.get_legend_handler(legend_handler_map, orig_handle) if handler is None: warnings.warn( "Legend does not support {!r} instances.\nA proxy artist " "may be used instead.\nSee: " "http://matplotlib.org/users/legend_guide.html" "#using-proxy-artist".format(orig_handle) ) # We don't have a handle for this artist, so we just defer # to None. handle_list.append(None) else: textbox = TextArea(lab, textprops=label_prop, multilinebaseline=True, minimumdescent=True) text_list.append(textbox._text) labelboxes.append(textbox) handlebox = DrawingArea(width=self.handlelength * fontsize, height=height, xdescent=0., ydescent=descent) handleboxes.append(handlebox) # Create the artist for the legend which represents the # original artist/handle. handle_list.append(handler.legend_artist(self, orig_handle, fontsize, handlebox)) if handleboxes: # We calculate number of rows in each column. The first # (num_largecol) columns will have (nrows+1) rows, and remaining # (num_smallcol) columns will have (nrows) rows. ncol = min(self._ncol, len(handleboxes)) nrows, num_largecol = divmod(len(handleboxes), ncol) num_smallcol = ncol - num_largecol # starting index of each column and number of rows in it. rows_per_col = [nrows + 1] * num_largecol + [nrows] * num_smallcol start_idxs = np.concatenate([[0], np.cumsum(rows_per_col)[:-1]]) cols = zip(start_idxs, rows_per_col) else: cols = [] handle_label = list(zip(handleboxes, labelboxes)) columnbox = [] for i0, di in cols: # pack handleBox and labelBox into itemBox itemBoxes = [HPacker(pad=0, sep=self.handletextpad * fontsize, children=[h, t] if markerfirst else [t, h], align="baseline") for h, t in handle_label[i0:i0 + di]] # minimumdescent=False for the text of the last row of the column if markerfirst: itemBoxes[-1].get_children()[1].set_minimumdescent(False) else: itemBoxes[-1].get_children()[0].set_minimumdescent(False) # pack columnBox alignment = "baseline" if markerfirst else "right" columnbox.append(VPacker(pad=0, sep=self.labelspacing * fontsize, align=alignment, children=itemBoxes)) mode = "expand" if self._mode == "expand" else "fixed" sep = self.columnspacing * fontsize self._legend_handle_box = HPacker(pad=0, sep=sep, align="baseline", mode=mode, children=columnbox) self._legend_title_box = TextArea("") self._legend_box = VPacker(pad=self.borderpad * fontsize, sep=self.labelspacing * fontsize, align="center", children=[self._legend_title_box, self._legend_handle_box]) self._legend_box.set_figure(self.figure) self._legend_box.set_offset(self._findoffset) self.texts = text_list self.legendHandles = handle_list def _auto_legend_data(self): """ Returns list of vertices and extents covered by the plot. Returns a two long list. First element is a list of (x, y) vertices (in display-coordinates) covered by all the lines and line collections, in the legend's handles. Second element is a list of bounding boxes for all the patches in the legend's handles. """ # should always hold because function is only called internally assert self.isaxes ax = self.parent bboxes = [] lines = [] offsets = [] for handle in ax.lines: assert isinstance(handle, Line2D) path = handle.get_path() trans = handle.get_transform() tpath = trans.transform_path(path) lines.append(tpath) for handle in ax.patches: assert isinstance(handle, Patch) if isinstance(handle, Rectangle): transform = handle.get_data_transform() bboxes.append(handle.get_bbox().transformed(transform)) else: transform = handle.get_transform() bboxes.append(handle.get_path().get_extents(transform)) for handle in ax.collections: transform, transOffset, hoffsets, paths = handle._prepare_points() if len(hoffsets): for offset in transOffset.transform(hoffsets): offsets.append(offset) try: vertices = np.concatenate([l.vertices for l in lines]) except ValueError: vertices = np.array([]) return [vertices, bboxes, lines, offsets] def draw_frame(self, b): 'b is a boolean. Set draw frame to b' self.set_frame_on(b) def get_children(self): 'return a list of child artists' children = [] if self._legend_box: children.append(self._legend_box) children.append(self.get_frame()) return children def get_frame(self): 'return the Rectangle instance used to frame the legend' return self.legendPatch def get_lines(self): 'return a list of lines.Line2D instances in the legend' return [h for h in self.legendHandles if isinstance(h, Line2D)] def get_patches(self): 'return a list of patch instances in the legend' return silent_list('Patch', [h for h in self.legendHandles if isinstance(h, Patch)]) def get_texts(self): 'return a list of text.Text instance in the legend' return silent_list('Text', self.texts) def set_title(self, title, prop=None): """ set the legend title. Fontproperties can be optionally set with *prop* parameter. """ self._legend_title_box._text.set_text(title) if prop is not None: if isinstance(prop, dict): prop = FontProperties(**prop) self._legend_title_box._text.set_fontproperties(prop) if title: self._legend_title_box.set_visible(True) else: self._legend_title_box.set_visible(False) self.stale = True def get_title(self): 'return Text instance for the legend title' return self._legend_title_box._text def get_window_extent(self, *args, **kwargs): 'return a extent of the legend' return self.legendPatch.get_window_extent(*args, **kwargs) def get_frame_on(self): """ Get whether the legend box patch is drawn """ return self._drawFrame def set_frame_on(self, b): """ Set whether the legend box patch is drawn ACCEPTS: [ *True* | *False* ] """ self._drawFrame = b self.stale = True def get_bbox_to_anchor(self): """ return the bbox that the legend will be anchored """ if self._bbox_to_anchor is None: return self.parent.bbox else: return self._bbox_to_anchor def set_bbox_to_anchor(self, bbox, transform=None): """ set the bbox that the legend will be anchored. *bbox* can be a BboxBase instance, a tuple of [left, bottom, width, height] in the given transform (normalized axes coordinate if None), or a tuple of [left, bottom] where the width and height will be assumed to be zero. """ if bbox is None: self._bbox_to_anchor = None return elif isinstance(bbox, BboxBase): self._bbox_to_anchor = bbox else: try: l = len(bbox) except TypeError: raise ValueError("Invalid argument for bbox : %s" % str(bbox)) if l == 2: bbox = [bbox[0], bbox[1], 0, 0] self._bbox_to_anchor = Bbox.from_bounds(*bbox) if transform is None: transform = BboxTransformTo(self.parent.bbox) self._bbox_to_anchor = TransformedBbox(self._bbox_to_anchor, transform) self.stale = True def _get_anchored_bbox(self, loc, bbox, parentbbox, renderer): """ Place the *bbox* inside the *parentbbox* according to a given location code. Return the (x,y) coordinate of the bbox. - loc: a location code in range(1, 11). This corresponds to the possible values for self._loc, excluding "best". - bbox: bbox to be placed, display coodinate units. - parentbbox: a parent box which will contain the bbox. In display coordinates. """ assert loc in range(1, 11) # called only internally BEST, UR, UL, LL, LR, R, CL, CR, LC, UC, C = list(xrange(11)) anchor_coefs = {UR: "NE", UL: "NW", LL: "SW", LR: "SE", R: "E", CL: "W", CR: "E", LC: "S", UC: "N", C: "C"} c = anchor_coefs[loc] fontsize = renderer.points_to_pixels(self._fontsize) container = parentbbox.padded(-(self.borderaxespad) * fontsize) anchored_box = bbox.anchored(c, container=container) return anchored_box.x0, anchored_box.y0 def _find_best_position(self, width, height, renderer, consider=None): """ Determine the best location to place the legend. `consider` is a list of (x, y) pairs to consider as a potential lower-left corner of the legend. All are display coords. """ # should always hold because function is only called internally assert self.isaxes verts, bboxes, lines, offsets = self._auto_legend_data() bbox = Bbox.from_bounds(0, 0, width, height) if consider is None: consider = [self._get_anchored_bbox(x, bbox, self.get_bbox_to_anchor(), renderer) for x in range(1, len(self.codes))] candidates = [] for idx, (l, b) in enumerate(consider): legendBox = Bbox.from_bounds(l, b, width, height) badness = 0 # XXX TODO: If markers are present, it would be good to # take them into account when checking vertex overlaps in # the next line. badness = (legendBox.count_contains(verts) + legendBox.count_contains(offsets) + legendBox.count_overlaps(bboxes) + sum(line.intersects_bbox(legendBox, filled=False) for line in lines)) if badness == 0: return l, b # Include the index to favor lower codes in case of a tie. candidates.append((badness, idx, (l, b))) _, _, (l, b) = min(candidates) return l, b def contains(self, event): return self.legendPatch.contains(event) def draggable(self, state=None, use_blit=False, update="loc"): """ Set the draggable state -- if state is * None : toggle the current state * True : turn draggable on * False : turn draggable off If draggable is on, you can drag the legend on the canvas with the mouse. The DraggableLegend helper instance is returned if draggable is on. The update parameter control which parameter of the legend changes when dragged. If update is "loc", the *loc* parameter of the legend is changed. If "bbox", the *bbox_to_anchor* parameter is changed. """ is_draggable = self._draggable is not None # if state is None we'll toggle if state is None: state = not is_draggable if state: if self._draggable is None: self._draggable = DraggableLegend(self, use_blit, update=update) else: if self._draggable is not None: self._draggable.disconnect() self._draggable = None return self._draggable
def plot_parallel_dist(): fig = plt.figure(figsize=(fig_size_x, fig_size_y), dpi=dpi, constrained_layout=True) fig.set_constrained_layout_pads(w_pad=0.1, h_pad=0) ax = plt.gca() size_lvl = [] left_lower = [] boxes = [] colors = list(mcolors.BASE_COLORS) if 'k' in colors: colors.remove('k') if time_procs > len(colors): for i in range(time_procs - len(colors)): colors.append(list(np.random.rand(3))) nt = len(self.problem[0].t) split = self.split_into(number_points=nt, number_processes=time_procs) for i in range(len(split)): left_lower.append([np.sum(split[:i]), -0.2]) for i in range(self.lvl_max): for j in range(len(split)): if i != 0: boxes.append( FancyBboxPatch( [ a + b for a, b in zip(left_lower[j], [-0.25, i]) ], split[j] - 0.25, 0.4, boxstyle= "round,pad=-0.0040,rounding_size=0.015", fc=colors[j], label='Process ' + str(j))) else: boxes.append( FancyBboxPatch( [ a + b for a, b in zip(left_lower[j], [-0.25, i]) ], np.sum(split) - 0.25, 0.4, boxstyle= "round,pad=-0.0040,rounding_size=0.015", fc=colors[j], label='Process ' + str(j))) break for item in boxes: ax.add_patch(item) legend = plt.legend(handles=boxes[-time_procs:], fontsize=text_size, loc='center left', bbox_to_anchor=(1, 0.5), prop={ 'weight': 'bold', 'size': text_size }) for i in range(self.lvl_max): if i != self.lvl_max - 1: tmp_coarse = np.where( np.in1d(self.problem[0].t, self.problem[i + 1].t))[0] tmp_fine = np.where( np.in1d(self.problem[0].t, self.problem[i].t))[0] plt.plot([0, nt - 1], [i, i], color='k') for j in range(nt): if j in tmp_coarse: plt.plot([j, j], [((self.lvl_max - 1) - i) + 0.1, ((self.lvl_max - 1) - i) - 0.1], color='k') elif j in tmp_fine: plt.plot([j, j], [((self.lvl_max - 1) - i) + 0.05, ((self.lvl_max - 1) - i) - 0.05], color='k') [s.set_visible(False) for s in ax.spines.values()] ax.set_xticks([]) ax.set_yticks(np.arange(0, self.lvl_max, 1).tolist()) ax.set_yticklabels( ['Level ' + str(i) for i in range(self.lvl_max - 1, -1, -1)]) plt.yticks(size=text_size, weight='bold') if save_name is not None: plt.savefig(save_name, dpi=dpi) plt.show()
class Legend(Artist): """ Place a legend on the axes at location loc. Labels are a sequence of strings and loc can be a string or an integer specifying the legend location The location codes are:: 'best' : 0, (only implemented for axis legends) 'upper right' : 1, 'upper left' : 2, 'lower left' : 3, 'lower right' : 4, 'right' : 5, 'center left' : 6, 'center right' : 7, 'lower center' : 8, 'upper center' : 9, 'center' : 10, loc can be a tuple of the noramilzed coordinate values with respect its parent. Return value is a sequence of text, line instances that make up the legend """ codes = {'best' : 0, # only implemented for axis legends 'upper right' : 1, 'upper left' : 2, 'lower left' : 3, 'lower right' : 4, 'right' : 5, 'center left' : 6, 'center right' : 7, 'lower center' : 8, 'upper center' : 9, 'center' : 10, } zorder = 5 def __str__(self): return "Legend" def __init__(self, parent, handles, labels, loc = None, numpoints = None, # the number of points in the legend line markerscale = None, # the relative size of legend markers vs. original scatterpoints = 3, # TODO: may be an rcParam scatteryoffsets=None, prop = None, # properties for the legend texts pad = None, # deprecated; use borderpad labelsep = None, # deprecated; use labelspacing handlelen = None, # deprecated; use handlelength handletextsep = None, # deprecated; use handletextpad axespad = None, # deprecated; use borderaxespad borderpad = None, # the whitespace inside the legend border labelspacing=None, #the vertical space between the legend entries handlelength=None, # the length of the legend handles handletextpad=None, # the pad between the legend handle and text borderaxespad=None, # the pad between the axes and legend border columnspacing=None, # spacing between columns ncol=1, # number of columns mode=None, # mode for horizontal distribution of columns. None, "expand" fancybox=None, # True use a fancy box, false use a rounded box, none use rc shadow = None, title = None, # set a title for the legend bbox_to_anchor = None, # bbox that the legend will be anchored. bbox_transform = None, # transform for the bbox frameon = True, # draw frame ): """ - *parent* : the artist that contains the legend - *handles* : a list of artists (lines, patches) to add to the legend - *labels* : a list of strings to label the legend Optional keyword arguments: ================ ================================================================== Keyword Description ================ ================================================================== loc a location code prop the font property markerscale the relative size of legend markers vs. original numpoints the number of points in the legend for line scatterpoints the number of points in the legend for scatter plot scatteryoffsets a list of yoffsets for scatter symbols in legend frameon if True, draw a frame (default is True) fancybox if True, draw a frame with a round fancybox. If None, use rc shadow if True, draw a shadow behind legend ncol number of columns borderpad the fractional whitespace inside the legend border labelspacing the vertical space between the legend entries handlelength the length of the legend handles handletextpad the pad between the legend handle and text borderaxespad the pad between the axes and legend border columnspacing the spacing between columns title the legend title bbox_to_anchor the bbox that the legend will be anchored. bbox_transform the transform for the bbox. transAxes if None. ================ ================================================================== The pad and spacing parameters are measure in font-size units. E.g., a fontsize of 10 points and a handlelength=5 implies a handlelength of 50 points. Values from rcParams will be used if None. Users can specify any arbitrary location for the legend using the *bbox_to_anchor* keyword argument. bbox_to_anchor can be an instance of BboxBase(or its derivatives) or a tuple of 2 or 4 floats. See :meth:`set_bbox_to_anchor` for more detail. The legend location can be specified by setting *loc* with a tuple of 2 floats, which is interpreted as the lower-left corner of the legend in the normalized axes coordinate. """ from matplotlib.axes import Axes # local import only to avoid circularity from matplotlib.figure import Figure # local import only to avoid circularity Artist.__init__(self) if prop is None: self.prop=FontProperties(size=rcParams["legend.fontsize"]) elif isinstance(prop, dict): self.prop=FontProperties(**prop) if "size" not in prop: self.prop.set_size(rcParams["legend.fontsize"]) else: self.prop=prop self._fontsize = self.prop.get_size_in_points() propnames=['numpoints', 'markerscale', 'shadow', "columnspacing", "scatterpoints"] self.texts = [] self.legendHandles = [] self._legend_title_box = None localdict = locals() for name in propnames: if localdict[name] is None: value = rcParams["legend."+name] else: value = localdict[name] setattr(self, name, value) deprecated_kwds = {"pad":"borderpad", "labelsep":"labelspacing", "handlelen":"handlelength", "handletextsep":"handletextpad", "axespad":"borderaxespad"} bbox = parent.bbox axessize_fontsize = min(bbox.width, bbox.height)/self._fontsize for k, v in deprecated_kwds.items(): if localdict[k] is not None and localdict[v] is None: warnings.warn("Use '%s' instead of '%s'." % (v, k), DeprecationWarning) setattr(self, v, localdict[k]*axessize_fontsize) continue if localdict[v] is None: setattr(self, v, rcParams["legend."+v]) else: setattr(self, v, localdict[v]) del localdict handles = list(handles) if len(handles)<2: ncol = 1 self._ncol = ncol if self.numpoints <= 0: raise ValueError("numpoints must be >= 0; it was %d"% numpoints) if scatteryoffsets is None: self._scatteryoffsets = np.array([3./8., 4./8., 2.5/8.]) else: self._scatteryoffsets = np.asarray(scatteryoffsets) reps = int(self.numpoints / len(self._scatteryoffsets)) + 1 self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[:self.scatterpoints] self._legend_box = None if isinstance(parent,Axes): self.isaxes = True self.set_axes(parent) self.set_figure(parent.figure) elif isinstance(parent,Figure): self.isaxes = False self.set_figure(parent) else: raise TypeError("Legend needs either Axes or Figure as parent") self.parent = parent if loc is None: loc = rcParams["legend.loc"] if not self.isaxes and loc in [0,'best']: loc = 'upper right' if is_string_like(loc): if loc not in self.codes: if self.isaxes: warnings.warn('Unrecognized location "%s". Falling back on "best"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes.keys()))) loc = 0 else: warnings.warn('Unrecognized location "%s". Falling back on "upper right"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes.keys()))) loc = 1 else: loc = self.codes[loc] if not self.isaxes and loc == 0: warnings.warn('Automatic legend placement (loc="best") not implemented for figure legend. ' 'Falling back on "upper right".') loc = 1 self._mode = mode self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform) self.legendPatch = FancyBboxPatch( xy=(0.0, 0.0), width=1., height=1., facecolor=rcParams["axes.facecolor"], edgecolor=rcParams["axes.edgecolor"], mutation_scale=self._fontsize, snap=True ) if fancybox is None: fancybox = rcParams["legend.fancybox"] if fancybox == True: self.legendPatch.set_boxstyle("round",pad=0, rounding_size=0.2) else: self.legendPatch.set_boxstyle("square",pad=0) self._set_artist_props(self.legendPatch) self._drawFrame = frameon self._init_legend_box(handles, labels) self._loc = loc self.set_title(title) self._last_fontsize_points = self._fontsize self._draggable = None def _set_artist_props(self, a): """ set the boilerplate props for artists added to axes """ a.set_figure(self.figure) if self.isaxes: a.set_axes(self.axes) a.set_transform(self.get_transform()) def _set_loc(self, loc): self._loc_real = loc if loc == 0: _findoffset = self._findoffset_best else: _findoffset = self._findoffset_loc self._legend_box.set_offset(_findoffset) self._loc_real = loc def _get_loc(self): return self._loc_real _loc = property(_get_loc, _set_loc) def _findoffset_best(self, width, height, xdescent, ydescent, renderer): "Helper function to locate the legend at its best position" ox, oy = self._find_best_position(width, height, renderer) return ox+xdescent, oy+ydescent def _findoffset_loc(self, width, height, xdescent, ydescent, renderer): "Heper function to locate the legend using the location code" if iterable(self._loc) and len(self._loc)==2: fx, fy = self._loc bbox = self.get_bbox_to_anchor() x, y = bbox.x0 + bbox.width * fx, bbox.y0 + bbox.height * fy else: bbox = Bbox.from_bounds(0, 0, width, height) x, y = self._get_anchored_bbox(self._loc, bbox, self.get_bbox_to_anchor(), renderer) return x+xdescent, y+ydescent @allow_rasterization def draw(self, renderer): "Draw everything that belongs to the legend" if not self.get_visible(): return renderer.open_group('legend') fontsize = renderer.points_to_pixels(self._fontsize) if self._mode in ["expand"]: pad = 2*(self.borderaxespad+self.borderpad)*fontsize self._legend_box.set_width(self.get_bbox_to_anchor().width-pad) if self._drawFrame: bbox = self._legend_box.get_window_extent(renderer) self.legendPatch.set_bounds(bbox.x0, bbox.y0, bbox.width, bbox.height) self.legendPatch.set_mutation_scale(fontsize) if self.shadow: shadow = Shadow(self.legendPatch, 2, -2) shadow.draw(renderer) self.legendPatch.draw(renderer) self._legend_box.draw(renderer) renderer.close_group('legend') def _approx_text_height(self, renderer=None): """ Return the approximate height of the text. This is used to place the legend handle. """ if renderer is None: return self._fontsize else: return renderer.points_to_pixels(self._fontsize) def _init_legend_box(self, handles, labels): """ Initiallize the legend_box. The legend_box is an instance of the OffsetBox, which is packed with legend handles and texts. Once packed, their location is calculated during the drawing time. """ fontsize = self._fontsize text_list = [] # the list of text instances handle_list = [] # the list of text instances label_prop = dict(verticalalignment='baseline', horizontalalignment='left', fontproperties=self.prop, ) labelboxes = [] handleboxes = [] height = self._approx_text_height() * 0.7 descent = 0. for handle, lab in zip(handles, labels): if isinstance(handle, RegularPolyCollection) or \ isinstance(handle, CircleCollection): npoints = self.scatterpoints else: npoints = self.numpoints if npoints > 1: xdata = np.linspace(0.3*fontsize, (self.handlelength-0.3)*fontsize, npoints) xdata_marker = xdata elif npoints == 1: xdata = np.linspace(0, self.handlelength*fontsize, 2) xdata_marker = [0.5*self.handlelength*fontsize] if isinstance(handle, Line2D): ydata = ((height-descent)/2.)*np.ones(xdata.shape, float) legline = Line2D(xdata, ydata) legline.update_from(handle) self._set_artist_props(legline) # after update legline.set_clip_box(None) legline.set_clip_path(None) legline.set_drawstyle('default') legline.set_marker('None') handle_list.append(legline) legline_marker = Line2D(xdata_marker, ydata[:len(xdata_marker)]) legline_marker.update_from(handle) self._set_artist_props(legline_marker) legline_marker.set_clip_box(None) legline_marker.set_clip_path(None) legline_marker.set_linestyle('None') legline._legmarker = legline_marker elif isinstance(handle, Patch): p = Rectangle(xy=(0., 0.), width = self.handlelength*fontsize, height=(height-descent), ) p.update_from(handle) self._set_artist_props(p) p.set_clip_box(None) p.set_clip_path(None) handle_list.append(p) elif isinstance(handle, LineCollection): ydata = ((height-descent)/2.)*np.ones(xdata.shape, float) legline = Line2D(xdata, ydata) self._set_artist_props(legline) legline.set_clip_box(None) legline.set_clip_path(None) lw = handle.get_linewidth()[0] dashes = handle.get_dashes()[0] color = handle.get_colors()[0] legline.set_color(color) legline.set_linewidth(lw) if dashes[0] is not None: # dashed line legline.set_dashes(dashes[1]) handle_list.append(legline) elif isinstance(handle, RegularPolyCollection): ydata = height*self._scatteryoffsets size_max, size_min = max(handle.get_sizes()),\ min(handle.get_sizes()) if self.scatterpoints < 4: sizes = [.5*(size_max+size_min), size_max, size_min] else: sizes = (size_max-size_min)*np.linspace(0,1,self.scatterpoints)+size_min p = type(handle)(handle.get_numsides(), rotation=handle.get_rotation(), sizes=sizes, offsets=zip(xdata_marker,ydata), transOffset=self.get_transform(), ) p.update_from(handle) p.set_figure(self.figure) p.set_clip_box(None) p.set_clip_path(None) handle_list.append(p) elif isinstance(handle, CircleCollection): ydata = height*self._scatteryoffsets size_max, size_min = max(handle.get_sizes()),\ min(handle.get_sizes()) if self.scatterpoints < 4: sizes = [.5*(size_max+size_min), size_max, size_min] else: sizes = (size_max-size_min)*np.linspace(0,1,self.scatterpoints)+size_min p = type(handle)(sizes, offsets=zip(xdata_marker,ydata), transOffset=self.get_transform(), ) p.update_from(handle) p.set_figure(self.figure) p.set_clip_box(None) p.set_clip_path(None) handle_list.append(p) else: handle_type = type(handle) warnings.warn("Legend does not support %s\nUse proxy artist instead.\n\nhttp://matplotlib.sourceforge.net/users/legend_guide.html#using-proxy-artist\n" % (str(handle_type),)) handle_list.append(None) handle = handle_list[-1] if handle is not None: # handle is None is the artist is not supproted textbox = TextArea(lab, textprops=label_prop, multilinebaseline=True, minimumdescent=True) text_list.append(textbox._text) labelboxes.append(textbox) handlebox = DrawingArea(width=self.handlelength*fontsize, height=height, xdescent=0., ydescent=descent) handlebox.add_artist(handle) if isinstance(handle, RegularPolyCollection) or \ isinstance(handle, CircleCollection): handle._transOffset = handlebox.get_transform() handle.set_transform(None) if hasattr(handle, "_legmarker"): handlebox.add_artist(handle._legmarker) handleboxes.append(handlebox) if len(handleboxes) > 0: ncol = min(self._ncol, len(handleboxes)) nrows, num_largecol = divmod(len(handleboxes), ncol) num_smallcol = ncol-num_largecol largecol = safezip(range(0, num_largecol*(nrows+1), (nrows+1)), [nrows+1] * num_largecol) smallcol = safezip(range(num_largecol*(nrows+1), len(handleboxes), nrows), [nrows] * num_smallcol) else: largecol, smallcol = [], [] handle_label = safezip(handleboxes, labelboxes) columnbox = [] for i0, di in largecol+smallcol: itemBoxes = [HPacker(pad=0, sep=self.handletextpad*fontsize, children=[h, t], align="baseline") for h, t in handle_label[i0:i0+di]] itemBoxes[-1].get_children()[1].set_minimumdescent(False) columnbox.append(VPacker(pad=0, sep=self.labelspacing*fontsize, align="baseline", children=itemBoxes)) if self._mode == "expand": mode = "expand" else: mode = "fixed" sep = self.columnspacing*fontsize self._legend_handle_box = HPacker(pad=0, sep=sep, align="baseline", mode=mode, children=columnbox) self._legend_title_box = TextArea("") self._legend_box = VPacker(pad=self.borderpad*fontsize, sep=self.labelspacing*fontsize, align="center", children=[self._legend_title_box, self._legend_handle_box]) self._legend_box.set_figure(self.figure) self.texts = text_list self.legendHandles = handle_list def _auto_legend_data(self): """ Returns list of vertices and extents covered by the plot. Returns a two long list. First element is a list of (x, y) vertices (in display-coordinates) covered by all the lines and line collections, in the legend's handles. Second element is a list of bounding boxes for all the patches in the legend's handles. """ assert self.isaxes # should always hold because function is only called internally ax = self.parent vertices = [] bboxes = [] lines = [] for handle in ax.lines: assert isinstance(handle, Line2D) path = handle.get_path() trans = handle.get_transform() tpath = trans.transform_path(path) lines.append(tpath) for handle in ax.patches: assert isinstance(handle, Patch) if isinstance(handle, Rectangle): transform = handle.get_data_transform() bboxes.append(handle.get_bbox().transformed(transform)) else: transform = handle.get_transform() bboxes.append(handle.get_path().get_extents(transform)) return [vertices, bboxes, lines] def draw_frame(self, b): 'b is a boolean. Set draw frame to b' self.set_frame_on(b) def get_children(self): 'return a list of child artists' children = [] if self._legend_box: children.append(self._legend_box) children.extend(self.get_lines()) children.extend(self.get_patches()) children.extend(self.get_texts()) children.append(self.get_frame()) if self._legend_title_box: children.append(self.get_title()) return children def get_frame(self): 'return the Rectangle instance used to frame the legend' return self.legendPatch def get_lines(self): 'return a list of lines.Line2D instances in the legend' return [h for h in self.legendHandles if isinstance(h, Line2D)] def get_patches(self): 'return a list of patch instances in the legend' return silent_list('Patch', [h for h in self.legendHandles if isinstance(h, Patch)]) def get_texts(self): 'return a list of text.Text instance in the legend' return silent_list('Text', self.texts) def set_title(self, title): 'set the legend title' self._legend_title_box._text.set_text(title) if title: self._legend_title_box.set_visible(True) else: self._legend_title_box.set_visible(False) def get_title(self): 'return Text instance for the legend title' return self._legend_title_box._text def get_window_extent(self): 'return a extent of the the legend' return self.legendPatch.get_window_extent() def get_frame_on(self): """ Get whether the legend box patch is drawn """ return self._drawFrame def set_frame_on(self, b): """ Set whether the legend box patch is drawn ACCEPTS: [ *True* | *False* ] """ self._drawFrame = b def get_bbox_to_anchor(self): """ return the bbox that the legend will be anchored """ if self._bbox_to_anchor is None: return self.parent.bbox else: return self._bbox_to_anchor def set_bbox_to_anchor(self, bbox, transform=None): """ set the bbox that the legend will be anchored. *bbox* can be a BboxBase instance, a tuple of [left, bottom, width, height] in the given transform (normalized axes coordinate if None), or a tuple of [left, bottom] where the width and height will be assumed to be zero. """ if bbox is None: self._bbox_to_anchor = None return elif isinstance(bbox, BboxBase): self._bbox_to_anchor = bbox else: try: l = len(bbox) except TypeError: raise ValueError("Invalid argument for bbox : %s" % str(bbox)) if l == 2: bbox = [bbox[0], bbox[1], 0, 0] self._bbox_to_anchor = Bbox.from_bounds(*bbox) if transform is None: transform = BboxTransformTo(self.parent.bbox) self._bbox_to_anchor = TransformedBbox(self._bbox_to_anchor, transform) def _get_anchored_bbox(self, loc, bbox, parentbbox, renderer): """ Place the *bbox* inside the *parentbbox* according to a given location code. Return the (x,y) coordinate of the bbox. - loc: a location code in range(1, 11). This corresponds to the possible values for self._loc, excluding "best". - bbox: bbox to be placed, display coodinate units. - parentbbox: a parent box which will contain the bbox. In display coordinates. """ assert loc in range(1,11) # called only internally BEST, UR, UL, LL, LR, R, CL, CR, LC, UC, C = range(11) anchor_coefs={UR:"NE", UL:"NW", LL:"SW", LR:"SE", R:"E", CL:"W", CR:"E", LC:"S", UC:"N", C:"C"} c = anchor_coefs[loc] fontsize = renderer.points_to_pixels(self._fontsize) container = parentbbox.padded(-(self.borderaxespad) * fontsize) anchored_box = bbox.anchored(c, container=container) return anchored_box.x0, anchored_box.y0 def _find_best_position(self, width, height, renderer, consider=None): """ Determine the best location to place the legend. `consider` is a list of (x, y) pairs to consider as a potential lower-left corner of the legend. All are display coords. """ assert self.isaxes # should always hold because function is only called internally verts, bboxes, lines = self._auto_legend_data() bbox = Bbox.from_bounds(0, 0, width, height) consider = [self._get_anchored_bbox(x, bbox, self.get_bbox_to_anchor(), renderer) for x in range(1, len(self.codes))] candidates = [] for l, b in consider: legendBox = Bbox.from_bounds(l, b, width, height) badness = 0 badness = legendBox.count_contains(verts) badness += legendBox.count_overlaps(bboxes) for line in lines: if line.intersects_bbox(legendBox): badness += 1 ox, oy = l, b if badness == 0: return ox, oy candidates.append((badness, (l, b))) minCandidate = candidates[0] for candidate in candidates: if candidate[0] < minCandidate[0]: minCandidate = candidate ox, oy = minCandidate[1] return ox, oy def draggable(self, state=None, use_blit=False): """ Set the draggable state -- if state is * None : toggle the current state * True : turn draggable on * False : turn draggable off If draggable is on, you can drag the legend on the canvas with the mouse. The DraggableLegend helper instance is returned if draggable is on. """ is_draggable = self._draggable is not None if state is None: state = not is_draggable if state: if self._draggable is None: self._draggable = DraggableLegend(self, use_blit) else: if self._draggable is not None: self._draggable.disconnect() self._draggable = None return self._draggable
def __init__(self, loc, pad=0.4, borderpad=0.5, child=None, prop=None, frameon=True, bbox_to_anchor=None, bbox_transform=None, **kwargs): """ loc is a string or an integer specifying the legend location. The valid location codes are:: 'upper right' : 1, 'upper left' : 2, 'lower left' : 3, 'lower right' : 4, 'right' : 5, 'center left' : 6, 'center right' : 7, 'lower center' : 8, 'upper center' : 9, 'center' : 10, pad : pad around the child for drawing a frame. given in fraction of fontsize. borderpad : pad between offsetbox frame and the bbox_to_anchor, child : OffsetBox instance that will be anchored. prop : font property. This is only used as a reference for paddings. frameon : draw a frame box if True. bbox_to_anchor : bbox to anchor. Use self.axes.bbox if None. bbox_transform : with which the bbox_to_anchor will be transformed. """ super(AnchoredOffsetbox, self).__init__(**kwargs) self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform) self.set_child(child) self.loc = loc self.borderpad = borderpad self.pad = pad if prop is None: self.prop = FontProperties(size=rcParams["legend.fontsize"]) elif isinstance(prop, dict): self.prop = FontProperties(**prop) if "size" not in prop: self.prop.set_size(rcParams["legend.fontsize"]) else: self.prop = prop self.patch = FancyBboxPatch( xy=(0.0, 0.0), width=1., height=1., facecolor='w', edgecolor='k', mutation_scale=self.prop.get_size_in_points(), snap=True) self.patch.set_boxstyle("square", pad=0) self._drawFrame = frameon
def __init__( self, parent, handles, labels, loc=None, numpoints=None, # the number of points in the legend line markerscale=None, # the relative size of legend markers # vs. original markerfirst=True, # controls ordering (left-to-right) of # legend marker and label scatterpoints=None, # number of scatter points scatteryoffsets=None, prop=None, # properties for the legend texts fontsize=None, # keyword to set font size directly # spacing & pad defined as a fraction of the font-size borderpad=None, # the whitespace inside the legend border labelspacing=None, # the vertical space between the legend # entries handlelength=None, # the length of the legend handles handleheight=None, # the height of the legend handles handletextpad=None, # the pad between the legend handle # and text borderaxespad=None, # the pad between the axes and legend # border columnspacing=None, # spacing between columns ncol=1, # number of columns mode=None, # mode for horizontal distribution of columns. # None, "expand" fancybox=None, # True use a fancy box, false use a rounded # box, none use rc shadow=None, title=None, # set a title for the legend framealpha=None, # set frame alpha edgecolor=None, # frame patch edgecolor facecolor=None, # frame patch facecolor bbox_to_anchor=None, # bbox that the legend will be anchored. bbox_transform=None, # transform for the bbox frameon=None, # draw frame handler_map=None, ): """ - *parent*: the artist that contains the legend - *handles*: a list of artists (lines, patches) to be added to the legend - *labels*: a list of strings to label the legend Optional keyword arguments: ================ ==================================================== Keyword Description ================ ==================================================== loc Location code string, or tuple (see below). prop the font property fontsize the font size (used only if prop is not specified) markerscale the relative size of legend markers vs. original markerfirst If True (default), marker is to left of the label. numpoints the number of points in the legend for line scatterpoints the number of points in the legend for scatter plot scatteryoffsets a list of yoffsets for scatter symbols in legend frameon If True, draw the legend on a patch (frame). fancybox If True, draw the frame with a round fancybox. shadow If True, draw a shadow behind legend. framealpha Transparency of the frame. edgecolor Frame edgecolor. facecolor Frame facecolor. ncol number of columns borderpad the fractional whitespace inside the legend border labelspacing the vertical space between the legend entries handlelength the length of the legend handles handleheight the height of the legend handles handletextpad the pad between the legend handle and text borderaxespad the pad between the axes and legend border columnspacing the spacing between columns title the legend title bbox_to_anchor the bbox that the legend will be anchored. bbox_transform the transform for the bbox. transAxes if None. ================ ==================================================== The pad and spacing parameters are measured in font-size units. e.g., a fontsize of 10 points and a handlelength=5 implies a handlelength of 50 points. Values from rcParams will be used if None. Users can specify any arbitrary location for the legend using the *bbox_to_anchor* keyword argument. bbox_to_anchor can be an instance of BboxBase(or its derivatives) or a tuple of 2 or 4 floats. See :meth:`set_bbox_to_anchor` for more detail. The legend location can be specified by setting *loc* with a tuple of 2 floats, which is interpreted as the lower-left corner of the legend in the normalized axes coordinate. """ # local import only to avoid circularity from matplotlib.axes import Axes from matplotlib.figure import Figure Artist.__init__(self) if prop is None: if fontsize is not None: self.prop = FontProperties(size=fontsize) else: self.prop = FontProperties(size=rcParams["legend.fontsize"]) elif isinstance(prop, dict): self.prop = FontProperties(**prop) if "size" not in prop: self.prop.set_size(rcParams["legend.fontsize"]) else: self.prop = prop self._fontsize = self.prop.get_size_in_points() self.texts = [] self.legendHandles = [] self._legend_title_box = None #: A dictionary with the extra handler mappings for this Legend #: instance. self._custom_handler_map = handler_map locals_view = locals() for name in [ "numpoints", "markerscale", "shadow", "columnspacing", "scatterpoints", "handleheight", 'borderpad', 'labelspacing', 'handlelength', 'handletextpad', 'borderaxespad' ]: if locals_view[name] is None: value = rcParams["legend." + name] else: value = locals_view[name] setattr(self, name, value) del locals_view handles = list(handles) if len(handles) < 2: ncol = 1 self._ncol = ncol if self.numpoints <= 0: raise ValueError("numpoints must be > 0; it was %d" % numpoints) # introduce y-offset for handles of the scatter plot if scatteryoffsets is None: self._scatteryoffsets = np.array([3. / 8., 4. / 8., 2.5 / 8.]) else: self._scatteryoffsets = np.asarray(scatteryoffsets) reps = int(self.scatterpoints / len(self._scatteryoffsets)) + 1 self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[:self.scatterpoints] # _legend_box is an OffsetBox instance that contains all # legend items and will be initialized from _init_legend_box() # method. self._legend_box = None if isinstance(parent, Axes): self.isaxes = True self.axes = parent self.set_figure(parent.figure) elif isinstance(parent, Figure): self.isaxes = False self.set_figure(parent) else: raise TypeError("Legend needs either Axes or Figure as parent") self.parent = parent if loc is None: loc = rcParams["legend.loc"] if not self.isaxes and loc in [0, 'best']: loc = 'upper right' if is_string_like(loc): if loc not in self.codes: if self.isaxes: warnings.warn('Unrecognized location "%s". Falling back ' 'on "best"; valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes))) loc = 0 else: warnings.warn('Unrecognized location "%s". Falling back ' 'on "upper right"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes))) loc = 1 else: loc = self.codes[loc] if not self.isaxes and loc == 0: warnings.warn('Automatic legend placement (loc="best") not ' 'implemented for figure legend. ' 'Falling back on "upper right".') loc = 1 self._mode = mode self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform) # We use FancyBboxPatch to draw a legend frame. The location # and size of the box will be updated during the drawing time. if facecolor is None: facecolor = rcParams["legend.facecolor"] if facecolor == 'inherit': facecolor = rcParams["axes.facecolor"] if edgecolor is None: edgecolor = rcParams["legend.edgecolor"] if edgecolor == 'inherit': edgecolor = rcParams["axes.edgecolor"] self.legendPatch = FancyBboxPatch(xy=(0.0, 0.0), width=1., height=1., facecolor=facecolor, edgecolor=edgecolor, mutation_scale=self._fontsize, snap=True) # The width and height of the legendPatch will be set (in the # draw()) to the length that includes the padding. Thus we set # pad=0 here. if fancybox is None: fancybox = rcParams["legend.fancybox"] if fancybox: self.legendPatch.set_boxstyle("round", pad=0, rounding_size=0.2) else: self.legendPatch.set_boxstyle("square", pad=0) self._set_artist_props(self.legendPatch) self._drawFrame = frameon if frameon is None: self._drawFrame = rcParams["legend.frameon"] # init with null renderer self._init_legend_box(handles, labels, markerfirst) if framealpha is None: self.get_frame().set_alpha(rcParams["legend.framealpha"]) else: self.get_frame().set_alpha(framealpha) self._loc = loc self.set_title(title) self._last_fontsize_points = self._fontsize self._draggable = None
class Legend(Artist): """ Place a legend on the axes at location loc. """ codes = { 'best': 0, # only implemented for axes legends 'upper right': 1, 'upper left': 2, 'lower left': 3, 'lower right': 4, 'right': 5, 'center left': 6, 'center right': 7, 'lower center': 8, 'upper center': 9, 'center': 10, } zorder = 5 def __str__(self): return "Legend" @docstring.dedent_interpd def __init__( self, parent, handles, labels, loc=None, numpoints=None, # the number of points in the legend line markerscale=None, # the relative size of legend markers # vs. original markerfirst=True, # controls ordering (left-to-right) of # legend marker and label scatterpoints=None, # number of scatter points scatteryoffsets=None, prop=None, # properties for the legend texts fontsize=None, # keyword to set font size directly labelcolor=None, # keyword to set the text color # spacing & pad defined as a fraction of the font-size borderpad=None, # the whitespace inside the legend border labelspacing=None, # the vertical space between the legend # entries handlelength=None, # the length of the legend handles handleheight=None, # the height of the legend handles handletextpad=None, # the pad between the legend handle # and text borderaxespad=None, # the pad between the axes and legend # border columnspacing=None, # spacing between columns ncol=1, # number of columns mode=None, # mode for horizontal distribution of columns. # None, "expand" fancybox=None, # True use a fancy box, false use a rounded # box, none use rc shadow=None, title=None, # set a title for the legend title_fontsize=None, # the font size for the title framealpha=None, # set frame alpha edgecolor=None, # frame patch edgecolor facecolor=None, # frame patch facecolor bbox_to_anchor=None, # bbox that the legend will be anchored. bbox_transform=None, # transform for the bbox frameon=None, # draw frame handler_map=None, ): """ Parameters ---------- parent : `~matplotlib.axes.Axes` or `.Figure` The artist that contains the legend. handles : list of `.Artist` A list of Artists (lines, patches) to be added to the legend. labels : list of str A list of labels to show next to the artists. The length of handles and labels should be the same. If they are not, they are truncated to the smaller of both lengths. Other Parameters ---------------- %(_legend_kw_doc)s Notes ----- Users can specify any arbitrary location for the legend using the *bbox_to_anchor* keyword argument. *bbox_to_anchor* can be a `.BboxBase` (or derived therefrom) or a tuple of 2 or 4 floats. See `set_bbox_to_anchor` for more detail. The legend location can be specified by setting *loc* with a tuple of 2 floats, which is interpreted as the lower-left corner of the legend in the normalized axes coordinate. """ # local import only to avoid circularity from matplotlib.axes import Axes from matplotlib.figure import Figure Artist.__init__(self) if prop is None: if fontsize is not None: self.prop = FontProperties(size=fontsize) else: self.prop = FontProperties( size=mpl.rcParams["legend.fontsize"]) else: self.prop = FontProperties._from_any(prop) if isinstance(prop, dict) and "size" not in prop: self.prop.set_size(mpl.rcParams["legend.fontsize"]) self._fontsize = self.prop.get_size_in_points() self.texts = [] self.legendHandles = [] self._legend_title_box = None #: A dictionary with the extra handler mappings for this Legend #: instance. self._custom_handler_map = handler_map locals_view = locals() for name in [ "numpoints", "markerscale", "shadow", "columnspacing", "scatterpoints", "handleheight", 'borderpad', 'labelspacing', 'handlelength', 'handletextpad', 'borderaxespad' ]: if locals_view[name] is None: value = mpl.rcParams["legend." + name] else: value = locals_view[name] setattr(self, name, value) del locals_view # trim handles and labels if illegal label... _lab, _hand = [], [] for label, handle in zip(labels, handles): if isinstance(label, str) and label.startswith('_'): cbook._warn_external('The label {!r} of handle {!r} cannot be ' 'a string starting with ' '"_"'.format(label, handle)) else: _lab.append(label) _hand.append(handle) labels, handles = _lab, _hand handles = list(handles) if len(handles) < 2: ncol = 1 self._ncol = ncol if self.numpoints <= 0: raise ValueError("numpoints must be > 0; it was %d" % numpoints) # introduce y-offset for handles of the scatter plot if scatteryoffsets is None: self._scatteryoffsets = np.array([3. / 8., 4. / 8., 2.5 / 8.]) else: self._scatteryoffsets = np.asarray(scatteryoffsets) reps = self.scatterpoints // len(self._scatteryoffsets) + 1 self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[:self.scatterpoints] # _legend_box is a VPacker instance that contains all # legend items and will be initialized from _init_legend_box() # method. self._legend_box = None if isinstance(parent, Axes): self.isaxes = True self.axes = parent self.set_figure(parent.figure) elif isinstance(parent, Figure): self.isaxes = False self.set_figure(parent) else: raise TypeError("Legend needs either Axes or Figure as parent") self.parent = parent self._loc_used_default = loc is None if loc is None: loc = mpl.rcParams["legend.loc"] if not self.isaxes and loc in [0, 'best']: loc = 'upper right' if isinstance(loc, str): if loc not in self.codes: raise ValueError( "Unrecognized location {!r}. Valid locations are\n\t{}\n". format(loc, '\n\t'.join(self.codes))) else: loc = self.codes[loc] if not self.isaxes and loc == 0: raise ValueError( "Automatic legend placement (loc='best') not implemented for " "figure legend.") self._mode = mode self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform) # We use FancyBboxPatch to draw a legend frame. The location # and size of the box will be updated during the drawing time. if facecolor is None: facecolor = mpl.rcParams["legend.facecolor"] if facecolor == 'inherit': facecolor = mpl.rcParams["axes.facecolor"] if edgecolor is None: edgecolor = mpl.rcParams["legend.edgecolor"] if edgecolor == 'inherit': edgecolor = mpl.rcParams["axes.edgecolor"] if fancybox is None: fancybox = mpl.rcParams["legend.fancybox"] self.legendPatch = FancyBboxPatch( xy=(0, 0), width=1, height=1, facecolor=facecolor, edgecolor=edgecolor, # If shadow is used, default to alpha=1 (#8943). alpha=(framealpha if framealpha is not None else 1 if shadow else mpl.rcParams["legend.framealpha"]), # The width and height of the legendPatch will be set (in draw()) # to the length that includes the padding. Thus we set pad=0 here. boxstyle=("round,pad=0,rounding_size=0.2" if fancybox else "square,pad=0"), mutation_scale=self._fontsize, snap=True, visible=(frameon if frameon is not None else mpl.rcParams["legend.frameon"])) self._set_artist_props(self.legendPatch) # init with null renderer self._init_legend_box(handles, labels, markerfirst) tmp = self._loc_used_default self._set_loc(loc) self._loc_used_default = tmp # ignore changes done by _set_loc # figure out title fontsize: if title_fontsize is None: title_fontsize = mpl.rcParams['legend.title_fontsize'] tprop = FontProperties(size=title_fontsize) self.set_title(title, prop=tprop) self._draggable = None # set the text color color_getters = { # getter function depends on line or patch 'linecolor': ['get_color', 'get_facecolor'], 'markerfacecolor': ['get_markerfacecolor', 'get_facecolor'], 'mfc': ['get_markerfacecolor', 'get_facecolor'], 'markeredgecolor': ['get_markeredgecolor', 'get_edgecolor'], 'mec': ['get_markeredgecolor', 'get_edgecolor'], } if labelcolor is None: pass elif isinstance(labelcolor, str) and labelcolor in color_getters: getter_names = color_getters[labelcolor] for handle, text in zip(self.legendHandles, self.texts): for getter_name in getter_names: try: color = getattr(handle, getter_name)() text.set_color(color) break except AttributeError: pass elif np.iterable(labelcolor): for text, color in zip( self.texts, itertools.cycle(colors.to_rgba_array(labelcolor))): text.set_color(color) else: raise ValueError("Invalid argument for labelcolor : %s" % str(labelcolor)) def _set_artist_props(self, a): """ Set the boilerplate props for artists added to axes. """ a.set_figure(self.figure) if self.isaxes: # a.set_axes(self.axes) a.axes = self.axes a.set_transform(self.get_transform()) def _set_loc(self, loc): # find_offset function will be provided to _legend_box and # _legend_box will draw itself at the location of the return # value of the find_offset. self._loc_used_default = False self._loc_real = loc self.stale = True self._legend_box.set_offset(self._findoffset) def _get_loc(self): return self._loc_real _loc = property(_get_loc, _set_loc) def _findoffset(self, width, height, xdescent, ydescent, renderer): """Helper function to locate the legend.""" if self._loc == 0: # "best". x, y = self._find_best_position(width, height, renderer) elif self._loc in Legend.codes.values(): # Fixed location. bbox = Bbox.from_bounds(0, 0, width, height) x, y = self._get_anchored_bbox(self._loc, bbox, self.get_bbox_to_anchor(), renderer) else: # Axes or figure coordinates. fx, fy = self._loc bbox = self.get_bbox_to_anchor() x, y = bbox.x0 + bbox.width * fx, bbox.y0 + bbox.height * fy return x + xdescent, y + ydescent @allow_rasterization def draw(self, renderer): # docstring inherited if not self.get_visible(): return renderer.open_group('legend', gid=self.get_gid()) fontsize = renderer.points_to_pixels(self._fontsize) # if mode == fill, set the width of the legend_box to the # width of the parent (minus pads) if self._mode in ["expand"]: pad = 2 * (self.borderaxespad + self.borderpad) * fontsize self._legend_box.set_width(self.get_bbox_to_anchor().width - pad) # update the location and size of the legend. This needs to # be done in any case to clip the figure right. bbox = self._legend_box.get_window_extent(renderer) self.legendPatch.set_bounds(bbox.x0, bbox.y0, bbox.width, bbox.height) self.legendPatch.set_mutation_scale(fontsize) if self.shadow: Shadow(self.legendPatch, 2, -2).draw(renderer) self.legendPatch.draw(renderer) self._legend_box.draw(renderer) renderer.close_group('legend') self.stale = False # _default_handler_map defines the default mapping between plot # elements and the legend handlers. _default_handler_map = { StemContainer: legend_handler.HandlerStem(), ErrorbarContainer: legend_handler.HandlerErrorbar(), Line2D: legend_handler.HandlerLine2D(), Patch: legend_handler.HandlerPatch(), LineCollection: legend_handler.HandlerLineCollection(), RegularPolyCollection: legend_handler.HandlerRegularPolyCollection(), CircleCollection: legend_handler.HandlerCircleCollection(), BarContainer: legend_handler.HandlerPatch( update_func=legend_handler.update_from_first_child), tuple: legend_handler.HandlerTuple(), PathCollection: legend_handler.HandlerPathCollection(), PolyCollection: legend_handler.HandlerPolyCollection() } # (get|set|update)_default_handler_maps are public interfaces to # modify the default handler map. @classmethod def get_default_handler_map(cls): """ A class method that returns the default handler map. """ return cls._default_handler_map @classmethod def set_default_handler_map(cls, handler_map): """ A class method to set the default handler map. """ cls._default_handler_map = handler_map @classmethod def update_default_handler_map(cls, handler_map): """ A class method to update the default handler map. """ cls._default_handler_map.update(handler_map) def get_legend_handler_map(self): """ Return the handler map. """ default_handler_map = self.get_default_handler_map() if self._custom_handler_map: hm = default_handler_map.copy() hm.update(self._custom_handler_map) return hm else: return default_handler_map @staticmethod def get_legend_handler(legend_handler_map, orig_handle): """ Return a legend handler from *legend_handler_map* that corresponds to *orig_handler*. *legend_handler_map* should be a dictionary object (that is returned by the get_legend_handler_map method). It first checks if the *orig_handle* itself is a key in the *legend_handler_map* and return the associated value. Otherwise, it checks for each of the classes in its method-resolution-order. If no matching key is found, it returns ``None``. """ try: return legend_handler_map[orig_handle] except (TypeError, KeyError): # TypeError if unhashable. pass for handle_type in type(orig_handle).mro(): try: return legend_handler_map[handle_type] except KeyError: pass return None def _init_legend_box(self, handles, labels, markerfirst=True): """ Initialize the legend_box. The legend_box is an instance of the OffsetBox, which is packed with legend handles and texts. Once packed, their location is calculated during the drawing time. """ fontsize = self._fontsize # legend_box is a HPacker, horizontally packed with # columns. Each column is a VPacker, vertically packed with # legend items. Each legend item is HPacker packed with # legend handleBox and labelBox. handleBox is an instance of # offsetbox.DrawingArea which contains legend handle. labelBox # is an instance of offsetbox.TextArea which contains legend # text. text_list = [] # the list of text instances handle_list = [] # the list of text instances handles_and_labels = [] label_prop = dict( verticalalignment='baseline', horizontalalignment='left', fontproperties=self.prop, ) # The approximate height and descent of text. These values are # only used for plotting the legend handle. descent = 0.35 * fontsize * (self.handleheight - 0.7) # 0.35 and 0.7 are just heuristic numbers and may need to be improved. height = fontsize * self.handleheight - descent # each handle needs to be drawn inside a box of (x, y, w, h) = # (0, -descent, width, height). And their coordinates should # be given in the display coordinates. # The transformation of each handle will be automatically set # to self.get_transform(). If the artist does not use its # default transform (e.g., Collections), you need to # manually set their transform to the self.get_transform(). legend_handler_map = self.get_legend_handler_map() for orig_handle, lab in zip(handles, labels): handler = self.get_legend_handler(legend_handler_map, orig_handle) if handler is None: cbook._warn_external( "Legend does not support {!r} instances.\nA proxy artist " "may be used instead.\nSee: " "https://matplotlib.org/users/legend_guide.html" "#creating-artists-specifically-for-adding-to-the-legend-" "aka-proxy-artists".format(orig_handle)) # We don't have a handle for this artist, so we just defer # to None. handle_list.append(None) else: textbox = TextArea(lab, textprops=label_prop, multilinebaseline=True, minimumdescent=True) handlebox = DrawingArea(width=self.handlelength * fontsize, height=height, xdescent=0., ydescent=descent) text_list.append(textbox._text) # Create the artist for the legend which represents the # original artist/handle. handle_list.append( handler.legend_artist(self, orig_handle, fontsize, handlebox)) handles_and_labels.append((handlebox, textbox)) if handles_and_labels: # We calculate number of rows in each column. The first # (num_largecol) columns will have (nrows+1) rows, and remaining # (num_smallcol) columns will have (nrows) rows. ncol = min(self._ncol, len(handles_and_labels)) nrows, num_largecol = divmod(len(handles_and_labels), ncol) num_smallcol = ncol - num_largecol # starting index of each column and number of rows in it. rows_per_col = [nrows + 1] * num_largecol + [nrows] * num_smallcol start_idxs = np.concatenate([[0], np.cumsum(rows_per_col)[:-1]]) cols = zip(start_idxs, rows_per_col) else: cols = [] columnbox = [] for i0, di in cols: # pack handleBox and labelBox into itemBox itemBoxes = [ HPacker(pad=0, sep=self.handletextpad * fontsize, children=[h, t] if markerfirst else [t, h], align="baseline") for h, t in handles_and_labels[i0:i0 + di] ] # minimumdescent=False for the text of the last row of the column if markerfirst: itemBoxes[-1].get_children()[1].set_minimumdescent(False) else: itemBoxes[-1].get_children()[0].set_minimumdescent(False) # pack columnBox alignment = "baseline" if markerfirst else "right" columnbox.append( VPacker(pad=0, sep=self.labelspacing * fontsize, align=alignment, children=itemBoxes)) mode = "expand" if self._mode == "expand" else "fixed" sep = self.columnspacing * fontsize self._legend_handle_box = HPacker(pad=0, sep=sep, align="baseline", mode=mode, children=columnbox) self._legend_title_box = TextArea("") self._legend_box = VPacker( pad=self.borderpad * fontsize, sep=self.labelspacing * fontsize, align="center", children=[self._legend_title_box, self._legend_handle_box]) self._legend_box.set_figure(self.figure) self.texts = text_list self.legendHandles = handle_list def _auto_legend_data(self): """ Return display coordinates for hit testing for "best" positioning. Returns ------- bboxes List of bounding boxes of all patches. lines List of `.Path` corresponding to each line. offsets List of (x, y) offsets of all collection. """ assert self.isaxes # always holds, as this is only called internally ax = self.parent lines = [ line.get_transform().transform_path(line.get_path()) for line in ax.lines ] bboxes = [ patch.get_bbox().transformed(patch.get_data_transform()) if isinstance(patch, Rectangle) else patch.get_path().get_extents( patch.get_transform()) for patch in ax.patches ] offsets = [] for handle in ax.collections: _, transOffset, hoffsets, _ = handle._prepare_points() for offset in transOffset.transform(hoffsets): offsets.append(offset) return bboxes, lines, offsets def get_children(self): # docstring inherited return [self._legend_box, self.get_frame()] def get_frame(self): """Return the `~.patches.Rectangle` used to frame the legend.""" return self.legendPatch def get_lines(self): r"""Return the list of `~.lines.Line2D`\s in the legend.""" return [h for h in self.legendHandles if isinstance(h, Line2D)] def get_patches(self): r"""Return the list of `~.patches.Patch`\s in the legend.""" return silent_list( 'Patch', [h for h in self.legendHandles if isinstance(h, Patch)]) def get_texts(self): r"""Return the list of `~.text.Text`\s in the legend.""" return silent_list('Text', self.texts) def set_title(self, title, prop=None): """ Set the legend title. Fontproperties can be optionally set with *prop* parameter. """ self._legend_title_box._text.set_text(title) if title: self._legend_title_box._text.set_visible(True) self._legend_title_box.set_visible(True) else: self._legend_title_box._text.set_visible(False) self._legend_title_box.set_visible(False) if prop is not None: self._legend_title_box._text.set_fontproperties(prop) self.stale = True def get_title(self): """Return the `.Text` instance for the legend title.""" return self._legend_title_box._text def get_window_extent(self, renderer=None): # docstring inherited if renderer is None: renderer = self.figure._cachedRenderer return self._legend_box.get_window_extent(renderer=renderer) def get_tightbbox(self, renderer): """ Like `.Legend.get_window_extent`, but uses the box for the legend. Parameters ---------- renderer : `.RendererBase` subclass renderer that will be used to draw the figures (i.e. ``fig.canvas.get_renderer()``) Returns ------- `.BboxBase` The bounding box in figure pixel coordinates. """ return self._legend_box.get_window_extent(renderer) def get_frame_on(self): """Get whether the legend box patch is drawn.""" return self.legendPatch.get_visible() def set_frame_on(self, b): """ Set whether the legend box patch is drawn. Parameters ---------- b : bool """ self.legendPatch.set_visible(b) self.stale = True draw_frame = set_frame_on # Backcompat alias. def get_bbox_to_anchor(self): """Return the bbox that the legend will be anchored to.""" if self._bbox_to_anchor is None: return self.parent.bbox else: return self._bbox_to_anchor def set_bbox_to_anchor(self, bbox, transform=None): """ Set the bbox that the legend will be anchored to. *bbox* can be - A `.BboxBase` instance - A tuple of ``(left, bottom, width, height)`` in the given transform (normalized axes coordinate if None) - A tuple of ``(left, bottom)`` where the width and height will be assumed to be zero. """ if bbox is None: self._bbox_to_anchor = None return elif isinstance(bbox, BboxBase): self._bbox_to_anchor = bbox else: try: l = len(bbox) except TypeError as err: raise ValueError("Invalid argument for bbox : %s" % str(bbox)) from err if l == 2: bbox = [bbox[0], bbox[1], 0, 0] self._bbox_to_anchor = Bbox.from_bounds(*bbox) if transform is None: transform = BboxTransformTo(self.parent.bbox) self._bbox_to_anchor = TransformedBbox(self._bbox_to_anchor, transform) self.stale = True def _get_anchored_bbox(self, loc, bbox, parentbbox, renderer): """ Place the *bbox* inside the *parentbbox* according to a given location code. Return the (x, y) coordinate of the bbox. - loc: a location code in range(1, 11). This corresponds to the possible values for self._loc, excluding "best". - bbox: bbox to be placed, display coordinate units. - parentbbox: a parent box which will contain the bbox. In display coordinates. """ assert loc in range(1, 11) # called only internally BEST, UR, UL, LL, LR, R, CL, CR, LC, UC, C = range(11) anchor_coefs = { UR: "NE", UL: "NW", LL: "SW", LR: "SE", R: "E", CL: "W", CR: "E", LC: "S", UC: "N", C: "C" } c = anchor_coefs[loc] fontsize = renderer.points_to_pixels(self._fontsize) container = parentbbox.padded(-self.borderaxespad * fontsize) anchored_box = bbox.anchored(c, container=container) return anchored_box.x0, anchored_box.y0 def _find_best_position(self, width, height, renderer, consider=None): """ Determine the best location to place the legend. *consider* is a list of ``(x, y)`` pairs to consider as a potential lower-left corner of the legend. All are display coords. """ assert self.isaxes # always holds, as this is only called internally start_time = time.perf_counter() bboxes, lines, offsets = self._auto_legend_data() bbox = Bbox.from_bounds(0, 0, width, height) if consider is None: consider = [ self._get_anchored_bbox(x, bbox, self.get_bbox_to_anchor(), renderer) for x in range(1, len(self.codes)) ] candidates = [] for idx, (l, b) in enumerate(consider): legendBox = Bbox.from_bounds(l, b, width, height) badness = 0 # XXX TODO: If markers are present, it would be good to take them # into account when checking vertex overlaps in the next line. badness = ( sum(legendBox.count_contains(line.vertices) for line in lines) + legendBox.count_contains(offsets) + legendBox.count_overlaps(bboxes) + sum( line.intersects_bbox(legendBox, filled=False) for line in lines)) if badness == 0: return l, b # Include the index to favor lower codes in case of a tie. candidates.append((badness, idx, (l, b))) _, _, (l, b) = min(candidates) if self._loc_used_default and time.perf_counter() - start_time > 1: cbook._warn_external( 'Creating legend with loc="best" can be slow with large ' 'amounts of data.') return l, b def contains(self, event): inside, info = self._default_contains(event) if inside is not None: return inside, info return self.legendPatch.contains(event) def set_draggable(self, state, use_blit=False, update='loc'): """ Enable or disable mouse dragging support of the legend. Parameters ---------- state : bool Whether mouse dragging is enabled. use_blit : bool, optional Use blitting for faster image composition. For details see :ref:`func-animation`. update : {'loc', 'bbox'}, optional The legend parameter to be changed when dragged: - 'loc': update the *loc* parameter of the legend - 'bbox': update the *bbox_to_anchor* parameter of the legend Returns ------- If *state* is ``True`` this returns the `~.DraggableLegend` helper instance. Otherwise this returns ``None``. """ if state: if self._draggable is None: self._draggable = DraggableLegend(self, use_blit, update=update) else: if self._draggable is not None: self._draggable.disconnect() self._draggable = None return self._draggable def get_draggable(self): """Return ``True`` if the legend is draggable, ``False`` otherwise.""" return self._draggable is not None
def __init__( self, parent, handles, labels, loc=None, numpoints=None, # the number of points in the legend line markerscale=None, # the relative size of legend markers # vs. original markerfirst=True, # controls ordering (left-to-right) of # legend marker and label scatterpoints=None, # number of scatter points scatteryoffsets=None, prop=None, # properties for the legend texts fontsize=None, # keyword to set font size directly labelcolor=None, # keyword to set the text color # spacing & pad defined as a fraction of the font-size borderpad=None, # the whitespace inside the legend border labelspacing=None, # the vertical space between the legend # entries handlelength=None, # the length of the legend handles handleheight=None, # the height of the legend handles handletextpad=None, # the pad between the legend handle # and text borderaxespad=None, # the pad between the axes and legend # border columnspacing=None, # spacing between columns ncol=1, # number of columns mode=None, # mode for horizontal distribution of columns. # None, "expand" fancybox=None, # True use a fancy box, false use a rounded # box, none use rc shadow=None, title=None, # set a title for the legend title_fontsize=None, # the font size for the title framealpha=None, # set frame alpha edgecolor=None, # frame patch edgecolor facecolor=None, # frame patch facecolor bbox_to_anchor=None, # bbox that the legend will be anchored. bbox_transform=None, # transform for the bbox frameon=None, # draw frame handler_map=None, ): """ Parameters ---------- parent : `~matplotlib.axes.Axes` or `.Figure` The artist that contains the legend. handles : list of `.Artist` A list of Artists (lines, patches) to be added to the legend. labels : list of str A list of labels to show next to the artists. The length of handles and labels should be the same. If they are not, they are truncated to the smaller of both lengths. Other Parameters ---------------- %(_legend_kw_doc)s Notes ----- Users can specify any arbitrary location for the legend using the *bbox_to_anchor* keyword argument. *bbox_to_anchor* can be a `.BboxBase` (or derived therefrom) or a tuple of 2 or 4 floats. See `set_bbox_to_anchor` for more detail. The legend location can be specified by setting *loc* with a tuple of 2 floats, which is interpreted as the lower-left corner of the legend in the normalized axes coordinate. """ # local import only to avoid circularity from matplotlib.axes import Axes from matplotlib.figure import Figure Artist.__init__(self) if prop is None: if fontsize is not None: self.prop = FontProperties(size=fontsize) else: self.prop = FontProperties( size=mpl.rcParams["legend.fontsize"]) else: self.prop = FontProperties._from_any(prop) if isinstance(prop, dict) and "size" not in prop: self.prop.set_size(mpl.rcParams["legend.fontsize"]) self._fontsize = self.prop.get_size_in_points() self.texts = [] self.legendHandles = [] self._legend_title_box = None #: A dictionary with the extra handler mappings for this Legend #: instance. self._custom_handler_map = handler_map locals_view = locals() for name in [ "numpoints", "markerscale", "shadow", "columnspacing", "scatterpoints", "handleheight", 'borderpad', 'labelspacing', 'handlelength', 'handletextpad', 'borderaxespad' ]: if locals_view[name] is None: value = mpl.rcParams["legend." + name] else: value = locals_view[name] setattr(self, name, value) del locals_view # trim handles and labels if illegal label... _lab, _hand = [], [] for label, handle in zip(labels, handles): if isinstance(label, str) and label.startswith('_'): cbook._warn_external('The label {!r} of handle {!r} cannot be ' 'a string starting with ' '"_"'.format(label, handle)) else: _lab.append(label) _hand.append(handle) labels, handles = _lab, _hand handles = list(handles) if len(handles) < 2: ncol = 1 self._ncol = ncol if self.numpoints <= 0: raise ValueError("numpoints must be > 0; it was %d" % numpoints) # introduce y-offset for handles of the scatter plot if scatteryoffsets is None: self._scatteryoffsets = np.array([3. / 8., 4. / 8., 2.5 / 8.]) else: self._scatteryoffsets = np.asarray(scatteryoffsets) reps = self.scatterpoints // len(self._scatteryoffsets) + 1 self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[:self.scatterpoints] # _legend_box is a VPacker instance that contains all # legend items and will be initialized from _init_legend_box() # method. self._legend_box = None if isinstance(parent, Axes): self.isaxes = True self.axes = parent self.set_figure(parent.figure) elif isinstance(parent, Figure): self.isaxes = False self.set_figure(parent) else: raise TypeError("Legend needs either Axes or Figure as parent") self.parent = parent self._loc_used_default = loc is None if loc is None: loc = mpl.rcParams["legend.loc"] if not self.isaxes and loc in [0, 'best']: loc = 'upper right' if isinstance(loc, str): if loc not in self.codes: raise ValueError( "Unrecognized location {!r}. Valid locations are\n\t{}\n". format(loc, '\n\t'.join(self.codes))) else: loc = self.codes[loc] if not self.isaxes and loc == 0: raise ValueError( "Automatic legend placement (loc='best') not implemented for " "figure legend.") self._mode = mode self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform) # We use FancyBboxPatch to draw a legend frame. The location # and size of the box will be updated during the drawing time. if facecolor is None: facecolor = mpl.rcParams["legend.facecolor"] if facecolor == 'inherit': facecolor = mpl.rcParams["axes.facecolor"] if edgecolor is None: edgecolor = mpl.rcParams["legend.edgecolor"] if edgecolor == 'inherit': edgecolor = mpl.rcParams["axes.edgecolor"] if fancybox is None: fancybox = mpl.rcParams["legend.fancybox"] self.legendPatch = FancyBboxPatch( xy=(0, 0), width=1, height=1, facecolor=facecolor, edgecolor=edgecolor, # If shadow is used, default to alpha=1 (#8943). alpha=(framealpha if framealpha is not None else 1 if shadow else mpl.rcParams["legend.framealpha"]), # The width and height of the legendPatch will be set (in draw()) # to the length that includes the padding. Thus we set pad=0 here. boxstyle=("round,pad=0,rounding_size=0.2" if fancybox else "square,pad=0"), mutation_scale=self._fontsize, snap=True, visible=(frameon if frameon is not None else mpl.rcParams["legend.frameon"])) self._set_artist_props(self.legendPatch) # init with null renderer self._init_legend_box(handles, labels, markerfirst) tmp = self._loc_used_default self._set_loc(loc) self._loc_used_default = tmp # ignore changes done by _set_loc # figure out title fontsize: if title_fontsize is None: title_fontsize = mpl.rcParams['legend.title_fontsize'] tprop = FontProperties(size=title_fontsize) self.set_title(title, prop=tprop) self._draggable = None # set the text color color_getters = { # getter function depends on line or patch 'linecolor': ['get_color', 'get_facecolor'], 'markerfacecolor': ['get_markerfacecolor', 'get_facecolor'], 'mfc': ['get_markerfacecolor', 'get_facecolor'], 'markeredgecolor': ['get_markeredgecolor', 'get_edgecolor'], 'mec': ['get_markeredgecolor', 'get_edgecolor'], } if labelcolor is None: pass elif isinstance(labelcolor, str) and labelcolor in color_getters: getter_names = color_getters[labelcolor] for handle, text in zip(self.legendHandles, self.texts): for getter_name in getter_names: try: color = getattr(handle, getter_name)() text.set_color(color) break except AttributeError: pass elif np.iterable(labelcolor): for text, color in zip( self.texts, itertools.cycle(colors.to_rgba_array(labelcolor))): text.set_color(color) else: raise ValueError("Invalid argument for labelcolor : %s" % str(labelcolor))
class Legend(Artist): """ Place a legend on the axes at location loc. Labels are a sequence of strings and loc can be a string or an integer specifying the legend location The location codes are:: 'best' : 0, (only implemented for axis legends) 'upper right' : 1, 'upper left' : 2, 'lower left' : 3, 'lower right' : 4, 'right' : 5, 'center left' : 6, 'center right' : 7, 'lower center' : 8, 'upper center' : 9, 'center' : 10, loc can be a tuple of the noramilzed coordinate values with respect its parent. """ codes = {'best' : 0, # only implemented for axis legends 'upper right' : 1, 'upper left' : 2, 'lower left' : 3, 'lower right' : 4, 'right' : 5, 'center left' : 6, 'center right' : 7, 'lower center' : 8, 'upper center' : 9, 'center' : 10, } zorder = 5 def __str__(self): return "Legend" def __init__(self, parent, handles, labels, loc = None, numpoints = None, # the number of points in the legend line markerscale = None, # the relative size of legend markers vs. original scatterpoints = 3, # TODO: may be an rcParam scatteryoffsets=None, prop = None, # properties for the legend texts # the following dimensions are in axes coords pad = None, # deprecated; use borderpad labelsep = None, # deprecated; use labelspacing handlelen = None, # deprecated; use handlelength handletextsep = None, # deprecated; use handletextpad axespad = None, # deprecated; use borderaxespad # spacing & pad defined as a fraction of the font-size borderpad = None, # the whitespace inside the legend border labelspacing=None, #the vertical space between the legend entries handlelength=None, # the length of the legend handles handleheight=None, # the height of the legend handles handletextpad=None, # the pad between the legend handle and text borderaxespad=None, # the pad between the axes and legend border columnspacing=None, # spacing between columns ncol=1, # number of columns mode=None, # mode for horizontal distribution of columns. None, "expand" fancybox=None, # True use a fancy box, false use a rounded box, none use rc shadow = None, title = None, # set a title for the legend bbox_to_anchor = None, # bbox that the legend will be anchored. bbox_transform = None, # transform for the bbox frameon = None, # draw frame handler_map = None, ): """ - *parent* : the artist that contains the legend - *handles* : a list of artists (lines, patches) to be added to the legend - *labels* : a list of strings to label the legend Optional keyword arguments: ================ ================================================================== Keyword Description ================ ================================================================== loc a location code prop the font property markerscale the relative size of legend markers vs. original numpoints the number of points in the legend for line scatterpoints the number of points in the legend for scatter plot scatteryoffsets a list of yoffsets for scatter symbols in legend frameon if True, draw a frame around the legend. If None, use rc fancybox if True, draw a frame with a round fancybox. If None, use rc shadow if True, draw a shadow behind legend ncol number of columns borderpad the fractional whitespace inside the legend border labelspacing the vertical space between the legend entries handlelength the length of the legend handles handleheight the length of the legend handles handletextpad the pad between the legend handle and text borderaxespad the pad between the axes and legend border columnspacing the spacing between columns title the legend title bbox_to_anchor the bbox that the legend will be anchored. bbox_transform the transform for the bbox. transAxes if None. ================ ================================================================== The pad and spacing parameters are measured in font-size units. E.g., a fontsize of 10 points and a handlelength=5 implies a handlelength of 50 points. Values from rcParams will be used if None. Users can specify any arbitrary location for the legend using the *bbox_to_anchor* keyword argument. bbox_to_anchor can be an instance of BboxBase(or its derivatives) or a tuple of 2 or 4 floats. See :meth:`set_bbox_to_anchor` for more detail. The legend location can be specified by setting *loc* with a tuple of 2 floats, which is interpreted as the lower-left corner of the legend in the normalized axes coordinate. """ from matplotlib.axes import Axes # local import only to avoid circularity from matplotlib.figure import Figure # local import only to avoid circularity Artist.__init__(self) if prop is None: self.prop=FontProperties(size=rcParams["legend.fontsize"]) elif isinstance(prop, dict): self.prop=FontProperties(**prop) if "size" not in prop: self.prop.set_size(rcParams["legend.fontsize"]) else: self.prop=prop self._fontsize = self.prop.get_size_in_points() propnames=['numpoints', 'markerscale', 'shadow', "columnspacing", "scatterpoints", "handleheight"] self.texts = [] self.legendHandles = [] self._legend_title_box = None self._handler_map = handler_map localdict = locals() for name in propnames: if localdict[name] is None: value = rcParams["legend."+name] else: value = localdict[name] setattr(self, name, value) # Take care the deprecated keywords deprecated_kwds = {"pad":"borderpad", "labelsep":"labelspacing", "handlelen":"handlelength", "handletextsep":"handletextpad", "axespad":"borderaxespad"} # convert values of deprecated keywords (ginve in axes coords) # to new vaules in a fraction of the font size # conversion factor bbox = parent.bbox axessize_fontsize = min(bbox.width, bbox.height)/self._fontsize for k, v in deprecated_kwds.iteritems(): # use deprecated value if not None and if their newer # counter part is None. if localdict[k] is not None and localdict[v] is None: warnings.warn("Use '%s' instead of '%s'." % (v, k), DeprecationWarning) setattr(self, v, localdict[k]*axessize_fontsize) continue # Otherwise, use new keywords if localdict[v] is None: setattr(self, v, rcParams["legend."+v]) else: setattr(self, v, localdict[v]) del localdict handles = list(handles) if len(handles)<2: ncol = 1 self._ncol = ncol if self.numpoints <= 0: raise ValueError("numpoints must be > 0; it was %d"% numpoints) # introduce y-offset for handles of the scatter plot if scatteryoffsets is None: self._scatteryoffsets = np.array([3./8., 4./8., 2.5/8.]) else: self._scatteryoffsets = np.asarray(scatteryoffsets) reps = int(self.scatterpoints / len(self._scatteryoffsets)) + 1 self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[:self.scatterpoints] # _legend_box is an OffsetBox instance that contains all # legend items and will be initialized from _init_legend_box() # method. self._legend_box = None if isinstance(parent,Axes): self.isaxes = True self.set_axes(parent) self.set_figure(parent.figure) elif isinstance(parent,Figure): self.isaxes = False self.set_figure(parent) else: raise TypeError("Legend needs either Axes or Figure as parent") self.parent = parent if loc is None: loc = rcParams["legend.loc"] if not self.isaxes and loc in [0,'best']: loc = 'upper right' if is_string_like(loc): if loc not in self.codes: if self.isaxes: warnings.warn('Unrecognized location "%s". Falling back on "best"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes.iterkeys()))) loc = 0 else: warnings.warn('Unrecognized location "%s". Falling back on "upper right"; ' 'valid locations are\n\t%s\n' % (loc, '\n\t'.join(self.codes.iterkeys()))) loc = 1 else: loc = self.codes[loc] if not self.isaxes and loc == 0: warnings.warn('Automatic legend placement (loc="best") not implemented for figure legend. ' 'Falling back on "upper right".') loc = 1 self._mode = mode self.set_bbox_to_anchor(bbox_to_anchor, bbox_transform) # We use FancyBboxPatch to draw a legend frame. The location # and size of the box will be updated during the drawing time. self.legendPatch = FancyBboxPatch( xy=(0.0, 0.0), width=1., height=1., facecolor=rcParams["axes.facecolor"], edgecolor=rcParams["axes.edgecolor"], mutation_scale=self._fontsize, snap=True ) # The width and height of the legendPatch will be set (in the # draw()) to the length that includes the padding. Thus we set # pad=0 here. if fancybox is None: fancybox = rcParams["legend.fancybox"] if fancybox == True: self.legendPatch.set_boxstyle("round",pad=0, rounding_size=0.2) else: self.legendPatch.set_boxstyle("square",pad=0) self._set_artist_props(self.legendPatch) self._drawFrame = frameon if frameon is None: self._drawFrame = rcParams["legend.frameon"] # init with null renderer self._init_legend_box(handles, labels) self._loc = loc self.set_title(title) self._last_fontsize_points = self._fontsize self._draggable = None def _set_artist_props(self, a): """ set the boilerplate props for artists added to axes """ a.set_figure(self.figure) if self.isaxes: a.set_axes(self.axes) a.set_transform(self.get_transform()) def _set_loc(self, loc): # find_offset function will be provided to _legend_box and # _legend_box will draw itself at the location of the return # value of the find_offset. self._loc_real = loc if loc == 0: _findoffset = self._findoffset_best else: _findoffset = self._findoffset_loc #def findoffset(width, height, xdescent, ydescent): # return _findoffset(width, height, xdescent, ydescent, renderer) self._legend_box.set_offset(_findoffset) self._loc_real = loc def _get_loc(self): return self._loc_real _loc = property(_get_loc, _set_loc) def _findoffset_best(self, width, height, xdescent, ydescent, renderer): "Helper function to locate the legend at its best position" ox, oy = self._find_best_position(width, height, renderer) return ox+xdescent, oy+ydescent def _findoffset_loc(self, width, height, xdescent, ydescent, renderer): "Heper function to locate the legend using the location code" if iterable(self._loc) and len(self._loc)==2: # when loc is a tuple of axes(or figure) coordinates. fx, fy = self._loc bbox = self.get_bbox_to_anchor() x, y = bbox.x0 + bbox.width * fx, bbox.y0 + bbox.height * fy else: bbox = Bbox.from_bounds(0, 0, width, height) x, y = self._get_anchored_bbox(self._loc, bbox, self.get_bbox_to_anchor(), renderer) return x+xdescent, y+ydescent @allow_rasterization def draw(self, renderer): "Draw everything that belongs to the legend" if not self.get_visible(): return renderer.open_group('legend') fontsize = renderer.points_to_pixels(self._fontsize) # if mode == fill, set the width of the legend_box to the # width of the paret (minus pads) if self._mode in ["expand"]: pad = 2*(self.borderaxespad+self.borderpad)*fontsize self._legend_box.set_width(self.get_bbox_to_anchor().width-pad) if self._drawFrame: # update the location and size of the legend bbox = self._legend_box.get_window_extent(renderer) self.legendPatch.set_bounds(bbox.x0, bbox.y0, bbox.width, bbox.height) self.legendPatch.set_mutation_scale(fontsize) if self.shadow: shadow = Shadow(self.legendPatch, 2, -2) shadow.draw(renderer) self.legendPatch.draw(renderer) self._legend_box.draw(renderer) renderer.close_group('legend') def _approx_text_height(self, renderer=None): """ Return the approximate height of the text. This is used to place the legend handle. """ if renderer is None: return self._fontsize else: return renderer.points_to_pixels(self._fontsize) # _default_handler_map defines the default mapping between plot # elements and the legend handlers. _default_handler_map = { StemContainer:legend_handler.HandlerStem(), ErrorbarContainer:legend_handler.HandlerErrorbar(), Line2D:legend_handler.HandlerLine2D(), Patch:legend_handler.HandlerPatch(), LineCollection:legend_handler.HandlerLineCollection(), RegularPolyCollection:legend_handler.HandlerRegularPolyCollection(), CircleCollection:legend_handler.HandlerCircleCollection(), BarContainer:legend_handler.HandlerPatch(update_func=legend_handler.update_from_first_child), tuple:legend_handler.HandlerTuple(), PathCollection:legend_handler.HandlerPathCollection() } # (get|set|update)_default_handler_maps are public interfaces to # modify the defalut handler map. @classmethod def get_default_handler_map(cls): """ A class method that returns the default handler map. """ return cls._default_handler_map @classmethod def set_default_handler_map(cls, handler_map): """ A class method to set the default handler map. """ cls._default_handler_map = handler_map @classmethod def update_default_handler_map(cls, handler_map): """ A class method to update the default handler map. """ cls._default_handler_map.update(handler_map) def get_legend_handler_map(self): """ return the handler map. """ default_handler_map = self.get_default_handler_map() if self._handler_map: hm = default_handler_map.copy() hm.update(self._handler_map) return hm else: return default_handler_map @staticmethod def get_legend_handler(legend_handler_map, orig_handle): """ return a legend handler from *legend_handler_map* that corresponds to *orig_handler*. *legend_handler_map* should be a dictionary object (that is returned by the get_legend_handler_map method). It first checks if the *orig_handle* itself is a key in the *legend_hanler_map* and return the associated value. Otherwised, it checks for each of the classes in its method-resolution-order. If no matching key is found, it returns None. """ legend_handler_keys = legend_handler_map.keys() if orig_handle in legend_handler_keys: handler = legend_handler_map[orig_handle] else: for handle_type in type(orig_handle).mro(): if handle_type in legend_handler_map: handler = legend_handler_map[handle_type] break else: handler = None return handler def _init_legend_box(self, handles, labels): """ Initiallize the legend_box. The legend_box is an instance of the OffsetBox, which is packed with legend handles and texts. Once packed, their location is calculated during the drawing time. """ fontsize = self._fontsize # legend_box is a HPacker, horizontally packed with # columns. Each column is a VPacker, vertically packed with # legend items. Each legend item is HPacker packed with # legend handleBox and labelBox. handleBox is an instance of # offsetbox.DrawingArea which contains legend handle. labelBox # is an instance of offsetbox.TextArea which contains legend # text. text_list = [] # the list of text instances handle_list = [] # the list of text instances label_prop = dict(verticalalignment='baseline', horizontalalignment='left', fontproperties=self.prop, ) labelboxes = [] handleboxes = [] # The approximate height and descent of text. These values are # only used for plotting the legend handle. descent = 0.35*self._approx_text_height()*(self.handleheight - 0.7) # 0.35 and 0.7 are just heuristic numbers. this may need to be improbed height = self._approx_text_height() * self.handleheight - descent # each handle needs to be drawn inside a box of (x, y, w, h) = # (0, -descent, width, height). And their corrdinates should # be given in the display coordinates. # The transformation of each handle will be automatically set # to self.get_trasnform(). If the artist does not uses its # default trasnform (eg, Collections), you need to # manually set their transform to the self.get_transform(). legend_handler_map = self.get_legend_handler_map() for orig_handle, lab in zip(handles, labels): handler = self.get_legend_handler(legend_handler_map, orig_handle) if handler is None: warnings.warn("Legend does not support %s\nUse proxy artist instead.\n\nhttp://matplotlib.sourceforge.net/users/legend_guide.html#using-proxy-artist\n" % (str(orig_handle),)) handle_list.append(None) continue textbox = TextArea(lab, textprops=label_prop, multilinebaseline=True, minimumdescent=True) text_list.append(textbox._text) labelboxes.append(textbox) handlebox = DrawingArea(width=self.handlelength*fontsize, height=height, xdescent=0., ydescent=descent) handle = handler(self, orig_handle, \ #xdescent, ydescent, width, height, fontsize, handlebox) handle_list.append(handle) handleboxes.append(handlebox) if len(handleboxes) > 0: # We calculate number of lows in each column. The first # (num_largecol) columns will have (nrows+1) rows, and remaing # (num_smallcol) columns will have (nrows) rows. ncol = min(self._ncol, len(handleboxes)) nrows, num_largecol = divmod(len(handleboxes), ncol) num_smallcol = ncol-num_largecol # starting index of each column and number of rows in it. largecol = safezip(range(0, num_largecol*(nrows+1), (nrows+1)), [nrows+1] * num_largecol) smallcol = safezip(range(num_largecol*(nrows+1), len(handleboxes), nrows), [nrows] * num_smallcol) else: largecol, smallcol = [], [] handle_label = safezip(handleboxes, labelboxes) columnbox = [] for i0, di in largecol+smallcol: # pack handleBox and labelBox into itemBox itemBoxes = [HPacker(pad=0, sep=self.handletextpad*fontsize, children=[h, t], align="baseline") for h, t in handle_label[i0:i0+di]] # minimumdescent=False for the text of the last row of the column itemBoxes[-1].get_children()[1].set_minimumdescent(False) # pack columnBox columnbox.append(VPacker(pad=0, sep=self.labelspacing*fontsize, align="baseline", children=itemBoxes)) if self._mode == "expand": mode = "expand" else: mode = "fixed" sep = self.columnspacing*fontsize self._legend_handle_box = HPacker(pad=0, sep=sep, align="baseline", mode=mode, children=columnbox) self._legend_title_box = TextArea("") self._legend_box = VPacker(pad=self.borderpad*fontsize, sep=self.labelspacing*fontsize, align="center", children=[self._legend_title_box, self._legend_handle_box]) self._legend_box.set_figure(self.figure) self.texts = text_list self.legendHandles = handle_list def _auto_legend_data(self): """ Returns list of vertices and extents covered by the plot. Returns a two long list. First element is a list of (x, y) vertices (in display-coordinates) covered by all the lines and line collections, in the legend's handles. Second element is a list of bounding boxes for all the patches in the legend's handles. """ assert self.isaxes # should always hold because function is only called internally ax = self.parent vertices = [] bboxes = [] lines = [] for handle in ax.lines: assert isinstance(handle, Line2D) path = handle.get_path() trans = handle.get_transform() tpath = trans.transform_path(path) lines.append(tpath) for handle in ax.patches: assert isinstance(handle, Patch) if isinstance(handle, Rectangle): transform = handle.get_data_transform() bboxes.append(handle.get_bbox().transformed(transform)) else: transform = handle.get_transform() bboxes.append(handle.get_path().get_extents(transform)) return [vertices, bboxes, lines] def draw_frame(self, b): 'b is a boolean. Set draw frame to b' self.set_frame_on(b) def get_children(self): 'return a list of child artists' children = [] if self._legend_box: children.append(self._legend_box) children.extend(self.get_lines()) children.extend(self.get_patches()) children.extend(self.get_texts()) children.append(self.get_frame()) if self._legend_title_box: children.append(self.get_title()) return children def get_frame(self): 'return the Rectangle instance used to frame the legend' return self.legendPatch def get_lines(self): 'return a list of lines.Line2D instances in the legend' return [h for h in self.legendHandles if isinstance(h, Line2D)] def get_patches(self): 'return a list of patch instances in the legend' return silent_list('Patch', [h for h in self.legendHandles if isinstance(h, Patch)]) def get_texts(self): 'return a list of text.Text instance in the legend' return silent_list('Text', self.texts) def set_title(self, title, prop=None): """ set the legend title. Fontproperties can be optionally set with *prop* parameter. """ self._legend_title_box._text.set_text(title) if prop is not None: if isinstance(prop, dict): prop = FontProperties(**prop) self._legend_title_box._text.set_fontproperties(prop) if title: self._legend_title_box.set_visible(True) else: self._legend_title_box.set_visible(False) def get_title(self): 'return Text instance for the legend title' return self._legend_title_box._text def get_window_extent(self, *args, **kwargs) : 'return a extent of the the legend' return self.legendPatch.get_window_extent(*args, **kwargs) def get_frame_on(self): """ Get whether the legend box patch is drawn """ return self._drawFrame def set_frame_on(self, b): """ Set whether the legend box patch is drawn ACCEPTS: [ *True* | *False* ] """ self._drawFrame = b def get_bbox_to_anchor(self): """ return the bbox that the legend will be anchored """ if self._bbox_to_anchor is None: return self.parent.bbox else: return self._bbox_to_anchor def set_bbox_to_anchor(self, bbox, transform=None): """ set the bbox that the legend will be anchored. *bbox* can be a BboxBase instance, a tuple of [left, bottom, width, height] in the given transform (normalized axes coordinate if None), or a tuple of [left, bottom] where the width and height will be assumed to be zero. """ if bbox is None: self._bbox_to_anchor = None return elif isinstance(bbox, BboxBase): self._bbox_to_anchor = bbox else: try: l = len(bbox) except TypeError: raise ValueError("Invalid argument for bbox : %s" % str(bbox)) if l == 2: bbox = [bbox[0], bbox[1], 0, 0] self._bbox_to_anchor = Bbox.from_bounds(*bbox) if transform is None: transform = BboxTransformTo(self.parent.bbox) self._bbox_to_anchor = TransformedBbox(self._bbox_to_anchor, transform) def _get_anchored_bbox(self, loc, bbox, parentbbox, renderer): """ Place the *bbox* inside the *parentbbox* according to a given location code. Return the (x,y) coordinate of the bbox. - loc: a location code in range(1, 11). This corresponds to the possible values for self._loc, excluding "best". - bbox: bbox to be placed, display coodinate units. - parentbbox: a parent box which will contain the bbox. In display coordinates. """ assert loc in range(1,11) # called only internally BEST, UR, UL, LL, LR, R, CL, CR, LC, UC, C = range(11) anchor_coefs={UR:"NE", UL:"NW", LL:"SW", LR:"SE", R:"E", CL:"W", CR:"E", LC:"S", UC:"N", C:"C"} c = anchor_coefs[loc] fontsize = renderer.points_to_pixels(self._fontsize) container = parentbbox.padded(-(self.borderaxespad) * fontsize) anchored_box = bbox.anchored(c, container=container) return anchored_box.x0, anchored_box.y0 def _find_best_position(self, width, height, renderer, consider=None): """ Determine the best location to place the legend. `consider` is a list of (x, y) pairs to consider as a potential lower-left corner of the legend. All are display coords. """ assert self.isaxes # should always hold because function is only called internally verts, bboxes, lines = self._auto_legend_data() bbox = Bbox.from_bounds(0, 0, width, height) consider = [self._get_anchored_bbox(x, bbox, self.get_bbox_to_anchor(), renderer) for x in range(1, len(self.codes))] #tx, ty = self.legendPatch.get_x(), self.legendPatch.get_y() candidates = [] for l, b in consider: legendBox = Bbox.from_bounds(l, b, width, height) badness = 0 badness = legendBox.count_contains(verts) badness += legendBox.count_overlaps(bboxes) for line in lines: if line.intersects_bbox(legendBox): badness += 1 ox, oy = l, b if badness == 0: return ox, oy candidates.append((badness, (l, b))) # rather than use min() or list.sort(), do this so that we are assured # that in the case of two equal badnesses, the one first considered is # returned. # NOTE: list.sort() is stable.But leave as it is for now. -JJL minCandidate = candidates[0] for candidate in candidates: if candidate[0] < minCandidate[0]: minCandidate = candidate ox, oy = minCandidate[1] return ox, oy def contains(self, event): return self.legendPatch.contains(event) def draggable(self, state=None, use_blit=False, update="loc"): """ Set the draggable state -- if state is * None : toggle the current state * True : turn draggable on * False : turn draggable off If draggable is on, you can drag the legend on the canvas with the mouse. The DraggableLegend helper instance is returned if draggable is on. The update parameter control which parameter of the legend changes when dragged. If update is "loc", the *loc* paramter of the legend is changed. If "bbox", the *bbox_to_anchor* parameter is changed. """ is_draggable = self._draggable is not None # if state is None we'll toggle if state is None: state = not is_draggable if state: if self._draggable is None: self._draggable = DraggableLegend(self, use_blit, update=update) else: if self._draggable is not None: self._draggable.disconnect() self._draggable = None return self._draggable
from figpptx.comparer import Comparer styles = [ "circle", "square", "round", "larrow", "rarrow", "darrow", ] # fig = plt.figure() fig, ax = plt.subplots() for i, style in enumerate(styles): ys = 1 / (len(styles)) * (i) height = 1 / (2 * len(styles)) boxstyle = BoxStyle(style, pad=0.01) p_bbox = FancyBboxPatch( (0.2, ys), 0.3, height, boxstyle=boxstyle, ec="black", fc=f"C{i}", ) ax.add_patch(p_bbox) fig.tight_layout() comparer = Comparer() comparer.compare(fig)