diff options
Diffstat (limited to 'numpy/lib/arraysetops.py')
| -rw-r--r-- | numpy/lib/arraysetops.py | 82 |
1 files changed, 70 insertions, 12 deletions
diff --git a/numpy/lib/arraysetops.py b/numpy/lib/arraysetops.py index 5880ea154..2f8c07114 100644 --- a/numpy/lib/arraysetops.py +++ b/numpy/lib/arraysetops.py @@ -28,6 +28,7 @@ To do: Optionally return indices analogously to unique for all functions. from __future__ import division, absolute_import, print_function import numpy as np +from numpy.core.overrides import array_function_dispatch __all__ = [ @@ -36,6 +37,11 @@ __all__ = [ ] +def _ediff1d_dispatcher(ary, to_end=None, to_begin=None): + return (ary, to_end, to_begin) + + +@array_function_dispatch(_ediff1d_dispatcher) def ediff1d(ary, to_end=None, to_begin=None): """ The differences between consecutive elements of an array. @@ -82,6 +88,11 @@ def ediff1d(ary, to_end=None, to_begin=None): # force a 1d array ary = np.asanyarray(ary).ravel() + # we have unit tests enforcing + # propagation of the dtype of input + # ary to returned result + dtype_req = ary.dtype + # fast track default case if to_begin is None and to_end is None: return ary[1:] - ary[:-1] @@ -89,13 +100,23 @@ def ediff1d(ary, to_end=None, to_begin=None): if to_begin is None: l_begin = 0 else: - to_begin = np.asanyarray(to_begin).ravel() + to_begin = np.asanyarray(to_begin) + if not np.can_cast(to_begin, dtype_req): + raise TypeError("dtype of to_begin must be compatible " + "with input ary") + + to_begin = to_begin.ravel() l_begin = len(to_begin) if to_end is None: l_end = 0 else: - to_end = np.asanyarray(to_end).ravel() + to_end = np.asanyarray(to_end) + if not np.can_cast(to_end, dtype_req): + raise TypeError("dtype of to_end must be compatible " + "with input ary") + + to_end = to_end.ravel() l_end = len(to_end) # do the calculation in place and copy to_begin and to_end @@ -118,6 +139,12 @@ def _unpack_tuple(x): return x +def _unique_dispatcher(ar, return_index=None, return_inverse=None, + return_counts=None, axis=None): + return (ar,) + + +@array_function_dispatch(_unique_dispatcher) def unique(ar, return_index=False, return_inverse=False, return_counts=False, axis=None): """ @@ -298,6 +325,12 @@ def _unique1d(ar, return_index=False, return_inverse=False, return ret +def _intersect1d_dispatcher( + ar1, ar2, assume_unique=None, return_indices=None): + return (ar1, ar2) + + +@array_function_dispatch(_intersect1d_dispatcher) def intersect1d(ar1, ar2, assume_unique=False, return_indices=False): """ Find the intersection of two arrays. @@ -312,12 +345,12 @@ def intersect1d(ar1, ar2, assume_unique=False, return_indices=False): If True, the input arrays are both assumed to be unique, which can speed up the calculation. Default is False. return_indices : bool - If True, the indices which correspond to the intersection of the - two arrays are returned. The first instance of a value is used - if there are multiple. Default is False. - - .. versionadded:: 1.15.0 - + If True, the indices which correspond to the intersection of the two + arrays are returned. The first instance of a value is used if there are + multiple. Default is False. + + .. versionadded:: 1.15.0 + Returns ------- intersect1d : ndarray @@ -326,7 +359,7 @@ def intersect1d(ar1, ar2, assume_unique=False, return_indices=False): The indices of the first occurrences of the common values in `ar1`. Only provided if `return_indices` is True. comm2 : ndarray - The indices of the first occurrences of the common values in `ar2`. + The indices of the first occurrences of the common values in `ar2`. Only provided if `return_indices` is True. @@ -345,7 +378,7 @@ def intersect1d(ar1, ar2, assume_unique=False, return_indices=False): >>> from functools import reduce >>> reduce(np.intersect1d, ([1, 3, 4, 3], [3, 1, 2, 1], [6, 3, 4, 2])) array([3]) - + To return the indices of the values common to the input arrays along with the intersected values: >>> x = np.array([1, 1, 2, 3, 4]) @@ -355,8 +388,11 @@ def intersect1d(ar1, ar2, assume_unique=False, return_indices=False): (array([0, 2, 4]), array([1, 0, 2])) >>> xy, x[x_ind], y[y_ind] (array([1, 2, 4]), array([1, 2, 4]), array([1, 2, 4])) - + """ + ar1 = np.asanyarray(ar1) + ar2 = np.asanyarray(ar2) + if not assume_unique: if return_indices: ar1, ind1 = unique(ar1, return_index=True) @@ -367,7 +403,7 @@ def intersect1d(ar1, ar2, assume_unique=False, return_indices=False): else: ar1 = ar1.ravel() ar2 = ar2.ravel() - + aux = np.concatenate((ar1, ar2)) if return_indices: aux_sort_indices = np.argsort(aux, kind='mergesort') @@ -389,6 +425,12 @@ def intersect1d(ar1, ar2, assume_unique=False, return_indices=False): else: return int1d + +def _setxor1d_dispatcher(ar1, ar2, assume_unique=None): + return (ar1, ar2) + + +@array_function_dispatch(_setxor1d_dispatcher) def setxor1d(ar1, ar2, assume_unique=False): """ Find the set exclusive-or of two arrays. @@ -543,6 +585,11 @@ def in1d(ar1, ar2, assume_unique=False, invert=False): return ret[rev_idx] +def _isin_dispatcher(element, test_elements, assume_unique=None, invert=None): + return (element, test_elements) + + +@array_function_dispatch(_isin_dispatcher) def isin(element, test_elements, assume_unique=False, invert=False): """ Calculates `element in test_elements`, broadcasting over `element` only. @@ -641,6 +688,11 @@ def isin(element, test_elements, assume_unique=False, invert=False): invert=invert).reshape(element.shape) +def _union1d_dispatcher(ar1, ar2): + return (ar1, ar2) + + +@array_function_dispatch(_union1d_dispatcher) def union1d(ar1, ar2): """ Find the union of two arrays. @@ -676,6 +728,12 @@ def union1d(ar1, ar2): """ return unique(np.concatenate((ar1, ar2), axis=None)) + +def _setdiff1d_dispatcher(ar1, ar2, assume_unique=None): + return (ar1, ar2) + + +@array_function_dispatch(_setdiff1d_dispatcher) def setdiff1d(ar1, ar2, assume_unique=False): """ Find the set difference of two arrays. |