Repository: pydata/sparse

Issue:
Add any and all methods.
@stsievert Willing to take a crack at this? It's mostly docs/tests/copy-paste with a one-line implementation and can be easily copied from `min` or `max`:

https://github.com/pydata/sparse/blob/6210da31b9fc63c01b083c121421f335d73a580d/sparse/coo/core.py#L692-L939

Discussion:
hameerabbasi: Oh. Apparently you're using an old NumPy version... The docstrings need to be produced with newer versions to pass the tests, unfortunately. (1.14.0 or newer I think).

Patch:
diff --git a/docs/generated/sparse.COO.all.rst b/docs/generated/sparse.COO.all.rst
new file mode 100644
index 0000000..f79e99c
--- /dev/null
+++ b/docs/generated/sparse.COO.all.rst
@@ -0,0 +1,6 @@
+COO.all
+=======
+
+.. currentmodule:: sparse
+
+.. automethod:: COO.all
\ No newline at end of file
diff --git a/docs/generated/sparse.COO.any.rst b/docs/generated/sparse.COO.any.rst
new file mode 100644
index 0000000..1bdad68
--- /dev/null
+++ b/docs/generated/sparse.COO.any.rst
@@ -0,0 +1,6 @@
+COO.any
+=======
+
+.. currentmodule:: sparse
+
+.. automethod:: COO.any
\ No newline at end of file
diff --git a/docs/generated/sparse.COO.rst b/docs/generated/sparse.COO.rst
index d0a1c9f..cdd16d4 100644
--- a/docs/generated/sparse.COO.rst
+++ b/docs/generated/sparse.COO.rst
@@ -40,10 +40,13 @@ COO
       :toctree:
 
       COO.reduce
+
       COO.sum
-      COO.max
-      COO.min
       COO.prod
+      COO.min
+      COO.max
+      COO.any
+      COO.all
 
    .. rubric:: :ref:`Converting to other formats <converting>`
    .. autosummary::
diff --git a/docs/operations.rst b/docs/operations.rst
index 9ca3183..524f2af 100644
--- a/docs/operations.rst
+++ b/docs/operations.rst
@@ -217,7 +217,21 @@ All of the following will raise an :obj:`IndexError`, like in Numpy 1.13 and lat
    z[1, 4, 8]
    z[-6]
 
-.. note:: Numpy advanced indexing is currently not supported.
+
+Advanced Indexing
+~~~~~~~~~~~~~~~~~
+
+Advanced indexing (indexing arrays with other arrays) is supported, but only for indexing
+with a *single array*. Indexing a single array with multiple arrays is not supported at
+this time. As above, if  :code:`z.shape` is :code:`(5, 6, 7)`, all of the following will
+work like NumPy:
+
+.. code-block:: python
+
+   z[[0, 1, 2]]
+   z[1, [3]]
+   z[1, 4, [3, 6]]
+   z[:3, :2, [1, 5]]
 
 .. _operations-other:
 
diff --git a/requirements.txt b/requirements.txt
index e0f833c..bc7f77c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,3 @@
 numpy >= 1.13
 scipy >= 0.19
