Coverage for /home/martinb/.local/share/virtualenvs/camcops/lib/python3.6/site-packages/matplotlib/quiver.py : 12%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2Support for plotting vector fields.
4Presently this contains Quiver and Barb. Quiver plots an arrow in the
5direction of the vector, with the size of the arrow related to the
6magnitude of the vector.
8Barbs are like quiver in that they point along a vector, but
9the magnitude of the vector is given schematically by the presence of barbs
10or flags on the barb.
12This will also become a home for things such as standard
13deviation ellipses, which can and will be derived very easily from
14the Quiver code.
15"""
17import math
18import weakref
20import numpy as np
21from numpy import ma
23from matplotlib import cbook, docstring, font_manager
24import matplotlib.artist as martist
25import matplotlib.collections as mcollections
26from matplotlib.patches import CirclePolygon
27import matplotlib.text as mtext
28import matplotlib.transforms as transforms
31_quiver_doc = """
32Plot a 2D field of arrows.
34Call signature::
36 quiver([X, Y], U, V, [C], **kw)
38Where *X*, *Y* define the arrow locations, *U*, *V* define the arrow
39directions, and *C* optionally sets the color.
41**Arrow size**
43The default settings auto-scales the length of the arrows to a reasonable size.
44To change this behavior see the *scale* and *scale_units* parameters.
46**Arrow shape**
48The defaults give a slightly swept-back arrow; to make the head a
49triangle, make *headaxislength* the same as *headlength*. To make the
50arrow more pointed, reduce *headwidth* or increase *headlength* and
51*headaxislength*. To make the head smaller relative to the shaft,
52scale down all the head parameters. You will probably do best to leave
53minshaft alone.
55**Arrow outline**
57*linewidths* and *edgecolors* can be used to customize the arrow
58outlines.
60Parameters
61----------
62X, Y : 1D or 2D array-like, optional
63 The x and y coordinates of the arrow locations.
65 If not given, they will be generated as a uniform integer meshgrid based
66 on the dimensions of *U* and *V*.
68 If *X* and *Y* are 1D but *U*, *V* are 2D, *X*, *Y* are expanded to 2D
69 using ``X, Y = np.meshgrid(X, Y)``. In this case ``len(X)`` and ``len(Y)``
70 must match the column and row dimensions of *U* and *V*.
72U, V : 1D or 2D array-like
73 The x and y direction components of the arrow vectors.
75 They must have the same number of elements, matching the number of arrow
76 locations. *U* and *V* may be masked. Only locations unmasked in
77 *U*, *V*, and *C* will be drawn.
79C : 1D or 2D array-like, optional
80 Numeric data that defines the arrow colors by colormapping via *norm* and
81 *cmap*.
83 This does not support explicit colors. If you want to set colors directly,
84 use *color* instead. The size of *C* must match the number of arrow
85 locations.
87units : {'width', 'height', 'dots', 'inches', 'x', 'y' 'xy'}, default: 'width'
88 The arrow dimensions (except for *length*) are measured in multiples of
89 this unit.
91 The following values are supported:
93 - 'width', 'height': The width or height of the axis.
94 - 'dots', 'inches': Pixels or inches based on the figure dpi.
95 - 'x', 'y', 'xy': *X*, *Y* or :math:`\\sqrt{X^2 + Y^2}` in data units.
97 The arrows scale differently depending on the units. For
98 'x' or 'y', the arrows get larger as one zooms in; for other
99 units, the arrow size is independent of the zoom state. For
100 'width or 'height', the arrow size increases with the width and
101 height of the axes, respectively, when the window is resized;
102 for 'dots' or 'inches', resizing does not change the arrows.
104angles : {'uv', 'xy'} or array-like, optional, default: 'uv'
105 Method for determining the angle of the arrows.
107 - 'uv': The arrow axis aspect ratio is 1 so that
108 if *U* == *V* the orientation of the arrow on the plot is 45 degrees
109 counter-clockwise from the horizontal axis (positive to the right).
111 Use this if the arrows symbolize a quantity that is not based on
112 *X*, *Y* data coordinates.
114 - 'xy': Arrows point from (x, y) to (x+u, y+v).
115 Use this for plotting a gradient field, for example.
117 - Alternatively, arbitrary angles may be specified explicitly as an array
118 of values in degrees, counter-clockwise from the horizontal axis.
120 In this case *U*, *V* is only used to determine the length of the
121 arrows.
123 Note: inverting a data axis will correspondingly invert the
124 arrows only with ``angles='xy'``.
126scale : float, optional
127 Number of data units per arrow length unit, e.g., m/s per plot width; a
128 smaller scale parameter makes the arrow longer. Default is *None*.
130 If *None*, a simple autoscaling algorithm is used, based on the average
131 vector length and the number of vectors. The arrow length unit is given by
132 the *scale_units* parameter.
134scale_units : {'width', 'height', 'dots', 'inches', 'x', 'y', 'xy'}, optional
135 If the *scale* kwarg is *None*, the arrow length unit. Default is *None*.
137 e.g. *scale_units* is 'inches', *scale* is 2.0, and ``(u, v) = (1, 0)``,
138 then the vector will be 0.5 inches long.
140 If *scale_units* is 'width' or 'height', then the vector will be half the
141 width/height of the axes.
143 If *scale_units* is 'x' then the vector will be 0.5 x-axis
144 units. To plot vectors in the x-y plane, with u and v having
145 the same units as x and y, use
146 ``angles='xy', scale_units='xy', scale=1``.
148width : float, optional
149 Shaft width in arrow units; default depends on choice of units,
150 above, and number of vectors; a typical starting value is about
151 0.005 times the width of the plot.
153headwidth : float, optional, default: 3
154 Head width as multiple of shaft width.
156headlength : float, optional, default: 5
157 Head length as multiple of shaft width.
159headaxislength : float, optional, default: 4.5
160 Head length at shaft intersection.
162minshaft : float, optional, default: 1
163 Length below which arrow scales, in units of head length. Do not
164 set this to less than 1, or small arrows will look terrible!
166minlength : float, optional, default: 1
167 Minimum length as a multiple of shaft width; if an arrow length
168 is less than this, plot a dot (hexagon) of this diameter instead.
170pivot : {'tail', 'mid', 'middle', 'tip'}, optional, default: 'tail'
171 The part of the arrow that is anchored to the *X*, *Y* grid. The arrow
172 rotates about this point.
174 'mid' is a synonym for 'middle'.
176color : color or color sequence, optional
177 Explicit color(s) for the arrows. If *C* has been set, *color* has no
178 effect.
180 This is a synonym for the `~.PolyCollection` *facecolor* parameter.
182Other Parameters
183----------------
184**kwargs : `~matplotlib.collections.PolyCollection` properties, optional
185 All other keyword arguments are passed on to `.PolyCollection`:
187 %(PolyCollection)s
189See Also
190--------
191quiverkey : Add a key to a quiver plot.
192""" % docstring.interpd.params
195class QuiverKey(martist.Artist):
196 """Labelled arrow for use as a quiver plot scale key."""
197 halign = {'N': 'center', 'S': 'center', 'E': 'left', 'W': 'right'}
198 valign = {'N': 'bottom', 'S': 'top', 'E': 'center', 'W': 'center'}
199 pivot = {'N': 'middle', 'S': 'middle', 'E': 'tip', 'W': 'tail'}
201 def __init__(self, Q, X, Y, U, label,
202 *, angle=0, coordinates='axes', color=None, labelsep=0.1,
203 labelpos='N', labelcolor=None, fontproperties=None,
204 **kw):
205 """
206 Add a key to a quiver plot.
208 The positioning of the key depends on *X*, *Y*, *coordinates*, and
209 *labelpos*. If *labelpos* is 'N' or 'S', *X*, *Y* give the position of
210 the middle of the key arrow. If *labelpos* is 'E', *X*, *Y* positions
211 the head, and if *labelpos* is 'W', *X*, *Y* positions the tail; in
212 either of these two cases, *X*, *Y* is somewhere in the middle of the
213 arrow+label key object.
215 Parameters
216 ----------
217 Q : `matplotlib.quiver.Quiver`
218 A `.Quiver` object as returned by a call to `~.Axes.quiver()`.
219 X, Y : float
220 The location of the key.
221 U : float
222 The length of the key.
223 label : str
224 The key label (e.g., length and units of the key).
225 angle : float, default: 0
226 The angle of the key arrow, in degrees anti-clockwise from the
227 x-axis.
228 coordinates : {'axes', 'figure', 'data', 'inches'}, default: 'axes'
229 Coordinate system and units for *X*, *Y*: 'axes' and 'figure' are
230 normalized coordinate systems with (0, 0) in the lower left and
231 (1, 1) in the upper right; 'data' are the axes data coordinates
232 (used for the locations of the vectors in the quiver plot itself);
233 'inches' is position in the figure in inches, with (0, 0) at the
234 lower left corner.
235 color : color
236 Overrides face and edge colors from *Q*.
237 labelpos : {'N', 'S', 'E', 'W'}
238 Position the label above, below, to the right, to the left of the
239 arrow, respectively.
240 labelsep : float, default: 0.1
241 Distance in inches between the arrow and the label.
242 labelcolor : color, default: :rc:`text.color`
243 Label color.
244 fontproperties : dict, optional
245 A dictionary with keyword arguments accepted by the
246 `~matplotlib.font_manager.FontProperties` initializer:
247 *family*, *style*, *variant*, *size*, *weight*.
248 **kwargs
249 Any additional keyword arguments are used to override vector
250 properties taken from *Q*.
251 """
252 martist.Artist.__init__(self)
253 self.Q = Q
254 self.X = X
255 self.Y = Y
256 self.U = U
257 self.angle = angle
258 self.coord = coordinates
259 self.color = color
260 self.label = label
261 self._labelsep_inches = labelsep
262 self.labelsep = (self._labelsep_inches * Q.ax.figure.dpi)
264 # try to prevent closure over the real self
265 weak_self = weakref.ref(self)
267 def on_dpi_change(fig):
268 self_weakref = weak_self()
269 if self_weakref is not None:
270 self_weakref.labelsep = self_weakref._labelsep_inches * fig.dpi
271 # simple brute force update works because _init is called at
272 # the start of draw.
273 self_weakref._initialized = False
275 self._cid = Q.ax.figure.callbacks.connect('dpi_changed',
276 on_dpi_change)
278 self.labelpos = labelpos
279 self.labelcolor = labelcolor
280 self.fontproperties = fontproperties or dict()
281 self.kw = kw
282 _fp = self.fontproperties
283 # boxprops = dict(facecolor='red')
284 self.text = mtext.Text(
285 text=label, # bbox=boxprops,
286 horizontalalignment=self.halign[self.labelpos],
287 verticalalignment=self.valign[self.labelpos],
288 fontproperties=font_manager.FontProperties(**_fp))
290 if self.labelcolor is not None:
291 self.text.set_color(self.labelcolor)
292 self._initialized = False
293 self.zorder = Q.zorder + 0.1
295 def remove(self):
296 """
297 Overload the remove method
298 """
299 self.Q.ax.figure.callbacks.disconnect(self._cid)
300 self._cid = None
301 # pass the remove call up the stack
302 martist.Artist.remove(self)
304 def _init(self):
305 if True: # not self._initialized:
306 if not self.Q._initialized:
307 self.Q._init()
308 self._set_transform()
309 _pivot = self.Q.pivot
310 self.Q.pivot = self.pivot[self.labelpos]
311 # Hack: save and restore the Umask
312 _mask = self.Q.Umask
313 self.Q.Umask = ma.nomask
314 u = self.U * np.cos(np.radians(self.angle))
315 v = self.U * np.sin(np.radians(self.angle))
316 angle = self.Q.angles if isinstance(self.Q.angles, str) else 'uv'
317 self.verts = self.Q._make_verts(
318 np.array([u]), np.array([v]), angle)
319 self.Q.Umask = _mask
320 self.Q.pivot = _pivot
321 kw = self.Q.polykw
322 kw.update(self.kw)
323 self.vector = mcollections.PolyCollection(
324 self.verts,
325 offsets=[(self.X, self.Y)],
326 transOffset=self.get_transform(),
327 **kw)
328 if self.color is not None:
329 self.vector.set_color(self.color)
330 self.vector.set_transform(self.Q.get_transform())
331 self.vector.set_figure(self.get_figure())
332 self._initialized = True
334 def _text_x(self, x):
335 if self.labelpos == 'E':
336 return x + self.labelsep
337 elif self.labelpos == 'W':
338 return x - self.labelsep
339 else:
340 return x
342 def _text_y(self, y):
343 if self.labelpos == 'N':
344 return y + self.labelsep
345 elif self.labelpos == 'S':
346 return y - self.labelsep
347 else:
348 return y
350 @martist.allow_rasterization
351 def draw(self, renderer):
352 self._init()
353 self.vector.draw(renderer)
354 x, y = self.get_transform().transform((self.X, self.Y))
355 self.text.set_x(self._text_x(x))
356 self.text.set_y(self._text_y(y))
357 self.text.draw(renderer)
358 self.stale = False
360 def _set_transform(self):
361 if self.coord == 'data':
362 self.set_transform(self.Q.ax.transData)
363 elif self.coord == 'axes':
364 self.set_transform(self.Q.ax.transAxes)
365 elif self.coord == 'figure':
366 self.set_transform(self.Q.ax.figure.transFigure)
367 elif self.coord == 'inches':
368 self.set_transform(self.Q.ax.figure.dpi_scale_trans)
369 else:
370 raise ValueError('unrecognized coordinates')
372 def set_figure(self, fig):
373 martist.Artist.set_figure(self, fig)
374 self.text.set_figure(fig)
376 def contains(self, mouseevent):
377 inside, info = self._default_contains(mouseevent)
378 if inside is not None:
379 return inside, info
380 # Maybe the dictionary should allow one to
381 # distinguish between a text hit and a vector hit.
382 if (self.text.contains(mouseevent)[0] or
383 self.vector.contains(mouseevent)[0]):
384 return True, {}
385 return False, {}
387 @cbook.deprecated("3.2")
388 @property
389 def quiverkey_doc(self):
390 return self.__init__.__doc__
393def _parse_args(*args, caller_name='function'):
394 """
395 Helper function to parse positional parameters for colored vector plots.
397 This is currently used for Quiver and Barbs.
399 Parameters
400 ----------
401 *args : list
402 list of 2-5 arguments. Depending on their number they are parsed to::
404 U, V
405 U, V, C
406 X, Y, U, V
407 X, Y, U, V, C
409 caller_name : str
410 Name of the calling method (used in error messages).
411 """
412 X = Y = C = None
414 len_args = len(args)
415 if len_args == 2:
416 # The use of atleast_1d allows for handling scalar arguments while also
417 # keeping masked arrays
418 U, V = np.atleast_1d(*args)
419 elif len_args == 3:
420 U, V, C = np.atleast_1d(*args)
421 elif len_args == 4:
422 X, Y, U, V = np.atleast_1d(*args)
423 elif len_args == 5:
424 X, Y, U, V, C = np.atleast_1d(*args)
425 else:
426 raise TypeError(f'{caller_name} takes 2-5 positional arguments but '
427 f'{len_args} were given')
429 nr, nc = (1, U.shape[0]) if U.ndim == 1 else U.shape
431 if X is not None:
432 X = X.ravel()
433 Y = Y.ravel()
434 if len(X) == nc and len(Y) == nr:
435 X, Y = [a.ravel() for a in np.meshgrid(X, Y)]
436 elif len(X) != len(Y):
437 raise ValueError('X and Y must be the same size, but '
438 f'X.size is {X.size} and Y.size is {Y.size}.')
439 else:
440 indexgrid = np.meshgrid(np.arange(nc), np.arange(nr))
441 X, Y = [np.ravel(a) for a in indexgrid]
442 # Size validation for U, V, C is left to the set_UVC method.
443 return X, Y, U, V, C
446def _check_consistent_shapes(*arrays):
447 all_shapes = {a.shape for a in arrays}
448 if len(all_shapes) != 1:
449 raise ValueError('The shapes of the passed in arrays do not match')
452class Quiver(mcollections.PolyCollection):
453 """
454 Specialized PolyCollection for arrows.
456 The only API method is set_UVC(), which can be used
457 to change the size, orientation, and color of the
458 arrows; their locations are fixed when the class is
459 instantiated. Possibly this method will be useful
460 in animations.
462 Much of the work in this class is done in the draw()
463 method so that as much information as possible is available
464 about the plot. In subsequent draw() calls, recalculation
465 is limited to things that might have changed, so there
466 should be no performance penalty from putting the calculations
467 in the draw() method.
468 """
470 _PIVOT_VALS = ('tail', 'middle', 'tip')
472 @docstring.Substitution(_quiver_doc)
473 def __init__(self, ax, *args,
474 scale=None, headwidth=3, headlength=5, headaxislength=4.5,
475 minshaft=1, minlength=1, units='width', scale_units=None,
476 angles='uv', width=None, color='k', pivot='tail', **kw):
477 """
478 The constructor takes one required argument, an Axes
479 instance, followed by the args and kwargs described
480 by the following pyplot interface documentation:
481 %s
482 """
483 self.ax = ax
484 X, Y, U, V, C = _parse_args(*args, caller_name='quiver()')
485 self.X = X
486 self.Y = Y
487 self.XY = np.column_stack((X, Y))
488 self.N = len(X)
489 self.scale = scale
490 self.headwidth = headwidth
491 self.headlength = float(headlength)
492 self.headaxislength = headaxislength
493 self.minshaft = minshaft
494 self.minlength = minlength
495 self.units = units
496 self.scale_units = scale_units
497 self.angles = angles
498 self.width = width
500 if pivot.lower() == 'mid':
501 pivot = 'middle'
502 self.pivot = pivot.lower()
503 cbook._check_in_list(self._PIVOT_VALS, pivot=self.pivot)
505 self.transform = kw.pop('transform', ax.transData)
506 kw.setdefault('facecolors', color)
507 kw.setdefault('linewidths', (0,))
508 mcollections.PolyCollection.__init__(self, [], offsets=self.XY,
509 transOffset=self.transform,
510 closed=False,
511 **kw)
512 self.polykw = kw
513 self.set_UVC(U, V, C)
514 self._initialized = False
516 # try to prevent closure over the real self
517 weak_self = weakref.ref(self)
519 def on_dpi_change(fig):
520 self_weakref = weak_self()
521 if self_weakref is not None:
522 # vertices depend on width, span which in turn depend on dpi
523 self_weakref._new_UV = True
524 # simple brute force update works because _init is called at
525 # the start of draw.
526 self_weakref._initialized = False
528 self._cid = self.ax.figure.callbacks.connect('dpi_changed',
529 on_dpi_change)
531 @cbook.deprecated("3.1", alternative="get_facecolor()")
532 @property
533 def color(self):
534 return self.get_facecolor()
536 @cbook.deprecated("3.1")
537 @property
538 def keyvec(self):
539 return None
541 @cbook.deprecated("3.1")
542 @property
543 def keytext(self):
544 return None
546 def remove(self):
547 """
548 Overload the remove method
549 """
550 # disconnect the call back
551 self.ax.figure.callbacks.disconnect(self._cid)
552 self._cid = None
553 # pass the remove call up the stack
554 mcollections.PolyCollection.remove(self)
556 def _init(self):
557 """
558 Initialization delayed until first draw;
559 allow time for axes setup.
560 """
561 # It seems that there are not enough event notifications
562 # available to have this work on an as-needed basis at present.
563 if True: # not self._initialized:
564 trans = self._set_transform()
565 ax = self.ax
566 self.span = trans.inverted().transform_bbox(ax.bbox).width
567 if self.width is None:
568 sn = np.clip(math.sqrt(self.N), 8, 25)
569 self.width = 0.06 * self.span / sn
571 # _make_verts sets self.scale if not already specified
572 if not self._initialized and self.scale is None:
573 self._make_verts(self.U, self.V, self.angles)
575 self._initialized = True
577 def get_datalim(self, transData):
578 trans = self.get_transform()
579 transOffset = self.get_offset_transform()
580 full_transform = (trans - transData) + (transOffset - transData)
581 XY = full_transform.transform(self.XY)
582 bbox = transforms.Bbox.null()
583 bbox.update_from_data_xy(XY, ignore=True)
584 return bbox
586 @martist.allow_rasterization
587 def draw(self, renderer):
588 self._init()
589 verts = self._make_verts(self.U, self.V, self.angles)
590 self.set_verts(verts, closed=False)
591 self._new_UV = False
592 mcollections.PolyCollection.draw(self, renderer)
593 self.stale = False
595 def set_UVC(self, U, V, C=None):
596 # We need to ensure we have a copy, not a reference
597 # to an array that might change before draw().
598 U = ma.masked_invalid(U, copy=True).ravel()
599 V = ma.masked_invalid(V, copy=True).ravel()
600 if C is not None:
601 C = ma.masked_invalid(C, copy=True).ravel()
602 for name, var in zip(('U', 'V', 'C'), (U, V, C)):
603 if not (var is None or var.size == self.N or var.size == 1):
604 raise ValueError(f'Argument {name} has a size {var.size}'
605 f' which does not match {self.N},'
606 ' the number of arrow positions')
608 mask = ma.mask_or(U.mask, V.mask, copy=False, shrink=True)
609 if C is not None:
610 mask = ma.mask_or(mask, C.mask, copy=False, shrink=True)
611 if mask is ma.nomask:
612 C = C.filled()
613 else:
614 C = ma.array(C, mask=mask, copy=False)
615 self.U = U.filled(1)
616 self.V = V.filled(1)
617 self.Umask = mask
618 if C is not None:
619 self.set_array(C)
620 self._new_UV = True
621 self.stale = True
623 def _dots_per_unit(self, units):
624 """
625 Return a scale factor for converting from units to pixels
626 """
627 ax = self.ax
628 if units in ('x', 'y', 'xy'):
629 if units == 'x':
630 dx0 = ax.viewLim.width
631 dx1 = ax.bbox.width
632 elif units == 'y':
633 dx0 = ax.viewLim.height
634 dx1 = ax.bbox.height
635 else: # 'xy' is assumed
636 dxx0 = ax.viewLim.width
637 dxx1 = ax.bbox.width
638 dyy0 = ax.viewLim.height
639 dyy1 = ax.bbox.height
640 dx1 = np.hypot(dxx1, dyy1)
641 dx0 = np.hypot(dxx0, dyy0)
642 dx = dx1 / dx0
643 else:
644 if units == 'width':
645 dx = ax.bbox.width
646 elif units == 'height':
647 dx = ax.bbox.height
648 elif units == 'dots':
649 dx = 1.0
650 elif units == 'inches':
651 dx = ax.figure.dpi
652 else:
653 raise ValueError('unrecognized units')
654 return dx
656 def _set_transform(self):
657 """
658 Sets the PolygonCollection transform to go
659 from arrow width units to pixels.
660 """
661 dx = self._dots_per_unit(self.units)
662 self._trans_scale = dx # pixels per arrow width unit
663 trans = transforms.Affine2D().scale(dx)
664 self.set_transform(trans)
665 return trans
667 def _angles_lengths(self, U, V, eps=1):
668 xy = self.ax.transData.transform(self.XY)
669 uv = np.column_stack((U, V))
670 xyp = self.ax.transData.transform(self.XY + eps * uv)
671 dxy = xyp - xy
672 angles = np.arctan2(dxy[:, 1], dxy[:, 0])
673 lengths = np.hypot(*dxy.T) / eps
674 return angles, lengths
676 def _make_verts(self, U, V, angles):
677 uv = (U + V * 1j)
678 str_angles = angles if isinstance(angles, str) else ''
679 if str_angles == 'xy' and self.scale_units == 'xy':
680 # Here eps is 1 so that if we get U, V by diffing
681 # the X, Y arrays, the vectors will connect the
682 # points, regardless of the axis scaling (including log).
683 angles, lengths = self._angles_lengths(U, V, eps=1)
684 elif str_angles == 'xy' or self.scale_units == 'xy':
685 # Calculate eps based on the extents of the plot
686 # so that we don't end up with roundoff error from
687 # adding a small number to a large.
688 eps = np.abs(self.ax.dataLim.extents).max() * 0.001
689 angles, lengths = self._angles_lengths(U, V, eps=eps)
690 if str_angles and self.scale_units == 'xy':
691 a = lengths
692 else:
693 a = np.abs(uv)
694 if self.scale is None:
695 sn = max(10, math.sqrt(self.N))
696 if self.Umask is not ma.nomask:
697 amean = a[~self.Umask].mean()
698 else:
699 amean = a.mean()
700 # crude auto-scaling
701 # scale is typical arrow length as a multiple of the arrow width
702 scale = 1.8 * amean * sn / self.span
703 if self.scale_units is None:
704 if self.scale is None:
705 self.scale = scale
706 widthu_per_lenu = 1.0
707 else:
708 if self.scale_units == 'xy':
709 dx = 1
710 else:
711 dx = self._dots_per_unit(self.scale_units)
712 widthu_per_lenu = dx / self._trans_scale
713 if self.scale is None:
714 self.scale = scale * widthu_per_lenu
715 length = a * (widthu_per_lenu / (self.scale * self.width))
716 X, Y = self._h_arrows(length)
717 if str_angles == 'xy':
718 theta = angles
719 elif str_angles == 'uv':
720 theta = np.angle(uv)
721 else:
722 theta = ma.masked_invalid(np.deg2rad(angles)).filled(0)
723 theta = theta.reshape((-1, 1)) # for broadcasting
724 xy = (X + Y * 1j) * np.exp(1j * theta) * self.width
725 XY = np.stack((xy.real, xy.imag), axis=2)
726 if self.Umask is not ma.nomask:
727 XY = ma.array(XY)
728 XY[self.Umask] = ma.masked
729 # This might be handled more efficiently with nans, given
730 # that nans will end up in the paths anyway.
732 return XY
734 def _h_arrows(self, length):
735 """Length is in arrow width units."""
736 # It might be possible to streamline the code
737 # and speed it up a bit by using complex (x, y)
738 # instead of separate arrays; but any gain would be slight.
739 minsh = self.minshaft * self.headlength
740 N = len(length)
741 length = length.reshape(N, 1)
742 # This number is chosen based on when pixel values overflow in Agg
743 # causing rendering errors
744 # length = np.minimum(length, 2 ** 16)
745 np.clip(length, 0, 2 ** 16, out=length)
746 # x, y: normal horizontal arrow
747 x = np.array([0, -self.headaxislength,
748 -self.headlength, 0],
749 np.float64)
750 x = x + np.array([0, 1, 1, 1]) * length
751 y = 0.5 * np.array([1, 1, self.headwidth, 0], np.float64)
752 y = np.repeat(y[np.newaxis, :], N, axis=0)
753 # x0, y0: arrow without shaft, for short vectors
754 x0 = np.array([0, minsh - self.headaxislength,
755 minsh - self.headlength, minsh], np.float64)
756 y0 = 0.5 * np.array([1, 1, self.headwidth, 0], np.float64)
757 ii = [0, 1, 2, 3, 2, 1, 0, 0]
758 X = x[:, ii]
759 Y = y[:, ii]
760 Y[:, 3:-1] *= -1
761 X0 = x0[ii]
762 Y0 = y0[ii]
763 Y0[3:-1] *= -1
764 shrink = length / minsh if minsh != 0. else 0.
765 X0 = shrink * X0[np.newaxis, :]
766 Y0 = shrink * Y0[np.newaxis, :]
767 short = np.repeat(length < minsh, 8, axis=1)
768 # Now select X0, Y0 if short, otherwise X, Y
769 np.copyto(X, X0, where=short)
770 np.copyto(Y, Y0, where=short)
771 if self.pivot == 'middle':
772 X -= 0.5 * X[:, 3, np.newaxis]
773 elif self.pivot == 'tip':
774 # numpy bug? using -= does not work here unless we multiply by a
775 # float first, as with 'mid'.
776 X = X - X[:, 3, np.newaxis]
777 elif self.pivot != 'tail':
778 cbook._check_in_list(["middle", "tip", "tail"], pivot=self.pivot)
780 tooshort = length < self.minlength
781 if tooshort.any():
782 # Use a heptagonal dot:
783 th = np.arange(0, 8, 1, np.float64) * (np.pi / 3.0)
784 x1 = np.cos(th) * self.minlength * 0.5
785 y1 = np.sin(th) * self.minlength * 0.5
786 X1 = np.repeat(x1[np.newaxis, :], N, axis=0)
787 Y1 = np.repeat(y1[np.newaxis, :], N, axis=0)
788 tooshort = np.repeat(tooshort, 8, 1)
789 np.copyto(X, X1, where=tooshort)
790 np.copyto(Y, Y1, where=tooshort)
791 # Mask handling is deferred to the caller, _make_verts.
792 return X, Y
794 quiver_doc = _quiver_doc
797_barbs_doc = r"""
798Plot a 2D field of barbs.
800Call signature::
802 barbs([X, Y], U, V, [C], **kw)
804Where *X*, *Y* define the barb locations, *U*, *V* define the barb
805directions, and *C* optionally sets the color.
807All arguments may be 1D or 2D. *U*, *V*, *C* may be masked arrays, but masked
808*X*, *Y* are not supported at present.
810Barbs are traditionally used in meteorology as a way to plot the speed
811and direction of wind observations, but can technically be used to
812plot any two dimensional vector quantity. As opposed to arrows, which
813give vector magnitude by the length of the arrow, the barbs give more
814quantitative information about the vector magnitude by putting slanted
815lines or a triangle for various increments in magnitude, as show
816schematically below::
818 : /\ \
819 : / \ \
820 : / \ \ \
821 : / \ \ \
822 : ------------------------------
825The largest increment is given by a triangle (or "flag"). After those
826come full lines (barbs). The smallest increment is a half line. There
827is only, of course, ever at most 1 half line. If the magnitude is
828small and only needs a single half-line and no full lines or
829triangles, the half-line is offset from the end of the barb so that it
830can be easily distinguished from barbs with a single full line. The
831magnitude for the barb shown above would nominally be 65, using the
832standard increments of 50, 10, and 5.
834See also https://en.wikipedia.org/wiki/Wind_barb.
838Parameters
839----------
840X, Y : 1D or 2D array-like, optional
841 The x and y coordinates of the barb locations. See *pivot* for how the
842 barbs are drawn to the x, y positions.
844 If not given, they will be generated as a uniform integer meshgrid based
845 on the dimensions of *U* and *V*.
847 If *X* and *Y* are 1D but *U*, *V* are 2D, *X*, *Y* are expanded to 2D
848 using ``X, Y = np.meshgrid(X, Y)``. In this case ``len(X)`` and ``len(Y)``
849 must match the column and row dimensions of *U* and *V*.
851U, V : 1D or 2D array-like
852 The x and y components of the barb shaft.
854C : 1D or 2D array-like, optional
855 Numeric data that defines the barb colors by colormapping via *norm* and
856 *cmap*.
858 This does not support explicit colors. If you want to set colors directly,
859 use *barbcolor* instead.
861length : float, default: 7
862 Length of the barb in points; the other parts of the barb
863 are scaled against this.
865pivot : {'tip', 'middle'} or float, default: 'tip'
866 The part of the arrow that is anchored to the *X*, *Y* grid. The barb
867 rotates about this point. This can also be a number, which shifts the
868 start of the barb that many points away from grid point.
870barbcolor : color or color sequence
871 Specifies the color of all parts of the barb except for the flags. This
872 parameter is analogous to the *edgecolor* parameter for polygons,
873 which can be used instead. However this parameter will override
874 facecolor.
876flagcolor : color or color sequence
877 Specifies the color of any flags on the barb. This parameter is
878 analogous to the *facecolor* parameter for polygons, which can be
879 used instead. However, this parameter will override facecolor. If
880 this is not set (and *C* has not either) then *flagcolor* will be
881 set to match *barbcolor* so that the barb has a uniform color. If
882 *C* has been set, *flagcolor* has no effect.
884sizes : dict, optional
885 A dictionary of coefficients specifying the ratio of a given
886 feature to the length of the barb. Only those values one wishes to
887 override need to be included. These features include:
889 - 'spacing' - space between features (flags, full/half barbs)
890 - 'height' - height (distance from shaft to top) of a flag or full barb
891 - 'width' - width of a flag, twice the width of a full barb
892 - 'emptybarb' - radius of the circle used for low magnitudes
894fill_empty : bool, default: False
895 Whether the empty barbs (circles) that are drawn should be filled with
896 the flag color. If they are not filled, the center is transparent.
898rounding : bool, default: True
899 Whether the vector magnitude should be rounded when allocating barb
900 components. If True, the magnitude is rounded to the nearest multiple
901 of the half-barb increment. If False, the magnitude is simply truncated
902 to the next lowest multiple.
904barb_increments : dict, optional
905 A dictionary of increments specifying values to associate with
906 different parts of the barb. Only those values one wishes to
907 override need to be included.
909 - 'half' - half barbs (Default is 5)
910 - 'full' - full barbs (Default is 10)
911 - 'flag' - flags (default is 50)
913flip_barb : bool or array-like of bool, default: False
914 Whether the lines and flags should point opposite to normal.
915 Normal behavior is for the barbs and lines to point right (comes from wind
916 barbs having these features point towards low pressure in the Northern
917 Hemisphere).
919 A single value is applied to all barbs. Individual barbs can be flipped by
920 passing a bool array of the same size as *U* and *V*.
922Returns
923-------
924barbs : `~matplotlib.quiver.Barbs`
926Other Parameters
927----------------
928**kwargs
929 The barbs can further be customized using `.PolyCollection` keyword
930 arguments:
932 %(PolyCollection)s
933""" % docstring.interpd.params
935docstring.interpd.update(barbs_doc=_barbs_doc)
938class Barbs(mcollections.PolyCollection):
939 '''
940 Specialized PolyCollection for barbs.
942 The only API method is :meth:`set_UVC`, which can be used to
943 change the size, orientation, and color of the arrows. Locations
944 are changed using the :meth:`set_offsets` collection method.
945 Possibly this method will be useful in animations.
947 There is one internal function :meth:`_find_tails` which finds
948 exactly what should be put on the barb given the vector magnitude.
949 From there :meth:`_make_barbs` is used to find the vertices of the
950 polygon to represent the barb based on this information.
951 '''
952 # This may be an abuse of polygons here to render what is essentially maybe
953 # 1 triangle and a series of lines. It works fine as far as I can tell
954 # however.
955 @docstring.interpd
956 def __init__(self, ax, *args,
957 pivot='tip', length=7, barbcolor=None, flagcolor=None,
958 sizes=None, fill_empty=False, barb_increments=None,
959 rounding=True, flip_barb=False, **kw):
960 """
961 The constructor takes one required argument, an Axes
962 instance, followed by the args and kwargs described
963 by the following pyplot interface documentation:
964 %(barbs_doc)s
965 """
966 self.sizes = sizes or dict()
967 self.fill_empty = fill_empty
968 self.barb_increments = barb_increments or dict()
969 self.rounding = rounding
970 self.flip = np.atleast_1d(flip_barb)
971 transform = kw.pop('transform', ax.transData)
972 self._pivot = pivot
973 self._length = length
974 barbcolor = barbcolor
975 flagcolor = flagcolor
977 # Flagcolor and barbcolor provide convenience parameters for
978 # setting the facecolor and edgecolor, respectively, of the barb
979 # polygon. We also work here to make the flag the same color as the
980 # rest of the barb by default
982 if None in (barbcolor, flagcolor):
983 kw['edgecolors'] = 'face'
984 if flagcolor:
985 kw['facecolors'] = flagcolor
986 elif barbcolor:
987 kw['facecolors'] = barbcolor
988 else:
989 # Set to facecolor passed in or default to black
990 kw.setdefault('facecolors', 'k')
991 else:
992 kw['edgecolors'] = barbcolor
993 kw['facecolors'] = flagcolor
995 # Explicitly set a line width if we're not given one, otherwise
996 # polygons are not outlined and we get no barbs
997 if 'linewidth' not in kw and 'lw' not in kw:
998 kw['linewidth'] = 1
1000 # Parse out the data arrays from the various configurations supported
1001 x, y, u, v, c = _parse_args(*args, caller_name='barbs()')
1002 self.x = x
1003 self.y = y
1004 xy = np.column_stack((x, y))
1006 # Make a collection
1007 barb_size = self._length ** 2 / 4 # Empirically determined
1008 mcollections.PolyCollection.__init__(self, [], (barb_size,),
1009 offsets=xy,
1010 transOffset=transform, **kw)
1011 self.set_transform(transforms.IdentityTransform())
1013 self.set_UVC(u, v, c)
1015 def _find_tails(self, mag, rounding=True, half=5, full=10, flag=50):
1016 '''
1017 Find how many of each of the tail pieces is necessary. Flag
1018 specifies the increment for a flag, barb for a full barb, and half for
1019 half a barb. Mag should be the magnitude of a vector (i.e., >= 0).
1021 This returns a tuple of:
1023 (*number of flags*, *number of barbs*, *half_flag*, *empty_flag*)
1025 *half_flag* is a boolean whether half of a barb is needed,
1026 since there should only ever be one half on a given
1027 barb. *empty_flag* flag is an array of flags to easily tell if
1028 a barb is empty (too low to plot any barbs/flags.
1029 '''
1031 # If rounding, round to the nearest multiple of half, the smallest
1032 # increment
1033 if rounding:
1034 mag = half * (mag / half + 0.5).astype(int)
1036 num_flags = np.floor(mag / flag).astype(int)
1037 mag = mag % flag
1039 num_barb = np.floor(mag / full).astype(int)
1040 mag = mag % full
1042 half_flag = mag >= half
1043 empty_flag = ~(half_flag | (num_flags > 0) | (num_barb > 0))
1045 return num_flags, num_barb, half_flag, empty_flag
1047 def _make_barbs(self, u, v, nflags, nbarbs, half_barb, empty_flag, length,
1048 pivot, sizes, fill_empty, flip):
1049 '''
1050 This function actually creates the wind barbs. *u* and *v*
1051 are components of the vector in the *x* and *y* directions,
1052 respectively.
1054 *nflags*, *nbarbs*, and *half_barb*, empty_flag* are,
1055 *respectively, the number of flags, number of barbs, flag for
1056 *half a barb, and flag for empty barb, ostensibly obtained
1057 *from :meth:`_find_tails`.
1059 *length* is the length of the barb staff in points.
1061 *pivot* specifies the point on the barb around which the
1062 entire barb should be rotated. Right now, valid options are
1063 'tip' and 'middle'. Can also be a number, which shifts the start
1064 of the barb that many points from the origin.
1066 *sizes* is a dictionary of coefficients specifying the ratio
1067 of a given feature to the length of the barb. These features
1068 include:
1070 - *spacing*: space between features (flags, full/half
1071 barbs)
1073 - *height*: distance from shaft of top of a flag or full
1074 barb
1076 - *width* - width of a flag, twice the width of a full barb
1078 - *emptybarb* - radius of the circle used for low
1079 magnitudes
1081 *fill_empty* specifies whether the circle representing an
1082 empty barb should be filled or not (this changes the drawing
1083 of the polygon).
1085 *flip* is a flag indicating whether the features should be flipped to
1086 the other side of the barb (useful for winds in the southern
1087 hemisphere).
1089 This function returns list of arrays of vertices, defining a polygon
1090 for each of the wind barbs. These polygons have been rotated to
1091 properly align with the vector direction.
1092 '''
1094 # These control the spacing and size of barb elements relative to the
1095 # length of the shaft
1096 spacing = length * sizes.get('spacing', 0.125)
1097 full_height = length * sizes.get('height', 0.4)
1098 full_width = length * sizes.get('width', 0.25)
1099 empty_rad = length * sizes.get('emptybarb', 0.15)
1101 # Controls y point where to pivot the barb.
1102 pivot_points = dict(tip=0.0, middle=-length / 2.)
1104 endx = 0.0
1105 try:
1106 endy = float(pivot)
1107 except ValueError:
1108 endy = pivot_points[pivot.lower()]
1110 # Get the appropriate angle for the vector components. The offset is
1111 # due to the way the barb is initially drawn, going down the y-axis.
1112 # This makes sense in a meteorological mode of thinking since there 0
1113 # degrees corresponds to north (the y-axis traditionally)
1114 angles = -(ma.arctan2(v, u) + np.pi / 2)
1116 # Used for low magnitude. We just get the vertices, so if we make it
1117 # out here, it can be reused. The center set here should put the
1118 # center of the circle at the location(offset), rather than at the
1119 # same point as the barb pivot; this seems more sensible.
1120 circ = CirclePolygon((0, 0), radius=empty_rad).get_verts()
1121 if fill_empty:
1122 empty_barb = circ
1123 else:
1124 # If we don't want the empty one filled, we make a degenerate
1125 # polygon that wraps back over itself
1126 empty_barb = np.concatenate((circ, circ[::-1]))
1128 barb_list = []
1129 for index, angle in np.ndenumerate(angles):
1130 # If the vector magnitude is too weak to draw anything, plot an
1131 # empty circle instead
1132 if empty_flag[index]:
1133 # We can skip the transform since the circle has no preferred
1134 # orientation
1135 barb_list.append(empty_barb)
1136 continue
1138 poly_verts = [(endx, endy)]
1139 offset = length
1141 # Handle if this barb should be flipped
1142 barb_height = -full_height if flip[index] else full_height
1144 # Add vertices for each flag
1145 for i in range(nflags[index]):
1146 # The spacing that works for the barbs is a little to much for
1147 # the flags, but this only occurs when we have more than 1
1148 # flag.
1149 if offset != length:
1150 offset += spacing / 2.
1151 poly_verts.extend(
1152 [[endx, endy + offset],
1153 [endx + barb_height, endy - full_width / 2 + offset],
1154 [endx, endy - full_width + offset]])
1156 offset -= full_width + spacing
1158 # Add vertices for each barb. These really are lines, but works
1159 # great adding 3 vertices that basically pull the polygon out and
1160 # back down the line
1161 for i in range(nbarbs[index]):
1162 poly_verts.extend(
1163 [(endx, endy + offset),
1164 (endx + barb_height, endy + offset + full_width / 2),
1165 (endx, endy + offset)])
1167 offset -= spacing
1169 # Add the vertices for half a barb, if needed
1170 if half_barb[index]:
1171 # If the half barb is the first on the staff, traditionally it
1172 # is offset from the end to make it easy to distinguish from a
1173 # barb with a full one
1174 if offset == length:
1175 poly_verts.append((endx, endy + offset))
1176 offset -= 1.5 * spacing
1177 poly_verts.extend(
1178 [(endx, endy + offset),
1179 (endx + barb_height / 2, endy + offset + full_width / 4),
1180 (endx, endy + offset)])
1182 # Rotate the barb according the angle. Making the barb first and
1183 # then rotating it made the math for drawing the barb really easy.
1184 # Also, the transform framework makes doing the rotation simple.
1185 poly_verts = transforms.Affine2D().rotate(-angle).transform(
1186 poly_verts)
1187 barb_list.append(poly_verts)
1189 return barb_list
1191 def set_UVC(self, U, V, C=None):
1192 self.u = ma.masked_invalid(U, copy=False).ravel()
1193 self.v = ma.masked_invalid(V, copy=False).ravel()
1195 # Flip needs to have the same number of entries as everything else.
1196 # Use broadcast_to to avoid a bloated array of identical values.
1197 # (can't rely on actual broadcasting)
1198 if len(self.flip) == 1:
1199 flip = np.broadcast_to(self.flip, self.u.shape)
1200 else:
1201 flip = self.flip
1203 if C is not None:
1204 c = ma.masked_invalid(C, copy=False).ravel()
1205 x, y, u, v, c, flip = cbook.delete_masked_points(
1206 self.x.ravel(), self.y.ravel(), self.u, self.v, c,
1207 flip.ravel())
1208 _check_consistent_shapes(x, y, u, v, c, flip)
1209 else:
1210 x, y, u, v, flip = cbook.delete_masked_points(
1211 self.x.ravel(), self.y.ravel(), self.u, self.v, flip.ravel())
1212 _check_consistent_shapes(x, y, u, v, flip)
1214 magnitude = np.hypot(u, v)
1215 flags, barbs, halves, empty = self._find_tails(magnitude,
1216 self.rounding,
1217 **self.barb_increments)
1219 # Get the vertices for each of the barbs
1221 plot_barbs = self._make_barbs(u, v, flags, barbs, halves, empty,
1222 self._length, self._pivot, self.sizes,
1223 self.fill_empty, flip)
1224 self.set_verts(plot_barbs)
1226 # Set the color array
1227 if C is not None:
1228 self.set_array(c)
1230 # Update the offsets in case the masked data changed
1231 xy = np.column_stack((x, y))
1232 self._offsets = xy
1233 self.stale = True
1235 def set_offsets(self, xy):
1236 """
1237 Set the offsets for the barb polygons. This saves the offsets passed
1238 in and masks them as appropriate for the existing U/V data.
1240 Parameters
1241 ----------
1242 xy : sequence of pairs of floats
1243 """
1244 self.x = xy[:, 0]
1245 self.y = xy[:, 1]
1246 x, y, u, v = cbook.delete_masked_points(
1247 self.x.ravel(), self.y.ravel(), self.u, self.v)
1248 _check_consistent_shapes(x, y, u, v)
1249 xy = np.column_stack((x, y))
1250 mcollections.PolyCollection.set_offsets(self, xy)
1251 self.stale = True
1253 barbs_doc = _barbs_doc