diff options
Diffstat (limited to 'numpy/lib')
| -rw-r--r-- | numpy/lib/function_base.py | 54 | ||||
| -rw-r--r-- | numpy/lib/tests/test_function_base.py | 53 |
2 files changed, 84 insertions, 23 deletions
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py index dda3e021b..f5e9ff2a5 100644 --- a/numpy/lib/function_base.py +++ b/numpy/lib/function_base.py @@ -15,7 +15,7 @@ from numpy.core.numeric import ( ) from numpy.core.umath import ( pi, multiply, add, arctan2, frompyfunc, cos, less_equal, sqrt, sin, - mod, exp, log10 + mod, exp, log10, not_equal, subtract ) from numpy.core.fromnumeric import ( ravel, nonzero, sort, partition, mean, any, sum @@ -1827,7 +1827,7 @@ def gradient(f, *varargs, **kwargs): def diff(a, n=1, axis=-1): """ - Calculate the n-th discrete difference along given axis. + Calculate the n-th discrete difference along the given axis. The first difference is given by ``out[n] = a[n+1] - a[n]`` along the given axis, higher differences are calculated by using `diff` @@ -1838,16 +1838,21 @@ def diff(a, n=1, axis=-1): a : array_like Input array n : int, optional - The number of times values are differenced. + 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. + The axis along which the difference is taken, default is the + last 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 that of the input. + 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 -------- @@ -1855,13 +1860,13 @@ def diff(a, n=1, axis=-1): Notes ----- - For boolean arrays, the preservation of type means that the result - will contain `False` when consecutive elements are the same and - `True` when they differ. + 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: + 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) @@ -1869,8 +1874,8 @@ def diff(a, n=1, axis=-1): >>> u8_arr[1,...] - u8_arr[0,...] array(255, np.uint8) - If this is not desirable, then the array should be cast to a larger integer - type first: + 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) @@ -1891,30 +1896,33 @@ def diff(a, n=1, axis=-1): >>> 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]') + """ if n == 0: return a if n < 0: raise ValueError( "order must be non-negative but got " + repr(n)) + a = asanyarray(a) nd = a.ndim - slice1 = [slice(None)]*nd - slice2 = [slice(None)]*nd + axis = normalize_axis_index(axis, nd) + + slice1 = [slice(None)] * nd + slice2 = [slice(None)] * nd slice1[axis] = slice(1, None) slice2[axis] = slice(None, -1) slice1 = tuple(slice1) slice2 = tuple(slice2) - if a.dtype == np.bool_: - da = a[slice1] != a[slice2] - else: - da = a[slice1] - a[slice2] + op = not_equal if a.dtype == np.bool_ else subtract + for _ in range(n): + a = op(a[slice1], a[slice2]) - if n > 1: - return diff(da, n-1, axis=axis) - else: - return da + return a def interp(x, xp, fp, left=None, right=None, period=None): diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py index 099a2d407..4ecb02821 100644 --- a/numpy/lib/tests/test_function_base.py +++ b/numpy/lib/tests/test_function_base.py @@ -6,6 +6,7 @@ import sys import decimal import numpy as np +from numpy import ma from numpy.testing import ( run_module_suite, assert_, assert_equal, assert_array_equal, assert_almost_equal, assert_array_almost_equal, assert_raises, @@ -647,6 +648,19 @@ class TestDiff(object): assert_array_equal(diff(x), out) assert_array_equal(diff(x, n=2), out2) + def test_axis(self): + x = np.zeros((10, 20, 30)) + x[:, 1::2, :] = 1 + exp = np.ones((10, 19, 30)) + exp[:, 1::2, :] = -1 + assert_array_equal(diff(x), np.zeros((10, 20, 29))) + assert_array_equal(diff(x, axis=-1), np.zeros((10, 20, 29))) + assert_array_equal(diff(x, axis=0), np.zeros((9, 20, 30))) + assert_array_equal(diff(x, axis=1), exp) + assert_array_equal(diff(x, axis=-2), exp) + assert_raises(np.AxisError, diff, x, axis=3) + assert_raises(np.AxisError, diff, x, axis=-4) + def test_nd(self): x = 20 * rand(10, 20, 30) out1 = x[:, :, 1:] - x[:, :, :-1] @@ -658,6 +672,45 @@ class TestDiff(object): assert_array_equal(diff(x, axis=0), out3) assert_array_equal(diff(x, n=2, axis=0), out4) + def test_n(self): + x = list(range(3)) + assert_raises(ValueError, diff, x, n=-1) + output = [diff(x, n=n) for n in range(1, 5)] + expected = [[1, 1], [0], [], []] + assert_(diff(x, n=0) is x) + for n, (expected, out) in enumerate(zip(expected, output), start=1): + assert_(type(out) is np.ndarray) + assert_array_equal(out, expected) + assert_equal(out.dtype, np.int_) + assert_equal(len(out), max(0, len(x) - n)) + + def test_times(self): + x = np.arange('1066-10-13', '1066-10-16', dtype=np.datetime64) + expected = [ + np.array([1, 1], dtype='timedelta64[D]'), + np.array([0], dtype='timedelta64[D]'), + ] + expected.extend([np.array([], dtype='timedelta64[D]')] * 3) + for n, exp in enumerate(expected, start=1): + out = diff(x, n=n) + assert_array_equal(out, exp) + assert_equal(out.dtype, exp.dtype) + + def test_subclass(self): + x = ma.array([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]], + mask=[[False, False], [True, False], + [False, True], [True, True], [False, False]]) + out = diff(x) + assert_array_equal(out.data, [[1], [1], [1], [1], [1]]) + assert_array_equal(out.mask, [[False], [True], + [True], [True], [False]]) + assert_(type(out) is type(x)) + + out3 = diff(x, n=3) + assert_array_equal(out3.data, [[], [], [], [], []]) + assert_array_equal(out3.mask, [[], [], [], [], []]) + assert_(type(out3) is type(x)) + class TestDelete(object): |