-numba
+numba >= 0.39
diff --git a/sparse/coo/core.py b/sparse/coo/core.py
index 1ec395b..881a721 100644
--- a/sparse/coo/core.py
+++ b/sparse/coo/core.py
@@ -811,6 +811,114 @@ class COO(SparseArray, NDArrayOperatorsMixin):
         """
         return np.maximum.reduce(self, out=out, axis=axis, keepdims=keepdims)
 
+    def any(self, axis=None, keepdims=False, out=None):
+        """
+        See if any values along array are ``True``. Uses all axes by default.
+
+        Parameters
+        ----------
+        axis : Union[int, Iterable[int]], optional
+            The axes along which to minimize. Uses all axes by default.
+        keepdims : bool, optional
+            Whether or not to keep the dimensions of the original array.
+
+        Returns
+        -------
+        COO
+            The reduced output sparse array.
+
+        See Also
+        --------
+        :obj:`numpy.all` : Equivalent numpy function.
+
+        Notes
+        -----
+        * This function internally calls :obj:`COO.sum_duplicates` to bring the array into
+          canonical form.
+        * The :code:`out` parameter is provided just for compatibility with Numpy and
+          isn't actually supported.
+
+        Examples
+        --------
+        You can use :obj:`COO.min` to minimize an array across any dimension.
+
+        >>> x = np.array([[False, False],
+        ...               [False, True ],
+        ...               [True,  False],
+        ...               [True,  True ]])
+        >>> s = COO.from_numpy(x)
+        >>> s2 = s.any(axis=1)
+        >>> s2.todense()  # doctest: +SKIP
+        array([False,  True,  True,  True])
+
+        You can also use the :code:`keepdims` argument to keep the dimensions after the
+        minimization.
+
+        >>> s3 = s.any(axis=1, keepdims=True)
+        >>> s3.shape
+        (4, 1)
+
+        By default, this reduces the array down to one number, minimizing along all axes.
+
+        >>> s.any()
+        True
+        """
+        return np.logical_or.reduce(self, out=out, axis=axis, keepdims=keepdims)
+
+    def all(self, axis=None, keepdims=False, out=None):
+        """
+        See if all values in an array are ``True``. Uses all axes by default.
+
+        Parameters
+        ----------
+        axis : Union[int, Iterable[int]], optional
+            The axes along which to minimize. Uses all axes by default.
+        keepdims : bool, optional
+            Whether or not to keep the dimensions of the original array.
+
+        Returns
+        -------
+        COO
+            The reduced output sparse array.
+
+        See Also
+        --------
+        :obj:`numpy.all` : Equivalent numpy function.
+
+        Notes
+        -----
+        * This function internally calls :obj:`COO.sum_duplicates` to bring the array into
+          canonical form.
+        * The :code:`out` parameter is provided just for compatibility with Numpy and
+          isn't actually supported.
+
+        Examples
+        --------
+        You can use :obj:`COO.min` to minimize an array across any dimension.
+
+        >>> x = np.array([[False, False],
+        ...               [False, True ],
+        ...               [True,  False],
+        ...               [True,  True ]])
+        >>> s = COO.from_numpy(x)
+        >>> s2 = s.all(axis=1)
+        >>> s2.todense()  # doctest: +SKIP
+        array([False, False, False,  True])
+
+        You can also use the :code:`keepdims` argument to keep the dimensions after the
+        minimization.
+
+        >>> s3 = s.all(axis=1, keepdims=True)
+        >>> s3.shape
+        (4, 1)
+
+        By default, this reduces the array down to one boolean, minimizing along all axes.
+
+        >>> s.all()
+        False
+        """
+        return np.logical_and.reduce(self, out=out, axis=axis, keepdims=keepdims)
+
     def min(self, axis=None, keepdims=False, out=None):
         """
         Minimize along the given axes. Uses all axes by default.
@@ -865,7 +973,7 @@ class COO(SparseArray, NDArrayOperatorsMixin):
         >>> s3.shape
         (5, 1)
 
-        By default, this reduces the array down to one number, minimizing along all axes.
+        By default, this reduces the array down to one boolean, minimizing along all axes.
 
         >>> s.min()
         0
diff --git a/sparse/coo/indexing.py b/sparse/coo/indexing.py
index a3eaf0f..7ee9688 100644
--- a/sparse/coo/indexing.py
+++ b/sparse/coo/indexing.py
@@ -1,4 +1,3 @@
-from collections import Iterable
 from numbers import Integral
 
 import numba
@@ -60,11 +59,12 @@ def getitem(x, index):
     index = normalize_index(index, x.shape)
 
     # zip_longest so things like x[..., None] are picked up.
