summaryrefslogtreecommitdiff
path: root/numpy
diff options
context:
space:
mode:
Diffstat (limited to 'numpy')
-rw-r--r--numpy/core/code_generators/numpy_api.py1
-rw-r--r--numpy/core/include/numpy/ufuncobject.h2
-rw-r--r--numpy/core/setup.py7
-rw-r--r--numpy/core/src/umath/struct_ufunc_test.c.src122
-rw-r--r--numpy/core/src/umath/test_rational.c.src33
-rw-r--r--numpy/core/src/umath/ufunc_object.c163
-rw-r--r--numpy/core/src/umath/ufunc_type_resolution.c47
-rw-r--r--numpy/core/tests/test_ufunc.py20
8 files changed, 375 insertions, 20 deletions
diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py
index 07a87f98d..152d0f948 100644
--- a/numpy/core/code_generators/numpy_api.py
+++ b/numpy/core/code_generators/numpy_api.py
@@ -384,6 +384,7 @@ ufunc_funcs_api = {
# End 1.6 API
'PyUFunc_DefaultTypeResolver': 39,
'PyUFunc_ValidateCasting': 40,
+ 'PyUFunc_RegisterLoopForDescr': 41,
}
# List of all the dicts which define the C API
diff --git a/numpy/core/include/numpy/ufuncobject.h b/numpy/core/include/numpy/ufuncobject.h
index 686d12c38..75611426c 100644
--- a/numpy/core/include/numpy/ufuncobject.h
+++ b/numpy/core/include/numpy/ufuncobject.h
@@ -319,6 +319,8 @@ typedef struct _loop1d_info {
void *data;
int *arg_types;
struct _loop1d_info *next;
+ int nargs;
+ PyArray_Descr **arg_dtypes;
} PyUFunc_Loop1d;
diff --git a/numpy/core/setup.py b/numpy/core/setup.py
index 3b08d6edd..437f75c7b 100644
--- a/numpy/core/setup.py
+++ b/numpy/core/setup.py
@@ -929,6 +929,13 @@ def configuration(parent_package='',top_path=None):
sources = [join('src','umath', 'test_rational.c.src')])
#######################################################################
+ # struct_ufunc_test module #
+ #######################################################################
+
+ config.add_extension('struct_ufunc_test',
+ sources = [join('src','umath', 'struct_ufunc_test.c.src')])
+
+ #######################################################################
# multiarray_tests module #
#######################################################################
diff --git a/numpy/core/src/umath/struct_ufunc_test.c.src b/numpy/core/src/umath/struct_ufunc_test.c.src
new file mode 100644
index 000000000..4bd24559f
--- /dev/null
+++ b/numpy/core/src/umath/struct_ufunc_test.c.src
@@ -0,0 +1,122 @@
+#include "Python.h"
+#include "math.h"
+#include "numpy/ndarraytypes.h"
+#include "numpy/ufuncobject.h"
+#include "numpy/npy_3kcompat.h"
+
+
+/*
+ * struct_ufunc_test.c
+ * This is the C code for creating your own
+ * Numpy ufunc for a structured array dtype.
+ *
+ * Details explaining the Python-C API can be found under
+ * 'Extending and Embedding' and 'Python/C API' at
+ * docs.python.org .
+ */
+
+static PyMethodDef StructUfuncTestMethods[] = {
+ {NULL, NULL, 0, NULL}
+};
+
+/* The loop definition must precede the PyMODINIT_FUNC. */
+
+static void add_uint64_triplet(char **args, npy_intp *dimensions,
+ npy_intp* steps, void* data)
+{
+ npy_intp i;
+ npy_intp is1=steps[0];
+ npy_intp is2=steps[1];
+ npy_intp os=steps[2];
+ npy_intp n=dimensions[0];
+ uint64_t *x, *y, *z;
+
+ char *i1=args[0];
+ char *i2=args[1];
+ char *op=args[2];
+
+ for (i = 0; i < n; i++) {
+
+ x = (uint64_t*)i1;
+ y = (uint64_t*)i2;
+ z = (uint64_t*)op;
+
+ z[0] = x[0] + y[0];
+ z[1] = x[1] + y[1];
+ z[2] = x[2] + y[2];
+
+ i1 += is1;
+ i2 += is2;
+ op += os;
+ }
+}
+
+#if defined(NPY_PY3K)
+static struct PyModuleDef moduledef = {
+ PyModuleDef_HEAD_INIT,
+ "struct_ufunc_test",
+ NULL,
+ -1,
+ StructUfuncTestMethods,
+ NULL,
+ NULL,
+ NULL,
+ NULL
+};
+#endif
+
+#if defined(NPY_PY3K)
+PyMODINIT_FUNC PyInit_struct_ufunc_test(void)
+#else
+PyMODINIT_FUNC initstruct_ufunc_test(void)
+#endif
+{
+ PyObject *m, *add_triplet, *d;
+ PyObject *dtype_dict;
+ PyArray_Descr *dtype;
+ PyArray_Descr *dtypes[3];
+
+#if defined(NPY_PY3K)
+ m = PyModule_Create(&moduledef);
+#else
+ m = Py_InitModule("struct_ufunc_test", StructUfuncTestMethods);
+#endif
+
+ if (m == NULL) {
+#if defined(NPY_PY3K)
+ return NULL;
+#else
+ return;
+#endif
+ }
+
+ import_array();
+ import_umath();
+
+ add_triplet = PyUFunc_FromFuncAndData(NULL, NULL, NULL, 0, 2, 1,
+ PyUFunc_None, "add_triplet",
+ "add_triplet_docstring", 0);
+
+ dtype_dict = Py_BuildValue("[(s, s), (s, s), (s, s)]",
+ "f0", "u8", "f1", "u8", "f2", "u8");
+ PyArray_DescrConverter(dtype_dict, &dtype);
+ Py_DECREF(dtype_dict);
+
+ dtypes[0] = dtype;
+ dtypes[1] = dtype;
+ dtypes[2] = dtype;
+
+ PyUFunc_RegisterLoopForDescr(add_triplet,
+ dtype,
+ &add_uint64_triplet,
+ dtypes,
+ NULL);
+
+ d = PyModule_GetDict(m);
+
+ PyDict_SetItemString(d, "add_triplet", add_triplet);
+ Py_DECREF(add_triplet);
+#if defined(NPY_PY3K)
+ return m;
+#endif
+}
diff --git a/numpy/core/src/umath/test_rational.c.src b/numpy/core/src/umath/test_rational.c.src
index 6e1f6a2aa..f9153b87c 100644
--- a/numpy/core/src/umath/test_rational.c.src
+++ b/numpy/core/src/umath/test_rational.c.src
@@ -1093,6 +1093,21 @@ rational_ufunc_test_add(char** args, npy_intp* dimensions,
}
+static void
+rational_ufunc_test_add_rationals(char** args, npy_intp* dimensions,
+ npy_intp* steps, void* data) {
+ npy_intp is0 = steps[0], is1 = steps[1], os = steps[2], n = *dimensions;
+ char *i0 = args[0], *i1 = args[1], *o = args[2];
+ int k;
+ for (k = 0; k < n; k++) {
+ rational x = *(rational*)i0;
+ rational y = *(rational*)i1;
+ *(rational*)o = rational_add(x, y);
+ i0 += is0; i1 += is1; o += os;
+ }
+}
+
+
PyMethodDef module_methods[] = {
{0} /* sentinel */
};
@@ -1307,6 +1322,24 @@ PyMODINIT_FUNC inittest_rational(void) {
PyModule_AddObject(m,"test_add",(PyObject*)ufunc);
}
+ /* Create test ufunc with rational types using RegisterLoopForDescr */
+ {
+ PyObject* ufunc = PyUFunc_FromFuncAndData(0,0,0,0,2,1,
+ PyUFunc_None,(char*)"test_add_rationals",
+ (char*)"add two matrices of rationals and return rational matrix",0);
+ if (!ufunc) {
+ goto fail;
+ }
+ PyArray_Descr* types[3] = {&npyrational_descr,
+ &npyrational_descr,
+ &npyrational_descr};
+ if (PyUFunc_RegisterLoopForDescr((PyUFuncObject*)ufunc, &npyrational_descr,
+ rational_ufunc_test_add_rationals, types, 0) < 0) {
+ goto fail;
+ }
+ PyModule_AddObject(m,"test_add_rationals",(PyObject*)ufunc);
+ }
+
/* Create numerator and denominator ufuncs */
#define NEW_UNARY_UFUNC(name,type,doc) { \
int types[2] = {npy_rational,type}; \
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index a3a164731..8a028096a 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -724,8 +724,9 @@ static int get_ufunc_arguments(PyUFuncObject *ufunc,
PyObject *obj, *context;
PyObject *str_key_obj = NULL;
char *ufunc_name;
+ int type_num;
- int any_flexible = 0, any_object = 0;
+ int any_flexible = 0, any_object = 0, any_flexible_userloops = 0;
ufunc_name = ufunc->name ? ufunc->name : "<unnamed ufunc>";
@@ -764,23 +765,55 @@ static int get_ufunc_arguments(PyUFuncObject *ufunc,
if (out_op[i] == NULL) {
return -1;
}
+
+ type_num = PyArray_DESCR(out_op[i])->type_num;
if (!any_flexible &&
- PyTypeNum_ISFLEXIBLE(PyArray_DESCR(out_op[i])->type_num)) {
+ PyTypeNum_ISFLEXIBLE(type_num)) {
any_flexible = 1;
}
if (!any_object &&
- PyTypeNum_ISOBJECT(PyArray_DESCR(out_op[i])->type_num)) {
+ PyTypeNum_ISOBJECT(type_num)) {
any_object = 1;
}
+
+ /*
+ * If any operand is a flexible dtype, check to see if any
+ * struct dtype ufuncs are registered. A ufunc has been registered
+ * for a struct dtype if ufunc's arg_dtypes array is not NULL.
+ */
+ if (PyTypeNum_ISFLEXIBLE(type_num) &&
+ !any_flexible_userloops &&
+ ufunc->userloops != NULL) {
+ PyUFunc_Loop1d *funcdata;
+ PyObject *key, *obj;
+ key = PyInt_FromLong(type_num);
+ if (key == NULL) {
+ continue;
+ }
+ obj = PyDict_GetItem(ufunc->userloops, key);
+ Py_DECREF(key);
+ if (obj == NULL) {
+ continue;
+ }
+ funcdata = (PyUFunc_Loop1d *)NpyCapsule_AsVoidPtr(obj);
+ while (funcdata != NULL) {
+ if (funcdata->arg_dtypes != NULL) {
+ any_flexible_userloops = 1;
+ break;
+ }
+ funcdata = funcdata->next;
+ }
+ }
}
/*
* Indicate not implemented if there are flexible objects (structured
- * type or string) but no object types.
+ * type or string) but no object types and no registered struct
+ * dtype ufuncs.
*
* Not sure - adding this increased to 246 errors, 150 failures.
*/
- if (any_flexible && !any_object) {
+ if (any_flexible && !any_flexible_userloops && !any_object) {
return -2;
}
@@ -4356,9 +4389,19 @@ cmp_arg_types(int *arg1, int *arg2, int n)
static NPY_INLINE void
_free_loop1d_list(PyUFunc_Loop1d *data)
{
+ int i;
+
while (data != NULL) {
PyUFunc_Loop1d *next = data->next;
PyArray_free(data->arg_types);
+
+ if (data->arg_dtypes != NULL) {
+ for (i = 0; i < data->nargs; i++) {
+ Py_DECREF(data->arg_dtypes[i]);
+ }
+ PyArray_free(data->arg_dtypes);
+ }
+
PyArray_free(data);
data = next;
}
@@ -4381,6 +4424,112 @@ _loop1d_list_free(void *ptr)
#endif
+/*
+ * This function allows the user to register a 1-d loop with an already
+ * created ufunc. This function is similar to RegisterLoopForType except
+ * that it allows a 1-d loop to be registered with PyArray_Descr objects
+ * instead of dtype type num values. This allows a 1-d loop to be registered
+ * for a structured array dtype or a custom dtype. The ufunc is called
+ * whenever any of it's input arguments match the user_dtype argument.
+ * ufunc - ufunc object created from call to PyUFunc_FromFuncAndData
+ * user_dtype - dtype that ufunc will be registered with
+ * function - 1-d loop function pointer
+ * arg_dtypes - array of dtype objects describing the ufunc operands
+ * data - arbitrary data pointer passed in to loop function
+ */
+/*UFUNC_API*/
+NPY_NO_EXPORT int
+PyUFunc_RegisterLoopForDescr(PyUFuncObject *ufunc,
+ PyArray_Descr *user_dtype,
+ PyUFuncGenericFunction function,
+ PyArray_Descr **arg_dtypes,
+ void *data)
+{
+ int i;
+ int result = 0;
+ int *arg_typenums;
+ PyObject *key, *cobj;
+
+ if (user_dtype == NULL) {
+ PyErr_SetString(PyExc_TypeError,
+ "unknown user defined struct dtype");
+ return -1;
+ }
+
+ key = PyInt_FromLong((long) user_dtype->type_num);
+ if (key == NULL) {
+ return -1;
+ }
+
+ arg_typenums = PyArray_malloc(ufunc->nargs * sizeof(int));
+ if (arg_typenums == NULL) {
+ PyErr_NoMemory();
+ return -1;
+ }
+ if (arg_dtypes != NULL) {
+ for (i = 0; i < ufunc->nargs; i++) {
+ arg_typenums[i] = arg_dtypes[i]->type_num;
+ }
+ }
+ else {
+ for (i = 0; i < ufunc->nargs; i++) {
+ arg_typenums[i] = user_dtype->type_num;
+ }
+ }
+
+ result = PyUFunc_RegisterLoopForType(ufunc, user_dtype->type_num,
+ function, arg_typenums, data);
+
+ if (result == 0) {
+ cobj = PyDict_GetItem(ufunc->userloops, key);
+ if (cobj == NULL) {
+ PyErr_SetString(PyExc_KeyError,
+ "userloop for user dtype not found");
+ result = -1;
+ }
+ else {
+ PyUFunc_Loop1d *current, *prev = NULL;
+ int cmp = 1;
+ current = (PyUFunc_Loop1d *)NpyCapsule_AsVoidPtr(cobj);
+ while (current != NULL) {
+ cmp = cmp_arg_types(current->arg_types,
+ arg_typenums, ufunc->nargs);
+ if (cmp >= 0 && current->arg_dtypes == NULL) {
+ break;
+ }
+ prev = current;
+ current = current->next;
+ }
+ if (cmp == 0 && current->arg_dtypes == NULL) {
+ current->arg_dtypes = PyArray_malloc(ufunc->nargs *
+ sizeof(PyArray_Descr*));
+ if (arg_dtypes != NULL) {
+ for (i = 0; i < ufunc->nargs; i++) {
+ current->arg_dtypes[i] = arg_dtypes[i];
+ Py_INCREF(current->arg_dtypes[i]);
+ }
+ }
+ else {
+ for (i = 0; i < ufunc->nargs; i++) {
+ current->arg_dtypes[i] = user_dtype;
+ Py_INCREF(current->arg_dtypes[i]);
+ }
+ }
+ current->nargs = ufunc->nargs;
+ }
+ else {
+ result = -1;
+ }
+ }
+ }
+
+ PyArray_free(arg_typenums);
+
+ Py_DECREF(key);
+
+ return result;
+}
+
/*UFUNC_API*/
NPY_NO_EXPORT int
PyUFunc_RegisterLoopForType(PyUFuncObject *ufunc,
@@ -4396,7 +4545,7 @@ PyUFunc_RegisterLoopForType(PyUFuncObject *ufunc,
int *newtypes=NULL;
descr=PyArray_DescrFromType(usertype);
- if ((usertype < NPY_USERDEF) || (descr==NULL)) {
+ if ((usertype < NPY_USERDEF && usertype != NPY_VOID) || (descr==NULL)) {
PyErr_SetString(PyExc_TypeError, "unknown user-defined type");
return -1;
}
@@ -4432,6 +4581,8 @@ PyUFunc_RegisterLoopForType(PyUFuncObject *ufunc,
funcdata->arg_types = newtypes;
funcdata->data = data;
funcdata->next = NULL;
+ funcdata->arg_dtypes = NULL;
+ funcdata->nargs = 0;
/* Get entry for this user-defined type*/
cobj = PyDict_GetItem(ufunc->userloops, key);
diff --git a/numpy/core/src/umath/ufunc_type_resolution.c b/numpy/core/src/umath/ufunc_type_resolution.c
index 7e3fb5633..6a16692b0 100644
--- a/numpy/core/src/umath/ufunc_type_resolution.c
+++ b/numpy/core/src/umath/ufunc_type_resolution.c
@@ -1188,7 +1188,8 @@ find_userloop(PyUFuncObject *ufunc,
}
type_num = dtypes[i]->type_num;
- if (type_num != last_userdef && PyTypeNum_ISUSERDEF(type_num)) {
+ if (type_num != last_userdef &&
+ (PyTypeNum_ISUSERDEF(type_num) || type_num == NPY_VOID)) {
PyObject *key, *obj;
last_userdef = type_num;
@@ -1436,7 +1437,7 @@ ufunc_loop_matches(PyUFuncObject *self,
NPY_CASTING output_casting,
int any_object,
int use_min_scalar,
- int *types,
+ int *types, PyArray_Descr **dtypes,
int *out_no_castable_output,
char *out_err_src_typecode,
char *out_err_dst_typecode)
@@ -1463,7 +1464,18 @@ ufunc_loop_matches(PyUFuncObject *self,
return 0;
}
- tmp = PyArray_DescrFromType(types[i]);
+ /*
+ * If type num is NPY_VOID and struct dtypes have been passed in,
+ * use struct dtype object. Otherwise create new dtype object
+ * from type num.
+ */
+ if (types[i] == NPY_VOID && dtypes != NULL) {
+ tmp = dtypes[i];
+ Py_INCREF(tmp);
+ }
+ else {
+ tmp = PyArray_DescrFromType(types[i]);
+ }
if (tmp == NULL) {
return -1;
}
@@ -1525,7 +1537,7 @@ ufunc_loop_matches(PyUFuncObject *self,
static int
set_ufunc_loop_data_types(PyUFuncObject *self, PyArrayObject **op,
PyArray_Descr **out_dtypes,
- int *type_nums)
+ int *type_nums, PyArray_Descr **dtypes)
{
int i, nin = self->nin, nop = nin + self->nout;
@@ -1536,11 +1548,16 @@ set_ufunc_loop_data_types(PyUFuncObject *self, PyArrayObject **op,
* instead of creating a new one, similarly to preserve metadata.
**/
for (i = 0; i < nop; ++i) {
+ if (dtypes != NULL) {
+ out_dtypes[i] = dtypes[i];
+ Py_XINCREF(out_dtypes[i]);
/*
* Copy the dtype from 'op' if the type_num matches,
* to preserve metadata.
*/
- if (op[i] != NULL && PyArray_DESCR(op[i])->type_num == type_nums[i]) {
+ }
+ else if (op[i] != NULL &&
+ PyArray_DESCR(op[i])->type_num == type_nums[i]) {
out_dtypes[i] = ensure_dtype_nbo(PyArray_DESCR(op[i]));
Py_XINCREF(out_dtypes[i]);
}
@@ -1603,7 +1620,8 @@ linear_search_userloop_type_resolver(PyUFuncObject *self,
}
type_num = PyArray_DESCR(op[i])->type_num;
- if (type_num != last_userdef && PyTypeNum_ISUSERDEF(type_num)) {
+ if (type_num != last_userdef &&
+ (PyTypeNum_ISUSERDEF(type_num) || type_num == NPY_VOID)) {
PyObject *key, *obj;
last_userdef = type_num;
@@ -1623,7 +1641,7 @@ linear_search_userloop_type_resolver(PyUFuncObject *self,
switch (ufunc_loop_matches(self, op,
input_casting, output_casting,
any_object, use_min_scalar,
- types,
+ types, funcdata->arg_dtypes,
out_no_castable_output, out_err_src_typecode,
out_err_dst_typecode)) {
/* Error */
@@ -1631,7 +1649,7 @@ linear_search_userloop_type_resolver(PyUFuncObject *self,
return -1;
/* Found a match */
case 1:
- set_ufunc_loop_data_types(self, op, out_dtype, types);
+ set_ufunc_loop_data_types(self, op, out_dtype, types, funcdata->arg_dtypes);
return 1;
}
@@ -1707,12 +1725,13 @@ type_tuple_userloop_type_resolver(PyUFuncObject *self,
switch (ufunc_loop_matches(self, op,
casting, casting,
any_object, use_min_scalar,
- types,
+ types, NULL,
&no_castable_output, &err_src_typecode,
&err_dst_typecode)) {
/* It works */
case 1:
- set_ufunc_loop_data_types(self, op, out_dtype, types);
+ set_ufunc_loop_data_types(self, op,
+ out_dtype, types, NULL);
return 1;
/* Didn't match */
case 0:
@@ -1875,7 +1894,7 @@ linear_search_type_resolver(PyUFuncObject *self,
switch (ufunc_loop_matches(self, op,
input_casting, output_casting,
any_object, use_min_scalar,
- types,
+ types, NULL,
&no_castable_output, &err_src_typecode,
&err_dst_typecode)) {
/* Error */
@@ -1883,7 +1902,7 @@ linear_search_type_resolver(PyUFuncObject *self,
return -1;
/* Found a match */
case 1:
- set_ufunc_loop_data_types(self, op, out_dtype, types);
+ set_ufunc_loop_data_types(self, op, out_dtype, types, NULL);
return 0;
}
}
@@ -2081,7 +2100,7 @@ type_tuple_type_resolver(PyUFuncObject *self,
switch (ufunc_loop_matches(self, op,
casting, casting,
any_object, use_min_scalar,
- types,
+ types, NULL,
&no_castable_output, &err_src_typecode,
&err_dst_typecode)) {
/* Error */
@@ -2089,7 +2108,7 @@ type_tuple_type_resolver(PyUFuncObject *self,
return -1;
/* It worked */
case 1:
- set_ufunc_loop_data_types(self, op, out_dtype, types);
+ set_ufunc_loop_data_types(self, op, out_dtype, types, NULL);
return 0;
/* Didn't work */
case 0:
diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py
index 3005da8da..ad489124e 100644
--- a/numpy/core/tests/test_ufunc.py
+++ b/numpy/core/tests/test_ufunc.py
@@ -813,6 +813,26 @@ class TestUfunc(TestCase):
assert_equal(a, 3)
opflag_tests.inplace_add(a, [3, 4])
assert_equal(a, 10)
+
+ def test_struct_ufunc(self):
+ import numpy.core.struct_ufunc_test as struct_ufunc
+
+ a = np.array([(1,2,3)], dtype='u8,u8,u8')
+ b = np.array([(1,2,3)], dtype='u8,u8,u8')
+
+ result = struct_ufunc.add_triplet(a, b)
+ assert_equal(result, np.array([(2, 4, 6)], dtype='u8,u8,u8'))
+
+ def test_custom_ufunc(self):
+ a = np.array([rational(1,2), rational(1,3), rational(1,4)],
+ dtype=rational);
+ b = np.array([rational(1,2), rational(1,3), rational(1,4)],
+ dtype=rational);
+
+ result = test_add_rationals(a, b)
+ expected = np.array([rational(1), rational(2,3), rational(1,2)],
+ dtype=rational);
+ assert_equal(result, expected);
if __name__ == "__main__":
run_module_suite()
View Lorry Controller Status