[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

sparse/__init__.py

@@ -49,11 +49,9 @@
     roll,
     tril,
     triu,
-    where,
-    unique_all,
     unique_counts,
-    unique_inverse,
     unique_values,
+    where,
 )
 from ._dok import DOK
 from ._io import load_npz, save_npz
@@ -118,9 +116,7 @@
     "min",
     "max",
     "nanreduce",
-    "unique_all",
     "unique_counts",
-    "unique_inverse",
     "unique_values",
 ]
 

sparse/_coo/__init__.py

@@ -20,11 +20,9 @@
     stack,
     tril,
     triu,
-    where,
-    unique_all,
     unique_counts,
-    unique_inverse,
     unique_values,
+    where,
 )
 from .core import COO, as_coo
 
@@ -53,8 +51,6 @@
     "result_type",
     "diagonal",
     "diagonalize",
-    "unique_all",
     "unique_counts",
-    "unique_inverse",
     "unique_values",
 ]

sparse/_coo/common.py

@@ -2,7 +2,7 @@
 import warnings
 from collections.abc import Iterable
 from functools import reduce
-from typing import Optional, Tuple
+from typing import NamedTuple, Optional, Tuple
 
 import numba
 
@@ -1059,50 +1059,44 @@ def clip(a, a_min=None, a_max=None, out=None):
     return a.clip(a_min, a_max)
 
 
-def unique_all(x, /):
-    """
-    Returns the unique elements of an input array x, the first occurring indices
-    for each unique element in x, the indices from the set of unique elements that
-    reconstruct x, and the corresponding counts for each unique element in x.
-    """
-    from .core import COO
-    if not isinstance(x, COO):
-        raise ValueError(f"Only COO arrays are supported but {type(x)} was passed.")
+# Array API set functions
 
-    x = x.flatten()
-    values, index, inverse, counts = np.unique(
-        x.data, return_index=True, return_inverse=True, return_counts=True
-    )
-    index = x.coords.squeeze()[index]
-    if x.nnz < x.size:
-        # find the first occurence of the fill value
-        last_idx = -1
-        first_fill_value = x.coords.max() + 1 if x.coords.size > 0 else 0
-        for idx in np.nditer(x.coords, flags=["zerosize_ok"]):
-            if idx - last_idx > 1:
-                first_fill_value = last_idx + 1
-                break
-            else:
-                last_idx = idx
+class UniqueCountsResult(NamedTuple):
+    values: np.ndarray
+    counts: np.ndarray
 
-        values = np.concatenate([[x.fill_value], values])
-        index = np.concatenate([[first_fill_value], index])
-        inverse = inverse + 1
-        counts = np.concatenate([[x.size - x.nnz], counts])
 
-    from .._dok import DOK
-    result_inverse = DOK(shape=x.size, dtype=np.intp, fill_value=np.intp(0))
-    result_inverse[x.coords.squeeze()] = inverse
+def unique_counts(x, /):
+    """
+    Returns the unique elements of an input array `x`, and the corresponding
+    counts for each unique element in `x`.
+
+    Parameters
+    ----------
+    x : COO
+        Input COO array. It will be flattened if it is not already 1-D.
 
-    return values, index, result_inverse, counts
+    Returns
+    -------
+    out : namedtuple
+        The result containing:
+        * values - The unique elements of an input array.
+        * counts - The corresponding counts for each unique element.
 
+    Raises
+    ------
+    ValueError
+        If the input array is in a different format than COO.
 
-def unique_counts(x, /):
-    """
-    Returns the unique elements of an input array x and the corresponding
-    counts for each unique element in x.
+    Examples
+    --------
+    >>> import sparse
+    >>> x = sparse.COO.from_numpy([1, 0, 2, 1, 2, -3])
+    >>> sparse.unique_counts(x)
+    UniqueCountsResult(values=array([-3,  0,  1,  2]), counts=array([1, 1, 2, 2]))
     """
     from .core import COO
+
     if not isinstance(x, COO):
         raise ValueError(f"Only COO arrays are supported but {type(x)} was passed.")
 
@@ -1111,43 +1105,48 @@ def unique_counts(x, /):
     if x.nnz < x.size:
         values = np.concatenate([[x.fill_value], values])
         counts = np.concatenate([[x.size - x.nnz], counts])
-    return values, counts
+        sorted_indices = np.argsort(values)
+        values[sorted_indices] = values.copy()
+        counts[sorted_indices] = counts.copy()
 
+    return UniqueCountsResult(values, counts)
 
-def unique_inverse(x, /):
-    """
-    Returns the unique elements of an input array x and the indices from
-    the set of unique elements that reconstruct x.
-    """
-    from .core import COO
-    if not isinstance(x, COO):
-        raise ValueError(f"Only COO arrays are supported but {type(x)} was passed.")
 
-    x = x.flatten()
-    values, inverse = np.unique(x.data, return_inverse=True)
-    if x.nnz < x.size:
-        values = np.concatenate([[x.fill_value], values])
-        inverse = inverse + 1
+def unique_values(x, /):
+    """
+    Returns the unique elements of an input array `x`.
 
