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