summaryrefslogtreecommitdiff
path: root/numpy/core
diff options
context:
space:
mode:
Diffstat (limited to 'numpy/core')
-rw-r--r--numpy/core/src/umath/_umath_tests.c.src46
-rw-r--r--numpy/core/src/umath/override.c12
-rw-r--r--numpy/core/src/umath/ufunc_object.c144
-rw-r--r--numpy/core/tests/test_ufunc.py63
-rw-r--r--numpy/core/tests/test_umath.py1
5 files changed, 246 insertions, 20 deletions
diff --git a/numpy/core/src/umath/_umath_tests.c.src b/numpy/core/src/umath/_umath_tests.c.src
index 2a74c1aaa..fcbdbe330 100644
--- a/numpy/core/src/umath/_umath_tests.c.src
+++ b/numpy/core/src/umath/_umath_tests.c.src
@@ -253,6 +253,38 @@ static void
/**end repeat**/
+char *cumsum_signature = "(i)->(i)";
+
+/*
+ * This implements the function
+ * out[n] = sum_i^n in[i]
+ */
+
+/**begin repeat
+
+ #TYPE=LONG,DOUBLE#
+ #typ=npy_long,npy_double#
+*/
+
+static void
+@TYPE@_cumsum(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
+{
+ INIT_OUTER_LOOP_2
+ npy_intp di = dimensions[0];
+ npy_intp i;
+ npy_intp is=steps[0], os=steps[1];
+ BEGIN_OUTER_LOOP_2
+ char *ip=args[0], *op=args[1];
+ @typ@ cumsum = 0;
+ for (i = 0; i < di; i++, ip += is, op += os) {
+ cumsum += (*(@typ@ *)ip);
+ *(@typ@ *)op = cumsum;
+ }
+ END_OUTER_LOOP
+}
+
+/**end repeat**/
+
static PyUFuncGenericFunction inner1d_functions[] = { LONG_inner1d, DOUBLE_inner1d };
static void * inner1d_data[] = { (void *)NULL, (void *)NULL };
@@ -270,6 +302,10 @@ static void *eucldiean_pdist_data[] = { (void *)NULL, (void *)NULL };
static char euclidean_pdist_signatures[] = { NPY_FLOAT, NPY_FLOAT,
NPY_DOUBLE, NPY_DOUBLE };
+static PyUFuncGenericFunction cumsum_functions[] = { LONG_cumsum, DOUBLE_cumsum };
+static void * cumsum_data[] = { (void *)NULL, (void *)NULL };
+static char cumsum_signatures[] = { NPY_LONG, NPY_LONG, NPY_DOUBLE, NPY_DOUBLE };
+
static int
addUfuncs(PyObject *dictionary) {
@@ -321,6 +357,16 @@ addUfuncs(PyObject *dictionary) {
}
PyDict_SetItemString(dictionary, "euclidean_pdist", f);
Py_DECREF(f);
+ f = PyUFunc_FromFuncAndDataAndSignature(cumsum_functions,
+ cumsum_data, cumsum_signatures,
+ 2, 1, 1, PyUFunc_None, "cumsum",
+ "Cumulative sum of the input (n)->(n)\n",
+ 0, cumsum_signature);
+ if (f == NULL) {
+ return -1;
+ }
+ PyDict_SetItemString(dictionary, "cumsum", f);
+ Py_DECREF(f);
f = PyUFunc_FromFuncAndDataAndSignature(inner1d_functions, inner1d_data,
inner1d_signatures, 2, 2, 1, PyUFunc_None, "inner1d_no_doc",
NULL,
diff --git a/numpy/core/src/umath/override.c b/numpy/core/src/umath/override.c
index c298fe315..c0bc47b7b 100644
--- a/numpy/core/src/umath/override.c
+++ b/numpy/core/src/umath/override.c
@@ -51,6 +51,7 @@ normalize___call___args(PyUFuncObject *ufunc, PyObject *args,
npy_intp nin = ufunc->nin;
npy_intp nout = ufunc->nout;
npy_intp nargs = PyTuple_GET_SIZE(args);
+ npy_intp nkwds = PyDict_Size(*normal_kwds);
PyObject *obj;
if (nargs < nin) {
@@ -74,7 +75,7 @@ normalize___call___args(PyUFuncObject *ufunc, PyObject *args,
/* If we have more args than nin, they must be the output variables.*/
if (nargs > nin) {
- if(PyDict_GetItemString(*normal_kwds, "out")) {
+ if(nkwds > 0 && PyDict_GetItemString(*normal_kwds, "out")) {
PyErr_Format(PyExc_TypeError,
"argument given by name ('out') and position "
"(%"NPY_INTP_FMT")", nin);
@@ -112,8 +113,15 @@ normalize___call___args(PyUFuncObject *ufunc, PyObject *args,
Py_DECREF(obj);
}
}
+ /* gufuncs accept either 'axes' or 'axis', but not both */
+ if (nkwds >= 2 && (PyDict_GetItemString(*normal_kwds, "axis") &&
+ PyDict_GetItemString(*normal_kwds, "axes"))) {
+ PyErr_SetString(PyExc_TypeError,
+ "cannot specify both 'axis' and 'axes'");
+ return -1;
+ }
/* finally, ufuncs accept 'sig' or 'signature' normalize to 'signature' */
- return normalize_signature_keyword(*normal_kwds);
+ return nkwds == 0 ? 0 : normalize_signature_keyword(*normal_kwds);
}
static int
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 631999bc1..9b03a7916 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -566,11 +566,12 @@ get_ufunc_arguments(PyUFuncObject *ufunc,
NPY_ORDER *out_order,
NPY_CASTING *out_casting,
PyObject **out_extobj,
- PyObject **out_typetup,
- int *out_subok,
- PyArrayObject **out_wheremask,
- PyObject **out_axes,
- int *out_keepdims)
+ PyObject **out_typetup, /* type: Tuple[np.dtype] */
+ int *out_subok, /* bool */
+ PyArrayObject **out_wheremask, /* PyArray of bool */
+ PyObject **out_axes, /* type: List[Tuple[T]] */
+ PyObject **out_axis, /* type: T */
+ int *out_keepdims) /* bool */
{
int i, nargs;
int nin = ufunc->nin;
@@ -588,6 +589,9 @@ get_ufunc_arguments(PyUFuncObject *ufunc,
if (out_axes != NULL) {
*out_axes = NULL;
}
+ if (out_axis != NULL) {
+ *out_axis = NULL;
+ }
if (out_wheremask != NULL) {
*out_wheremask = NULL;
}
@@ -827,9 +831,23 @@ get_ufunc_arguments(PyUFuncObject *ufunc,
case 'a':
/* possible axes argument for generalized ufunc */
if (out_axes != NULL && strcmp(str, "axes") == 0) {
+ if (out_axis != NULL && *out_axis != NULL) {
+ PyErr_SetString(PyExc_TypeError,
+ "cannot specify both 'axis' and 'axes'");
+ goto fail;
+ }
Py_INCREF(value);
*out_axes = value;
-
+ bad_arg = 0;
+ }
+ else if (out_axis != NULL && strcmp(str, "axis") == 0) {
+ if (out_axes != NULL && *out_axes != NULL) {
+ PyErr_SetString(PyExc_TypeError,
+ "cannot specify both 'axis' and 'axes'");
+ goto fail;
+ }
+ Py_INCREF(value);
+ *out_axis = value;
bad_arg = 0;
}
break;
@@ -1045,6 +1063,10 @@ fail:
Py_XDECREF(*out_axes);
*out_axes = NULL;
}
+ if (out_axis != NULL) {
+ Py_XDECREF(*out_axis);
+ *out_axis = NULL;
+ }
return -1;
}
@@ -1891,6 +1913,27 @@ _has_output_coredims(PyUFuncObject *ufunc) {
}
/*
+ * Check whether the gufunc can be used with axis, i.e., that there is only
+ * a single, shared core dimension (which means that operands either have
+ * that dimension, or have no core dimensions). Returns 0 if all is fine,
+ * and sets an error and returns -1 if not.
+ */
+static int
+_check_axis_support(PyUFuncObject *ufunc) {
+ if (ufunc->core_num_dim_ix != 1) {
+ PyErr_Format(PyExc_TypeError,
+ "%s: axis can only be used with a single shared core "
+ "dimension, not with the %d distinct ones implied by "
+ "signature %s.",
+ ufunc_get_name_cstr(ufunc),
+ ufunc->core_num_dim_ix,
+ ufunc->core_signature);
+ return -1;
+ }
+ return 0;
+}
+
+/*
* Check whether the gufunc can be used with keepdims, i.e., that all its
* input arguments have the same number of core dimension, and all output
* arguments have no core dimensions. Returns 0 if all is fine, and sets
@@ -1905,7 +1948,7 @@ _check_keepdims_support(PyUFuncObject *ufunc) {
if (ufunc->core_num_dims[i] != (i < nin ? input_core_dims : 0)) {
PyErr_Format(PyExc_TypeError,
"%s does not support keepdims: its signature %s requires "
- "that %s %d has %d core dimensions, but keepdims can only "
+ "%s %d to have %d core dimensions, but keepdims can only "
"be used when all inputs have the same number of core "
"dimensions and all outputs have no core dimensions.",
ufunc_get_name_cstr(ufunc),
@@ -1931,8 +1974,7 @@ static int
_parse_axes_arg(PyUFuncObject *ufunc, int core_num_dims[], PyObject *axes,
PyArrayObject **op, int broadcast_ndim, int **remap_axis) {
int nin = ufunc->nin;
- int nout = ufunc->nout;
- int nop = nin + nout;
+ int nop = ufunc->nargs;
int iop, list_size;
if (!PyList_Check(axes)) {
@@ -2050,6 +2092,59 @@ _parse_axes_arg(PyUFuncObject *ufunc, int core_num_dims[], PyObject *axes,
return 0;
}
+/*
+ * Simplified version of the above, using axis to fill the remap_axis
+ * array, which maps default to actual axes for each operand, indexed as
+ * as remap_axis[iop][iaxis]. The default axis order has first all broadcast
+ * axes and then the core axes the gufunc operates on.
+ *
+ * Returns 0 on success, and -1 on failure
+ */
+static int
+_parse_axis_arg(PyUFuncObject *ufunc, int core_num_dims[], PyObject *axis,
+ PyArrayObject **op, int broadcast_ndim, int **remap_axis) {
+ int nop = ufunc->nargs;
+ int iop, axis_int;
+
+ axis_int = PyArray_PyIntAsInt(axis);
+ if (error_converting(axis_int)) {
+ return -1;
+ }
+
+ for (iop = 0; iop < nop; ++iop) {
+ int axis, op_ndim, op_axis;
+
+ /* _check_axis_support ensures core_num_dims is 0 or 1 */
+ if (core_num_dims[iop] == 0) {
+ remap_axis[iop] = NULL;
+ continue;
+ }
+ if (op[iop]) {
+ op_ndim = PyArray_NDIM(op[iop]);
+ }
+ else {
+ op_ndim = broadcast_ndim + 1;
+ }
+ op_axis = axis_int; /* ensure we don't modify axis_int */
+ if (check_and_adjust_axis(&op_axis, op_ndim) < 0) {
+ return -1;
+ }
+ /* Are we actually remapping away from last axis? */
+ if (op_axis == op_ndim - 1) {
+ remap_axis[iop] = NULL;
+ continue;
+ }
+ remap_axis[iop][op_ndim - 1] = op_axis;
+ for (axis = 0; axis < op_axis; axis++) {
+ remap_axis[iop][axis] = axis;
+ }
+ for (axis = op_axis; axis < op_ndim - 1; axis++) {
+ remap_axis[iop][axis] = axis + 1;
+ }
+ } /* end of for(iop) loop over operands */
+ return 0;
+}
+
#define REMAP_AXIS(iop, axis) ((remap_axis != NULL && \
remap_axis[iop] != NULL)? \
remap_axis[iop][axis] : axis)
@@ -2245,7 +2340,8 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
NPY_ORDER order = NPY_KEEPORDER;
/* Use the default assignment casting rule */
NPY_CASTING casting = NPY_DEFAULT_ASSIGN_CASTING;
- PyObject *extobj = NULL, *type_tup = NULL, *axes = NULL;
+ /* other possible keyword arguments */
+ PyObject *extobj = NULL, *type_tup = NULL, *axes = NULL, *axis = NULL;
int keepdims = -1;
if (ufunc == NULL) {
@@ -2270,10 +2366,12 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
NPY_UF_DBG_PRINT("Getting arguments\n");
- /* Get all the arguments */
+ /*
+ * Get all the arguments.
+ */
retval = get_ufunc_arguments(ufunc, args, kwds,
op, &order, &casting, &extobj,
- &type_tup, &subok, NULL, &axes, &keepdims);
+ &type_tup, &subok, NULL, &axes, &axis, &keepdims);
if (retval < 0) {
goto fail;
}
@@ -2288,6 +2386,12 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
goto fail;
}
}
+ if (axis != NULL) {
+ retval = _check_axis_support(ufunc);
+ if (retval < 0) {
+ goto fail;
+ }
+ }
/*
* If keepdims is set and true, signal all dimensions will be the same.
*/
@@ -2356,7 +2460,7 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
}
/* Possibly remap axes. */
- if (axes) {
+ if (axes != NULL || axis != NULL) {
remap_axis = PyArray_malloc(sizeof(remap_axis[0]) * nop);
remap_axis_memory = PyArray_malloc(sizeof(remap_axis_memory[0]) *
nop * NPY_MAXDIMS);
@@ -2367,8 +2471,14 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
for (i=0; i < nop; i++) {
remap_axis[i] = remap_axis_memory + i * NPY_MAXDIMS;
}
- retval = _parse_axes_arg(ufunc, core_num_dims, axes, op, broadcast_ndim,
- remap_axis);
+ if (axis) {
+ retval = _parse_axis_arg(ufunc, core_num_dims, axis, op,
+ broadcast_ndim, remap_axis);
+ }
+ else {
+ retval = _parse_axes_arg(ufunc, core_num_dims, axes, op,
+ broadcast_ndim, remap_axis);
+ }
if(retval < 0) {
goto fail;
}
@@ -2714,6 +2824,7 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
Py_XDECREF(type_tup);
Py_XDECREF(extobj);
Py_XDECREF(axes);
+ Py_XDECREF(axis);
Py_XDECREF(full_args.in);
Py_XDECREF(full_args.out);
@@ -2735,6 +2846,7 @@ fail:
Py_XDECREF(type_tup);
Py_XDECREF(extobj);
Py_XDECREF(axes);
+ Py_XDECREF(axis);
Py_XDECREF(full_args.in);
Py_XDECREF(full_args.out);
PyArray_free(remap_axis_memory);
@@ -2812,7 +2924,7 @@ PyUFunc_GenericFunction(PyUFuncObject *ufunc,
/* Get all the arguments */
retval = get_ufunc_arguments(ufunc, args, kwds,
op, &order, &casting, &extobj,
- &type_tup, &subok, &wheremask, NULL, NULL);
+ &type_tup, &subok, &wheremask, NULL, NULL, NULL);
if (retval < 0) {
goto fail;
}
diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py
index ac4346461..ef9ced354 100644
--- a/numpy/core/tests/test_ufunc.py
+++ b/numpy/core/tests/test_ufunc.py
@@ -726,6 +726,50 @@ class TestUfunc(object):
# should be able to deal with bad unrelated kwargs.
assert_raises(TypeError, mm, z, z, axes=[0, 1], parrot=True)
+ def test_axis_argument(self):
+ # inner1d signature: '(i),(i)->()'
+ inner1d = umt.inner1d
+ a = np.arange(27.).reshape((3, 3, 3))
+ b = np.arange(10., 19.).reshape((3, 1, 3))
+ c = inner1d(a, b)
+ assert_array_equal(c, (a * b).sum(-1))
+ c = inner1d(a, b, axis=-1)
+ assert_array_equal(c, (a * b).sum(-1))
+ out = np.zeros_like(c)
+ d = inner1d(a, b, axis=-1, out=out)
+ assert_(d is out)
+ assert_array_equal(d, c)
+ c = inner1d(a, b, axis=0)
+ assert_array_equal(c, (a * b).sum(0))
+ # Sanity checks on innerwt and cumsum.
+ a = np.arange(6).reshape((2, 3))
+ b = np.arange(10, 16).reshape((2, 3))
+ w = np.arange(20, 26).reshape((2, 3))
+ assert_array_equal(umt.innerwt(a, b, w, axis=0),
+ np.sum(a * b * w, axis=0))
+ assert_array_equal(umt.cumsum(a, axis=0), np.cumsum(a, axis=0))
+ assert_array_equal(umt.cumsum(a, axis=-1), np.cumsum(a, axis=-1))
+ out = np.empty_like(a)
+ b = umt.cumsum(a, out=out, axis=0)
+ assert_(out is b)
+ assert_array_equal(b, np.cumsum(a, axis=0))
+ b = umt.cumsum(a, out=out, axis=1)
+ assert_(out is b)
+ assert_array_equal(b, np.cumsum(a, axis=-1))
+ # Check errors.
+ # Cannot pass in both axis and axes.
+ assert_raises(TypeError, inner1d, a, b, axis=0, axes=[0, 0])
+ # Not an integer.
+ assert_raises(TypeError, inner1d, a, b, axis=[0])
+ # more than 1 core dimensions.
+ mm = umt.matrix_multiply
+ assert_raises(TypeError, mm, a, b, axis=1)
+ # Output wrong size in axis.
+ out = np.empty((1, 2, 3), dtype=a.dtype)
+ assert_raises(ValueError, umt.cumsum, a, out=out, axis=0)
+ # Regular ufuncs should not accept axis.
+ assert_raises(TypeError, np.add, 1., 1., axis=0)
+
def test_keepdims_argument(self):
# inner1d signature: '(i),(i)->()'
inner1d = umt.inner1d
@@ -741,7 +785,15 @@ class TestUfunc(object):
d = inner1d(a, b, keepdims=True, out=out)
assert_(d is out)
assert_array_equal(d, c)
- # Now combined with axes.
+ # Now combined with axis and axes.
+ c = inner1d(a, b, axis=-1, keepdims=False)
+ assert_array_equal(c, (a * b).sum(-1, keepdims=False))
+ c = inner1d(a, b, axis=-1, keepdims=True)
+ assert_array_equal(c, (a * b).sum(-1, keepdims=True))
+ c = inner1d(a, b, axis=0, keepdims=False)
+ assert_array_equal(c, (a * b).sum(0, keepdims=False))
+ c = inner1d(a, b, axis=0, keepdims=True)
+ assert_array_equal(c, (a * b).sum(0, keepdims=True))
c = inner1d(a, b, axes=[(-1,), (-1,), ()], keepdims=False)
assert_array_equal(c, (a * b).sum(-1))
c = inner1d(a, b, axes=[(-1,), (-1,), (-1,)], keepdims=True)
@@ -785,10 +837,12 @@ class TestUfunc(object):
w = np.arange(20, 26).reshape((2, 3))
assert_array_equal(umt.innerwt(a, b, w, keepdims=True),
np.sum(a * b * w, axis=-1, keepdims=True))
+ assert_array_equal(umt.innerwt(a, b, w, axis=0, keepdims=True),
+ np.sum(a * b * w, axis=0, keepdims=True))
# Check errors.
# Not a boolean
assert_raises(TypeError, inner1d, a, b, keepdims='true')
- # 1 core dimension only.
+ # More than 1 core dimension, and core output dimensions.
mm = umt.matrix_multiply
assert_raises(TypeError, mm, a, b, keepdims=True)
assert_raises(TypeError, mm, a, b, keepdims=False)
@@ -886,6 +940,11 @@ class TestUfunc(object):
# An output array is required to determine p with signature (n,d)->(p)
assert_raises(ValueError, umt.euclidean_pdist, a)
+ def test_cumsum(self):
+ a = np.arange(10)
+ result = umt.cumsum(a)
+ assert_array_equal(result, a.cumsum())
+
def test_object_logical(self):
a = np.array([3, None, True, False, "test", ""], dtype=object)
assert_equal(np.logical_or(a, None),
diff --git a/numpy/core/tests/test_umath.py b/numpy/core/tests/test_umath.py
index 93ec73094..3c0d1759a 100644
--- a/numpy/core/tests/test_umath.py
+++ b/numpy/core/tests/test_umath.py
@@ -1810,6 +1810,7 @@ class TestSpecialMethods(object):
assert_raises(TypeError, np.multiply, a)
assert_raises(TypeError, np.multiply, a, a, a, a)
assert_raises(TypeError, np.multiply, a, a, sig='a', signature='a')
+ assert_raises(TypeError, ncu_tests.inner1d, a, a, axis=0, axes=[0, 0])
# reduce, positional args
res = np.multiply.reduce(a, 'axis0', 'dtype0', 'out0', 'keep0')
View Lorry Controller Status