Implement diff (#781)

Author: tomwhite

api_status.md

@@ -90,7 +90,7 @@ This table shows which parts of the the [Array API](https://data-apis.org/array-
 |                          | `var`               | :white_check_mark: |            |                              |
 | Utility Functions        | `all`               | :white_check_mark: |            |                              |
 |                          | `any`               | :white_check_mark: |            |                              |
-|                          | `diff`              | :x:                | 2024.12    |                              |
+|                          | `diff`              | :white_check_mark: | 2024.12    |                              |
 
 ### Linear Algebra Extension
 

cubed/__init__.py

@@ -353,9 +353,9 @@
 
 __all__ += ["cumulative_sum", "max", "mean", "min", "prod", "std", "sum", "var"]
 
-from .array_api.utility_functions import all, any
+from .array_api.utility_functions import all, any, diff
 
-__all__ += ["all", "any"]
+__all__ += ["all", "any", "diff"]
 
 # extensions
 

cubed/array/overlap.py

@@ -127,10 +127,16 @@ def get_item_with_depth(
 ) -> Tuple[slice, ...]:
     """Convert a chunk index to a tuple of slices with depth offsets."""
     starts = tuple(_cumsum(c, initial_zero=True) for c in chunks)
+
+    def depth_offsets(d):
+        if isinstance(d, int):
+            return -d, d
+        return d
+
     loc = tuple(
         (
-            _clamp(0, start[i] - depth[ax], start[-1]),
-            _clamp(0, start[i + 1] + depth[ax], start[-1]),
+            _clamp(0, start[i] + depth_offsets(depth[ax])[0], start[-1]),
+            _clamp(0, start[i + 1] + depth_offsets(depth[ax])[1], start[-1]),
         )
         for ax, (i, start) in enumerate(zip(idx, starts))
     )
@@ -140,7 +146,7 @@ def get_item_with_depth(
 def _pad_boundaries(x, depth, boundary, numblocks, block_id):
     for i in range(x.ndim):
         d = depth.get(i, 0)
-        if d == 0 or block_id[i] not in (0, numblocks[i] - 1):
+        if d == 0 or block_id[i] not in (0, numblocks[i] - 1) or boundary[i] == "none":
             continue
         pad_shape = list(x.shape)
         pad_shape[i] = d

cubed/array_api/__init__.py

@@ -276,6 +276,6 @@
 
 __all__ += ["cumulative_sum", "max", "mean", "min", "prod", "std", "sum", "var"]
 
-from .utility_functions import all, any
+from .utility_functions import all, any, diff
 
-__all__ += ["all", "any"]
+__all__ += ["all", "any", "diff"]

cubed/array_api/utility_functions.py

@@ -1,6 +1,10 @@
+from cubed.array.overlap import map_overlap
 from cubed.array_api.creation_functions import asarray
+from cubed.array_api.manipulation_functions import concat
 from cubed.backend_array_api import namespace as nxp
 from cubed.core import reduction
+from cubed.vendor.dask.array.core import normalize_chunks
+from cubed.vendor.dask.array.utils import validate_axis
 
 
 def all(x, /, *, axis=None, keepdims=False, split_every=None):
@@ -27,3 +31,34 @@ def any(x, /, *, axis=None, keepdims=False, split_every=None):
         keepdims=keepdims,
         split_every=split_every,
     )
+
+
+def diff(x, /, *, axis=-1, n=1, prepend=None, append=None):
+    axis = validate_axis(axis, x.ndim)
+
+    if n < 0:
+        raise ValueError(f"order of diff must be non-negative, but was {n}")
+    if n == 0:
+        return x
+
+    combined = []
+    if prepend is not None:
+        combined.append(prepend)
+    combined.append(x)
+    if append is not None:
+        combined.append(append)
+    if len(combined) > 1:
+        x = concat(combined, axis=axis, chunks=x.chunksize)
+
+    shape = tuple(s - n if i == axis else s for i, s in enumerate(x.shape))
+    chunks = normalize_chunks(x.chunksize, shape, dtype=x.dtype)
+    depth = {axis: (0, n)}  # only need look-ahead values for differencing
+    return map_overlap(
+        nxp.diff,
+        x,
+        dtype=x.dtype,
+        chunks=chunks,
+        depth=depth,
+        axis=axis,
+        n=n,
+    )

cubed/tests/test_array_api.py

@@ -970,3 +970,24 @@ def test_all_zero_dimension(spec, executor):
     assert b.ndim == 0
     assert b.size == 1
     assert b.compute(executor=executor)
+
+
+@pytest.mark.parametrize("n", [1, 2])
+def test_diff(n):
+    x = np.array([1, 5, 3, 8, 7, 2, 6, 9])
+    a = xp.asarray(x, chunks=(3,))
+    b = xp.diff(a, n=n)
+
+    assert_array_equal(b.compute(), np.diff(x, n=n))
+
+
+@pytest.mark.parametrize(
+    ("shape", "axis"),
+    [((10, 15, 20), 0), ((10, 15, 20), 1), ((10, 15, 20), 2), ((10, 15, 20), -1)],
+)
+@pytest.mark.parametrize("n", [0, 1, 2])
+def test_diff_3d(shape, n, axis):
+    x = np.random.default_rng().integers(0, 10, shape)
+    a = xp.asarray(x, chunks=(len(shape) * (5,)))
+
+    assert_array_equal(xp.diff(a, axis=axis, n=n), np.diff(x, axis=axis, n=n))