quiver.py 45 KB

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