-    if len(index) != 0 and all(ind == slice(0, dim, 1) for ind, dim in zip_longest(index, x.shape)):
+    if len(index) != 0 and all(isinstance(ind, slice) and ind == slice(0, dim, 1)
+                               for ind, dim in zip_longest(index, x.shape)):
         return x
 
     # Get the mask
-    mask = _mask(x.coords, index, x.shape)
+    mask, adv_idx = _mask(x.coords, index, x.shape)
 
     # Get the length of the mask
     if isinstance(mask, slice):
@@ -76,8 +76,7 @@ def getitem(x, index):
     shape = []
     i = 0
 
-    sorted = True
-
+    sorted = adv_idx is None or adv_idx.pos == 0
     for ind in index:
         # Nothing is added to shape or coords if the index is an integer.
         if isinstance(ind, Integral):
@@ -88,11 +87,13 @@ def getitem(x, index):
             shape.append(len(range(ind.start, ind.stop, ind.step)))
             coords.append((x.coords[i, mask] - ind.start) // ind.step)
             i += 1
-
             if ind.step < 0:
                 sorted = False
-        elif isinstance(ind, Iterable):
-            raise NotImplementedError('Advanced indexing is not yet supported.')
+        # Add the index and shape for the advanced index.
+        elif isinstance(ind, np.ndarray):
+            shape.append(adv_idx.length)
+            coords.append(adv_idx.idx)
+            i += 1
         # Add a dimension for None.
         elif ind is None:
             coords.append(np.zeros(n, dtype=np.intp))
@@ -122,21 +123,61 @@ def getitem(x, index):
 
 def _mask(coords, indices, shape):
     indices = _prune_indices(indices, shape)
+    indices, adv_idx, adv_idx_pos = _separate_adv_indices(indices)
+
+    if len(adv_idx) != 0:
+        if len(adv_idx) != 1:
+            raise IndexError('Only indices with at most one iterable index are supported.')
+
+        adv_idx = adv_idx[0]
+        adv_idx_pos = adv_idx_pos[0]
+
+        if adv_idx.ndim != 1:
+            raise IndexError('Only one-dimensional iterable indices supported.')
+
+        mask, aidxs = _compute_multi_mask(coords, _ind_ar_from_indices(indices), adv_idx, adv_idx_pos)
+        return mask, _AdvIdxInfo(aidxs, adv_idx_pos, len(adv_idx))
+
+    mask, is_slice = _compute_mask(coords, _ind_ar_from_indices(indices))
 
+    if is_slice:
+        return slice(mask[0], mask[1], 1), None
+    else:
+        return mask, None
+
+
+def _ind_ar_from_indices(indices):
+    """
+    Computes an index "array" from indices, such that ``indices[i]`` is
+    transformed to ``ind_ar[i]`` and ``ind_ar[i].shape == (3,)``. It has the
+    format ``[start, stop, step]``. Integers are converted into steps as well.
+
+    Parameters
+    ----------
+    indices : Iterable
+        Input indices (slices and integers)
+
+    Returns
+    -------
+    ind_ar : np.ndarray
+        The output array.
+
+    Examples
+    --------
+    >>> _ind_ar_from_indices([1])
+    array([[1, 2, 1]])
+    >>> _ind_ar_from_indices([slice(5, 7, 2)])
+    array([[5, 7, 2]])
+    """
     ind_ar = np.empty((len(indices), 3), dtype=np.intp)
 
     for i, idx in enumerate(indices):
         if isinstance(idx, slice):
             ind_ar[i] = [idx.start, idx.stop, idx.step]
-        else:  # idx is an integer
+        elif isinstance(idx, Integral):
             ind_ar[i] = [idx, idx + 1, 1]
 
-    mask, is_slice = _compute_mask(coords, ind_ar)
-
-    if is_slice:
-        return slice(mask[0], mask[1], 1)
-    else:
-        return mask
+    return ind_ar
 
 
 def _prune_indices(indices, shape, prune_none=True):
@@ -185,6 +226,85 @@ def _prune_indices(indices, shape, prune_none=True):
     return indices
 
 
+def _separate_adv_indices(indices):
+    """
+    Separates advanced from normal indices.
+
+    Parameters
+    ----------
+    indices : list
+        The input indices
+
+    Returns
+    -------
+    new_idx : list
+        The normal indices.
+    adv_idx : list
+        The advanced indices.
+    adv_idx_pos : list
+        The positions of the advanced indices.
+    """
+    adv_idx_pos = []
+    new_idx = []
+    adv_idx = []
+
+    for i, idx in enumerate(indices):
+        if isinstance(idx, np.ndarray):
+            adv_idx.append(idx)
+            adv_idx_pos.append(i)
+        else:
+            new_idx.append(idx)
+
+    return new_idx, adv_idx, adv_idx_pos
+
+
+@numba.jit(nopython=True, nogil=True)
+def _compute_multi_mask(coords, indices, adv_idx, adv_idx_pos):  # pragma: no cover
+    """
+    Computes a mask with the advanced index, and also returns the advanced index
+    dimension.
+
+    Parameters
+    ----------
+    coords : np.ndarray
+        Coordinates of the input array.
+    indices : np.ndarray
+        The indices in slice format.
+    adv_idx : int
+        The advanced index.
+    adv_idx_pos : int
+        The position of the advanced index.
+
+    Returns
+    -------
+    mask : np.ndarray
+        The mask.
+    aidxs : np.ndarray
+        The advanced array index.
+    """
+    mask = []
+    a_indices = []
+    full_idx = np.empty((len(indices) + 1, 3), dtype=np.intp)
+
+    full_idx[:adv_idx_pos] = indices[:adv_idx_pos]
+    full_idx[adv_idx_pos + 1:] = indices[adv_idx_pos:]
+
+    for i, aidx in enumerate(adv_idx):
+        full_idx[adv_idx_pos] = [aidx, aidx + 1, 1]
+        partial_mask, is_slice = _compute_mask(coords, full_idx)
+        if is_slice:
+            slice_mask = []
+            for j in range(partial_mask[0], partial_mask[1]):
+                slice_mask.append(j)
+            partial_mask = np.array(slice_mask)
+
+        mask.extend(partial_mask)
+        for _ in range(len(partial_mask)):
+            a_indices.append(i)
+
+    return np.array(mask), np.array(a_indices)
+
+
 @numba.jit(nopython=True, nogil=True)
 def _compute_mask(coords, indices):  # pragma: no cover
     """
@@ -327,8 +447,8 @@ def _get_mask_pairs(starts_old, stops_old, c, idx):  # pragma: no cover
         # For each matching "integer" in the slice, search within the "sub-coords"
         # Using binary search.
         for p_match in range(idx[0], idx[1], idx[2]):
-            start = np.searchsorted(c[starts_old[j]:stops_old[j]], p_match) + starts_old[j]
-            stop = np.searchsorted(c[starts_old[j]:stops_old[j]], p_match + 1) + starts_old[j]
+            start = np.searchsorted(c[starts_old[j]:stops_old[j]], p_match, side='left') + starts_old[j]
+            stop = np.searchsorted(c[starts_old[j]:stops_old[j]], p_match, side='right') + starts_old[j]
 
             if start != stop:
                 starts.append(start)
@@ -413,7 +533,7 @@ def _filter_pairs(starts, stops, coords, indices):  # pragma: no cover
 
                 match &= ((elem - idx[0]) % idx[2] == 0 and
                           ((idx[2] > 0 and idx[0] <= elem < idx[1]) or
-                          (idx[2] < 0 and idx[0] >= elem > idx[1])))
+                           (idx[2] < 0 and idx[0] >= elem > idx[1])))
 
             # and append to the mask if so.
             if match:
@@ -460,3 +580,10 @@ def _join_adjacent_pairs(starts_old, stops_old):  # pragma: no cover
     stops.append(stops_old[-1])
 
     return starts, stops
+
+
+class _AdvIdxInfo(object):
+    def __init__(self, idx, pos, length):
+        self.idx = idx
+        self.pos = pos
+        self.length = length
diff --git a/sparse/slicing.py b/sparse/slicing.py
index 56e5b25..087bd52 100644
--- a/sparse/slicing.py
+++ b/sparse/slicing.py
@@ -3,6 +3,7 @@
 
 import math
 from numbers import Integral, Number
+from collections import Iterable
 
 import numpy as np
 
@@ -103,18 +104,18 @@ def check_index(ind, dimension):
     >>> check_index([6, 3], 5)
     Traceback (most recent call last):
     ...
-    IndexError: Index out of bounds 5
+    IndexError: Index out of bounds for dimension 5
     >>> check_index(slice(0, 3), 5)
     """
     # unknown dimension, assumed to be in bounds
-    if isinstance(ind, (list, np.ndarray)):
+    if isinstance(ind, Iterable):
         x = np.asanyarray(ind)
         if np.issubdtype(x.dtype, np.integer) and \
-                ((x >= dimension).any() or (x < -dimension).any()):
-            raise IndexError("Index out of bounds %s" % dimension)
+                ((x >= dimension) | (x < -dimension)).any():
+            raise IndexError("Index out of bounds for dimension {:d}".format(dimension))
         elif x.dtype == bool and len(x) != dimension:
-            raise IndexError("boolean index did not match indexed array; dimension is %s "
-                             "but corresponding boolean dimension is %s", (dimension, len(x)))
+            raise IndexError("boolean index did not match indexed array; dimension is {:d} "
+                             "but corresponding boolean dimension is {:d}".format(dimension, len(x)))
     elif isinstance(ind, slice):
         return
     elif not isinstance(ind, Integral):
@@ -122,12 +123,11 @@ def check_index(ind, dimension):
                          "integer or boolean arrays are valid indices")
 
     elif ind >= dimension:
