ENH: add np.divmod ufunc

9 files changed, 156 insertions, 0 deletions

diff --git a/doc/release/1.13.0-notes.rst b/doc/release/1.13.0-notes.rst
index 2f32ddb28..3945639da 100644
--- a/doc/release/1.13.0-notes.rst
+++ b/doc/release/1.13.0-notes.rst
@@ -344,6 +344,12 @@ New ``positive`` ufunc
 This ufunc corresponds to unary `+`, but unlike `+` on an ndarray it will raise
 an error if array values do not support numeric operations.
 
+New ``divmod`` ufunc
+--------------------
+This ufunc corresponds to the Python builtin `divmod`. ``np.divmod(x, y)``
+calculates a result equivalent to
+``(np.floor_divide(x, y), np.remainder(x, y))`` but faster.
+
 Better ``repr`` of object arrays
 --------------------------------
 Object arrays that contain themselves no longer cause a recursion error.
diff --git a/doc/source/reference/routines.math.rst b/doc/source/reference/routines.math.rst
index 72e1462b0..a2fb06958 100644
--- a/doc/source/reference/routines.math.rst
+++ b/doc/source/reference/routines.math.rst
@@ -121,6 +121,7 @@ Arithmetic operations
    mod
    modf
    remainder
+   divmod
 
 Handling complex numbers
 ------------------------
diff --git a/doc/source/reference/ufuncs.rst b/doc/source/reference/ufuncs.rst
index 94663a141..e82571f6c 100644
--- a/doc/source/reference/ufuncs.rst
+++ b/doc/source/reference/ufuncs.rst
@@ -510,6 +510,7 @@ Math operations
     remainder
     mod
     fmod
+    divmod
     absolute
     fabs
     rint
diff --git a/numpy/core/code_generators/generate_umath.py b/numpy/core/code_generators/generate_umath.py
index dfba04c18..7eb362f80 100644
--- a/numpy/core/code_generators/generate_umath.py
+++ b/numpy/core/code_generators/generate_umath.py
@@ -786,6 +786,13 @@ defdict = {
           TD(intflt),
           TD(O, f='PyNumber_Remainder'),
           ),
