Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.9/dist-packages/matplotlib/quiver.py: 13%
Shortcuts on this page
r m x 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
19import numpy as np
20from numpy import ma
22from matplotlib import _api, cbook, _docstring
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 : :mpltype:`color` or list :mpltype:`color`, 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 horizontal 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 : :mpltype:`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 : :mpltype:`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 self.text = mtext.Text(
307 text=label,
308 horizontalalignment=self.halign[self.labelpos],
309 verticalalignment=self.valign[self.labelpos],
310 fontproperties=self.fontproperties)
311 if self.labelcolor is not None:
312 self.text.set_color(self.labelcolor)
313 self._dpi_at_last_init = None
314 self.zorder = Q.zorder + 0.1
316 @property
317 def labelsep(self):
318 return self._labelsep_inches * self.Q.axes.figure.dpi
320 def _init(self):
321 if True: # self._dpi_at_last_init != self.axes.figure.dpi
322 if self.Q._dpi_at_last_init != self.Q.axes.figure.dpi:
323 self.Q._init()
324 self._set_transform()
325 with cbook._setattr_cm(self.Q, pivot=self.pivot[self.labelpos],
326 # Hack: save and restore the Umask
327 Umask=ma.nomask):
328 u = self.U * np.cos(np.radians(self.angle))
329 v = self.U * np.sin(np.radians(self.angle))
330 self.verts = self.Q._make_verts([[0., 0.]],
331 np.array([u]), np.array([v]), 'uv')
332 kwargs = self.Q.polykw
333 kwargs.update(self.kw)
334 self.vector = mcollections.PolyCollection(
335 self.verts,
336 offsets=[(self.X, self.Y)],
337 offset_transform=self.get_transform(),
338 **kwargs)
339 if self.color is not None:
340 self.vector.set_color(self.color)
341 self.vector.set_transform(self.Q.get_transform())
342 self.vector.set_figure(self.get_figure())
343 self._dpi_at_last_init = self.Q.axes.figure.dpi
345 def _text_shift(self):
346 return {
347 "N": (0, +self.labelsep),
348 "S": (0, -self.labelsep),
349 "E": (+self.labelsep, 0),
350 "W": (-self.labelsep, 0),
351 }[self.labelpos]
353 @martist.allow_rasterization
354 def draw(self, renderer):
355 self._init()
356 self.vector.draw(renderer)
357 pos = self.get_transform().transform((self.X, self.Y))
358 self.text.set_position(pos + self._text_shift())
359 self.text.draw(renderer)
360 self.stale = False
362 def _set_transform(self):
363 self.set_transform(_api.check_getitem({
364 "data": self.Q.axes.transData,
365 "axes": self.Q.axes.transAxes,
366 "figure": self.Q.axes.figure.transFigure,
367 "inches": self.Q.axes.figure.dpi_scale_trans,
368 }, coordinates=self.coord))
370 def set_figure(self, fig):
371 super().set_figure(fig)
372 self.text.set_figure(fig)
374 def contains(self, mouseevent):
375 if self._different_canvas(mouseevent):
376 return False, {}
377 # Maybe the dictionary should allow one to
378 # distinguish between a text hit and a vector hit.
379 if (self.text.contains(mouseevent)[0] or
380 self.vector.contains(mouseevent)[0]):
381 return True, {}
382 return False, {}
385def _parse_args(*args, caller_name='function'):
386 """
387 Helper function to parse positional parameters for colored vector plots.
389 This is currently used for Quiver and Barbs.
391 Parameters
392 ----------
393 *args : list
394 list of 2-5 arguments. Depending on their number they are parsed to::
396 U, V
397 U, V, C
398 X, Y, U, V
399 X, Y, U, V, C
401 caller_name : str
402 Name of the calling method (used in error messages).
403 """
404 X = Y = C = None
406 nargs = len(args)
407 if nargs == 2:
408 # The use of atleast_1d allows for handling scalar arguments while also
409 # keeping masked arrays
410 U, V = np.atleast_1d(*args)
411 elif nargs == 3:
412 U, V, C = np.atleast_1d(*args)
413 elif nargs == 4:
414 X, Y, U, V = np.atleast_1d(*args)
415 elif nargs == 5:
416 X, Y, U, V, C = np.atleast_1d(*args)
417 else:
418 raise _api.nargs_error(caller_name, takes="from 2 to 5", given=nargs)
420 nr, nc = (1, U.shape[0]) if U.ndim == 1 else U.shape
422 if X is not None:
423 X = X.ravel()
424 Y = Y.ravel()
425 if len(X) == nc and len(Y) == nr:
426 X, Y = [a.ravel() for a in np.meshgrid(X, Y)]
427 elif len(X) != len(Y):
428 raise ValueError('X and Y must be the same size, but '
429 f'X.size is {X.size} and Y.size is {Y.size}.')
430 else:
431 indexgrid = np.meshgrid(np.arange(nc), np.arange(nr))
432 X, Y = [np.ravel(a) for a in indexgrid]
433 # Size validation for U, V, C is left to the set_UVC method.
434 return X, Y, U, V, C
437def _check_consistent_shapes(*arrays):
438 all_shapes = {a.shape for a in arrays}
439 if len(all_shapes) != 1:
440 raise ValueError('The shapes of the passed in arrays do not match')
443class Quiver(mcollections.PolyCollection):
444 """
445 Specialized PolyCollection for arrows.
447 The only API method is set_UVC(), which can be used
448 to change the size, orientation, and color of the
449 arrows; their locations are fixed when the class is
450 instantiated. Possibly this method will be useful
451 in animations.
453 Much of the work in this class is done in the draw()
454 method so that as much information as possible is available
455 about the plot. In subsequent draw() calls, recalculation
456 is limited to things that might have changed, so there
457 should be no performance penalty from putting the calculations
458 in the draw() method.
459 """
461 _PIVOT_VALS = ('tail', 'middle', 'tip')
463 @_docstring.Substitution(_quiver_doc)
464 def __init__(self, ax, *args,
465 scale=None, headwidth=3, headlength=5, headaxislength=4.5,
466 minshaft=1, minlength=1, units='width', scale_units=None,
467 angles='uv', width=None, color='k', pivot='tail', **kwargs):
468 """
469 The constructor takes one required argument, an Axes
470 instance, followed by the args and kwargs described
471 by the following pyplot interface documentation:
472 %s
473 """
474 self._axes = ax # The attr actually set by the Artist.axes property.
475 X, Y, U, V, C = _parse_args(*args, caller_name='quiver')
476 self.X = X
477 self.Y = Y
478 self.XY = np.column_stack((X, Y))
479 self.N = len(X)
480 self.scale = scale
481 self.headwidth = headwidth
482 self.headlength = float(headlength)
483 self.headaxislength = headaxislength
484 self.minshaft = minshaft
485 self.minlength = minlength
486 self.units = units
487 self.scale_units = scale_units
488 self.angles = angles
489 self.width = width
491 if pivot.lower() == 'mid':
492 pivot = 'middle'
493 self.pivot = pivot.lower()
494 _api.check_in_list(self._PIVOT_VALS, pivot=self.pivot)
496 self.transform = kwargs.pop('transform', ax.transData)
497 kwargs.setdefault('facecolors', color)
498 kwargs.setdefault('linewidths', (0,))
499 super().__init__([], offsets=self.XY, offset_transform=self.transform,
500 closed=False, **kwargs)
501 self.polykw = kwargs
502 self.set_UVC(U, V, C)
503 self._dpi_at_last_init = None
505 def _init(self):
506 """
507 Initialization delayed until first draw;
508 allow time for axes setup.
509 """
510 # It seems that there are not enough event notifications
511 # available to have this work on an as-needed basis at present.
512 if True: # self._dpi_at_last_init != self.axes.figure.dpi
513 trans = self._set_transform()
514 self.span = trans.inverted().transform_bbox(self.axes.bbox).width
515 if self.width is None:
516 sn = np.clip(math.sqrt(self.N), 8, 25)
517 self.width = 0.06 * self.span / sn
519 # _make_verts sets self.scale if not already specified
520 if (self._dpi_at_last_init != self.axes.figure.dpi
521 and self.scale is None):
522 self._make_verts(self.XY, self.U, self.V, self.angles)
524 self._dpi_at_last_init = self.axes.figure.dpi
526 def get_datalim(self, transData):
527 trans = self.get_transform()
528 offset_trf = self.get_offset_transform()
529 full_transform = (trans - transData) + (offset_trf - transData)
530 XY = full_transform.transform(self.XY)
531 bbox = transforms.Bbox.null()
532 bbox.update_from_data_xy(XY, ignore=True)
533 return bbox
535 @martist.allow_rasterization
536 def draw(self, renderer):
537 self._init()
538 verts = self._make_verts(self.XY, self.U, self.V, self.angles)
539 self.set_verts(verts, closed=False)
540 super().draw(renderer)
541 self.stale = False
543 def set_UVC(self, U, V, C=None):
544 # We need to ensure we have a copy, not a reference
545 # to an array that might change before draw().
546 U = ma.masked_invalid(U, copy=True).ravel()
547 V = ma.masked_invalid(V, copy=True).ravel()
548 if C is not None:
549 C = ma.masked_invalid(C, copy=True).ravel()
550 for name, var in zip(('U', 'V', 'C'), (U, V, C)):
551 if not (var is None or var.size == self.N or var.size == 1):
552 raise ValueError(f'Argument {name} has a size {var.size}'
553 f' which does not match {self.N},'
554 ' the number of arrow positions')
556 mask = ma.mask_or(U.mask, V.mask, copy=False, shrink=True)
557 if C is not None:
558 mask = ma.mask_or(mask, C.mask, copy=False, shrink=True)
559 if mask is ma.nomask:
560 C = C.filled()
561 else:
562 C = ma.array(C, mask=mask, copy=False)
563 self.U = U.filled(1)
564 self.V = V.filled(1)
565 self.Umask = mask
566 if C is not None:
567 self.set_array(C)
568 self.stale = True
570 def _dots_per_unit(self, units):
571 """Return a scale factor for converting from units to pixels."""
572 bb = self.axes.bbox
573 vl = self.axes.viewLim
574 return _api.check_getitem({
575 'x': bb.width / vl.width,
576 'y': bb.height / vl.height,
577 'xy': np.hypot(*bb.size) / np.hypot(*vl.size),
578 'width': bb.width,
579 'height': bb.height,
580 'dots': 1.,
581 'inches': self.axes.figure.dpi,
582 }, units=units)
584 def _set_transform(self):
585 """
586 Set the PolyCollection transform to go
587 from arrow width units to pixels.
588 """
589 dx = self._dots_per_unit(self.units)
590 self._trans_scale = dx # pixels per arrow width unit
591 trans = transforms.Affine2D().scale(dx)
592 self.set_transform(trans)
593 return trans
595 # Calculate angles and lengths for segment between (x, y), (x+u, y+v)
596 def _angles_lengths(self, XY, U, V, eps=1):
597 xy = self.axes.transData.transform(XY)
598 uv = np.column_stack((U, V))
599 xyp = self.axes.transData.transform(XY + eps * uv)
600 dxy = xyp - xy
601 angles = np.arctan2(dxy[:, 1], dxy[:, 0])
602 lengths = np.hypot(*dxy.T) / eps
603 return angles, lengths
605 # XY is stacked [X, Y].
606 # See quiver() doc for meaning of X, Y, U, V, angles.
607 def _make_verts(self, XY, U, V, angles):
608 uv = (U + V * 1j)
609 str_angles = angles if isinstance(angles, str) else ''
610 if str_angles == 'xy' and self.scale_units == 'xy':
611 # Here eps is 1 so that if we get U, V by diffing
612 # the X, Y arrays, the vectors will connect the
613 # points, regardless of the axis scaling (including log).
614 angles, lengths = self._angles_lengths(XY, U, V, eps=1)
615 elif str_angles == 'xy' or self.scale_units == 'xy':
616 # Calculate eps based on the extents of the plot
617 # so that we don't end up with roundoff error from
618 # adding a small number to a large.
619 eps = np.abs(self.axes.dataLim.extents).max() * 0.001
620 angles, lengths = self._angles_lengths(XY, U, V, eps=eps)
622 if str_angles and self.scale_units == 'xy':
623 a = lengths
624 else:
625 a = np.abs(uv)
627 if self.scale is None:
628 sn = max(10, math.sqrt(self.N))
629 if self.Umask is not ma.nomask:
630 amean = a[~self.Umask].mean()
631 else:
632 amean = a.mean()
633 # crude auto-scaling
634 # scale is typical arrow length as a multiple of the arrow width
635 scale = 1.8 * amean * sn / self.span
637 if self.scale_units is None:
638 if self.scale is None:
639 self.scale = scale
640 widthu_per_lenu = 1.0
641 else:
642 if self.scale_units == 'xy':
643 dx = 1
644 else:
645 dx = self._dots_per_unit(self.scale_units)
646 widthu_per_lenu = dx / self._trans_scale
647 if self.scale is None:
648 self.scale = scale * widthu_per_lenu
649 length = a * (widthu_per_lenu / (self.scale * self.width))
650 X, Y = self._h_arrows(length)
651 if str_angles == 'xy':
652 theta = angles
653 elif str_angles == 'uv':
654 theta = np.angle(uv)
655 else:
656 theta = ma.masked_invalid(np.deg2rad(angles)).filled(0)
657 theta = theta.reshape((-1, 1)) # for broadcasting
658 xy = (X + Y * 1j) * np.exp(1j * theta) * self.width
659 XY = np.stack((xy.real, xy.imag), axis=2)
660 if self.Umask is not ma.nomask:
661 XY = ma.array(XY)
662 XY[self.Umask] = ma.masked
663 # This might be handled more efficiently with nans, given
664 # that nans will end up in the paths anyway.
666 return XY
668 def _h_arrows(self, length):
669 """Length is in arrow width units."""
670 # It might be possible to streamline the code
671 # and speed it up a bit by using complex (x, y)
672 # instead of separate arrays; but any gain would be slight.
673 minsh = self.minshaft * self.headlength
674 N = len(length)
675 length = length.reshape(N, 1)
676 # This number is chosen based on when pixel values overflow in Agg
677 # causing rendering errors
678 # length = np.minimum(length, 2 ** 16)
679 np.clip(length, 0, 2 ** 16, out=length)
680 # x, y: normal horizontal arrow
681 x = np.array([0, -self.headaxislength,
682 -self.headlength, 0],
683 np.float64)
684 x = x + np.array([0, 1, 1, 1]) * length
685 y = 0.5 * np.array([1, 1, self.headwidth, 0], np.float64)
686 y = np.repeat(y[np.newaxis, :], N, axis=0)
687 # x0, y0: arrow without shaft, for short vectors
688 x0 = np.array([0, minsh - self.headaxislength,
689 minsh - self.headlength, minsh], np.float64)
690 y0 = 0.5 * np.array([1, 1, self.headwidth, 0], np.float64)
691 ii = [0, 1, 2, 3, 2, 1, 0, 0]
692 X = x[:, ii]
693 Y = y[:, ii]
694 Y[:, 3:-1] *= -1
695 X0 = x0[ii]
696 Y0 = y0[ii]
697 Y0[3:-1] *= -1
698 shrink = length / minsh if minsh != 0. else 0.
699 X0 = shrink * X0[np.newaxis, :]
700 Y0 = shrink * Y0[np.newaxis, :]
701 short = np.repeat(length < minsh, 8, axis=1)
702 # Now select X0, Y0 if short, otherwise X, Y
703 np.copyto(X, X0, where=short)
704 np.copyto(Y, Y0, where=short)
705 if self.pivot == 'middle':
706 X -= 0.5 * X[:, 3, np.newaxis]
707 elif self.pivot == 'tip':
708 # numpy bug? using -= does not work here unless we multiply by a
709 # float first, as with 'mid'.
710 X = X - X[:, 3, np.newaxis]
711 elif self.pivot != 'tail':
712 _api.check_in_list(["middle", "tip", "tail"], pivot=self.pivot)
714 tooshort = length < self.minlength
715 if tooshort.any():
716 # Use a heptagonal dot:
717 th = np.arange(0, 8, 1, np.float64) * (np.pi / 3.0)
718 x1 = np.cos(th) * self.minlength * 0.5
719 y1 = np.sin(th) * self.minlength * 0.5
720 X1 = np.repeat(x1[np.newaxis, :], N, axis=0)
721 Y1 = np.repeat(y1[np.newaxis, :], N, axis=0)
722 tooshort = np.repeat(tooshort, 8, 1)
723 np.copyto(X, X1, where=tooshort)
724 np.copyto(Y, Y1, where=tooshort)
725 # Mask handling is deferred to the caller, _make_verts.
726 return X, Y
729_barbs_doc = r"""
730Plot a 2D field of wind barbs.
732Call signature::
734 barbs([X, Y], U, V, [C], **kwargs)
736Where *X*, *Y* define the barb locations, *U*, *V* define the barb
737directions, and *C* optionally sets the color.
739All arguments may be 1D or 2D. *U*, *V*, *C* may be masked arrays, but masked
740*X*, *Y* are not supported at present.
742Barbs are traditionally used in meteorology as a way to plot the speed
743and direction of wind observations, but can technically be used to
744plot any two dimensional vector quantity. As opposed to arrows, which
745give vector magnitude by the length of the arrow, the barbs give more
746quantitative information about the vector magnitude by putting slanted
747lines or a triangle for various increments in magnitude, as show
748schematically below::
750 : /\ \
751 : / \ \
752 : / \ \ \
753 : / \ \ \
754 : ------------------------------
756The largest increment is given by a triangle (or "flag"). After those
757come full lines (barbs). The smallest increment is a half line. There
758is only, of course, ever at most 1 half line. If the magnitude is
759small and only needs a single half-line and no full lines or
760triangles, the half-line is offset from the end of the barb so that it
761can be easily distinguished from barbs with a single full line. The
762magnitude for the barb shown above would nominally be 65, using the
763standard increments of 50, 10, and 5.
765See also https://en.wikipedia.org/wiki/Wind_barb.
767Parameters
768----------
769X, Y : 1D or 2D array-like, optional
770 The x and y coordinates of the barb locations. See *pivot* for how the
771 barbs are drawn to the x, y positions.
773 If not given, they will be generated as a uniform integer meshgrid based
774 on the dimensions of *U* and *V*.
776 If *X* and *Y* are 1D but *U*, *V* are 2D, *X*, *Y* are expanded to 2D
777 using ``X, Y = np.meshgrid(X, Y)``. In this case ``len(X)`` and ``len(Y)``
778 must match the column and row dimensions of *U* and *V*.
780U, V : 1D or 2D array-like
781 The x and y components of the barb shaft.
783C : 1D or 2D array-like, optional
784 Numeric data that defines the barb colors by colormapping via *norm* and
785 *cmap*.
787 This does not support explicit colors. If you want to set colors directly,
788 use *barbcolor* instead.
790length : float, default: 7
791 Length of the barb in points; the other parts of the barb
792 are scaled against this.
794pivot : {'tip', 'middle'} or float, default: 'tip'
795 The part of the arrow that is anchored to the *X*, *Y* grid. The barb
796 rotates about this point. This can also be a number, which shifts the
797 start of the barb that many points away from grid point.
799barbcolor : :mpltype:`color` or color sequence
800 The color of all parts of the barb except for the flags. This parameter
801 is analogous to the *edgecolor* parameter for polygons, which can be used
802 instead. However this parameter will override facecolor.
804flagcolor : :mpltype:`color` or color sequence
805 The color of any flags on the barb. This parameter is analogous to the
806 *facecolor* parameter for polygons, which can be used instead. However,
807 this parameter will override facecolor. If this is not set (and *C* has
808 not either) then *flagcolor* will be set to match *barbcolor* so that the
809 barb has a uniform color. If *C* has been set, *flagcolor* has no effect.
811sizes : dict, optional
812 A dictionary of coefficients specifying the ratio of a given
813 feature to the length of the barb. Only those values one wishes to
814 override need to be included. These features include:
816 - 'spacing' - space between features (flags, full/half barbs)
817 - 'height' - height (distance from shaft to top) of a flag or full barb
818 - 'width' - width of a flag, twice the width of a full barb
819 - 'emptybarb' - radius of the circle used for low magnitudes
821fill_empty : bool, default: False
822 Whether the empty barbs (circles) that are drawn should be filled with
823 the flag color. If they are not filled, the center is transparent.
825rounding : bool, default: True
826 Whether the vector magnitude should be rounded when allocating barb
827 components. If True, the magnitude is rounded to the nearest multiple
828 of the half-barb increment. If False, the magnitude is simply truncated
829 to the next lowest multiple.
831barb_increments : dict, optional
832 A dictionary of increments specifying values to associate with
833 different parts of the barb. Only those values one wishes to
834 override need to be included.
836 - 'half' - half barbs (Default is 5)
837 - 'full' - full barbs (Default is 10)
838 - 'flag' - flags (default is 50)
840flip_barb : bool or array-like of bool, default: False
841 Whether the lines and flags should point opposite to normal.
842 Normal behavior is for the barbs and lines to point right (comes from wind
843 barbs having these features point towards low pressure in the Northern
844 Hemisphere).
846 A single value is applied to all barbs. Individual barbs can be flipped by
847 passing a bool array of the same size as *U* and *V*.
849Returns
850-------
851barbs : `~matplotlib.quiver.Barbs`
853Other Parameters
854----------------
855data : indexable object, optional
856 DATA_PARAMETER_PLACEHOLDER
858**kwargs
859 The barbs can further be customized using `.PolyCollection` keyword
860 arguments:
862 %(PolyCollection:kwdoc)s
863""" % _docstring.interpd.params
865_docstring.interpd.update(barbs_doc=_barbs_doc)
868class Barbs(mcollections.PolyCollection):
869 """
870 Specialized PolyCollection for barbs.
872 The only API method is :meth:`set_UVC`, which can be used to
873 change the size, orientation, and color of the arrows. Locations
874 are changed using the :meth:`set_offsets` collection method.
875 Possibly this method will be useful in animations.
877 There is one internal function :meth:`_find_tails` which finds
878 exactly what should be put on the barb given the vector magnitude.
879 From there :meth:`_make_barbs` is used to find the vertices of the
880 polygon to represent the barb based on this information.
881 """
883 # This may be an abuse of polygons here to render what is essentially maybe
884 # 1 triangle and a series of lines. It works fine as far as I can tell
885 # however.
887 @_docstring.interpd
888 def __init__(self, ax, *args,
889 pivot='tip', length=7, barbcolor=None, flagcolor=None,
890 sizes=None, fill_empty=False, barb_increments=None,
891 rounding=True, flip_barb=False, **kwargs):
892 """
893 The constructor takes one required argument, an Axes
894 instance, followed by the args and kwargs described
895 by the following pyplot interface documentation:
896 %(barbs_doc)s
897 """
898 self.sizes = sizes or dict()
899 self.fill_empty = fill_empty
900 self.barb_increments = barb_increments or dict()
901 self.rounding = rounding
902 self.flip = np.atleast_1d(flip_barb)
903 transform = kwargs.pop('transform', ax.transData)
904 self._pivot = pivot
905 self._length = length
907 # Flagcolor and barbcolor provide convenience parameters for
908 # setting the facecolor and edgecolor, respectively, of the barb
909 # polygon. We also work here to make the flag the same color as the
910 # rest of the barb by default
912 if None in (barbcolor, flagcolor):
913 kwargs['edgecolors'] = 'face'
914 if flagcolor:
915 kwargs['facecolors'] = flagcolor
916 elif barbcolor:
917 kwargs['facecolors'] = barbcolor
918 else:
919 # Set to facecolor passed in or default to black
920 kwargs.setdefault('facecolors', 'k')
921 else:
922 kwargs['edgecolors'] = barbcolor
923 kwargs['facecolors'] = flagcolor
925 # Explicitly set a line width if we're not given one, otherwise
926 # polygons are not outlined and we get no barbs
927 if 'linewidth' not in kwargs and 'lw' not in kwargs:
928 kwargs['linewidth'] = 1
930 # Parse out the data arrays from the various configurations supported
931 x, y, u, v, c = _parse_args(*args, caller_name='barbs')
932 self.x = x
933 self.y = y
934 xy = np.column_stack((x, y))
936 # Make a collection
937 barb_size = self._length ** 2 / 4 # Empirically determined
938 super().__init__(
939 [], (barb_size,), offsets=xy, offset_transform=transform, **kwargs)
940 self.set_transform(transforms.IdentityTransform())
942 self.set_UVC(u, v, c)
944 def _find_tails(self, mag, rounding=True, half=5, full=10, flag=50):
945 """
946 Find how many of each of the tail pieces is necessary.
948 Parameters
949 ----------
950 mag : `~numpy.ndarray`
951 Vector magnitudes; must be non-negative (and an actual ndarray).
952 rounding : bool, default: True
953 Whether to round or to truncate to the nearest half-barb.
954 half, full, flag : float, defaults: 5, 10, 50
955 Increments for a half-barb, a barb, and a flag.
957 Returns
958 -------
959 n_flags, n_barbs : int array
960 For each entry in *mag*, the number of flags and barbs.
961 half_flag : bool array
962 For each entry in *mag*, whether a half-barb is needed.
963 empty_flag : bool array
964 For each entry in *mag*, whether nothing is drawn.
965 """
966 # If rounding, round to the nearest multiple of half, the smallest
967 # increment
968 if rounding:
969 mag = half * np.around(mag / half)
970 n_flags, mag = divmod(mag, flag)
971 n_barb, mag = divmod(mag, full)
972 half_flag = mag >= half
973 empty_flag = ~(half_flag | (n_flags > 0) | (n_barb > 0))
974 return n_flags.astype(int), n_barb.astype(int), half_flag, empty_flag
976 def _make_barbs(self, u, v, nflags, nbarbs, half_barb, empty_flag, length,
977 pivot, sizes, fill_empty, flip):
978 """
979 Create the wind barbs.
981 Parameters
982 ----------
983 u, v
984 Components of the vector in the x and y directions, respectively.
986 nflags, nbarbs, half_barb, empty_flag
987 Respectively, the number of flags, number of barbs, flag for
988 half a barb, and flag for empty barb, ostensibly obtained from
989 :meth:`_find_tails`.
991 length
992 The length of the barb staff in points.
994 pivot : {"tip", "middle"} or number
995 The point on the barb around which the entire barb should be
996 rotated. If a number, the start of the barb is shifted by that
997 many points from the origin.
999 sizes : dict
1000 Coefficients specifying the ratio of a given feature to the length
1001 of the barb. These features include:
1003 - *spacing*: space between features (flags, full/half barbs).
1004 - *height*: distance from shaft of top of a flag or full barb.
1005 - *width*: width of a flag, twice the width of a full barb.
1006 - *emptybarb*: radius of the circle used for low magnitudes.
1008 fill_empty : bool
1009 Whether the circle representing an empty barb should be filled or
1010 not (this changes the drawing of the polygon).
1012 flip : list of bool
1013 Whether the features should be flipped to the other side of the
1014 barb (useful for winds in the southern hemisphere).
1016 Returns
1017 -------
1018 list of arrays of vertices
1019 Polygon vertices for each of the wind barbs. These polygons have
1020 been rotated to properly align with the vector direction.
1021 """
1023 # These control the spacing and size of barb elements relative to the
1024 # length of the shaft
1025 spacing = length * sizes.get('spacing', 0.125)
1026 full_height = length * sizes.get('height', 0.4)
1027 full_width = length * sizes.get('width', 0.25)
1028 empty_rad = length * sizes.get('emptybarb', 0.15)
1030 # Controls y point where to pivot the barb.
1031 pivot_points = dict(tip=0.0, middle=-length / 2.)
1033 endx = 0.0
1034 try:
1035 endy = float(pivot)
1036 except ValueError:
1037 endy = pivot_points[pivot.lower()]
1039 # Get the appropriate angle for the vector components. The offset is
1040 # due to the way the barb is initially drawn, going down the y-axis.
1041 # This makes sense in a meteorological mode of thinking since there 0
1042 # degrees corresponds to north (the y-axis traditionally)
1043 angles = -(ma.arctan2(v, u) + np.pi / 2)
1045 # Used for low magnitude. We just get the vertices, so if we make it
1046 # out here, it can be reused. The center set here should put the
1047 # center of the circle at the location(offset), rather than at the
1048 # same point as the barb pivot; this seems more sensible.
1049 circ = CirclePolygon((0, 0), radius=empty_rad).get_verts()
1050 if fill_empty:
1051 empty_barb = circ
1052 else:
1053 # If we don't want the empty one filled, we make a degenerate
1054 # polygon that wraps back over itself
1055 empty_barb = np.concatenate((circ, circ[::-1]))
1057 barb_list = []
1058 for index, angle in np.ndenumerate(angles):
1059 # If the vector magnitude is too weak to draw anything, plot an
1060 # empty circle instead
1061 if empty_flag[index]:
1062 # We can skip the transform since the circle has no preferred
1063 # orientation
1064 barb_list.append(empty_barb)
1065 continue
1067 poly_verts = [(endx, endy)]
1068 offset = length
1070 # Handle if this barb should be flipped
1071 barb_height = -full_height if flip[index] else full_height
1073 # Add vertices for each flag
1074 for i in range(nflags[index]):
1075 # The spacing that works for the barbs is a little to much for
1076 # the flags, but this only occurs when we have more than 1
1077 # flag.
1078 if offset != length:
1079 offset += spacing / 2.
1080 poly_verts.extend(
1081 [[endx, endy + offset],
1082 [endx + barb_height, endy - full_width / 2 + offset],
1083 [endx, endy - full_width + offset]])
1085 offset -= full_width + spacing
1087 # Add vertices for each barb. These really are lines, but works
1088 # great adding 3 vertices that basically pull the polygon out and
1089 # back down the line
1090 for i in range(nbarbs[index]):
1091 poly_verts.extend(
1092 [(endx, endy + offset),
1093 (endx + barb_height, endy + offset + full_width / 2),
1094 (endx, endy + offset)])
1096 offset -= spacing
1098 # Add the vertices for half a barb, if needed
1099 if half_barb[index]:
1100 # If the half barb is the first on the staff, traditionally it
1101 # is offset from the end to make it easy to distinguish from a
1102 # barb with a full one
1103 if offset == length:
1104 poly_verts.append((endx, endy + offset))
1105 offset -= 1.5 * spacing
1106 poly_verts.extend(
1107 [(endx, endy + offset),
1108 (endx + barb_height / 2, endy + offset + full_width / 4),
1109 (endx, endy + offset)])
1111 # Rotate the barb according the angle. Making the barb first and
1112 # then rotating it made the math for drawing the barb really easy.
1113 # Also, the transform framework makes doing the rotation simple.
1114 poly_verts = transforms.Affine2D().rotate(-angle).transform(
1115 poly_verts)
1116 barb_list.append(poly_verts)
1118 return barb_list
1120 def set_UVC(self, U, V, C=None):
1121 # We need to ensure we have a copy, not a reference to an array that
1122 # might change before draw().
1123 self.u = ma.masked_invalid(U, copy=True).ravel()
1124 self.v = ma.masked_invalid(V, copy=True).ravel()
1126 # Flip needs to have the same number of entries as everything else.
1127 # Use broadcast_to to avoid a bloated array of identical values.
1128 # (can't rely on actual broadcasting)
1129 if len(self.flip) == 1:
1130 flip = np.broadcast_to(self.flip, self.u.shape)
1131 else:
1132 flip = self.flip
1134 if C is not None:
1135 c = ma.masked_invalid(C, copy=True).ravel()
1136 x, y, u, v, c, flip = cbook.delete_masked_points(
1137 self.x.ravel(), self.y.ravel(), self.u, self.v, c,
1138 flip.ravel())
1139 _check_consistent_shapes(x, y, u, v, c, flip)
1140 else:
1141 x, y, u, v, flip = cbook.delete_masked_points(
1142 self.x.ravel(), self.y.ravel(), self.u, self.v, flip.ravel())
1143 _check_consistent_shapes(x, y, u, v, flip)
1145 magnitude = np.hypot(u, v)
1146 flags, barbs, halves, empty = self._find_tails(
1147 magnitude, self.rounding, **self.barb_increments)
1149 # Get the vertices for each of the barbs
1151 plot_barbs = self._make_barbs(u, v, flags, barbs, halves, empty,
1152 self._length, self._pivot, self.sizes,
1153 self.fill_empty, flip)
1154 self.set_verts(plot_barbs)
1156 # Set the color array
1157 if C is not None:
1158 self.set_array(c)
1160 # Update the offsets in case the masked data changed
1161 xy = np.column_stack((x, y))
1162 self._offsets = xy
1163 self.stale = True
1165 def set_offsets(self, xy):
1166 """
1167 Set the offsets for the barb polygons. This saves the offsets passed
1168 in and masks them as appropriate for the existing U/V data.
1170 Parameters
1171 ----------
1172 xy : sequence of pairs of floats
1173 """
1174 self.x = xy[:, 0]
1175 self.y = xy[:, 1]
1176 x, y, u, v = cbook.delete_masked_points(
1177 self.x.ravel(), self.y.ravel(), self.u, self.v)
1178 _check_consistent_shapes(x, y, u, v)
1179 xy = np.column_stack((x, y))
1180 super().set_offsets(xy)
1181 self.stale = True