[pypy-commit] pypy default: add array_equiv

[bdkearns]

Author: Brian Kearns <bdkearns at gmail.com>
Branch: 
Changeset: r67383:075052ce29b7
Date: 2013-10-15 03:11 -0400
http://bitbucket.org/pypy/pypy/changeset/075052ce29b7/

Log:	add array_equiv

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
@@ -4,7 +4,7 @@
            'newaxis', 'ufunc',
            'asarray', 'asanyarray', 'base_repr',
            'array_repr', 'array_str', 'set_string_function',
-           'array_equal', 'outer', 'vdot', 'identity', 'little_endian',
+           'array_equal', 'array_equiv', 'outer', 'vdot', 'identity', 'little_endian',
            'Inf', 'inf', 'infty', 'Infinity', 'nan', 'NaN', 'False_', 'True_',
            'seterr',
           ]
@@ -453,6 +453,50 @@
         return False
     return bool((a1 == a2).all())
 
+def array_equiv(a1, a2):
+    """
+    Returns True if input arrays are shape consistent and all elements equal.
+
+    Shape consistent means they are either the same shape, or one input array
+    can be broadcasted to create the same shape as the other one.
+
+    Parameters
+    ----------
+    a1, a2 : array_like
+        Input arrays.
+
+    Returns
+    -------
+    out : bool
+        True if equivalent, False otherwise.
+
+    Examples
+    --------
+    >>> np.array_equiv([1, 2], [1, 2])
+    True
+    >>> np.array_equiv([1, 2], [1, 3])
+    False
+
+    Showing the shape equivalence:
+
+    >>> np.array_equiv([1, 2], [[1, 2], [1, 2]])
+    True
+    >>> np.array_equiv([1, 2], [[1, 2, 1, 2], [1, 2, 1, 2]])
+    False
+
+    >>> np.array_equiv([1, 2], [[1, 2], [1, 3]])
+    False
+
+    """
+    try:
+        a1, a2 = asarray(a1), asarray(a2)
+    except:
+        return False
+    try:
+        return bool(asarray(a1 == a2).all())
+    except ValueError:
+        return False
+
 def outer(a,b):
     """
     Compute the outer product of two vectors.
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
@@ -177,6 +177,15 @@
         assert not array_equal(array(a), b)
         assert not array_equal(array(a), array(b))
 
+    def test_equiv(self):
+        import numpypy as np
+
+        assert np.array_equiv([1, 2], [1, 2])
+        assert not np.array_equiv([1, 2], [1, 3])
+        assert np.array_equiv([1, 2], [[1, 2], [1, 2]])
+        assert not np.array_equiv([1, 2], [[1, 2, 1, 2], [1, 2, 1, 2]])
+        assert not np.array_equiv([1, 2], [[1, 2], [1, 3]])
+
 
 class AppTestNumeric(BaseNumpyAppTest):
     def test_outer(self):