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 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 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
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 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)
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
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
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)