-        raise IndexError("Index is not smaller than dimension %d >= %d" %
-                         (ind, dimension))
+        raise IndexError("Index is not smaller than dimension {:d} >= {:d}".format(ind, dimension))
 
     elif ind < -dimension:
-        msg = "Negative index is not greater than negative dimension %d <= -%d"
-        raise IndexError(msg % (ind, dimension))
+        msg = "Negative index is not greater than negative dimension {:d} <= -{:d}"
+        raise IndexError(msg.format(ind, dimension))
 
 
 def sanitize_index(ind):
@@ -157,7 +157,7 @@ def sanitize_index(ind):
     elif isinstance(ind, Number):
         return _sanitize_index_element(ind)
     index_array = np.asanyarray(ind)
-    if index_array.dtype == bool:
+    if index_array.dtype == np.bool_:
         nonzero = np.nonzero(index_array)
         if len(nonzero) == 1:
             # If a 1-element tuple, unwrap the element
diff --git a/sparse/utils.py b/sparse/utils.py
index db16256..f7602ea 100644
--- a/sparse/utils.py
+++ b/sparse/utils.py
@@ -12,7 +12,7 @@ def assert_eq(x, y, check_nnz=True, compare_dtype=True, **kwargs):
     if compare_dtype:
         assert x.dtype == y.dtype
 
