API: Add set functions

sparse/__init__.py

@@ -50,6 +50,10 @@
     tril,
     triu,
     where,
+    unique_all,
+    unique_counts,
+    unique_inverse,
+    unique_values,
 )
 from ._dok import DOK
 from ._io import load_npz, save_npz
@@ -114,6 +118,10 @@
     "min",
     "max",
     "nanreduce",
+    "unique_all",
+    "unique_counts",
+    "unique_inverse",
+    "unique_values",
 ]
 
 __array_api_version__ = "2022.12"

sparse/_coo/__init__.py

@@ -21,6 +21,10 @@
     tril,
     triu,
     where,
+    unique_all,
+    unique_counts,
+    unique_inverse,
+    unique_values,
 )
 from .core import COO, as_coo
 
@@ -49,4 +53,8 @@
     "result_type",
     "diagonal",
     "diagonalize",
+    "unique_all",
+    "unique_counts",
+    "unique_inverse",
+    "unique_values",
 ]

sparse/_coo/common.py

@@ -1059,6 +1059,98 @@ 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.")
+
+    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
+
+        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
+
+    return values, index, result_inverse, counts
+
+
+def unique_counts(x, /):
+    """
+    Returns the unique elements of an input array 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.")
+
+    x = x.flatten()
+    values, counts = np.unique(x.data, return_counts=True)
+    if x.nnz < x.size:
+        values = np.concatenate([[x.fill_value], values])
+        counts = np.concatenate([[x.size - x.nnz], counts])
+    return 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
+
+    from .._dok import DOK
+    result_inverse = DOK(shape=x.size, dtype=np.intp, fill_value=np.intp(0))
+    result_inverse[x.coords.squeeze()] = inverse
+
+    return values, result_inverse
+
+
+def unique_values(x, /):
+    """
+    Returns the unique elements of an input array 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 = np.unique(x.data)
+    if x.nnz < x.size:
+        values = np.concatenate([[x.fill_value], values])
+    return values
+
+
 @numba.jit(nopython=True, nogil=True)
 def _compute_minmax_args(
     coords: np.ndarray,

sparse/tests/test_coo.py

@@ -1745,3 +1745,59 @@ def test_squeeze_validation(self):
 
         with pytest.raises(ValueError, match="Specified axis `0` has a size greater than one: 3"):
             s_arr.squeeze(0)
+
+
+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_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)
+
+        result_values, result_counts = sparse.unique_counts(s_arr)
+        expected_values, expected_counts = np.unique(arr, return_counts=True)
+
+        np.testing.assert_equal(result_values, expected_values)
+        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)
+
+        result = sparse.unique_values(s_arr)
+        expected = np.unique(arr)
+
+        np.testing.assert_equal(result, expected)