"""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': 'middle', 'S': 'middle', 'E': 'tip', 'W': 'tail'}

def __init__(self, Q, X, Y, U, label,

*, angle=0, coordinates='axes', color=None, labelsep=0.1,

labelpos='N', labelcolor=None, fontproperties=None,

**kw):

"""

martist.Artist.__init__(self)

self.Q = Q

self.X = X

self.Y = Y

self.U = U

self.angle = angle

self.coord = coordinates

self.color = color

self.label = label

self._labelsep_inches = labelsep

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

# try to prevent closure over the real self

weak_self = weakref.ref(self)

def on_dpi_change(fig):

self_weakref = weak_self()

if self_weakref is not None:

self_weakref.labelsep = self_weakref._labelsep_inches * fig.dpi

# simple brute force update works because _init is called at

# the start of draw.

self_weakref._initialized = False

self._cid = Q.axes.figure.callbacks.connect(

'dpi_changed', on_dpi_change)

self.labelpos = labelpos

self.labelcolor = labelcolor

self.fontproperties = fontproperties or 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._from_any(_fp))

if self.labelcolor is not None:

self.text.set_color(self.labelcolor)

self._initialized = False

self.zorder = Q.zorder + 0.1

def remove(self):

# docstring inherited

self.Q.axes.figure.callbacks.disconnect(self._cid)

self._cid = None

super().remove() # pass the remove call up the stack

def _init(self):

if True: # not self._initialized:

if not self.Q._initialized:

self.Q._init()

self._set_transform()

with cbook._setattr_cm(self.Q, pivot=self.pivot[self.labelpos],

# Hack: save and restore the Umask

Umask=ma.nomask):

u = self.U * np.cos(np.radians(self.angle))

v = self.U * np.sin(np.radians(self.angle))

angle = (self.Q.angles if isinstance(self.Q.angles, str)

else 'uv')

self.verts = self.Q._make_verts(

np.array([u]), np.array([v]), angle)

kw = self.Q.polykw

kw.update(self.kw)

self.vector = mcollections.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.vector.set_figure(self.get_figure())

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

@martist.allow_rasterization

def draw(self, renderer):

self._init()

self.vector.draw(renderer)

x, y = self.get_transform().transform((self.X, self.Y))

self.text.set_x(self._text_x(x))

self.text.set_y(self._text_y(y))

self.text.draw(renderer)

self.stale = False

def _set_transform(self):

self.set_transform(cbook._check_getitem({

"data": self.Q.axes.transData,

"axes": self.Q.axes.transAxes,

"figure": self.Q.axes.figure.transFigure,

"inches": self.Q.axes.figure.dpi_scale_trans,

}, coordinates=self.coord))

def set_figure(self, fig):

martist.Artist.set_figure(self, fig)

self.text.set_figure(fig)

def contains(self, mouseevent):

inside, info = self._default_contains(mouseevent)

if inside is not None:

return inside, info

# 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, {}

@cbook.deprecated("3.2")

@property

def quiverkey_doc(self):

"""

X = Y = C = None

len_args = len(args)

if len_args == 2:

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

# keeping masked arrays

U, V = np.atleast_1d(*args)

elif len_args == 3:

U, V, C = np.atleast_1d(*args)

elif len_args == 4:

X, Y, U, V = np.atleast_1d(*args)

elif len_args == 5:

X, Y, U, V, C = np.atleast_1d(*args)

else:

raise TypeError(f'{caller_name} takes 2-5 positional arguments but '

f'{len_args} were given')

nr, nc = (1, U.shape[0]) if U.ndim == 1 else U.shape

if X is not None:

X = X.ravel()

Y = Y.ravel()

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

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

elif len(X) != len(Y):

raise ValueError('X and Y must be the same size, but '

f'X.size is {X.size} and Y.size is {Y.size}.')

else:

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

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

# Size validation for U, V, C is left to the set_UVC method.

"""

_PIVOT_VALS = ('tail', 'middle', 'tip')

@docstring.Substitution(_quiver_doc)

def __init__(self, ax, *args,

scale=None, headwidth=3, headlength=5, headaxislength=4.5,

minshaft=1, minlength=1, units='width', scale_units=None,

angles='uv', width=None, color='k', pivot='tail', **kw):

