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

martist.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_inches = kw.pop('labelsep', 0.1)

self.labelsep = (self._labelsep_inches * Q.ax.figure.dpi)

def on_dpi_change(fig):

self.labelsep = (self._labelsep_inches * fig.dpi)

self._initialized = False # simple brute force update

# works because _init is called

# at the start of draw.

Q.ax.figure.callbacks.connect('dpi_changed', on_dpi_change)

self.labelpos = kw.pop('labelpos', 'N')

self.labelcolor = kw.pop('labelcolor', None)

self.fontproperties = kw.pop('fontproperties', dict())

self.kw = kw

_fp = self.fontproperties

#boxprops = dict(facecolor='red')

self.text = mtext.Text(text=label, # bbox=boxprops,

horizontalalignment=self.halign[self.labelpos],

verticalalignment=self.valign[self.labelpos],

fontproperties=font_manager.FontProperties(**_fp))

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 True: ##not self._initialized:

self._set_transform()

_pivot = self.Q.pivot

self.Q.pivot = self.pivot[self.labelpos]

# Hack: save and restore the Umask

_mask = self.Q.Umask

self.Q.Umask = ma.nomask

self.verts = self.Q._make_verts(np.array([self.U]),

np.zeros((1,)))

self.Q.Umask = _mask

self.Q.pivot = _pivot

kw = self.Q.polykw

kw.update(self.kw)

self.vector = collections.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

elif self.labelpos == 'W':

return x - self.labelsep

else:

return x

def _text_y(self, y):

if self.labelpos == 'N':

return y + self.labelsep

elif self.labelpos == 'S':

return y - self.labelsep

else:

return y

@allow_rasterization

def draw(self, renderer):

self._init()

self.vector.draw(renderer)

