diff --git a/ci/Numba-array-api-xfails.txt b/ci/Numba-array-api-xfails.txt index 94a55c29..340469f8 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 9226da18..a4111c25 100644 --- a/sparse/numba_backend/__init__.py +++ b/sparse/numba_backend/__init__.py @@ -95,6 +95,7 @@ can_cast, concat, concatenate, + diff, dot, einsum, empty, @@ -341,6 +342,7 @@ "repeat", "tile", "unstack", + "diff", ] diff --git a/sparse/numba_backend/_common.py b/sparse/numba_backend/_common.py index 24b17686..04bd1a32 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 diff --git a/sparse/numba_backend/tests/test_coo.py b/sparse/numba_backend/tests/test_coo.py index c5882e17..168f2028 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 40ce4fe1..2744ebe2 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",