diff --git a/doc/release/1.17.0-notes.rst b/doc/release/1.17.0-notes.rst index 5750ef016441..c97f9a58e442 100644 --- a/doc/release/1.17.0-notes.rst +++ b/doc/release/1.17.0-notes.rst @@ -35,6 +35,15 @@ C API changes New Features ============ +``np.ufunc.reduce`` and related functions now accept a ``where`` mask +--------------------------------------------------------------------- +``np.ufunc.reduce``, ``np.sum``, ``np.prod``, ``np.min``, ``np.max`` all +now accept a ``where`` keyword argument, which can be used to tell which +elements to include in the reduction. For reductions that do not have an +identity, it is necessary to also pass in an initial value (e.g., +``initial=np.inf`` for ``np.min``). For instance, the equivalent of +``nansum`` would be, ``np.sum(a, where=~np.isnan(a))``. + Improvements ============ diff --git a/numpy/core/_add_newdocs.py b/numpy/core/_add_newdocs.py index 513415e0909b..5f5029aa06a8 100644 --- a/numpy/core/_add_newdocs.py +++ b/numpy/core/_add_newdocs.py @@ -3178,7 +3178,7 @@ add_newdoc('numpy.core.multiarray', 'ndarray', ('max', """ - a.max(axis=None, out=None, keepdims=False) + a.max(axis=None, out=None, keepdims=False, initial=, where=True) Return the maximum along a given axis. @@ -3208,7 +3208,7 @@ add_newdoc('numpy.core.multiarray', 'ndarray', ('min', """ - a.min(axis=None, out=None, keepdims=False) + a.min(axis=None, out=None, keepdims=False, initial=, where=True) Return the minimum along a given axis. @@ -3361,7 +3361,7 @@ add_newdoc('numpy.core.multiarray', 'ndarray', ('prod', """ - a.prod(axis=None, dtype=None, out=None, keepdims=False) + a.prod(axis=None, dtype=None, out=None, keepdims=False, initial=1, where=True) Return the product of the array elements over the given axis @@ -3930,7 +3930,7 @@ add_newdoc('numpy.core.multiarray', 'ndarray', ('sum', """ - a.sum(axis=None, dtype=None, out=None, keepdims=False) + a.sum(axis=None, dtype=None, out=None, keepdims=False, initial=0, where=True) Return the sum of the array elements over the given axis. @@ -4876,7 +4876,7 @@ add_newdoc('numpy.core', 'ufunc', ('reduce', """ - reduce(a, axis=0, dtype=None, out=None, keepdims=False, initial) + reduce(a, axis=0, dtype=None, out=None, keepdims=False, initial=, where=True) Reduces `a`'s dimension by one, by applying ufunc along one axis. @@ -4941,6 +4941,14 @@ .. versionadded:: 1.15.0 + where : array_like of bool, optional + A boolean array which is broadcasted to match the dimensions + of `a`, and selects elements to include in the reduction. Note + that for ufuncs like ``minimum`` that do not have an identity + defined, one has to pass in also ``initial``. + + .. versionadded:: 1.17.0 + Returns ------- r : ndarray @@ -4972,19 +4980,24 @@ array([[ 1, 5], [ 9, 13]]) - You can use the ``initial`` keyword argument to initialize the reduction with a - different value. + You can use the ``initial`` keyword argument to initialize the reduction + with a different value, and ``where`` to select specific elements to include: >>> np.add.reduce([10], initial=5) 15 >>> np.add.reduce(np.ones((2, 2, 2)), axis=(0, 2), initial=10) array([14., 14.]) + >>> a = np.array([10., np.nan, 10]) + >>> np.add.reduce(a, where=~np.isnan(a)) + 20.0 Allows reductions of empty arrays where they would normally fail, i.e. for ufuncs without an identity. >>> np.minimum.reduce([], initial=np.inf) inf + >>> np.minimum.reduce([[1., 2.], [3., 4.]], initial=10., where=[True, False]) + array([ 1., 10.]) >>> np.minimum.reduce([]) Traceback (most recent call last): ... @@ -6721,25 +6734,25 @@ add_newdoc('numpy.core.numerictypes', 'number', """ Abstract base class of all numeric scalar types. - + """) add_newdoc('numpy.core.numerictypes', 'integer', """ Abstract base class of all integer scalar types. - + """) add_newdoc('numpy.core.numerictypes', 'signedinteger', """ Abstract base class of all signed integer scalar types. - + """) add_newdoc('numpy.core.numerictypes', 'unsignedinteger', """ Abstract base class of all unsigned integer scalar types. - + """) add_newdoc('numpy.core.numerictypes', 'inexact', @@ -6747,20 +6760,20 @@ Abstract base class of all numeric scalar types with a (potentially) inexact representation of the values in its range, such as floating-point numbers. - + """) add_newdoc('numpy.core.numerictypes', 'floating', """ Abstract base class of all floating-point scalar types. - + """) add_newdoc('numpy.core.numerictypes', 'complexfloating', """ Abstract base class of all complex number scalar types that are made up of floating-point numbers. - + """) add_newdoc('numpy.core.numerictypes', 'flexible', @@ -6768,13 +6781,13 @@ Abstract base class of all scalar types without predefined length. The actual size of these types depends on the specific `np.dtype` instantiation. - + """) add_newdoc('numpy.core.numerictypes', 'character', """ Abstract base class of all character string scalar types. - + """) diff --git a/numpy/core/_methods.py b/numpy/core/_methods.py index 33f6d01a89c3..51362c7614ee 100644 --- a/numpy/core/_methods.py +++ b/numpy/core/_methods.py @@ -24,20 +24,20 @@ # avoid keyword arguments to speed up parsing, saves about 15%-20% for very # small reductions def _amax(a, axis=None, out=None, keepdims=False, - initial=_NoValue): - return umr_maximum(a, axis, None, out, keepdims, initial) + initial=_NoValue, where=True): + return umr_maximum(a, axis, None, out, keepdims, initial, where) def _amin(a, axis=None, out=None, keepdims=False, - initial=_NoValue): - return umr_minimum(a, axis, None, out, keepdims, initial) + initial=_NoValue, where=True): + return umr_minimum(a, axis, None, out, keepdims, initial, where) def _sum(a, axis=None, dtype=None, out=None, keepdims=False, - initial=_NoValue): - return umr_sum(a, axis, dtype, out, keepdims, initial) + initial=_NoValue, where=True): + return umr_sum(a, axis, dtype, out, keepdims, initial, where) def _prod(a, axis=None, dtype=None, out=None, keepdims=False, - initial=_NoValue): - return umr_prod(a, axis, dtype, out, keepdims, initial) + initial=_NoValue, where=True): + return umr_prod(a, axis, dtype, out, keepdims, initial, where) def _any(a, axis=None, dtype=None, out=None, keepdims=False): return umr_any(a, axis, dtype, out, keepdims) diff --git a/numpy/core/fromnumeric.py b/numpy/core/fromnumeric.py index 240eac6cef7c..d9437298698c 100644 --- a/numpy/core/fromnumeric.py +++ b/numpy/core/fromnumeric.py @@ -243,7 +243,7 @@ def reshape(a, newshape, order='C'): # A transpose makes the array non-contiguous >>> b = a.T - + # Taking a view makes it possible to modify the shape without modifying # the initial object. >>> c = b.view() @@ -1452,7 +1452,7 @@ def diagonal(a, offset=0, axis1=0, axis2=1): same type as `a` is returned unless `a` is a `matrix`, in which case a 1-D array rather than a (2-D) `matrix` is returned in order to maintain backward compatibility. - + If ``a.ndim > 2``, then the dimensions specified by `axis1` and `axis2` are removed, and a new axis inserted at the end corresponding to the diagonal. @@ -1963,12 +1963,13 @@ def clip(a, a_min, a_max, out=None): def _sum_dispatcher(a, axis=None, dtype=None, out=None, keepdims=None, - initial=None): + initial=None, where=None): return (a, out) @array_function_dispatch(_sum_dispatcher) -def sum(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, initial=np._NoValue): +def sum(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, + initial=np._NoValue, where=np._NoValue): """ Sum of array elements over a given axis. @@ -2012,6 +2013,11 @@ def sum(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, initial=np._No .. versionadded:: 1.15.0 + where : array_like of bool, optional + Elements to include in the sum. See `~numpy.ufunc.reduce` for details. + + .. versionadded:: 1.17.0 + Returns ------- sum_along_axis : ndarray @@ -2052,6 +2058,8 @@ def sum(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, initial=np._No array([0, 6]) >>> np.sum([[0, 1], [0, 5]], axis=1) array([1, 5]) + >>> np.sum([[0, 1], [np.nan, 5]], where=[False, True], axis=1) + array([1., 5.]) If the accumulator is too small, overflow occurs: @@ -2077,7 +2085,7 @@ def sum(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, initial=np._No return res return _wrapreduction(a, np.add, 'sum', axis, dtype, out, keepdims=keepdims, - initial=initial) + initial=initial, where=where) def _any_dispatcher(a, axis=None, out=None, keepdims=None): @@ -2394,12 +2402,14 @@ def ptp(a, axis=None, out=None, keepdims=np._NoValue): return _methods._ptp(a, axis=axis, out=out, **kwargs) -def _amax_dispatcher(a, axis=None, out=None, keepdims=None, initial=None): +def _amax_dispatcher(a, axis=None, out=None, keepdims=None, initial=None, + where=None): return (a, out) @array_function_dispatch(_amax_dispatcher) -def amax(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue): +def amax(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue, + where=np._NoValue): """ Return the maximum of an array or maximum along an axis. @@ -2437,6 +2447,11 @@ def amax(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue): .. versionadded:: 1.15.0 + where : array_like of bool, optional + Elements to compare for the maximum. See `~numpy.ufunc.reduce` + for details. + + .. versionadded:: 1.17.0 Returns ------- @@ -2482,11 +2497,14 @@ def amax(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue): array([2, 3]) >>> np.amax(a, axis=1) # Maxima along the second axis array([1, 3]) - + >>> np.amax(a, where=[False, True], initial=-1, axis=0) + array([-1, 3]) >>> b = np.arange(5, dtype=float) >>> b[2] = np.NaN >>> np.amax(b) nan + >>> np.amax(b, where=~np.isnan(b), initial=-1) + 4.0 >>> np.nanmax(b) 4.0 @@ -2505,16 +2523,18 @@ def amax(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue): >>> max([5], default=6) 5 """ - return _wrapreduction(a, np.maximum, 'max', axis, None, out, keepdims=keepdims, - initial=initial) + return _wrapreduction(a, np.maximum, 'max', axis, None, out, + keepdims=keepdims, initial=initial, where=where) -def _amin_dispatcher(a, axis=None, out=None, keepdims=None, initial=None): +def _amin_dispatcher(a, axis=None, out=None, keepdims=None, initial=None, + where=None): return (a, out) @array_function_dispatch(_amin_dispatcher) -def amin(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue): +def amin(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue, + where=np._NoValue): """ Return the minimum of an array or minimum along an axis. @@ -2552,6 +2572,12 @@ def amin(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue): .. versionadded:: 1.15.0 + where : array_like of bool, optional + Elements to compare for the minimum. See `~numpy.ufunc.reduce` + for details. + + .. versionadded:: 1.17.0 + Returns ------- amin : ndarray or scalar @@ -2596,11 +2622,15 @@ def amin(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue): array([0, 1]) >>> np.amin(a, axis=1) # Minima along the second axis array([0, 2]) + >>> np.amin(a, where=[False, True], initial=10, axis=0) + array([10, 1]) >>> b = np.arange(5, dtype=float) >>> b[2] = np.NaN >>> np.amin(b) nan + >>> np.amin(b, where=~np.isnan(b), initial=10) + 0.0 >>> np.nanmin(b) 0.0 @@ -2618,8 +2648,8 @@ def amin(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue): >>> min([6], default=5) 6 """ - return _wrapreduction(a, np.minimum, 'min', axis, None, out, keepdims=keepdims, - initial=initial) + return _wrapreduction(a, np.minimum, 'min', axis, None, out, + keepdims=keepdims, initial=initial, where=where) def _alen_dispathcer(a): @@ -2660,13 +2690,14 @@ def alen(a): return len(array(a, ndmin=1)) -def _prod_dispatcher( - a, axis=None, dtype=None, out=None, keepdims=None, initial=None): +def _prod_dispatcher(a, axis=None, dtype=None, out=None, keepdims=None, + initial=None, where=None): return (a, out) @array_function_dispatch(_prod_dispatcher) -def prod(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, initial=np._NoValue): +def prod(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, + initial=np._NoValue, where=np._NoValue): """ Return the product of array elements over a given axis. @@ -2711,6 +2742,11 @@ def prod(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, initial=np._N .. versionadded:: 1.15.0 + where : array_like of bool, optional + Elements to include in the product. See `~numpy.ufunc.reduce` for details. + + .. versionadded:: 1.17.0 + Returns ------- product_along_axis : ndarray, see `dtype` parameter above. @@ -2753,6 +2789,11 @@ def prod(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, initial=np._N >>> np.prod([[1.,2.],[3.,4.]], axis=1) array([ 2., 12.]) + Or select specific elements to include: + + >>> np.prod([1., np.nan, 3.], where=[True, False, True]) + 3.0 + If the type of `x` is unsigned, then the output type is the unsigned platform integer: @@ -2772,8 +2813,8 @@ def prod(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, initial=np._N >>> np.prod([1, 2], initial=5) 10 """ - return _wrapreduction(a, np.multiply, 'prod', axis, dtype, out, keepdims=keepdims, - initial=initial) + return _wrapreduction(a, np.multiply, 'prod', axis, dtype, out, + keepdims=keepdims, initial=initial, where=where) def _cumprod_dispatcher(a, axis=None, dtype=None, out=None): diff --git a/numpy/core/src/umath/override.c b/numpy/core/src/umath/override.c index 2ea23311b155..8d67f96ac2d1 100644 --- a/numpy/core/src/umath/override.c +++ b/numpy/core/src/umath/override.c @@ -226,14 +226,14 @@ normalize_reduce_args(PyUFuncObject *ufunc, PyObject *args, PyObject *obj; static PyObject *NoValue = NULL; static char *kwlist[] = {"array", "axis", "dtype", "out", "keepdims", - "initial"}; + "initial", "where"}; npy_cache_import("numpy", "_NoValue", &NoValue); if (NoValue == NULL) return -1; - if (nargs < 1 || nargs > 6) { + if (nargs < 1 || nargs > 7) { PyErr_Format(PyExc_TypeError, - "ufunc.reduce() takes from 1 to 6 positional " + "ufunc.reduce() takes from 1 to 7 positional " "arguments but %"NPY_INTP_FMT" were given", nargs); return -1; } diff --git a/numpy/core/src/umath/reduction.c b/numpy/core/src/umath/reduction.c index 791d3693f3c7..4174e69a8ae1 100644 --- a/numpy/core/src/umath/reduction.c +++ b/numpy/core/src/umath/reduction.c @@ -443,9 +443,9 @@ PyUFunc_ReduceWrapper(PyArrayObject *operand, PyArrayObject *out, /* Iterator parameters */ NpyIter *iter = NULL; - PyArrayObject *op[2]; - PyArray_Descr *op_dtypes[2]; - npy_uint32 flags, op_flags[2]; + PyArrayObject *op[3]; + PyArray_Descr *op_dtypes[3]; + npy_uint32 flags, op_flags[3]; /* More than one axis means multiple orders are possible */ if (!reorderable && count_axes(PyArray_NDIM(operand), axis_flags) > 1) { @@ -455,13 +455,12 @@ PyUFunc_ReduceWrapper(PyArrayObject *operand, PyArrayObject *out, funcname); return NULL; } - - - /* Validate that the parameters for future expansion are NULL */ - if (wheremask != NULL) { - PyErr_SetString(PyExc_RuntimeError, - "Reduce operations in NumPy do not yet support " - "a where mask"); + /* Can only use where with an initial ( from identity or argument) */ + if (wheremask != NULL && identity == Py_None) { + PyErr_Format(PyExc_ValueError, + "reduction operation '%s' does not have an identity, " + "so to use a where mask one has to specify 'initial'", + funcname); return NULL; } @@ -523,8 +522,16 @@ PyUFunc_ReduceWrapper(PyArrayObject *operand, PyArrayObject *out, NPY_ITER_NO_SUBTYPE; op_flags[1] = NPY_ITER_READONLY | NPY_ITER_ALIGNED; + if (wheremask != NULL) { + op[2] = wheremask; + op_dtypes[2] = PyArray_DescrFromType(NPY_BOOL); + if (op_dtypes[2] == NULL) { + goto fail; + } + op_flags[2] = NPY_ITER_READONLY; + } - iter = NpyIter_AdvancedNew(2, op, flags, + iter = NpyIter_AdvancedNew(wheremask == NULL ? 2 : 3, op, flags, NPY_KEEPORDER, casting, op_flags, op_dtypes, @@ -567,7 +574,7 @@ PyUFunc_ReduceWrapper(PyArrayObject *operand, PyArrayObject *out, goto fail; } } - + /* Check whether any errors occurred during the loop */ if (PyErr_Occurred() || _check_ufunc_fperr(errormask, NULL, "reduce") < 0) { diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c index ba7a11fa3b70..f950915f5d1b 100644 --- a/numpy/core/src/umath/ufunc_object.c +++ b/numpy/core/src/umath/ufunc_object.c @@ -3468,12 +3468,15 @@ reduce_loop(NpyIter *iter, char **dataptrs, npy_intp *strides, PyUFuncObject *ufunc = (PyUFuncObject *)data; char *dataptrs_copy[3]; npy_intp strides_copy[3]; + npy_bool masked; /* The normal selected inner loop */ PyUFuncGenericFunction innerloop = NULL; void *innerloopdata = NULL; NPY_BEGIN_THREADS_DEF; + /* Get the number of operands, to determine whether "where" is used */ + masked = (NpyIter_GetNOp(iter) == 3); /* Get the inner loop */ iter_dtypes = NpyIter_GetDescrArray(iter); @@ -3533,8 +3536,36 @@ reduce_loop(NpyIter *iter, char **dataptrs, npy_intp *strides, strides_copy[0] = strides[0]; strides_copy[1] = strides[1]; strides_copy[2] = strides[0]; - innerloop(dataptrs_copy, countptr, - strides_copy, innerloopdata); + + if (!masked) { + innerloop(dataptrs_copy, countptr, + strides_copy, innerloopdata); + } + else { + npy_intp count = *countptr; + char *maskptr = dataptrs[2]; + npy_intp mask_stride = strides[2]; + /* Optimization for when the mask is broadcast */ + npy_intp n = mask_stride == 0 ? count : 1; + while (count) { + char mask = *maskptr; + maskptr += mask_stride; + while (n < count && mask == *maskptr) { + n++; + maskptr += mask_stride; + } + /* If mask set, apply inner loop on this contiguous region */ + if (mask) { + innerloop(dataptrs_copy, &n, + strides_copy, innerloopdata); + } + dataptrs_copy[0] += n * strides[0]; + dataptrs_copy[1] += n * strides[1]; + dataptrs_copy[2] = dataptrs_copy[0]; + count -= n; + n = 1; + } + } } while (iternext(iter)); finish_loop: @@ -3563,7 +3594,7 @@ reduce_loop(NpyIter *iter, char **dataptrs, npy_intp *strides, static PyArrayObject * PyUFunc_Reduce(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *out, int naxes, int *axes, PyArray_Descr *odtype, int keepdims, - PyObject *initial) + PyObject *initial, PyArrayObject *wheremask) { int iaxes, ndim; npy_bool reorderable; @@ -3629,7 +3660,7 @@ PyUFunc_Reduce(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *out, return NULL; } - result = PyUFunc_ReduceWrapper(arr, out, NULL, dtype, dtype, + result = PyUFunc_ReduceWrapper(arr, out, wheremask, dtype, dtype, NPY_UNSAFE_CASTING, axis_flags, reorderable, keepdims, 0, @@ -4386,7 +4417,7 @@ PyUFunc_GenericReduction(PyUFuncObject *ufunc, PyObject *args, int i, naxes=0, ndim; int axes[NPY_MAXDIMS]; PyObject *axes_in = NULL; - PyArrayObject *mp = NULL, *ret = NULL; + PyArrayObject *mp = NULL, *wheremask = NULL, *ret = NULL; PyObject *op; PyObject *obj_ind, *context; PyArrayObject *indices = NULL; @@ -4395,7 +4426,7 @@ PyUFunc_GenericReduction(PyUFuncObject *ufunc, PyObject *args, int keepdims = 0; PyObject *initial = NULL; static char *reduce_kwlist[] = { - "array", "axis", "dtype", "out", "keepdims", "initial", NULL}; + "array", "axis", "dtype", "out", "keepdims", "initial", "where", NULL}; static char *accumulate_kwlist[] = { "array", "axis", "dtype", "out", NULL}; static char *reduceat_kwlist[] = { @@ -4458,22 +4489,23 @@ PyUFunc_GenericReduction(PyUFuncObject *ufunc, PyObject *args, } else if (operation == UFUNC_ACCUMULATE) { if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO&O&:accumulate", - accumulate_kwlist, - &op, - &axes_in, - PyArray_DescrConverter2, &otype, - PyArray_OutputConverter, &out)) { + accumulate_kwlist, + &op, + &axes_in, + PyArray_DescrConverter2, &otype, + PyArray_OutputConverter, &out)) { goto fail; } } else { - if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO&O&iO:reduce", - reduce_kwlist, - &op, - &axes_in, - PyArray_DescrConverter2, &otype, - PyArray_OutputConverter, &out, - &keepdims, &initial)) { + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO&O&iOO&:reduce", + reduce_kwlist, + &op, + &axes_in, + PyArray_DescrConverter2, &otype, + PyArray_OutputConverter, &out, + &keepdims, &initial, + _wheremask_converter, &wheremask)) { goto fail; } } @@ -4604,7 +4636,8 @@ PyUFunc_GenericReduction(PyUFuncObject *ufunc, PyObject *args, switch(operation) { case UFUNC_REDUCE: ret = PyUFunc_Reduce(ufunc, mp, out, naxes, axes, - otype, keepdims, initial); + otype, keepdims, initial, wheremask); + Py_XDECREF(wheremask); break; case UFUNC_ACCUMULATE: if (naxes != 1) { @@ -4662,6 +4695,7 @@ PyUFunc_GenericReduction(PyUFuncObject *ufunc, PyObject *args, fail: Py_XDECREF(otype); Py_XDECREF(mp); + Py_XDECREF(wheremask); return NULL; } diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py index b83b8ccffafc..b521522431eb 100644 --- a/numpy/core/tests/test_ufunc.py +++ b/numpy/core/tests/test_ufunc.py @@ -3,6 +3,8 @@ import warnings import itertools +import pytest + import numpy as np import numpy.core._umath_tests as umt import numpy.linalg._umath_linalg as uml @@ -596,6 +598,12 @@ def test_sum_initial(self): assert_equal(np.sum(np.ones((2, 3, 5), dtype=np.int64), axis=(0, 2), initial=2), [12, 12, 12]) + def test_sum_where(self): + # More extensive tests done in test_reduction_with_where. + assert_equal(np.sum([[1., 2.], [3., 4.]], where=[True, False]), 4.) + assert_equal(np.sum([[1., 2.], [3., 4.]], axis=0, initial=5., + where=[True, False]), [9., 5.]) + def test_inner1d(self): a = np.arange(6).reshape((2, 3)) assert_array_equal(umt.inner1d(a, a), np.sum(a*a, axis=-1)) @@ -1162,6 +1170,8 @@ def test_object_array_reduction(self): assert_equal(np.array([[1]], dtype=object).sum(), 1) assert_equal(np.array([[[1, 2]]], dtype=object).sum((0, 1)), [1, 2]) assert_equal(np.array([1], dtype=object).sum(initial=1), 2) + assert_equal(np.array([[1], [2, 3]], dtype=object) + .sum(initial=[0], where=[False, True]), [0, 2, 3]) def test_object_array_accumulate_inplace(self): # Checks that in-place accumulates work, see also gh-7402 @@ -1396,6 +1406,44 @@ def test_initial_reduction(self): res = np.add.reduce(a, initial=5) assert_equal(res, 15) + @pytest.mark.parametrize('axis', (0, 1, None)) + @pytest.mark.parametrize('where', (np.array([False, True, True]), + np.array([[True], [False], [True]]), + np.array([[True, False, False], + [False, True, False], + [False, True, True]]))) + def test_reduction_with_where(self, axis, where): + a = np.arange(9.).reshape(3, 3) + a_copy = a.copy() + a_check = np.zeros_like(a) + np.positive(a, out=a_check, where=where) + + res = np.add.reduce(a, axis=axis, where=where) + check = a_check.sum(axis) + assert_equal(res, check) + # Check we do not overwrite elements of a internally. + assert_array_equal(a, a_copy) + + @pytest.mark.parametrize(('axis', 'where'), + ((0, np.array([True, False, True])), + (1, [True, True, False]), + (None, True))) + @pytest.mark.parametrize('initial', (-np.inf, 5.)) + def test_reduction_with_where_and_initial(self, axis, where, initial): + a = np.arange(9.).reshape(3, 3) + a_copy = a.copy() + a_check = np.full(a.shape, -np.inf) + np.positive(a, out=a_check, where=where) + + res = np.maximum.reduce(a, axis=axis, where=where, initial=initial) + check = a_check.max(axis, initial=initial) + assert_equal(res, check) + + def test_reduction_where_initial_needed(self): + a = np.arange(9.).reshape(3, 3) + m = [False, True, False] + assert_raises(ValueError, np.maximum.reduce, a, where=m) + def test_identityless_reduction_nonreorderable(self): a = np.array([[8.0, 2.0, 2.0], [1.0, 0.5, 0.25]]) @@ -1749,16 +1797,19 @@ def test_reduce_arguments(self): assert_equal(f(d, 0, None, None, True), r.reshape((1,) + r.shape)) assert_equal(f(d, 0, None, None, False, 0), r) assert_equal(f(d, 0, None, None, False, initial=0), r) + assert_equal(f(d, 0, None, None, False, 0, True), r) + assert_equal(f(d, 0, None, None, False, 0, where=True), r) # multiple keywords assert_equal(f(d, axis=0, dtype=None, out=None, keepdims=False), r) assert_equal(f(d, 0, dtype=None, out=None, keepdims=False), r) assert_equal(f(d, 0, None, out=None, keepdims=False), r) - assert_equal(f(d, 0, None, out=None, keepdims=False, initial=0), r) + assert_equal(f(d, 0, None, out=None, keepdims=False, initial=0, + where=True), r) # too little assert_raises(TypeError, f) # too much - assert_raises(TypeError, f, d, 0, None, None, False, 0, 1) + assert_raises(TypeError, f, d, 0, None, None, False, 0, True, 1) # invalid axis assert_raises(TypeError, f, d, "invalid") assert_raises(TypeError, f, d, axis="invalid") diff --git a/numpy/core/tests/test_umath.py b/numpy/core/tests/test_umath.py index 2f8edebc0f9c..21097244ff6b 100644 --- a/numpy/core/tests/test_umath.py +++ b/numpy/core/tests/test_umath.py @@ -1894,7 +1894,8 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): # reduce, kwargs res = np.multiply.reduce(a, axis='axis0', dtype='dtype0', out='out0', - keepdims='keep0', initial='init0') + keepdims='keep0', initial='init0', + where='where0') assert_equal(res[0], a) assert_equal(res[1], np.multiply) assert_equal(res[2], 'reduce') @@ -1903,7 +1904,8 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): 'out': ('out0',), 'keepdims': 'keep0', 'axis': 'axis0', - 'initial': 'init0'}) + 'initial': 'init0', + 'where': 'where0'}) # reduce, output equal to None removed, but not other explicit ones, # even if they are at their default value. @@ -1913,14 +1915,18 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): assert_equal(res[4], {'axis': 0, 'keepdims': True}) res = np.multiply.reduce(a, None, out=(None,), dtype=None) assert_equal(res[4], {'axis': None, 'dtype': None}) - res = np.multiply.reduce(a, 0, None, None, False, 2) - assert_equal(res[4], {'axis': 0, 'dtype': None, 'keepdims': False, 'initial': 2}) - # np._NoValue ignored for initial. - res = np.multiply.reduce(a, 0, None, None, False, np._NoValue) - assert_equal(res[4], {'axis': 0, 'dtype': None, 'keepdims': False}) - # None kept for initial. - res = np.multiply.reduce(a, 0, None, None, False, None) - assert_equal(res[4], {'axis': 0, 'dtype': None, 'keepdims': False, 'initial': None}) + res = np.multiply.reduce(a, 0, None, None, False, 2, True) + assert_equal(res[4], {'axis': 0, 'dtype': None, 'keepdims': False, + 'initial': 2, 'where': True}) + # np._NoValue ignored for initial + res = np.multiply.reduce(a, 0, None, None, False, + np._NoValue, True) + assert_equal(res[4], {'axis': 0, 'dtype': None, 'keepdims': False, + 'where': True}) + # None kept for initial, True for where. + res = np.multiply.reduce(a, 0, None, None, False, None, True) + assert_equal(res[4], {'axis': 0, 'dtype': None, 'keepdims': False, + 'initial': None, 'where': True}) # reduce, wrong args assert_raises(ValueError, np.multiply.reduce, a, out=())