-    checking_method = np.array_equal \
+    check_equal = np.array_equal \
         if np.issubdtype(x.dtype, np.integer) and np.issubdtype(y.dtype, np.integer) \
         else functools.partial(np.allclose, equal_nan=True)
 
@@ -22,7 +22,7 @@ def assert_eq(x, y, check_nnz=True, compare_dtype=True, **kwargs):
         assert is_canonical(y)
 
     if isinstance(x, COO) and isinstance(y, COO) and check_nnz:
-        assert np.array_equal(x.coords, y.coords) and checking_method(x.data, y.data, **kwargs)
+        assert np.array_equal(x.coords, y.coords) and check_equal(x.data, y.data, **kwargs)
         return
 
     if hasattr(x, 'todense'):
@@ -37,7 +37,7 @@ def assert_eq(x, y, check_nnz=True, compare_dtype=True, **kwargs):
             assert_nnz(y, yy)
     else:
         yy = y
-    assert checking_method(xx, yy, **kwargs)
+    assert check_equal(xx, yy, **kwargs)
 
 
 def assert_nnz(s, x):


Test patch:
diff --git a/sparse/tests/test_coo.py b/sparse/tests/test_coo.py
index abca280..3940dea 100644
--- a/sparse/tests/test_coo.py
+++ b/sparse/tests/test_coo.py
@@ -27,6 +27,20 @@ def test_reductions(reduction, axis, keepdims, kwargs, eqkwargs):
     assert_eq(xx, yy, **eqkwargs)
 
 
