Coverage for /usr/lib/python3/dist-packages/matplotlib/quiver.py: 13%
406 statements
« prev ^ index » next coverage.py v7.4.4, created at 2025-06-14 15:25 +0200
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
19import numpy as np
20from numpy import ma
22from matplotlib import _api, cbook, _docstring, font_manager
23import matplotlib.artist as martist
24import matplotlib.collections as mcollections
25from matplotlib.patches import CirclePolygon
26import matplotlib.text as mtext
27import matplotlib.transforms as transforms
30_quiver_doc = """
31Plot a 2D field of arrows.
33Call signature::
35 quiver([X, Y], U, V, [C], **kwargs)
37*X*, *Y* define the arrow locations, *U*, *V* define the arrow directions, and
38*C* optionally sets the color.
40**Arrow length**
42The default settings auto-scales the length of the arrows to a reasonable size.
43To change this behavior see the *scale* and *scale_units* parameters.
45**Arrow shape**
47The arrow shape is determined by *width*, *headwidth*, *headlength* and
48*headaxislength*. See the notes below.
50**Arrow styling**
52Each arrow is internally represented by a filled polygon with a default edge
53linewidth of 0. As a result, an arrow is rather a filled area, not a line with
54a head, and `.PolyCollection` properties like *linewidth*, *edgecolor*,
55*facecolor*, etc. act accordingly.
58Parameters
59----------
60X, Y : 1D or 2D array-like, optional
61 The x and y coordinates of the arrow locations.
63 If not given, they will be generated as a uniform integer meshgrid based
64 on the dimensions of *U* and *V*.
66 If *X* and *Y* are 1D but *U*, *V* are 2D, *X*, *Y* are expanded to 2D
67 using ``X, Y = np.meshgrid(X, Y)``. In this case ``len(X)`` and ``len(Y)``
68 must match the column and row dimensions of *U* and *V*.
70U, V : 1D or 2D array-like
71 The x and y direction components of the arrow vectors. The interpretation
72 of these components (in data or in screen space) depends on *angles*.
74 *U* and *V* must have the same number of elements, matching the number of
75 arrow locations in *X*, *Y*. *U* and *V* may be masked. Locations masked
76 in any of *U*, *V*, and *C* will not be drawn.
78C : 1D or 2D array-like, optional
79 Numeric data that defines the arrow colors by colormapping via *norm* and
80 *cmap*.
82 This does not support explicit colors. If you want to set colors directly,
83 use *color* instead. The size of *C* must match the number of arrow
84 locations.
86angles : {'uv', 'xy'} or array-like, default: 'uv'
87 Method for determining the angle of the arrows.
89 - 'uv': Arrow direction in screen coordinates. Use this if the arrows
90 symbolize a quantity that is not based on *X*, *Y* data coordinates.
92 If *U* == *V* the orientation of the arrow on the plot is 45 degrees
93 counter-clockwise from the horizontal axis (positive to the right).
95 - 'xy': Arrow direction in data coordinates, i.e. the arrows point from
96 (x, y) to (x+u, y+v). Use this e.g. for plotting a gradient field.
98 - Arbitrary angles may be specified explicitly as an array of values
99 in degrees, counter-clockwise from the horizontal axis.
101 In this case *U*, *V* is only used to determine the length of the
102 arrows.
104 Note: inverting a data axis will correspondingly invert the
105 arrows only with ``angles='xy'``.
107pivot : {'tail', 'mid', 'middle', 'tip'}, default: 'tail'
108 The part of the arrow that is anchored to the *X*, *Y* grid. The arrow
109 rotates about this point.
111 'mid' is a synonym for 'middle'.
113scale : float, optional
114 Scales the length of the arrow inversely.
116 Number of data units per arrow length unit, e.g., m/s per plot width; a
117 smaller scale parameter makes the arrow longer. Default is *None*.
119 If *None*, a simple autoscaling algorithm is used, based on the average
120 vector length and the number of vectors. The arrow length unit is given by
121 the *scale_units* parameter.
123scale_units : {'width', 'height', 'dots', 'inches', 'x', 'y', 'xy'}, optional
124 If the *scale* kwarg is *None*, the arrow length unit. Default is *None*.
126 e.g. *scale_units* is 'inches', *scale* is 2.0, and ``(u, v) = (1, 0)``,
127 then the vector will be 0.5 inches long.
129 If *scale_units* is 'width' or 'height', then the vector will be half the
130 width/height of the axes.
132 If *scale_units* is 'x' then the vector will be 0.5 x-axis
133 units. To plot vectors in the x-y plane, with u and v having
134 the same units as x and y, use
135 ``angles='xy', scale_units='xy', scale=1``.
137units : {'width', 'height', 'dots', 'inches', 'x', 'y', 'xy'}, default: 'width'
138 Affects the arrow size (except for the length). In particular, the shaft
139 *width* is measured in multiples of this unit.
141 Supported values are:
143 - 'width', 'height': The width or height of the Axes.
144 - 'dots', 'inches': Pixels or inches based on the figure dpi.
145 - 'x', 'y', 'xy': *X*, *Y* or :math:`\\sqrt{X^2 + Y^2}` in data units.
147 The following table summarizes how these values affect the visible arrow
148 size under zooming and figure size changes:
150 ================= ================= ==================
151 units zoom figure size change
152 ================= ================= ==================
153 'x', 'y', 'xy' arrow size scales -
154 'width', 'height' - arrow size scales
155 'dots', 'inches' - -
156 ================= ================= ==================
158width : float, optional
159 Shaft width in arrow units. All head parameters are relative to *width*.
161 The default depends on choice of *units* above, and number of vectors;
162 a typical starting value is about 0.005 times the width of the plot.
164headwidth : float, default: 3
165 Head width as multiple of shaft *width*. See the notes below.
167headlength : float, default: 5
168 Head length as multiple of shaft *width*. See the notes below.
170headaxislength : float, default: 4.5
171 Head length at shaft intersection as multiple of shaft *width*.
172 See the notes below.
174minshaft : float, default: 1
175 Length below which arrow scales, in units of head length. Do not
176 set this to less than 1, or small arrows will look terrible!
178minlength : float, default: 1
179 Minimum length as a multiple of shaft width; if an arrow length
180 is less than this, plot a dot (hexagon) of this diameter instead.
182color : color or color sequence, optional
183 Explicit color(s) for the arrows. If *C* has been set, *color* has no
184 effect.
186 This is a synonym for the `.PolyCollection` *facecolor* parameter.
188Other Parameters
189----------------
190data : indexable object, optional
191 DATA_PARAMETER_PLACEHOLDER
193**kwargs : `~matplotlib.collections.PolyCollection` properties, optional
194 All other keyword arguments are passed on to `.PolyCollection`:
196 %(PolyCollection:kwdoc)s
198Returns
199-------
200`~matplotlib.quiver.Quiver`
202See Also
203--------
204.Axes.quiverkey : Add a key to a quiver plot.
206Notes
207-----
209**Arrow shape**
211The arrow is drawn as a polygon using the nodes as shown below. The values
212*headwidth*, *headlength*, and *headaxislength* are in units of *width*.
214.. image:: /_static/quiver_sizes.svg
215 :width: 500px
217The defaults give a slightly swept-back arrow. Here are some guidelines how to
218get other head shapes:
220- To make the head a triangle, make *headaxislength* the same as *headlength*.
221- To make the arrow more pointed, reduce *headwidth* or increase *headlength*
222 and *headaxislength*.
223- To make the head smaller relative to the shaft, scale down all the head
224 parameters proportionally.
225- To remove the head completely, set all *head* parameters to 0.
226- To get a diamond-shaped head, make *headaxislength* larger than *headlength*.
227- Warning: For *headaxislength* < (*headlength* / *headwidth*), the "headaxis"
228 nodes (i.e. the ones connecting the head with the shaft) will protrude out
229 of the head in forward direction so that the arrow head looks broken.
230""" % _docstring.interpd.params
232_docstring.interpd.update(quiver_doc=_quiver_doc)
235class QuiverKey(martist.Artist):
236 """Labelled arrow for use as a quiver plot scale key."""
237 halign = {'N': 'center', 'S': 'center', 'E': 'left', 'W': 'right'}
238 valign = {'N': 'bottom', 'S': 'top', 'E': 'center', 'W': 'center'}
239 pivot = {'N': 'middle', 'S': 'middle', 'E': 'tip', 'W': 'tail'}
241 def __init__(self, Q, X, Y, U, label,
242 *, angle=0, coordinates='axes', color=None, labelsep=0.1,
243 labelpos='N', labelcolor=None, fontproperties=None, **kwargs):
244 """
245 Add a key to a quiver plot.
247 The positioning of the key depends on *X*, *Y*, *coordinates*, and
248 *labelpos*. If *labelpos* is 'N' or 'S', *X*, *Y* give the position of
249 the middle of the key arrow. If *labelpos* is 'E', *X*, *Y* positions
250 the head, and if *labelpos* is 'W', *X*, *Y* positions the tail; in
251 either of these two cases, *X*, *Y* is somewhere in the middle of the
252 arrow+label key object.
254 Parameters
255 ----------
256 Q : `matplotlib.quiver.Quiver`
257 A `.Quiver` object as returned by a call to `~.Axes.quiver()`.
258 X, Y : float
259 The location of the key.
260 U : float
261 The length of the key.
262 label : str
263 The key label (e.g., length and units of the key).
264 angle : float, default: 0
265 The angle of the key arrow, in degrees anti-clockwise from the
266 x-axis.
267 coordinates : {'axes', 'figure', 'data', 'inches'}, default: 'axes'
268 Coordinate system and units for *X*, *Y*: 'axes' and 'figure' are
269 normalized coordinate systems with (0, 0) in the lower left and
270 (1, 1) in the upper right; 'data' are the axes data coordinates
271 (used for the locations of the vectors in the quiver plot itself);
272 'inches' is position in the figure in inches, with (0, 0) at the
273 lower left corner.
274 color : color
275 Overrides face and edge colors from *Q*.
276 labelpos : {'N', 'S', 'E', 'W'}
277 Position the label above, below, to the right, to the left of the
278 arrow, respectively.
279 labelsep : float, default: 0.1
280 Distance in inches between the arrow and the label.
281 labelcolor : color, default: :rc:`text.color`
282 Label color.
283 fontproperties : dict, optional
284 A dictionary with keyword arguments accepted by the
285 `~matplotlib.font_manager.FontProperties` initializer:
286 *family*, *style*, *variant*, *size*, *weight*.
287 **kwargs
288 Any additional keyword arguments are used to override vector
289 properties taken from *Q*.
290 """
291 super().__init__()
292 self.Q = Q
293 self.X = X
294 self.Y = Y
295 self.U = U
296 self.angle = angle
297 self.coord = coordinates
298 self.color = color
299 self.label = label
300 self._labelsep_inches = labelsep
302 self.labelpos = labelpos
303 self.labelcolor = labelcolor
304 self.fontproperties = fontproperties or dict()
305 self.kw = kwargs
306 _fp = self.fontproperties
307 self.text = mtext.Text(
308 text=label,
309 horizontalalignment=self.halign[self.labelpos],
310 verticalalignment=self.valign[self.labelpos],
311 fontproperties=font_manager.FontProperties._from_any(_fp))
313 if self.labelcolor is not None:
314 self.text.set_color(self.labelcolor)
315 self._dpi_at_last_init = None
316 self.zorder = Q.zorder + 0.1
318 @property
319 def labelsep(self):
320 return self._labelsep_inches * self.Q.axes.figure.dpi
322 def _init(self):
323 if True: # self._dpi_at_last_init != self.axes.figure.dpi
324 if self.Q._dpi_at_last_init != self.Q.axes.figure.dpi:
325 self.Q._init()
326 self._set_transform()
327 with cbook._setattr_cm(self.Q, pivot=self.pivot[self.labelpos],
328 # Hack: save and restore the Umask
329 Umask=ma.nomask):
330 u = self.U * np.cos(np.radians(self.angle))
331 v = self.U * np.sin(np.radians(self.angle))
332 angle = (self.Q.angles if isinstance(self.Q.angles, str)
333 else 'uv')
334 self.verts = self.Q._make_verts(
335 np.array([u]), np.array([v]), angle)
336 kwargs = self.Q.polykw
337 kwargs.update(self.kw)
338 self.vector = mcollections.PolyCollection(
339 self.verts,
340 offsets=[(self.X, self.Y)],
341 offset_transform=self.get_transform(),
342 **kwargs)
343 if self.color is not None:
344 self.vector.set_color(self.color)
345 self.vector.set_transform(self.Q.get_transform())
346 self.vector.set_figure(self.get_figure())
347 self._dpi_at_last_init = self.Q.axes.figure.dpi
349 def _text_x(self, x):
350 if self.labelpos == 'E':
351 return x + self.labelsep
352 elif self.labelpos == 'W':
353 return x - self.labelsep
354 else:
355 return x
357 def _text_y(self, y):
358 if self.labelpos == 'N':
359 return y + self.labelsep
360 elif self.labelpos == 'S':
361 return y - self.labelsep
362 else:
363 return y
365 @martist.allow_rasterization
366 def draw(self, renderer):
367 self._init()
368 self.vector.draw(renderer)
369 x, y = self.get_transform().transform((self.X, self.Y))
370 self.text.set_x(self._text_x(x))
371 self.text.set_y(self._text_y(y))
372 self.text.draw(renderer)
373 self.stale = False
375 def _set_transform(self):
376 self.set_transform(_api.check_getitem({
377 "data": self.Q.axes.transData,
378 "axes": self.Q.axes.transAxes,
379 "figure": self.Q.axes.figure.transFigure,
380 "inches": self.Q.axes.figure.dpi_scale_trans,
381 }, coordinates=self.coord))
383 def set_figure(self, fig):
384 super().set_figure(fig)
385 self.text.set_figure(fig)
387 def contains(self, mouseevent):
388 inside, info = self._default_contains(mouseevent)
389 if inside is not None:
390 return inside, info
391 # Maybe the dictionary should allow one to
392 # distinguish between a text hit and a vector hit.
393 if (self.text.contains(mouseevent)[0] or
394 self.vector.contains(mouseevent)[0]):
395 return True, {}
396 return False, {}
399def _parse_args(*args, caller_name='function'):
400 """
401 Helper function to parse positional parameters for colored vector plots.
403 This is currently used for Quiver and Barbs.
405 Parameters
406 ----------
407 *args : list
408 list of 2-5 arguments. Depending on their number they are parsed to::
410 U, V
411 U, V, C
412 X, Y, U, V
413 X, Y, U, V, C
415 caller_name : str
416 Name of the calling method (used in error messages).
417 """
418 X = Y = C = None
420 len_args = len(args)
421 if len_args == 2:
422 # The use of atleast_1d allows for handling scalar arguments while also
423 # keeping masked arrays
424 U, V = np.atleast_1d(*args)
425 elif len_args == 3:
426 U, V, C = np.atleast_1d(*args)
427 elif len_args == 4:
428 X, Y, U, V = np.atleast_1d(*args)
429 elif len_args == 5:
430 X, Y, U, V, C = np.atleast_1d(*args)
431 else:
432 raise TypeError(f'{caller_name} takes 2-5 positional arguments but '
433 f'{len_args} were given')
435 nr, nc = (1, U.shape[0]) if U.ndim == 1 else U.shape
437 if X is not None:
438 X = X.ravel()
439 Y = Y.ravel()
440 if len(X) == nc and len(Y) == nr:
441 X, Y = [a.ravel() for a in np.meshgrid(X, Y)]
442 elif len(X) != len(Y):
443 raise ValueError('X and Y must be the same size, but '
444 f'X.size is {X.size} and Y.size is {Y.size}.')
445 else:
446 indexgrid = np.meshgrid(np.arange(nc), np.arange(nr))
447 X, Y = [np.ravel(a) for a in indexgrid]
448 # Size validation for U, V, C is left to the set_UVC method.
449 return X, Y, U, V, C
452def _check_consistent_shapes(*arrays):
453 all_shapes = {a.shape for a in arrays}
454 if len(all_shapes) != 1:
455 raise ValueError('The shapes of the passed in arrays do not match')
458class Quiver(mcollections.PolyCollection):
459 """
460 Specialized PolyCollection for arrows.
462 The only API method is set_UVC(), which can be used
463 to change the size, orientation, and color of the
464 arrows; their locations are fixed when the class is
465 instantiated. Possibly this method will be useful
466 in animations.
468 Much of the work in this class is done in the draw()
469 method so that as much information as possible is available
470 about the plot. In subsequent draw() calls, recalculation
471 is limited to things that might have changed, so there
472 should be no performance penalty from putting the calculations
473 in the draw() method.
474 """
476 _PIVOT_VALS = ('tail', 'middle', 'tip')
478 @_docstring.Substitution(_quiver_doc)
479 def __init__(self, ax, *args,
480 scale=None, headwidth=3, headlength=5, headaxislength=4.5,
481 minshaft=1, minlength=1, units='width', scale_units=None,
482 angles='uv', width=None, color='k', pivot='tail', **kwargs):
483 """
484 The constructor takes one required argument, an Axes
485 instance, followed by the args and kwargs described
486 by the following pyplot interface documentation:
487 %s
488 """
489 self._axes = ax # The attr actually set by the Artist.axes property.
490 X, Y, U, V, C = _parse_args(*args, caller_name='quiver()')
491 self.X = X
492 self.Y = Y
493 self.XY = np.column_stack((X, Y))
494 self.N = len(X)
495 self.scale = scale
496 self.headwidth = headwidth
497 self.headlength = float(headlength)
498 self.headaxislength = headaxislength
499 self.minshaft = minshaft
500 self.minlength = minlength
501 self.units = units
502 self.scale_units = scale_units
503 self.angles = angles
504 self.width = width
506 if pivot.lower() == 'mid':
507 pivot = 'middle'
508 self.pivot = pivot.lower()
509 _api.check_in_list(self._PIVOT_VALS, pivot=self.pivot)
511 self.transform = kwargs.pop('transform', ax.transData)
512 kwargs.setdefault('facecolors', color)
513 kwargs.setdefault('linewidths', (0,))
514 super().__init__([], offsets=self.XY, offset_transform=self.transform,
515 closed=False, **kwargs)
516 self.polykw = kwargs
517 self.set_UVC(U, V, C)
518 self._dpi_at_last_init = None
520 def _init(self):
521 """
522 Initialization delayed until first draw;
523 allow time for axes setup.
524 """
525 # It seems that there are not enough event notifications
526 # available to have this work on an as-needed basis at present.
527 if True: # self._dpi_at_last_init != self.axes.figure.dpi
528 trans = self._set_transform()
529 self.span = trans.inverted().transform_bbox(self.axes.bbox).width
530 if self.width is None:
531 sn = np.clip(math.sqrt(self.N), 8, 25)
532 self.width = 0.06 * self.span / sn
534 # _make_verts sets self.scale if not already specified
535 if (self._dpi_at_last_init != self.axes.figure.dpi
536 and self.scale is None):
537 self._make_verts(self.U, self.V, self.angles)
539 self._dpi_at_last_init = self.axes.figure.dpi
541 def get_datalim(self, transData):
542 trans = self.get_transform()
543 offset_trf = self.get_offset_transform()
544 full_transform = (trans - transData) + (offset_trf - transData)
545 XY = full_transform.transform(self.XY)
546 bbox = transforms.Bbox.null()
547 bbox.update_from_data_xy(XY, ignore=True)
548 return bbox
550 @martist.allow_rasterization
551 def draw(self, renderer):
552 self._init()
553 verts = self._make_verts(self.U, self.V, self.angles)
554 self.set_verts(verts, closed=False)
555 super().draw(renderer)
556 self.stale = False
558 def set_UVC(self, U, V, C=None):
559 # We need to ensure we have a copy, not a reference
560 # to an array that might change before draw().
561 U = ma.masked_invalid(U, copy=True).ravel()
562 V = ma.masked_invalid(V, copy=True).ravel()
563 if C is not None:
564 C = ma.masked_invalid(C, copy=True).ravel()
565 for name, var in zip(('U', 'V', 'C'), (U, V, C)):
566 if not (var is None or var.size == self.N or var.size == 1):
567 raise ValueError(f'Argument {name} has a size {var.size}'
568 f' which does not match {self.N},'
569 ' the number of arrow positions')
571 mask = ma.mask_or(U.mask, V.mask, copy=False, shrink=True)
572 if C is not None:
573 mask = ma.mask_or(mask, C.mask, copy=False, shrink=True)
574 if mask is ma.nomask:
575 C = C.filled()
576 else:
577 C = ma.array(C, mask=mask, copy=False)
578 self.U = U.filled(1)
579 self.V = V.filled(1)
580 self.Umask = mask
581 if C is not None:
582 self.set_array(C)
583 self.stale = True
585 def _dots_per_unit(self, units):
586 """Return a scale factor for converting from units to pixels."""
587 bb = self.axes.bbox
588 vl = self.axes.viewLim
589 return _api.check_getitem({
590 'x': bb.width / vl.width,
591 'y': bb.height / vl.height,
592 'xy': np.hypot(*bb.size) / np.hypot(*vl.size),
593 'width': bb.width,
594 'height': bb.height,
595 'dots': 1.,
596 'inches': self.axes.figure.dpi,
597 }, units=units)
599 def _set_transform(self):
600 """
601 Set the PolyCollection transform to go
602 from arrow width units to pixels.
603 """
604 dx = self._dots_per_unit(self.units)
605 self._trans_scale = dx # pixels per arrow width unit
606 trans = transforms.Affine2D().scale(dx)
607 self.set_transform(trans)
608 return trans
610 def _angles_lengths(self, U, V, eps=1):
611 xy = self.axes.transData.transform(self.XY)
612 uv = np.column_stack((U, V))
613 xyp = self.axes.transData.transform(self.XY + eps * uv)
614 dxy = xyp - xy
615 angles = np.arctan2(dxy[:, 1], dxy[:, 0])
616 lengths = np.hypot(*dxy.T) / eps
617 return angles, lengths
619 def _make_verts(self, U, V, angles):
620 uv = (U + V * 1j)
621 str_angles = angles if isinstance(angles, str) else ''
622 if str_angles == 'xy' and self.scale_units == 'xy':
623 # Here eps is 1 so that if we get U, V by diffing
624 # the X, Y arrays, the vectors will connect the
625 # points, regardless of the axis scaling (including log).
626 angles, lengths = self._angles_lengths(U, V, eps=1)
627 elif str_angles == 'xy' or self.scale_units == 'xy':
628 # Calculate eps based on the extents of the plot
629 # so that we don't end up with roundoff error from
630 # adding a small number to a large.
631 eps = np.abs(self.axes.dataLim.extents).max() * 0.001
632 angles, lengths = self._angles_lengths(U, V, eps=eps)
633 if str_angles and self.scale_units == 'xy':
634 a = lengths
635 else:
636 a = np.abs(uv)
637 if self.scale is None:
638 sn = max(10, math.sqrt(self.N))
639 if self.Umask is not ma.nomask:
640 amean = a[~self.Umask].mean()
641 else:
642 amean = a.mean()
643 # crude auto-scaling
644 # scale is typical arrow length as a multiple of the arrow width
645 scale = 1.8 * amean * sn / self.span
646 if self.scale_units is None:
647 if self.scale is None:
648 self.scale = scale
649 widthu_per_lenu = 1.0
650 else:
651 if self.scale_units == 'xy':
652 dx = 1
653 else:
654 dx = self._dots_per_unit(self.scale_units)
655 widthu_per_lenu = dx / self._trans_scale
656 if self.scale is None:
657 self.scale = scale * widthu_per_lenu
658 length = a * (widthu_per_lenu / (self.scale * self.width))
659 X, Y = self._h_arrows(length)
660 if str_angles == 'xy':
661 theta = angles
662 elif str_angles == 'uv':
663 theta = np.angle(uv)
664 else:
665 theta = ma.masked_invalid(np.deg2rad(angles)).filled(0)
666 theta = theta.reshape((-1, 1)) # for broadcasting
667 xy = (X + Y * 1j) * np.exp(1j * theta) * self.width
668 XY = np.stack((xy.real, xy.imag), axis=2)
669 if self.Umask is not ma.nomask:
670 XY = ma.array(XY)
671 XY[self.Umask] = ma.masked
672 # This might be handled more efficiently with nans, given
673 # that nans will end up in the paths anyway.
675 return XY
677 def _h_arrows(self, length):
678 """Length is in arrow width units."""
679 # It might be possible to streamline the code
680 # and speed it up a bit by using complex (x, y)
681 # instead of separate arrays; but any gain would be slight.
682 minsh = self.minshaft * self.headlength
683 N = len(length)
684 length = length.reshape(N, 1)
685 # This number is chosen based on when pixel values overflow in Agg
686 # causing rendering errors
687 # length = np.minimum(length, 2 ** 16)
688 np.clip(length, 0, 2 ** 16, out=length)
689 # x, y: normal horizontal arrow
690 x = np.array([0, -self.headaxislength,
691 -self.headlength, 0],
692 np.float64)
693 x = x + np.array([0, 1, 1, 1]) * length
694 y = 0.5 * np.array([1, 1, self.headwidth, 0], np.float64)
695 y = np.repeat(y[np.newaxis, :], N, axis=0)
696 # x0, y0: arrow without shaft, for short vectors
697 x0 = np.array([0, minsh - self.headaxislength,
698 minsh - self.headlength, minsh], np.float64)
699 y0 = 0.5 * np.array([1, 1, self.headwidth, 0], np.float64)
700 ii = [0, 1, 2, 3, 2, 1, 0, 0]
701 X = x[:, ii]
702 Y = y[:, ii]
703 Y[:, 3:-1] *= -1
704 X0 = x0[ii]
705 Y0 = y0[ii]
706 Y0[3:-1] *= -1
707 shrink = length / minsh if minsh != 0. else 0.
708 X0 = shrink * X0[np.newaxis, :]
709 Y0 = shrink * Y0[np.newaxis, :]
710 short = np.repeat(length < minsh, 8, axis=1)
711 # Now select X0, Y0 if short, otherwise X, Y
712 np.copyto(X, X0, where=short)
713 np.copyto(Y, Y0, where=short)
714 if self.pivot == 'middle':
715 X -= 0.5 * X[:, 3, np.newaxis]
716 elif self.pivot == 'tip':
717 # numpy bug? using -= does not work here unless we multiply by a
718 # float first, as with 'mid'.
719 X = X - X[:, 3, np.newaxis]
720 elif self.pivot != 'tail':
721 _api.check_in_list(["middle", "tip", "tail"], pivot=self.pivot)
723 tooshort = length < self.minlength
724 if tooshort.any():
725 # Use a heptagonal dot:
726 th = np.arange(0, 8, 1, np.float64) * (np.pi / 3.0)
727 x1 = np.cos(th) * self.minlength * 0.5
728 y1 = np.sin(th) * self.minlength * 0.5
729 X1 = np.repeat(x1[np.newaxis, :], N, axis=0)
730 Y1 = np.repeat(y1[np.newaxis, :], N, axis=0)
731 tooshort = np.repeat(tooshort, 8, 1)
732 np.copyto(X, X1, where=tooshort)
733 np.copyto(Y, Y1, where=tooshort)
734 # Mask handling is deferred to the caller, _make_verts.
735 return X, Y
737 quiver_doc = _quiver_doc
740_barbs_doc = r"""
741Plot a 2D field of barbs.
743Call signature::
745 barbs([X, Y], U, V, [C], **kwargs)
747Where *X*, *Y* define the barb locations, *U*, *V* define the barb
748directions, and *C* optionally sets the color.
750All arguments may be 1D or 2D. *U*, *V*, *C* may be masked arrays, but masked
751*X*, *Y* are not supported at present.
753Barbs are traditionally used in meteorology as a way to plot the speed
754and direction of wind observations, but can technically be used to
755plot any two dimensional vector quantity. As opposed to arrows, which
756give vector magnitude by the length of the arrow, the barbs give more
757quantitative information about the vector magnitude by putting slanted
758lines or a triangle for various increments in magnitude, as show
759schematically below::
761 : /\ \
762 : / \ \
763 : / \ \ \
764 : / \ \ \
765 : ------------------------------
767The largest increment is given by a triangle (or "flag"). After those
768come full lines (barbs). The smallest increment is a half line. There
769is only, of course, ever at most 1 half line. If the magnitude is
770small and only needs a single half-line and no full lines or
771triangles, the half-line is offset from the end of the barb so that it
772can be easily distinguished from barbs with a single full line. The
773magnitude for the barb shown above would nominally be 65, using the
774standard increments of 50, 10, and 5.
776See also https://en.wikipedia.org/wiki/Wind_barb.
778Parameters
779----------
780X, Y : 1D or 2D array-like, optional
781 The x and y coordinates of the barb locations. See *pivot* for how the
782 barbs are drawn to the x, y positions.
784 If not given, they will be generated as a uniform integer meshgrid based
785 on the dimensions of *U* and *V*.
787 If *X* and *Y* are 1D but *U*, *V* are 2D, *X*, *Y* are expanded to 2D
788 using ``X, Y = np.meshgrid(X, Y)``. In this case ``len(X)`` and ``len(Y)``
789 must match the column and row dimensions of *U* and *V*.
791U, V : 1D or 2D array-like
792 The x and y components of the barb shaft.
794C : 1D or 2D array-like, optional
795 Numeric data that defines the barb colors by colormapping via *norm* and
796 *cmap*.
798 This does not support explicit colors. If you want to set colors directly,
799 use *barbcolor* instead.
801length : float, default: 7
802 Length of the barb in points; the other parts of the barb
803 are scaled against this.
805pivot : {'tip', 'middle'} or float, default: 'tip'
806 The part of the arrow that is anchored to the *X*, *Y* grid. The barb
807 rotates about this point. This can also be a number, which shifts the
808 start of the barb that many points away from grid point.
810barbcolor : color or color sequence
811 The color of all parts of the barb except for the flags. This parameter
812 is analogous to the *edgecolor* parameter for polygons, which can be used
813 instead. However this parameter will override facecolor.
815flagcolor : color or color sequence
816 The color of any flags on the barb. This parameter is analogous to the
817 *facecolor* parameter for polygons, which can be used instead. However,
818 this parameter will override facecolor. If this is not set (and *C* has
819 not either) then *flagcolor* will be set to match *barbcolor* so that the
820 barb has a uniform color. If *C* has been set, *flagcolor* has no effect.
822sizes : dict, optional
823 A dictionary of coefficients specifying the ratio of a given
824 feature to the length of the barb. Only those values one wishes to
825 override need to be included. These features include:
827 - 'spacing' - space between features (flags, full/half barbs)
828 - 'height' - height (distance from shaft to top) of a flag or full barb
829 - 'width' - width of a flag, twice the width of a full barb
830 - 'emptybarb' - radius of the circle used for low magnitudes
832fill_empty : bool, default: False
833 Whether the empty barbs (circles) that are drawn should be filled with
834 the flag color. If they are not filled, the center is transparent.
836rounding : bool, default: True
837 Whether the vector magnitude should be rounded when allocating barb
838 components. If True, the magnitude is rounded to the nearest multiple
839 of the half-barb increment. If False, the magnitude is simply truncated
840 to the next lowest multiple.
842barb_increments : dict, optional
843 A dictionary of increments specifying values to associate with
844 different parts of the barb. Only those values one wishes to
845 override need to be included.
847 - 'half' - half barbs (Default is 5)
848 - 'full' - full barbs (Default is 10)
849 - 'flag' - flags (default is 50)
851flip_barb : bool or array-like of bool, default: False
852 Whether the lines and flags should point opposite to normal.
853 Normal behavior is for the barbs and lines to point right (comes from wind
854 barbs having these features point towards low pressure in the Northern
855 Hemisphere).
857 A single value is applied to all barbs. Individual barbs can be flipped by
858 passing a bool array of the same size as *U* and *V*.
860Returns
861-------
862barbs : `~matplotlib.quiver.Barbs`
864Other Parameters
865----------------
866data : indexable object, optional
867 DATA_PARAMETER_PLACEHOLDER
869**kwargs
870 The barbs can further be customized using `.PolyCollection` keyword
871 arguments:
873 %(PolyCollection:kwdoc)s
874""" % _docstring.interpd.params
876_docstring.interpd.update(barbs_doc=_barbs_doc)
879class Barbs(mcollections.PolyCollection):
880 """
881 Specialized PolyCollection for barbs.
883 The only API method is :meth:`set_UVC`, which can be used to
884 change the size, orientation, and color of the arrows. Locations
885 are changed using the :meth:`set_offsets` collection method.
886 Possibly this method will be useful in animations.
888 There is one internal function :meth:`_find_tails` which finds
889 exactly what should be put on the barb given the vector magnitude.
890 From there :meth:`_make_barbs` is used to find the vertices of the
891 polygon to represent the barb based on this information.
892 """
893 # This may be an abuse of polygons here to render what is essentially maybe
894 # 1 triangle and a series of lines. It works fine as far as I can tell
895 # however.
896 @_docstring.interpd
897 def __init__(self, ax, *args,
898 pivot='tip', length=7, barbcolor=None, flagcolor=None,
899 sizes=None, fill_empty=False, barb_increments=None,
900 rounding=True, flip_barb=False, **kwargs):
901 """
902 The constructor takes one required argument, an Axes
903 instance, followed by the args and kwargs described
904 by the following pyplot interface documentation:
905 %(barbs_doc)s
906 """
907 self.sizes = sizes or dict()
908 self.fill_empty = fill_empty
909 self.barb_increments = barb_increments or dict()
910 self.rounding = rounding
911 self.flip = np.atleast_1d(flip_barb)
912 transform = kwargs.pop('transform', ax.transData)
913 self._pivot = pivot
914 self._length = length
916 # Flagcolor and barbcolor provide convenience parameters for
917 # setting the facecolor and edgecolor, respectively, of the barb
918 # polygon. We also work here to make the flag the same color as the
919 # rest of the barb by default
921 if None in (barbcolor, flagcolor):
922 kwargs['edgecolors'] = 'face'
923 if flagcolor:
924 kwargs['facecolors'] = flagcolor
925 elif barbcolor:
926 kwargs['facecolors'] = barbcolor
927 else:
928 # Set to facecolor passed in or default to black
929 kwargs.setdefault('facecolors', 'k')
930 else:
931 kwargs['edgecolors'] = barbcolor
932 kwargs['facecolors'] = flagcolor
934 # Explicitly set a line width if we're not given one, otherwise
935 # polygons are not outlined and we get no barbs
936 if 'linewidth' not in kwargs and 'lw' not in kwargs:
937 kwargs['linewidth'] = 1
939 # Parse out the data arrays from the various configurations supported
940 x, y, u, v, c = _parse_args(*args, caller_name='barbs()')
941 self.x = x
942 self.y = y
943 xy = np.column_stack((x, y))
945 # Make a collection
946 barb_size = self._length ** 2 / 4 # Empirically determined
947 super().__init__(
948 [], (barb_size,), offsets=xy, offset_transform=transform, **kwargs)
949 self.set_transform(transforms.IdentityTransform())
951 self.set_UVC(u, v, c)
953 def _find_tails(self, mag, rounding=True, half=5, full=10, flag=50):
954 """
955 Find how many of each of the tail pieces is necessary. Flag
956 specifies the increment for a flag, barb for a full barb, and half for
957 half a barb. Mag should be the magnitude of a vector (i.e., >= 0).
959 This returns a tuple of:
961 (*number of flags*, *number of barbs*, *half_flag*, *empty_flag*)
963 The bool *half_flag* indicates whether half of a barb is needed,
964 since there should only ever be one half on a given
965 barb. *empty_flag* flag is an array of flags to easily tell if
966 a barb is empty (too low to plot any barbs/flags.
967 """
969 # If rounding, round to the nearest multiple of half, the smallest
970 # increment
971 if rounding:
972 mag = half * (mag / half + 0.5).astype(int)
974 num_flags = np.floor(mag / flag).astype(int)
975 mag = mag % flag
977 num_barb = np.floor(mag / full).astype(int)
978 mag = mag % full
980 half_flag = mag >= half
981 empty_flag = ~(half_flag | (num_flags > 0) | (num_barb > 0))
983 return num_flags, num_barb, half_flag, empty_flag
985 def _make_barbs(self, u, v, nflags, nbarbs, half_barb, empty_flag, length,
986 pivot, sizes, fill_empty, flip):
987 """
988 Create the wind barbs.
990 Parameters
991 ----------
992 u, v
993 Components of the vector in the x and y directions, respectively.
995 nflags, nbarbs, half_barb, empty_flag
996 Respectively, the number of flags, number of barbs, flag for
997 half a barb, and flag for empty barb, ostensibly obtained from
998 :meth:`_find_tails`.
1000 length
1001 The length of the barb staff in points.
1003 pivot : {"tip", "middle"} or number
1004 The point on the barb around which the entire barb should be
1005 rotated. If a number, the start of the barb is shifted by that
1006 many points from the origin.
1008 sizes : dict
1009 Coefficients specifying the ratio of a given feature to the length
1010 of the barb. These features include:
1012 - *spacing*: space between features (flags, full/half barbs).
1013 - *height*: distance from shaft of top of a flag or full barb.
1014 - *width*: width of a flag, twice the width of a full barb.
1015 - *emptybarb*: radius of the circle used for low magnitudes.
1017 fill_empty : bool
1018 Whether the circle representing an empty barb should be filled or
1019 not (this changes the drawing of the polygon).
1021 flip : list of bool
1022 Whether the features should be flipped to the other side of the
1023 barb (useful for winds in the southern hemisphere).
1025 Returns
1026 -------
1027 list of arrays of vertices
1028 Polygon vertices for each of the wind barbs. These polygons have
1029 been rotated to properly align with the vector direction.
1030 """
1032 # These control the spacing and size of barb elements relative to the
1033 # length of the shaft
1034 spacing = length * sizes.get('spacing', 0.125)
1035 full_height = length * sizes.get('height', 0.4)
1036 full_width = length * sizes.get('width', 0.25)
1037 empty_rad = length * sizes.get('emptybarb', 0.15)
1039 # Controls y point where to pivot the barb.
1040 pivot_points = dict(tip=0.0, middle=-length / 2.)
1042 endx = 0.0
1043 try:
1044 endy = float(pivot)
1045 except ValueError:
1046 endy = pivot_points[pivot.lower()]
1048 # Get the appropriate angle for the vector components. The offset is
1049 # due to the way the barb is initially drawn, going down the y-axis.
1050 # This makes sense in a meteorological mode of thinking since there 0
1051 # degrees corresponds to north (the y-axis traditionally)
1052 angles = -(ma.arctan2(v, u) + np.pi / 2)
1054 # Used for low magnitude. We just get the vertices, so if we make it
1055 # out here, it can be reused. The center set here should put the
1056 # center of the circle at the location(offset), rather than at the
1057 # same point as the barb pivot; this seems more sensible.
1058 circ = CirclePolygon((0, 0), radius=empty_rad).get_verts()
1059 if fill_empty:
1060 empty_barb = circ
1061 else:
1062 # If we don't want the empty one filled, we make a degenerate
1063 # polygon that wraps back over itself
1064 empty_barb = np.concatenate((circ, circ[::-1]))
1066 barb_list = []
1067 for index, angle in np.ndenumerate(angles):
1068 # If the vector magnitude is too weak to draw anything, plot an
1069 # empty circle instead
1070 if empty_flag[index]:
1071 # We can skip the transform since the circle has no preferred
1072 # orientation
1073 barb_list.append(empty_barb)
1074 continue
1076 poly_verts = [(endx, endy)]
1077 offset = length
1079 # Handle if this barb should be flipped
1080 barb_height = -full_height if flip[index] else full_height
1082 # Add vertices for each flag
1083 for i in range(nflags[index]):
1084 # The spacing that works for the barbs is a little to much for
1085 # the flags, but this only occurs when we have more than 1
1086 # flag.
1087 if offset != length:
1088 offset += spacing / 2.
1089 poly_verts.extend(
1090 [[endx, endy + offset],
1091 [endx + barb_height, endy - full_width / 2 + offset],
1092 [endx, endy - full_width + offset]])
1094 offset -= full_width + spacing
1096 # Add vertices for each barb. These really are lines, but works
1097 # great adding 3 vertices that basically pull the polygon out and
1098 # back down the line
1099 for i in range(nbarbs[index]):
1100 poly_verts.extend(
1101 [(endx, endy + offset),
1102 (endx + barb_height, endy + offset + full_width / 2),
1103 (endx, endy + offset)])
1105 offset -= spacing
1107 # Add the vertices for half a barb, if needed
1108 if half_barb[index]:
1109 # If the half barb is the first on the staff, traditionally it
1110 # is offset from the end to make it easy to distinguish from a
1111 # barb with a full one
1112 if offset == length:
1113 poly_verts.append((endx, endy + offset))
1114 offset -= 1.5 * spacing
1115 poly_verts.extend(
1116 [(endx, endy + offset),
1117 (endx + barb_height / 2, endy + offset + full_width / 4),
1118 (endx, endy + offset)])
1120 # Rotate the barb according the angle. Making the barb first and
1121 # then rotating it made the math for drawing the barb really easy.
1122 # Also, the transform framework makes doing the rotation simple.
1123 poly_verts = transforms.Affine2D().rotate(-angle).transform(
1124 poly_verts)
1125 barb_list.append(poly_verts)
1127 return barb_list
1129 def set_UVC(self, U, V, C=None):
1130 # We need to ensure we have a copy, not a reference to an array that
1131 # might change before draw().
1132 self.u = ma.masked_invalid(U, copy=True).ravel()
1133 self.v = ma.masked_invalid(V, copy=True).ravel()
1135 # Flip needs to have the same number of entries as everything else.
1136 # Use broadcast_to to avoid a bloated array of identical values.
1137 # (can't rely on actual broadcasting)
1138 if len(self.flip) == 1:
1139 flip = np.broadcast_to(self.flip, self.u.shape)
1140 else:
1141 flip = self.flip
1143 if C is not None:
1144 c = ma.masked_invalid(C, copy=True).ravel()
1145 x, y, u, v, c, flip = cbook.delete_masked_points(
1146 self.x.ravel(), self.y.ravel(), self.u, self.v, c,
1147 flip.ravel())
1148 _check_consistent_shapes(x, y, u, v, c, flip)
1149 else:
1150 x, y, u, v, flip = cbook.delete_masked_points(
1151 self.x.ravel(), self.y.ravel(), self.u, self.v, flip.ravel())
1152 _check_consistent_shapes(x, y, u, v, flip)
1154 magnitude = np.hypot(u, v)
1155 flags, barbs, halves, empty = self._find_tails(magnitude,
1156 self.rounding,
1157 **self.barb_increments)
1159 # Get the vertices for each of the barbs
1161 plot_barbs = self._make_barbs(u, v, flags, barbs, halves, empty,
1162 self._length, self._pivot, self.sizes,
1163 self.fill_empty, flip)
1164 self.set_verts(plot_barbs)
1166 # Set the color array
1167 if C is not None:
1168 self.set_array(c)
1170 # Update the offsets in case the masked data changed
1171 xy = np.column_stack((x, y))
1172 self._offsets = xy
1173 self.stale = True
1175 def set_offsets(self, xy):
1176 """
1177 Set the offsets for the barb polygons. This saves the offsets passed
1178 in and masks them as appropriate for the existing U/V data.
1180 Parameters
1181 ----------
1182 xy : sequence of pairs of floats
1183 """
1184 self.x = xy[:, 0]
1185 self.y = xy[:, 1]
1186 x, y, u, v = cbook.delete_masked_points(
1187 self.x.ravel(), self.y.ravel(), self.u, self.v)
1188 _check_consistent_shapes(x, y, u, v)
1189 xy = np.column_stack((x, y))
1190 super().set_offsets(xy)
1191 self.stale = True
1193 barbs_doc = _barbs_doc