test_quiver.py 7.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261
  1. import numpy as np
  2. import pytest
  3. import sys
  4. from matplotlib import pyplot as plt
  5. from matplotlib.testing.decorators import image_comparison
  6. def draw_quiver(ax, **kw):
  7. X, Y = np.meshgrid(np.arange(0, 2 * np.pi, 1),
  8. np.arange(0, 2 * np.pi, 1))
  9. U = np.cos(X)
  10. V = np.sin(Y)
  11. Q = ax.quiver(U, V, **kw)
  12. return Q
  13. def test_quiver_memory_leak():
  14. fig, ax = plt.subplots()
  15. Q = draw_quiver(ax)
  16. ttX = Q.X
  17. Q.remove()
  18. del Q
  19. assert sys.getrefcount(ttX) == 2
  20. def test_quiver_key_memory_leak():
  21. fig, ax = plt.subplots()
  22. Q = draw_quiver(ax)
  23. qk = ax.quiverkey(Q, 0.5, 0.92, 2, r'$2 \frac{m}{s}$',
  24. labelpos='W',
  25. fontproperties={'weight': 'bold'})
  26. assert sys.getrefcount(qk) == 3
  27. qk.remove()
  28. assert sys.getrefcount(qk) == 2
  29. def test_quiver_number_of_args():
  30. X = [1, 2]
  31. with pytest.raises(
  32. TypeError,
  33. match='takes 2-5 positional arguments but 1 were given'):
  34. plt.quiver(X)
  35. with pytest.raises(
  36. TypeError,
  37. match='takes 2-5 positional arguments but 6 were given'):
  38. plt.quiver(X, X, X, X, X, X)
  39. def test_quiver_arg_sizes():
  40. X2 = [1, 2]
  41. X3 = [1, 2, 3]
  42. with pytest.raises(
  43. ValueError, match=('X and Y must be the same size, but '
  44. 'X.size is 2 and Y.size is 3.')):
  45. plt.quiver(X2, X3, X2, X2)
  46. with pytest.raises(
  47. ValueError, match=('Argument U has a size 3 which does not match '
  48. '2, the number of arrow positions')):
  49. plt.quiver(X2, X2, X3, X2)
  50. with pytest.raises(
  51. ValueError, match=('Argument V has a size 3 which does not match '
  52. '2, the number of arrow positions')):
  53. plt.quiver(X2, X2, X2, X3)
  54. with pytest.raises(
  55. ValueError, match=('Argument C has a size 3 which does not match '
  56. '2, the number of arrow positions')):
  57. plt.quiver(X2, X2, X2, X2, X3)
  58. def test_no_warnings():
  59. fig, ax = plt.subplots()
  60. X, Y = np.meshgrid(np.arange(15), np.arange(10))
  61. U = V = np.ones_like(X)
  62. phi = (np.random.rand(15, 10) - .5) * 150
  63. ax.quiver(X, Y, U, V, angles=phi)
  64. fig.canvas.draw() # Check that no warning is emitted.
  65. def test_zero_headlength():
  66. # Based on report by Doug McNeil:
  67. # http://matplotlib.1069221.n5.nabble.com/quiver-warnings-td28107.html
  68. fig, ax = plt.subplots()
  69. X, Y = np.meshgrid(np.arange(10), np.arange(10))
  70. U, V = np.cos(X), np.sin(Y)
  71. ax.quiver(U, V, headlength=0, headaxislength=0)
  72. fig.canvas.draw() # Check that no warning is emitted.
  73. @image_comparison(['quiver_animated_test_image.png'])
  74. def test_quiver_animate():
  75. # Tests fix for #2616
  76. fig, ax = plt.subplots()
  77. Q = draw_quiver(ax, animated=True)
  78. ax.quiverkey(Q, 0.5, 0.92, 2, r'$2 \frac{m}{s}$',
  79. labelpos='W', fontproperties={'weight': 'bold'})
  80. @image_comparison(['quiver_with_key_test_image.png'])
  81. def test_quiver_with_key():
  82. fig, ax = plt.subplots()
  83. ax.margins(0.1)
  84. Q = draw_quiver(ax)
  85. ax.quiverkey(Q, 0.5, 0.95, 2,
  86. r'$2\, \mathrm{m}\, \mathrm{s}^{-1}$',
  87. angle=-10,
  88. coordinates='figure',
  89. labelpos='W',
  90. fontproperties={'weight': 'bold', 'size': 'large'})
  91. @image_comparison(['quiver_single_test_image.png'], remove_text=True)
  92. def test_quiver_single():
  93. fig, ax = plt.subplots()
  94. ax.margins(0.1)
  95. ax.quiver([1], [1], [2], [2])
  96. def test_quiver_copy():
  97. fig, ax = plt.subplots()
  98. uv = dict(u=np.array([1.1]), v=np.array([2.0]))
  99. q0 = ax.quiver([1], [1], uv['u'], uv['v'])
  100. uv['v'][0] = 0
  101. assert q0.V[0] == 2.0
  102. @image_comparison(['quiver_key_pivot.png'], remove_text=True)
  103. def test_quiver_key_pivot():
  104. fig, ax = plt.subplots()
  105. u, v = np.mgrid[0:2*np.pi:10j, 0:2*np.pi:10j]
  106. q = ax.quiver(np.sin(u), np.cos(v))
  107. ax.set_xlim(-2, 11)
  108. ax.set_ylim(-2, 11)
  109. ax.quiverkey(q, 0.5, 1, 1, 'N', labelpos='N')
  110. ax.quiverkey(q, 1, 0.5, 1, 'E', labelpos='E')
  111. ax.quiverkey(q, 0.5, 0, 1, 'S', labelpos='S')
  112. ax.quiverkey(q, 0, 0.5, 1, 'W', labelpos='W')
  113. @image_comparison(['quiver_key_xy.png'], remove_text=True)
  114. def test_quiver_key_xy():
  115. # With scale_units='xy', ensure quiverkey still matches its quiver.
  116. # Note that the quiver and quiverkey lengths depend on the axes aspect
  117. # ratio, and that with angles='xy' their angles also depend on the axes
  118. # aspect ratio.
  119. X = np.arange(8)
  120. Y = np.zeros(8)
  121. angles = X * (np.pi / 4)
  122. uv = np.exp(1j * angles)
  123. U = uv.real
  124. V = uv.imag
  125. fig, axs = plt.subplots(2)
  126. for ax, angle_str in zip(axs, ('uv', 'xy')):
  127. ax.set_xlim(-1, 8)
  128. ax.set_ylim(-0.2, 0.2)
  129. q = ax.quiver(X, Y, U, V, pivot='middle',
  130. units='xy', width=0.05,
  131. scale=2, scale_units='xy',
  132. angles=angle_str)
  133. for x, angle in zip((0.2, 0.5, 0.8), (0, 45, 90)):
  134. ax.quiverkey(q, X=x, Y=0.8, U=1, angle=angle, label='', color='b')
  135. @image_comparison(['barbs_test_image.png'], remove_text=True)
  136. def test_barbs():
  137. x = np.linspace(-5, 5, 5)
  138. X, Y = np.meshgrid(x, x)
  139. U, V = 12*X, 12*Y
  140. fig, ax = plt.subplots()
  141. ax.barbs(X, Y, U, V, np.hypot(U, V), fill_empty=True, rounding=False,
  142. sizes=dict(emptybarb=0.25, spacing=0.2, height=0.3),
  143. cmap='viridis')
  144. @image_comparison(['barbs_pivot_test_image.png'], remove_text=True)
  145. def test_barbs_pivot():
  146. x = np.linspace(-5, 5, 5)
  147. X, Y = np.meshgrid(x, x)
  148. U, V = 12*X, 12*Y
  149. fig, ax = plt.subplots()
  150. ax.barbs(X, Y, U, V, fill_empty=True, rounding=False, pivot=1.7,
  151. sizes=dict(emptybarb=0.25, spacing=0.2, height=0.3))
  152. ax.scatter(X, Y, s=49, c='black')
  153. @image_comparison(['barbs_test_flip.png'], remove_text=True)
  154. def test_barbs_flip():
  155. """Test barbs with an array for flip_barb."""
  156. x = np.linspace(-5, 5, 5)
  157. X, Y = np.meshgrid(x, x)
  158. U, V = 12*X, 12*Y
  159. fig, ax = plt.subplots()
  160. ax.barbs(X, Y, U, V, fill_empty=True, rounding=False, pivot=1.7,
  161. sizes=dict(emptybarb=0.25, spacing=0.2, height=0.3),
  162. flip_barb=Y < 0)
  163. def test_bad_masked_sizes():
  164. """Test error handling when given differing sized masked arrays."""
  165. x = np.arange(3)
  166. y = np.arange(3)
  167. u = np.ma.array(15. * np.ones((4,)))
  168. v = np.ma.array(15. * np.ones_like(u))
  169. u[1] = np.ma.masked
  170. v[1] = np.ma.masked
  171. fig, ax = plt.subplots()
  172. with pytest.raises(ValueError):
  173. ax.barbs(x, y, u, v)
  174. def test_angles_and_scale():
  175. # angles array + scale_units kwarg
  176. fig, ax = plt.subplots()
  177. X, Y = np.meshgrid(np.arange(15), np.arange(10))
  178. U = V = np.ones_like(X)
  179. phi = (np.random.rand(15, 10) - .5) * 150
  180. ax.quiver(X, Y, U, V, angles=phi, scale_units='xy')
  181. @image_comparison(['quiver_xy.png'], remove_text=True)
  182. def test_quiver_xy():
  183. # simple arrow pointing from SW to NE
  184. fig, ax = plt.subplots(subplot_kw=dict(aspect='equal'))
  185. ax.quiver(0, 0, 1, 1, angles='xy', scale_units='xy', scale=1)
  186. ax.set_xlim(0, 1.1)
  187. ax.set_ylim(0, 1.1)
  188. ax.grid()
  189. def test_quiverkey_angles():
  190. # Check that only a single arrow is plotted for a quiverkey when an array
  191. # of angles is given to the original quiver plot
  192. fig, ax = plt.subplots()
  193. X, Y = np.meshgrid(np.arange(2), np.arange(2))
  194. U = V = angles = np.ones_like(X)
  195. q = ax.quiver(X, Y, U, V, angles=angles)
  196. qk = ax.quiverkey(q, 1, 1, 2, 'Label')
  197. # The arrows are only created when the key is drawn
  198. fig.canvas.draw()
  199. assert len(qk.verts) == 1
  200. def test_quiver_setuvc_numbers():
  201. """Check that it is possible to set all arrow UVC to the same numbers"""
  202. fig, ax = plt.subplots()
  203. X, Y = np.meshgrid(np.arange(2), np.arange(2))
  204. U = V = np.ones_like(X)
  205. q = ax.quiver(X, Y, U, V)
  206. q.set_UVC(0, 1)