def draw(self, renderer): self._init() if self._new_UV: verts = self._make_verts(self.U, self.V) self.set_verts(verts) self._new_UV = False PolyCollection.draw(self, renderer)
class QuiverKey(Artist): ''' Labelled arrow for use as a quiver plot scale key. ''' halign = {'N': 'center', 'S': 'center', 'E': 'left', 'W': 'right'} valign = {'N': 'bottom', 'S': 'top', 'E': 'center', 'W': 'center'} pivot = {'N': 'mid', 'S': 'mid', 'E': 'tip', 'W': 'tail'} def __init__(self, Q, X, Y, U, label, **kw): Artist.__init__(self) self.Q = Q self.X = X self.Y = Y self.U = U self.coord = kw.pop('coordinates', 'axes') self.color = kw.pop('color', None) self.label = label self.labelsep = T.Value(kw.pop('labelsep', 0.1)) * Q.ax.figure.dpi self.labelpos = kw.pop('labelpos', 'N') self.labelcolor = kw.pop('labelcolor', None) self.fontproperties = kw.pop('fontproperties', dict()) self.kw = kw self.text = Text(text=label, horizontalalignment=self.halign[self.labelpos], verticalalignment=self.valign[self.labelpos], fontproperties=FontProperties(**self.fontproperties)) if self.labelcolor is not None: self.text.set_color(self.labelcolor) self._initialized = False self.zorder = Q.zorder + 0.1 __init__.__doc__ = _quiverkey_doc def _init(self): if not self._initialized: self._set_transform() _pivot = self.Q.pivot self.Q.pivot = self.pivot[self.labelpos] self.verts = self.Q._make_verts(nx.array([self.U]), nx.zeros((1,))) self.Q.pivot = _pivot kw = self.Q.polykw kw.update(self.kw) self.vector = PolyCollection(self.verts, offsets=[(self.X,self.Y)], transOffset=self.get_transform(), **kw) if self.color is not None: self.vector.set_color(self.color) self.vector.set_transform(self.Q.get_transform()) self._initialized = True def _text_x(self, x): if self.labelpos == 'E': return x + self.labelsep.get() elif self.labelpos == 'W': return x - self.labelsep.get() else: return x def _text_y(self, y): if self.labelpos == 'N': return y + self.labelsep.get() elif self.labelpos == 'S': return y - self.labelsep.get() else: return y def draw(self, renderer): self._init() self.vector.draw(renderer) x, y = self.get_transform().xy_tup((self.X, self.Y)) self.text.set_x(self._text_x(x)) self.text.set_y(self._text_y(y)) self.text.draw(renderer) def _set_transform(self): if self.coord == 'data': self.set_transform(self.Q.ax.transData) elif self.coord == 'axes': self.set_transform(self.Q.ax.transAxes) elif self.coord == 'figure': self.set_transform(self.Q.ax.figure.transFigure) elif self.coord == 'inches': dx = ax.figure.dpi bb = T.Bbox(T.origin(), T.Point(dx, dx)) trans = T.get_bbox_transform(T.unit_bbox(), bb) self.set_transform(trans) else: raise ValueError('unrecognized coordinates') quiverkey_doc = _quiverkey_doc