summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
Diffstat
-rw-r--r--doc/release/1.17.0-notes.rst9
-rw-r--r--numpy/core/_add_newdocs.py45
-rw-r--r--numpy/core/_methods.py16
-rw-r--r--numpy/core/fromnumeric.py79
-rw-r--r--numpy/core/src/umath/override.c6
-rw-r--r--numpy/core/src/umath/reduction.c31
-rw-r--r--numpy/core/src/umath/ufunc_object.c72
-rw-r--r--numpy/core/tests/test_ufunc.py55
-rw-r--r--numpy/core/tests/test_umath.py26
9 files changed, 250 insertions, 89 deletions
diff --git a/doc/release/1.17.0-notes.rst b/doc/release/1.17.0-notes.rst
index 5750ef016..c97f9a58e 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 513415e09..5f5029aa0 100644
--- a/numpy/core/_add_newdocs.py
+++ b/numpy/core/_add_newdocs.py
@@ -3178,7 +3178,7 @@ add_newdoc('numpy.core.multiarray', 'ndarray', ('itemset',
add_newdoc('numpy.core.multiarray', 'ndarray', ('max',
"""
- a.max(axis=None, out=None, keepdims=False)
+ a.max(axis=None, out=None, keepdims=False, initial=<no value>, where=True)
Return the maximum along a given axis.
@@ -3208,7 +3208,7 @@ add_newdoc('numpy.core.multiarray', 'ndarray', ('mean',
add_newdoc('numpy.core.multiarray', 'ndarray', ('min',
"""
- a.min(axis=None, out=None, keepdims=False)
+ a.min(axis=None, out=None, keepdims=False, initial=<no value>, where=True)
Return the minimum along a given axis.
@@ -3361,7 +3361,7 @@ add_newdoc('numpy.core.multiarray', 'ndarray', ('nonzero',
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', ('std',
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', ('signature',
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=<no value>, where=True)
Reduces `a`'s dimension by one, by applying ufunc along one axis.
@@ -4941,6 +4941,14 @@ add_newdoc('numpy.core', 'ufunc', ('reduce',
.. 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 @@ add_newdoc('numpy.core', 'ufunc', ('reduce',
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', 'generic', ('view',
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 @@ add_newdoc('numpy.core.numerictypes', 'inexact',
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 @@ add_newdoc('numpy.core.numerictypes', 'flexible',
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 33f6d01a8..51362c761 100644
--- a/numpy/core/_methods.py
+++ b/numpy/core/_methods.py
@@ -24,20 +24,20 @@ umr_all = um.logical_and.reduce
# 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 240eac6ce..d94372986 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 2ea23311b..8d67f96ac 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 791d3693f..4174e69a8 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 ba7a11fa3..f950915f5 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 @@ finish_loop:
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 b83b8ccff..b52152243 100644
--- a/numpy/core/tests/test_ufunc.py
+++ b/numpy/core/tests/test_ufunc.py
@@ -3,6 +3,8 @@ from __future__ import division, absolute_import, print_function
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 @@ class TestUfunc(object):
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 @@ class TestUfunc(object):
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 @@ class TestUfunc(object):
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 @@ class TestUfunc(object):
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 2f8edebc0..21097244f 100644
--- a/numpy/core/tests/test_umath.py
+++ b/numpy/core/tests/test_umath.py
@@ -1894,7 +1894,8 @@ class TestSpecialMethods(object):
# 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 @@ class TestSpecialMethods(object):
'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 @@ class TestSpecialMethods(object):
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=())
View Lorry Controller Status