codekingpro/portable-devtools
114k
1"""Test functions for 1D array set operations.
2
3"""
4import pytest
5
6import numpy as np
7from numpy import ediff1d, intersect1d, isin, setdiff1d, setxor1d, union1d, unique
8from numpy.dtypes import StringDType
9from numpy.exceptions import AxisError
10from numpy.testing import (
11 assert_array_equal,
12 assert_equal,
13 assert_raises,
14 assert_raises_regex,
15)
16
17
18class TestSetOps:
19
20 def test_intersect1d(self):
21 # unique inputs
22 a = np.array([5, 7, 1, 2])
23 b = np.array([2, 4, 3, 1, 5])
24
25 ec = np.array([1, 2, 5])
26 c = intersect1d(a, b, assume_unique=True)
27 assert_array_equal(c, ec)
28
29 # non-unique inputs
30 a = np.array([5, 5, 7, 1, 2])
31 b = np.array([2, 1, 4, 3, 3, 1, 5])
32
33 ed = np.array([1, 2, 5])
34 c = intersect1d(a, b)
35 assert_array_equal(c, ed)
36 assert_array_equal([], intersect1d([], []))
37
38 def test_intersect1d_array_like(self):
39 # See gh-11772
40 class Test:
41 def __array__(self, dtype=None, copy=None):
42 return np.arange(3)
43
44 a = Test()
45 res = intersect1d(a, a)
46 assert_array_equal(res, a)
47 res = intersect1d([1, 2, 3], [1, 2, 3])
48 assert_array_equal(res, [1, 2, 3])
49
50 def test_intersect1d_indices(self):
51 # unique inputs
52 a = np.array([1, 2, 3, 4])
53 b = np.array([2, 1, 4, 6])
54 c, i1, i2 = intersect1d(a, b, assume_unique=True, return_indices=True)
55 ee = np.array([1, 2, 4])
56 assert_array_equal(c, ee)
57 assert_array_equal(a[i1], ee)
58 assert_array_equal(b[i2], ee)
59
60 # non-unique inputs
61 a = np.array([1, 2, 2, 3, 4, 3, 2])
62 b = np.array([1, 8, 4, 2, 2, 3, 2, 3])
63 c, i1, i2 = intersect1d(a, b, return_indices=True)
64 ef = np.array([1, 2, 3, 4])
65 assert_array_equal(c, ef)
66 assert_array_equal(a[i1], ef)
67 assert_array_equal(b[i2], ef)
68
69 # non1d, unique inputs
70 a = np.array([[2, 4, 5, 6], [7, 8, 1, 15]])
71 b = np.array([[3, 2, 7, 6], [10, 12, 8, 9]])
72 c, i1, i2 = intersect1d(a, b, assume_unique=True, return_indices=True)
73 ui1 = np.unravel_index(i1, a.shape)
74 ui2 = np.unravel_index(i2, b.shape)
75 ea = np.array([2, 6, 7, 8])
76 assert_array_equal(ea, a[ui1])
77 assert_array_equal(ea, b[ui2])
78
79 # non1d, not assumed to be uniqueinputs
80 a = np.array([[2, 4, 5, 6, 6], [4, 7, 8, 7, 2]])
81 b = np.array([[3, 2, 7, 7], [10, 12, 8, 7]])
82 c, i1, i2 = intersect1d(a, b, return_indices=True)
83 ui1 = np.unravel_index(i1, a.shape)
84 ui2 = np.unravel_index(i2, b.shape)
85 ea = np.array([2, 7, 8])
86 assert_array_equal(ea, a[ui1])
87 assert_array_equal(ea, b[ui2])
88
89 def test_setxor1d(self):
90 a = np.array([5, 7, 1, 2])
91 b = np.array([2, 4, 3, 1, 5])
92
93 ec = np.array([3, 4, 7])
94 c = setxor1d(a, b)
95 assert_array_equal(c, ec)
96
97 a = np.array([1, 2, 3])
98 b = np.array([6, 5, 4])
99
100 ec = np.array([1, 2, 3, 4, 5, 6])
101 c = setxor1d(a, b)
102 assert_array_equal(c, ec)
103
104 a = np.array([1, 8, 2, 3])
105 b = np.array([6, 5, 4, 8])
106
107 ec = np.array([1, 2, 3, 4, 5, 6])
108 c = setxor1d(a, b)
109 assert_array_equal(c, ec)
110
111 assert_array_equal([], setxor1d([], []))
112
113 def test_setxor1d_unique(self):
114 a = np.array([1, 8, 2, 3])
115 b = np.array([6, 5, 4, 8])
116
117 ec = np.array([1, 2, 3, 4, 5, 6])
118 c = setxor1d(a, b, assume_unique=True)
119 assert_array_equal(c, ec)
120
121 a = np.array([[1], [8], [2], [3]])
122 b = np.array([[6, 5], [4, 8]])
123
124 ec = np.array([1, 2, 3, 4, 5, 6])
125 c = setxor1d(a, b, assume_unique=True)
126 assert_array_equal(c, ec)
127
128 def test_ediff1d(self):
129 zero_elem = np.array([])
130 one_elem = np.array([1])
131 two_elem = np.array([1, 2])
132
133 assert_array_equal([], ediff1d(zero_elem))
134 assert_array_equal([0], ediff1d(zero_elem, to_begin=0))
135 assert_array_equal([0], ediff1d(zero_elem, to_end=0))
136 assert_array_equal([-1, 0], ediff1d(zero_elem, to_begin=-1, to_end=0))
137 assert_array_equal([], ediff1d(one_elem))
138 assert_array_equal([1], ediff1d(two_elem))
139 assert_array_equal([7, 1, 9], ediff1d(two_elem, to_begin=7, to_end=9))
140 assert_array_equal([5, 6, 1, 7, 8],
141 ediff1d(two_elem, to_begin=[5, 6], to_end=[7, 8]))
142 assert_array_equal([1, 9], ediff1d(two_elem, to_end=9))
143 assert_array_equal([1, 7, 8], ediff1d(two_elem, to_end=[7, 8]))
144 assert_array_equal([7, 1], ediff1d(two_elem, to_begin=7))
145 assert_array_equal([5, 6, 1], ediff1d(two_elem, to_begin=[5, 6]))
146
147 @pytest.mark.parametrize("ary, prepend, append, expected", [
148 # should fail because trying to cast
149 # np.nan standard floating point value
150 # into an integer array:
151 (np.array([1, 2, 3], dtype=np.int64),
152 None,
153 np.nan,
154 'to_end'),
155 # should fail because attempting
156 # to downcast to int type:
157 (np.array([1, 2, 3], dtype=np.int64),
158 np.array([5, 7, 2], dtype=np.float32),
159 None,
160 'to_begin'),
161 # should fail because attempting to cast
162 # two special floating point values
163 # to integers (on both sides of ary),
164 # `to_begin` is in the error message as the impl checks this first:
165 (np.array([1., 3., 9.], dtype=np.int8),
166 np.nan,
167 np.nan,
168 'to_begin'),
169 ])
170 def test_ediff1d_forbidden_type_casts(self, ary, prepend, append, expected):
171 # verify resolution of gh-11490
172
173 # specifically, raise an appropriate
174 # Exception when attempting to append or
175 # prepend with an incompatible type
176 msg = f'dtype of `{expected}` must be compatible'
177 with assert_raises_regex(TypeError, msg):
178 ediff1d(ary=ary,
179 to_end=append,
180 to_begin=prepend)
181
182 @pytest.mark.parametrize(
183 "ary,prepend,append,expected",
184 [
185 (np.array([1, 2, 3], dtype=np.int16),
186 2**16, # will be cast to int16 under same kind rule.
187 2**16 + 4,
188 np.array([0, 1, 1, 4], dtype=np.int16)),
189 (np.array([1, 2, 3], dtype=np.float32),
190 np.array([5], dtype=np.float64),
191 None,
192 np.array([5, 1, 1], dtype=np.float32)),
193 (np.array([1, 2, 3], dtype=np.int32),
194 0,
195 0,
196 np.array([0, 1, 1, 0], dtype=np.int32)),
197 (np.array([1, 2, 3], dtype=np.int64),
198 3,
199 -9,
200 np.array([3, 1, 1, -9], dtype=np.int64)),
201 ]
202 )
203 def test_ediff1d_scalar_handling(self,
204 ary,
205 prepend,
206 append,
207 expected):
208 # maintain backwards-compatibility
209 # of scalar prepend / append behavior
210 # in ediff1d following fix for gh-11490
211 actual = np.ediff1d(ary=ary,
212 to_end=append,
213 to_begin=prepend)
214 assert_equal(actual, expected)
215 assert actual.dtype == expected.dtype
216
217 @pytest.mark.parametrize("kind", [None, "sort", "table"])
218 def test_isin(self, kind):
219 def _isin_slow(a, b):
220 b = np.asarray(b).flatten().tolist()
221 return a in b
222 isin_slow = np.vectorize(_isin_slow, otypes=[bool], excluded={1})
223
224 def assert_isin_equal(a, b):
225 x = isin(a, b, kind=kind)
226 y = isin_slow(a, b)
227 assert_array_equal(x, y)
228
229 # multidimensional arrays in both arguments
230 a = np.arange(24).reshape([2, 3, 4])
231 b = np.array([[10, 20, 30], [0, 1, 3], [11, 22, 33]])
232 assert_isin_equal(a, b)
233
234 # array-likes as both arguments
235 c = [(9, 8), (7, 6)]
236 d = (9, 7)
237 assert_isin_equal(c, d)
238
239 # zero-d array:
240 f = np.array(3)
241 assert_isin_equal(f, b)
242 assert_isin_equal(a, f)
243 assert_isin_equal(f, f)
244
245 # scalar:
246 assert_isin_equal(5, b)
247 assert_isin_equal(a, 6)
248 assert_isin_equal(5, 6)
249
250 # empty array-like:
251 if kind != "table":
252 # An empty list will become float64,
253 # which is invalid for kind="table"
254 x = []
255 assert_isin_equal(x, b)
256 assert_isin_equal(a, x)
257 assert_isin_equal(x, x)
258
259 # empty array with various types:
260 for dtype in [bool, np.int64, np.float64]:
261 if kind == "table" and dtype == np.float64:
262 continue
263
264 if dtype in {np.int64, np.float64}:
265 ar = np.array([10, 20, 30], dtype=dtype)
266 elif dtype in {bool}:
267 ar = np.array([True, False, False])
268
269 empty_array = np.array([], dtype=dtype)
270
271 assert_isin_equal(empty_array, ar)
272 assert_isin_equal(ar, empty_array)
273 assert_isin_equal(empty_array, empty_array)
274
275 @pytest.mark.parametrize("kind", [None, "sort", "table"])
276 def test_isin_additional(self, kind):
277 # we use two different sizes for the b array here to test the
278 # two different paths in isin().
279 for mult in (1, 10):
280 # One check without np.array to make sure lists are handled correct
281 a = [5, 7, 1, 2]
282 b = [2, 4, 3, 1, 5] * mult
283 ec = np.array([True, False, True, True])
284 c = isin(a, b, assume_unique=True, kind=kind)
285 assert_array_equal(c, ec)
286
287 a[0] = 8
288 ec = np.array([False, False, True, True])
289 c = isin(a, b, assume_unique=True, kind=kind)
290 assert_array_equal(c, ec)
291
292 a[0], a[3] = 4, 8
293 ec = np.array([True, False, True, False])
294 c = isin(a, b, assume_unique=True, kind=kind)
295 assert_array_equal(c, ec)
296
297 a = np.array([5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5])
298 b = [2, 3, 4] * mult
299 ec = [False, True, False, True, True, True, True, True, True,
300 False, True, False, False, False]
301 c = isin(a, b, kind=kind)
302 assert_array_equal(c, ec)
303
304 b = b + [5, 5, 4] * mult
305 ec = [True, True, True, True, True, True, True, True, True, True,
306 True, False, True, True]
307 c = isin(a, b, kind=kind)
308 assert_array_equal(c, ec)
309
310 a = np.array([5, 7, 1, 2])
311 b = np.array([2, 4, 3, 1, 5] * mult)
312 ec = np.array([True, False, True, True])
313 c = isin(a, b, kind=kind)
314 assert_array_equal(c, ec)
315
316 a = np.array([5, 7, 1, 1, 2])
317 b = np.array([2, 4, 3, 3, 1, 5] * mult)
318 ec = np.array([True, False, True, True, True])
319 c = isin(a, b, kind=kind)
320 assert_array_equal(c, ec)
321
322 a = np.array([5, 5])
323 b = np.array([2, 2] * mult)
324 ec = np.array([False, False])
325 c = isin(a, b, kind=kind)
326 assert_array_equal(c, ec)
327
328 a = np.array([5])
329 b = np.array([2])
330 ec = np.array([False])
331 c = isin(a, b, kind=kind)
332 assert_array_equal(c, ec)
333
334 if kind in {None, "sort"}:
335 assert_array_equal(isin([], [], kind=kind), [])
336
337 def test_isin_char_array(self):
338 a = np.array(['a', 'b', 'c', 'd', 'e', 'c', 'e', 'b'])
339 b = np.array(['a', 'c'])
340
341 ec = np.array([True, False, True, False, False, True, False, False])
342 c = isin(a, b)
343
344 assert_array_equal(c, ec)
345
346 @pytest.mark.parametrize("kind", [None, "sort", "table"])
347 def test_isin_invert(self, kind):
348 "Test isin's invert parameter"
349 # We use two different sizes for the b array here to test the
350 # two different paths in isin().
351 for mult in (1, 10):
352 a = np.array([5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5])
353 b = [2, 3, 4] * mult
354 assert_array_equal(np.invert(isin(a, b, kind=kind)),
355 isin(a, b, invert=True, kind=kind))
356
357 # float:
358 if kind in {None, "sort"}:
359 for mult in (1, 10):
360 a = np.array([5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5],
361 dtype=np.float32)
362 b = [2, 3, 4] * mult
363 b = np.array(b, dtype=np.float32)
364 assert_array_equal(np.invert(isin(a, b, kind=kind)),
365 isin(a, b, invert=True, kind=kind))
366
367 def test_isin_hit_alternate_algorithm(self):
368 """Hit the standard isin code with integers"""
369 # Need extreme range to hit standard code
370 # This hits it without the use of kind='table'
371 a = np.array([5, 4, 5, 3, 4, 4, 1e9], dtype=np.int64)
372 b = np.array([2, 3, 4, 1e9], dtype=np.int64)
373 expected = np.array([0, 1, 0, 1, 1, 1, 1], dtype=bool)
374 assert_array_equal(expected, isin(a, b))
375 assert_array_equal(np.invert(expected), isin(a, b, invert=True))
376
377 a = np.array([5, 7, 1, 2], dtype=np.int64)
378 b = np.array([2, 4, 3, 1, 5, 1e9], dtype=np.int64)
379 ec = np.array([True, False, True, True])
380 c = isin(a, b, assume_unique=True)
381 assert_array_equal(c, ec)
382
383 @pytest.mark.parametrize("kind", [None, "sort", "table"])
384 def test_isin_boolean(self, kind):
385 """Test that isin works for boolean input"""
386 a = np.array([True, False])
387 b = np.array([False, False, False])
388 expected = np.array([False, True])
389 assert_array_equal(expected,
390 isin(a, b, kind=kind))
391 assert_array_equal(np.invert(expected),
392 isin(a, b, invert=True, kind=kind))
393
394 @pytest.mark.parametrize("kind", [None, "sort"])
395 def test_isin_timedelta(self, kind):
396 """Test that isin works for timedelta input"""
397 rstate = np.random.RandomState(0)
398 a = rstate.randint(0, 100, size=10)
399 b = rstate.randint(0, 100, size=10)
400 truth = isin(a, b)
401 a_timedelta = a.astype("timedelta64[s]")
402 b_timedelta = b.astype("timedelta64[s]")
403 assert_array_equal(truth, isin(a_timedelta, b_timedelta, kind=kind))
404
405 def test_isin_table_timedelta_fails(self):
406 a = np.array([0, 1, 2], dtype="timedelta64[s]")
407 b = a
408 # Make sure it raises a value error:
409 with pytest.raises(ValueError):
410 isin(a, b, kind="table")
411
412 @pytest.mark.parametrize(
413 "dtype1,dtype2",
414 [
415 (np.int8, np.int16),
416 (np.int16, np.int8),
417 (np.uint8, np.uint16),
418 (np.uint16, np.uint8),
419 (np.uint8, np.int16),
420 (np.int16, np.uint8),
421 (np.uint64, np.int64),
422 ]
423 )
424 @pytest.mark.parametrize("kind", [None, "sort", "table"])
425 def test_isin_mixed_dtype(self, dtype1, dtype2, kind):
426 """Test that isin works as expected for mixed dtype input."""
427 is_dtype2_signed = np.issubdtype(dtype2, np.signedinteger)
428 ar1 = np.array([0, 0, 1, 1], dtype=dtype1)
429
430 if is_dtype2_signed:
431 ar2 = np.array([-128, 0, 127], dtype=dtype2)
432 else:
433 ar2 = np.array([127, 0, 255], dtype=dtype2)
434
435 expected = np.array([True, True, False, False])
436
437 expect_failure = kind == "table" and (
438 dtype1 == np.int16 and dtype2 == np.int8)
439
440 if expect_failure:
441 with pytest.raises(RuntimeError, match="exceed the maximum"):
442 isin(ar1, ar2, kind=kind)
443 else:
444 assert_array_equal(isin(ar1, ar2, kind=kind), expected)
445
446 @pytest.mark.parametrize("data", [
447 np.array([2**63, 2**63 + 1], dtype=np.uint64),
448 np.array([-2**62, -2**62 - 1], dtype=np.int64),
449 ])
450 @pytest.mark.parametrize("kind", [None, "sort", "table"])
451 def test_isin_mixed_huge_vals(self, kind, data):
452 """Test values outside intp range (negative ones if 32bit system)"""
453 query = data[1]
454 res = np.isin(data, query, kind=kind)
455 assert_array_equal(res, [False, True])
456 # Also check that nothing weird happens for values can't possibly
457 # in range.
458 data = data.astype(np.int32) # clearly different values
459 res = np.isin(data, query, kind=kind)
460 assert_array_equal(res, [False, False])
461
462 @pytest.mark.parametrize("kind", [None, "sort", "table"])
463 def test_isin_mixed_boolean(self, kind):
464 """Test that isin works as expected for bool/int input."""
465 for dtype in np.typecodes["AllInteger"]:
466 a = np.array([True, False, False], dtype=bool)
467 b = np.array([0, 0, 0, 0], dtype=dtype)
468 expected = np.array([False, True, True], dtype=bool)
469 assert_array_equal(isin(a, b, kind=kind), expected)
470
471 a, b = b, a
472 expected = np.array([True, True, True, True], dtype=bool)
473 assert_array_equal(isin(a, b, kind=kind), expected)
474
475 def test_isin_first_array_is_object(self):
476 ar1 = [None]
477 ar2 = np.array([1] * 10)
478 expected = np.array([False])
479 result = np.isin(ar1, ar2)
480 assert_array_equal(result, expected)
481
482 def test_isin_second_array_is_object(self):
483 ar1 = 1
484 ar2 = np.array([None] * 10)
485 expected = np.array([False])
486 result = np.isin(ar1, ar2)
487 assert_array_equal(result, expected)
488
489 def test_isin_both_arrays_are_object(self):
490 ar1 = [None]
491 ar2 = np.array([None] * 10)
492 expected = np.array([True])
493 result = np.isin(ar1, ar2)
494 assert_array_equal(result, expected)
495
496 def test_isin_both_arrays_have_structured_dtype(self):
497 # Test arrays of a structured data type containing an integer field
498 # and a field of dtype `object` allowing for arbitrary Python objects
499 dt = np.dtype([('field1', int), ('field2', object)])
500 ar1 = np.array([(1, None)], dtype=dt)
501 ar2 = np.array([(1, None)] * 10, dtype=dt)
502 expected = np.array([True])
503 result = np.isin(ar1, ar2)
504 assert_array_equal(result, expected)
505
506 def test_isin_with_arrays_containing_tuples(self):
507 ar1 = np.array([(1,), 2], dtype=object)
508 ar2 = np.array([(1,), 2], dtype=object)
509 expected = np.array([True, True])
510 result = np.isin(ar1, ar2)
511 assert_array_equal(result, expected)
512 result = np.isin(ar1, ar2, invert=True)
513 assert_array_equal(result, np.invert(expected))
514
515 # An integer is added at the end of the array to make sure
516 # that the array builder will create the array with tuples
517 # and after it's created the integer is removed.
518 # There's a bug in the array constructor that doesn't handle
519 # tuples properly and adding the integer fixes that.
520 ar1 = np.array([(1,), (2, 1), 1], dtype=object)
521 ar1 = ar1[:-1]
522 ar2 = np.array([(1,), (2, 1), 1], dtype=object)
523 ar2 = ar2[:-1]
524 expected = np.array([True, True])
525 result = np.isin(ar1, ar2)
526 assert_array_equal(result, expected)
527 result = np.isin(ar1, ar2, invert=True)
528 assert_array_equal(result, np.invert(expected))
529
530 ar1 = np.array([(1,), (2, 3), 1], dtype=object)
531 ar1 = ar1[:-1]
532 ar2 = np.array([(1,), 2], dtype=object)
533 expected = np.array([True, False])
534 result = np.isin(ar1, ar2)
535 assert_array_equal(result, expected)
536 result = np.isin(ar1, ar2, invert=True)
537 assert_array_equal(result, np.invert(expected))
538
539 def test_isin_errors(self):
540 """Test that isin raises expected errors."""
541
542 # Error 1: `kind` is not one of 'sort' 'table' or None.
543 ar1 = np.array([1, 2, 3, 4, 5])
544 ar2 = np.array([2, 4, 6, 8, 10])
545 assert_raises(ValueError, isin, ar1, ar2, kind='quicksort')
546
547 # Error 2: `kind="table"` does not work for non-integral arrays.
548 obj_ar1 = np.array([1, 'a', 3, 'b', 5], dtype=object)
549 obj_ar2 = np.array([1, 'a', 3, 'b', 5], dtype=object)
550 assert_raises(ValueError, isin, obj_ar1, obj_ar2, kind='table')
551
552 for dtype in [np.int32, np.int64]:
553 ar1 = np.array([-1, 2, 3, 4, 5], dtype=dtype)
554 # The range of this array will overflow:
555 overflow_ar2 = np.array([-1, np.iinfo(dtype).max], dtype=dtype)
556
557 # Error 3: `kind="table"` will trigger a runtime error
558 # if there is an integer overflow expected when computing the
559 # range of ar2
560 assert_raises(
561 RuntimeError,
562 isin, ar1, overflow_ar2, kind='table'
563 )
564
565 # Non-error: `kind=None` will *not* trigger a runtime error
566 # if there is an integer overflow, it will switch to
567 # the `sort` algorithm.
568 result = np.isin(ar1, overflow_ar2, kind=None)
569 assert_array_equal(result, [True] + [False] * 4)
570 result = np.isin(ar1, overflow_ar2, kind='sort')
571 assert_array_equal(result, [True] + [False] * 4)
572
573 def test_union1d(self):
574 a = np.array([5, 4, 7, 1, 2])
575 b = np.array([2, 4, 3, 3, 2, 1, 5])
576
577 ec = np.array([1, 2, 3, 4, 5, 7])
578 c = union1d(a, b)
579 assert_array_equal(c, ec)
580
581 # Tests gh-10340, arguments to union1d should be
582 # flattened if they are not already 1D
583 x = np.array([[0, 1, 2], [3, 4, 5]])
584 y = np.array([0, 1, 2, 3, 4])
585 ez = np.array([0, 1, 2, 3, 4, 5])
586 z = union1d(x, y)
587 assert_array_equal(z, ez)
588
589 assert_array_equal([], union1d([], []))
590
591 def test_setdiff1d(self):
592 a = np.array([6, 5, 4, 7, 1, 2, 7, 4])
593 b = np.array([2, 4, 3, 3, 2, 1, 5])
594
595 ec = np.array([6, 7])
596 c = setdiff1d(a, b)
597 assert_array_equal(c, ec)
598
599 a = np.arange(21)
600 b = np.arange(19)
601 ec = np.array([19, 20])
602 c = setdiff1d(a, b)
603 assert_array_equal(c, ec)
604
605 assert_array_equal([], setdiff1d([], []))
606 a = np.array((), np.uint32)
607 assert_equal(setdiff1d(a, []).dtype, np.uint32)
608
609 def test_setdiff1d_unique(self):
610 a = np.array([3, 2, 1])
611 b = np.array([7, 5, 2])
612 expected = np.array([3, 1])
613 actual = setdiff1d(a, b, assume_unique=True)
614 assert_equal(actual, expected)
615
616 def test_setdiff1d_char_array(self):
617 a = np.array(['a', 'b', 'c'])
618 b = np.array(['a', 'b', 's'])
619 assert_array_equal(setdiff1d(a, b), np.array(['c']))
620
621 def test_manyways(self):
622 a = np.array([5, 7, 1, 2, 8])
623 b = np.array([9, 8, 2, 4, 3, 1, 5])
624
625 c1 = setxor1d(a, b)
626 aux1 = intersect1d(a, b)
627 aux2 = union1d(a, b)
628 c2 = setdiff1d(aux2, aux1)
629 assert_array_equal(c1, c2)
630
631
632class TestUnique:
633
634 def check_all(self, a, b, i1, i2, c, dt):
635 base_msg = 'check {0} failed for type {1}'
636
637 msg = base_msg.format('values', dt)
638 v = unique(a)
639 assert_array_equal(v, b, msg)
640 assert type(v) == type(b)
641
642 msg = base_msg.format('return_index', dt)
643 v, j = unique(a, True, False, False)
644 assert_array_equal(v, b, msg)
645 assert_array_equal(j, i1, msg)
646 assert type(v) == type(b)
647
648 msg = base_msg.format('return_inverse', dt)
649 v, j = unique(a, False, True, False)
650 assert_array_equal(v, b, msg)
651 assert_array_equal(j, i2, msg)
652 assert type(v) == type(b)
653
654 msg = base_msg.format('return_counts', dt)
655 v, j = unique(a, False, False, True)
656 assert_array_equal(v, b, msg)
657 assert_array_equal(j, c, msg)
658 assert type(v) == type(b)
659
660 msg = base_msg.format('return_index and return_inverse', dt)
661 v, j1, j2 = unique(a, True, True, False)
662 assert_array_equal(v, b, msg)
663 assert_array_equal(j1, i1, msg)
664 assert_array_equal(j2, i2, msg)
665 assert type(v) == type(b)
666
667 msg = base_msg.format('return_index and return_counts', dt)
668 v, j1, j2 = unique(a, True, False, True)
669 assert_array_equal(v, b, msg)
670 assert_array_equal(j1, i1, msg)
671 assert_array_equal(j2, c, msg)
672 assert type(v) == type(b)
673
674 msg = base_msg.format('return_inverse and return_counts', dt)
675 v, j1, j2 = unique(a, False, True, True)
676 assert_array_equal(v, b, msg)
677 assert_array_equal(j1, i2, msg)
678 assert_array_equal(j2, c, msg)
679 assert type(v) == type(b)
680
681 msg = base_msg.format(('return_index, return_inverse '
682 'and return_counts'), dt)
683 v, j1, j2, j3 = unique(a, True, True, True)
684 assert_array_equal(v, b, msg)
685 assert_array_equal(j1, i1, msg)
686 assert_array_equal(j2, i2, msg)
687 assert_array_equal(j3, c, msg)
688 assert type(v) == type(b)
689
690 def get_types(self):
691 types = []
692 types.extend(np.typecodes['AllInteger'])
693 types.extend(np.typecodes['AllFloat'])
694 types.append('datetime64[D]')
695 types.append('timedelta64[D]')
696 return types
697
698 def test_unique_1d(self):
699
700 a = [5, 7, 1, 2, 1, 5, 7] * 10
701 b = [1, 2, 5, 7]
702 i1 = [2, 3, 0, 1]
703 i2 = [2, 3, 0, 1, 0, 2, 3] * 10
704 c = np.multiply([2, 1, 2, 2], 10)
705
706 # test for numeric arrays
707 types = self.get_types()
708 for dt in types:
709 aa = np.array(a, dt)
710 bb = np.array(b, dt)
711 self.check_all(aa, bb, i1, i2, c, dt)
712
713 # test for object arrays
714 dt = 'O'
715 aa = np.empty(len(a), dt)
716 aa[:] = a
717 bb = np.empty(len(b), dt)
718 bb[:] = b
719 self.check_all(aa, bb, i1, i2, c, dt)
720
721 # test for structured arrays
722 dt = [('', 'i'), ('', 'i')]
723 aa = np.array(list(zip(a, a)), dt)
724 bb = np.array(list(zip(b, b)), dt)
725 self.check_all(aa, bb, i1, i2, c, dt)
726
727 # test for ticket #2799
728 aa = [1. + 0.j, 1 - 1.j, 1]
729 assert_array_equal(
730 np.sort(np.unique(aa)),
731 [1. - 1.j, 1.],
732 )
733
734 # test for ticket #4785
735 a = [(1, 2), (1, 2), (2, 3)]
736 unq = [1, 2, 3]
737 inv = [[0, 1], [0, 1], [1, 2]]
738 a1 = unique(a)
739 assert_array_equal(a1, unq)
740 a2, a2_inv = unique(a, return_inverse=True)
741 assert_array_equal(a2, unq)
742 assert_array_equal(a2_inv, inv)
743
744 # test for chararrays with return_inverse (gh-5099)
745 a = np.char.chararray(5)
746 a[...] = ''
747 a2, a2_inv = np.unique(a, return_inverse=True)
748 assert_array_equal(a2_inv, np.zeros(5))
749
750 # test for ticket #9137
751 a = []
752 a1_idx = np.unique(a, return_index=True)[1]
753 a2_inv = np.unique(a, return_inverse=True)[1]
754 a3_idx, a3_inv = np.unique(a, return_index=True,
755 return_inverse=True)[1:]
756 assert_equal(a1_idx.dtype, np.intp)
757 assert_equal(a2_inv.dtype, np.intp)
758 assert_equal(a3_idx.dtype, np.intp)
759 assert_equal(a3_inv.dtype, np.intp)
760
761 # test for ticket 2111 - float
762 a = [2.0, np.nan, 1.0, np.nan]
763 ua = [1.0, 2.0, np.nan]
764 ua_idx = [2, 0, 1]
765 ua_inv = [1, 2, 0, 2]
766 ua_cnt = [1, 1, 2]
767 # order of unique values is not guaranteed
768 assert_equal(np.sort(np.unique(a)), np.sort(ua))
769 assert_equal(np.unique(a, return_index=True), (ua, ua_idx))
770 assert_equal(np.unique(a, return_inverse=True), (ua, ua_inv))
771 assert_equal(np.unique(a, return_counts=True), (ua, ua_cnt))
772
773 # test for ticket 2111 - complex
774 a = [2.0 - 1j, np.nan, 1.0 + 1j, complex(0.0, np.nan), complex(1.0, np.nan)]
775 ua = [1.0 + 1j, 2.0 - 1j, complex(0.0, np.nan)]
776 ua_idx = [2, 0, 3]
777 ua_inv = [1, 2, 0, 2, 2]
778 ua_cnt = [1, 1, 3]
779 # order of unique values is not guaranteed
780 assert_equal(np.sort(np.unique(a)), np.sort(ua))
781 assert_equal(np.unique(a, return_index=True), (ua, ua_idx))
782 assert_equal(np.unique(a, return_inverse=True), (ua, ua_inv))
783 assert_equal(np.unique(a, return_counts=True), (ua, ua_cnt))
784
785 # test for ticket 2111 - datetime64
786 nat = np.datetime64('nat')
787 a = [np.datetime64('2020-12-26'), nat, np.datetime64('2020-12-24'), nat]
788 ua = [np.datetime64('2020-12-24'), np.datetime64('2020-12-26'), nat]
789 ua_idx = [2, 0, 1]
790 ua_inv = [1, 2, 0, 2]
791 ua_cnt = [1, 1, 2]
792 assert_equal(np.unique(a), ua)
793 assert_equal(np.unique(a, return_index=True), (ua, ua_idx))
794 assert_equal(np.unique(a, return_inverse=True), (ua, ua_inv))
795 assert_equal(np.unique(a, return_counts=True), (ua, ua_cnt))
796
797 # test for ticket 2111 - timedelta
798 nat = np.timedelta64('nat')
799 a = [np.timedelta64(1, 'D'), nat, np.timedelta64(1, 'h'), nat]
800 ua = [np.timedelta64(1, 'h'), np.timedelta64(1, 'D'), nat]
801 ua_idx = [2, 0, 1]
802 ua_inv = [1, 2, 0, 2]
803 ua_cnt = [1, 1, 2]
804 assert_equal(np.unique(a), ua)
805 assert_equal(np.unique(a, return_index=True), (ua, ua_idx))
806 assert_equal(np.unique(a, return_inverse=True), (ua, ua_inv))
807 assert_equal(np.unique(a, return_counts=True), (ua, ua_cnt))
808
809 # test for gh-19300
810 all_nans = [np.nan] * 4
811 ua = [np.nan]
812 ua_idx = [0]
813 ua_inv = [0, 0, 0, 0]
814 ua_cnt = [4]
815 assert_equal(np.unique(all_nans), ua)
816 assert_equal(np.unique(all_nans, return_index=True), (ua, ua_idx))
817 assert_equal(np.unique(all_nans, return_inverse=True), (ua, ua_inv))
818 assert_equal(np.unique(all_nans, return_counts=True), (ua, ua_cnt))
819
820 def test_unique_zero_sized(self):
821 # test for zero-sized arrays
822 types = self.get_types()
823 types.extend('SU')
824 for dt in types:
825 a = np.array([], dt)
826 b = np.array([], dt)
827 i1 = np.array([], np.int64)
828 i2 = np.array([], np.int64)
829 c = np.array([], np.int64)
830 self.check_all(a, b, i1, i2, c, dt)
831
832 def test_unique_subclass(self):
833 class Subclass(np.ndarray):
834 pass
835
836 i1 = [2, 3, 0, 1]
837 i2 = [2, 3, 0, 1, 0, 2, 3] * 10
838 c = np.multiply([2, 1, 2, 2], 10)
839
840 # test for numeric arrays
841 types = self.get_types()
842 for dt in types:
843 a = np.array([5, 7, 1, 2, 1, 5, 7] * 10, dtype=dt)
844 b = np.array([1, 2, 5, 7], dtype=dt)
845 aa = Subclass(a.shape, dtype=dt, buffer=a)
846 bb = Subclass(b.shape, dtype=dt, buffer=b)
847 self.check_all(aa, bb, i1, i2, c, dt)
848
849 def test_unique_byte_string_hash_based(self):
850 # test for byte string arrays
851 arr = ['apple', 'banana', 'apple', 'cherry', 'date', 'banana', 'fig', 'grape']
852 unq_sorted = ['apple', 'banana', 'cherry', 'date', 'fig', 'grape']
853
854 a1 = unique(arr, sorted=False)
855 # the result varies depending on the impl of std::unordered_set,
856 # so we check them by sorting
857 assert_array_equal(sorted(a1.tolist()), unq_sorted)
858
859 def test_unique_unicode_string_hash_based(self):
860 # test for unicode string arrays
861 arr = [
862 'café', 'cafe', 'café', 'naïve', 'naive',
863 'résumé', 'naïve', 'resume', 'résumé',
864 ]
865 unq_sorted = ['cafe', 'café', 'naive', 'naïve', 'resume', 'résumé']
866
867 a1 = unique(arr, sorted=False)
868 # the result varies depending on the impl of std::unordered_set,
869 # so we check them by sorting
870 assert_array_equal(sorted(a1.tolist()), unq_sorted)
871
872 def test_unique_vstring_hash_based_equal_nan(self):
873 # test for unicode and nullable string arrays (equal_nan=True)
874 a = np.array([
875 # short strings
876 'straße',
877 None,
878 'strasse',
879 'straße',
880 None,
881 'niño',
882 'nino',
883 'élève',
884 'eleve',
885 'niño',
886 'élève',
887 # medium strings
888 'b' * 20,
889 'ß' * 30,
890 None,
891 'é' * 30,
892 'e' * 20,
893 'ß' * 30,
894 'n' * 30,
895 'ñ' * 20,
896 None,
897 'e' * 20,
898 'ñ' * 20,
899 # long strings
900 'b' * 300,
901 'ß' * 400,
902 None,
903 'é' * 400,
904 'e' * 300,
905 'ß' * 400,
906 'n' * 400,
907 'ñ' * 300,
908 None,
909 'e' * 300,
910 'ñ' * 300,
911 ],
912 dtype=StringDType(na_object=None)
913 )
914 unq_sorted_wo_none = [
915 'b' * 20,
916 'b' * 300,
917 'e' * 20,
918 'e' * 300,
919 'eleve',
920 'nino',
921 'niño',
922 'n' * 30,
923 'n' * 400,
924 'strasse',
925 'straße',
926 'ß' * 30,
927 'ß' * 400,
928 'élève',
929 'é' * 30,
930 'é' * 400,
931 'ñ' * 20,
932 'ñ' * 300,
933 ]
934
935 a1 = unique(a, sorted=False, equal_nan=True)
936 # the result varies depending on the impl of std::unordered_set,
937 # so we check them by sorting
938
939 # a1 should have exactly one None
940 count_none = sum(x is None for x in a1)
941 assert_equal(count_none, 1)
942
943 a1_wo_none = sorted(x for x in a1 if x is not None)
944 assert_array_equal(a1_wo_none, unq_sorted_wo_none)
945
946 def test_unique_vstring_hash_based_not_equal_nan(self):
947 # test for unicode and nullable string arrays (equal_nan=False)
948 a = np.array([
949 # short strings
950 'straße',
951 None,
952 'strasse',
953 'straße',
954 None,
955 'niño',
956 'nino',
957 'élève',
958 'eleve',
959 'niño',
960 'élève',
961 # medium strings
962 'b' * 20,
963 'ß' * 30,
964 None,
965 'é' * 30,
966 'e' * 20,
967 'ß' * 30,
968 'n' * 30,
969 'ñ' * 20,
970 None,
971 'e' * 20,
972 'ñ' * 20,
973 # long strings
974 'b' * 300,
975 'ß' * 400,
976 None,
977 'é' * 400,
978 'e' * 300,
979 'ß' * 400,
980 'n' * 400,
981 'ñ' * 300,
982 None,
983 'e' * 300,
984 'ñ' * 300,
985 ],
986 dtype=StringDType(na_object=None)
987 )
988 unq_sorted_wo_none = [
989 'b' * 20,
990 'b' * 300,
991 'e' * 20,
992 'e' * 300,
993 'eleve',
994 'nino',
995 'niño',
996 'n' * 30,
997 'n' * 400,
998 'strasse',
999 'straße',
1000 'ß' * 30,
1001 'ß' * 400,
1002 'élève',
1003 'é' * 30,
1004 'é' * 400,
1005 'ñ' * 20,
1006 'ñ' * 300,
1007 ]
1008
1009 a1 = unique(a, sorted=False, equal_nan=False)
1010 # the result varies depending on the impl of std::unordered_set,
1011 # so we check them by sorting
1012
1013 # a1 should have exactly one None
1014 count_none = sum(x is None for x in a1)
1015 assert_equal(count_none, 6)
1016
1017 a1_wo_none = sorted(x for x in a1 if x is not None)
1018 assert_array_equal(a1_wo_none, unq_sorted_wo_none)
1019
1020 def test_unique_vstring_errors(self):
1021 a = np.array(
1022 [
1023 'apple', 'banana', 'apple', None, 'cherry',
1024 'date', 'banana', 'fig', None, 'grape',
1025 ] * 2,
1026 dtype=StringDType(na_object=None)
1027 )
1028 assert_raises(ValueError, unique, a, equal_nan=False)
1029
1030 @pytest.mark.parametrize("arg", ["return_index", "return_inverse", "return_counts"])
1031 def test_unsupported_hash_based(self, arg):
1032 """These currently never use the hash-based solution. However,
1033 it seems easier to just allow it.
1034
1035 When the hash-based solution is added, this test should fail and be
1036 replaced with something more comprehensive.
1037 """
1038 a = np.array([1, 5, 2, 3, 4, 8, 199, 1, 3, 5])
1039
1040 res_not_sorted = np.unique([1, 1], sorted=False, **{arg: True})
1041 res_sorted = np.unique([1, 1], sorted=True, **{arg: True})
1042 # The following should fail without first sorting `res_not_sorted`.
1043 for arr, expected in zip(res_not_sorted, res_sorted):
1044 assert_array_equal(arr, expected)
1045
1046 def test_unique_axis_errors(self):
1047 assert_raises(TypeError, self._run_axis_tests, object)
1048 assert_raises(TypeError, self._run_axis_tests,
1049 [('a', int), ('b', object)])
1050
1051 assert_raises(AxisError, unique, np.arange(10), axis=2)
1052 assert_raises(AxisError, unique, np.arange(10), axis=-2)
1053
1054 def test_unique_axis_list(self):
1055 msg = "Unique failed on list of lists"
1056 inp = [[0, 1, 0], [0, 1, 0]]
1057 inp_arr = np.asarray(inp)
1058 assert_array_equal(unique(inp, axis=0), unique(inp_arr, axis=0), msg)
1059 assert_array_equal(unique(inp, axis=1), unique(inp_arr, axis=1), msg)
1060
1061 def test_unique_axis(self):
1062 types = []
1063 types.extend(np.typecodes['AllInteger'])
1064 types.extend(np.typecodes['AllFloat'])
1065 types.append('datetime64[D]')
1066 types.append('timedelta64[D]')
1067 types.append([('a', int), ('b', int)])
1068 types.append([('a', int), ('b', float)])
1069
1070 for dtype in types:
1071 self._run_axis_tests(dtype)
1072
1073 msg = 'Non-bitwise-equal booleans test failed'
1074 data = np.arange(10, dtype=np.uint8).reshape(-1, 2).view(bool)
1075 result = np.array([[False, True], [True, True]], dtype=bool)
1076 assert_array_equal(unique(data, axis=0), result, msg)
1077
1078 msg = 'Negative zero equality test failed'
1079 data = np.array([[-0.0, 0.0], [0.0, -0.0], [-0.0, 0.0], [0.0, -0.0]])
1080 result = np.array([[-0.0, 0.0]])
1081 assert_array_equal(unique(data, axis=0), result, msg)
1082
1083 @pytest.mark.parametrize("axis", [0, -1])
1084 def test_unique_1d_with_axis(self, axis):
1085 x = np.array([4, 3, 2, 3, 2, 1, 2, 2])
1086 uniq = unique(x, axis=axis)
1087 assert_array_equal(uniq, [1, 2, 3, 4])
1088
1089 @pytest.mark.parametrize("axis", [None, 0, -1])
1090 def test_unique_inverse_with_axis(self, axis):
1091 x = np.array([[4, 4, 3], [2, 2, 1], [2, 2, 1], [4, 4, 3]])
1092 uniq, inv = unique(x, return_inverse=True, axis=axis)
1093 assert_equal(inv.ndim, x.ndim if axis is None else 1)
1094 assert_array_equal(x, np.take(uniq, inv, axis=axis))
1095
1096 def test_unique_axis_zeros(self):
1097 # issue 15559
1098 single_zero = np.empty(shape=(2, 0), dtype=np.int8)
1099 uniq, idx, inv, cnt = unique(single_zero, axis=0, return_index=True,
1100 return_inverse=True, return_counts=True)
1101
1102 # there's 1 element of shape (0,) along axis 0
1103 assert_equal(uniq.dtype, single_zero.dtype)
1104 assert_array_equal(uniq, np.empty(shape=(1, 0)))
1105 assert_array_equal(idx, np.array([0]))
1106 assert_array_equal(inv, np.array([0, 0]))
1107 assert_array_equal(cnt, np.array([2]))
1108
1109 # there's 0 elements of shape (2,) along axis 1
1110 uniq, idx, inv, cnt = unique(single_zero, axis=1, return_index=True,
1111 return_inverse=True, return_counts=True)
1112
1113 assert_equal(uniq.dtype, single_zero.dtype)
1114 assert_array_equal(uniq, np.empty(shape=(2, 0)))
1115 assert_array_equal(idx, np.array([]))
1116 assert_array_equal(inv, np.array([]))
1117 assert_array_equal(cnt, np.array([]))
1118
1119 # test a "complicated" shape
1120 shape = (0, 2, 0, 3, 0, 4, 0)
1121 multiple_zeros = np.empty(shape=shape)
1122 for axis in range(len(shape)):
1123 expected_shape = list(shape)
1124 if shape[axis] == 0:
1125 expected_shape[axis] = 0
1126 else:
1127 expected_shape[axis] = 1
1128
1129 assert_array_equal(unique(multiple_zeros, axis=axis),
1130 np.empty(shape=expected_shape))
1131
1132 def test_unique_masked(self):
1133 # issue 8664
1134 x = np.array([64, 0, 1, 2, 3, 63, 63, 0, 0, 0, 1, 2, 0, 63, 0],
1135 dtype='uint8')
1136 y = np.ma.masked_equal(x, 0)
1137
1138 v = np.unique(y)
1139 v2, i, c = np.unique(y, return_index=True, return_counts=True)
1140
1141 msg = 'Unique returned different results when asked for index'
1142 assert_array_equal(v.data, v2.data, msg)
1143 assert_array_equal(v.mask, v2.mask, msg)
1144
1145 def test_unique_sort_order_with_axis(self):
1146 # These tests fail if sorting along axis is done by treating subarrays
1147 # as unsigned byte strings. See gh-10495.
1148 fmt = "sort order incorrect for integer type '%s'"
1149 for dt in 'bhilq':
1150 a = np.array([[-1], [0]], dt)
1151 b = np.unique(a, axis=0)
1152 assert_array_equal(a, b, fmt % dt)
1153
1154 def _run_axis_tests(self, dtype):
1155 data = np.array([[0, 1, 0, 0],
1156 [1, 0, 0, 0],
1157 [0, 1, 0, 0],
1158 [1, 0, 0, 0]]).astype(dtype)
1159
1160 msg = 'Unique with 1d array and axis=0 failed'
1161 result = np.array([0, 1])
1162 assert_array_equal(unique(data), result.astype(dtype), msg)
1163
1164 msg = 'Unique with 2d array and axis=0 failed'
1165 result = np.array([[0, 1, 0, 0], [1, 0, 0, 0]])
1166 assert_array_equal(unique(data, axis=0), result.astype(dtype), msg)
1167
1168 msg = 'Unique with 2d array and axis=1 failed'
1169 result = np.array([[0, 0, 1], [0, 1, 0], [0, 0, 1], [0, 1, 0]])
1170 assert_array_equal(unique(data, axis=1), result.astype(dtype), msg)
1171
1172 msg = 'Unique with 3d array and axis=2 failed'
1173 data3d = np.array([[[1, 1],
1174 [1, 0]],
1175 [[0, 1],
1176 [0, 0]]]).astype(dtype)
1177 result = np.take(data3d, [1, 0], axis=2)
1178 assert_array_equal(unique(data3d, axis=2), result, msg)
1179
1180 uniq, idx, inv, cnt = unique(data, axis=0, return_index=True,
1181 return_inverse=True, return_counts=True)
1182 msg = "Unique's return_index=True failed with axis=0"
1183 assert_array_equal(data[idx], uniq, msg)
1184 msg = "Unique's return_inverse=True failed with axis=0"
1185 assert_array_equal(np.take(uniq, inv, axis=0), data)
1186 msg = "Unique's return_counts=True failed with axis=0"
1187 assert_array_equal(cnt, np.array([2, 2]), msg)
1188
1189 uniq, idx, inv, cnt = unique(data, axis=1, return_index=True,
1190 return_inverse=True, return_counts=True)
1191 msg = "Unique's return_index=True failed with axis=1"
1192 assert_array_equal(data[:, idx], uniq)
1193 msg = "Unique's return_inverse=True failed with axis=1"
1194 assert_array_equal(np.take(uniq, inv, axis=1), data)
1195 msg = "Unique's return_counts=True failed with axis=1"
1196 assert_array_equal(cnt, np.array([2, 1, 1]), msg)
1197
1198 def test_unique_nanequals(self):
1199 # issue 20326
1200 a = np.array([1, 1, np.nan, np.nan, np.nan])