JezK
Edit File: test_arraysetops.py
"""Test functions for 1D array set operations. """ import numpy as np from numpy.testing import (assert_array_equal, assert_equal, assert_raises, assert_raises_regex) from numpy.lib.arraysetops import ( ediff1d, intersect1d, setxor1d, union1d, setdiff1d, unique, in1d, isin ) import pytest class TestSetOps: def test_intersect1d(self): # unique inputs a = np.array([5, 7, 1, 2]) b = np.array([2, 4, 3, 1, 5]) ec = np.array([1, 2, 5]) c = intersect1d(a, b, assume_unique=True) assert_array_equal(c, ec) # non-unique inputs a = np.array([5, 5, 7, 1, 2]) b = np.array([2, 1, 4, 3, 3, 1, 5]) ed = np.array([1, 2, 5]) c = intersect1d(a, b) assert_array_equal(c, ed) assert_array_equal([], intersect1d([], [])) def test_intersect1d_array_like(self): # See gh-11772 class Test: def __array__(self): return np.arange(3) a = Test() res = intersect1d(a, a) assert_array_equal(res, a) res = intersect1d([1, 2, 3], [1, 2, 3]) assert_array_equal(res, [1, 2, 3]) def test_intersect1d_indices(self): # unique inputs a = np.array([1, 2, 3, 4]) b = np.array([2, 1, 4, 6]) c, i1, i2 = intersect1d(a, b, assume_unique=True, return_indices=True) ee = np.array([1, 2, 4]) assert_array_equal(c, ee) assert_array_equal(a[i1], ee) assert_array_equal(b[i2], ee) # non-unique inputs a = np.array([1, 2, 2, 3, 4, 3, 2]) b = np.array([1, 8, 4, 2, 2, 3, 2, 3]) c, i1, i2 = intersect1d(a, b, return_indices=True) ef = np.array([1, 2, 3, 4]) assert_array_equal(c, ef) assert_array_equal(a[i1], ef) assert_array_equal(b[i2], ef) # non1d, unique inputs a = np.array([[2, 4, 5, 6], [7, 8, 1, 15]]) b = np.array([[3, 2, 7, 6], [10, 12, 8, 9]]) c, i1, i2 = intersect1d(a, b, assume_unique=True, return_indices=True) ui1 = np.unravel_index(i1, a.shape) ui2 = np.unravel_index(i2, b.shape) ea = np.array([2, 6, 7, 8]) assert_array_equal(ea, a[ui1]) assert_array_equal(ea, b[ui2]) # non1d, not assumed to be uniqueinputs a = np.array([[2, 4, 5, 6, 6], [4, 7, 8, 7, 2]]) b = np.array([[3, 2, 7, 7], [10, 12, 8, 7]]) c, i1, i2 = intersect1d(a, b, return_indices=True) ui1 = np.unravel_index(i1, a.shape) ui2 = np.unravel_index(i2, b.shape) ea = np.array([2, 7, 8]) assert_array_equal(ea, a[ui1]) assert_array_equal(ea, b[ui2]) def test_setxor1d(self): a = np.array([5, 7, 1, 2]) b = np.array([2, 4, 3, 1, 5]) ec = np.array([3, 4, 7]) c = setxor1d(a, b) assert_array_equal(c, ec) a = np.array([1, 2, 3]) b = np.array([6, 5, 4]) ec = np.array([1, 2, 3, 4, 5, 6]) c = setxor1d(a, b) assert_array_equal(c, ec) a = np.array([1, 8, 2, 3]) b = np.array([6, 5, 4, 8]) ec = np.array([1, 2, 3, 4, 5, 6]) c = setxor1d(a, b) assert_array_equal(c, ec) assert_array_equal([], setxor1d([], [])) def test_ediff1d(self): zero_elem = np.array([]) one_elem = np.array([1]) two_elem = np.array([1, 2]) assert_array_equal([], ediff1d(zero_elem)) assert_array_equal([0], ediff1d(zero_elem, to_begin=0)) assert_array_equal([0], ediff1d(zero_elem, to_end=0)) assert_array_equal([-1, 0], ediff1d(zero_elem, to_begin=-1, to_end=0)) assert_array_equal([], ediff1d(one_elem)) assert_array_equal([1], ediff1d(two_elem)) assert_array_equal([7, 1, 9], ediff1d(two_elem, to_begin=7, to_end=9)) assert_array_equal([5, 6, 1, 7, 8], ediff1d(two_elem, to_begin=[5, 6], to_end=[7, 8])) assert_array_equal([1, 9], ediff1d(two_elem, to_end=9)) assert_array_equal([1, 7, 8], ediff1d(two_elem, to_end=[7, 8])) assert_array_equal([7, 1], ediff1d(two_elem, to_begin=7)) assert_array_equal([5, 6, 1], ediff1d(two_elem, to_begin=[5, 6])) @pytest.mark.parametrize("ary, prepend, append, expected", [ # should fail because trying to cast # np.nan standard floating point value # into an integer array: (np.array([1, 2, 3], dtype=np.int64), None, np.nan, 'to_end'), # should fail because attempting # to downcast to int type: (np.array([1, 2, 3], dtype=np.int64), np.array([5, 7, 2], dtype=np.float32), None, 'to_begin'), # should fail because attempting to cast # two special floating point values # to integers (on both sides of ary), # `to_begin` is in the error message as the impl checks this first: (np.array([1., 3., 9.], dtype=np.int8), np.nan, np.nan, 'to_begin'), ]) def test_ediff1d_forbidden_type_casts(self, ary, prepend, append, expected): # verify resolution of gh-11490 # specifically, raise an appropriate # Exception when attempting to append or # prepend with an incompatible type msg = 'dtype of `{}` must be compatible'.format(expected) with assert_raises_regex(TypeError, msg): ediff1d(ary=ary, to_end=append, to_begin=prepend) @pytest.mark.parametrize( "ary,prepend,append,expected", [ (np.array([1, 2, 3], dtype=np.int16), 2**16, # will be cast to int16 under same kind rule. 2**16 + 4, np.array([0, 1, 1, 4], dtype=np.int16)), (np.array([1, 2, 3], dtype=np.float32), np.array([5], dtype=np.float64), None, np.array([5, 1, 1], dtype=np.float32)), (np.array([1, 2, 3], dtype=np.int32), 0, 0, np.array([0, 1, 1, 0], dtype=np.int32)), (np.array([1, 2, 3], dtype=np.int64), 3, -9, np.array([3, 1, 1, -9], dtype=np.int64)), ] ) def test_ediff1d_scalar_handling(self, ary, prepend, append, expected): # maintain backwards-compatibility # of scalar prepend / append behavior # in ediff1d following fix for gh-11490 actual = np.ediff1d(ary=ary, to_end=append, to_begin=prepend) assert_equal(actual, expected) assert actual.dtype == expected.dtype @pytest.mark.parametrize("kind", [None, "sort", "table"]) def test_isin(self, kind): # the tests for in1d cover most of isin's behavior # if in1d is removed, would need to change those tests to test # isin instead. def _isin_slow(a, b): b = np.asarray(b).flatten().tolist() return a in b isin_slow = np.vectorize(_isin_slow, otypes=[bool], excluded={1}) def assert_isin_equal(a, b): x = isin(a, b, kind=kind) y = isin_slow(a, b) assert_array_equal(x, y) # multidimensional arrays in both arguments a = np.arange(24).reshape([2, 3, 4]) b = np.array([[10, 20, 30], [0, 1, 3], [11, 22, 33]]) assert_isin_equal(a, b) # array-likes as both arguments c = [(9, 8), (7, 6)] d = (9, 7) assert_isin_equal(c, d) # zero-d array: f = np.array(3) assert_isin_equal(f, b) assert_isin_equal(a, f) assert_isin_equal(f, f) # scalar: assert_isin_equal(5, b) assert_isin_equal(a, 6) assert_isin_equal(5, 6) # empty array-like: if kind != "table": # An empty list will become float64, # which is invalid for kind="table" x = [] assert_isin_equal(x, b) assert_isin_equal(a, x) assert_isin_equal(x, x) # empty array with various types: for dtype in [bool, np.int64, np.float64]: if kind == "table" and dtype == np.float64: continue if dtype in {np.int64, np.float64}: ar = np.array([10, 20, 30], dtype=dtype) elif dtype in {bool}: ar = np.array([True, False, False]) empty_array = np.array([], dtype=dtype) assert_isin_equal(empty_array, ar) assert_isin_equal(ar, empty_array) assert_isin_equal(empty_array, empty_array) @pytest.mark.parametrize("kind", [None, "sort", "table"]) def test_in1d(self, kind): # we use two different sizes for the b array here to test the # two different paths in in1d(). for mult in (1, 10): # One check without np.array to make sure lists are handled correct a = [5, 7, 1, 2] b = [2, 4, 3, 1, 5] * mult ec = np.array([True, False, True, True]) c = in1d(a, b, assume_unique=True, kind=kind) assert_array_equal(c, ec) a[0] = 8 ec = np.array([False, False, True, True]) c = in1d(a, b, assume_unique=True, kind=kind) assert_array_equal(c, ec) a[0], a[3] = 4, 8 ec = np.array([True, False, True, False]) c = in1d(a, b, assume_unique=True, kind=kind) assert_array_equal(c, ec) a = np.array([5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5]) b = [2, 3, 4] * mult ec = [False, True, False, True, True, True, True, True, True, False, True, False, False, False] c = in1d(a, b, kind=kind) assert_array_equal(c, ec) b = b + [5, 5, 4] * mult ec = [True, True, True, True, True, True, True, True, True, True, True, False, True, True] c = in1d(a, b, kind=kind) assert_array_equal(c, ec) a = np.array([5, 7, 1, 2]) b = np.array([2, 4, 3, 1, 5] * mult) ec = np.array([True, False, True, True]) c = in1d(a, b, kind=kind) assert_array_equal(c, ec) a = np.array([5, 7, 1, 1, 2]) b = np.array([2, 4, 3, 3, 1, 5] * mult) ec = np.array([True, False, True, True, True]) c = in1d(a, b, kind=kind) assert_array_equal(c, ec) a = np.array([5, 5]) b = np.array([2, 2] * mult) ec = np.array([False, False]) c = in1d(a, b, kind=kind) assert_array_equal(c, ec) a = np.array([5]) b = np.array([2]) ec = np.array([False]) c = in1d(a, b, kind=kind) assert_array_equal(c, ec) if kind in {None, "sort"}: assert_array_equal(in1d([], [], kind=kind), []) def test_in1d_char_array(self): a = np.array(['a', 'b', 'c', 'd', 'e', 'c', 'e', 'b']) b = np.array(['a', 'c']) ec = np.array([True, False, True, False, False, True, False, False]) c = in1d(a, b) assert_array_equal(c, ec) @pytest.mark.parametrize("kind", [None, "sort", "table"]) def test_in1d_invert(self, kind): "Test in1d's invert parameter" # We use two different sizes for the b array here to test the # two different paths in in1d(). for mult in (1, 10): a = np.array([5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5]) b = [2, 3, 4] * mult assert_array_equal(np.invert(in1d(a, b, kind=kind)), in1d(a, b, invert=True, kind=kind)) # float: if kind in {None, "sort"}: for mult in (1, 10): a = np.array([5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5], dtype=np.float32) b = [2, 3, 4] * mult b = np.array(b, dtype=np.float32) assert_array_equal(np.invert(in1d(a, b, kind=kind)), in1d(a, b, invert=True, kind=kind)) @pytest.mark.parametrize("kind", [None, "sort", "table"]) def test_in1d_ravel(self, kind): # Test that in1d ravels its input arrays. This is not documented # behavior however. The test is to ensure consistentency. a = np.arange(6).reshape(2, 3) b = np.arange(3, 9).reshape(3, 2) long_b = np.arange(3, 63).reshape(30, 2) ec = np.array([False, False, False, True, True, True]) assert_array_equal(in1d(a, b, assume_unique=True, kind=kind), ec) assert_array_equal(in1d(a, b, assume_unique=False, kind=kind), ec) assert_array_equal(in1d(a, long_b, assume_unique=True, kind=kind), ec) assert_array_equal(in1d(a, long_b, assume_unique=False, kind=kind), ec) def test_in1d_hit_alternate_algorithm(self): """Hit the standard isin code with integers""" # Need extreme range to hit standard code # This hits it without the use of kind='table' a = np.array([5, 4, 5, 3, 4, 4, 1e9], dtype=np.int64) b = np.array([2, 3, 4, 1e9], dtype=np.int64) expected = np.array([0, 1, 0, 1, 1, 1, 1], dtype=bool) assert_array_equal(expected, in1d(a, b)) assert_array_equal(np.invert(expected), in1d(a, b, invert=True)) a = np.array([5, 7, 1, 2], dtype=np.int64) b = np.array([2, 4, 3, 1, 5, 1e9], dtype=np.int64) ec = np.array([True, False, True, True]) c = in1d(a, b, assume_unique=True) assert_array_equal(c, ec) @pytest.mark.parametrize("kind", [None, "sort", "table"]) def test_in1d_boolean(self, kind): """Test that in1d works for boolean input""" a = np.array([True, False]) b = np.array([False, False, False]) expected = np.array([False, True]) assert_array_equal(expected, in1d(a, b, kind=kind)) assert_array_equal(np.invert(expected), in1d(a, b, invert=True, kind=kind)) @pytest.mark.parametrize("kind", [None, "sort"]) def test_in1d_timedelta(self, kind): """Test that in1d works for timedelta input""" rstate = np.random.RandomState(0) a = rstate.randint(0, 100, size=10) b = rstate.randint(0, 100, size=10) truth = in1d(a, b) a_timedelta = a.astype("timedelta64[s]") b_timedelta = b.astype("timedelta64[s]") assert_array_equal(truth, in1d(a_timedelta, b_timedelta, kind=kind)) def test_in1d_table_timedelta_fails(self): a = np.array([0, 1, 2], dtype="timedelta64[s]") b = a # Make sure it raises a value error: with pytest.raises(ValueError): in1d(a, b, kind="table") @pytest.mark.parametrize( "dtype1,dtype2", [ (np.int8, np.int16), (np.int16, np.int8), (np.uint8, np.uint16), (np.uint16, np.uint8), (np.uint8, np.int16), (np.int16, np.uint8), ] ) @pytest.mark.parametrize("kind", [None, "sort", "table"]) def test_in1d_mixed_dtype(self, dtype1, dtype2, kind): """Test that in1d works as expected for mixed dtype input.""" is_dtype2_signed = np.issubdtype(dtype2, np.signedinteger) ar1 = np.array([0, 0, 1, 1], dtype=dtype1) if is_dtype2_signed: ar2 = np.array([-128, 0, 127], dtype=dtype2) else: ar2 = np.array([127, 0, 255], dtype=dtype2) expected = np.array([True, True, False, False]) expect_failure = kind == "table" and any(( dtype1 == np.int8 and dtype2 == np.int16, dtype1 == np.int16 and dtype2 == np.int8 )) if expect_failure: with pytest.raises(RuntimeError, match="exceed the maximum"): in1d(ar1, ar2, kind=kind) else: assert_array_equal(in1d(ar1, ar2, kind=kind), expected) @pytest.mark.parametrize("kind", [None, "sort", "table"]) def test_in1d_mixed_boolean(self, kind): """Test that in1d works as expected for bool/int input.""" for dtype in np.typecodes["AllInteger"]: a = np.array([True, False, False], dtype=bool) b = np.array([0, 0, 0, 0], dtype=dtype) expected = np.array([False, True, True], dtype=bool) assert_array_equal(in1d(a, b, kind=kind), expected) a, b = b, a expected = np.array([True, True, True, True], dtype=bool) assert_array_equal(in1d(a, b, kind=kind), expected) def test_in1d_first_array_is_object(self): ar1 = [None] ar2 = np.array([1]*10) expected = np.array([False]) result = np.in1d(ar1, ar2) assert_array_equal(result, expected) def test_in1d_second_array_is_object(self): ar1 = 1 ar2 = np.array([None]*10) expected = np.array([False]) result = np.in1d(ar1, ar2) assert_array_equal(result, expected) def test_in1d_both_arrays_are_object(self): ar1 = [None] ar2 = np.array([None]*10) expected = np.array([True]) result = np.in1d(ar1, ar2) assert_array_equal(result, expected) def test_in1d_both_arrays_have_structured_dtype(self): # Test arrays of a structured data type containing an integer field # and a field of dtype `object` allowing for arbitrary Python objects dt = np.dtype([('field1', int), ('field2', object)]) ar1 = np.array([(1, None)], dtype=dt) ar2 = np.array([(1, None)]*10, dtype=dt) expected = np.array([True]) result = np.in1d(ar1, ar2) assert_array_equal(result, expected) def test_in1d_with_arrays_containing_tuples(self): ar1 = np.array([(1,), 2], dtype=object) ar2 = np.array([(1,), 2], dtype=object) expected = np.array([True, True]) result = np.in1d(ar1, ar2) assert_array_equal(result, expected) result = np.in1d(ar1, ar2, invert=True) assert_array_equal(result, np.invert(expected)) # An integer is added at the end of the array to make sure # that the array builder will create the array with tuples # and after it's created the integer is removed. # There's a bug in the array constructor that doesn't handle # tuples properly and adding the integer fixes that. ar1 = np.array([(1,), (2, 1), 1], dtype=object) ar1 = ar1[:-1] ar2 = np.array([(1,), (2, 1), 1], dtype=object) ar2 = ar2[:-1] expected = np.array([True, True]) result = np.in1d(ar1, ar2) assert_array_equal(result, expected) result = np.in1d(ar1, ar2, invert=True) assert_array_equal(result, np.invert(expected)) ar1 = np.array([(1,), (2, 3), 1], dtype=object) ar1 = ar1[:-1] ar2 = np.array([(1,), 2], dtype=object) expected = np.array([True, False]) result = np.in1d(ar1, ar2) assert_array_equal(result, expected) result = np.in1d(ar1, ar2, invert=True) assert_array_equal(result, np.invert(expected)) def test_in1d_errors(self): """Test that in1d raises expected errors.""" # Error 1: `kind` is not one of 'sort' 'table' or None. ar1 = np.array([1, 2, 3, 4, 5]) ar2 = np.array([2, 4, 6, 8, 10]) assert_raises(ValueError, in1d, ar1, ar2, kind='quicksort') # Error 2: `kind="table"` does not work for non-integral arrays. obj_ar1 = np.array([1, 'a', 3, 'b', 5], dtype=object) obj_ar2 = np.array([1, 'a', 3, 'b', 5], dtype=object) assert_raises(ValueError, in1d, obj_ar1, obj_ar2, kind='table') for dtype in [np.int32, np.int64]: ar1 = np.array([-1, 2, 3, 4, 5], dtype=dtype) # The range of this array will overflow: overflow_ar2 = np.array([-1, np.iinfo(dtype).max], dtype=dtype) # Error 3: `kind="table"` will trigger a runtime error # if there is an integer overflow expected when computing the # range of ar2 assert_raises( RuntimeError, in1d, ar1, overflow_ar2, kind='table' ) # Non-error: `kind=None` will *not* trigger a runtime error # if there is an integer overflow, it will switch to # the `sort` algorithm. result = np.in1d(ar1, overflow_ar2, kind=None) assert_array_equal(result, [True] + [False] * 4) result = np.in1d(ar1, overflow_ar2, kind='sort') assert_array_equal(result, [True] + [False] * 4) def test_union1d(self): a = np.array([5, 4, 7, 1, 2]) b = np.array([2, 4, 3, 3, 2, 1, 5]) ec = np.array([1, 2, 3, 4, 5, 7]) c = union1d(a, b) assert_array_equal(c, ec) # Tests gh-10340, arguments to union1d should be # flattened if they are not already 1D x = np.array([[0, 1, 2], [3, 4, 5]]) y = np.array([0, 1, 2, 3, 4]) ez = np.array([0, 1, 2, 3, 4, 5]) z = union1d(x, y) assert_array_equal(z, ez) assert_array_equal([], union1d([], [])) def test_setdiff1d(self): a = np.array([6, 5, 4, 7, 1, 2, 7, 4]) b = np.array([2, 4, 3, 3, 2, 1, 5]) ec = np.array([6, 7]) c = setdiff1d(a, b) assert_array_equal(c, ec) a = np.arange(21) b = np.arange(19) ec = np.array([19, 20]) c = setdiff1d(a, b) assert_array_equal(c, ec) assert_array_equal([], setdiff1d([], [])) a = np.array((), np.uint32) assert_equal(setdiff1d(a, []).dtype, np.uint32) def test_setdiff1d_unique(self): a = np.array([3, 2, 1]) b = np.array([7, 5, 2]) expected = np.array([3, 1]) actual = setdiff1d(a, b, assume_unique=True) assert_equal(actual, expected) def test_setdiff1d_char_array(self): a = np.array(['a', 'b', 'c']) b = np.array(['a', 'b', 's']) assert_array_equal(setdiff1d(a, b), np.array(['c'])) def test_manyways(self): a = np.array([5, 7, 1, 2, 8]) b = np.array([9, 8, 2, 4, 3, 1, 5]) c1 = setxor1d(a, b) aux1 = intersect1d(a, b) aux2 = union1d(a, b) c2 = setdiff1d(aux2, aux1) assert_array_equal(c1, c2) class TestUnique: def test_unique_1d(self): def check_all(a, b, i1, i2, c, dt): base_msg = 'check {0} failed for type {1}' msg = base_msg.format('values', dt) v = unique(a) assert_array_equal(v, b, msg) msg = base_msg.format('return_index', dt) v, j = unique(a, True, False, False) assert_array_equal(v, b, msg) assert_array_equal(j, i1, msg) msg = base_msg.format('return_inverse', dt) v, j = unique(a, False, True, False) assert_array_equal(v, b, msg) assert_array_equal(j, i2, msg) msg = base_msg.format('return_counts', dt) v, j = unique(a, False, False, True) assert_array_equal(v, b, msg) assert_array_equal(j, c, msg) msg = base_msg.format('return_index and return_inverse', dt) v, j1, j2 = unique(a, True, True, False) assert_array_equal(v, b, msg) assert_array_equal(j1, i1, msg) assert_array_equal(j2, i2, msg) msg = base_msg.format('return_index and return_counts', dt) v, j1, j2 = unique(a, True, False, True) assert_array_equal(v, b, msg) assert_array_equal(j1, i1, msg) assert_array_equal(j2, c, msg) msg = base_msg.format('return_inverse and return_counts', dt) v, j1, j2 = unique(a, False, True, True) assert_array_equal(v, b, msg) assert_array_equal(j1, i2, msg) assert_array_equal(j2, c, msg) msg = base_msg.format(('return_index, return_inverse ' 'and return_counts'), dt) v, j1, j2, j3 = unique(a, True, True, True) assert_array_equal(v, b, msg) assert_array_equal(j1, i1, msg) assert_array_equal(j2, i2, msg) assert_array_equal(j3, c, msg) a = [5, 7, 1, 2, 1, 5, 7]*10 b = [1, 2, 5, 7] i1 = [2, 3, 0, 1] i2 = [2, 3, 0, 1, 0, 2, 3]*10 c = np.multiply([2, 1, 2, 2], 10) # test for numeric arrays types = [] types.extend(np.typecodes['AllInteger']) types.extend(np.typecodes['AllFloat']) types.append('datetime64[D]') types.append('timedelta64[D]') for dt in types: aa = np.array(a, dt) bb = np.array(b, dt) check_all(aa, bb, i1, i2, c, dt) # test for object arrays dt = 'O' aa = np.empty(len(a), dt) aa[:] = a bb = np.empty(len(b), dt) bb[:] = b check_all(aa, bb, i1, i2, c, dt) # test for structured arrays dt = [('', 'i'), ('', 'i')] aa = np.array(list(zip(a, a)), dt) bb = np.array(list(zip(b, b)), dt) check_all(aa, bb, i1, i2, c, dt) # test for ticket #2799 aa = [1. + 0.j, 1 - 1.j, 1] assert_array_equal(np.unique(aa), [1. - 1.j, 1. + 0.j]) # test for ticket #4785 a = [(1, 2), (1, 2), (2, 3)] unq = [1, 2, 3] inv = [0, 1, 0, 1, 1, 2] a1 = unique(a) assert_array_equal(a1, unq) a2, a2_inv = unique(a, return_inverse=True) assert_array_equal(a2, unq) assert_array_equal(a2_inv, inv) # test for chararrays with return_inverse (gh-5099) a = np.chararray(5) a[...] = '' a2, a2_inv = np.unique(a, return_inverse=True) assert_array_equal(a2_inv, np.zeros(5)) # test for ticket #9137 a = [] a1_idx = np.unique(a, return_index=True)[1] a2_inv = np.unique(a, return_inverse=True)[1] a3_idx, a3_inv = np.unique(a, return_index=True, return_inverse=True)[1:] assert_equal(a1_idx.dtype, np.intp) assert_equal(a2_inv.dtype, np.intp) assert_equal(a3_idx.dtype, np.intp) assert_equal(a3_inv.dtype, np.intp) # test for ticket 2111 - float a = [2.0, np.nan, 1.0, np.nan] ua = [1.0, 2.0, np.nan] ua_idx = [2, 0, 1] ua_inv = [1, 2, 0, 2] ua_cnt = [1, 1, 2] assert_equal(np.unique(a), ua) assert_equal(np.unique(a, return_index=True), (ua, ua_idx)) assert_equal(np.unique(a, return_inverse=True), (ua, ua_inv)) assert_equal(np.unique(a, return_counts=True), (ua, ua_cnt)) # test for ticket 2111 - complex a = [2.0-1j, np.nan, 1.0+1j, complex(0.0, np.nan), complex(1.0, np.nan)] ua = [1.0+1j, 2.0-1j, complex(0.0, np.nan)] ua_idx = [2, 0, 3] ua_inv = [1, 2, 0, 2, 2] ua_cnt = [1, 1, 3] assert_equal(np.unique(a), ua) assert_equal(np.unique(a, return_index=True), (ua, ua_idx)) assert_equal(np.unique(a, return_inverse=True), (ua, ua_inv)) assert_equal(np.unique(a, return_counts=True), (ua, ua_cnt)) # test for ticket 2111 - datetime64 nat = np.datetime64('nat') a = [np.datetime64('2020-12-26'), nat, np.datetime64('2020-12-24'), nat] ua = [np.datetime64('2020-12-24'), np.datetime64('2020-12-26'), nat] ua_idx = [2, 0, 1] ua_inv = [1, 2, 0, 2] ua_cnt = [1, 1, 2] assert_equal(np.unique(a), ua) assert_equal(np.unique(a, return_index=True), (ua, ua_idx)) assert_equal(np.unique(a, return_inverse=True), (ua, ua_inv)) assert_equal(np.unique(a, return_counts=True), (ua, ua_cnt)) # test for ticket 2111 - timedelta nat = np.timedelta64('nat') a = [np.timedelta64(1, 'D'), nat, np.timedelta64(1, 'h'), nat] ua = [np.timedelta64(1, 'h'), np.timedelta64(1, 'D'), nat] ua_idx = [2, 0, 1] ua_inv = [1, 2, 0, 2] ua_cnt = [1, 1, 2] assert_equal(np.unique(a), ua) assert_equal(np.unique(a, return_index=True), (ua, ua_idx)) assert_equal(np.unique(a, return_inverse=True), (ua, ua_inv)) assert_equal(np.unique(a, return_counts=True), (ua, ua_cnt)) # test for gh-19300 all_nans = [np.nan] * 4 ua = [np.nan] ua_idx = [0] ua_inv = [0, 0, 0, 0] ua_cnt = [4] assert_equal(np.unique(all_nans), ua) assert_equal(np.unique(all_nans, return_index=True), (ua, ua_idx)) assert_equal(np.unique(all_nans, return_inverse=True), (ua, ua_inv)) assert_equal(np.unique(all_nans, return_counts=True), (ua, ua_cnt)) def test_unique_axis_errors(self): assert_raises(TypeError, self._run_axis_tests, object) assert_raises(TypeError, self._run_axis_tests, [('a', int), ('b', object)]) assert_raises(np.AxisError, unique, np.arange(10), axis=2) assert_raises(np.AxisError, unique, np.arange(10), axis=-2) def test_unique_axis_list(self): msg = "Unique failed on list of lists" inp = [[0, 1, 0], [0, 1, 0]] inp_arr = np.asarray(inp) assert_array_equal(unique(inp, axis=0), unique(inp_arr, axis=0), msg) assert_array_equal(unique(inp, axis=1), unique(inp_arr, axis=1), msg) def test_unique_axis(self): types = [] types.extend(np.typecodes['AllInteger']) types.extend(np.typecodes['AllFloat']) types.append('datetime64[D]') types.append('timedelta64[D]') types.append([('a', int), ('b', int)]) types.append([('a', int), ('b', float)]) for dtype in types: self._run_axis_tests(dtype) msg = 'Non-bitwise-equal booleans test failed' data = np.arange(10, dtype=np.uint8).reshape(-1, 2).view(bool) result = np.array([[False, True], [True, True]], dtype=bool) assert_array_equal(unique(data, axis=0), result, msg) msg = 'Negative zero equality test failed' data = np.array([[-0.0, 0.0], [0.0, -0.0], [-0.0, 0.0], [0.0, -0.0]]) result = np.array([[-0.0, 0.0]]) assert_array_equal(unique(data, axis=0), result, msg) @pytest.mark.parametrize("axis", [0, -1]) def test_unique_1d_with_axis(self, axis): x = np.array([4, 3, 2, 3, 2, 1, 2, 2]) uniq = unique(x, axis=axis) assert_array_equal(uniq, [1, 2, 3, 4]) def test_unique_axis_zeros(self): # issue 15559 single_zero = np.empty(shape=(2, 0), dtype=np.int8) uniq, idx, inv, cnt = unique(single_zero, axis=0, return_index=True, return_inverse=True, return_counts=True) # there's 1 element of shape (0,) along axis 0 assert_equal(uniq.dtype, single_zero.dtype) assert_array_equal(uniq, np.empty(shape=(1, 0))) assert_array_equal(idx, np.array([0])) assert_array_equal(inv, np.array([0, 0])) assert_array_equal(cnt, np.array([2])) # there's 0 elements of shape (2,) along axis 1 uniq, idx, inv, cnt = unique(single_zero, axis=1, return_index=True, return_inverse=True, return_counts=True) assert_equal(uniq.dtype, single_zero.dtype) assert_array_equal(uniq, np.empty(shape=(2, 0))) assert_array_equal(idx, np.array([])) assert_array_equal(inv, np.array([])) assert_array_equal(cnt, np.array([])) # test a "complicated" shape shape = (0, 2, 0, 3, 0, 4, 0) multiple_zeros = np.empty(shape=shape) for axis in range(len(shape)): expected_shape = list(shape) if shape[axis] == 0: expected_shape[axis] = 0 else: expected_shape[axis] = 1 assert_array_equal(unique(multiple_zeros, axis=axis), np.empty(shape=expected_shape)) def test_unique_masked(self): # issue 8664 x = np.array([64, 0, 1, 2, 3, 63, 63, 0, 0, 0, 1, 2, 0, 63, 0], dtype='uint8') y = np.ma.masked_equal(x, 0) v = np.unique(y) v2, i, c = np.unique(y, return_index=True, return_counts=True) msg = 'Unique returned different results when asked for index' assert_array_equal(v.data, v2.data, msg) assert_array_equal(v.mask, v2.mask, msg) def test_unique_sort_order_with_axis(self): # These tests fail if sorting along axis is done by treating subarrays # as unsigned byte strings. See gh-10495. fmt = "sort order incorrect for integer type '%s'" for dt in 'bhilq': a = np.array([[-1], [0]], dt) b = np.unique(a, axis=0) assert_array_equal(a, b, fmt % dt) def _run_axis_tests(self, dtype): data = np.array([[0, 1, 0, 0], [1, 0, 0, 0], [0, 1, 0, 0], [1, 0, 0, 0]]).astype(dtype) msg = 'Unique with 1d array and axis=0 failed' result = np.array([0, 1]) assert_array_equal(unique(data), result.astype(dtype), msg) msg = 'Unique with 2d array and axis=0 failed' result = np.array([[0, 1, 0, 0], [1, 0, 0, 0]]) assert_array_equal(unique(data, axis=0), result.astype(dtype), msg) msg = 'Unique with 2d array and axis=1 failed' result = np.array([[0, 0, 1], [0, 1, 0], [0, 0, 1], [0, 1, 0]]) assert_array_equal(unique(data, axis=1), result.astype(dtype), msg) msg = 'Unique with 3d array and axis=2 failed' data3d = np.array([[[1, 1], [1, 0]], [[0, 1], [0, 0]]]).astype(dtype) result = np.take(data3d, [1, 0], axis=2) assert_array_equal(unique(data3d, axis=2), result, msg) uniq, idx, inv, cnt = unique(data, axis=0, return_index=True, return_inverse=True, return_counts=True) msg = "Unique's return_index=True failed with axis=0" assert_array_equal(data[idx], uniq, msg) msg = "Unique's return_inverse=True failed with axis=0" assert_array_equal(uniq[inv], data) msg = "Unique's return_counts=True failed with axis=0" assert_array_equal(cnt, np.array([2, 2]), msg) uniq, idx, inv, cnt = unique(data, axis=1, return_index=True, return_inverse=True, return_counts=True) msg = "Unique's return_index=True failed with axis=1" assert_array_equal(data[:, idx], uniq) msg = "Unique's return_inverse=True failed with axis=1" assert_array_equal(uniq[:, inv], data) msg = "Unique's return_counts=True failed with axis=1" assert_array_equal(cnt, np.array([2, 1, 1]), msg) def test_unique_nanequals(self): # issue 20326 a = np.array([1, 1, np.nan, np.nan, np.nan]) unq = np.unique(a) not_unq = np.unique(a, equal_nan=False) assert_array_equal(unq, np.array([1, np.nan])) assert_array_equal(not_unq, np.array([1, np.nan, np.nan, np.nan]))