"""

self._axes = ax # The attr actually set by the Artist.axes property.

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

self.X = X

self.Y = Y

self.XY = np.column_stack((X, Y))

self.N = len(X)

self.scale = scale

self.headwidth = headwidth

self.headlength = float(headlength)

self.headaxislength = headaxislength

self.minshaft = minshaft

self.minlength = minlength

self.units = units

self.scale_units = scale_units

self.angles = angles

self.width = width

if pivot.lower() == 'mid':

pivot = 'middle'

self.pivot = pivot.lower()

cbook._check_in_list(self._PIVOT_VALS, pivot=self.pivot)

self.transform = kw.pop('transform', ax.transData)

kw.setdefault('facecolors', color)

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

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

transOffset=self.transform,

closed=False,

**kw)

self.polykw = kw

self.set_UVC(U, V, C)

self._initialized = False

weak_self = weakref.ref(self) # Prevent closure over the real self.

def on_dpi_change(fig):

self_weakref = weak_self()

if self_weakref is not None:

# vertices depend on width, span which in turn depend on dpi

self_weakref._new_UV = True

# simple brute force update works because _init is called at

# the start of draw.

self_weakref._initialized = False

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

@cbook.deprecated("3.3", alternative="axes")

def ax(self):

return self.axes

def remove(self):

# docstring inherited

self.axes.figure.callbacks.disconnect(self._cid)

self._cid = None

super().remove() # pass the remove call up the stack

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

self.span = trans.inverted().transform_bbox(self.axes.bbox).width

if self.width is None:

sn = np.clip(math.sqrt(self.N), 8, 25)

self.width = 0.06 * self.span / sn

# _make_verts sets self.scale if not already specified

if not self._initialized and self.scale is None:

self._make_verts(self.U, self.V, self.angles)

self._initialized = True

def get_datalim(self, transData):

trans = self.get_transform()

transOffset = self.get_offset_transform()

full_transform = (trans - transData) + (transOffset - transData)

XY = full_transform.transform(self.XY)

bbox = transforms.Bbox.null()

bbox.update_from_data_xy(XY, ignore=True)

return bbox

@martist.allow_rasterization

def draw(self, renderer):

self._init()

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

self.set_verts(verts, closed=False)

self._new_UV = False

mcollections.PolyCollection.draw(self, renderer)

self.stale = False

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

# We need to ensure we have a copy, not a reference

# to an array that might change before draw().

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

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

if C is not None:

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

for name, var in zip(('U', 'V', 'C'), (U, V, C)):

if not (var is None or var.size == self.N or var.size == 1):

raise ValueError(f'Argument {name} has a size {var.size}'

f' which does not match {self.N},'

' the number of arrow positions')

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

if C is not None:

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

self.stale = True

def _dots_per_unit(self, units):

"""

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

if units == 'x':

dx0 = self.axes.viewLim.width

dx1 = self.axes.bbox.width

elif units == 'y':

dx0 = self.axes.viewLim.height

dx1 = self.axes.bbox.height

else: # 'xy' is assumed

dxx0 = self.axes.viewLim.width

dxx1 = self.axes.bbox.width

dyy0 = self.axes.viewLim.height

dyy1 = self.axes.bbox.height

dx1 = np.hypot(dxx1, dyy1)

dx0 = np.hypot(dxx0, dyy0)

dx = dx1 / dx0

else:

if units == 'width':

dx = self.axes.bbox.width

elif units == 'height':

dx = self.axes.bbox.height

elif units == 'dots':

dx = 1.0

elif units == 'inches':

dx = self.axes.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.axes.transData.transform(self.XY)

uv = np.column_stack((U, V))

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

dxy = xyp - xy

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

lengths = np.hypot(*dxy.T) / eps

return angles, lengths

def _make_verts(self, U, V, angles):

uv = (U + V * 1j)

str_angles = angles if isinstance(angles, str) else ''

if str_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 str_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.

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

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

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

a = lengths

else:

a = np.abs(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()

# crude auto-scaling

# scale is typical arrow length as a multiple of the arrow width

scale = 1.8 * amean * sn / self.span

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 str_angles == 'xy':

theta = angles

elif str_angles == 'uv':

theta = np.angle(uv)

else:

theta = ma.masked_invalid(np.deg2rad(angles)).filled(0)

theta = theta.reshape((-1, 1)) # for broadcasting

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

XY = np.stack((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[:, ii]

Y = y[:, ii]

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

X0 = x0[ii]

Y0 = y0[ii]

Y0[3:-1] *= -1

shrink = length / minsh if minsh != 0. else 0.

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

np.copyto(X, X0, where=short)

np.copyto(Y, Y0, where=short)

if self.pivot == 'middle':

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

elif self.pivot == 'tip':

# numpy bug? using -= does not work here unless we multiply by a

# float first, as with 'mid'.

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

elif self.pivot != 'tail':

cbook._check_in_list(["middle", "tip", "tail"], pivot=self.pivot)

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)

np.copyto(X, X1, where=tooshort)

np.copyto(Y, Y1, where=tooshort)

# 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,

pivot='tip', length=7, barbcolor=None, flagcolor=None,

sizes=None, fill_empty=False, barb_increments=None,

rounding=True, flip_barb=False, **kw):

"""

self.sizes = sizes or dict()

self.fill_empty = fill_empty

self.barb_increments = barb_increments or dict()

self.rounding = rounding

self.flip = np.atleast_1d(flip_barb)

transform = kw.pop('transform', ax.transData)

self._pivot = pivot

self._length = length

barbcolor = barbcolor

flagcolor = flagcolor

# Flagcolor 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

# Explicitly set a line width if we're not given one, otherwise

# polygons are not outlined and we get no barbs

if 'linewidth' not in kw and 'lw' not in kw:

kw['linewidth'] = 1

# Parse out the data arrays from the various configurations supported

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

self.x = x

self.y = y

xy = np.column_stack((x, y))

# Make a collection

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

mcollections.PolyCollection.__init__(self, [], (barb_size,),

offsets=xy,

transOffset=transform, **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(int)