diff options
-rw-r--r-- | numpy/core/fromnumeric.py | 84 | ||||
-rw-r--r-- | numpy/core/src/multiarray/methods.c | 2 | ||||
-rw-r--r-- | numpy/core/src/multiarray/shape.c | 19 | ||||
-rw-r--r-- | numpy/core/tests/test_maskna.py | 69 |
4 files changed, 136 insertions, 38 deletions
diff --git a/numpy/core/fromnumeric.py b/numpy/core/fromnumeric.py index 03cb427cd..4d66b51b2 100644 --- a/numpy/core/fromnumeric.py +++ b/numpy/core/fromnumeric.py @@ -879,7 +879,7 @@ def squeeze(a, axis=None): axis : None or int or tuple of ints, optional Selects a subset of the single-dimensional entries in the shape. If an axis is selected with shape entry greater than - one, that axis is left untouched. + one, an error is raised. Returns ------- @@ -895,6 +895,8 @@ def squeeze(a, axis=None): (1, 3, 1) >>> np.squeeze(x).shape (3,) + >>> np.squeeze(x, axis=(2,)).shape + (1, 3) """ try: @@ -1383,7 +1385,7 @@ def clip(a, a_min, a_max, out=None): return clip(a_min, a_max, out) -def sum(a, axis=None, dtype=None, out=None, skipna=False): +def sum(a, axis=None, dtype=None, out=None, skipna=False, keepdims=False): """ Sum of array elements over a given axis. @@ -1416,6 +1418,10 @@ def sum(a, axis=None, dtype=None, out=None, skipna=False): skipna : bool, optional If this is set to True, skips any NA values during summation instead of propagating them. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the original `arr`. Returns ------- @@ -1470,14 +1476,14 @@ def sum(a, axis=None, dtype=None, out=None, skipna=False): sum = a.sum except AttributeError: return um.add.reduce(a, axis=axis, dtype=dtype, - out=out, skipna=skipna) - # NOTE: Dropping the skipna parameter here... + out=out, skipna=skipna, keepdims=keepdims) + # NOTE: Dropping the skipna and keepdims parameters here... return sum(axis=axis, dtype=dtype, out=out) else: return um.add.reduce(a, axis=axis, dtype=dtype, - out=out, skipna=skipna) + out=out, skipna=skipna, keepdims=keepdims) -def product (a, axis=None, dtype=None, out=None): +def product (a, axis=None, dtype=None, out=None, skipna=False, keepdims=False): """ Return the product of array elements over a given axis. @@ -1486,10 +1492,10 @@ def product (a, axis=None, dtype=None, out=None): prod : equivalent function; see for details. """ - return um.multiply.reduce(a, axis=axis, dtype=dtype, out=out) + return um.multiply.reduce(a, axis=axis, dtype=dtype, out=out, skipna=skipna, keepdims=keepdims) -def sometrue(a, axis=None, out=None): +def sometrue(a, axis=None, out=None, skipna=False, keepdims=False): """ Check whether some values are true. @@ -1500,10 +1506,10 @@ def sometrue(a, axis=None, out=None): any : equivalent function """ - return um.logical_or.reduce(a, axis=axis, out=out) + return um.logical_or.reduce(a, axis=axis, out=out, skipna=skipna, keepdims=keepdims) -def alltrue (a, axis=None, out=None): +def alltrue (a, axis=None, out=None, skipna=False, keepdims=False): """ Check if all elements of input array are true. @@ -1512,9 +1518,9 @@ def alltrue (a, axis=None, out=None): numpy.all : Equivalent function; see for details. """ - return um.logical_and.reduce(a, axis=axis, out=out) + return um.logical_and.reduce(a, axis=axis, out=out, skipna=skipna, keepdims=keepdims) -def any(a, axis=None, out=None, skipna=False): +def any(a, axis=None, out=None, skipna=False, keepdims=False): """ Test whether any array element along a given axis evaluates to True. @@ -1543,6 +1549,10 @@ def any(a, axis=None, out=None, skipna=False): skipna : bool, optional If this is set to True, skips any NA values during summation instead of propagating them. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the original `arr`. Returns ------- @@ -1586,9 +1596,9 @@ def any(a, axis=None, out=None, skipna=False): (191614240, 191614240) """ - return um.logical_or.reduce(a, axis=axis, out=out, skipna=skipna) + return um.logical_or.reduce(a, axis=axis, out=out, skipna=skipna, keepdims=keepdims) -def all(a, axis=None, out=None, skipna=False): +def all(a, axis=None, out=None, skipna=False, keepdims=False): """ Test whether all array elements along a given axis evaluate to True. @@ -1615,6 +1625,10 @@ def all(a, axis=None, out=None, skipna=False): skipna : bool, optional If this is set to True, skips any NA values during summation instead of propagating them. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the original `arr`. Returns ------- @@ -1653,7 +1667,7 @@ def all(a, axis=None, out=None, skipna=False): (28293632, 28293632, array([ True], dtype=bool)) """ - return um.logical_and.reduce(a, axis=axis, out=out, skipna=skipna) + return um.logical_and.reduce(a, axis=axis, out=out, skipna=skipna, keepdims=keepdims) def cumsum (a, axis=None, dtype=None, out=None): """ @@ -1786,7 +1800,7 @@ def ptp(a, axis=None, out=None): return ptp(axis, out) -def amax(a, axis=None, out=None, skipna=False): +def amax(a, axis=None, out=None, skipna=False, keepdims=False): """ Return the maximum of an array or maximum along an axis. @@ -1801,8 +1815,12 @@ def amax(a, axis=None, out=None, skipna=False): the same shape and buffer length as the expected output. See `doc.ufuncs` (Section "Output arguments") for more details. skipna : bool, optional - If this is set to True, skips any NA values during summation + If this is set to True, skips any NA values during reduction instead of propagating them. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the original `arr`. Returns ------- @@ -1849,14 +1867,14 @@ def amax(a, axis=None, out=None, skipna=False): amax = a.max except AttributeError: return um.maximum.reduce(a, axis=axis, - out=out, skipna=skipna) - # NOTE: Dropping the skipna parameter + out=out, skipna=skipna, keepdims=keepdims) + # NOTE: Dropping the skipna and keepdims parameters return amax(axis=axis, out=out) else: return um.maximum.reduce(a, axis=axis, - out=out, skipna=skipna) + out=out, skipna=skipna, keepdims=keepdims) -def amin(a, axis=None, out=None, skipna=False): +def amin(a, axis=None, out=None, skipna=False, keepdims=False): """ Return the minimum of an array or minimum along an axis. @@ -1871,8 +1889,12 @@ def amin(a, axis=None, out=None, skipna=False): be of the same shape and buffer length as the expected output. See `doc.ufuncs` (Section "Output arguments") for more details. skipna : bool, optional - If this is set to True, skips any NA values during summation + If this is set to True, skips any NA values during reduction instead of propagating them. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the original `arr`. Returns ------- @@ -1919,12 +1941,12 @@ def amin(a, axis=None, out=None, skipna=False): amin = a.min except AttributeError: return um.minimum.reduce(a, axis=axis, - out=out, skipna=skipna) - # NOTE: Dropping the skipna parameter + out=out, skipna=skipna, keepdims=keepdims) + # NOTE: Dropping the skipna and keepdims parameters return amin(axis=axis, out=out) else: return um.minimum.reduce(a, axis=axis, - out=out, skipna=skipna) + out=out, skipna=skipna, keepdims=keepdims) def alen(a): """ @@ -1959,7 +1981,7 @@ def alen(a): return len(array(a,ndmin=1)) -def prod(a, axis=None, dtype=None, out=None, skipna=False): +def prod(a, axis=None, dtype=None, out=None, skipna=False, keepdims=False): """ Return the product of array elements over a given axis. @@ -1988,8 +2010,12 @@ def prod(a, axis=None, dtype=None, out=None, skipna=False): the same shape as the expected output, but the type of the output values will be cast if necessary. skipna : bool, optional - If this is set to True, skips any NA values during summation + If this is set to True, skips any NA values during reduction instead of propagating them. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the original `arr`. Returns ------- @@ -2048,11 +2074,11 @@ def prod(a, axis=None, dtype=None, out=None, skipna=False): prod = a.prod except AttributeError: return um.multiply.reduce(a, axis=axis, dtype=dtype, - out=out, skipna=skipna) + out=out, skipna=skipna, keepdims=keepdims) return prod(axis=axis, dtype=dtype, out=out) else: return um.multiply.reduce(a, axis=axis, dtype=dtype, - out=out, skipna=skipna) + out=out, skipna=skipna, keepdims=keepdims) def cumprod(a, axis=None, dtype=None, out=None): """ diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c index 6a5a5620b..754d7161f 100644 --- a/numpy/core/src/multiarray/methods.c +++ b/numpy/core/src/multiarray/methods.c @@ -151,7 +151,7 @@ array_squeeze(PyArrayObject *self, PyObject *args, PyObject *kwds) return NULL; } - if (axis_in == NULL) { + if (axis_in == NULL || axis_in == Py_None) { return PyArray_Squeeze(self); } else { diff --git a/numpy/core/src/multiarray/shape.c b/numpy/core/src/multiarray/shape.c index 97b0b204a..cbbdbea75 100644 --- a/numpy/core/src/multiarray/shape.c +++ b/numpy/core/src/multiarray/shape.c @@ -716,7 +716,6 @@ NPY_NO_EXPORT PyObject * PyArray_SqueezeSelected(PyArrayObject *self, npy_bool *axis_flags) { PyArrayObject *ret; - npy_bool unit_dims[NPY_MAXDIMS]; int idim, ndim, any_ones; npy_intp *shape; @@ -726,12 +725,16 @@ PyArray_SqueezeSelected(PyArrayObject *self, npy_bool *axis_flags) /* Verify that the axes requested are all of size one */ any_ones = 0; for (idim = 0; idim < ndim; ++idim) { - if (axis_flags[idim] != 0 && shape[idim] == 1) { - unit_dims[idim] = 1; - any_ones = 1; - } - else { - unit_dims[idim] = 0; + if (axis_flags[idim] != 0) { + if (shape[idim] == 1) { + any_ones = 1; + } + else { + PyErr_SetString(PyExc_ValueError, + "cannot select an axis to squeeze out " + "which has size greater than one"); + return NULL; + } } } @@ -746,7 +749,7 @@ PyArray_SqueezeSelected(PyArrayObject *self, npy_bool *axis_flags) return NULL; } - PyArray_RemoveAxesInPlace(ret, unit_dims); + PyArray_RemoveAxesInPlace(ret, axis_flags); /* * If self isn't not a base class ndarray, call its diff --git a/numpy/core/tests/test_maskna.py b/numpy/core/tests/test_maskna.py index f946e8073..8454db4f9 100644 --- a/numpy/core/tests/test_maskna.py +++ b/numpy/core/tests/test_maskna.py @@ -693,6 +693,13 @@ def test_maskna_ufunc_1D(): # NA support assert_raises(ValueError, np.add, a, b, out=c_orig) + # Divide in-place with NA + a = np.array([[np.NA], [12.]], maskna=True) + assert_array_equal(a, [[np.NA], [666.]]) + # The assert_array_equal should have caught this... + assert_equal(np.isna(a), [[1], [0]]) + assert_equal(a[~np.isna(a)], [4.]) + def test_maskna_ufunc_sum_1D(): check_maskna_ufunc_sum_1D(np.sum) @@ -1055,9 +1062,71 @@ def test_array_maskna_squeeze(): # np.squeeze a = np.zeros((1,3,1,1,4,2,1), maskna=True) a[0,1,0,0,3,0,0] = np.NA + res = np.squeeze(a) assert_equal(res.shape, (3,4,2)) assert_(np.isna(res[1,3,0])) + res = np.squeeze(a, axis=(0,2,6)) + assert_equal(res.shape, (3,1,4,2)) + assert_(np.isna(res[1,0,3,0])) + +def test_array_maskna_mean(): + # np.mean + + # With an NA mask, but no NA + a = np.arange(6, maskna=True).reshape(2,3) + + res = np.mean(a) + assert_equal(res, 2.5) + res = np.mean(a, axis=0) + assert_equal(res, [1.5, 2.5, 3.5]) + + # With an NA and skipna=False + a = np.arange(6, maskna=True).reshape(2,3) + a[0,1] = np.NA + + res = np.mean(a) + assert_(type(res) is np.NAType) + + res = np.mean(a, axis=0) + assert_array_equal(res, [1.5, np.NA, 3.5]) + + res = np.mean(a, axis=1) + assert_array_equal(res, [np.NA, 4.0]) + + # With an NA and skipna=True + res = np.mean(a, skipna=True) + assert_almost_equal(res, 2.8) + + res = np.mean(a, axis=0, skipna=True) + assert_array_equal(res, [1.5, 4.0, 3.5]) + + res = np.mean(a, axis=1, skipna=True) + assert_array_equal(res, [1.0, 4.0]) + +def test_array_maskna_var_std(): + # np.var, np.std + + # With an NA and skipna=False + a = np.arange(6, maskna=True).reshape(2,3) + a[0,1] = np.NA + + res = np.var(a) + assert_(type(res) is np.NAType) + res = np.std(a) + assert_(type(res) is np.NAType) + + res = np.var(a, axis=0) + assert_array_equal(res, [2.25, np.NA, 2.25]) + res = np.std(a, axis=0) + assert_array_equal(res, [1.5, np.NA, 1.5]) + + res = np.var(a, axis=1) + assert_array_almost_equal(res, [np.NA, 0.66666666666666663]) + res = np.std(a, axis=1) + assert_array_almost_equal(res, [np.NA, 0.81649658092772603]) + + if __name__ == "__main__": run_module_suite() |