Merge pull request #10635 from hameerabbasi/ufunc-reduce-identity

ENH: Implement initial kwarg for ufunc.add.reduce

8 files changed, 214 insertions, 41 deletions

diff --git a/doc/release/1.15.0-notes.rst b/doc/release/1.15.0-notes.rst
index 49e8ab22d..a6b23b892 100644
--- a/doc/release/1.15.0-notes.rst
+++ b/doc/release/1.15.0-notes.rst
@@ -158,6 +158,12 @@ Added experimental support for the 64-bit RISC-V architecture.
 Improvements
 ============
 
+``np.ufunc.reduce`` and related functions now accept an initial value
+---------------------------------------------------------------------
+``np.ufunc.reduce``, ``np.sum``, ``np.prod``, ``np.min`` and ``np.max`` all
+now accept an ``initial`` keyword argument that specifies the value to start
+the reduction with.
+
 ``np.flip`` can operate over multiple axes
 ------------------------------------------
 ``np.flip`` now accepts None, or tuples of int, in its ``axis`` argument. If
diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py
index 93a521658..c187f8e31 100644
--- a/numpy/add_newdocs.py
+++ b/numpy/add_newdocs.py
@@ -5850,7 +5850,7 @@ add_newdoc('numpy.core', 'ufunc', ('signature',
 
 add_newdoc('numpy.core', 'ufunc', ('reduce',
     """
-    reduce(a, axis=0, dtype=None, out=None, keepdims=False)
+    reduce(a, axis=0, dtype=None, out=None, keepdims=False, initial)
 
     Reduces `a`'s dimension by one, by applying ufunc along one axis.
 
@@ -5906,6 +5906,14 @@ add_newdoc('numpy.core', 'ufunc', ('reduce',
         the result will broadcast correctly against the original `arr`.
 
         .. versionadded:: 1.7.0
+    initial : scalar, optional
+        The value with which to start the reduction.
+        If the ufunc has no identity or the dtype is object, this defaults
+        to None - otherwise it defaults to ufunc.identity.
+        If ``None`` is given, the first element of the reduction is used,
+        and an error is thrown if the reduction is empty.
+        
+        .. versionadded:: 1.15.0
 
     Returns
     -------
@@ -5937,7 +5945,24 @@ add_newdoc('numpy.core', 'ufunc', ('reduce',
     >>> np.add.reduce(X, 2)
     array([[ 1,  5],
            [ 9, 13]])
-
+           
+    You can use the ``initial`` keyword argument to initialize the reduction with a
+    different value.
+    
+    >>> np.add.reduce([10], initial=5)
+    15
+    >>> np.add.reduce(np.ones((2, 2, 2)), axis=(0, 2), initializer=10)
+    array([14., 14.])
+    
+    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([])
+    Traceback (most recent call last):
+        ...
+    ValueError: zero-size array to reduction operation minimum which has no identity
     """))
 
 add_newdoc('numpy.core', 'ufunc', ('accumulate',
diff --git a/numpy/core/_methods.py b/numpy/core/_methods.py
index 0f928676b..33f6d01a8 100644
--- a/numpy/core/_methods.py
+++ b/numpy/core/_methods.py
@@ -11,6 +11,7 @@ from numpy.core import multiarray as mu
 from numpy.core import umath as um
 from numpy.core.numeric import asanyarray
 from numpy.core import numerictypes as nt
+from numpy._globals import _NoValue
 
 # save those O(100) nanoseconds!
 umr_maximum = um.maximum.reduce
@@ -22,17 +23,21 @@ 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):
-    return umr_maximum(a, axis, None, out, keepdims)
+def _amax(a, axis=None, out=None, keepdims=False,
+          initial=_NoValue):
+    return umr_maximum(a, axis, None, out, keepdims, initial)
 
-def _amin(a, axis=None, out=None, keepdims=False):
-    return umr_minimum(a, axis, None, out, keepdims)
+def _amin(a, axis=None, out=None, keepdims=False,
+          initial=_NoValue):
+    return umr_minimum(a, axis, None, out, keepdims, initial)
 
-def _sum(a, axis=None, dtype=None, out=None, keepdims=False):
-    return umr_sum(a, axis, dtype, out, keepdims)
+def _sum(a, axis=None, dtype=None, out=None, keepdims=False,
+         initial=_NoValue):
+    return umr_sum(a, axis, dtype, out, keepdims, initial)
 
-def _prod(a, axis=None, dtype=None, out=None, keepdims=False):
-    return umr_prod(a, axis, dtype, out, keepdims)
+def _prod(a, axis=None, dtype=None, out=None, keepdims=False,
+          initial=_NoValue):
+    return umr_prod(a, axis, dtype, out, keepdims, initial)
 
 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 948c2139d..75bcedd81 100644
--- a/numpy/core/fromnumeric.py
+++ b/numpy/core/fromnumeric.py
@@ -1812,7 +1812,7 @@ def clip(a, a_min, a_max, out=None):
     return _wrapfunc(a, 'clip', a_min, a_max, out=out)
 
 
-def sum(a, axis=None, dtype=None, out=None, keepdims=np._NoValue):
+def sum(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, initial=np._NoValue):
     """
     Sum of array elements over a given axis.
 
@@ -1851,6 +1851,10 @@ def sum(a, axis=None, dtype=None, out=None, keepdims=np._NoValue):
         `ndarray`, however any non-default value will be.  If the
         sub-class' method does not implement `keepdims` any
         exceptions will be raised.
+    initial : scalar, optional
+        Starting value for the sum. See `~numpy.ufunc.reduce` for details.
+
+        .. versionadded:: 1.15.0
 
     Returns
     -------
@@ -1898,6 +1902,10 @@ def sum(a, axis=None, dtype=None, out=None, keepdims=np._NoValue):
     >>> np.ones(128, dtype=np.int8).sum(dtype=np.int8)
     -128
 
+    You can also start the sum with a value other than zero:
+
+    >>> np.sum([10], initial=5)
+    15
     """
     if isinstance(a, _gentype):
         # 2018-02-25, 1.15.0
@@ -1912,7 +1920,8 @@ def sum(a, axis=None, dtype=None, out=None, keepdims=np._NoValue):
             return out
         return res
 
-    return _wrapreduction(a, np.add, 'sum', axis, dtype, out, keepdims=keepdims)
+    return _wrapreduction(a, np.add, 'sum', axis, dtype, out, keepdims=keepdims,
+                          initial=initial)
 
 
 def any(a, axis=None, out=None, keepdims=np._NoValue):
@@ -2209,7 +2218,7 @@ def ptp(a, axis=None, out=None, keepdims=np._NoValue):
     return _methods._ptp(a, axis=axis, out=out, **kwargs)
 
 
-def amax(a, axis=None, out=None, keepdims=np._NoValue):
+def amax(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue):
     """
     Return the maximum of an array or maximum along an axis.
 
@@ -2241,6 +2250,13 @@ def amax(a, axis=None, out=None, keepdims=np._NoValue):
         sub-class' method does not implement `keepdims` any
         exceptions will be raised.
 
+    initial : scalar, optional
+        The minimum value of an output element. Must be present to allow
+        computation on empty slice. See `~numpy.ufunc.reduce` for details.
+
+        .. versionadded:: 1.15.0
+
+
     Returns
     -------
     amax : ndarray or scalar
@@ -2293,11 +2309,26 @@ def amax(a, axis=None, out=None, keepdims=np._NoValue):
     >>> np.nanmax(b)
     4.0
 
+    You can use an initial value to compute the maximum of an empty slice, or
+    to initialize it to a different value:
+
+    >>> np.max([[-50], [10]], axis=-1, initial=0)
+    array([ 0, 10])
+
+    Notice that the initial value is used as one of the elements for which the
+    maximum is determined, unlike for the default argument Python's max
+    function, which is only used for empty iterables.
+
+    >>> np.max([5], initial=6)
+    6
+    >>> max([5], default=6)
+    5
     """
-    return _wrapreduction(a, np.maximum, 'max', axis, None, out, keepdims=keepdims)
+    return _wrapreduction(a, np.maximum, 'max', axis, None, out, keepdims=keepdims,
+                          initial=initial)
 
 
-def amin(a, axis=None, out=None, keepdims=np._NoValue):
+def amin(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue):
     """
     Return the minimum of an array or minimum along an axis.
 
@@ -2329,6 +2360,12 @@ def amin(a, axis=None, out=None, keepdims=np._NoValue):
         sub-class' method does not implement `keepdims` any
         exceptions will be raised.
 
+    initial : scalar, optional
+        The maximum value of an output element. Must be present to allow
+        computation on empty slice. See `~numpy.ufunc.reduce` for details.
+
+        .. versionadded:: 1.15.0
+
     Returns
     -------
     amin : ndarray or scalar
@@ -2381,8 +2418,22 @@ def amin(a, axis=None, out=None, keepdims=np._NoValue):
     >>> np.nanmin(b)
     0.0
 
+    >>> np.min([[-50], [10]], axis=-1, initial=0)
+    array([-50,   0])
+
+    Notice that the initial value is used as one of the elements for which the
+    minimum is determined, unlike for the default argument Python's max
+    function, which is only used for empty iterables.
+
+    Notice that this isn't the same as Python's ``default`` argument.
+
+    >>> np.min([6], initial=5)
+    5
+    >>> min([6], default=5)
+    6
     """
-    return _wrapreduction(a, np.minimum, 'min', axis, None, out, keepdims=keepdims)
+    return _wrapreduction(a, np.minimum, 'min', axis, None, out, keepdims=keepdims,
+                          initial=initial)
 
 
 def alen(a):
@@ -2418,7 +2469,7 @@ def alen(a):
         return len(array(a, ndmin=1))
 
 
-def prod(a, axis=None, dtype=None, out=None, keepdims=np._NoValue):
+def prod(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, initial=np._NoValue):
     """
     Return the product of array elements over a given axis.
 
@@ -2458,6 +2509,10 @@ def prod(a, axis=None, dtype=None, out=None, keepdims=np._NoValue):
         `ndarray`, however any non-default value will be.  If the
         sub-class' method does not implement `keepdims` any
         exceptions will be raised.
+    initial : scalar, optional
+        The starting value for this product. See `~numpy.ufunc.reduce` for details.
+
+        .. versionadded:: 1.15.0
 
     Returns
     -------
@@ -2515,8 +2570,13 @@ def prod(a, axis=None, dtype=None, out=None, keepdims=np._NoValue):
     >>> np.prod(x).dtype == int
     True
 
+    You can also start the product with a value other than one:
+
+    >>> np.prod([1, 2], initial=5)
+    10
     """
-    return _wrapreduction(a, np.multiply, 'prod', axis, dtype, out, keepdims=keepdims)
+    return _wrapreduction(a, np.multiply, 'prod', axis, dtype, out, keepdims=keepdims,
+                          initial=initial)
 
 
 def cumprod(a, axis=None, dtype=None, out=None):
diff --git a/numpy/core/src/umath/override.c b/numpy/core/src/umath/override.c
index 0aef093b0..123d9af87 100644
--- a/numpy/core/src/umath/override.c
+++ b/numpy/core/src/umath/override.c
@@ -123,11 +123,16 @@ normalize_reduce_args(PyUFuncObject *ufunc, PyObject *args,
     npy_intp nargs = PyTuple_GET_SIZE(args);
     npy_intp i;
     PyObject *obj;
-    static char *kwlist[] = {"array", "axis", "dtype", "out", "keepdims"};
+    static PyObject *NoValue = NULL;
+    static char *kwlist[] = {"array", "axis", "dtype", "out", "keepdims",
+        "initial"};
+
+    npy_cache_import("numpy", "_NoValue", &NoValue);
+    if (NoValue == NULL) return -1;
 
-    if (nargs < 1 || nargs > 5) {
+    if (nargs < 1 || nargs > 6) {
         PyErr_Format(PyExc_TypeError,
-                     "ufunc.reduce() takes from 1 to 5 positional "
+                     "ufunc.reduce() takes from 1 to 6 positional "
                      "arguments but %"NPY_INTP_FMT" were given", nargs);
         return -1;
     }
@@ -151,6 +156,10 @@ normalize_reduce_args(PyUFuncObject *ufunc, PyObject *args,
             }
             obj = PyTuple_GetSlice(args, 3, 4);
         }
+        /* Remove initial=np._NoValue */
+        if (i == 5 && obj == NoValue) {
+            continue;
+        }
         PyDict_SetItemString(*normal_kwds, kwlist[i], obj);
         if (i == 3) {
             Py_DECREF(obj);
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index e0423630b..f7f35fab9 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -3019,20 +3019,25 @@ finish_loop:
  */
 static PyArrayObject *
 PyUFunc_Reduce(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *out,
-        int naxes, int *axes, PyArray_Descr *odtype, int keepdims)
+        int naxes, int *axes, PyArray_Descr *odtype, int keepdims,
+        PyObject *initial)
 {
     int iaxes, ndim;
     npy_bool reorderable;
     npy_bool axis_flags[NPY_MAXDIMS];
     PyArray_Descr *dtype;
     PyArrayObject *result;
-    PyObject *identity = NULL;
+    PyObject *identity;
     const char *ufunc_name = ufunc_get_name_cstr(ufunc);
     /* These parameters come from a TLS global */
     int buffersize = 0, errormask = 0;
+    static PyObject *NoValue = NULL;
 
     NPY_UF_DBG_PRINT1("\nEvaluating ufunc %s.reduce\n", ufunc_name);
 
+    npy_cache_import("numpy", "_NoValue", &NoValue);
+    if (NoValue == NULL) return NULL;
+
     ndim = PyArray_NDIM(arr);
 
     /* Create an array of flags for reduction */
@@ -3056,19 +3061,28 @@ PyUFunc_Reduce(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *out,
     if (identity == NULL) {
         return NULL;
     }
-    /*
-     * The identity for a dynamic dtype like
-     * object arrays can't be used in general
-     */
-    if (identity != Py_None && PyArray_ISOBJECT(arr) && PyArray_SIZE(arr) != 0) {
+
+    /* Get the initial value */
+    if (initial == NULL || initial == NoValue) {
+        initial = identity;
+
+        /*
+        * The identity for a dynamic dtype like
+        * object arrays can't be used in general
+        */
+        if (initial != Py_None && PyArray_ISOBJECT(arr) && PyArray_SIZE(arr) != 0) {
+            Py_DECREF(initial);
+            initial = Py_None;
+            Py_INCREF(initial);
+        }
+    } else {
         Py_DECREF(identity);
-        identity = Py_None;
-        Py_INCREF(identity);
+        Py_INCREF(initial);  /* match the reference count in the if above */
     }
 
     /* Get the reduction dtype */
     if (reduce_type_resolver(ufunc, arr, odtype, &dtype) < 0) {
-        Py_DECREF(identity);
+        Py_DECREF(initial);
         return NULL;
     }
 
@@ -3076,12 +3090,12 @@ PyUFunc_Reduce(PyUFuncObject *ufunc, PyArrayObject *arr, PyArrayObject *out,
                                    NPY_UNSAFE_CASTING,
                                    axis_flags, reorderable,
                                    keepdims, 0,
-                                   identity,
+                                   initial,
                                    reduce_loop,
                                    ufunc, buffersize, ufunc_name, errormask);
 
     Py_DECREF(dtype);
-    Py_DECREF(identity);
+    Py_DECREF(initial);
     return result;
 }
 
@@ -3845,8 +3859,9 @@ PyUFunc_GenericReduction(PyUFuncObject *ufunc, PyObject *args,
     PyArray_Descr *otype = NULL;
     PyArrayObject *out = NULL;
     int keepdims = 0;
+    PyObject *initial = NULL;
     static char *reduce_kwlist[] = {
-            "array", "axis", "dtype", "out", "keepdims", NULL};
+            "array", "axis", "dtype", "out", "keepdims", "initial", NULL};
     static char *accumulate_kwlist[] = {
             "array", "axis", "dtype", "out", NULL};
     static char *reduceat_kwlist[] = {
@@ -3918,13 +3933,13 @@ PyUFunc_GenericReduction(PyUFuncObject *ufunc, PyObject *args,
         }
     }
     else {
-        if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO&O&i:reduce",
+        if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO&O&iO:reduce",
                                         reduce_kwlist,
                                         &op,
                                         &axes_in,
                                         PyArray_DescrConverter2, &otype,
                                         PyArray_OutputConverter, &out,
-                                        &keepdims)) {
+                                        &keepdims, &initial)) {
             goto fail;
         }
     }
@@ -4055,7 +4070,7 @@ PyUFunc_GenericReduction(PyUFuncObject *ufunc, PyObject *args,
     switch(operation) {
     case UFUNC_REDUCE:
         ret = PyUFunc_Reduce(ufunc, mp, out, naxes, axes,
-                                          otype, keepdims);
+                                          otype, keepdims, initial);
         break;
     case UFUNC_ACCUMULATE:
         if (naxes != 1) {
diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py
index 7a276c04d..fe40456d5 100644
--- a/numpy/core/tests/test_ufunc.py
+++ b/numpy/core/tests/test_ufunc.py
@@ -494,6 +494,17 @@ class TestUfunc(object):
             d += d
             assert_almost_equal(d, 2. + 2j)
 
+    def test_sum_initial(self):
+        # Integer, single axis
+        assert_equal(np.sum([3], initial=2), 5)
+
+        # Floating point
+        assert_almost_equal(np.sum([0.2], initial=0.1), 0.3)
+
+        # Multiple non-adjacent axes
+        assert_equal(np.sum(np.ones((2, 3, 5), dtype=np.int64), axis=(0, 2), initial=2),
+                     [12, 12, 12])
+
     def test_inner1d(self):
         a = np.arange(6).reshape((2, 3))
         assert_array_equal(umt.inner1d(a, a), np.sum(a*a, axis=-1))
@@ -844,6 +855,7 @@ class TestUfunc(object):
         assert_equal(np.min(a), False)
         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)
 
     def test_object_array_accumulate_inplace(self):
         # Checks that in-place accumulates work, see also gh-7402
@@ -987,7 +999,7 @@ class TestUfunc(object):
         assert_equal(np.sqrt(a, where=m), [1])
 
     def check_identityless_reduction(self, a):
-        # np.minimum.reduce is a identityless reduction
+        # np.minimum.reduce is an identityless reduction
 
         # Verify that it sees the zero at various positions
         a[...] = 1
@@ -1056,6 +1068,35 @@ class TestUfunc(object):
         a = a[1:, 1:, 1:]
         self.check_identityless_reduction(a)
 
+    def test_initial_reduction(self):
+        # np.minimum.reduce is an identityless reduction
+
+        # For cases like np.maximum(np.abs(...), initial=0)
+        # More generally, a supremum over non-negative numbers.
+        assert_equal(np.maximum.reduce([], initial=0), 0)
+
+        # For cases like reduction of an empty array over the reals.
+        assert_equal(np.minimum.reduce([], initial=np.inf), np.inf)
+        assert_equal(np.maximum.reduce([], initial=-np.inf), -np.inf)
+
+        # Random tests
+        assert_equal(np.minimum.reduce([5], initial=4), 4)
+        assert_equal(np.maximum.reduce([4], initial=5), 5)
+        assert_equal(np.maximum.reduce([5], initial=4), 5)
+        assert_equal(np.minimum.reduce([4], initial=5), 4)
+
+        # Check initial=None raises ValueError for both types of ufunc reductions
+        assert_raises(ValueError, np.minimum.reduce, [], initial=None)
+        assert_raises(ValueError, np.add.reduce, [], initial=None)
+
+        # Check that np._NoValue gives default behavior.
+        assert_equal(np.add.reduce([], initial=np._NoValue), 0)
+
+        # Check that initial kwarg behaves as intended for dtype=object
+        a = np.array([10], dtype=object)
+        res = np.add.reduce(a, initial=5)
+        assert_equal(res, 15)
+
     def test_identityless_reduction_nonreorderable(self):
         a = np.array([[8.0, 2.0, 2.0], [1.0, 0.5, 0.25]])
 
@@ -1407,15 +1448,18 @@ class TestUfunc(object):
         assert_equal(f(d, 0, None, None), r)
         assert_equal(f(d, 0, None, None, keepdims=False), r)
         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)
         # 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)
 
         # too little
         assert_raises(TypeError, f)
         # too much
-        assert_raises(TypeError, f, d, 0, None, None, False, 1)
+        assert_raises(TypeError, f, d, 0, None, None, False, 0, 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 9da6abd4b..042835650 100644
--- a/numpy/core/tests/test_umath.py
+++ b/numpy/core/tests/test_umath.py
@@ -1759,7 +1759,7 @@ class TestSpecialMethods(object):
 
         # reduce, kwargs
         res = np.multiply.reduce(a, axis='axis0', dtype='dtype0', out='out0',
-                                 keepdims='keep0')
+                                 keepdims='keep0', initial='init0')
         assert_equal(res[0], a)
         assert_equal(res[1], np.multiply)
         assert_equal(res[2], 'reduce')
@@ -1767,7 +1767,8 @@ class TestSpecialMethods(object):
         assert_equal(res[4], {'dtype':'dtype0',
                               'out': ('out0',),
                               'keepdims': 'keep0',
-                              'axis': 'axis0'})
+                              'axis': 'axis0',
+                              'initial': 'init0'})
 
         # reduce, output equal to None removed, but not other explicit ones,
         # even if they are at their default value.
@@ -1777,6 +1778,14 @@ 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})
 
         # reduce, wrong args
         assert_raises(ValueError, np.multiply.reduce, a, out=())