-    from .._dok import DOK
-    result_inverse = DOK(shape=x.size, dtype=np.intp, fill_value=np.intp(0))
-    result_inverse[x.coords.squeeze()] = inverse
+    Parameters
+    ----------
+    x : COO
+        Input COO array. It will be flattened if it is not already 1-D.
 
-    return values, result_inverse
+    Returns
+    -------
+    out : ndarray
+        The unique elements of an input array.
 
+    Raises
+    ------
+    ValueError
+        If the input array is in a different format than COO.
 
-def unique_values(x, /):
-    """
-    Returns the unique elements of an input array x.
+    Examples
+    --------
+    >>> import sparse
+    >>> x = sparse.COO.from_numpy([1, 0, 2, 1, 2, -3])
+    >>> sparse.unique_values(x)
+    array([-3, 0, 1, 2])
     """
     from .core import COO
+
     if not isinstance(x, COO):
         raise ValueError(f"Only COO arrays are supported but {type(x)} was passed.")
 
     x = x.flatten()
     values = np.unique(x.data)
     if x.nnz < x.size:
-        values = np.concatenate([[x.fill_value], values])
+        values = np.sort(np.concatenate([[x.fill_value], values]))
     return values
 
 

sparse/tests/test_coo.py

@@ -1748,34 +1748,14 @@ def test_squeeze_validation(self):
 
 
 class TestUnique:
-    arr = np.array(
-        [[0, 0, 1, 5, 3, 0],
-         [1, 0, 4, 0, 3, 0],
-         [0, 1, 0, 1, 1, 0]],
-        dtype=np.int64
-    )
-    arr_empty = np.zeros((5,5))
+    arr = np.array([[0, 0, 1, 5, 3, 0], [1, 0, 4, 0, 3, 0], [0, 1, 0, 1, 1, 0]], dtype=np.int64)
+    arr_empty = np.zeros((5, 5))
     arr_full = np.arange(1, 10)
 
     @pytest.mark.parametrize("arr", [arr, arr_empty, arr_full])
-    def test_unique_all(self, arr):
-        s_arr = sparse.COO.from_numpy(arr)
-
-        result_values, result_indices, result_inverse, result_count = (
-            sparse.unique_all(s_arr)
-        )
-        expected_values, expected_indices, expected_inverse, expected_count = np.unique(
-            arr, return_index=True, return_inverse=True, return_counts=True
-        )
-
-        np.testing.assert_equal(result_values, expected_values)
-        np.testing.assert_equal(result_indices, expected_indices)
-        np.testing.assert_equal(result_inverse.todense(), expected_inverse)
-        np.testing.assert_equal(result_count, expected_count)
-
-    @pytest.mark.parametrize("arr", [arr, arr_empty, arr_full])
-    def test_unique_counts(self, arr):
-        s_arr = sparse.COO.from_numpy(arr)
+    @pytest.mark.parametrize("fill_value", [-1, 0, 1])
+    def test_unique_counts(self, arr, fill_value):
+        s_arr = sparse.COO.from_numpy(arr, fill_value)
 
         result_values, result_counts = sparse.unique_counts(s_arr)
         expected_values, expected_counts = np.unique(arr, return_counts=True)
@@ -1784,20 +1764,18 @@ def test_unique_counts(self, arr):
         np.testing.assert_equal(result_counts, expected_counts)
 
     @pytest.mark.parametrize("arr", [arr, arr_empty, arr_full])
-    def test_unique_inverse(self, arr):
-        s_arr = sparse.COO.from_numpy(arr)
-
-        result_values, result_inverse = sparse.unique_inverse(s_arr)
-        expected_values, expected_inverse = np.unique(arr, return_inverse=True)
-
-        np.testing.assert_equal(result_values, expected_values)
-        np.testing.assert_equal(result_inverse.todense(), expected_inverse)
-
-    @pytest.mark.parametrize("arr", [arr, arr_empty, arr_full])
-    def test_unique_values(self, arr):
-        s_arr = sparse.COO.from_numpy(arr)
+    @pytest.mark.parametrize("fill_value", [-1, 0, 1])
+    def test_unique_values(self, arr, fill_value):
+        s_arr = sparse.COO.from_numpy(arr, fill_value)
 
         result = sparse.unique_values(s_arr)
         expected = np.unique(arr)
 
         np.testing.assert_equal(result, expected)
+
+    @pytest.mark.parametrize(
+        "func", [sparse.unique_counts, sparse.unique_values]
+    )
+    def test_input_validation(self, func):
+        with pytest.raises(ValueError, match=r"Only COO arrays are supported"):
+            func(self.arr)