+'divmod':
+    Ufunc(2, 2, None,
+          docstrings.get('numpy.core.umath.divmod'),
+          None,
+          TD(intflt),
+          TD(O, f='PyNumber_Divmod'),
+          ),
 'hypot':
     Ufunc(2, 1, Zero,
           docstrings.get('numpy.core.umath.hypot'),
diff --git a/numpy/core/code_generators/ufunc_docstrings.py b/numpy/core/code_generators/ufunc_docstrings.py
index ed9e05b15..8677e4328 100644
--- a/numpy/core/code_generators/ufunc_docstrings.py
+++ b/numpy/core/code_generators/ufunc_docstrings.py
@@ -1290,6 +1290,7 @@ add_newdoc('numpy.core.umath', 'floor_divide',
     See Also
     --------
     remainder : Remainder complementary to floor_divide.
+    divmod : Simultaneous floor division and remainder.
     divide : Standard division.
     floor : Round a number to the nearest integer toward minus infinity.
     ceil : Round a number to the nearest integer toward infinity.
@@ -2474,6 +2475,11 @@ add_newdoc('numpy.core.umath', 'modf',
     -----
     For integer input the return values are floats.
 
+    See Also
+    --------
+    divmod : ``divmod(x, 1)`` is equivalent to ``modf`` with the return values
+             switched, except it always has a positive remainder.
+
     Examples
     --------
     >>> np.modf([0, 3.5])
@@ -2541,6 +2547,8 @@ add_newdoc('numpy.core.umath', 'positive',
     """
     Numerical positive, element-wise.
 
+    .. versionadded:: 1.13.0
+
     Parameters
     ----------
     x : array_like or scalar
@@ -2843,6 +2851,7 @@ add_newdoc('numpy.core.umath', 'remainder',
     See Also
     --------
     floor_divide : Equivalent of Python ``//`` operator.
+    divmod : Simultaneous floor division and remainder.
     fmod : Equivalent of the Matlab(TM) ``rem`` function.
     divide, floor
 
@@ -2860,6 +2869,47 @@ add_newdoc('numpy.core.umath', 'remainder',
 
     """)
 
+add_newdoc('numpy.core.umath', 'divmod',
+    """
+    Return element-wise quotient and remainder simultaneously.
+
+    .. versionadded:: 1.13.0
+
+    ``np.divmod(x, y)`` is equivalent to ``(x // y, x % y)``, but faster
+    because it avoids redundant work. It is used to implement the Python
+    built-in function ``divmod`` on NumPy arrays.
+
+    Parameters
+    ----------
+    x1 : array_like
+        Dividend array.
+    x2 : array_like
+        Divisor array.
+    out : tuple of ndarray, optional
+        Arrays into which the output is placed. Their types are preserved and
+        must be of the right shape to hold the output.
+
+    Returns
+    -------
+    out1 : ndarray
+        Element-wise quotient resulting from floor division.
+    out2 : ndarray
+        Element-wise remainder from division.
+
+    See Also
+    --------
+    floor_divide : Equivalent to Python's ``//`` operator.
+    remainder : Equivalent to Python's ``%`` operator.
+    modf : Equivalent to ``divmod(x, 1)`` for positive ``x`` with the return
+           values switched.
+
+    Examples
+    --------
+    >>> np.divmod(np.arange(5), 3)
+    (array([0, 0, 0, 1, 1]), array([0, 1, 2, 0, 1]))
+
+    """)
+
 add_newdoc('numpy.core.umath', 'right_shift',
     """
     Shift the bits of an integer to the right.
diff --git a/numpy/core/src/umath/loops.c.src b/numpy/core/src/umath/loops.c.src
index 47faaf180..596b50a1c 100644
--- a/numpy/core/src/umath/loops.c.src
+++ b/numpy/core/src/umath/loops.c.src
@@ -1114,6 +1114,34 @@ NPY_NO_EXPORT void
     }
 }
 
+NPY_NO_EXPORT void
+@TYPE@_divmod(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
+{
+    BINARY_LOOP_TWO_OUT {
+        const @type@ in1 = *(@type@ *)ip1;
+        const @type@ in2 = *(@type@ *)ip2;
+        /* see FIXME note for divide above */
+        if (in2 == 0 || (in1 == NPY_MIN_@TYPE@ && in2 == -1)) {
+            npy_set_floatstatus_divbyzero();
+            *((@type@ *)op1) = 0;
+            *((@type@ *)op2) = 0;
+        }
+        else {
+            /* handle mixed case the way Python does */
+            const @type@ quo = in1 / in2;
+            const @type@ rem = in1 % in2;
+            if ((in1 > 0) == (in2 > 0) || rem == 0) {
+                *((@type@ *)op1) = quo;
+                *((@type@ *)op2) = rem;
+            }
+            else {
+                *((@type@ *)op1) = quo - 1;
+                *((@type@ *)op2) = rem + in2;
+            }
+        }
+    }
+}
+
 /**end repeat**/
 
 /**begin repeat
@@ -1168,6 +1196,24 @@ NPY_NO_EXPORT void
     }
 }
 
+NPY_NO_EXPORT void
+@TYPE@_divmod(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
+{
+    BINARY_LOOP_TWO_OUT {
+        const @type@ in1 = *(@type@ *)ip1;
+        const @type@ in2 = *(@type@ *)ip2;
+        if (in2 == 0) {
+            npy_set_floatstatus_divbyzero();
+            *((@type@ *)op1) = 0;
+            *((@type@ *)op2) = 0;
+        }
+        else {
+            *((@type@ *)op1)= in1/in2;
+            *((@type@ *)op2) = in1 % in2;
+        }
+    }
+}
+
 /**end repeat**/
 
 /*
@@ -1832,6 +1878,16 @@ NPY_NO_EXPORT void
 }
 
 NPY_NO_EXPORT void
+@TYPE@_divmod(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
+{
+    BINARY_LOOP_TWO_OUT {
+        const @type@ in1 = *(@type@ *)ip1;
+        const @type@ in2 = *(@type@ *)ip2;
+        *((@type@ *)op1) = npy_divmod@c@(in1, in2, (@type@ *)op2);
+    }
+}
+
+NPY_NO_EXPORT void
 @TYPE@_square(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(data))
 {
     char * margs[] = {args[0], args[0], args[1]};
@@ -2160,6 +2216,16 @@ HALF_remainder(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNU
 }
 
 NPY_NO_EXPORT void
+HALF_divmod(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
+{
+    BINARY_LOOP_TWO_OUT {
+        const npy_half in1 = *(npy_half *)ip1;
+        const npy_half in2 = *(npy_half *)ip2;
+        *((npy_half *)op1) = npy_half_divmod(in1, in2, (npy_half *)op2);
+    }
+}
+
+NPY_NO_EXPORT void
 HALF_square(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(data))
 {
     UNARY_LOOP {
diff --git a/numpy/core/src/umath/loops.h.src b/numpy/core/src/umath/loops.h.src
index d20b776a0..a3d81e908 100644
--- a/numpy/core/src/umath/loops.h.src
+++ b/numpy/core/src/umath/loops.h.src
@@ -140,6 +140,9 @@ NPY_NO_EXPORT void
 NPY_NO_EXPORT void
 @S@@TYPE@_remainder(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func));
 
+NPY_NO_EXPORT void
+@S@@TYPE@_divmod(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func));
+
 /**end repeat1**/
 
 /**end repeat**/
@@ -220,6 +223,9 @@ NPY_NO_EXPORT void
 @TYPE@_remainder(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func));
 
 NPY_NO_EXPORT void
+@TYPE@_divmod(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func));
+
+NPY_NO_EXPORT void
 @TYPE@_square(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(data));
 
 NPY_NO_EXPORT void
diff --git a/numpy/core/tests/test_half.py b/numpy/core/tests/test_half.py
index 75b74f407..7a4d36333 100644
--- a/numpy/core/tests/test_half.py
+++ b/numpy/core/tests/test_half.py
@@ -317,6 +317,7 @@ class TestHalf(TestCase):
 
         assert_equal(np.floor_divide(a, b), [0, 0, 2, 1, 0])
         assert_equal(np.remainder(a, b), [0, 1, 0, 0, 2])
+        assert_equal(np.divmod(a, b), ([0, 0, 2, 1, 0], [0, 1, 0, 0, 2]))
         assert_equal(np.square(b), [4, 25, 1, 16, 9])
         assert_equal(np.reciprocal(b), [-0.5, 0.199951171875, 1, 0.25, 0.333251953125])
         assert_equal(np.ones_like(b), [1, 1, 1, 1, 1])
diff --git a/numpy/core/tests/test_umath.py b/numpy/core/tests/test_umath.py
index f13c056c3..7d147f727 100644
--- a/numpy/core/tests/test_umath.py
+++ b/numpy/core/tests/test_umath.py
@@ -286,6 +286,13 @@ class TestRemainder(TestCase):
                 else:
                     assert_(b > rem >= 0, msg)
 
+                div, rem = np.divmod(a, b)
+                assert_equal(div*b + rem, a, err_msg=msg)
+                if sg2 == -1:
+                    assert_(b < rem <= 0, msg)
+                else:
+                    assert_(b > rem >= 0, msg)
+
     def test_float_remainder_exact(self):
         # test that float results are exact for small integers. This also
         # holds for the same integers scaled by powers of two.
@@ -312,6 +319,10 @@ class TestRemainder(TestCase):
             assert_equal(div, tgtdiv, err_msg=msg)
             assert_equal(rem, tgtrem, err_msg=msg)
 
+            div, rem = np.divmod(fa, fb)
+            assert_equal(div, tgtdiv, err_msg=msg)
+            assert_equal(rem, tgtrem, err_msg=msg)
+
     def test_float_remainder_roundoff(self):
         # gh-6127
         dt = np.typecodes['Float']
@@ -330,6 +341,13 @@ class TestRemainder(TestCase):
                 else:
                     assert_(b > rem >= 0, msg)
 
+                div, rem = np.divmod(a, b)
+                assert_equal(div*b + rem, a, err_msg=msg)
+                if sg2 == -1:
+                    assert_(b < rem <= 0, msg)
+                else:
+                    assert_(b > rem >= 0, msg)
+
     def test_float_remainder_corner_cases(self):
         # Check remainder magnitude.
         for dt in np.typecodes['Float']: