Added axis argument to numpy.unique

1 files changed, 186 insertions, 105 deletions

diff --git a/numpy/lib/tests/test_arraysetops.py b/numpy/lib/tests/test_arraysetops.py
index 75918fbee..8b142c264 100644
--- a/numpy/lib/tests/test_arraysetops.py
+++ b/numpy/lib/tests/test_arraysetops.py
@@ -5,7 +5,7 @@ from __future__ import division, absolute_import, print_function
 
 import numpy as np
 from numpy.testing import (
-    run_module_suite, TestCase, assert_array_equal, assert_equal
+    run_module_suite, TestCase, assert_array_equal, assert_equal, assert_raises
     )
 from numpy.lib.arraysetops import (
     ediff1d, intersect1d, setxor1d, union1d, setdiff1d, unique, in1d
@@ -14,107 +14,6 @@ from numpy.lib.arraysetops import (
 
 class TestSetOps(TestCase):
 
-    def test_unique(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, 1, 0, 0)
-            assert_array_equal(v, b, msg)
-            assert_array_equal(j, i1, msg)
-
-            msg = base_msg.format('return_inverse', dt)
-            v, j = unique(a, 0, 1, 0)
-            assert_array_equal(v, b, msg)
-            assert_array_equal(j, i2, msg)
-
-            msg = base_msg.format('return_counts', dt)
-            v, j = unique(a, 0, 0, 1)
-            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, 1, 1, 0)
-            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, 1, 0, 1)
-            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, 0, 1, 1)
-            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, 1, 1, 1)
-            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))
-
     def test_intersect1d(self):
         # unique inputs
         a = np.array([5, 7, 1, 2])
@@ -182,7 +81,7 @@ class TestSetOps(TestCase):
         # 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
+            # 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])
@@ -201,8 +100,8 @@ class TestSetOps(TestCase):
 
             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]
+            ec = [False, True, False, True, True, True, True, True, True,
+                  False, True, False, False, False]
             c = in1d(a, b)
             assert_array_equal(c, ec)
 
@@ -313,5 +212,187 @@ class TestSetOps(TestCase):
         assert_array_equal(c1, c2)
 
 
+class TestUnique(TestCase):
+
+    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, 1, 0, 0)
+            assert_array_equal(v, b, msg)
+            assert_array_equal(j, i1, msg)
+
+            msg = base_msg.format('return_inverse', dt)
+            v, j = unique(a, 0, 1, 0)
+            assert_array_equal(v, b, msg)
+            assert_array_equal(j, i2, msg)
+
+            msg = base_msg.format('return_counts', dt)
+            v, j = unique(a, 0, 0, 1)
+            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, 1, 1, 0)
+            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, 1, 0, 1)
+            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, 0, 1, 1)
+            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, 1, 1, 1)
+            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))
+
+    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(ValueError, unique, np.arange(10), axis=2)
+        assert_raises(ValueError, 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)
+
+    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.dstack([data] * 3)
+        result = data3d[..., :1]
+        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)
+
+
 if __name__ == "__main__":
     run_module_suite()