comparison DEPENDENCIES/mingw32/Python27/Lib/site-packages/numpy/lib/tests/test_recfunctions.py @ 87:2a2c65a20a8b

Add Python libs and headers
author Chris Cannam
date Wed, 25 Feb 2015 14:05:22 +0000
parents
children
comparison
equal deleted inserted replaced
86:413a9d26189e 87:2a2c65a20a8b
1 from __future__ import division, absolute_import, print_function
2
3 import numpy as np
4 import numpy.ma as ma
5 from numpy.ma.mrecords import MaskedRecords
6 from numpy.ma.testutils import (
7 run_module_suite, TestCase, assert_, assert_equal
8 )
9 from numpy.lib.recfunctions import (
10 drop_fields, rename_fields, get_fieldstructure, recursive_fill_fields,
11 find_duplicates, merge_arrays, append_fields, stack_arrays, join_by
12 )
13 get_names = np.lib.recfunctions.get_names
14 get_names_flat = np.lib.recfunctions.get_names_flat
15 zip_descr = np.lib.recfunctions.zip_descr
16
17
18 class TestRecFunctions(TestCase):
19 # Misc tests
20
21 def setUp(self):
22 x = np.array([1, 2, ])
23 y = np.array([10, 20, 30])
24 z = np.array([('A', 1.), ('B', 2.)],
25 dtype=[('A', '|S3'), ('B', float)])
26 w = np.array([(1, (2, 3.0)), (4, (5, 6.0))],
27 dtype=[('a', int), ('b', [('ba', float), ('bb', int)])])
28 self.data = (w, x, y, z)
29
30 def test_zip_descr(self):
31 # Test zip_descr
32 (w, x, y, z) = self.data
33
34 # Std array
35 test = zip_descr((x, x), flatten=True)
36 assert_equal(test,
37 np.dtype([('', int), ('', int)]))
38 test = zip_descr((x, x), flatten=False)
39 assert_equal(test,
40 np.dtype([('', int), ('', int)]))
41
42 # Std & flexible-dtype
43 test = zip_descr((x, z), flatten=True)
44 assert_equal(test,
45 np.dtype([('', int), ('A', '|S3'), ('B', float)]))
46 test = zip_descr((x, z), flatten=False)
47 assert_equal(test,
48 np.dtype([('', int),
49 ('', [('A', '|S3'), ('B', float)])]))
50
51 # Standard & nested dtype
52 test = zip_descr((x, w), flatten=True)
53 assert_equal(test,
54 np.dtype([('', int),
55 ('a', int),
56 ('ba', float), ('bb', int)]))
57 test = zip_descr((x, w), flatten=False)
58 assert_equal(test,
59 np.dtype([('', int),
60 ('', [('a', int),
61 ('b', [('ba', float), ('bb', int)])])]))
62
63 def test_drop_fields(self):
64 # Test drop_fields
65 a = np.array([(1, (2, 3.0)), (4, (5, 6.0))],
66 dtype=[('a', int), ('b', [('ba', float), ('bb', int)])])
67
68 # A basic field
69 test = drop_fields(a, 'a')
70 control = np.array([((2, 3.0),), ((5, 6.0),)],
71 dtype=[('b', [('ba', float), ('bb', int)])])
72 assert_equal(test, control)
73
74 # Another basic field (but nesting two fields)
75 test = drop_fields(a, 'b')
76 control = np.array([(1,), (4,)], dtype=[('a', int)])
77 assert_equal(test, control)
78
79 # A nested sub-field
80 test = drop_fields(a, ['ba', ])
81 control = np.array([(1, (3.0,)), (4, (6.0,))],
82 dtype=[('a', int), ('b', [('bb', int)])])
83 assert_equal(test, control)
84
85 # All the nested sub-field from a field: zap that field
86 test = drop_fields(a, ['ba', 'bb'])
87 control = np.array([(1,), (4,)], dtype=[('a', int)])
88 assert_equal(test, control)
89
90 test = drop_fields(a, ['a', 'b'])
91 assert_(test is None)
92
93 def test_rename_fields(self):
94 # Test rename fields
95 a = np.array([(1, (2, [3.0, 30.])), (4, (5, [6.0, 60.]))],
96 dtype=[('a', int),
97 ('b', [('ba', float), ('bb', (float, 2))])])
98 test = rename_fields(a, {'a': 'A', 'bb': 'BB'})
99 newdtype = [('A', int), ('b', [('ba', float), ('BB', (float, 2))])]
100 control = a.view(newdtype)
101 assert_equal(test.dtype, newdtype)
102 assert_equal(test, control)
103
104 def test_get_names(self):
105 # Test get_names
106 ndtype = np.dtype([('A', '|S3'), ('B', float)])
107 test = get_names(ndtype)
108 assert_equal(test, ('A', 'B'))
109
110 ndtype = np.dtype([('a', int), ('b', [('ba', float), ('bb', int)])])
111 test = get_names(ndtype)
112 assert_equal(test, ('a', ('b', ('ba', 'bb'))))
113
114 def test_get_names_flat(self):
115 # Test get_names_flat
116 ndtype = np.dtype([('A', '|S3'), ('B', float)])
117 test = get_names_flat(ndtype)
118 assert_equal(test, ('A', 'B'))
119
120 ndtype = np.dtype([('a', int), ('b', [('ba', float), ('bb', int)])])
121 test = get_names_flat(ndtype)
122 assert_equal(test, ('a', 'b', 'ba', 'bb'))
123
124 def test_get_fieldstructure(self):
125 # Test get_fieldstructure
126
127 # No nested fields
128 ndtype = np.dtype([('A', '|S3'), ('B', float)])
129 test = get_fieldstructure(ndtype)
130 assert_equal(test, {'A': [], 'B': []})
131
132 # One 1-nested field
133 ndtype = np.dtype([('A', int), ('B', [('BA', float), ('BB', '|S1')])])
134 test = get_fieldstructure(ndtype)
135 assert_equal(test, {'A': [], 'B': [], 'BA': ['B', ], 'BB': ['B']})
136
137 # One 2-nested fields
138 ndtype = np.dtype([('A', int),
139 ('B', [('BA', int),
140 ('BB', [('BBA', int), ('BBB', int)])])])
141 test = get_fieldstructure(ndtype)
142 control = {'A': [], 'B': [], 'BA': ['B'], 'BB': ['B'],
143 'BBA': ['B', 'BB'], 'BBB': ['B', 'BB']}
144 assert_equal(test, control)
145
146 def test_find_duplicates(self):
147 # Test find_duplicates
148 a = ma.array([(2, (2., 'B')), (1, (2., 'B')), (2, (2., 'B')),
149 (1, (1., 'B')), (2, (2., 'B')), (2, (2., 'C'))],
150 mask=[(0, (0, 0)), (0, (0, 0)), (0, (0, 0)),
151 (0, (0, 0)), (1, (0, 0)), (0, (1, 0))],
152 dtype=[('A', int), ('B', [('BA', float), ('BB', '|S1')])])
153
154 test = find_duplicates(a, ignoremask=False, return_index=True)
155 control = [0, 2]
156 assert_equal(sorted(test[-1]), control)
157 assert_equal(test[0], a[test[-1]])
158
159 test = find_duplicates(a, key='A', return_index=True)
160 control = [0, 1, 2, 3, 5]
161 assert_equal(sorted(test[-1]), control)
162 assert_equal(test[0], a[test[-1]])
163
164 test = find_duplicates(a, key='B', return_index=True)
165 control = [0, 1, 2, 4]
166 assert_equal(sorted(test[-1]), control)
167 assert_equal(test[0], a[test[-1]])
168
169 test = find_duplicates(a, key='BA', return_index=True)
170 control = [0, 1, 2, 4]
171 assert_equal(sorted(test[-1]), control)
172 assert_equal(test[0], a[test[-1]])
173
174 test = find_duplicates(a, key='BB', return_index=True)
175 control = [0, 1, 2, 3, 4]
176 assert_equal(sorted(test[-1]), control)
177 assert_equal(test[0], a[test[-1]])
178
179 def test_find_duplicates_ignoremask(self):
180 # Test the ignoremask option of find_duplicates
181 ndtype = [('a', int)]
182 a = ma.array([1, 1, 1, 2, 2, 3, 3],
183 mask=[0, 0, 1, 0, 0, 0, 1]).view(ndtype)
184 test = find_duplicates(a, ignoremask=True, return_index=True)
185 control = [0, 1, 3, 4]
186 assert_equal(sorted(test[-1]), control)
187 assert_equal(test[0], a[test[-1]])
188
189 test = find_duplicates(a, ignoremask=False, return_index=True)
190 control = [0, 1, 2, 3, 4, 6]
191 assert_equal(sorted(test[-1]), control)
192 assert_equal(test[0], a[test[-1]])
193
194
195 class TestRecursiveFillFields(TestCase):
196 # Test recursive_fill_fields.
197 def test_simple_flexible(self):
198 # Test recursive_fill_fields on flexible-array
199 a = np.array([(1, 10.), (2, 20.)], dtype=[('A', int), ('B', float)])
200 b = np.zeros((3,), dtype=a.dtype)
201 test = recursive_fill_fields(a, b)
202 control = np.array([(1, 10.), (2, 20.), (0, 0.)],
203 dtype=[('A', int), ('B', float)])
204 assert_equal(test, control)
205
206 def test_masked_flexible(self):
207 # Test recursive_fill_fields on masked flexible-array
208 a = ma.array([(1, 10.), (2, 20.)], mask=[(0, 1), (1, 0)],
209 dtype=[('A', int), ('B', float)])
210 b = ma.zeros((3,), dtype=a.dtype)
211 test = recursive_fill_fields(a, b)
212 control = ma.array([(1, 10.), (2, 20.), (0, 0.)],
213 mask=[(0, 1), (1, 0), (0, 0)],
214 dtype=[('A', int), ('B', float)])
215 assert_equal(test, control)
216
217
218 class TestMergeArrays(TestCase):
219 # Test merge_arrays
220
221 def setUp(self):
222 x = np.array([1, 2, ])
223 y = np.array([10, 20, 30])
224 z = np.array(
225 [('A', 1.), ('B', 2.)], dtype=[('A', '|S3'), ('B', float)])
226 w = np.array(
227 [(1, (2, 3.0)), (4, (5, 6.0))],
228 dtype=[('a', int), ('b', [('ba', float), ('bb', int)])])
229 self.data = (w, x, y, z)
230
231 def test_solo(self):
232 # Test merge_arrays on a single array.
233 (_, x, _, z) = self.data
234
235 test = merge_arrays(x)
236 control = np.array([(1,), (2,)], dtype=[('f0', int)])
237 assert_equal(test, control)
238 test = merge_arrays((x,))
239 assert_equal(test, control)
240
241 test = merge_arrays(z, flatten=False)
242 assert_equal(test, z)
243 test = merge_arrays(z, flatten=True)
244 assert_equal(test, z)
245
246 def test_solo_w_flatten(self):
247 # Test merge_arrays on a single array w & w/o flattening
248 w = self.data[0]
249 test = merge_arrays(w, flatten=False)
250 assert_equal(test, w)
251
252 test = merge_arrays(w, flatten=True)
253 control = np.array([(1, 2, 3.0), (4, 5, 6.0)],
254 dtype=[('a', int), ('ba', float), ('bb', int)])
255 assert_equal(test, control)
256
257 def test_standard(self):
258 # Test standard & standard
259 # Test merge arrays
260 (_, x, y, _) = self.data
261 test = merge_arrays((x, y), usemask=False)
262 control = np.array([(1, 10), (2, 20), (-1, 30)],
263 dtype=[('f0', int), ('f1', int)])
264 assert_equal(test, control)
265
266 test = merge_arrays((x, y), usemask=True)
267 control = ma.array([(1, 10), (2, 20), (-1, 30)],
268 mask=[(0, 0), (0, 0), (1, 0)],
269 dtype=[('f0', int), ('f1', int)])
270 assert_equal(test, control)
271 assert_equal(test.mask, control.mask)
272
273 def test_flatten(self):
274 # Test standard & flexible
275 (_, x, _, z) = self.data
276 test = merge_arrays((x, z), flatten=True)
277 control = np.array([(1, 'A', 1.), (2, 'B', 2.)],
278 dtype=[('f0', int), ('A', '|S3'), ('B', float)])
279 assert_equal(test, control)
280
281 test = merge_arrays((x, z), flatten=False)
282 control = np.array([(1, ('A', 1.)), (2, ('B', 2.))],
283 dtype=[('f0', int),
284 ('f1', [('A', '|S3'), ('B', float)])])
285 assert_equal(test, control)
286
287 def test_flatten_wflexible(self):
288 # Test flatten standard & nested
289 (w, x, _, _) = self.data
290 test = merge_arrays((x, w), flatten=True)
291 control = np.array([(1, 1, 2, 3.0), (2, 4, 5, 6.0)],
292 dtype=[('f0', int),
293 ('a', int), ('ba', float), ('bb', int)])
294 assert_equal(test, control)
295
296 test = merge_arrays((x, w), flatten=False)
297 controldtype = [('f0', int),
298 ('f1', [('a', int),
299 ('b', [('ba', float), ('bb', int)])])]
300 control = np.array([(1., (1, (2, 3.0))), (2, (4, (5, 6.0)))],
301 dtype=controldtype)
302 assert_equal(test, control)
303
304 def test_wmasked_arrays(self):
305 # Test merge_arrays masked arrays
306 (_, x, _, _) = self.data
307 mx = ma.array([1, 2, 3], mask=[1, 0, 0])
308 test = merge_arrays((x, mx), usemask=True)
309 control = ma.array([(1, 1), (2, 2), (-1, 3)],
310 mask=[(0, 1), (0, 0), (1, 0)],
311 dtype=[('f0', int), ('f1', int)])
312 assert_equal(test, control)
313 test = merge_arrays((x, mx), usemask=True, asrecarray=True)
314 assert_equal(test, control)
315 assert_(isinstance(test, MaskedRecords))
316
317 def test_w_singlefield(self):
318 # Test single field
319 test = merge_arrays((np.array([1, 2]).view([('a', int)]),
320 np.array([10., 20., 30.])),)
321 control = ma.array([(1, 10.), (2, 20.), (-1, 30.)],
322 mask=[(0, 0), (0, 0), (1, 0)],
323 dtype=[('a', int), ('f1', float)])
324 assert_equal(test, control)
325
326 def test_w_shorter_flex(self):
327 # Test merge_arrays w/ a shorter flexndarray.
328 z = self.data[-1]
329
330 # Fixme, this test looks incomplete and broken
331 #test = merge_arrays((z, np.array([10, 20, 30]).view([('C', int)])))
332 #control = np.array([('A', 1., 10), ('B', 2., 20), ('-1', -1, 20)],
333 # dtype=[('A', '|S3'), ('B', float), ('C', int)])
334 #assert_equal(test, control)
335
336 # Hack to avoid pyflakes warnings about unused variables
337 merge_arrays((z, np.array([10, 20, 30]).view([('C', int)])))
338 np.array([('A', 1., 10), ('B', 2., 20), ('-1', -1, 20)],
339 dtype=[('A', '|S3'), ('B', float), ('C', int)])
340
341 def test_singlerecord(self):
342 (_, x, y, z) = self.data
343 test = merge_arrays((x[0], y[0], z[0]), usemask=False)
344 control = np.array([(1, 10, ('A', 1))],
345 dtype=[('f0', int),
346 ('f1', int),
347 ('f2', [('A', '|S3'), ('B', float)])])
348 assert_equal(test, control)
349
350
351 class TestAppendFields(TestCase):
352 # Test append_fields
353
354 def setUp(self):
355 x = np.array([1, 2, ])
356 y = np.array([10, 20, 30])
357 z = np.array(
358 [('A', 1.), ('B', 2.)], dtype=[('A', '|S3'), ('B', float)])
359 w = np.array([(1, (2, 3.0)), (4, (5, 6.0))],
360 dtype=[('a', int), ('b', [('ba', float), ('bb', int)])])
361 self.data = (w, x, y, z)
362
363 def test_append_single(self):
364 # Test simple case
365 (_, x, _, _) = self.data
366 test = append_fields(x, 'A', data=[10, 20, 30])
367 control = ma.array([(1, 10), (2, 20), (-1, 30)],
368 mask=[(0, 0), (0, 0), (1, 0)],
369 dtype=[('f0', int), ('A', int)],)
370 assert_equal(test, control)
371
372 def test_append_double(self):
373 # Test simple case
374 (_, x, _, _) = self.data
375 test = append_fields(x, ('A', 'B'), data=[[10, 20, 30], [100, 200]])
376 control = ma.array([(1, 10, 100), (2, 20, 200), (-1, 30, -1)],
377 mask=[(0, 0, 0), (0, 0, 0), (1, 0, 1)],
378 dtype=[('f0', int), ('A', int), ('B', int)],)
379 assert_equal(test, control)
380
381 def test_append_on_flex(self):
382 # Test append_fields on flexible type arrays
383 z = self.data[-1]
384 test = append_fields(z, 'C', data=[10, 20, 30])
385 control = ma.array([('A', 1., 10), ('B', 2., 20), (-1, -1., 30)],
386 mask=[(0, 0, 0), (0, 0, 0), (1, 1, 0)],
387 dtype=[('A', '|S3'), ('B', float), ('C', int)],)
388 assert_equal(test, control)
389
390 def test_append_on_nested(self):
391 # Test append_fields on nested fields
392 w = self.data[0]
393 test = append_fields(w, 'C', data=[10, 20, 30])
394 control = ma.array([(1, (2, 3.0), 10),
395 (4, (5, 6.0), 20),
396 (-1, (-1, -1.), 30)],
397 mask=[(
398 0, (0, 0), 0), (0, (0, 0), 0), (1, (1, 1), 0)],
399 dtype=[('a', int),
400 ('b', [('ba', float), ('bb', int)]),
401 ('C', int)],)
402 assert_equal(test, control)
403
404
405 class TestStackArrays(TestCase):
406 # Test stack_arrays
407 def setUp(self):
408 x = np.array([1, 2, ])
409 y = np.array([10, 20, 30])
410 z = np.array(
411 [('A', 1.), ('B', 2.)], dtype=[('A', '|S3'), ('B', float)])
412 w = np.array([(1, (2, 3.0)), (4, (5, 6.0))],
413 dtype=[('a', int), ('b', [('ba', float), ('bb', int)])])
414 self.data = (w, x, y, z)
415
416 def test_solo(self):
417 # Test stack_arrays on single arrays
418 (_, x, _, _) = self.data
419 test = stack_arrays((x,))
420 assert_equal(test, x)
421 self.assertTrue(test is x)
422
423 test = stack_arrays(x)
424 assert_equal(test, x)
425 self.assertTrue(test is x)
426
427 def test_unnamed_fields(self):
428 # Tests combinations of arrays w/o named fields
429 (_, x, y, _) = self.data
430
431 test = stack_arrays((x, x), usemask=False)
432 control = np.array([1, 2, 1, 2])
433 assert_equal(test, control)
434
435 test = stack_arrays((x, y), usemask=False)
436 control = np.array([1, 2, 10, 20, 30])
437 assert_equal(test, control)
438
439 test = stack_arrays((y, x), usemask=False)
440 control = np.array([10, 20, 30, 1, 2])
441 assert_equal(test, control)
442
443 def test_unnamed_and_named_fields(self):
444 # Test combination of arrays w/ & w/o named fields
445 (_, x, _, z) = self.data
446
447 test = stack_arrays((x, z))
448 control = ma.array([(1, -1, -1), (2, -1, -1),
449 (-1, 'A', 1), (-1, 'B', 2)],
450 mask=[(0, 1, 1), (0, 1, 1),
451 (1, 0, 0), (1, 0, 0)],
452 dtype=[('f0', int), ('A', '|S3'), ('B', float)])
453 assert_equal(test, control)
454 assert_equal(test.mask, control.mask)
455
456 test = stack_arrays((z, x))
457 control = ma.array([('A', 1, -1), ('B', 2, -1),
458 (-1, -1, 1), (-1, -1, 2), ],
459 mask=[(0, 0, 1), (0, 0, 1),
460 (1, 1, 0), (1, 1, 0)],
461 dtype=[('A', '|S3'), ('B', float), ('f2', int)])
462 assert_equal(test, control)
463 assert_equal(test.mask, control.mask)
464
465 test = stack_arrays((z, z, x))
466 control = ma.array([('A', 1, -1), ('B', 2, -1),
467 ('A', 1, -1), ('B', 2, -1),
468 (-1, -1, 1), (-1, -1, 2), ],
469 mask=[(0, 0, 1), (0, 0, 1),
470 (0, 0, 1), (0, 0, 1),
471 (1, 1, 0), (1, 1, 0)],
472 dtype=[('A', '|S3'), ('B', float), ('f2', int)])
473 assert_equal(test, control)
474
475 def test_matching_named_fields(self):
476 # Test combination of arrays w/ matching field names
477 (_, x, _, z) = self.data
478 zz = np.array([('a', 10., 100.), ('b', 20., 200.), ('c', 30., 300.)],
479 dtype=[('A', '|S3'), ('B', float), ('C', float)])
480 test = stack_arrays((z, zz))
481 control = ma.array([('A', 1, -1), ('B', 2, -1),
482 (
483 'a', 10., 100.), ('b', 20., 200.), ('c', 30., 300.)],
484 dtype=[('A', '|S3'), ('B', float), ('C', float)],
485 mask=[(0, 0, 1), (0, 0, 1),
486 (0, 0, 0), (0, 0, 0), (0, 0, 0)])
487 assert_equal(test, control)
488 assert_equal(test.mask, control.mask)
489
490 test = stack_arrays((z, zz, x))
491 ndtype = [('A', '|S3'), ('B', float), ('C', float), ('f3', int)]
492 control = ma.array([('A', 1, -1, -1), ('B', 2, -1, -1),
493 ('a', 10., 100., -1), ('b', 20., 200., -1),
494 ('c', 30., 300., -1),
495 (-1, -1, -1, 1), (-1, -1, -1, 2)],
496 dtype=ndtype,
497 mask=[(0, 0, 1, 1), (0, 0, 1, 1),
498 (0, 0, 0, 1), (0, 0, 0, 1), (0, 0, 0, 1),
499 (1, 1, 1, 0), (1, 1, 1, 0)])
500 assert_equal(test, control)
501 assert_equal(test.mask, control.mask)
502
503 def test_defaults(self):
504 # Test defaults: no exception raised if keys of defaults are not fields.
505 (_, _, _, z) = self.data
506 zz = np.array([('a', 10., 100.), ('b', 20., 200.), ('c', 30., 300.)],
507 dtype=[('A', '|S3'), ('B', float), ('C', float)])
508 defaults = {'A': '???', 'B': -999., 'C': -9999., 'D': -99999.}
509 test = stack_arrays((z, zz), defaults=defaults)
510 control = ma.array([('A', 1, -9999.), ('B', 2, -9999.),
511 (
512 'a', 10., 100.), ('b', 20., 200.), ('c', 30., 300.)],
513 dtype=[('A', '|S3'), ('B', float), ('C', float)],
514 mask=[(0, 0, 1), (0, 0, 1),
515 (0, 0, 0), (0, 0, 0), (0, 0, 0)])
516 assert_equal(test, control)
517 assert_equal(test.data, control.data)
518 assert_equal(test.mask, control.mask)
519
520 def test_autoconversion(self):
521 # Tests autoconversion
522 adtype = [('A', int), ('B', bool), ('C', float)]
523 a = ma.array([(1, 2, 3)], mask=[(0, 1, 0)], dtype=adtype)
524 bdtype = [('A', int), ('B', float), ('C', float)]
525 b = ma.array([(4, 5, 6)], dtype=bdtype)
526 control = ma.array([(1, 2, 3), (4, 5, 6)], mask=[(0, 1, 0), (0, 0, 0)],
527 dtype=bdtype)
528 test = stack_arrays((a, b), autoconvert=True)
529 assert_equal(test, control)
530 assert_equal(test.mask, control.mask)
531 try:
532 test = stack_arrays((a, b), autoconvert=False)
533 except TypeError:
534 pass
535 else:
536 raise AssertionError
537
538 def test_checktitles(self):
539 # Test using titles in the field names
540 adtype = [(('a', 'A'), int), (('b', 'B'), bool), (('c', 'C'), float)]
541 a = ma.array([(1, 2, 3)], mask=[(0, 1, 0)], dtype=adtype)
542 bdtype = [(('a', 'A'), int), (('b', 'B'), bool), (('c', 'C'), float)]
543 b = ma.array([(4, 5, 6)], dtype=bdtype)
544 test = stack_arrays((a, b))
545 control = ma.array([(1, 2, 3), (4, 5, 6)], mask=[(0, 1, 0), (0, 0, 0)],
546 dtype=bdtype)
547 assert_equal(test, control)
548 assert_equal(test.mask, control.mask)
549
550
551 class TestJoinBy(TestCase):
552 def setUp(self):
553 self.a = np.array(list(zip(np.arange(10), np.arange(50, 60),
554 np.arange(100, 110))),
555 dtype=[('a', int), ('b', int), ('c', int)])
556 self.b = np.array(list(zip(np.arange(5, 15), np.arange(65, 75),
557 np.arange(100, 110))),
558 dtype=[('a', int), ('b', int), ('d', int)])
559
560 def test_inner_join(self):
561 # Basic test of join_by
562 a, b = self.a, self.b
563
564 test = join_by('a', a, b, jointype='inner')
565 control = np.array([(5, 55, 65, 105, 100), (6, 56, 66, 106, 101),
566 (7, 57, 67, 107, 102), (8, 58, 68, 108, 103),
567 (9, 59, 69, 109, 104)],
568 dtype=[('a', int), ('b1', int), ('b2', int),
569 ('c', int), ('d', int)])
570 assert_equal(test, control)
571
572 def test_join(self):
573 a, b = self.a, self.b
574
575 # Fixme, this test is broken
576 #test = join_by(('a', 'b'), a, b)
577 #control = np.array([(5, 55, 105, 100), (6, 56, 106, 101),
578 # (7, 57, 107, 102), (8, 58, 108, 103),
579 # (9, 59, 109, 104)],
580 # dtype=[('a', int), ('b', int),
581 # ('c', int), ('d', int)])
582 #assert_equal(test, control)
583
584 # Hack to avoid pyflakes unused variable warnings
585 join_by(('a', 'b'), a, b)
586 np.array([(5, 55, 105, 100), (6, 56, 106, 101),
587 (7, 57, 107, 102), (8, 58, 108, 103),
588 (9, 59, 109, 104)],
589 dtype=[('a', int), ('b', int),
590 ('c', int), ('d', int)])
591
592 def test_outer_join(self):
593 a, b = self.a, self.b
594
595 test = join_by(('a', 'b'), a, b, 'outer')
596 control = ma.array([(0, 50, 100, -1), (1, 51, 101, -1),
597 (2, 52, 102, -1), (3, 53, 103, -1),
598 (4, 54, 104, -1), (5, 55, 105, -1),
599 (5, 65, -1, 100), (6, 56, 106, -1),
600 (6, 66, -1, 101), (7, 57, 107, -1),
601 (7, 67, -1, 102), (8, 58, 108, -1),
602 (8, 68, -1, 103), (9, 59, 109, -1),
603 (9, 69, -1, 104), (10, 70, -1, 105),
604 (11, 71, -1, 106), (12, 72, -1, 107),
605 (13, 73, -1, 108), (14, 74, -1, 109)],
606 mask=[(0, 0, 0, 1), (0, 0, 0, 1),
607 (0, 0, 0, 1), (0, 0, 0, 1),
608 (0, 0, 0, 1), (0, 0, 0, 1),
609 (0, 0, 1, 0), (0, 0, 0, 1),
610 (0, 0, 1, 0), (0, 0, 0, 1),
611 (0, 0, 1, 0), (0, 0, 0, 1),
612 (0, 0, 1, 0), (0, 0, 0, 1),
613 (0, 0, 1, 0), (0, 0, 1, 0),
614 (0, 0, 1, 0), (0, 0, 1, 0),
615 (0, 0, 1, 0), (0, 0, 1, 0)],
616 dtype=[('a', int), ('b', int),
617 ('c', int), ('d', int)])
618 assert_equal(test, control)
619
620 def test_leftouter_join(self):
621 a, b = self.a, self.b
622
623 test = join_by(('a', 'b'), a, b, 'leftouter')
624 control = ma.array([(0, 50, 100, -1), (1, 51, 101, -1),
625 (2, 52, 102, -1), (3, 53, 103, -1),
626 (4, 54, 104, -1), (5, 55, 105, -1),
627 (6, 56, 106, -1), (7, 57, 107, -1),
628 (8, 58, 108, -1), (9, 59, 109, -1)],
629 mask=[(0, 0, 0, 1), (0, 0, 0, 1),
630 (0, 0, 0, 1), (0, 0, 0, 1),
631 (0, 0, 0, 1), (0, 0, 0, 1),
632 (0, 0, 0, 1), (0, 0, 0, 1),
633 (0, 0, 0, 1), (0, 0, 0, 1)],
634 dtype=[('a', int), ('b', int), ('c', int), ('d', int)])
635 assert_equal(test, control)
636
637
638 class TestJoinBy2(TestCase):
639 @classmethod
640 def setUp(cls):
641 cls.a = np.array(list(zip(np.arange(10), np.arange(50, 60),
642 np.arange(100, 110))),
643 dtype=[('a', int), ('b', int), ('c', int)])
644 cls.b = np.array(list(zip(np.arange(10), np.arange(65, 75),
645 np.arange(100, 110))),
646 dtype=[('a', int), ('b', int), ('d', int)])
647
648 def test_no_r1postfix(self):
649 # Basic test of join_by no_r1postfix
650 a, b = self.a, self.b
651
652 test = join_by(
653 'a', a, b, r1postfix='', r2postfix='2', jointype='inner')
654 control = np.array([(0, 50, 65, 100, 100), (1, 51, 66, 101, 101),
655 (2, 52, 67, 102, 102), (3, 53, 68, 103, 103),
656 (4, 54, 69, 104, 104), (5, 55, 70, 105, 105),
657 (6, 56, 71, 106, 106), (7, 57, 72, 107, 107),
658 (8, 58, 73, 108, 108), (9, 59, 74, 109, 109)],
659 dtype=[('a', int), ('b', int), ('b2', int),
660 ('c', int), ('d', int)])
661 assert_equal(test, control)
662
663 def test_no_postfix(self):
664 self.assertRaises(ValueError, join_by, 'a', self.a, self.b,
665 r1postfix='', r2postfix='')
666
667 def test_no_r2postfix(self):
668 # Basic test of join_by no_r2postfix
669 a, b = self.a, self.b
670
671 test = join_by(
672 'a', a, b, r1postfix='1', r2postfix='', jointype='inner')
673 control = np.array([(0, 50, 65, 100, 100), (1, 51, 66, 101, 101),
674 (2, 52, 67, 102, 102), (3, 53, 68, 103, 103),
675 (4, 54, 69, 104, 104), (5, 55, 70, 105, 105),
676 (6, 56, 71, 106, 106), (7, 57, 72, 107, 107),
677 (8, 58, 73, 108, 108), (9, 59, 74, 109, 109)],
678 dtype=[('a', int), ('b1', int), ('b', int),
679 ('c', int), ('d', int)])
680 assert_equal(test, control)
681
682 def test_two_keys_two_vars(self):
683 a = np.array(list(zip(np.tile([10, 11], 5), np.repeat(np.arange(5), 2),
684 np.arange(50, 60), np.arange(10, 20))),
685 dtype=[('k', int), ('a', int), ('b', int), ('c', int)])
686
687 b = np.array(list(zip(np.tile([10, 11], 5), np.repeat(np.arange(5), 2),
688 np.arange(65, 75), np.arange(0, 10))),
689 dtype=[('k', int), ('a', int), ('b', int), ('c', int)])
690
691 control = np.array([(10, 0, 50, 65, 10, 0), (11, 0, 51, 66, 11, 1),
692 (10, 1, 52, 67, 12, 2), (11, 1, 53, 68, 13, 3),
693 (10, 2, 54, 69, 14, 4), (11, 2, 55, 70, 15, 5),
694 (10, 3, 56, 71, 16, 6), (11, 3, 57, 72, 17, 7),
695 (10, 4, 58, 73, 18, 8), (11, 4, 59, 74, 19, 9)],
696 dtype=[('k', int), ('a', int), ('b1', int),
697 ('b2', int), ('c1', int), ('c2', int)])
698 test = join_by(
699 ['a', 'k'], a, b, r1postfix='1', r2postfix='2', jointype='inner')
700 assert_equal(test.dtype, control.dtype)
701 assert_equal(test, control)
702
703
704 if __name__ == '__main__':
705 run_module_suite()