Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 540a77d

Browse files
toddrjentoddrjen
committed
add collections test file with EventCollection tests
1 parent 9e908d9 commit 540a77d

File tree

1 file changed

+392
-0
lines changed

1 file changed

+392
-0
lines changed

‎lib/matplotlib/tests/test_collections.py

Lines changed: 392 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,392 @@
1+
"""
2+
Tests specific to the collections module.
3+
"""
4+
5+
import nose.tools
6+
import numpy as np
7+
import matplotlib.pyplot as plt
8+
from matplotlib.collections import EventCollection
9+
from matplotlib.testing.decorators import cleanup, image_comparison
10+
11+
12+
def generate_EventCollection_plot():
13+
'''
14+
generate the initial collection and plot it
15+
'''
16+
positions = np.array([0., 1., 2., 3., 5., 8., 13., 21.])
17+
extra_positions = np.array([34., 55., 89.])
18+
orientation = 'horizontal'
19+
lineoffset = 1
20+
linelength = .5
21+
linewidth = 2
22+
color = [1, 0, 0, 1]
23+
linestyle = 'solid'
24+
antialiased = True
25+
26+
coll = EventCollection(positions,
27+
orientation=orientation,
28+
lineoffset=lineoffset,
29+
linelength=linelength,
30+
linewidth=linewidth,
31+
color=color,
32+
linestyle=linestyle,
33+
antialiased=antialiased
34+
)
35+
36+
fig = plt.figure()
37+
splt = fig.add_subplot(1, 1, 1)
38+
splt.add_collection(coll)
39+
splt.set_title('EventCollection: default')
40+
props = {'positions': positions,
41+
'extra_positions': extra_positions,
42+
'orientation': orientation,
43+
'lineoffset': lineoffset,
44+
'linelength': linelength,
45+
'linewidth': linewidth,
46+
'color': color,
47+
'linestyle': linestyle,
48+
'antialiased': antialiased
49+
}
50+
splt.set_xlim(-1, 22)
51+
splt.set_ylim(0, 2)
52+
return splt, coll, props
53+
54+
55+
@image_comparison(baseline_images=['EventCollection_plot__default'])
56+
def test__EventCollection__get_segments():
57+
'''
58+
check to make sure the default segments have the correct coordinates
59+
'''
60+
_, coll, props = generate_EventCollection_plot()
61+
check_segments(coll,
62+
props['positions'],
63+
props['linelength'],
64+
props['lineoffset'],
65+
props['orientation'])
66+
67+
68+
@cleanup
69+
def test__EventCollection__get_positions():
70+
'''
71+
check to make sure the default positions match the input positions
72+
'''
73+
_, coll, props = generate_EventCollection_plot()
74+
np.testing.assert_array_equal(props['positions'], coll.get_positions())
75+
76+
77+
@cleanup
78+
def test__EventCollection__get_orientation():
79+
'''
80+
check to make sure the default orientation matches the input
81+
orientation
82+
'''
83+
_, coll, props = generate_EventCollection_plot()
84+
nose.tools.assert_equal(props['orientation'], coll.get_orientation())
85+
86+
87+
@cleanup
88+
def test__EventCollection__is_horizontal():
89+
'''
90+
check to make sure the default orientation matches the input
91+
orientation
92+
'''
93+
_, coll, _ = generate_EventCollection_plot()
94+
nose.tools.assert_equal(True, coll.is_horizontal())
95+
96+
97+
@cleanup
98+
def test__EventCollection__get_linelength():
99+
'''
100+
check to make sure the default linelength matches the input linelength
101+
'''
102+
_, coll, props = generate_EventCollection_plot()
103+
nose.tools.assert_equal(props['linelength'], coll.get_linelength())
104+
105+
106+
@cleanup
107+
def test__EventCollection__get_lineoffset():
108+
'''
109+
check to make sure the default lineoffset matches the input lineoffset
110+
'''
111+
_, coll, props = generate_EventCollection_plot()
112+
nose.tools.assert_equal(props['lineoffset'], coll.get_lineoffset())
113+
114+
115+
@cleanup
116+
def test__EventCollection__get_linestyle():
117+
'''
118+
check to make sure the default linestyle matches the input linestyle
119+
'''
120+
_, coll, _ = generate_EventCollection_plot()
121+
nose.tools.assert_equal(coll.get_linestyle(), [(None, None)])
122+
123+
124+
@cleanup
125+
def test__EventCollection__get_color():
126+
'''
127+
check to make sure the default color matches the input color
128+
'''
129+
_, coll, props = generate_EventCollection_plot()
130+
np.testing.assert_array_equal(props['color'], coll.get_color())
131+
check_allprop_array(coll.get_colors(), props['color'])
132+
133+
134+
@image_comparison(baseline_images=['EventCollection_plot__set_positions'])
135+
def test__EventCollection__set_positions():
136+
'''
137+
check to make sure set_positions works properly
138+
'''
139+
splt, coll, props = generate_EventCollection_plot()
140+
new_positions = np.hstack([props['positions'], props['extra_positions']])
141+
coll.set_positions(new_positions)
142+
np.testing.assert_array_equal(new_positions, coll.get_positions())
143+
check_segments(coll, new_positions,
144+
props['linelength'],
145+
props['lineoffset'],
146+
props['orientation'])
147+
splt.set_title('EventCollection: set_positions')
148+
splt.set_xlim(-1, 90)
149+
150+
151+
@image_comparison(baseline_images=['EventCollection_plot__add_positions'])
152+
def test__EventCollection__add_positions():
153+
'''
154+
check to make sure add_positions works properly
155+
'''
156+
splt, coll, props = generate_EventCollection_plot()
157+
new_positions = np.hstack([props['positions'],
158+
props['extra_positions'][0]])
159+
coll.add_positions(props['extra_positions'][0])
160+
np.testing.assert_array_equal(new_positions, coll.get_positions())
161+
check_segments(coll,
162+
new_positions,
163+
props['linelength'],
164+
props['lineoffset'],
165+
props['orientation'])
166+
splt.set_title('EventCollection: add_positions')
167+
splt.set_xlim(-1, 35)
168+
169+
170+
@image_comparison(baseline_images=['EventCollection_plot__append_positions'])
171+
def test__EventCollection__append_positions():
172+
'''
173+
check to make sure append_positions works properly
174+
'''
175+
splt, coll, props = generate_EventCollection_plot()
176+
new_positions = np.hstack([props['positions'],
177+
props['extra_positions'][2]])
178+
coll.append_positions(props['extra_positions'][2])
179+
np.testing.assert_array_equal(new_positions, coll.get_positions())
180+
check_segments(coll,
181+
new_positions,
182+
props['linelength'],
183+
props['lineoffset'],
184+
props['orientation'])
185+
splt.set_title('EventCollection: append_positions')
186+
splt.set_xlim(-1, 90)
187+
188+
189+
@image_comparison(baseline_images=['EventCollection_plot__extend_positions'])
190+
def test__EventCollection__extend_positions():
191+
'''
192+
check to make sure extend_positions works properly
193+
'''
194+
splt, coll, props = generate_EventCollection_plot()
195+
new_positions = np.hstack([props['positions'],
196+
props['extra_positions'][1:]])
197+
coll.extend_positions(props['extra_positions'][1:])
198+
np.testing.assert_array_equal(new_positions, coll.get_positions())
199+
check_segments(coll,
200+
new_positions,
201+
props['linelength'],
202+
props['lineoffset'],
203+
props['orientation'])
204+
splt.set_title('EventCollection: extend_positions')
205+
splt.set_xlim(-1, 90)
206+
207+
208+
@image_comparison(baseline_images=['EventCollection_plot__switch_orientation'])
209+
def test__EventCollection__switch_orientation():
210+
'''
211+
check to make sure switch_orientation works properly
212+
'''
213+
splt, coll, props = generate_EventCollection_plot()
214+
new_orientation = 'vertical'
215+
coll.switch_orientation()
216+
nose.tools.assert_equal(new_orientation, coll.get_orientation())
217+
nose.tools.assert_equal(False, coll.is_horizontal())
218+
new_positions = coll.get_positions()
219+
check_segments(coll,
220+
new_positions,
221+
props['linelength'],
222+
props['lineoffset'], new_orientation)
223+
splt.set_title('EventCollection: switch_orientation')
224+
splt.set_ylim(-1, 22)
225+
splt.set_xlim(0, 2)
226+
227+
228+
@image_comparison(baseline_images=
229+
['EventCollection_plot__switch_orientation__2x'])
230+
def test__EventCollection__switch_orientation_2x():
231+
'''
232+
check to make sure calling switch_orientation twice sets the
233+
orientation back to the default
234+
'''
235+
splt, coll, props = generate_EventCollection_plot()
236+
coll.switch_orientation()
237+
coll.switch_orientation()
238+
new_positions = coll.get_positions()
239+
nose.tools.assert_equal(props['orientation'], coll.get_orientation())
240+
nose.tools.assert_equal(True, coll.is_horizontal())
241+
np.testing.assert_array_equal(props['positions'], new_positions)
242+
check_segments(coll,
243+
new_positions,
244+
props['linelength'],
245+
props['lineoffset'],
246+
props['orientation'])
247+
splt.set_title('EventCollection: switch_orientation 2x')
248+
249+
250+
@image_comparison(baseline_images=['EventCollection_plot__set_orientation'])
251+
def test__EventCollection__set_orientation():
252+
'''
253+
check to make sure set_orientation works properly
254+
'''
255+
splt, coll, props = generate_EventCollection_plot()
256+
new_orientation = 'vertical'
257+
coll.set_orientation(new_orientation)
258+
nose.tools.assert_equal(new_orientation, coll.get_orientation())
259+
nose.tools.assert_equal(False, coll.is_horizontal())
260+
check_segments(coll,
261+
props['positions'],
262+
props['linelength'],
263+
props['lineoffset'],
264+
new_orientation)
265+
splt.set_title('EventCollection: set_orientation')
266+
splt.set_ylim(-1, 22)
267+
splt.set_xlim(0, 2)
268+
269+
270+
@image_comparison(baseline_images=['EventCollection_plot__set_linelength'])
271+
def test__EventCollection__set_linelength():
272+
'''
273+
check to make sure set_linelength works properly
274+
'''
275+
splt, coll, props = generate_EventCollection_plot()
276+
new_linelength = 15
277+
coll.set_linelength(new_linelength)
278+
nose.tools.assert_equal(new_linelength, coll.get_linelength())
279+
check_segments(coll,
280+
props['positions'],
281+
new_linelength,
282+
props['lineoffset'],
283+
props['orientation'])
284+
splt.set_title('EventCollection: set_linelength')
285+
splt.set_ylim(-20, 20)
286+
287+
288+
@image_comparison(baseline_images=['EventCollection_plot__set_lineoffset'])
289+
def test__EventCollection__set_lineoffset():
290+
'''
291+
check to make sure set_lineoffset works properly
292+
'''
293+
splt, coll, props = generate_EventCollection_plot()
294+
new_lineoffset = -5.
295+
coll.set_lineoffset(new_lineoffset)
296+
nose.tools.assert_equal(new_lineoffset, coll.get_lineoffset())
297+
check_segments(coll,
298+
props['positions'],
299+
props['linelength'],
300+
new_lineoffset,
301+
props['orientation'])
302+
splt.set_title('EventCollection: set_lineoffset')
303+
splt.set_ylim(-6, -4)
304+
305+
306+
@image_comparison(baseline_images=['EventCollection_plot__set_linestyle'])
307+
def test__EventCollection__set_linestyle():
308+
'''
309+
check to make sure set_linestyle works properly
310+
'''
311+
splt, coll, _ = generate_EventCollection_plot()
312+
new_linestyle = 'dashed'
313+
coll.set_linestyle(new_linestyle)
314+
nose.tools.assert_equal(coll.get_linestyle(), [(0, (6.0, 6.0))])
315+
splt.set_title('EventCollection: set_linestyle')
316+
317+
318+
@image_comparison(baseline_images=['EventCollection_plot__set_linewidth'])
319+
def test__EventCollection__set_linewidth():
320+
'''
321+
check to make sure set_linestyle works properly
322+
'''
323+
splt, coll, _ = generate_EventCollection_plot()
324+
new_linewidth = 5
325+
coll.set_linewidth(new_linewidth)
326+
nose.tools.assert_equal(coll.get_linewidth(), new_linewidth)
327+
splt.set_title('EventCollection: set_linewidth')
328+
329+
330+
@image_comparison(baseline_images=['EventCollection_plot__set_color'])
331+
def test__EventCollection__set_color():
332+
'''
333+
check to make sure set_color works properly
334+
'''
335+
splt, coll, _ = generate_EventCollection_plot()
336+
new_color = np.array([0, 1, 1, 1])
337+
coll.set_color(new_color)
338+
np.testing.assert_array_equal(new_color, coll.get_color())
339+
check_allprop_array(coll.get_colors(), new_color)
340+
splt.set_title('EventCollection: set_color')
341+
342+
343+
def check_segments(coll, positions, linelength, lineoffset, orientation):
344+
'''
345+
check to make sure all values in the segment are correct, given a
346+
particular set of inputs
347+
348+
note: this is not a test, it is used by tests
349+
'''
350+
segments = coll.get_segments()
351+
if (orientation.lower() == 'horizontal'
352+
or orientation.lower() == 'none' or orientation is None):
353+
# if horizontal, the position in is in the y-axis
354+
pos1 = 1
355+
pos2 = 0
356+
elif orientation.lower() == 'vertical':
357+
# if vertical, the position in is in the x-axis
358+
pos1 = 0
359+
pos2 = 1
360+
else:
361+
raise ValueError("orientation must be 'horizontal' or 'vertical'")
362+
363+
# test to make sure each segment is correct
364+
for i, segment in enumerate(segments):
365+
nose.tools.assert_equal(segment[0, pos1], lineoffset + linelength / 2.)
366+
nose.tools.assert_equal(segment[1, pos1], lineoffset - linelength / 2.)
367+
nose.tools.assert_equal(segment[0, pos2], positions[i])
368+
nose.tools.assert_equal(segment[1, pos2], positions[i])
369+
370+
371+
def check_allprop(values, target):
372+
'''
373+
check to make sure all values match the given target
374+
375+
note: this is not a test, it is used by tests
376+
'''
377+
for value in values:
378+
nose.tools.assert_equal(value, target)
379+
380+
381+
def check_allprop_array(values, target):
382+
'''
383+
check to make sure all values match the given target if arrays
384+
385+
note: this is not a test, it is used by tests
386+
'''
387+
for value in values:
388+
np.testing.assert_array_equal(value, target)
389+
390+
if __name__ == '_main_':
391+
import nose
392+
nose.runmodule(argv=['-s', '--with-doctest'], exit=False)

0 commit comments

Comments
 (0)