8000 TST: Add a new gufunc with (i)->(i) signature for testing with axis. · numpy/numpy@b02fa32 · GitHub
[go: up one dir, main page]
Skip to content

Commit b02fa32

Browse files
mhvkmhvk
committed
TST: Add a new gufunc with (i)->(i) signature for testing with axis.
1 parent c5ef411 commit b02fa32

File tree

2 files changed

+64
-1
lines changed

2 files changed

+64
-1
lines changed

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

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -253,6 +253,38 @@ static void
253253

254254
/**end repeat**/
255255

256+
char *cumsum_signature = "(i)->(i)";
257+
258+
/*
259+
* This implements the function
260+
* out[n] = sum_i^n in[i]
261+
*/
262+
263+
/**begin repeat
264+
265+
#TYPE=LONG,DOUBLE#
266+
#typ=npy_long,npy_double#
267+
*/
268+
269+
static void
270+
@TYPE@_cumsum(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
271+
{
272+
INIT_OUTER_LOOP_2
273+
npy_intp di = dimensions[0];
274+
npy_intp i;
275+
npy_intp is=steps[0], os=steps[1];
276+
BEGIN_OUTER_LOOP_2
277+
char *ip=args[0], *op=args[1];
278+
@typ@ cumsum = 0;
279+
for (i = 0; i < di; i++, ip += is, op += os) {
280+
cumsum += (*(@typ@ *)ip);
281+
*(@typ@ *)op = cumsum;
282+
}
283+
END_OUTER_LOOP
284+
}
285+
286+
/**end repeat**/
287+
256288

257289
static PyUFuncGenericFunction inner1d_functions[] = { LONG_inner1d, DOUBLE_inner1d };
258290
static void * inner1d_data[] = { (void *)NULL, (void *)NULL };
@@ -270,6 +302,10 @@ static void *eucldiean_pdist_data[] = { (void *)NULL, (void *)NULL };
270302
static char euclidean_pdist_signatures[] = { NPY_FLOAT, NPY_FLOAT,
271303
NPY_DOUBLE, NPY_DOUBLE };
272304

305+
static PyUFuncGenericFunction cumsum_functions[] = { LONG_cumsum, DOUBLE_cumsum };
306+
static void * cumsum_data[] = { (void *)NULL, (void *)NULL };
307+
static char cumsum_signatures[] = { NPY_LONG, NPY_LONG, NPY_DOUBLE, NPY_DOUBLE };
308+
273309

274310
static int
275311
addUfuncs(PyObject *dictionary) {
@@ -321,6 +357,16 @@ addUfuncs(PyObject *dictionary) {
321357
}
322358
PyDict_SetItemString(dictionary, "euclidean_pdist", f);
323359
Py_DECREF(f);
360+
f = PyUFunc_FromFuncAndDataAndSignature(cumsum_functions,
361+
cumsum_data, cumsum_signatures,
362+
2, 1, 1, PyUFunc_None, "cumsum",
363+
"Cumulative sum of the input (n)->(n)\n",
364+
0, cumsum_signature);
365+
if (f == NULL) {
366+
return -1;
367+
}
368+
PyDict_SetItemString(dictionary, "cumsum", f);
369+
Py_DECREF(f);
324370
f = PyUFunc_FromFuncAndDataAndSignature(inner1d_functions, inner1d_data,
325371
inner1d_signatures, 2, 2, 1, PyUFunc_None, "inner1d_no_doc",
326372
NULL,

‎numpy/core/tests/test_ufunc.py

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -741,12 +741,21 @@ def test_axis_argument(self):
741741
assert_array_equal(d, c)
742742
c = inner1d(a, b, axis=0)
743743
assert_array_equal(c, (a * b).sum(0))
744-
# Sanity check on innerwt.
744+
# Sanity checks on innerwt and cumsum.
745745
a = np.arange(6).reshape((2, 3))
746746
b = np.arange(10, 16).reshape((2, 3))
747747
w = np.arange(20, 26).reshape((2, 3))
748748
assert_array_equal(umt.innerwt(a, b, w, axis=0),
749749
np.sum(a * b * w, axis=0))
750+
assert_array_equal(umt.cumsum(a, axis=0), np.cumsum(a, axis=0))
751+
assert_array_equal(umt.cumsum(a, axis=-1), np.cumsum(a, axis=-1))
752+
out = np.empty_like(a)
753+
b = umt.cumsum(a, out=out, axis=0)
754+
assert_(out is b)
755+
assert_array_equal(b, np.cumsum(a, axis=0))
756+
b = umt.cumsum(a, out=out, axis=1)
757+
assert_(out is b)
758+
assert_array_equal(b, np.cumsum(a, axis=-1))
750759
# Check errors.
751760
# Cannot pass in both axis and axes.
752761
assert_raises(TypeError, inner1d, a, b, axis=0, axes=[0, 0])
@@ -755,6 +764,9 @@ def test_axis_argument(self):
755764
# more than 1 core dimensions.
756765
mm = umt.matrix_multiply
757766
assert_raises(TypeError, mm, a, b, axis=1)
767+
# Output wrong size in axis.
768+
out = np.empty((1, 2, 3), dtype=a.dtype)
769+
assert_raises(ValueError, umt.cumsum, a, out=out, axis=0)
758770
# Regular ufuncs should not accept axis.
759771
assert_raises(TypeError, np.add, 1., 1., axis=0)
760772

@@ -928,6 +940,11 @@ def test_euclidean_pdist(self):
928940
# An output array is required to determine p with signature (n,d)->(p)
929941
assert_raises(ValueError, umt.euclidean_pdist, a)
930942

943+
def test_cumsum(self):
944+
a = np.arange(10)
945+
result = umt.cumsum(a)
946+
assert_array_equal(result, a.cumsum())
947+
931948
def test_object_logical(self):
932949
a = np.array([3, None, True, False, "test", ""], dtype=object)
933950
assert_equal(np.logical_or(a, None),

0 commit comments

Comments
 (0)
0