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