diff options
Diffstat (limited to 'numpy/lib')
| -rw-r--r-- | numpy/lib/function_base.py | 30 | ||||
| -rw-r--r-- | numpy/lib/tests/test_function_base.py | 52 |
2 files changed, 81 insertions, 1 deletions
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py index 2992e92bb..704e56514 100644 --- a/numpy/lib/function_base.py +++ b/numpy/lib/function_base.py @@ -1090,7 +1090,7 @@ def gradient(f, *varargs, **kwargs): return outvals -def diff(a, n=1, axis=-1): +def diff(a, n=1, axis=-1, prepend=np._NoValue, append=np._NoValue): """ Calculate the n-th discrete difference along the given axis. @@ -1108,6 +1108,12 @@ def diff(a, n=1, axis=-1): 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 ------- @@ -1176,6 +1182,28 @@ def diff(a, n=1, axis=-1): nd = a.ndim axis = normalize_axis_index(axis, nd) + combined = [] + if prepend is not np._NoValue: + 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 np._NoValue: + 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) + slice1 = [slice(None)] * nd slice2 = [slice(None)] * nd slice1[axis] = slice(1, None) diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py index d5faed6ae..40cca1dbb 100644 --- a/numpy/lib/tests/test_function_base.py +++ b/numpy/lib/tests/test_function_base.py @@ -734,6 +734,58 @@ class TestDiff(object): assert_array_equal(out3.mask, [[], [], [], [], []]) assert_(type(out3) is type(x)) + def test_prepend(self): + x = np.arange(5) + 1 + assert_array_equal(diff(x, prepend=0), np.ones(5)) + assert_array_equal(diff(x, prepend=[0]), np.ones(5)) + assert_array_equal(np.cumsum(np.diff(x, prepend=0)), x) + assert_array_equal(diff(x, prepend=[-1, 0]), np.ones(6)) + + x = np.arange(4).reshape(2, 2) + result = np.diff(x, axis=1, prepend=0) + expected = [[0, 1], [2, 1]] + assert_array_equal(result, expected) + result = np.diff(x, axis=1, prepend=[[0], [0]]) + assert_array_equal(result, expected) + + result = np.diff(x, axis=0, prepend=0) + expected = [[0, 1], [2, 2]] + assert_array_equal(result, expected) + result = np.diff(x, axis=0, prepend=[[0, 0]]) + assert_array_equal(result, expected) + + assert_raises(ValueError, np.diff, x, prepend=np.zeros((3,3))) + + assert_raises(np.AxisError, diff, x, prepend=0, axis=3) + + def test_append(self): + x = np.arange(5) + result = diff(x, append=0) + expected = [1, 1, 1, 1, -4] + assert_array_equal(result, expected) + result = diff(x, append=[0]) + assert_array_equal(result, expected) + result = diff(x, append=[0, 2]) + expected = expected + [2] + assert_array_equal(result, expected) + + x = np.arange(4).reshape(2, 2) + result = np.diff(x, axis=1, append=0) + expected = [[1, -1], [1, -3]] + assert_array_equal(result, expected) + result = np.diff(x, axis=1, append=[[0], [0]]) + assert_array_equal(result, expected) + + result = np.diff(x, axis=0, append=0) + expected = [[2, 2], [-2, -3]] + assert_array_equal(result, expected) + result = np.diff(x, axis=0, append=[[0, 0]]) + assert_array_equal(result, expected) + + assert_raises(ValueError, np.diff, x, append=np.zeros((3,3))) + + assert_raises(np.AxisError, diff, x, append=0, axis=3) + class TestDelete(object): |