Repository: pydata/sparse (python)

Issue:
np.diff equivalent
It would be nice to have the equivalent of `np.diff()` for COO arrays.


Pull request:
Adding `diff` function
<!--
# Thanks for contributing a pull request!
## Please make sure you see our contribution guidelines: https://github.com/pydata/sparse/blob/main/docs/contributing.md

Your PR title should start with any of these abbreviatons: `build`, `chore`, `ci`, `depr`, `docs`, `feat`, `fix`, `perf`, `refactor`, `release`, `test`. Add a `!`at the end, if it is a breaking change.
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] 🪄 Feature
- [ ] 🐞 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📚 Documentation
- [x] 🧪 Test
- [ ] 🛠️ Other

## Related issues

- Closes #592

## Checklist

- [x] Code follows style guide
- [x] Tests added
- [ ] Documented the changes

***

## Please explain your changes below.
Added sparse equivalent of `np.diff`

Patch:
diff --git a/ci/Numba-array-api-xfails.txt b/ci/Numba-array-api-xfails.txt
index 94a55c2..340469f 100644
--- a/ci/Numba-array-api-xfails.txt
+++ b/ci/Numba-array-api-xfails.txt
@@ -69,13 +69,11 @@ array_api_tests/test_has_names.py::test_has_names[fft-irfftn]
 array_api_tests/test_creation_functions.py::test_empty_like
 array_api_tests/test_data_type_functions.py::test_finfo[complex64]
 array_api_tests/test_manipulation_functions.py::test_squeeze
-array_api_tests/test_has_names.py::test_has_names[utility-diff]
 array_api_tests/test_has_names.py::test_has_names[statistical-cumulative_sum]
 array_api_tests/test_has_names.py::test_has_names[statistical-cumulative_prod]
 array_api_tests/test_has_names.py::test_has_names[indexing-take_along_axis]
 array_api_tests/test_has_names.py::test_has_names[searching-count_nonzero]
 array_api_tests/test_has_names.py::test_has_names[searching-searchsorted]
-array_api_tests/test_signatures.py::test_func_signature[diff]
 array_api_tests/test_signatures.py::test_func_signature[take_along_axis]
 array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i > 0) -> +infinity]
 array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i < 0) -> -infinity]
diff --git a/sparse/numba_backend/__init__.py b/sparse/numba_backend/__init__.py
index 9226da1..a4111c2 100644
--- a/sparse/numba_backend/__init__.py
+++ b/sparse/numba_backend/__init__.py
@@ -95,6 +95,7 @@ from ._common import (
     can_cast,
     concat,
     concatenate,
+    diff,
     dot,
     einsum,
     empty,
@@ -341,6 +342,7 @@ __all__ = [
     "repeat",
     "tile",
     "unstack",
+    "diff",
 ]
 
 
diff --git a/sparse/numba_backend/_common.py b/sparse/numba_backend/_common.py
index 24b1768..04bd1a3 100644
--- a/sparse/numba_backend/_common.py
+++ b/sparse/numba_backend/_common.py
@@ -3217,3 +3217,36 @@ def unstack(x, axis=0):
     new_order = (axis,) + tuple(i for i in range(ndim) if i != axis)
     x = x.transpose(new_order)
     return (*x,)
+
+
+def diff(x, axis=-1, n=1, prepend=None, append=None):
+    """
+    Calculates the n-th discrete difference along the given axis.
+
+    Parameters
+    ----------
+    x : SparseArray
+        Input sparse arrays.
+    n : int
+        The number of times values are differenced. Default: 1.
+    axis : int
+        The axis along which the difference is taken. Default: -1.
+
+    Returns
+    -------
+    out : SparseArray
+        An array containing the n-th discrete difference along the given axis.
+    """
+    if not isinstance(x, SparseArray):
+        raise TypeError("`x` must be a SparseArray.")
+
+    if axis < 0:
+        axis = x.ndim + axis
+    if prepend is not None:
+        x = concatenate([prepend, x], axis=axis)
+    if append is not None:
+        x = concatenate([x, append], axis=axis)
+    result = x
+    for _ in range(n):
+        result = result[(slice(None),) * axis + (slice(1, None),)] - result[(slice(None),) * axis + (slice(None, -1),)]
+    return result


Test patch:
diff --git a/sparse/numba_backend/tests/test_coo.py b/sparse/numba_backend/tests/test_coo.py
index c5882e1..168f202 100644
--- a/sparse/numba_backend/tests/test_coo.py
+++ b/sparse/numba_backend/tests/test_coo.py
@@ -2030,3 +2030,43 @@ def test_unstack_invalid_type():
     a = np.arange(6).reshape(2, 3)  # not a sparse array
     with pytest.raises(TypeError, match="must be a SparseArray"):
         sparse.unstack(a, axis=0)
+
+
+@pytest.mark.parametrize("ndim", range(1, 4))
+@pytest.mark.parametrize("shape_range", [3])
+@pytest.mark.parametrize("n", [1, 2])
+@pytest.mark.parametrize("use_prepend, use_append", [(False, False), (True, False), (False, True), (True, True)])
+def test_diff_matches_numpy(ndim, shape_range, n, use_prepend, use_append):
+    rng = np.random.default_rng(42)
+    shape = tuple(rng.integers(2, shape_range + 2) for _ in range(ndim))
+    x = rng.integers(0, 10, size=shape)
+    sparse_x = COO.from_numpy(x)
+
+    for axis in range(-ndim, ndim):
+        prepend = rng.integers(0, 10, size=x.shape).astype(x.dtype) if use_prepend else None
+        append = rng.integers(0, 10, size=x.shape).astype(x.dtype) if use_append else None
+
+        sparse_prepend = COO.from_numpy(prepend) if prepend is not None else None
+        sparse_append = COO.from_numpy(append) if append is not None else None
+
+        sparse_result = sparse.diff(sparse_x, axis=axis, n=n, prepend=sparse_prepend, append=sparse_append)
+
+        kwargs = {}
+        if prepend is not None:
+            kwargs["prepend"] = prepend
+        if append is not None:
+            kwargs["append"] = append
+
+        dense_result = np.diff(x, axis=axis, n=n, **kwargs)
+
+        np.testing.assert_array_equal(
+            sparse_result.todense(),
+            dense_result,
+            err_msg=f"Mismatch at axis={axis}, n={n}, prepend={use_prepend}, append={use_append}",
+        )
+
+
+def test_diff_invalid_type():
+    a = np.arange(6).reshape(2, 3)
+    with pytest.raises(TypeError, match="must be a SparseArray"):
+        sparse.diff(a)
diff --git a/sparse/numba_backend/tests/test_namespace.py b/sparse/numba_backend/tests/test_namespace.py
index 40ce4fe..2744ebe 100644
--- a/sparse/numba_backend/tests/test_namespace.py
+++ b/sparse/numba_backend/tests/test_namespace.py
@@ -51,6 +51,7 @@ def test_namespace():
         "cosh",
         "diagonal",
         "diagonalize",
+        "diff",
         "divide",
         "dot",
         "e",