Merge pull request #15747 from WarrenWeckesser/unique-with-zero-len

BUG: lib: Handle axes with length 0 in np.unique.

1 files changed, 81 insertions, 37 deletions

diff --git a/numpy/lib/tests/test_arraysetops.py b/numpy/lib/tests/test_arraysetops.py
index 851fd31ea..81ba789e3 100644
--- a/numpy/lib/tests/test_arraysetops.py
+++ b/numpy/lib/tests/test_arraysetops.py
@@ -11,7 +11,6 @@ from numpy.lib.arraysetops import (
 import pytest
 
 
-
 class TestSetOps:
 
     def test_intersect1d(self):
@@ -118,12 +117,13 @@ class TestSetOps:
         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]))
+        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", [
         # should fail because trying to cast
@@ -156,27 +156,27 @@ class TestSetOps:
                     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)),
-         ])
+    @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,
@@ -191,7 +191,6 @@ class TestSetOps:
         assert_equal(actual, expected)
         assert actual.dtype == expected.dtype
 
-
     def test_isin(self):
         # the tests for in1d cover most of isin's behavior
         # if in1d is removed, would need to change those tests to test
@@ -200,33 +199,34 @@ class TestSetOps:
             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)
             y = isin_slow(a, b)
             assert_array_equal(x, y)
 
-        #multidimensional arrays in both arguments
+        # 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
+        # array-likes as both arguments
         c = [(9, 8), (7, 6)]
         d = (9, 7)
         assert_isin_equal(c, d)
 
-        #zero-d array:
+        # zero-d array:
         f = np.array(3)
         assert_isin_equal(f, b)
         assert_isin_equal(a, f)
         assert_isin_equal(f, f)
 
-        #scalar:
+        # scalar:
         assert_isin_equal(5, b)
         assert_isin_equal(a, 6)
         assert_isin_equal(5, 6)
 
-        #empty array-like:
+        # empty array-like:
         x = []
         assert_isin_equal(x, b)
         assert_isin_equal(a, x)
@@ -520,7 +520,8 @@ class TestUnique:
         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:]
+        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)
@@ -563,9 +564,52 @@ class TestUnique:
         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')
+        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)
@@ -580,7 +624,7 @@ class TestUnique:
         # 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)
+            a = np.array([[-1], [0]], dt)
             b = np.unique(a, axis=0)
             assert_array_equal(a, b, fmt % dt)