Chris@87: """Test functions for 1D array set operations. Chris@87: Chris@87: """ Chris@87: from __future__ import division, absolute_import, print_function Chris@87: Chris@87: import numpy as np Chris@87: from numpy.testing import ( Chris@87: run_module_suite, TestCase, assert_array_equal Chris@87: ) Chris@87: from numpy.lib.arraysetops import ( Chris@87: ediff1d, intersect1d, setxor1d, union1d, setdiff1d, unique, in1d Chris@87: ) Chris@87: Chris@87: Chris@87: class TestSetOps(TestCase): Chris@87: Chris@87: def test_unique(self): Chris@87: Chris@87: def check_all(a, b, i1, i2, c, dt): Chris@87: base_msg = 'check {0} failed for type {1}' Chris@87: Chris@87: msg = base_msg.format('values', dt) Chris@87: v = unique(a) Chris@87: assert_array_equal(v, b, msg) Chris@87: Chris@87: msg = base_msg.format('return_index', dt) Chris@87: v, j = unique(a, 1, 0, 0) Chris@87: assert_array_equal(v, b, msg) Chris@87: assert_array_equal(j, i1, msg) Chris@87: Chris@87: msg = base_msg.format('return_inverse', dt) Chris@87: v, j = unique(a, 0, 1, 0) Chris@87: assert_array_equal(v, b, msg) Chris@87: assert_array_equal(j, i2, msg) Chris@87: Chris@87: msg = base_msg.format('return_counts', dt) Chris@87: v, j = unique(a, 0, 0, 1) Chris@87: assert_array_equal(v, b, msg) Chris@87: assert_array_equal(j, c, msg) Chris@87: Chris@87: msg = base_msg.format('return_index and return_inverse', dt) Chris@87: v, j1, j2 = unique(a, 1, 1, 0) Chris@87: assert_array_equal(v, b, msg) Chris@87: assert_array_equal(j1, i1, msg) Chris@87: assert_array_equal(j2, i2, msg) Chris@87: Chris@87: msg = base_msg.format('return_index and return_counts', dt) Chris@87: v, j1, j2 = unique(a, 1, 0, 1) Chris@87: assert_array_equal(v, b, msg) Chris@87: assert_array_equal(j1, i1, msg) Chris@87: assert_array_equal(j2, c, msg) Chris@87: Chris@87: msg = base_msg.format('return_inverse and return_counts', dt) Chris@87: v, j1, j2 = unique(a, 0, 1, 1) Chris@87: assert_array_equal(v, b, msg) Chris@87: assert_array_equal(j1, i2, msg) Chris@87: assert_array_equal(j2, c, msg) Chris@87: Chris@87: msg = base_msg.format(('return_index, return_inverse ' Chris@87: 'and return_counts'), dt) Chris@87: v, j1, j2, j3 = unique(a, 1, 1, 1) Chris@87: assert_array_equal(v, b, msg) Chris@87: assert_array_equal(j1, i1, msg) Chris@87: assert_array_equal(j2, i2, msg) Chris@87: assert_array_equal(j3, c, msg) Chris@87: Chris@87: a = [5, 7, 1, 2, 1, 5, 7]*10 Chris@87: b = [1, 2, 5, 7] Chris@87: i1 = [2, 3, 0, 1] Chris@87: i2 = [2, 3, 0, 1, 0, 2, 3]*10 Chris@87: c = np.multiply([2, 1, 2, 2], 10) Chris@87: Chris@87: # test for numeric arrays Chris@87: types = [] Chris@87: types.extend(np.typecodes['AllInteger']) Chris@87: types.extend(np.typecodes['AllFloat']) Chris@87: types.append('datetime64[D]') Chris@87: types.append('timedelta64[D]') Chris@87: for dt in types: Chris@87: aa = np.array(a, dt) Chris@87: bb = np.array(b, dt) Chris@87: check_all(aa, bb, i1, i2, c, dt) Chris@87: Chris@87: # test for object arrays Chris@87: dt = 'O' Chris@87: aa = np.empty(len(a), dt) Chris@87: aa[:] = a Chris@87: bb = np.empty(len(b), dt) Chris@87: bb[:] = b Chris@87: check_all(aa, bb, i1, i2, c, dt) Chris@87: Chris@87: # test for structured arrays Chris@87: dt = [('', 'i'), ('', 'i')] Chris@87: aa = np.array(list(zip(a, a)), dt) Chris@87: bb = np.array(list(zip(b, b)), dt) Chris@87: check_all(aa, bb, i1, i2, c, dt) Chris@87: Chris@87: # test for ticket #2799 Chris@87: aa = [1. + 0.j, 1 - 1.j, 1] Chris@87: assert_array_equal(np.unique(aa), [1. - 1.j, 1. + 0.j]) Chris@87: Chris@87: # test for ticket #4785 Chris@87: a = [(1, 2), (1, 2), (2, 3)] Chris@87: unq = [1, 2, 3] Chris@87: inv = [0, 1, 0, 1, 1, 2] Chris@87: a1 = unique(a) Chris@87: assert_array_equal(a1, unq) Chris@87: a2, a2_inv = unique(a, return_inverse=True) Chris@87: assert_array_equal(a2, unq) Chris@87: assert_array_equal(a2_inv, inv) Chris@87: Chris@87: def test_intersect1d(self): Chris@87: # unique inputs Chris@87: a = np.array([5, 7, 1, 2]) Chris@87: b = np.array([2, 4, 3, 1, 5]) Chris@87: Chris@87: ec = np.array([1, 2, 5]) Chris@87: c = intersect1d(a, b, assume_unique=True) Chris@87: assert_array_equal(c, ec) Chris@87: Chris@87: # non-unique inputs Chris@87: a = np.array([5, 5, 7, 1, 2]) Chris@87: b = np.array([2, 1, 4, 3, 3, 1, 5]) Chris@87: Chris@87: ed = np.array([1, 2, 5]) Chris@87: c = intersect1d(a, b) Chris@87: assert_array_equal(c, ed) Chris@87: Chris@87: assert_array_equal([], intersect1d([], [])) Chris@87: Chris@87: def test_setxor1d(self): Chris@87: a = np.array([5, 7, 1, 2]) Chris@87: b = np.array([2, 4, 3, 1, 5]) Chris@87: Chris@87: ec = np.array([3, 4, 7]) Chris@87: c = setxor1d(a, b) Chris@87: assert_array_equal(c, ec) Chris@87: Chris@87: a = np.array([1, 2, 3]) Chris@87: b = np.array([6, 5, 4]) Chris@87: Chris@87: ec = np.array([1, 2, 3, 4, 5, 6]) Chris@87: c = setxor1d(a, b) Chris@87: assert_array_equal(c, ec) Chris@87: Chris@87: a = np.array([1, 8, 2, 3]) Chris@87: b = np.array([6, 5, 4, 8]) Chris@87: Chris@87: ec = np.array([1, 2, 3, 4, 5, 6]) Chris@87: c = setxor1d(a, b) Chris@87: assert_array_equal(c, ec) Chris@87: Chris@87: assert_array_equal([], setxor1d([], [])) Chris@87: Chris@87: def test_ediff1d(self): Chris@87: zero_elem = np.array([]) Chris@87: one_elem = np.array([1]) Chris@87: two_elem = np.array([1, 2]) Chris@87: Chris@87: assert_array_equal([], ediff1d(zero_elem)) Chris@87: assert_array_equal([0], ediff1d(zero_elem, to_begin=0)) Chris@87: assert_array_equal([0], ediff1d(zero_elem, to_end=0)) Chris@87: assert_array_equal([-1, 0], ediff1d(zero_elem, to_begin=-1, to_end=0)) Chris@87: assert_array_equal([], ediff1d(one_elem)) Chris@87: assert_array_equal([1], ediff1d(two_elem)) Chris@87: Chris@87: def test_in1d(self): Chris@87: # we use two different sizes for the b array here to test the Chris@87: # two different paths in in1d(). Chris@87: for mult in (1, 10): Chris@87: # One check without np.array, to make sure lists are handled correct Chris@87: a = [5, 7, 1, 2] Chris@87: b = [2, 4, 3, 1, 5] * mult Chris@87: ec = np.array([True, False, True, True]) Chris@87: c = in1d(a, b, assume_unique=True) Chris@87: assert_array_equal(c, ec) Chris@87: Chris@87: a[0] = 8 Chris@87: ec = np.array([False, False, True, True]) Chris@87: c = in1d(a, b, assume_unique=True) Chris@87: assert_array_equal(c, ec) Chris@87: Chris@87: a[0], a[3] = 4, 8 Chris@87: ec = np.array([True, False, True, False]) Chris@87: c = in1d(a, b, assume_unique=True) Chris@87: assert_array_equal(c, ec) Chris@87: Chris@87: a = np.array([5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5]) Chris@87: b = [2, 3, 4] * mult Chris@87: ec = [False, True, False, True, True, True, True, True, True, False, Chris@87: True, False, False, False] Chris@87: c = in1d(a, b) Chris@87: assert_array_equal(c, ec) Chris@87: Chris@87: b = b + [5, 5, 4] * mult Chris@87: ec = [True, True, True, True, True, True, True, True, True, True, Chris@87: True, False, True, True] Chris@87: c = in1d(a, b) Chris@87: assert_array_equal(c, ec) Chris@87: Chris@87: a = np.array([5, 7, 1, 2]) Chris@87: b = np.array([2, 4, 3, 1, 5] * mult) Chris@87: ec = np.array([True, False, True, True]) Chris@87: c = in1d(a, b) Chris@87: assert_array_equal(c, ec) Chris@87: Chris@87: a = np.array([5, 7, 1, 1, 2]) Chris@87: b = np.array([2, 4, 3, 3, 1, 5] * mult) Chris@87: ec = np.array([True, False, True, True, True]) Chris@87: c = in1d(a, b) Chris@87: assert_array_equal(c, ec) Chris@87: Chris@87: a = np.array([5, 5]) Chris@87: b = np.array([2, 2] * mult) Chris@87: ec = np.array([False, False]) Chris@87: c = in1d(a, b) Chris@87: assert_array_equal(c, ec) Chris@87: Chris@87: a = np.array([5]) Chris@87: b = np.array([2]) Chris@87: ec = np.array([False]) Chris@87: c = in1d(a, b) Chris@87: assert_array_equal(c, ec) Chris@87: Chris@87: assert_array_equal(in1d([], []), []) Chris@87: Chris@87: def test_in1d_char_array(self): Chris@87: a = np.array(['a', 'b', 'c', 'd', 'e', 'c', 'e', 'b']) Chris@87: b = np.array(['a', 'c']) Chris@87: Chris@87: ec = np.array([True, False, True, False, False, True, False, False]) Chris@87: c = in1d(a, b) Chris@87: Chris@87: assert_array_equal(c, ec) Chris@87: Chris@87: def test_in1d_invert(self): Chris@87: "Test in1d's invert parameter" Chris@87: # We use two different sizes for the b array here to test the Chris@87: # two different paths in in1d(). Chris@87: for mult in (1, 10): Chris@87: a = np.array([5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5]) Chris@87: b = [2, 3, 4] * mult Chris@87: assert_array_equal(np.invert(in1d(a, b)), in1d(a, b, invert=True)) Chris@87: Chris@87: def test_in1d_ravel(self): Chris@87: # Test that in1d ravels its input arrays. This is not documented Chris@87: # behavior however. The test is to ensure consistentency. Chris@87: a = np.arange(6).reshape(2, 3) Chris@87: b = np.arange(3, 9).reshape(3, 2) Chris@87: long_b = np.arange(3, 63).reshape(30, 2) Chris@87: ec = np.array([False, False, False, True, True, True]) Chris@87: Chris@87: assert_array_equal(in1d(a, b, assume_unique=True), ec) Chris@87: assert_array_equal(in1d(a, b, assume_unique=False), ec) Chris@87: assert_array_equal(in1d(a, long_b, assume_unique=True), ec) Chris@87: assert_array_equal(in1d(a, long_b, assume_unique=False), ec) Chris@87: Chris@87: def test_union1d(self): Chris@87: a = np.array([5, 4, 7, 1, 2]) Chris@87: b = np.array([2, 4, 3, 3, 2, 1, 5]) Chris@87: Chris@87: ec = np.array([1, 2, 3, 4, 5, 7]) Chris@87: c = union1d(a, b) Chris@87: assert_array_equal(c, ec) Chris@87: Chris@87: assert_array_equal([], union1d([], [])) Chris@87: Chris@87: def test_setdiff1d(self): Chris@87: a = np.array([6, 5, 4, 7, 1, 2, 7, 4]) Chris@87: b = np.array([2, 4, 3, 3, 2, 1, 5]) Chris@87: Chris@87: ec = np.array([6, 7]) Chris@87: c = setdiff1d(a, b) Chris@87: assert_array_equal(c, ec) Chris@87: Chris@87: a = np.arange(21) Chris@87: b = np.arange(19) Chris@87: ec = np.array([19, 20]) Chris@87: c = setdiff1d(a, b) Chris@87: assert_array_equal(c, ec) Chris@87: Chris@87: assert_array_equal([], setdiff1d([], [])) Chris@87: Chris@87: def test_setdiff1d_char_array(self): Chris@87: a = np.array(['a', 'b', 'c']) Chris@87: b = np.array(['a', 'b', 's']) Chris@87: assert_array_equal(setdiff1d(a, b), np.array(['c'])) Chris@87: Chris@87: def test_manyways(self): Chris@87: a = np.array([5, 7, 1, 2, 8]) Chris@87: b = np.array([9, 8, 2, 4, 3, 1, 5]) Chris@87: Chris@87: c1 = setxor1d(a, b) Chris@87: aux1 = intersect1d(a, b) Chris@87: aux2 = union1d(a, b) Chris@87: c2 = setdiff1d(aux2, aux1) Chris@87: assert_array_equal(c1, c2) Chris@87: Chris@87: Chris@87: if __name__ == "__main__": Chris@87: run_module_suite()