x, y = self.get_transform().transform_point((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':

self.set_transform(self.Q.ax.figure.dpi_scale_trans)

else:

raise ValueError('unrecognized coordinates')

def set_figure(self, fig):

martist.Artist.set_figure(self, fig)

self.text.set_figure(fig)

def contains(self, mouseevent):

# Maybe the dictionary should allow one to

# distinguish between a text hit and a vector hit.

if (self.text.contains(mouseevent)[0]

or self.vector.contains(mouseevent)[0]):

return True, {}

return False, {}

X, Y, U, V, C = [None]*5

args = list(args)

# The use of atleast_1d allows for handling scalar arguments while also

# keeping masked arrays

if len(args) == 3 or len(args) == 5:

C = np.atleast_1d(args.pop(-1))

V = np.atleast_1d(args.pop(-1))

U = np.atleast_1d(args.pop(-1))

if U.ndim == 1:

nr, nc = 1, U.shape[0]

else:

nr, nc = U.shape

if len(args) == 2: # remaining after removing U,V,C

X, Y = [np.array(a).ravel() for a in args]

if len(X) == nc and len(Y) == nr:

X, Y = [a.ravel() for a in np.meshgrid(X, Y)]

else:

indexgrid = np.meshgrid(np.arange(nc), np.arange(nr))

X, Y = [np.ravel(a) for a in indexgrid]

"""

@docstring.Substitution(_quiver_doc)

def __init__(self, ax, *args, **kw):

"""

self.ax = ax

X, Y, U, V, C = _parse_args(*args)

self.X = X

self.Y = Y

self.XY = np.hstack((X[:,np.newaxis], Y[:,np.newaxis]))

self.N = len(X)

self.scale = kw.pop('scale', None)

self.headwidth = kw.pop('headwidth', 3)

self.headlength = float(kw.pop('headlength', 5))

self.headaxislength = kw.pop('headaxislength', 4.5)

self.minshaft = kw.pop('minshaft', 1)

self.minlength = kw.pop('minlength', 1)

self.units = kw.pop('units', 'width')

self.scale_units = kw.pop('scale_units', None)

self.angles = kw.pop('angles', 'uv')

self.width = kw.pop('width', None)

self.color = kw.pop('color', 'k')

self.pivot = kw.pop('pivot', 'tail')

kw.setdefault('facecolors', self.color)

kw.setdefault('linewidths', (0,))

collections.PolyCollection.__init__(self, [], offsets=self.XY,

transOffset=ax.transData,

closed=False,

**kw)

self.polykw = kw

self.set_UVC(U, V, C)

self._initialized = False

self.keyvec = None

self.keytext = None

def on_dpi_change(fig):

self._new_UV = True # vertices depend on width, span

# which in turn depend on dpi

self._initialized = False # simple brute force update

# works because _init is called

# at the start of draw.

self.ax.figure.callbacks.connect('dpi_changed', on_dpi_change)

def _init(self):

"""

# It seems that there are not enough event notifications

# available to have this work on an as-needed basis at present.

if True: ##not self._initialized:

trans = self._set_transform()

ax = self.ax

sx, sy = trans.inverted().transform_point(

(ax.bbox.width, ax.bbox.height))

self.span = sx

if self.width is None:

sn = max(8, min(25, math.sqrt(self.N)))

self.width = 0.06 * self.span / sn

@allow_rasterization

def draw(self, renderer):

self._init()

if (self._new_UV or self.angles == 'xy'

or self.scale_units in ['x','y', 'xy']):

verts = self._make_verts(self.U, self.V)

self.set_verts(verts, closed=False)

self._new_UV = False

collections.PolyCollection.draw(self, renderer)

def set_UVC(self, U, V, C=None):

U = ma.masked_invalid(U, copy=False).ravel()

V = ma.masked_invalid(V, copy=False).ravel()

mask = ma.mask_or(U.mask, V.mask, copy=False, shrink=True)

if C is not None:

C = ma.masked_invalid(C, copy=False).ravel()

mask = ma.mask_or(mask, C.mask, copy=False, shrink=True)

if mask is ma.nomask:

C = C.filled()

else:

C = ma.array(C, mask=mask, copy=False)

self.U = U.filled(1)

self.V = V.filled(1)

self.Umask = mask

if C is not None:

self.set_array(C)

self._new_UV = True

def _dots_per_unit(self, units):

"""

ax = self.ax

if units in ('x', 'y', 'xy'):

if units == 'x':

dx0 = ax.viewLim.width

dx1 = ax.bbox.width

elif units == 'y':

dx0 = ax.viewLim.height

dx1 = ax.bbox.height

else: # 'xy' is assumed

dxx0 = ax.viewLim.width

dxx1 = ax.bbox.width

dyy0 = ax.viewLim.height

dyy1 = ax.bbox.height

dx1 = np.sqrt(dxx1*dxx1+dyy1*dyy1)

dx0 = np.sqrt(dxx0*dxx0+dyy0*dyy0)

dx = dx1/dx0

else:

if units == 'width':

dx = ax.bbox.width

elif units == 'height':

dx = ax.bbox.height

elif units == 'dots':

dx = 1.0

elif units == 'inches':

dx = ax.figure.dpi

else:

raise ValueError('unrecognized units')

return dx

def _set_transform(self):

"""

dx = self._dots_per_unit(self.units)

self._trans_scale = dx # pixels per arrow width unit

trans = transforms.Affine2D().scale(dx)

self.set_transform(trans)

return trans

def _angles_lengths(self, U, V, eps=1):

xy = self.ax.transData.transform(self.XY)

uv = np.hstack((U[:,np.newaxis], V[:,np.newaxis]))

xyp = self.ax.transData.transform(self.XY + eps * uv)

dxy = xyp - xy

angles = np.arctan2(dxy[:,1], dxy[:,0])

lengths = np.absolute(dxy[:,0] + dxy[:,1]*1j) / eps

return angles, lengths

def _make_verts(self, U, V):

uv = (U+V*1j)

if self.angles == 'xy' and self.scale_units == 'xy':

# Here eps is 1 so that if we get U, V by diffing

# the X, Y arrays, the vectors will connect the

# points, regardless of the axis scaling (including log).

angles, lengths = self._angles_lengths(U, V, eps=1)

elif self.angles == 'xy' or self.scale_units == 'xy':

# Calculate eps based on the extents of the plot

# so that we don't end up with roundoff error from

# adding a small number to a large.

angles, lengths = self._angles_lengths(U, V,

eps=np.abs(self.ax.dataLim.extents).max() * 0.001)

if self.scale_units == 'xy':

a = lengths

else:

a = np.absolute(uv)

if self.scale is None:

sn = max(10, math.sqrt(self.N))

if self.Umask is not ma.nomask:

amean = a[~self.Umask].mean()

else:

amean = a.mean()

scale = 1.8 * amean * sn / self.span # crude auto-scaling

# scale is typical arrow length as a multiple

# of the arrow width

if self.scale_units is None:

if self.scale is None:

self.scale = scale

widthu_per_lenu = 1.0

else:

if self.scale_units == 'xy':

dx = 1

else:

dx = self._dots_per_unit(self.scale_units)

widthu_per_lenu = dx/self._trans_scale

if self.scale is None:

self.scale = scale * widthu_per_lenu

length = a * (widthu_per_lenu / (self.scale * self.width))

X, Y = self._h_arrows(length)

if self.angles == 'xy':

theta = angles

elif self.angles == 'uv':

theta = np.angle(uv)

else:

# Make a copy to avoid changing the input array.

theta = ma.masked_invalid(self.angles, copy=True).filled(0)

theta = theta.ravel()

theta *= (np.pi/180.0)

theta.shape = (theta.shape[0], 1) # for broadcasting

xy = (X+Y*1j) * np.exp(1j*theta)*self.width

xy = xy[:,:,np.newaxis]

XY = np.concatenate((xy.real, xy.imag), axis=2)

if self.Umask is not ma.nomask:

XY = ma.array(XY)

XY[self.Umask] = ma.masked

# This might be handled more efficiently with nans, given

# that nans will end up in the paths anyway.

return XY

def _h_arrows(self, length):

""" length is in arrow width units """

# It might be possible to streamline the code

# and speed it up a bit by using complex (x,y)

# instead of separate arrays; but any gain would be slight.

minsh = self.minshaft * self.headlength

N = len(length)

length = length.reshape(N, 1)

# This number is chosen based on when pixel values overflow in Agg

# causing rendering errors

#length = np.minimum(length, 2 ** 16)

np.clip(length, 0, 2**16, out=length)

# x, y: normal horizontal arrow

x = np.array([0, -self.headaxislength,

-self.headlength, 0], np.float64)

x = x + np.array([0,1,1,1]) * length

y = 0.5 * np.array([1, 1, self.headwidth, 0], np.float64)

y = np.repeat(y[np.newaxis,:], N, axis=0)

# x0, y0: arrow without shaft, for short vectors

x0 = np.array([0, minsh-self.headaxislength,

minsh-self.headlength, minsh], np.float64)

y0 = 0.5 * np.array([1, 1, self.headwidth, 0], np.float64)

ii = [0,1,2,3,2,1,0,0]

X = x.take(ii, 1)

Y = y.take(ii, 1)

Y[:, 3:-1] *= -1

X0 = x0.take(ii)

Y0 = y0.take(ii)

Y0[3:-1] *= -1

shrink = length/minsh

X0 = shrink * X0[np.newaxis,:]

Y0 = shrink * Y0[np.newaxis,:]

short = np.repeat(length < minsh, 8, axis=1)

# Now select X0, Y0 if short, otherwise X, Y

cbook._putmask(X, short, X0)

cbook._putmask(Y, short, Y0)

if self.pivot[:3] == 'mid':

X -= 0.5 * X[:,3, np.newaxis]

elif self.pivot[:3] == 'tip':

X = X - X[:,3, np.newaxis] #numpy bug? using -= does not

# work here unless we multiply

# by a float first, as with 'mid'.

tooshort = length < self.minlength

if tooshort.any():

# Use a heptagonal dot:

th = np.arange(0,8,1, np.float64) * (np.pi/3.0)

x1 = np.cos(th) * self.minlength * 0.5

y1 = np.sin(th) * self.minlength * 0.5

X1 = np.repeat(x1[np.newaxis, :], N, axis=0)

Y1 = np.repeat(y1[np.newaxis, :], N, axis=0)

tooshort = np.repeat(tooshort, 8, 1)

cbook._putmask(X, tooshort, X1)

cbook._putmask(Y, tooshort, Y1)

# Mask handling is deferred to the caller, _make_verts.

return X, Y

'''

#This may be an abuse of polygons here to render what is essentially maybe

#1 triangle and a series of lines. It works fine as far as I can tell

#however.

@docstring.interpd

def __init__(self, ax, *args, **kw):

"""

self._pivot = kw.pop('pivot', 'tip')

self._length = kw.pop('length', 7)

barbcolor = kw.pop('barbcolor', None)

flagcolor = kw.pop('flagcolor', None)

self.sizes = kw.pop('sizes', dict())

self.fill_empty = kw.pop('fill_empty', False)

self.barb_increments = kw.pop('barb_increments', dict())

self.rounding = kw.pop('rounding', True)

self.flip = kw.pop('flip_barb', False)

#Flagcolor and and barbcolor provide convenience parameters for setting

#the facecolor and edgecolor, respectively, of the barb polygon. We

#also work here to make the flag the same color as the rest of the barb

#by default

if None in (barbcolor, flagcolor):

kw['edgecolors'] = 'face'

if flagcolor:

kw['facecolors'] = flagcolor

elif barbcolor:

kw['facecolors'] = barbcolor

else:

#Set to facecolor passed in or default to black

kw.setdefault('facecolors', 'k')

else:

kw['edgecolors'] = barbcolor

kw['facecolors'] = flagcolor

#Parse out the data arrays from the various configurations supported

x, y, u, v, c = _parse_args(*args)

self.x = x

self.y = y

xy = np.hstack((x[:,np.newaxis], y[:,np.newaxis]))

#Make a collection

barb_size = self._length**2 / 4 #Empirically determined

collections.PolyCollection.__init__(self, [], (barb_size,), offsets=xy,

transOffset=ax.transData, **kw)

self.set_transform(transforms.IdentityTransform())

self.set_UVC(u, v, c)

def _find_tails(self, mag, rounding=True, half=5, full=10, flag=50):

'''

#If rounding, round to the nearest multiple of half, the smallest

#increment

if rounding:

mag = half * (mag / half + 0.5).astype(np.int)

num_flags = np.floor(mag / flag).astype(np.int)

mag = np.mod(mag, flag)

num_barb = np.floor(mag / full).astype(np.int)

mag = np.mod(mag, full)

half_flag = mag >= half

empty_flag = ~(half_flag | (num_flags > 0) | (num_barb > 0))

return num_flags, num_barb, half_flag, empty_flag

def _make_barbs(self, u, v, nflags, nbarbs, half_barb, empty_flag, length,

pivot, sizes, fill_empty, flip):

'''

#These control the spacing and size of barb elements relative to the

#length of the shaft

spacing = length * sizes.get('spacing', 0.125)

full_height = length * sizes.get('height', 0.4)

full_width = length * sizes.get('width', 0.25)

empty_rad = length * sizes.get('emptybarb', 0.15)

#Controls y point where to pivot the barb.

pivot_points = dict(tip=0.0, middle=-length/2.)

#Check for flip

if flip: full_height = -full_height

endx = 0.0

endy = pivot_points[pivot.lower()]

#Get the appropriate angle for the vector components. The offset is due

#to the way the barb is initially drawn, going down the y-axis. This

#makes sense in a meteorological mode of thinking since there 0 degrees

#corresponds to north (the y-axis traditionally)

angles = -(ma.arctan2(v, u) + np.pi/2)

#Used for low magnitude. We just get the vertices, so if we make it

#out here, it can be reused. The center set here should put the

#center of the circle at the location(offset), rather than at the

#same point as the barb pivot; this seems more sensible.

circ = CirclePolygon((0,0), radius=empty_rad).get_verts()

if fill_empty:

empty_barb = circ

else:

#If we don't want the empty one filled, we make a degenerate polygon

#that wraps back over itself

empty_barb = np.concatenate((circ, circ[::-1]))

barb_list = []

for index, angle in np.ndenumerate(angles):

#If the vector magnitude is too weak to draw anything, plot an

#empty circle instead

if empty_flag[index]:

#We can skip the transform since the circle has no preferred

#orientation

barb_list.append(empty_barb)

continue

poly_verts = [(endx, endy)]

offset = length

#Add vertices for each flag

for i in range(nflags[index]):

#The spacing that works for the barbs is a little to much for

#the flags, but this only occurs when we have more than 1 flag.

if offset != length: offset += spacing / 2.

poly_verts.extend([[endx, endy + offset],

[endx + full_height, endy - full_width/2 + offset],

[endx, endy - full_width + offset]])

offset -= full_width + spacing

#Add vertices for each barb. These really are lines, but works

#great adding 3 vertices that basically pull the polygon out and

#back down the line

for i in range(nbarbs[index]):

poly_verts.extend([(endx, endy + offset),

(endx + full_height, endy + offset + full_width/2),