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