Merge pull request #13610 from eric-wieser/argwhere

seberg · web-flow · commit 94ae1c59d78e · 2019-09-06T09:45:21.000-05:00
ENH: Always produce a consistent shape in the result of `argwhere`
diff --git a/doc/release/upcoming_changes/13610.improvement.rst b/doc/release/upcoming_changes/13610.improvement.rst
@@ -0,0 +1,5 @@
+``argwhere`` now produces a consistent result on 0d arrays
+----------------------------------------------------------
+On N-d arrays, `numpy.argwhere` now always produces an array of shape
+``(n_non_zero, arr.ndim)``, even when ``arr.ndim == 0``. Previously, the
+last axis would have a dimension of 1 in this case.
diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py
@@ -26,6 +26,7 @@
 
 from . import overrides
 from . import umath
+from . import shape_base
 from .overrides import set_module
 from .umath import (multiply, invert, sin, PINF, NAN)
 from . import numerictypes
@@ -545,16 +546,19 @@ def argwhere(a):
 
     Returns
     -------
-    index_array : ndarray
+    index_array : (N, a.ndim) ndarray
         Indices of elements that are non-zero. Indices are grouped by element.
+        This array will have shape ``(N, a.ndim)`` where ``N`` is the number of
+        non-zero items.
 
     See Also
     --------
     where, nonzero
 
     Notes
     -----
-    ``np.argwhere(a)`` is the same as ``np.transpose(np.nonzero(a))``.
+    ``np.argwhere(a)`` is almost the same as ``np.transpose(np.nonzero(a))``,
+    but produces a result of the correct shape for a 0D array.
 
     The output of ``argwhere`` is not suitable for indexing arrays.
     For this purpose use ``nonzero(a)`` instead.
@@ -572,6 +576,11 @@ def argwhere(a):
            [1, 2]])
 
     """
+    # nonzero does not behave well on 0d, so promote to 1d
+    if np.ndim(a) == 0:
+        a = shape_base.atleast_1d(a)
+        # then remove the added dimension
+        return argwhere(a)[:,:0]
     return transpose(nonzero(a))
 
 
diff --git a/numpy/core/shape_base.py b/numpy/core/shape_base.py
@@ -9,8 +9,9 @@
 
 from . import numeric as _nx
 from . import overrides
-from .numeric import array, asanyarray, newaxis
+from ._asarray import array, asanyarray
 from .multiarray import normalize_axis_index
+from . import fromnumeric as _from_nx
 
 
 array_function_dispatch = functools.partial(
@@ -123,7 +124,7 @@ def atleast_2d(*arys):
         if ary.ndim == 0:
             result = ary.reshape(1, 1)
         elif ary.ndim == 1:
-            result = ary[newaxis, :]
+            result = ary[_nx.newaxis, :]
         else:
             result = ary
         res.append(result)
@@ -193,9 +194,9 @@ def atleast_3d(*arys):
         if ary.ndim == 0:
             result = ary.reshape(1, 1, 1)
         elif ary.ndim == 1:
-            result = ary[newaxis, :, newaxis]
+            result = ary[_nx.newaxis, :, _nx.newaxis]
         elif ary.ndim == 2:
-            result = ary[:, :, newaxis]
+            result = ary[:, :, _nx.newaxis]
         else:
             result = ary
         res.append(result)
@@ -435,9 +436,9 @@ def stack(arrays, axis=0, out=None):
 # Internal functions to eliminate the overhead of repeated dispatch in one of
 # the two possible paths inside np.block.
 # Use getattr to protect against __array_function__ being disabled.
-_size = getattr(_nx.size, '__wrapped__', _nx.size)
-_ndim = getattr(_nx.ndim, '__wrapped__', _nx.ndim)
-_concatenate = getattr(_nx.concatenate, '__wrapped__', _nx.concatenate)
+_size = getattr(_from_nx.size, '__wrapped__', _from_nx.size)
+_ndim = getattr(_from_nx.ndim, '__wrapped__', _from_nx.ndim)
+_concatenate = getattr(_from_nx.concatenate, '__wrapped__', _from_nx.concatenate)
 
 
 def _block_format_index(index):
diff --git a/numpy/core/tests/test_numeric.py b/numpy/core/tests/test_numeric.py
@@ -2583,6 +2583,30 @@ def test_no_overwrite(self):
 
 
 class TestArgwhere(object):
+
+    @pytest.mark.parametrize('nd', [0, 1, 2])
+    def test_nd(self, nd):
+        # get an nd array with multiple elements in every dimension
+        x = np.empty((2,)*nd, bool)
+
+        # none
+        x[...] = False
+        assert_equal(np.argwhere(x).shape, (0, nd))
+
+        # only one
+        x[...] = False
+        x.flat[0] = True
+        assert_equal(np.argwhere(x).shape, (1, nd))
+
+        # all but one
+        x[...] = True
+        x.flat[0] = False
+        assert_equal(np.argwhere(x).shape, (x.size - 1, nd))
+
+        # all
+        x[...] = True
+        assert_equal(np.argwhere(x).shape, (x.size, nd))
+
     def test_2D(self):
         x = np.arange(6).reshape((2, 3))
         assert_array_equal(np.argwhere(x > 1),
