10000 Merge pull request #9063 from shoyer/divmod · numpy/numpy@11f3ebf · GitHub
[go: up one dir, main page]
Skip to content

Commit 11f3ebf

Browse files
eric-wiesereric-wieser
authored
Merge pull request #9063 from shoyer/divmod
ENH: add np.divmod ufunc
2 parents d7d1b2a + 8fbf75e commit 11f3ebf

17 files changed

+348
-193
lines changed

‎doc/neps/ufunc-overrides.rst

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -662,15 +662,15 @@ Symbol Operator NumPy Ufunc(s)
662662
``/`` ``div`` :func:`divide`
663663
(Python 2)
664664
``//`` ``floordiv`` :func:`floor_divide`
665-
``%`` ``mod`` :func:`mod`
665+
``%`` ``mod`` :func:`remainder`
666+
NA ``divmod`` :func:`divmod`
666667
``**`` ``pow`` :func:`power`
667668
``<<`` ``lshift`` :func:`left_shift`
668669
``>>`` ``rshift`` :func:`right_shift`
669670
``&`` ``and_`` :func:`bitwise_and`
670671
``^`` ``xor_`` :func:`bitwise_xor`
671672
``|`` ``or_`` :func:`bitwise_or`
672-
NA ``divmod`` :func:`floor_divide`, :func:`mod` [10]_
673-
``@`` ``matmul`` Not yet implemented as a ufunc
673+
``@`` ``matmul`` Not yet implemented as a ufunc [10]_
674674
====== ============ =========================================
675675

676676
And here is the list of unary operators:
@@ -684,7 +684,10 @@ NA ``abs`` :func:`absolute`
684684
``~`` ``invert`` :func:`invert`
685685
====== ============ =========================================
686686

687-
.. [10] In the future, NumPy may switch to use a single ufunc ``divmod_`` instead.
687+
.. [10] Because NumPy's :func:`matmul` is not a ufunc, it is
688+
`currently not possible <https://github.com/numpy/numpy/issues/9028>`_
689+
to override ``numpy_array @ other`` with ``other`` taking precedence
690+
if ``other`` implements ``__array_func__``.
688691
.. [11] :class:`ndarray` currently does a copy instead of using this ufunc.
689692
690693
Future extensions to other functions

‎doc/release/1.13.0-notes.rst

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -368,6 +368,13 @@ New ``positive`` ufunc
368368
This ufunc corresponds to unary `+`, but unlike `+` on an ndarray it will raise
369369
an error if array values do not support numeric operations.
370370

371+
New ``divmod`` ufunc
372+
--------------------
373+
This ufunc corresponds to the Python builtin `divmod`, and is used to implement
374+
`divmod` when called on numpy arrays. ``np.divmod(x, y)`` calculates a result
375+
equivalent to ``(np.floor_divide(x, y), np.remainder(x, y))`` but is
376+
approximately twice as fast as calling the functions separately.
377+
371378
Better ``repr`` of object arrays
372379
--------------------------------
373380
Object arrays that contain themselves no longer cause a recursion error.

‎doc/source/reference/routines.math.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,7 @@ Arithmetic operations
121121
mod
122122
modf
123123
remainder
124+
divmod
124125

125126
Handling complex numbers
126127
------------------------

‎doc/source/reference/ufuncs.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -510,6 +510,7 @@ Math operations
510510
remainder
511511
mod
512512
fmod
513+
divmod
513514
absolute
514515
fabs
515516
rint

‎numpy/core/code_generators/generate_umath.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -786,6 +786,13 @@ def english_upper(s):
786786
TD(intflt),
787787
TD(O, f='PyNumber_Remainder'),
788788
),
789+
'divmod':
790+
Ufunc(2, 2, None,
791+
docstrings.get('numpy.core.umath.divmod'),
792+
None,
793+
TD(intflt),
794+
TD(O, f='PyNumber_Divmod'),
795+
),
789796
'hypot':
790797
Ufunc(2, 1, Zero,
791798
docstrings.get('numpy.core.umath.hypot'),

‎numpy/core/code_generators/ufunc_docstrings.py

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1290,6 +1290,7 @@ def add_newdoc(place, name, doc):
12901290
See Also
12911291
--------
12921292
remainder : Remainder complementary to floor_divide.
1293+
divmod : Simultaneous floor division and remainder.
12931294
divide : Standard division.
12941295
floor : Round a number to the nearest integer toward minus infinity.
12951296
ceil : Round a number to the nearest integer toward infinity.
@@ -2509,6 +2510,11 @@ def add_newdoc(place, name, doc):
25092510
-----
25102511
For integer input the return values are floats.
25112512
2513+
See Also
2514+
--------
2515+
divmod : ``divmod(x, 1)`` is equivalent to ``modf`` with the return values
2516+
switched, except it always has a positive remainder.
2517+
25122518
Examples
25132519
--------
25142520
>>> np.modf([0, 3.5])
@@ -2576,6 +2582,8 @@ def add_newdoc(place, name, doc):
25762582
"""
25772583
Numerical positive, element-wise.
25782584
2585+
.. versionadded:: 1.13.0
2586+
25792587
Parameters
25802588
----------
25812589
x : array_like or scalar
@@ -2878,6 +2886,7 @@ def add_newdoc(place, name, doc):
28782886
See Also
28792887
--------
28802888
floor_divide : Equivalent of Python ``//`` operator.
2889+
divmod : Simultaneous floor division and remainder.
28812890
fmod : Equivalent of the Matlab(TM) ``rem`` function.
28822891
divide, floor
28832892
@@ -2895,6 +2904,47 @@ def add_newdoc(place, name, doc):
28952904
28962905
""")
28972906

