Writing Your Own UFunc and Generalized UFunc For a Custom Data Type

The following are several examples for creating your own ufuncs and generalized ufuncs for custom data types.

All the examples use the custom dtype 'Rational' located in the numpy-dtypes repository on github.com

Extending Existing UFunc for Custom DType

The first example shows how to extend the existing 'add' ufunc for the Rational dtype. The add ufunc is extended

for Rational dtypes using Rational's 'rational_add' function which takes two rational numbers and returns a rational

object representing the sum of those two rational numbers:

static NPY_INLINE rational

rational_add(rational x, rational y) {

return make_rational_fast((int64_t)x.n*d(y) + (int64_t)d(x)*y.n, (int64_t)d(x)*d(y));

}

1. A 1-d loop function is created which loops over each pair of elements from two 1-d arrays of rationals and

calls rational_add for each pair:

void rational_ufunc_add(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;

}

}

The loop function must have the exact signature as above. The function parameters are:

char **args - array of pointers to the actual data for the input and output arrays. In this example there are

three pointers: two pointers pointing to blocks of memory for the two input arrays, and one pointer pointing to

the block of memory for the output array. The result of rational_add should be stored in the output array.

dimensions - a pointer to the size of the dimension over which this function is looping

steps - a pointer to the number of bytes to jump to get to the next element in this dimension

for each of the input and output arguments

data - arbitrary data (extra arguments, function names, etc.) that can be stored with the ufunc

and will be passed in when it is called

The Rational dtype has an example of a C macro which can be used to generate create the above function for

different rational ufuncs:

#define BINARY_UFUNC(name,intype0,intype1,outtype,exp) \

void name(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++) { \

intype0 x = *(intype0*)i0; \

intype1 y = *(intype1*)i1; \

*(outtype*)o = exp; \

i0 += is0; i1 += is1; o += os; \

} \

}

#define RATIONAL_BINARY_UFUNC(name, type, exp) BINARY_UFUNC(rational_ufunc_##name, rational, rational, type, exp)

which can be used like so:

RATIONAL_BINARY_UFUNC(add, rational, rational_add(x,y))

with the following arguments:

- name suffix of 1-d loop function (the generated loop function will have the name rational_ufunc_<name>'

- output type

- expression to calculate the output value for each pair of input elements. In this example the expression

is a call to the function rational_add.

2. In the 'initrational' function used to initialize the Rational dtype with numpy, a PyUFuncObject is obtained for

the existing 'add' ufunc in the numpy module:

PyUFuncObject* ufunc = (PyUFuncObject*)PyObject_GetAttrString(numpy,"add");

3. The 1-d loop function is registered using the PyUFuncObject obtained in step 2:

int types[] = {npy_rational,npy_rational,npy_rational};

if (PyUFunc_RegisterLoopForType((PyUFuncObject*)ufunc,npy_rational,rational_ufunc_biggest,types,0) < 0) {

return;

}

The function parameters are:

- pointer to PyUFuncObject obtained in step 2

- custom rational dtype id (obtained when dtype is registered with call to PyArray_RegisterDataType)

- 1-d loop function

- array of input and output type ids (in this case two input rational types and one output rational type)

- pointer to arbitrary data that will be passed to 1-d loop function

4. Steps 2-3 can also be accomplished by using a c MACRO similar to the one provided with Rational:

#define REGISTER_UFUNC(name,...) { \

PyUFuncObject* ufunc = (PyUFuncObject*)PyObject_GetAttrString(numpy,#name); \

if (!ufunc) { \

return; \

} \

int _types[] = __VA_ARGS__; \

if (sizeof(_types)/sizeof(int)!=ufunc->nargs) { \

PyErr_Format(PyExc_AssertionError,"ufunc %s takes %d arguments, our loop takes %ld",#name,ufunc->nargs,sizeof(_types)/sizeof(int)); \

return; \

} \

if (PyUFunc_RegisterLoopForType((PyUFuncObject*)ufunc,npy_rational,rational_ufunc_##name,_types,0)<0) { \

return; \

} \

}

#define REGISTER_UFUNC_BINARY_RATIONAL(name) REGISTER_UFUNC(name,{npy_rational,npy_rational,npy_rational})

REGISTER_UFUNC_BINARY_RATIONAL(add)

Creating New UFunc for Custom DType

The next example shows how to create a new ufunc for the Rational dtype. The ufunc example is called 'numerator'

and generates an array of numerator values based rational numbers from the input array.

1. A 1-d loop function is created as before which takes the numerator value from each element of the input array

and stores it in the output array:

void rational_ufunc_numerator(char** args, npy_intp* dimensions, npy_intp* steps, void* data) {

npy_intp is = steps[0], os = steps[1], n = *dimensions;

char *i = args[0], *o = args[1];

int k;

for (k = 0; k < n; k++) {

rational x = *(rational*)i;

*(int64_t*)o = x.n;

i += is; o += os;

}

}

You can also use the c MACRO provided in Rational for generating the above function:

#define UNARY_UFUNC(name,type,exp) \

void rational_ufunc_##name(char** args, npy_intp* dimensions, npy_intp* steps, void* data) { \

npy_intp is = steps[0], os = steps[1], n = *dimensions; \

char *i = args[0], *o = args[1]; \

int k; \

for (k = 0; k < n; k++) { \

rational x = *(rational*)i; \

*(type*)o = exp; \

i += is; o += os; \

} \

}

UNARY_UFUNC(numerator,int64_t,x.n)

2. In the 'initrational' function used to initialize the Rational dtype with numpy, a new PyUFuncObject is created

for the new 'numerator' ufunc using the PyUFunc_FromFuncAndData function:

PyObject* ufunc = PyUFunc_FromFuncAndData(0,0,0,0,1,1,PyUFunc_None,(char*)"numerator",(char*)"rational number numerator",0);

In this use case, the first four parameters should be set to zero since we're creating a new ufunc instead of

extending an existing one. The rest of the parameters:

- number of inputs to function that the loop function calls for each pair of elements

- number of outputs of loop function

- name of the ufunc

- documentation string describing the ufunc

- unused; present for backwards compatibility

3. The 1-d loop function is registered using the loop function and the PyUFuncObject created in step 2:

int _types[] = {npy_rational,NPY_INT64};

if (PyUFunc_RegisterLoopForType((PyUFuncObject*)ufunc,npy_rational,rational_ufunc_numerator,_types,0)<0) {

return;

}

4. Finally, a function called 'numerator' is added to the rational module which will call the numerator ufunc:

PyModule_AddObject(m,"numerator",(PyObject*)ufunc);

5. Steps 2-4 can also be accomplished by using a c MACRO similar to the one provided with Rational:

#define NEW_UNARY_UFUNC(name,type,doc) { \

PyObject* ufunc = PyUFunc_FromFuncAndData(0,0,0,0,1,1,PyUFunc_None,(char*)#name,(char*)doc,0); \

if (!ufunc) { \

return; \

} \

int types[2] = {npy_rational,type}; \

if (PyUFunc_RegisterLoopForType((PyUFuncObject*)ufunc,npy_rational,rational_ufunc_##name,types,0)<0) { \

return; \

} \

PyModule_AddObject(m,#name,(PyObject*)ufunc); \

}

NEW_UNARY_UFUNC(numerator,NPY_INT64,"rational number numerator");