Ejemplo n.º 1
0
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,  (currently not supported, defaults to upper right)
      'upper right'  : 1,  (default)
      '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,
 
    Return value is a sequence of text, line instances that make
    up the legend
    """


    codes = {'best'         : 0,
             'upper right'  : 1,  # default
             '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,
             }




    def __init__(self, parent, handles, labels, loc,
                 isaxes=True,
                 numpoints = 4,      # the number of points in the legend line
                 prop = FontProperties(size='smaller'),
                 pad = 0.2,          # the fractional whitespace inside the legend border
                 markerscale = 0.6,    # the relative size of legend markers vs. original
                 # the following dimensions are in axes coords
                 labelsep = 0.005,     # the vertical space between the legend entries
                 handlelen = 0.05,     # the length of the legend lines
                 handletextsep = 0.02, # the space between the legend line and legend text
                 axespad = 0.02,       # the border between the axes and legend edge

                 shadow=False,
                 ):
        """
  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 
  loc                   # a location code
  isaxes=True           # whether this is an axes legend
  numpoints = 4         # the number of points in the legend line
  fontprop = FontProperties('smaller')  # the font property
  pad = 0.2             # the fractional whitespace inside the legend border
  markerscale = 0.6     # the relative size of legend markers vs. original
  shadow                # if True, draw a shadow behind legend 

