quiver.py 48 KB

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