Chris@87
|
1 """Test functions for 1D array set operations.
|
Chris@87
|
2
|
Chris@87
|
3 """
|
Chris@87
|
4 from __future__ import division, absolute_import, print_function
|
Chris@87
|
5
|
Chris@87
|
6 import numpy as np
|
Chris@87
|
7 from numpy.testing import (
|
Chris@87
|
8 run_module_suite, TestCase, assert_array_equal
|
Chris@87
|
9 )
|
Chris@87
|
10 from numpy.lib.arraysetops import (
|
Chris@87
|
11 ediff1d, intersect1d, setxor1d, union1d, setdiff1d, unique, in1d
|
Chris@87
|
12 )
|
Chris@87
|
13
|
Chris@87
|
14
|
Chris@87
|
15 class TestSetOps(TestCase):
|
Chris@87
|
16
|
Chris@87
|
17 def test_unique(self):
|
Chris@87
|
18
|
Chris@87
|
19 def check_all(a, b, i1, i2, c, dt):
|
Chris@87
|
20 base_msg = 'check {0} failed for type {1}'
|
Chris@87
|
21
|
Chris@87
|
22 msg = base_msg.format('values', dt)
|
Chris@87
|
23 v = unique(a)
|
Chris@87
|
24 assert_array_equal(v, b, msg)
|
Chris@87
|
25
|
Chris@87
|
26 msg = base_msg.format('return_index', dt)
|
Chris@87
|
27 v, j = unique(a, 1, 0, 0)
|
Chris@87
|
28 assert_array_equal(v, b, msg)
|
Chris@87
|
29 assert_array_equal(j, i1, msg)
|
Chris@87
|
30
|
Chris@87
|
31 msg = base_msg.format('return_inverse', dt)
|
Chris@87
|
32 v, j = unique(a, 0, 1, 0)
|
Chris@87
|
33 assert_array_equal(v, b, msg)
|
Chris@87
|
34 assert_array_equal(j, i2, msg)
|
Chris@87
|
35
|
Chris@87
|
36 msg = base_msg.format('return_counts', dt)
|
Chris@87
|
37 v, j = unique(a, 0, 0, 1)
|
Chris@87
|
38 assert_array_equal(v, b, msg)
|
Chris@87
|
39 assert_array_equal(j, c, msg)
|
Chris@87
|
40
|
Chris@87
|
41 msg = base_msg.format('return_index and return_inverse', dt)
|
Chris@87
|
42 v, j1, j2 = unique(a, 1, 1, 0)
|
Chris@87
|
43 assert_array_equal(v, b, msg)
|
Chris@87
|
44 assert_array_equal(j1, i1, msg)
|
Chris@87
|
45 assert_array_equal(j2, i2, msg)
|
Chris@87
|
46
|
Chris@87
|
47 msg = base_msg.format('return_index and return_counts', dt)
|
Chris@87
|
48 v, j1, j2 = unique(a, 1, 0, 1)
|
Chris@87
|
49 assert_array_equal(v, b, msg)
|
Chris@87
|
50 assert_array_equal(j1, i1, msg)
|
Chris@87
|
51 assert_array_equal(j2, c, msg)
|
Chris@87
|
52
|
Chris@87
|
53 msg = base_msg.format('return_inverse and return_counts', dt)
|
Chris@87
|
54 v, j1, j2 = unique(a, 0, 1, 1)
|
Chris@87
|
55 assert_array_equal(v, b, msg)
|
Chris@87
|
56 assert_array_equal(j1, i2, msg)
|
Chris@87
|
57 assert_array_equal(j2, c, msg)
|
Chris@87
|
58
|
Chris@87
|
59 msg = base_msg.format(('return_index, return_inverse '
|
Chris@87
|
60 'and return_counts'), dt)
|
Chris@87
|
61 v, j1, j2, j3 = unique(a, 1, 1, 1)
|
Chris@87
|
62 assert_array_equal(v, b, msg)
|
Chris@87
|
63 assert_array_equal(j1, i1, msg)
|
Chris@87
|
64 assert_array_equal(j2, i2, msg)
|
Chris@87
|
65 assert_array_equal(j3, c, msg)
|
Chris@87
|
66
|
Chris@87
|
67 a = [5, 7, 1, 2, 1, 5, 7]*10
|
Chris@87
|
68 b = [1, 2, 5, 7]
|
Chris@87
|
69 i1 = [2, 3, 0, 1]
|
Chris@87
|
70 i2 = [2, 3, 0, 1, 0, 2, 3]*10
|
Chris@87
|
71 c = np.multiply([2, 1, 2, 2], 10)
|
Chris@87
|
72
|
Chris@87
|
73 # test for numeric arrays
|
Chris@87
|
74 types = []
|
Chris@87
|
75 types.extend(np.typecodes['AllInteger'])
|
Chris@87
|
76 types.extend(np.typecodes['AllFloat'])
|
Chris@87
|
77 types.append('datetime64[D]')
|
Chris@87
|
78 types.append('timedelta64[D]')
|
Chris@87
|
79 for dt in types:
|
Chris@87
|
80 aa = np.array(a, dt)
|
Chris@87
|
81 bb = np.array(b, dt)
|
Chris@87
|
82 check_all(aa, bb, i1, i2, c, dt)
|
Chris@87
|
83
|
Chris@87
|
84 # test for object arrays
|
Chris@87
|
85 dt = 'O'
|
Chris@87
|
86 aa = np.empty(len(a), dt)
|
Chris@87
|
87 aa[:] = a
|
Chris@87
|
88 bb = np.empty(len(b), dt)
|
Chris@87
|
89 bb[:] = b
|
Chris@87
|
90 check_all(aa, bb, i1, i2, c, dt)
|
Chris@87
|
91
|
Chris@87
|
92 # test for structured arrays
|
Chris@87
|
93 dt = [('', 'i'), ('', 'i')]
|
Chris@87
|
94 aa = np.array(list(zip(a, a)), dt)
|
Chris@87
|
95 bb = np.array(list(zip(b, b)), dt)
|
Chris@87
|
96 check_all(aa, bb, i1, i2, c, dt)
|
Chris@87
|
97
|
Chris@87
|
98 # test for ticket #2799
|
Chris@87
|
99 aa = [1. + 0.j, 1 - 1.j, 1]
|
Chris@87
|
100 assert_array_equal(np.unique(aa), [1. - 1.j, 1. + 0.j])
|
Chris@87
|
101
|
Chris@87
|
102 # test for ticket #4785
|
Chris@87
|
103 a = [(1, 2), (1, 2), (2, 3)]
|
Chris@87
|
104 unq = [1, 2, 3]
|
Chris@87
|
105 inv = [0, 1, 0, 1, 1, 2]
|
Chris@87
|
106 a1 = unique(a)
|
Chris@87
|
107 assert_array_equal(a1, unq)
|
Chris@87
|
108 a2, a2_inv = unique(a, return_inverse=True)
|
Chris@87
|
109 assert_array_equal(a2, unq)
|
Chris@87
|
110 assert_array_equal(a2_inv, inv)
|
Chris@87
|
111
|
Chris@87
|
112 def test_intersect1d(self):
|
Chris@87
|
113 # unique inputs
|
Chris@87
|
114 a = np.array([5, 7, 1, 2])
|
Chris@87
|
115 b = np.array([2, 4, 3, 1, 5])
|
Chris@87
|
116
|
Chris@87
|
117 ec = np.array([1, 2, 5])
|
Chris@87
|
118 c = intersect1d(a, b, assume_unique=True)
|
Chris@87
|
119 assert_array_equal(c, ec)
|
Chris@87
|
120
|
Chris@87
|
121 # non-unique inputs
|
Chris@87
|
122 a = np.array([5, 5, 7, 1, 2])
|
Chris@87
|
123 b = np.array([2, 1, 4, 3, 3, 1, 5])
|
Chris@87
|
124
|
Chris@87
|
125 ed = np.array([1, 2, 5])
|
Chris@87
|
126 c = intersect1d(a, b)
|
Chris@87
|
127 assert_array_equal(c, ed)
|
Chris@87
|
128
|
Chris@87
|
129 assert_array_equal([], intersect1d([], []))
|
Chris@87
|
130
|
Chris@87
|
131 def test_setxor1d(self):
|
Chris@87
|
132 a = np.array([5, 7, 1, 2])
|
Chris@87
|
133 b = np.array([2, 4, 3, 1, 5])
|
Chris@87
|
134
|
Chris@87
|
135 ec = np.array([3, 4, 7])
|
Chris@87
|
136 c = setxor1d(a, b)
|
Chris@87
|
137 assert_array_equal(c, ec)
|
Chris@87
|
138
|
Chris@87
|
139 a = np.array([1, 2, 3])
|
Chris@87
|
140 b = np.array([6, 5, 4])
|
Chris@87
|
141
|
Chris@87
|
142 ec = np.array([1, 2, 3, 4, 5, 6])
|
Chris@87
|
143 c = setxor1d(a, b)
|
Chris@87
|
144 assert_array_equal(c, ec)
|
Chris@87
|
145
|
Chris@87
|
146 a = np.array([1, 8, 2, 3])
|
Chris@87
|
147 b = np.array([6, 5, 4, 8])
|
Chris@87
|
148
|
Chris@87
|
149 ec = np.array([1, 2, 3, 4, 5, 6])
|
Chris@87
|
150 c = setxor1d(a, b)
|
Chris@87
|
151 assert_array_equal(c, ec)
|
Chris@87
|
152
|
Chris@87
|
153 assert_array_equal([], setxor1d([], []))
|
Chris@87
|
154
|
Chris@87
|
155 def test_ediff1d(self):
|
Chris@87
|
156 zero_elem = np.array([])
|
Chris@87
|
157 one_elem = np.array([1])
|
Chris@87
|
158 two_elem = np.array([1, 2])
|
Chris@87
|
159
|
Chris@87
|
160 assert_array_equal([], ediff1d(zero_elem))
|
Chris@87
|
161 assert_array_equal([0], ediff1d(zero_elem, to_begin=0))
|
Chris@87
|
162 assert_array_equal([0], ediff1d(zero_elem, to_end=0))
|
Chris@87
|
163 assert_array_equal([-1, 0], ediff1d(zero_elem, to_begin=-1, to_end=0))
|
Chris@87
|
164 assert_array_equal([], ediff1d(one_elem))
|
Chris@87
|
165 assert_array_equal([1], ediff1d(two_elem))
|
Chris@87
|
166
|
Chris@87
|
167 def test_in1d(self):
|
Chris@87
|
168 # we use two different sizes for the b array here to test the
|
Chris@87
|
169 # two different paths in in1d().
|
Chris@87
|
170 for mult in (1, 10):
|
Chris@87
|
171 # One check without np.array, to make sure lists are handled correct
|
Chris@87
|
172 a = [5, 7, 1, 2]
|
Chris@87
|
173 b = [2, 4, 3, 1, 5] * mult
|
Chris@87
|
174 ec = np.array([True, False, True, True])
|
Chris@87
|
175 c = in1d(a, b, assume_unique=True)
|
Chris@87
|
176 assert_array_equal(c, ec)
|
Chris@87
|
177
|
Chris@87
|
178 a[0] = 8
|
Chris@87
|
179 ec = np.array([False, False, True, True])
|
Chris@87
|
180 c = in1d(a, b, assume_unique=True)
|
Chris@87
|
181 assert_array_equal(c, ec)
|
Chris@87
|
182
|
Chris@87
|
183 a[0], a[3] = 4, 8
|
Chris@87
|
184 ec = np.array([True, False, True, False])
|
Chris@87
|
185 c = in1d(a, b, assume_unique=True)
|
Chris@87
|
186 assert_array_equal(c, ec)
|
Chris@87
|
187
|
Chris@87
|
188 a = np.array([5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5])
|
Chris@87
|
189 b = [2, 3, 4] * mult
|
Chris@87
|
190 ec = [False, True, False, True, True, True, True, True, True, False,
|
Chris@87
|
191 True, False, False, False]
|
Chris@87
|
192 c = in1d(a, b)
|
Chris@87
|
193 assert_array_equal(c, ec)
|
Chris@87
|
194
|
Chris@87
|
195 b = b + [5, 5, 4] * mult
|
Chris@87
|
196 ec = [True, True, True, True, True, True, True, True, True, True,
|
Chris@87
|
197 True, False, True, True]
|
Chris@87
|
198 c = in1d(a, b)
|
Chris@87
|
199 assert_array_equal(c, ec)
|
Chris@87
|
200
|
Chris@87
|
201 a = np.array([5, 7, 1, 2])
|
Chris@87
|
202 b = np.array([2, 4, 3, 1, 5] * mult)
|
Chris@87
|
203 ec = np.array([True, False, True, True])
|
Chris@87
|
204 c = in1d(a, b)
|
Chris@87
|
205 assert_array_equal(c, ec)
|
Chris@87
|
206
|
Chris@87
|
207 a = np.array([5, 7, 1, 1, 2])
|
Chris@87
|
208 b = np.array([2, 4, 3, 3, 1, 5] * mult)
|
Chris@87
|
209 ec = np.array([True, False, True, True, True])
|
Chris@87
|
210 c = in1d(a, b)
|
Chris@87
|
211 assert_array_equal(c, ec)
|
Chris@87
|
212
|
Chris@87
|
213 a = np.array([5, 5])
|
Chris@87
|
214 b = np.array([2, 2] * mult)
|
Chris@87
|
215 ec = np.array([False, False])
|
Chris@87
|
216 c = in1d(a, b)
|
Chris@87
|
217 assert_array_equal(c, ec)
|
Chris@87
|
218
|
Chris@87
|
219 a = np.array([5])
|
Chris@87
|
220 b = np.array([2])
|
Chris@87
|
221 ec = np.array([False])
|
Chris@87
|
222 c = in1d(a, b)
|
Chris@87
|
223 assert_array_equal(c, ec)
|
Chris@87
|
224
|
Chris@87
|
225 assert_array_equal(in1d([], []), [])
|
Chris@87
|
226
|
Chris@87
|
227 def test_in1d_char_array(self):
|
Chris@87
|
228 a = np.array(['a', 'b', 'c', 'd', 'e', 'c', 'e', 'b'])
|
Chris@87
|
229 b = np.array(['a', 'c'])
|
Chris@87
|
230
|
Chris@87
|
231 ec = np.array([True, False, True, False, False, True, False, False])
|
Chris@87
|
232 c = in1d(a, b)
|
Chris@87
|
233
|
Chris@87
|
234 assert_array_equal(c, ec)
|
Chris@87
|
235
|
Chris@87
|
236 def test_in1d_invert(self):
|
Chris@87
|
237 "Test in1d's invert parameter"
|
Chris@87
|
238 # We use two different sizes for the b array here to test the
|
Chris@87
|
239 # two different paths in in1d().
|
Chris@87
|
240 for mult in (1, 10):
|
Chris@87
|
241 a = np.array([5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5])
|
Chris@87
|
242 b = [2, 3, 4] * mult
|
Chris@87
|
243 assert_array_equal(np.invert(in1d(a, b)), in1d(a, b, invert=True))
|
Chris@87
|
244
|
Chris@87
|
245 def test_in1d_ravel(self):
|
Chris@87
|
246 # Test that in1d ravels its input arrays. This is not documented
|
Chris@87
|
247 # behavior however. The test is to ensure consistentency.
|
Chris@87
|
248 a = np.arange(6).reshape(2, 3)
|
Chris@87
|
249 b = np.arange(3, 9).reshape(3, 2)
|
Chris@87
|
250 long_b = np.arange(3, 63).reshape(30, 2)
|
Chris@87
|
251 ec = np.array([False, False, False, True, True, True])
|
Chris@87
|
252
|
Chris@87
|
253 assert_array_equal(in1d(a, b, assume_unique=True), ec)
|
Chris@87
|
254 assert_array_equal(in1d(a, b, assume_unique=False), ec)
|
Chris@87
|
255 assert_array_equal(in1d(a, long_b, assume_unique=True), ec)
|
Chris@87
|
256 assert_array_equal(in1d(a, long_b, assume_unique=False), ec)
|
Chris@87
|
257
|
Chris@87
|
258 def test_union1d(self):
|
Chris@87
|
259 a = np.array([5, 4, 7, 1, 2])
|
Chris@87
|
260 b = np.array([2, 4, 3, 3, 2, 1, 5])
|
Chris@87
|
261
|
Chris@87
|
262 ec = np.array([1, 2, 3, 4, 5, 7])
|
Chris@87
|
263 c = union1d(a, b)
|
Chris@87
|
264 assert_array_equal(c, ec)
|
Chris@87
|
265
|
Chris@87
|
266 assert_array_equal([], union1d([], []))
|
Chris@87
|
267
|
Chris@87
|
268 def test_setdiff1d(self):
|
Chris@87
|
269 a = np.array([6, 5, 4, 7, 1, 2, 7, 4])
|
Chris@87
|
270 b = np.array([2, 4, 3, 3, 2, 1, 5])
|
Chris@87
|
271
|
Chris@87
|
272 ec = np.array([6, 7])
|
Chris@87
|
273 c = setdiff1d(a, b)
|
Chris@87
|
274 assert_array_equal(c, ec)
|
Chris@87
|
275
|
Chris@87
|
276 a = np.arange(21)
|
Chris@87
|
277 b = np.arange(19)
|
Chris@87
|
278 ec = np.array([19, 20])
|
Chris@87
|
279 c = setdiff1d(a, b)
|
Chris@87
|
280 assert_array_equal(c, ec)
|
Chris@87
|
281
|
Chris@87
|
282 assert_array_equal([], setdiff1d([], []))
|
Chris@87
|
283
|
Chris@87
|
284 def test_setdiff1d_char_array(self):
|
Chris@87
|
285 a = np.array(['a', 'b', 'c'])
|
Chris@87
|
286 b = np.array(['a', 'b', 's'])
|
Chris@87
|
287 assert_array_equal(setdiff1d(a, b), np.array(['c']))
|
Chris@87
|
288
|
Chris@87
|
289 def test_manyways(self):
|
Chris@87
|
290 a = np.array([5, 7, 1, 2, 8])
|
Chris@87
|
291 b = np.array([9, 8, 2, 4, 3, 1, 5])
|
Chris@87
|
292
|
Chris@87
|
293 c1 = setxor1d(a, b)
|
Chris@87
|
294 aux1 = intersect1d(a, b)
|
Chris@87
|
295 aux2 = union1d(a, b)
|
Chris@87
|
296 c2 = setdiff1d(aux2, aux1)
|
Chris@87
|
297 assert_array_equal(c1, c2)
|
Chris@87
|
298
|
Chris@87
|
299
|
Chris@87
|
300 if __name__ == "__main__":
|
Chris@87
|
301 run_module_suite()
|