The following dimensions are in axes coords
  labelsep = 0.005     # the vertical space between the legend entries
  handlelen = 0.05     # the length of the legend lines
  handletextsep = 0.02 # the space between the legend line and legend text
  axespad = 0.02       # the border between the axes and legend edge
        """
        Artist.__init__(self)
        if is_string_like(loc) and not self.codes.has_key(loc):
            verbose.report_error('Unrecognized location %s. Falling back on upper right; valid locations are\n%s\t' %(loc, '\n\t'.join(self.codes.keys())))
        if is_string_like(loc): loc = self.codes.get(loc, 1)
        
        self.numpoints = numpoints
        self.prop = prop
        self.fontsize = prop.get_size_in_points()
        self.pad = pad
        self.markerscale = markerscale
        self.labelsep = labelsep
        self.handlelen = handlelen
        self.handletextsep = handletextsep
        self.axespad = axespad
        self.shadow = shadow
        
        if isaxes:  # parent is an Axes
            self.set_figure(parent.figure)
        else:        # parent is a Figure
            self.set_figure(parent)

        self.parent = parent
        self.set_transform( get_bbox_transform( unit_bbox(), parent.bbox) )
        self._loc = loc   

        # make a trial box in the middle of the axes.  relocate it
        # based on it's bbox
        left, upper = 0.5, 0.5
        if self.numpoints == 1:
            self._xdata = array([left + self.handlelen*0.5])
        else:
            self._xdata = linspace(left, left + self.handlelen, self.numpoints)
        textleft = left+ self.handlelen+self.handletextsep
        self.texts = self._get_texts(labels, textleft, upper)
        self.handles = self._get_handles(handles, self.texts)
        
        left, top = self.texts[-1].get_position()
        HEIGHT = self._approx_text_height()
        bottom = top-HEIGHT
        left -= self.handlelen + self.handletextsep + self.pad
        self.legendPatch = Rectangle(
            xy=(left, bottom), width=0.5, height=HEIGHT*len(self.texts),
            facecolor='w', edgecolor='k',
            )
        self._set_artist_props(self.legendPatch)
        self._drawFrame = True

    def _set_artist_props(self, a):
        a.set_figure(self.figure)
        a.set_transform(self._transform)
        
    def _approx_text_height(self):
        return self.fontsize/72.0*self.figure.dpi.get()/self.parent.bbox.height()

            
    def draw(self, renderer):
        if not self.get_visible(): return 
        renderer.open_group('legend')
        self._update_positions(renderer)
        if self._drawFrame:
            if self.shadow:
                shadow = Shadow(self.legendPatch, -0.005, -0.005)
                shadow.draw(renderer)
            self.legendPatch.draw(renderer)
        

        for h in self.handles:            
            if h is not None:
		h.draw(renderer)
        	if 0: bbox_artist(h, renderer)

        for t in self.texts:
            if 0: bbox_artist(t, renderer)
            t.draw(renderer)
        renderer.close_group('legend')
        #draw_bbox(self.save, renderer, 'g')
        #draw_bbox(self.ibox, renderer, 'r', self._transform)

    def _get_handle_text_bbox(self, renderer):
        'Get a bbox for the text and lines in axes coords'

        bboxesText = [t.get_window_extent(renderer) for t in self.texts]
        bboxesHandles = [h.get_window_extent(renderer) for h in self.handles if h is not None]


        bboxesAll = bboxesText
        bboxesAll.extend(bboxesHandles)
        bbox = bbox_all(bboxesAll)
        self.save = bbox

        ibox =  inverse_transform_bbox(self._transform, bbox)
        self.ibox = ibox

        return ibox
        
    def _get_handles(self, handles, texts):
        HEIGHT = self._approx_text_height()

        ret = []   # the returned legend lines

        for handle, label in zip(handles, texts):
            x, y = label.get_position()
            x -= self.handlelen + self.handletextsep
            if isinstance(handle, Line2D):
                ydata = (y-HEIGHT/2)*ones(self._xdata.shape, Float)
                legline = Line2D(self._xdata, ydata)
                self._set_artist_props(legline)
                legline.copy_properties(handle)
                legline.set_markersize(self.markerscale*legline.get_markersize())
                legline.set_data_clipping(False)
                ret.append(legline)
            elif isinstance(handle, Patch):

                p = Rectangle(xy=(min(self._xdata), y-3/4*HEIGHT),
                              width = self.handlelen, height=HEIGHT/2,
                              )
                p.copy_properties(handle)
                self._set_artist_props(p)
                ret.append(p)
            elif isinstance(handle, LineCollection):
                ydata = (y-HEIGHT/2)*ones(self._xdata.shape, Float)
                legline = Line2D(self._xdata, ydata)
                self._set_artist_props(legline)
                lw = handle.get_linewidths()[0]
                color = handle.get_colors()[0]                
                legline.set_color(color)
                legline.set_linewidth(lw)
                ret.append(legline)
                
	    else:
		ret.append(None)
                                               
        return ret

    def draw_frame(self, b):
        'b is a boolean.  Set draw frame to b'
        self._drawFrame = b

    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.handles 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.handles 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_texts(self, labels, left, upper):

        # height in axes coords
        HEIGHT = self._approx_text_height()
        pos = upper
        x = left 

        ret = []  # the returned list of text instances
        for l in labels:
            text = Text(
                x=x, y=pos,
                text=l,
                fontproperties=self.prop,
                verticalalignment='top',
                horizontalalignment='left',
                )
            self._set_artist_props(text)
            ret.append(text)
            pos -= HEIGHT
            
        return ret

            
    def get_window_extent(self):
        return self.legendPatch.get_window_extent()


    def _offset(self, ox, oy):
        'Move all the artists by ox,oy (axes coords)'
        for t in self.texts:
            x,y = t.get_position()
            t.set_position( (x+ox, y+oy) )

        for h in self.handles:
            if isinstance(h, Line2D):
                x,y = h.get_xdata(), h.get_ydata()
                h.set_data( x+ox, y+oy)
            elif isinstance(h, Rectangle):
                h.xy[0] = h.xy[0] + ox
                h.xy[1] = h.xy[1] + oy

        x, y = self.legendPatch.get_x(), self.legendPatch.get_y()
        self.legendPatch.set_x(x+ox)
        self.legendPatch.set_y(y+oy)

    def _update_positions(self, renderer):
        # called from renderer to allow more precise estimates of
        # widths and heights with get_window_extent

        def get_tbounds(text):  #get text bounds in axes coords
            bbox = text.get_window_extent(renderer)
            bboxa = inverse_transform_bbox(self._transform, bbox)
            return bboxa.get_bounds()
            
        hpos = []
        for t, tabove in zip(self.texts[1:], self.texts[:-1]):
            x,y = t.get_position()
            l,b,w,h = get_tbounds(tabove)
            hpos.append( (b,h) )
            t.set_position( (x, b-0.1*h) )

        # now do the same for last line
        l,b,w,h = get_tbounds(self.texts[-1])
        hpos.append( (b,h) )
        
        for handle, tup in zip(self.handles, hpos):
            y,h = tup
            if isinstance(handle, Line2D):
                ydata = y*ones(self._xdata.shape, Float)            
                handle.set_ydata(ydata+h/2)
            elif isinstance(handle, Rectangle):
                handle.set_y(y+1/4*h)
                handle.set_height(h/2)

        # Set the data for the legend patch
        bbox = self._get_handle_text_bbox(renderer).deepcopy()
        bbox.scale(1 + self.pad, 1 + self.pad)
        l,b,w,h = bbox.get_bounds()
        self.legendPatch.set_bounds(l,b,w,h)

        BEST, UR, UL, LL, LR, R, CL, CR, LC, UC, C = range(11)
        ox, oy = 0, 0                           # center


        if iterable(self._loc) and len(self._loc)==2:
            xo = self.legendPatch.get_x()
            yo = self.legendPatch.get_y()
            x, y = self._loc
            ox = x-xo
            oy = y-yo
            self._offset(ox, oy)
        else:
            if self._loc in (UL, LL, CL):           # left
                ox = self.axespad - l
            if self._loc in (BEST, UR, LR, R, CR):  # right
                ox = 1 - (l + w + self.axespad)
            if self._loc in (BEST, UR, UL, UC):     # upper
                oy = 1 - (b + h + self.axespad)
            if self._loc in (LL, LR, LC):           # lower
                oy = self.axespad - b
            if self._loc in (LC, UC, C):            # center x
                ox = (0.5-w/2)-l
            if self._loc in (CL, CR, C):            # center y
                oy = (0.5-h/2)-b
            self._offset(ox, oy)
Ejemplo n.º 2
0
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,

    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
                 prop = None,
                 pad = None,           # the fractional whitespace inside the legend border
                 markerscale = None,   # the relative size of legend markers vs. original
                 # the following dimensions are in axes coords
                 labelsep = None,      # the vertical space between the legend entries
                 handlelen = None,     # the length of the legend lines
                 handletextsep = None, # the space between the legend line and legend text
                 axespad = None,       # the border between the axes and legend edge

                 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
  loc                   # a location code
  numpoints = 4         # the number of points in the legend line
  prop = FontProperties(size='smaller')  # the font property
  pad = 0.2             # the fractional whitespace inside the legend border
  markerscale = 0.6     # the relative size of legend markers vs. original
  shadow                # if True, draw a shadow behind legend

The following dimensions are in axes coords
  labelsep = 0.005     # the vertical space between the legend entries
  handlelen = 0.05     # the length of the legend lines
  handletextsep = 0.02 # the space between the legend line and legend text
  axespad = 0.02       # the border between the axes and legend edge
        """
        from axes import Axes     # local import only to avoid circularity
        from figure import Figure # local import only to avoid circularity

        Artist.__init__(self)

        proplist=[numpoints, pad, markerscale, labelsep, handlelen, handletextsep, axespad, shadow]
        propnames=['numpoints', 'pad', 'markerscale', 'labelsep', 'handlelen', 'handletextsep', 'axespad', 'shadow']
        for name, value in safezip(propnames,proplist):
            if value is None:
                value=rcParams["legend."+name]
            setattr(self,name,value)
        if self.numpoints <= 0:
            raise ValueError("numpoints must be >= 0; it was %d"% numpoints)
        if prop is None:
            self.prop=FontProperties(size=rcParams["legend.fontsize"])
        else:
            self.prop=prop
        self.fontsize = self.prop.get_size_in_points()

        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
        self._offsetTransform = Affine2D()
        self._parentTransform = BboxTransformTo(parent.bbox)
        Artist.set_transform(self, self._offsetTransform + self._parentTransform)

        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 not self.codes.has_key(loc):
                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.legendPatch = Rectangle(
            xy=(0.0, 0.0), width=0.5, height=0.5,
            facecolor='w', edgecolor='k',
            )
        self._set_artist_props(self.legendPatch)

        # make a trial box in the middle of the axes.  relocate it
        # based on it's bbox
        left, top = 0.5, 0.5
        textleft = left+ self.handlelen+self.handletextsep
        self.texts = self._get_texts(labels, textleft, top)
        self.legendHandles = self._get_handles(handles, self.texts)

        self._drawFrame = True

    def _set_artist_props(self, a):
        a.set_figure(self.figure)
        a.set_transform(self.get_transform())

    def _approx_text_height(self):
        return self.fontsize/72.0*self.figure.dpi/self.parent.bbox.height


    def draw(self, renderer):
        if not self.get_visible(): return
        renderer.open_group('legend')
        self._update_positions(renderer)
        if self._drawFrame:
            if self.shadow:
                shadow = Shadow(self.legendPatch, -0.005, -0.005)
                shadow.draw(renderer)
            self.legendPatch.draw(renderer)


        if not len(self.legendHandles) and not len(self.texts): return
        for h in self.legendHandles:
            if h is not None:
                h.draw(renderer)
                if hasattr(h, '_legmarker'):
                    h._legmarker.draw(renderer)
                if 0: bbox_artist(h, renderer)

        for t in self.texts:
            if 0: bbox_artist(t, renderer)
            t.draw(renderer)
        renderer.close_group('legend')
        #draw_bbox(self.save, renderer, 'g')
        #draw_bbox(self.ibox, renderer, 'r', self.get_transform())

    def _get_handle_text_bbox(self, renderer):
        'Get a bbox for the text and lines in axes coords'

        bboxesText = [t.get_window_extent(renderer) for t in self.texts]
        bboxesHandles = [h.get_window_extent(renderer) for h in self.legendHandles if h is not None]

        bboxesAll = bboxesText
        bboxesAll.extend(bboxesHandles)
        bbox = Bbox.union(bboxesAll)

        self.save = bbox

        ibox = bbox.inverse_transformed(self.get_transform())
        self.ibox = ibox

        return ibox

    def _get_handles(self, handles, texts):
        handles = list(handles)
        texts = list(texts)
        HEIGHT = self._approx_text_height()
        left = 0.5

        ret = []   # the returned legend lines

        # we need to pad the text with empties for the numpoints=1
        # centered marker proxy

        for handle, label in safezip(handles, texts):
            if self.numpoints > 1:
                xdata = np.linspace(left, left + self.handlelen, self.numpoints)
                xdata_marker = xdata
            elif self.numpoints == 1:
                xdata = np.linspace(left, left + self.handlelen, 2)
                xdata_marker = [left + 0.5*self.handlelen]

            x, y = label.get_position()
            x -= self.handlelen + self.handletextsep
            if isinstance(handle, Line2D):
                ydata = (y-HEIGHT/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)
                ret.append(legline)
                legline.set_marker('None')

                legline_marker = Line2D(xdata_marker, ydata[:len(xdata_marker)])
                legline_marker.update_from(handle)
                legline_marker.set_linestyle('None')
                self._set_artist_props(legline_marker)
                # 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=(min(xdata), y-3/4*HEIGHT),
                              width = self.handlelen, height=HEIGHT/2,
                              )
                p.update_from(handle)
                self._set_artist_props(p)
                p.set_clip_box(None)
                p.set_clip_path(None)
                ret.append(p)
            elif isinstance(handle, LineCollection):
                ydata = (y-HEIGHT/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)
                ret.append(legline)

            elif isinstance(handle, RegularPolyCollection):
                if self.numpoints == 1:
                    xdata = np.array([left])
                p = Rectangle(xy=(min(xdata), y-3/4*HEIGHT),
                              width = self.handlelen, height=HEIGHT/2,
                              )
                p.set_facecolor(handle._facecolors[0])
                if handle._edgecolors != 'none' and len(handle._edgecolors):
                    p.set_edgecolor(handle._edgecolors[0])
                self._set_artist_props(p)
                p.set_clip_box(None)
                p.set_clip_path(None)
                ret.append(p)

            else:
                ret.append(None)

        return ret

    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
        axes-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 = []

        inverse_transform = ax.transAxes.inverted()

        for handle in ax.lines:
            assert isinstance(handle, Line2D)
            path = handle.get_path()
            trans = handle.get_transform()
            tpath = trans.transform_path(path)
            apath = inverse_transform.transform_path(tpath)
            lines.append(apath)

        for handle in ax.patches:
            assert isinstance(handle, Patch)

            if isinstance(handle, Rectangle):
                transform = handle.get_data_transform() + inverse_transform
                bboxes.append(handle.get_bbox().transformed(transform))
            else:
                transform = handle.get_transform() + inverse_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_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_texts(self, labels, left, upper):

        # height in axes coords
        HEIGHT = self._approx_text_height()
        pos = upper
        x = left

        ret = []  # the returned list of text instances
        for l in labels:
            text = Text(
                x=x, y=pos,
                text=l,
                fontproperties=self.prop,
                verticalalignment='top',
                horizontalalignment='left'
                )
            self._set_artist_props(text)
            ret.append(text)
            pos -= HEIGHT

        return ret


    def get_window_extent(self):
        return self.legendPatch.get_window_extent()


    def _offset(self, ox, oy):
        'Move all the artists by ox,oy (axes coords)'
        self._offsetTransform.clear().translate(ox, oy)

    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 axes coords.
        """

        assert self.isaxes # should always hold because function is only called internally

        verts, bboxes, lines = self._auto_legend_data()

        consider = [self._loc_to_axes_coords(x, width, height) 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-tx, b-ty
            if badness == 0:
                return ox, oy

            candidates.append((badness, (ox, oy)))

        # 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.
        minCandidate = candidates[0]
        for candidate in candidates:
            if candidate[0] < minCandidate[0]:
                minCandidate = candidate

        ox, oy = minCandidate[1]

        return ox, oy


    def _loc_to_axes_coords(self, loc, width, height):
        """Convert a location code to axes coordinates.

        - loc: a location code in range(1, 11).
          This corresponds to the possible values for self._loc, excluding "best".

        - width, height: the final size of the legend, axes units.
        """
        assert loc in range(1,11) # called only internally

        BEST, UR, UL, LL, LR, R, CL, CR, LC, UC, C = range(11)

        if loc in (UL, LL, CL):                # left
            x = self.axespad
        elif loc in (UR, LR, CR, R):           # right
            x = 1.0 - (width + self.axespad)
        elif loc in (LC, UC, C):               # center x
            x = (0.5 - width/2.0)

        if loc in (UR, UL, UC):                # upper
            y = 1.0 - (height + self.axespad)
        elif loc in (LL, LR, LC):              # lower
            y = self.axespad
        elif loc in (CL, CR, C, R):            # center y
            y = (0.5 - height/2.0)

        return x,y


    def _update_positions(self, renderer):
        # called from renderer to allow more precise estimates of
        # widths and heights with get_window_extent

        if not len(self.legendHandles) and not len(self.texts): return
        def get_tbounds(text):  #get text bounds in axes coords
            bbox = text.get_window_extent(renderer)
            bboxa = bbox.inverse_transformed(self.get_transform())
            return bboxa.bounds

        hpos = []
        for t, tabove in safezip(self.texts[1:], self.texts[:-1]):
            x,y = t.get_position()
            l,b,w,h = get_tbounds(tabove)
            b -= self.labelsep
            h += 2*self.labelsep
            hpos.append( (b,h) )
            t.set_position( (x, b-0.1*h) )

        # now do the same for last line

        l,b,w,h = get_tbounds(self.texts[-1])
        b -= self.labelsep
        h += 2*self.labelsep
        hpos.append( (b,h) )

        for handle, tup in safezip(self.legendHandles, hpos):
            y,h = tup
            if isinstance(handle, Line2D):
                ydata = y*np.ones(handle.get_xdata().shape, float)
                handle.set_ydata(ydata+h/2.)
                handle._legmarker.set_ydata(ydata+h/2.)
            elif isinstance(handle, Rectangle):
                handle.set_y(y+1/4*h)
                handle.set_height(h/2)

        # Set the data for the legend patch
        bbox = self._get_handle_text_bbox(renderer)

        bbox = bbox.expanded(1 + self.pad, 1 + self.pad)
        l, b, w, h = bbox.bounds
        self.legendPatch.set_bounds(l, b, w, h)

        ox, oy = 0, 0                           # center

        if iterable(self._loc) and len(self._loc)==2:
            xo = self.legendPatch.get_x()
            yo = self.legendPatch.get_y()
            x, y = self._loc
            ox, oy = x-xo, y-yo
        elif self._loc == 0:  # "best"
            ox, oy = self._find_best_position(w, h)
        else:
            x, y = self._loc_to_axes_coords(self._loc, w, h)
            ox, oy = x-l, y-b

        self._offset(ox, oy)
Ejemplo n.º 3
0
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,
      'upper right'  : 1,  (default)
      '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,

    Return value is a sequence of text, line instances that make
    up the legend
    """


    codes = {'best'         : 0,
             'upper right'  : 1,  # default
             '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,
             }




    def __init__(self, parent, handles, labels, loc,
                 isaxes= None,
                 numpoints = None,      # the number of points in the legend line
                 prop = None,
                 pad = None,          # the fractional whitespace inside the legend border
                 markerscale = None,    # the relative size of legend markers vs. original
                 # the following dimensions are in axes coords
                 labelsep = None,     # the vertical space between the legend entries
                 handlelen = None,     # the length of the legend lines
                 handletextsep = None, # the space between the legend line and legend text
                 axespad = None,       # the border between the axes and legend edge

                 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
  loc                   # a location code
  isaxes=True           # whether this is an axes legend
  numpoints = 4         # the number of points in the legend line
  fontprop = FontProperties(size='smaller')  # the font property
  pad = 0.2             # the fractional whitespace inside the legend border
  markerscale = 0.6     # the relative size of legend markers vs. original
  shadow                # if True, draw a shadow behind legend

The following dimensions are in axes coords
  labelsep = 0.005     # the vertical space between the legend entries
  handlelen = 0.05     # the length of the legend lines
  handletextsep = 0.02 # the space between the legend line and legend text
  axespad = 0.02       # the border between the axes and legend edge
        """
        Artist.__init__(self)
        if is_string_like(loc) and not self.codes.has_key(loc):
            warnings.warn('Unrecognized location %s. Falling back on upper right; valid locations are\n%s\t' %(loc, '\n\t'.join(self.codes.keys())))
        if is_string_like(loc): loc = self.codes.get(loc, 1)

        proplist=[numpoints, pad, markerscale, labelsep, handlelen, handletextsep, axespad, shadow, isaxes]
        propnames=['numpoints', 'pad', 'markerscale', 'labelsep', 'handlelen', 'handletextsep', 'axespad', 'shadow', 'isaxes']
        for name, value in zip(propnames,proplist):
            if value is None:
                value=rcParams["legend."+name]
            setattr(self,name,value)
        if prop is None:
            self.prop=FontProperties(size=rcParams["legend.fontsize"])
        else:
            self.prop=prop
        self.fontsize = self.prop.get_size_in_points()

        if self.isaxes:  # parent is an Axes
            self.set_figure(parent.figure)
        else:        # parent is a Figure
            self.set_figure(parent)

        self.parent = parent
        self.set_transform( get_bbox_transform( unit_bbox(), parent.bbox) )
        self._loc = loc

        # make a trial box in the middle of the axes.  relocate it
        # based on it's bbox
        left, top = 0.5, 0.5
        if self.numpoints == 1:
            self._xdata = array([left + self.handlelen*0.5])
        else:
            self._xdata = linspace(left, left + self.handlelen, self.numpoints)
        textleft = left+ self.handlelen+self.handletextsep
        self.texts = self._get_texts(labels, textleft, top)
        self.legendHandles = self._get_handles(handles, self.texts)


        if len(self.texts):
            left, top = self.texts[-1].get_position()
            HEIGHT = self._approx_text_height()*len(self.texts)
        else:
            HEIGHT = 0.2

        bottom = top-HEIGHT
        left -= self.handlelen + self.handletextsep + self.pad
        self.legendPatch = Rectangle(
            xy=(left, bottom), width=0.5, height=HEIGHT,
            facecolor='w', edgecolor='k',
            )
        self._set_artist_props(self.legendPatch)
        self._drawFrame = True

    def _set_artist_props(self, a):
        a.set_figure(self.figure)
        a.set_transform(self._transform)

    def _approx_text_height(self):
        return self.fontsize/72.0*self.figure.dpi.get()/self.parent.bbox.height()


    def draw(self, renderer):
        if not self.get_visible(): return
        renderer.open_group('legend')
        self._update_positions(renderer)
        if self._drawFrame:
            if self.shadow:
                shadow = Shadow(self.legendPatch, -0.005, -0.005)
                shadow.draw(renderer)
            self.legendPatch.draw(renderer)


        if not len(self.legendHandles) and not len(self.texts): return
        for h in self.legendHandles:
            if h is not None:
		h.draw(renderer)
        	if 0: bbox_artist(h, renderer)

        for t in self.texts:
            if 0: bbox_artist(t, renderer)
            t.draw(renderer)
        renderer.close_group('legend')
        #draw_bbox(self.save, renderer, 'g')
        #draw_bbox(self.ibox, renderer, 'r', self._transform)

    def _get_handle_text_bbox(self, renderer):
        'Get a bbox for the text and lines in axes coords'

        bboxesText = [t.get_window_extent(renderer) for t in self.texts]
        bboxesHandles = [h.get_window_extent(renderer) for h in self.legendHandles if h is not None]


        bboxesAll = bboxesText
        bboxesAll.extend(bboxesHandles)
        bbox = bbox_all(bboxesAll)
        self.save = bbox

        ibox =  inverse_transform_bbox(self._transform, bbox)
        self.ibox = ibox

        return ibox

    def _get_handles(self, handles, texts):
        HEIGHT = self._approx_text_height()

        ret = []   # the returned legend lines

        for handle, label in zip(handles, texts):
            x, y = label.get_position()
            x -= self.handlelen + self.handletextsep
            if isinstance(handle, Line2D):
                ydata = (y-HEIGHT/2)*ones(self._xdata.shape, Float)
                legline = Line2D(self._xdata, ydata)
                legline.update_from(handle)
                self._set_artist_props(legline) # after update
                legline.set_clip_box(None)
                legline.set_markersize(self.markerscale*legline.get_markersize())

                ret.append(legline)
            elif isinstance(handle, Patch):

                p = Rectangle(xy=(min(self._xdata), y-3/4*HEIGHT),
                              width = self.handlelen, height=HEIGHT/2,
                              )
                p.update_from(handle)
                self._set_artist_props(p)
                p.set_clip_box(None)
                ret.append(p)
            elif isinstance(handle, LineCollection):
                ydata = (y-HEIGHT/2)*ones(self._xdata.shape, Float)
                legline = Line2D(self._xdata, ydata)
                self._set_artist_props(legline)
                legline.set_clip_box(None)
                lw = handle.get_linewidth()[0]
                dashes = handle.get_dashes()
                color = handle.get_colors()[0]
                legline.set_color(color)
                legline.set_linewidth(lw)
                legline.set_dashes(dashes)
                ret.append(legline)

            elif isinstance(handle, RegularPolyCollection):

                p = Rectangle(xy=(min(self._xdata), y-3/4*HEIGHT),
                              width = self.handlelen, height=HEIGHT/2,
                              )
                p.set_facecolor(handle._facecolors[0])
                if handle._edgecolors != 'None':
                    p.set_edgecolor(handle._edgecolors[0])
                self._set_artist_props(p)
                p.set_clip_box(None)
                ret.append(p)

	    else:
		ret.append(None)

        return ret

    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
        axes-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.
        """

        if not self.isaxes:
            raise Exception, 'Auto legends not available for figure legends.'

        def get_handles(ax):
            handles = ax.lines
            handles.extend(ax.patches)
            handles.extend([c for c in ax.collections if isinstance(c, LineCollection)])

            return handles

        ax = self.parent
        handles = get_handles(ax)
        vertices = []
        bboxes = []
        lines = []

        inv = ax.transAxes.inverse_xy_tup
        for handle in handles:

            if isinstance(handle, Line2D):

                xdata = handle.get_xdata(valid_only = True)
                ydata = handle.get_ydata(valid_only = True)
                trans = handle.get_transform()
                xt, yt = trans.numerix_x_y(xdata, ydata)

                # XXX need a special method in transform to do a list of verts
                averts = [inv(v) for v in zip(xt, yt)]
                lines.append(averts)

            elif isinstance(handle, Patch):

                verts = handle.get_verts()
                trans = handle.get_transform()
                tverts = trans.seq_xy_tups(verts)

                averts = [inv(v) for v in tverts]

                bbox = unit_bbox()
                bbox.update(averts, True)
                bboxes.append(bbox)

            elif isinstance(handle, LineCollection):
                hlines = handle.get_lines()
                trans = handle.get_transform()
                for line in hlines:
                    tline = trans.seq_xy_tups(line)
                    aline = [inv(v) for v in tline]
                    lines.extend(line)
                    
        return [vertices, bboxes, lines]

    def draw_frame(self, b):
        'b is a boolean.  Set draw frame to b'
        self._drawFrame = b

    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_texts(self, labels, left, upper):

        # height in axes coords
        HEIGHT = self._approx_text_height()
        pos = upper
        x = left

        ret = []  # the returned list of text instances
        for l in labels:
            text = Text(
                x=x, y=pos,
                text=l,
                fontproperties=self.prop,
                verticalalignment='top',
                horizontalalignment='left',
                )
            self._set_artist_props(text)
            ret.append(text)
            pos -= HEIGHT

        return ret


    def get_window_extent(self):
        return self.legendPatch.get_window_extent()


    def _offset(self, ox, oy):
        'Move all the artists by ox,oy (axes coords)'
        for t in self.texts:
            x,y = t.get_position()
            t.set_position( (x+ox, y+oy) )

        for h in self.legendHandles:
            if isinstance(h, Line2D):
                x,y = h.get_xdata(valid_only = True), h.get_ydata(valid_only = True)
                h.set_data( x+ox, y+oy)
            elif isinstance(h, Rectangle):
                h.xy[0] = h.xy[0] + ox
                h.xy[1] = h.xy[1] + oy
            elif isinstance(h, RegularPolygon):
                h.verts = [(x + ox, y + oy) for x, y in h.verts]

        x, y = self.legendPatch.get_x(), self.legendPatch.get_y()
        self.legendPatch.set_x(x+ox)
        self.legendPatch.set_y(y+oy)

    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 axes coords.
        """

        verts, bboxes, lines = self._auto_legend_data()

        consider = [self._loc_to_axes_coords(x, width, height) for x in range(1, len(self.codes))]

        tx, ty = self.legendPatch.xy

        candidates = []
        for l, b in consider:
            legendBox = lbwh_to_bbox(l, b, width, height)
            badness = 0
            badness = legendBox.count_contains(verts)
            ox, oy = l-tx, b-ty
            for bbox in bboxes:
                if legendBox.overlaps(bbox):
                    badness += 1

            for line in lines:
                if line_cuts_bbox(line, legendBox):
                    badness += 1

            if badness == 0:
                return ox, oy

            candidates.append((badness, (ox, oy)))

        # 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.
        minCandidate = candidates[0]
        for candidate in candidates:
            if candidate[0] < minCandidate[0]:
                minCandidate = candidate

        ox, oy = minCandidate[1]

        return ox, oy


    def _loc_to_axes_coords(self, loc, width, height):
        """Convert a location code to axes coordinates.

        - loc: a location code, which may be a pair of literal axes coords, or
          in range(1, 11). This coresponds to the possible values for
          self._loc, excluding "best".

        - width, height: the final size of the legend, axes units.
        """
        BEST, UR, UL, LL, LR, R, CL, CR, LC, UC, C = range(11)

        left = self.axespad
        right = 1.0 - (self.axespad + width)
        upper = 1.0 - (self.axespad + height)
        lower = self.axespad
        centerx = 0.5 - (width/2.0)
        centery = 0.5 - (height/2.0)

        if loc == UR:
            return right, upper
        if loc == UL:
            return left, upper
        if loc == LL:
            return left, lower
        if loc == LR:
            return right, lower
        if loc == CL:
            return left, centery
        if loc in (CR, R):
            return right, centery
        if loc == LC:
            return centerx, lower
        if loc == UC:
            return centerx, upper
        if loc == C:
            return centerx, centery
        raise TypeError, "%r isn't an understood type code." % (loc,)

    def _update_positions(self, renderer):
        # called from renderer to allow more precise estimates of
        # widths and heights with get_window_extent

        if not len(self.legendHandles) and not len(self.texts): return
        def get_tbounds(text):  #get text bounds in axes coords
            bbox = text.get_window_extent(renderer)
            bboxa = inverse_transform_bbox(self._transform, bbox)
            return bboxa.get_bounds()

        hpos = []
        for t, tabove in zip(self.texts[1:], self.texts[:-1]):
            x,y = t.get_position()
            l,b,w,h = get_tbounds(tabove)
            b -= self.labelsep
            h += 2*self.labelsep
            hpos.append( (b,h) )
            t.set_position( (x, b-0.1*h) )

        # now do the same for last line
        l,b,w,h = get_tbounds(self.texts[-1])
        b -= self.labelsep
        h += 2*self.labelsep
        hpos.append( (b,h) )

        for handle, tup in zip(self.legendHandles, hpos):
            y,h = tup
            if isinstance(handle, Line2D):
                ydata = y*ones(self._xdata.shape, Float)
                handle.set_ydata(ydata+h/2)
            elif isinstance(handle, Rectangle):
                handle.set_y(y+1/4*h)
                handle.set_height(h/2)

        # Set the data for the legend patch
        bbox = self._get_handle_text_bbox(renderer).deepcopy()

        bbox.scale(1 + self.pad, 1 + self.pad)
        l,b,w,h = bbox.get_bounds()
        self.legendPatch.set_bounds(l,b,w,h)

        BEST, UR, UL, LL, LR, R, CL, CR, LC, UC, C = range(11)
        ox, oy = 0, 0                           # center


        if iterable(self._loc) and len(self._loc)==2:
            xo = self.legendPatch.get_x()
            yo = self.legendPatch.get_y()
            x, y = self._loc
            ox = x-xo
            oy = y-yo
            self._offset(ox, oy)
        else:
            if self._loc in (BEST,):
                ox, oy = self._find_best_position(w, h)
            if self._loc in (UL, LL, CL):           # left
                ox = self.axespad - l
            if self._loc in (UR, LR, R, CR):  # right
                ox = 1 - (l + w + self.axespad)
            if self._loc in (UR, UL, UC):     # upper
                oy = 1 - (b + h + self.axespad)
            if self._loc in (LL, LR, LC):           # lower
                oy = self.axespad - b
            if self._loc in (LC, UC, C):            # center x
                ox = (0.5-w/2)-l
            if self._loc in (CL, CR, C):            # center y
                oy = (0.5-h/2)-b
            self._offset(ox, oy)
Ejemplo n.º 4
0
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,
      'upper right'  : 1,  (default)
      '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,

    Return value is a sequence of text, line instances that make
    up the legend
    """


    codes = {'best'         : 0,
             'upper right'  : 1,  # default
             '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,
             }




    def __init__(self, parent, handles, labels, loc,
                 isaxes=True,
                 numpoints = 4,      # the number of points in the legend line
                 prop = FontProperties(size='smaller'),
                 pad = 0.2,          # the fractional whitespace inside the legend border
                 markerscale = 0.6,    # the relative size of legend markers vs. original
                 # the following dimensions are in axes coords
                 labelsep = 0.005,     # the vertical space between the legend entries
                 handlelen = 0.05,     # the length of the legend lines
                 handletextsep = 0.02, # the space between the legend line and legend text
                 axespad = 0.02,       # the border between the axes and legend edge

                 shadow=False,
                 ):
        """
  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
  loc                   # a location code
  isaxes=True           # whether this is an axes legend
  numpoints = 4         # the number of points in the legend line
  fontprop = FontProperties(size='smaller')  # the font property
  pad = 0.2             # the fractional whitespace inside the legend border
  markerscale = 0.6     # the relative size of legend markers vs. original
  shadow                # if True, draw a shadow behind legend

The following dimensions are in axes coords
  labelsep = 0.005     # the vertical space between the legend entries
  handlelen = 0.05     # the length of the legend lines
  handletextsep = 0.02 # the space between the legend line and legend text
  axespad = 0.02       # the border between the axes and legend edge
        """
        Artist.__init__(self)
        if is_string_like(loc) and not self.codes.has_key(loc):
            warnings.warn('Unrecognized location %s. Falling back on upper right; valid locations are\n%s\t' %(loc, '\n\t'.join(self.codes.keys())))
        if is_string_like(loc): loc = self.codes.get(loc, 1)

        self.numpoints = numpoints
        self.prop = prop
        self.fontsize = prop.get_size_in_points()
        self.pad = pad
        self.markerscale = markerscale
        self.labelsep = labelsep
        self.handlelen = handlelen
        self.handletextsep = handletextsep
        self.axespad = axespad
        self.shadow = shadow

        self.isaxes = isaxes
        if isaxes:  # parent is an Axes
            self.set_figure(parent.figure)
        else:        # parent is a Figure
            self.set_figure(parent)

        self.parent = parent
        self.set_transform( get_bbox_transform( unit_bbox(), parent.bbox) )
        self._loc = loc

        # make a trial box in the middle of the axes.  relocate it
        # based on it's bbox
        left, upper = 0.5, 0.5
        if self.numpoints == 1:
            self._xdata = array([left + self.handlelen*0.5])
        else:
            self._xdata = linspace(left, left + self.handlelen, self.numpoints)
        textleft = left+ self.handlelen+self.handletextsep
        self.texts = self._get_texts(labels, textleft, upper)
        self.legendHandles = self._get_handles(handles, self.texts)

        left, top = self.texts[-1].get_position()
        HEIGHT = self._approx_text_height()
        bottom = top-HEIGHT
        left -= self.handlelen + self.handletextsep + self.pad
        self.legendPatch = Rectangle(
            xy=(left, bottom), width=0.5, height=HEIGHT*len(self.texts),
            facecolor='w', edgecolor='k',
            )
        self._set_artist_props(self.legendPatch)
        self._drawFrame = True

    def _set_artist_props(self, a):
        a.set_figure(self.figure)
        a.set_transform(self._transform)

    def _approx_text_height(self):
        return self.fontsize/72.0*self.figure.dpi.get()/self.parent.bbox.height()


    def draw(self, renderer):
        if not self.get_visible(): return
        renderer.open_group('legend')
        self._update_positions(renderer)
        if self._drawFrame:
            if self.shadow:
                shadow = Shadow(self.legendPatch, -0.005, -0.005)
                shadow.draw(renderer)
            self.legendPatch.draw(renderer)


        for h in self.legendHandles:
            if h is not None:
		h.draw(renderer)
        	if 0: bbox_artist(h, renderer)

        for t in self.texts:
            if 0: bbox_artist(t, renderer)
            t.draw(renderer)
        renderer.close_group('legend')
        #draw_bbox(self.save, renderer, 'g')
        #draw_bbox(self.ibox, renderer, 'r', self._transform)

    def _get_handle_text_bbox(self, renderer):
        'Get a bbox for the text and lines in axes coords'

        bboxesText = [t.get_window_extent(renderer) for t in self.texts]
        bboxesHandles = [h.get_window_extent(renderer) for h in self.legendHandles if h is not None]


        bboxesAll = bboxesText
        bboxesAll.extend(bboxesHandles)
        bbox = bbox_all(bboxesAll)
        self.save = bbox

        ibox =  inverse_transform_bbox(self._transform, bbox)
        self.ibox = ibox

        return ibox

    def _get_handles(self, handles, texts):
        HEIGHT = self._approx_text_height()

        ret = []   # the returned legend lines

        for handle, label in zip(handles, texts):
            x, y = label.get_position()
            x -= self.handlelen + self.handletextsep
            if isinstance(handle, Line2D):
                ydata = (y-HEIGHT/2)*ones(self._xdata.shape, Float)
                legline = Line2D(self._xdata, ydata)
                legline.update_from(handle)
                self._set_artist_props(legline) # after update
                legline.set_clip_box(None)
                legline.set_markersize(self.markerscale*legline.get_markersize())
                legline.set_data_clipping(False)
                ret.append(legline)
            elif isinstance(handle, Patch):

                p = Rectangle(xy=(min(self._xdata), y-3/4*HEIGHT),
                              width = self.handlelen, height=HEIGHT/2,
                              )
                p.update_from(handle)
                self._set_artist_props(p)
                p.set_clip_box(None)
                ret.append(p)
            elif isinstance(handle, LineCollection):
                ydata = (y-HEIGHT/2)*ones(self._xdata.shape, Float)
                legline = Line2D(self._xdata, ydata)
                self._set_artist_props(legline)
                legline.set_clip_box(None)
                lw = handle.get_linewidths()[0]
                dashes = handle.get_dashes()
                color = handle.get_colors()[0]
                legline.set_color(color)
                legline.set_linewidth(lw)
                legline.set_dashes(dashes)
                ret.append(legline)

	    else:
		ret.append(None)

        return ret

    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
        axes-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.
        """

        if not self.isaxes:
            raise Exception, 'Auto legends not available for figure legends.'

        def get_handles(ax):
            handles = ax.lines
            handles.extend(ax.patches)
            handles.extend([c for c in ax.collections if isinstance(c, LineCollection)])
            return handles

        ax = self.parent
        handles = get_handles(ax)
        vertices = []
        bboxes = []
        lines = []

        inv = ax.transAxes.inverse_xy_tup
        for handle in handles:

            if isinstance(handle, Line2D):

                xdata = handle.get_xdata(valid_only = True)
                ydata = handle.get_ydata(valid_only = True)
                trans = handle.get_transform()
                xt, yt = trans.numerix_x_y(xdata, ydata)

                # XXX need a special method in transform to do a list of verts
                averts = [inv(v) for v in zip(xt, yt)]
                lines.append(averts)

            elif isinstance(handle, Patch):

                verts = handle.get_verts()
                trans = handle.get_transform()
                tverts = trans.seq_xy_tups(verts)

                averts = [inv(v) for v in tverts]

                bbox = unit_bbox()
                bbox.update(averts, True)
                bboxes.append(bbox)

            elif isinstance(handle, LineCollection):
                hlines = handle.get_lines()
                trans = handle.get_transform()
                for line in hlines:
                    tline = trans.seq_xy_tups(line)
                    aline = [inv(v) for v in tline]
                    lines.extend(line)
                    
        return [vertices, bboxes, lines]

    def draw_frame(self, b):
        'b is a boolean.  Set draw frame to b'
        self._drawFrame = b

    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_texts(self, labels, left, upper):

        # height in axes coords
        HEIGHT = self._approx_text_height()
        pos = upper
        x = left

        ret = []  # the returned list of text instances
        for l in labels:
            text = Text(
                x=x, y=pos,
                text=l,
                fontproperties=self.prop,
                verticalalignment='top',
                horizontalalignment='left',
                )
            self._set_artist_props(text)
            ret.append(text)
            pos -= HEIGHT

        return ret


    def get_window_extent(self):
        return self.legendPatch.get_window_extent()


    def _offset(self, ox, oy):
        'Move all the artists by ox,oy (axes coords)'
        for t in self.texts:
            x,y = t.get_position()
            t.set_position( (x+ox, y+oy) )

        for h in self.legendHandles:
            if isinstance(h, Line2D):
                x,y = h.get_xdata(valid_only = True), h.get_ydata(valid_only = True)
                h.set_data( x+ox, y+oy)
            elif isinstance(h, Rectangle):
                h.xy[0] = h.xy[0] + ox
                h.xy[1] = h.xy[1] + oy

        x, y = self.legendPatch.get_x(), self.legendPatch.get_y()
        self.legendPatch.set_x(x+ox)
        self.legendPatch.set_y(y+oy)

    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 axes coords.
        """

        verts, bboxes, lines = self._auto_legend_data()

        consider = [self._loc_to_axes_coords(x, width, height) for x in range(1, len(self.codes))]

        tx, ty = self.legendPatch.xy

        candidates = []
        for l, b in consider:
            legendBox = lbwh_to_bbox(l, b, width, height)
            badness = 0
            badness = legendBox.count_contains(verts)
            ox, oy = l-tx, b-ty
            for bbox in bboxes:
                if legendBox.overlaps(bbox):
                    badness += 1

            for line in lines:
                if line_cuts_bbox(line, legendBox):
                    badness += 1

            if badness == 0:
                return ox, oy

            candidates.append((badness, (ox, oy)))

        # 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.
        minCandidate = candidates[0]
        for candidate in candidates:
            if candidate[0] < minCandidate[0]:
                minCandidate = candidate

        ox, oy = minCandidate[1]

        return ox, oy


    def _loc_to_axes_coords(self, loc, width, height):
        """Convert a location code to axes coordinates.

        - loc: a location code, which may be a pair of literal axes coords, or
          in range(1, 11). This coresponds to the possible values for
          self._loc, excluding "best".

        - width, height: the final size of the legend, axes units.
        """
        BEST, UR, UL, LL, LR, R, CL, CR, LC, UC, C = range(11)

        left = self.axespad
        right = 1.0 - (self.axespad + width)
        upper = 1.0 - (self.axespad + height)
        lower = self.axespad
        centerx = 0.5 - (width/2.0)
        centery = 0.5 - (height/2.0)

        if loc == UR:
            return right, upper
        if loc == UL:
            return left, upper
        if loc == LL:
            return left, lower
        if loc == LR:
            return right, lower
        if loc == CL:
            return left, centery
        if loc in (CR, R):
            return right, centery
        if loc == LC:
            return centerx, lower
        if loc == UC:
            return centerx, upper
        if loc == C:
            return centerx, centery
        raise TypeError, "%r isn't an understood type code." % (loc,)

    def _update_positions(self, renderer):
        # called from renderer to allow more precise estimates of
        # widths and heights with get_window_extent

        def get_tbounds(text):  #get text bounds in axes coords
            bbox = text.get_window_extent(renderer)
            bboxa = inverse_transform_bbox(self._transform, bbox)
            return bboxa.get_bounds()

        hpos = []
        for t, tabove in zip(self.texts[1:], self.texts[:-1]):
            x,y = t.get_position()
            l,b,w,h = get_tbounds(tabove)
            b -= self.labelsep
            h += 2*self.labelsep
            hpos.append( (b,h) )
            t.set_position( (x, b-0.1*h) )

        # now do the same for last line
        l,b,w,h = get_tbounds(self.texts[-1])
        b -= self.labelsep
        h += 2*self.labelsep
        hpos.append( (b,h) )

        for handle, tup in zip(self.legendHandles, hpos):
            y,h = tup
            if isinstance(handle, Line2D):
                ydata = y*ones(self._xdata.shape, Float)
                handle.set_ydata(ydata+h/2)
            elif isinstance(handle, Rectangle):
                handle.set_y(y+1/4*h)
                handle.set_height(h/2)

        # Set the data for the legend patch
        bbox = self._get_handle_text_bbox(renderer).deepcopy()

        bbox.scale(1 + self.pad, 1 + self.pad)
        l,b,w,h = bbox.get_bounds()
        self.legendPatch.set_bounds(l,b,w,h)

        BEST, UR, UL, LL, LR, R, CL, CR, LC, UC, C = range(11)
        ox, oy = 0, 0                           # center


        if iterable(self._loc) and len(self._loc)==2:
            xo = self.legendPatch.get_x()
            yo = self.legendPatch.get_y()
            x, y = self._loc
            ox = x-xo
            oy = y-yo
            self._offset(ox, oy)
        else:
            if self._loc in (BEST,):
                ox, oy = self._find_best_position(w, h)
            if self._loc in (UL, LL, CL):           # left
                ox = self.axespad - l
            if self._loc in (UR, LR, R, CR):  # right
                ox = 1 - (l + w + self.axespad)
            if self._loc in (UR, UL, UC):     # upper
                oy = 1 - (b + h + self.axespad)
            if self._loc in (LL, LR, LC):           # lower
                oy = self.axespad - b
            if self._loc in (LC, UC, C):            # center x
                ox = (0.5-w/2)-l
            if self._loc in (CL, CR, C):            # center y
                oy = (0.5-h/2)-b
            self._offset(ox, oy)
Ejemplo n.º 5
0
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,  (currently not supported, defaults to upper right)
      'upper right'  : 1,  (default)
      '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,
 
    Return value is a sequence of text, line instances that make
    up the legend
    """

    codes = {
        "best": 0,
        "upper right": 1,  # default
        "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,
    }

    NUMPOINTS = 4  # the number of points in the legend line
    FONTSIZE = 10
    PAD = 0.2  # the fractional whitespace inside the legend border
    # the following dimensions are in axes coords
    LABELSEP = 0.005  # the vertical space between the legend entries
    HANDLELEN = 0.05  # the length of the legend lines
    HANDLETEXTSEP = 0.02  # the space between the legend line and legend text
    AXESPAD = 0.02  # the border between the axes and legend edge

    def __init__(self, parent, handles, labels, loc, isaxes=True):
        Artist.__init__(self)
        if is_string_like(loc) and not self.codes.has_key(loc):
            verbose.report_error(
                "Unrecognized location %s. Falling back on upper right; valid locations are\n%s\t"
                % (loc, "\n\t".join(self.codes.keys()))
            )
        if is_string_like(loc):
            loc = self.codes.get(loc, 1)

        if isaxes:  # parent is an Axes
            self.set_figure(parent.figure)
        else:  # parent is a Figure
            self.set_figure(parent)

        self.parent = parent
        self.set_transform(get_bbox_transform(unit_bbox(), parent.bbox))
        self._loc = loc

        # make a trial box in the middle of the axes.  relocate it
        # based on it's bbox
        left, upper = 0.5, 0.5
        if self.NUMPOINTS == 1:
            self._xdata = array([left + self.HANDLELEN * 0.5])
        else:
            self._xdata = linspace(left, left + self.HANDLELEN, self.NUMPOINTS)
        textleft = left + self.HANDLELEN + self.HANDLETEXTSEP
        self.texts = self._get_texts(labels, textleft, upper)
        self.handles = self._get_handles(handles, self.texts)

        left, top = self.texts[-1].get_position()
        HEIGHT = self._approx_text_height()
        bottom = top - HEIGHT
        left -= self.HANDLELEN + self.HANDLETEXTSEP + self.PAD
        self.legendPatch = Rectangle(
            xy=(left, bottom), width=0.5, height=HEIGHT * len(self.texts), facecolor="w", edgecolor="k"
        )
        self._set_artist_props(self.legendPatch)
        self._drawFrame = True

    def _set_artist_props(self, a):
        a.set_figure(self.figure)
        a.set_transform(self._transform)

    def _approx_text_height(self):
        return self.FONTSIZE / 72.0 * self.figure.dpi.get() / self.parent.bbox.height()

    def draw(self, renderer):
        renderer.open_group("legend")
        self._update_positions(renderer)
        if self._drawFrame:
            self.legendPatch.draw(renderer)
        for h in self.handles:
            if h is not None:
                h.draw(renderer)
                if 0:
                    bbox_artist(h, renderer)

        for t in self.texts:
            if 0:
                bbox_artist(t, renderer)
            t.draw(renderer)
        renderer.close_group("legend")
        # draw_bbox(self.save, renderer, 'g')
        # draw_bbox(self.ibox, renderer, 'r', self._transform)

    def _get_handle_text_bbox(self, renderer):
        "Get a bbox for the text and lines in axes coords"
        boxes = []
        bboxesText = [t.get_window_extent(renderer) for t in self.texts]
        bboxesHandles = [h.get_window_extent(renderer) for h in self.handles if h is not None]

        bboxesAll = bboxesText
        bboxesAll.extend(bboxesHandles)
        bbox = bbox_all(bboxesAll)
        self.save = bbox

        ibox = inverse_transform_bbox(self._transform, bbox)
        self.ibox = ibox

        return ibox

    def _get_handles(self, handles, texts):
        HEIGHT = self._approx_text_height()

        ret = []  # the returned legend lines
        for handle, label in zip(handles, texts):
            x, y = label.get_position()
            x -= self.HANDLELEN + self.HANDLETEXTSEP
            if isinstance(handle, Line2D):
                ydata = (y - HEIGHT / 2) * ones(self._xdata.shape, Float)
                legline = Line2D(self._xdata, ydata)
                self._set_artist_props(legline)
                legline.copy_properties(handle)
                legline.set_markersize(0.6 * legline.get_markersize())
                legline.set_data_clipping(False)
                ret.append(legline)
            elif isinstance(handle, Patch):

                p = Rectangle(xy=(min(self._xdata), y - 3 / 4 * HEIGHT), width=self.HANDLELEN, height=HEIGHT / 2)
                self._set_artist_props(p)
                p.copy_properties(handle)
                ret.append(p)
            else:
                ret.append(None)

        return ret

    def draw_frame(self, b):
        "b is a boolean.  Set draw frame to b"
        self._drawFrame = b

    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.handles if isinstance(h, Line2D)]

    def get_patches(self):
        "return a list of patch instances in the legend"
        return [h for h in self.handles if isinstance(h, Patch)]

    def get_texts(self):
        "return a list of text.Text instance in the legend"
        return self.texts

    def _get_texts(self, labels, left, upper):

        # height in axes coords
        HEIGHT = self._approx_text_height()
        pos = upper
        x = left

        ret = []  # the returned list of text instances
        for l in labels:
            text = Text(
                x=x,
                y=pos,
                text=l,
                fontproperties=FontProperties(size="smaller"),
                verticalalignment="top",
                horizontalalignment="left",
            )
            self._set_artist_props(text)
            ret.append(text)
            pos -= HEIGHT

        return ret

    def get_window_extent(self):
        return self.legendPatch.get_window_extent()

    def _offset(self, ox, oy):
        "Move all the artists by ox,oy (axes coords)"
        for t in self.texts:
            x, y = t.get_position()
            t.set_position((x + ox, y + oy))

        for h in self.handles:
            if isinstance(h, Line2D):
                x, y = h.get_xdata(), h.get_ydata()
                h.set_data(x + ox, y + oy)
            elif isinstance(h, Rectangle):
                h.xy[0] = h.xy[0] + ox
                h.xy[1] = h.xy[1] + oy

        x, y = self.legendPatch.get_x(), self.legendPatch.get_y()
        self.legendPatch.set_x(x + ox)
        self.legendPatch.set_y(y + oy)

    def _update_positions(self, renderer):
        # called from renderer to allow more precise estimates of
        # widths and heights with get_window_extent

        def get_tbounds(text):  # get text bounds in axes coords
            bbox = text.get_window_extent(renderer)
            bboxa = inverse_transform_bbox(self._transform, bbox)
            return bboxa.get_bounds()

        hpos = []
        for t, tabove in zip(self.texts[1:], self.texts[:-1]):
            x, y = t.get_position()
            l, b, w, h = get_tbounds(tabove)
            hpos.append((b, h))
            t.set_position((x, b - 0.1 * h))

        # now do the same for last line
        l, b, w, h = get_tbounds(self.texts[-1])
        hpos.append((b, h))

        for handle, tup in zip(self.handles, hpos):
            y, h = tup
            if isinstance(handle, Line2D):
                ydata = y * ones(self._xdata.shape, Float)
                handle.set_ydata(ydata + h / 2)
            elif isinstance(handle, Rectangle):
                handle.set_y(y + 1 / 4 * h)
                handle.set_height(h / 2)

        # Set the data for the legend patch
        bbox = self._get_handle_text_bbox(renderer).deepcopy()
        bbox.scale(1 + self.PAD, 1 + self.PAD)
        l, b, w, h = bbox.get_bounds()
        self.legendPatch.set_bounds(l, b, w, h)

        BEST, UR, UL, LL, LR, R, CL, CR, LC, UC, C = range(11)
        ox, oy = 0, 0  # center

        if iterable(self._loc) and len(self._loc) == 2:
            xo = self.legendPatch.get_x()
            yo = self.legendPatch.get_y()
            x, y = self._loc
            ox = x - xo
            oy = y - yo
            self._offset(ox, oy)
        else:
            if self._loc in (UL, LL, CL):  # left
                ox = self.AXESPAD - l
            if self._loc in (BEST, UR, LR, R, CR):  # right
                ox = 1 - (l + w + self.AXESPAD)
            if self._loc in (BEST, UR, UL, UC):  # upper
                oy = 1 - (b + h + self.AXESPAD)
            if self._loc in (LL, LR, LC):  # lower
                oy = self.AXESPAD - b
            if self._loc in (LC, UC, C):  # center x
                ox = (0.5 - w / 2) - l
            if self._loc in (CL, CR, C):  # center y
                oy = (0.5 - h / 2) - b
            self._offset(ox, oy)