[pypy-commit] pypy default: add argwhere, flatnonzero

[bdkearns]

Author: Brian Kearns <bdkearns at gmail.com>
Branch: 
Changeset: r67394:2add12c32263
Date: 2013-10-15 06:24 -0400
http://bitbucket.org/pypy/pypy/changeset/2add12c32263/

Log:	add argwhere, flatnonzero

diff --git a/lib_pypy/numpypy/core/numeric.py b/lib_pypy/numpypy/core/numeric.py
--- a/lib_pypy/numpypy/core/numeric.py
+++ b/lib_pypy/numpypy/core/numeric.py
@@ -1,12 +1,13 @@
 from __future__ import division, absolute_import, print_function
 
 __all__ = [
-           'newaxis', 'ufunc',
+           'newaxis', 'ufunc', 'argwhere',
            'asarray', 'asanyarray', 'base_repr',
            'array_repr', 'array_str', 'set_string_function',
            'array_equal', 'array_equiv', 'outer', 'vdot', 'identity', 'little_endian',
-           'Inf', 'inf', 'infty', 'Infinity', 'nan', 'NaN', 'False_', 'True_',
-           'seterr',
+           'seterr', 'flatnonzero',
+           'Inf', 'inf', 'infty', 'Infinity',
+           'nan', 'NaN', 'False_', 'True_',
           ]
 
 import sys
@@ -165,6 +166,85 @@
     """
     return array(a, dtype, copy=False, order=order, subok=True)
 
+def argwhere(a):
+    """
+    Find the indices of array elements that are non-zero, grouped by element.
+
+    Parameters
+    ----------
+    a : array_like
+        Input data.
+
+    Returns
+    -------
+    index_array : ndarray
+        Indices of elements that are non-zero. Indices are grouped by element.
+
+    See Also
+    --------
+    where, nonzero
+
+    Notes
+    -----
+    ``np.argwhere(a)`` is the same as ``np.transpose(np.nonzero(a))``.
+
+    The output of ``argwhere`` is not suitable for indexing arrays.
+    For this purpose use ``where(a)`` instead.
+
+    Examples
+    --------
+    >>> x = np.arange(6).reshape(2,3)
+    >>> x
+    array([[0, 1, 2],
+           [3, 4, 5]])
+    >>> np.argwhere(x>1)
+    array([[0, 2],
+           [1, 0],
+           [1, 1],
+           [1, 2]])
+
+    """
+    return transpose(asanyarray(a).nonzero())
+
+def flatnonzero(a):
+    """
+    Return indices that are non-zero in the flattened version of a.
+
+    This is equivalent to a.ravel().nonzero()[0].
+
+    Parameters
+    ----------
+    a : ndarray
+        Input array.
+
+    Returns
+    -------
+    res : ndarray
+        Output array, containing the indices of the elements of `a.ravel()`
+        that are non-zero.
+
+    See Also
+    --------
+    nonzero : Return the indices of the non-zero elements of the input array.
+    ravel : Return a 1-D array containing the elements of the input array.
+
+    Examples
+    --------
+    >>> x = np.arange(-2, 3)
+    >>> x
+    array([-2, -1,  0,  1,  2])
+    >>> np.flatnonzero(x)
+    array([0, 1, 3, 4])
+
+    Use the indices of the non-zero elements as an index array to extract
+    these elements:
+
+    >>> x.ravel()[np.flatnonzero(x)]
+    array([-2, -1,  1,  2])
+
+    """
+    return a.ravel().nonzero()[0]
+
 def base_repr(number, base=2, padding=0):
     """
     Return a string representation of a number in the given base system.
diff --git a/pypy/module/test_lib_pypy/numpypy/core/test_numeric.py b/pypy/module/test_lib_pypy/numpypy/core/test_numeric.py
--- a/pypy/module/test_lib_pypy/numpypy/core/test_numeric.py
+++ b/pypy/module/test_lib_pypy/numpypy/core/test_numeric.py
@@ -188,6 +188,23 @@
 
 
 class AppTestNumeric(BaseNumpyAppTest):
+    def test_argwhere(self):
+        import numpypy as np
+        x = np.arange(6).reshape(2,3)
+        a = np.argwhere(x>1)
+        assert np.array_equal(a,
+            [[0, 2],
+            [1, 0],
+            [1, 1],
+            [1, 2]]
+        )
+
+    def test_flatnonzero(self):
+        import numpypy as np
+        x = np.arange(-2, 3)
+        a = np.flatnonzero(x)
+        assert np.array_equal(a, [0, 1, 3, 4])
+
     def test_outer(self):
         from numpypy import array, outer
         a = [1, 2, 3]