Enhancements to arraysetops (ticket #1133, by Neil Crighton)

1 files changed, 100 insertions, 73 deletions

diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py
index 4454781c3..6041b7e2c 100644
--- a/numpy/ma/extras.py
+++ b/numpy/ma/extras.py
@@ -19,14 +19,14 @@ __all__ = ['apply_along_axis', 'atleast_1d', 'atleast_2d', 'atleast_3d',
            'ediff1d',
            'flatnotmasked_contiguous', 'flatnotmasked_edges',
            'hsplit', 'hstack',
-           'intersect1d', 'intersect1d_nu',
+           'in1d', 'intersect1d', 'intersect1d_nu',
            'mask_cols', 'mask_rowcols', 'mask_rows', 'masked_all',
            'masked_all_like', 'median', 'mr_',
            'notmasked_contiguous', 'notmasked_edges',
            'polyfit',
            'row_stack',
            'setdiff1d', 'setmember1d', 'setxor1d',
-           'unique1d', 'union1d',
+           'unique', 'unique1d', 'union1d',
            'vander', 'vstack',
            ]
 
@@ -45,6 +45,8 @@ import numpy.core.umath as umath
 from numpy.lib.index_tricks import AxisConcatenator
 from numpy.linalg import lstsq
 
+from numpy.lib.utils import deprecate_with_doc
+
 #...............................................................................
 def issequence(seq):
     """Is seq a sequence (ndarray, list or tuple)?"""
@@ -867,7 +869,7 @@ def ediff1d(arr, to_end=None, to_begin=None):
     return ed
 
 
-def unique1d(ar1, return_index=False, return_inverse=False):
+def unique(ar1, return_index=False, return_inverse=False):
     """
     Finds the unique elements of an array.
 
@@ -877,11 +879,11 @@ def unique1d(ar1, return_index=False, return_inverse=False):
 
     See Also
     --------
-    np.unique1d : equivalent function for ndarrays.
+    np.unique : equivalent function for ndarrays.
     """
-    output = np.unique1d(ar1,
-                         return_index=return_index,
-                         return_inverse=return_inverse)
+    output = np.unique(ar1,
+                       return_index=return_index,
+                       return_inverse=return_inverse)
     if isinstance(output, tuple):
         output = list(output)
         output[0] = output[0].view(MaskedArray)
@@ -891,33 +893,7 @@ def unique1d(ar1, return_index=False, return_inverse=False):
     return output
 
 
-def intersect1d(ar1, ar2):
-    """
-    Returns the repeated or unique elements belonging to the two arrays.
-
-    Masked values are assumed equals one to the other.
-    The output is always a masked array
-
-    See Also
-    --------
-    numpy.intersect1d : equivalent function for ndarrays.
-
-    Examples
-    --------
-    >>> x = array([1, 3, 3, 3], mask=[0, 0, 0, 1])
-    >>> y = array([3, 1, 1, 1], mask=[0, 0, 0, 1])
-    >>> intersect1d(x, y)
-    masked_array(data = [1 1 3 3 --],
-                 mask = [False False False False  True],
-           fill_value = 999999)
-    """
-    aux = ma.concatenate((ar1,ar2))
-    aux.sort()
-    return aux[aux[1:] == aux[:-1]]
-
-
-
-def intersect1d_nu(ar1, ar2):
+def intersect1d(ar1, ar2, assume_unique=False):
     """
     Returns the unique elements common to both arrays.
 
@@ -926,27 +902,28 @@ def intersect1d_nu(ar1, ar2):
 
     See Also
     --------
-    intersect1d : Returns repeated or unique common elements.
-    numpy.intersect1d_nu : equivalent function for ndarrays.
+    numpy.intersect1d : equivalent function for ndarrays.
 
     Examples
     --------
     >>> x = array([1, 3, 3, 3], mask=[0, 0, 0, 1])
     >>> y = array([3, 1, 1, 1], mask=[0, 0, 0, 1])
-    >>> intersect1d_nu(x, y)
+    >>> intersect1d(x, y)
     masked_array(data = [1 3 --],
                  mask = [False False  True],
            fill_value = 999999)
 
     """
-    # Might be faster than unique1d( intersect1d( ar1, ar2 ) )?
-    aux = ma.concatenate((unique1d(ar1), unique1d(ar2)))
+    if assume_unique:
+        aux = ma.concatenate((ar1, ar2))
+    else:
+        # Might be faster than unique1d( intersect1d( ar1, ar2 ) )?
+        aux = ma.concatenate((unique(ar1), unique(ar2)))
     aux.sort()
     return aux[aux[1:] == aux[:-1]]
 
 
-
-def setxor1d(ar1, ar2):
+def setxor1d(ar1, ar2, assume_unique=False):
     """
     Set exclusive-or of 1D arrays with unique elements.
 
@@ -955,6 +932,10 @@ def setxor1d(ar1, ar2):
     numpy.setxor1d : equivalent function for ndarrays
 
     """
+    if not assume_unique:
+        ar1 = unique(ar1)
+        ar2 = unique(ar2)
+
     aux = ma.concatenate((ar1, ar2))
     if aux.size == 0:
         return aux
@@ -966,54 +947,52 @@ def setxor1d(ar1, ar2):
     flag2 = (flag[1:] == flag[:-1])
     return aux[flag2]
 
-
-def setmember1d(ar1, ar2):
+def in1d(ar1, ar2, assume_unique=False):
     """
-    Return a boolean array set True where first element is in second array.
+    Test whether each element of an array is also present in a second
+    array.
 
     See Also
     --------
-    numpy.setmember1d : equivalent function for ndarrays.
-
-    """
-    ar1 = ma.asanyarray(ar1)
-    ar2 = ma.asanyarray( ar2 )
-    ar = ma.concatenate((ar1, ar2 ))
-    b1 = ma.zeros(ar1.shape, dtype = np.int8)
-    b2 = ma.ones(ar2.shape, dtype = np.int8)
-    tt = ma.concatenate((b1, b2))
+    numpy.in1d : equivalent function for ndarrays
 
-    # We need this to be a stable sort, so always use 'mergesort' here. The
-    # values from the first array should always come before the values from the
-    # second array.
-    perm = ar.argsort(kind='mergesort')
-    aux = ar[perm]
-    aux2 = tt[perm]
-#    flag = ediff1d( aux, 1 ) == 0
-    flag = ma.concatenate((aux[1:] == aux[:-1], [False]))
-    ii = ma.where( flag * aux2 )[0]
-    aux = perm[ii+1]
-    perm[ii+1] = perm[ii]
-    perm[ii] = aux
-    #
-    indx = perm.argsort(kind='mergesort')[:len( ar1 )]
-    #
-    return flag[indx]
+    Notes
+    -----
+    .. versionadded:: 1.4.0
+    """
+    if not assume_unique:
+        ar1, rev_idx = unique(ar1, return_inverse=True)
+        ar2 = unique(ar2)
+
+    ar = ma.concatenate( (ar1, ar2) )
+    # We need this to be a stable sort, so always use 'mergesort'
+    # here. The values from the first array should always come before
+    # the values from the second array.
+    order = ar.argsort(kind='mergesort')
+    sar = ar[order]
+    equal_adj = (sar[1:] == sar[:-1])
+    flag = ma.concatenate( (equal_adj, [False] ) )
+    indx = order.argsort(kind='mergesort')[:len( ar1 )]
+
+    if assume_unique:
+        return flag[indx]
+    else:
+        return flag[indx][rev_idx]
 
 
 def union1d(ar1, ar2):
     """
-    Union of 1D arrays with unique elements.
+    Union of two arrays.
 
     See also
     --------
     numpy.union1d : equivalent function for ndarrays.
 
     """
-    return unique1d(ma.concatenate((ar1, ar2)))
+    return unique(ma.concatenate((ar1, ar2)))
 
 
-def setdiff1d(ar1, ar2):
+def setdiff1d(ar1, ar2, assume_unique=False):
     """
     Set difference of 1D arrays with unique elements.
 
@@ -1022,12 +1001,60 @@ def setdiff1d(ar1, ar2):
     numpy.setdiff1d : equivalent function for ndarrays
 
     """
-    aux = setmember1d(ar1,ar2)
+    aux = in1d(ar1, ar2, assume_unique=assume_unique)
     if aux.size == 0:
         return aux
     else:
         return ma.asarray(ar1)[aux == 0]
 
+@deprecate_with_doc('')
+def unique1d(ar1, return_index=False, return_inverse=False):
+    """ This function is deprecated. Use ma.unique() instead. """
+    output = np.unique1d(ar1,
+                         return_index=return_index,
+                         return_inverse=return_inverse)
+    if isinstance(output, tuple):
+        output = list(output)
+        output[0] = output[0].view(MaskedArray)
+        output = tuple(output)
+    else:
+        output = output.view(MaskedArray)
+    return output
+
+@deprecate_with_doc('')
+def intersect1d_nu(ar1, ar2):
+    """ This function is deprecated. Use ma.intersect1d() instead."""
+    # Might be faster than unique1d( intersect1d( ar1, ar2 ) )?
+    aux = ma.concatenate((unique1d(ar1), unique1d(ar2)))
+    aux.sort()
+    return aux[aux[1:] == aux[:-1]]
+
+@deprecate_with_doc('')
+def setmember1d(ar1, ar2):
+    """ This function is deprecated. Use ma.in1d() instead."""
+    ar1 = ma.asanyarray(ar1)
+    ar2 = ma.asanyarray( ar2 )
+    ar = ma.concatenate((ar1, ar2 ))
+    b1 = ma.zeros(ar1.shape, dtype = np.int8)
+    b2 = ma.ones(ar2.shape, dtype = np.int8)
+    tt = ma.concatenate((b1, b2))
+
+    # We need this to be a stable sort, so always use 'mergesort' here. The
+    # values from the first array should always come before the values from the
+    # second array.
+    perm = ar.argsort(kind='mergesort')
+    aux = ar[perm]
+    aux2 = tt[perm]
+#    flag = ediff1d( aux, 1 ) == 0
+    flag = ma.concatenate((aux[1:] == aux[:-1], [False]))
+    ii = ma.where( flag * aux2 )[0]
+    aux = perm[ii+1]
+    perm[ii+1] = perm[ii]
+    perm[ii] = aux
+    #
+    indx = perm.argsort(kind='mergesort')[:len( ar1 )]
+    #
+    return flag[indx]
 
 
 #####--------------------------------------------------------------------------