2907+
add_newdoc('numpy.core.umath', 'divmod',
2908+
"""
2909+
Return element-wise quotient and remainder simultaneously.
2910+
2911+
.. versionadded:: 1.13.0
2912+
2913+
``np.divmod(x, y)`` is equivalent to ``(x // y, x % y)``, but faster
2914+
because it avoids redundant work. It is used to implement the Python
2915+
built-in function ``divmod`` on NumPy arrays.
2916+
2917+
Parameters
2918+
----------
2919+
x1 : array_like
2920+
Dividend array.
2921+
x2 : array_like
2922+
Divisor array.
2923+
out : tuple of ndarray, optional
2924+
Arrays into which the output is placed. Their types are preserved and
2925+
must be of the right shape to hold the output.
2926+
2927+
Returns
2928+
-------
2929+
out1 : ndarray
2930+
Element-wise quotient resulting from floor division.
2931+
out2 : ndarray
2932+
Element-wise remainder from floor division.
2933+
2934+
See Also
2935+
--------
2936+
floor_divide : Equivalent to Python's ``//`` operator.
2937+
remainder : Equivalent to Python's ``%`` operator.
2938+
modf : Equivalent to ``divmod(x, 1)`` for positive ``x`` with the return
2939+
values switched.
2940+
2941+
Examples
2942+
--------
2943+
>>> np.divmod(np.arange(5), 3)
2944+
(array([0, 0, 0, 1, 1]), array([0, 1, 2, 0, 1]))
2945+
2946+
""")
2947+
28982948
add_newdoc('numpy.core.umath', 'right_shift',
28992949
"""
29002950
Shift the bits of an integer to the right.

‎numpy/core/src/multiarray/number.c

Lines changed: 8 additions & 31 deletions
10000
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@ PyArray_SetNumericOps(PyObject *dict)
8383
SET(multiply);
8484
SET(divide);
8585
SET(remainder);
86+
SET(divmod);
8687
SET(power);
8788
SET(square);
8889
SET(reciprocal);
@@ -135,6 +136,7 @@ PyArray_GetNumericOps(void)
135136
GET(multiply);
136137
GET(divide);
137138
GET(remainder);
139+
GET(divmod);
138140
GET(power);
139141
GET(square);
140142
GET(reciprocal);
@@ -344,6 +346,12 @@ array_remainder(PyArrayObject *m1, PyObject *m2)
344346
return PyArray_GenericBinaryFunction(m1, m2, n_ops.remainder);
345347
}
346348

349+
static PyObject *
350+
array_divmod(PyArrayObject *m1, PyObject *m2)
351+
{
352+
BINOP_GIVE_UP_IF_NEEDED(m1, m2, nb_divmod, array_divmod);
353+
return PyArray_GenericBinaryFunction(m1, m2, n_ops.divmod);
354+
}
347355

348356
#if PY_VERSION_HEX >= 0x03050000
349357
/* Need this to be version dependent on account of the slot check */
@@ -796,37 +804,6 @@ _array_nonzero(PyArrayObject *mp)
796804
}
797805

798806

799-
800-
static PyObject *
801-
array_divmod(PyArrayObject *op1, PyObject *op2)
802-
{
803-
PyObject *divp, *modp, *result;
804-
805-
BINOP_GIVE_UP_IF_NEEDED(op1, op2, nb_divmod, array_divmod);
806-
807-
divp = array_floor_divide(op1, op2);
808-
if (divp == NULL) {
809-
return NULL;
810-
}
811-
else if(divp == Py_NotImplemented) {
812-
return divp;
813-
}
814-
modp = array_remainder(op1, op2);
815-
if (modp == NULL) {
816-
Py_DECREF(divp);
817-
return NULL;
818-
}
819-
else if(modp == Py_NotImplemented) {
820-
Py_DECREF(divp);
821-
return modp;
822-
}
823-
result = Py_BuildValue("OO", divp, modp);
824-
Py_DECREF(divp);
825-
Py_DECREF(modp);
826-
return result;
827-
}
828-
829-
830807
NPY_NO_EXPORT PyObject *
831808
array_int(PyArrayObject *v)
832809
{

‎numpy/core/src/multiarray/number.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ typedef struct {
77
PyObject *multiply;
88
PyObject *divide;
99
PyObject *remainder;
10+
PyObject *divmod;
1011
PyObject *power;
1112
PyObject *square;
1213
PyObject *reciprocal;

‎numpy/core/src/umath/loops.c.src

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1114,6 +1114,34 @@ NPY_NO_EXPORT void
11141114
}
11151115
}
11161116

1117+
NPY_NO_EXPORT void
1118+
@TYPE@_divmod(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
1119+
{
1120+
BINARY_LOOP_TWO_OUT {
1121+
const @type@ in1 = *(@type@ *)ip1;
1122+
const @type@ in2 = *(@type@ *)ip2;
1123+
/* see FIXME note for divide above */
1124+
if (in2 == 0 || (in1 == NPY_MIN_@TYPE@ && in2 == -1)) {
1125+
npy_set_floatstatus_divbyzero();
1126+
*((@type@ *)op1) = 0;
1127+
*((@type@ *)op2) = 0;
1128+
}
1129+
else {
1130+
/* handle mixed case the way Python does */
1131+
const @type@ quo = in1 / in2;
1132+
const @type@ rem = in1 % in2;
1133+
if ((in1 > 0) == (in2 > 0) || rem == 0) {
1134+
*((@type@ *)op1) = quo;
1135+
*((@type@ *)op2) = rem;
1136+
}
1137+
else {
1138+
*((@type@ *)op1) = quo - 1;
1139+
*((@type@ *)op2) = rem + in2;
1140+
}
1141+
}
1142+
}
1143+
}
1144+
11171145
/**end repeat**/
11181146

11191147
/**begin repeat
@@ -1168,6 +1196,24 @@ NPY_NO_EXPORT void
11681196
}
11691197
}
11701198

1199+
NPY_NO_EXPORT void
1200+
@TYPE@_divmod(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
1201+
{
1202+
BINARY_LOOP_TWO_OUT {
1203+
const @type@ in1 = *(@type@ *)ip1;
1204+
const @type@ in2 = *(@type@ *)ip2;
1205+
if (in2 == 0) {
1206+
npy_set_floatstatus_divbyzero();
1207+
*((@type@ *)op1) = 0;
1208+
*((@type@ *)op2) = 0;
1209+
}
1210+
else {
1211+
*((@type@ *)op1)= in1/in2;
1212+
*((@type@ *)op2) = in1 % in2;
1213+
}
1214+
}
1215+
}
1216+
11711217
/**end repeat**/
11721218

11731219
/*
@@ -1840,6 +1886,16 @@ NPY_NO_EXPORT void
18401886
}
18411887
}
18421888

1889+
NPY_NO_EXPORT void
1890+
@TYPE@_divmod(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
1891+
{
1892+
BINARY_LOOP_TWO_OUT {
1893+
const @type@ in1 = *(@type@ *)ip1;
1894+
const @type@ in2 = *(@type@ *)ip2;
1895+
*((@type@ *)op1) = npy_divmod@c@(in1, in2, (@type@ *)op2);
1896+
}
1897+
}
1898+
18431899
NPY_NO_EXPORT void
18441900
@TYPE@_square(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(data))
18451901
{
@@ -2168,6 +2224,16 @@ HALF_remainder(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNU
21682224
}
21692225
}
21702226

2227+
NPY_NO_EXPORT void
2228+
HALF_divmod(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
2229+
{
2230+
BINARY_LOOP_TWO_OUT {
2231+
const npy_half in1 = *(npy_half *)ip1;
2232+
const npy_half in2 = *(npy_half *)ip2;
2233+
*((npy_half *)op1) = npy_half_divmod(in1, in2, (npy_half *)op2);
2234+
}
2235+
}
2236+
21712237
NPY_NO_EXPORT void
21722238
HALF_square(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(data))
21732239
{

‎numpy/core/src/umath/loops.h.src

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,9 @@ NPY_NO_EXPORT void
140140
NPY_NO_EXPORT void
141141
@S@@TYPE@_remainder(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func));
142142

143+
NPY_NO_EXPORT void
144+
@S@@TYPE@_divmod(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func));
145+
143146
/**end repeat1**/
144147

145148
/**end repeat**/
@@ -219,6 +222,9 @@ NPY_NO_EXPORT void
219222
NPY_NO_EXPORT void
220223
@TYPE@_remainder(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func));
221224

225+
NPY_NO_EXPORT void
226+
@TYPE@_divmod(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func));
227+
222228
NPY_NO_EXPORT void
223229
@TYPE@_square(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(data));
224230

‎numpy/core/tests/test_half.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -317,6 +317,7 @@ def test_half_ufuncs(self):
317317

318318
assert_equal(np.floor_divide(a, b), [0, 0, 2, 1, 0])
319319
assert_equal(np.remainder(a, b), [0, 1, 0, 0, 2])
320+
assert_equal(np.divmod(a, b), ([0, 0, 2, 1, 0], [0, 1, 0, 0, 2]))
320321
assert_equal(np.square(b), [4, 25, 1, 16, 9])
321322
assert_equal(np.reciprocal(b), [-0.5, 0.199951171875, 1, 0.25, 0.333251953125])
322323
assert_equal(np.ones_like(b), [1, 1, 1, 1, 1])

0 commit comments

Comments
 (0)
0