diff --git a/cunumeric/module.py b/cunumeric/module.py index 0a8132a81..61e1111bc 100644 --- a/cunumeric/module.py +++ b/cunumeric/module.py @@ -44,7 +44,7 @@ from ._ufunc.comparison import maximum, minimum from ._ufunc.floating import floor, isnan -from ._ufunc.math import add, multiply +from ._ufunc.math import add, multiply, subtract from ._unary_red_utils import get_non_nan_unary_red_code from .array import ( add_boilerplate, @@ -6242,6 +6242,142 @@ def nansum( # Arithmetic operations +@add_boilerplate("a") +def diff( + a: ndarray, + n: int = 1, + axis: int = -1, + prepend: Any = None, + append: Any = None, +) -> ndarray: + """ + Calculate the n-th discrete difference along the given axis. + The first difference is given by ``out[i] = a[i+1] - a[i]`` along + the given axis, higher differences are calculated by using `diff` + recursively. + Parameters + ---------- + a : array_like + Input array + n : int, optional + The number of times values are differenced. If zero, the input + is returned as-is. + axis : int, optional + The axis along which the difference is taken, default is the + last axis. + prepend, append : array_like, optional + Values to prepend or append to `a` along axis prior to + performing the difference. Scalar values are expanded to + arrays with length 1 in the direction of axis and the shape + of the input array in along all other axes. Otherwise the + dimension and shape must match `a` except along axis. + Returns + ------- + diff : ndarray + The n-th differences. The shape of the output is the same as `a` + except along `axis` where the dimension is smaller by `n`. The + type of the output is the same as the type of the difference + between any two elements of `a`. This is the same as the type of + `a` in most cases. A notable exception is `datetime64`, which + results in a `timedelta64` output array. + See Also + -------- + gradient, ediff1d, cumsum + Notes + ----- + Type is preserved for boolean arrays, so the result will contain + `False` when consecutive elements are the same and `True` when they + differ. + For unsigned integer arrays, the results will also be unsigned. This + should not be surprising, as the result is consistent with + calculating the difference directly: + >>> u8_arr = np.array([1, 0], dtype=np.uint8) + >>> np.diff(u8_arr) + array([255], dtype=uint8) + >>> u8_arr[1,...] - u8_arr[0,...] + 255 + If this is not desirable, then the array should be cast to a larger + integer type first: + >>> i16_arr = u8_arr.astype(np.int16) + >>> np.diff(i16_arr) + array([-1], dtype=int16) + Examples + -------- + >>> x = np.array([1, 2, 4, 7, 0]) + >>> np.diff(x) + array([ 1, 2, 3, -7]) + >>> np.diff(x, n=2) + array([ 1, 1, -10]) + >>> x = np.array([[1, 3, 6, 10], [0, 5, 6, 8]]) + >>> np.diff(x) + array([[2, 3, 4], + [5, 1, 2]]) + >>> np.diff(x, axis=0) + array([[-1, 2, 0, -2]]) + >>> x = np.arange('1066-10-13', '1066-10-16', dtype=np.datetime64) + >>> np.diff(x) + array([1, 1], dtype='timedelta64[D]') + + Availability + -------- + Multiple GPUs, Multiple CPUs + """ + if n == 0: + return a + if n < 0: + raise ValueError("order must be non-negative but got " + repr(n)) + + nd = a.ndim + if nd == 0: + raise ValueError( + "diff requires input that is at least one dimensional" + ) + axis = normalize_axis_index(axis, nd) + + combined = [] + if prepend is not None: + prepend = np.asanyarray(prepend) + if prepend.ndim == 0: + shape = list(a.shape) + shape[axis] = 1 + prepend = np.broadcast_to(prepend, tuple(shape)) + combined.append(prepend) + + combined.append(a) + + if append is not None: + append = np.asanyarray(append) + if append.ndim == 0: + shape = list(a.shape) + shape[axis] = 1 + append = np.broadcast_to(append, tuple(shape)) + combined.append(append) + + if len(combined) > 1: + a = np.concatenate(combined, axis) + + # Diffing with n > shape results in an empty array. We have + # to handle this case explicitly as our slicing routines raise + # an exception with out-of-bounds slices, while NumPy's dont. + if a.shape[axis] <= n: + shape = list(a.shape) + shape[axis] = 0 + return np.empty(shape=shape, dtype=a.dtype) + + slice1l = [slice(None)] * nd + slice2l = [slice(None)] * nd + slice1l[axis] = slice(1, None) + slice2l[axis] = slice(None, -1) + slice1 = tuple(slice1l) + slice2 = tuple(slice2l) + + op = not_equal if a.dtype == np.bool_ else subtract + for _ in range(n): + a = op(a[slice1], a[slice2]) + + return a + + # Handling complex numbers diff --git a/tests/integration/test_diff.py b/tests/integration/test_diff.py new file mode 100644 index 000000000..29d2c598d --- /dev/null +++ b/tests/integration/test_diff.py @@ -0,0 +1,63 @@ +# Copyright 2022 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import numpy as np +import pytest +from utils.comparisons import allclose + +import cunumeric as cn + + +@pytest.mark.parametrize("args", [ + ((100,), 1, -1, None, None), + ((100,), 2, -1, None, None), + ((100,), 3, -1, None, None), + ((100,), 2, 0, None, None), + ((10, 10), 2, -1, None, None), + ((10, 10), 2, 0, None, None), + ((10, 10), 2, 1, None, None), + ((100,), 3, -1, [1.0, 2.0], None), + ((100,), 3, -1, None, [1.0, 2.0]), + ((100,), 3, -1, [1.0, 2.0], [1.0, 2.0]), + ((5,), 5, -1, None, None), + ((5,), 6, 0, None, None), + ((5,5), 5, 1, None, None), + ((5,5), 6, 1, None, None), +]) +def test_diff(args): + shape, n, axis, prepend, append = args + num = np.random.random(shape) + cun = cn.array(num) + + # We are not adopting the np._NoValue default arguments + # for this function, as no special behavior is needed on None. + n_prepend = np._NoValue if prepend is None else prepend + n_append = np._NoValue if append is None else append + res_num = np.diff(num, n=n, axis=axis, prepend=n_prepend, append=n_append) + res_cn = cn.diff(cun, n=n, axis=axis, prepend=prepend, append=append) + + assert allclose(res_num, res_cn) + + +def test_diff_nzero(): + a = cn.ones(100) + ad = cn.diff(a, n=0) + assert a is ad + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main(sys.argv))