+@pytest.mark.parametrize('reduction,kwargs,eqkwargs', [
+    ('any', {}, {}),
+    ('all', {}, {}),
+])
+@pytest.mark.parametrize('axis', [None, 0, 1, 2, (0, 2), -3, (1, -1)])
+@pytest.mark.parametrize('keepdims', [True, False])
+def test_reductions_bool(reduction, axis, keepdims, kwargs, eqkwargs):
+    x = sparse.random((2, 3, 4), density=.25).astype(bool)
+    y = x.todense()
+    xx = getattr(x, reduction)(axis=axis, keepdims=keepdims, **kwargs)
+    yy = getattr(y, reduction)(axis=axis, keepdims=keepdims, **kwargs)
+    assert_eq(xx, yy, **eqkwargs)
+
+
 @pytest.mark.parametrize('reduction,kwargs,eqkwargs', [
     (np.max, {}, {}),
     (np.sum, {}, {}),
@@ -433,25 +447,25 @@ def test_trinary_broadcasting(shapes, func):
 
 @pytest.mark.parametrize('shapes, func', [
     ([
-        (2,),
-        (3, 2),
-        (4, 3, 2),
-    ], lambda x, y, z: (x + y) * z),
+         (2,),
+         (3, 2),
+         (4, 3, 2),
+     ], lambda x, y, z: (x + y) * z),
     ([
-        (3,),
-        (2, 3),
-        (2, 2, 3),
-    ], lambda x, y, z: x * (y + z)),
+         (3,),
+         (2, 3),
+         (2, 2, 3),
+     ], lambda x, y, z: x * (y + z)),
     ([
-        (2,),
-        (2, 2),
-        (2, 2, 2),
-    ], lambda x, y, z: x * y * z),
+         (2,),
+         (2, 2),
+         (2, 2, 2),
+     ], lambda x, y, z: x * y * z),
     ([
-        (4,),
-        (4, 4),
-        (4, 4, 4),
-    ], lambda x, y, z: x + y + z),
+         (4,),
+         (4, 4),
+         (4, 4, 4),
+     ], lambda x, y, z: x + y + z),
 ])
 @pytest.mark.parametrize('value', [
     np.nan,
@@ -963,8 +977,11 @@ def test_slicing(index):
     (1, [2, 0], 0),
     ([True, False], slice(1, None), slice(-2, None)),
     (slice(1, None), slice(-2, None), [True, False, True, False]),
+    ([1, 0],),
+    (Ellipsis, [2, 1, 3],),
+    (slice(None), [2, 1, 2],),
+    (1, [2, 0, 1],),
 ])
-@pytest.mark.xfail(reason='Advanced indexing is temporarily broken.')
 def test_advanced_indexing(index):
     s = sparse.random((2, 3, 4), density=0.5)
     x = s.todense()
@@ -1001,6 +1018,8 @@ def test_custom_dtype_slicing():
     0.5,
     [0.5],
     {'potato': 'kartoffel'},
+    ([0, 1],) * 2,
+    ([[0, 1]],),
 ])
 def test_slicing_errors(index):
     s = sparse.random((2, 3, 4), density=0.5)