# # 编写您自己的ufunc

I have the Power!

— He-Man

## # 创建一个新的ufunc

umath模块是一个计算机生成的C模块，可以创建许多ufunc。它提供了许多如何创建通用函数的示例。使用ufunc机制创建自己的ufunc也不困难。假设您有一个函数，您想要在其输入上逐个元素地操作。通过创建一个新的ufunc，您将获得一个处理的函数

• 广播
• N维循环
• 自动类型转换，内存使用量最少
• 可选的输出数组

>>> logit(0)
-inf
>>> logit(1)
inf
>>> logit(2)
nan
>>> logit(-2)
nan


## # 示例非ufunc扩展名

#include <Python.h>
#include <math.h>

/*
* spammodule.c
* This is the C code for a non-numpy Python extension to
* define the logit function, where logit(p) = log(p/(1-p)).
* This function will not work on numpy arrays automatically.
* numpy.vectorize must be called in python to generate
* a numpy-friendly function.
*
* Details explaining the Python-C API can be found under
* 'Extending and Embedding' and 'Python/C API' at
* docs.python.org .
*/

/* This declares the logit function */
static PyObject* spam_logit(PyObject *self, PyObject *args);

/*
* This tells Python what methods this module has.
*/
static PyMethodDef SpamMethods[] = {
{"logit",
spam_logit,
METH_VARARGS, "compute logit"},
{NULL, NULL, 0, NULL}
};

/*
* This actually defines the logit function for
* input args from Python.
*/

static PyObject* spam_logit(PyObject *self, PyObject *args)
{
double p;

/* This parses the Python argument into a double */
if(!PyArg_ParseTuple(args, "d", &p)) {
return NULL;
}

/* THE ACTUAL LOGIT FUNCTION */
p = p/(1-p);
p = log(p);

/*This builds the answer back into a python object */
return Py_BuildValue("d", p);
}

/* This initiates the module using the above definitions. */
#if PY_VERSION_HEX >= 0x03000000
static struct PyModuleDef moduledef = {
"spam",
NULL,
-1,
SpamMethods,
NULL,
NULL,
NULL,
NULL
};

PyMODINIT_FUNC PyInit_spam(void)
{
PyObject *m;
m = PyModule_Create(&moduledef);
if (!m) {
return NULL;
}
return m;
}
#else
PyMODINIT_FUNC initspam(void)
{
PyObject *m;

m = Py_InitModule("spam", SpamMethods);
if (m == NULL) {
return;
}
}
#endif


'''
setup.py file for spammodule.c

Calling
$python setup.py build_ext --inplace will build the extension library in the current file. Calling$python setup.py build
will build a file that looks like ./build/lib*, where
lib* is a file that begins with lib. The library will
be in this file and end with a C library extension,
such as .so

Calling
$python setup.py install will install the module in your site-packages file. See the distutils section of 'Extending and Embedding the Python Interpreter' at docs.python.org for more information. ''' from distutils.core import setup, Extension module1 = Extension('spam', sources=['spammodule.c'], include_dirs=['/usr/local/lib']) setup(name = 'spam', version='1.0', description='This is my spam package', ext_modules = [module1])  将垃圾邮件模块导入python后，您可以通过spam.logit调用logit。请注意，上面使用的函数不能按原样应用于numpy数组。为此，我们必须在其上调用numpy.vectorize。例如，如果在包含垃圾邮件库或垃圾邮件的文件中打开了python解释器，则可以执行以下命令： >>> import numpy as np >>> import spam >>> spam.logit(0) -inf >>> spam.logit(1) inf >>> spam.logit(0.5) 0.0 >>> x = np.linspace(0,1,10) >>> spam.logit(x) TypeError: only length-1 arrays can be converted to Python scalars >>> f = np.vectorize(spam.logit) >>> f(x) array([ -inf, -2.07944154, -1.25276297, -0.69314718, -0.22314355, 0.22314355, 0.69314718, 1.25276297, 2.07944154, inf])  结果编辑功能并不快！numpy.vectorize只是循环遍历spam.logit。循环在C级完成，但numpy数组不断被解析并重新构建。这很贵。当作者将numpy.vectorize（spam.logit）与下面构造的logit ufuncs进行比较时，logit ufuncs几乎快4倍。当然，取决于功能的性质，可以实现更大或更小的加速。 ## # 一种dtype的NumPy ufunc示例 为简单起见，我们为单个dtype提供了一个ufunc，即'f8'双精度型。与前一节一样，我们首先给出.c文件，然后是用于创建包含ufunc的模块的setup.py文件。 代码中与ufunc的实际计算相对应的位置标有/ * BEGIN main ufunc computation * /和/ * END main ufunc computation * /。这些行之间的代码是必须更改以创建自己的ufunc的主要事物。 #include "Python.h" #include "math.h" #include "numpy/ndarraytypes.h" #include "numpy/ufuncobject.h" #include "numpy/npy_3kcompat.h" /* * single_type_logit.c * This is the C code for creating your own * NumPy ufunc for a logit function. * * In this code we only define the ufunc for * a single dtype. The computations that must * be replaced to create a ufunc for * a different function are marked with BEGIN * and END. * * Details explaining the Python-C API can be found under * 'Extending and Embedding' and 'Python/C API' at * docs.python.org . */ static PyMethodDef LogitMethods[] = { {NULL, NULL, 0, NULL} }; /* The loop definition must precede the PyMODINIT_FUNC. */ static void double_logit(char **args, npy_intp *dimensions, npy_intp* steps, void* data) { npy_intp i; npy_intp n = dimensions[0]; char *in = args[0], *out = args[1]; npy_intp in_step = steps[0], out_step = steps[1]; double tmp; for (i = 0; i < n; i++) { /*BEGIN main ufunc computation*/ tmp = *(double *)in; tmp /= 1-tmp; *((double *)out) = log(tmp); /*END main ufunc computation*/ in += in_step; out += out_step; } } /*This a pointer to the above function*/ PyUFuncGenericFunction funcs[1] = {&double_logit}; /* These are the input and return dtypes of logit.*/ static char types[2] = {NPY_DOUBLE, NPY_DOUBLE}; static void *data[1] = {NULL}; #if PY_VERSION_HEX >= 0x03000000 static struct PyModuleDef moduledef = { PyModuleDef_HEAD_INIT, "npufunc", NULL, -1, LogitMethods, NULL, NULL, NULL, NULL }; PyMODINIT_FUNC PyInit_npufunc(void) { PyObject *m, *logit, *d; m = PyModule_Create(&moduledef); if (!m) { return NULL; } import_array(); import_umath(); logit = PyUFunc_FromFuncAndData(funcs, data, types, 1, 1, 1, PyUFunc_None, "logit", "logit_docstring", 0); d = PyModule_GetDict(m); PyDict_SetItemString(d, "logit", logit); Py_DECREF(logit); return m; } #else PyMODINIT_FUNC initnpufunc(void) { PyObject *m, *logit, *d; m = Py_InitModule("npufunc", LogitMethods); if (m == NULL) { return; } import_array(); import_umath(); logit = PyUFunc_FromFuncAndData(funcs, data, types, 1, 1, 1, PyUFunc_None, "logit", "logit_docstring", 0); d = PyModule_GetDict(m); PyDict_SetItemString(d, "logit", logit); Py_DECREF(logit); } #endif  这是上面代码的setup.py文件。和以前一样，可以通过在命令提示符下调用python setup.py build来构建模块，也可以通过python setup.py install将其安装到site-packages。 ''' setup.py file for logit.c Note that since this is a numpy extension we use numpy.distutils instead of distutils from the python standard library. Calling$python setup.py build_ext --inplace
will build the extension library in the current file.

Calling
$python setup.py build will build a file that looks like ./build/lib*, where lib* is a file that begins with lib. The library will be in this file and end with a C library extension, such as .so Calling$python setup.py install
will install the module in your site-packages file.

See the distutils section of
'Extending and Embedding the Python Interpreter'
at docs.python.org  and the documentation
'''

def configuration(parent_package='', top_path=None):
import numpy
from numpy.distutils.misc_util import Configuration

config = Configuration('npufunc_directory',
parent_package,
top_path)

return config

if __name__ == "__main__":
from numpy.distutils.core import setup
setup(configuration=configuration)


>>> import numpy as np
>>> import npufunc
>>> npufunc.logit(0.5)
0.0
>>> a = np.linspace(0,1,5)
>>> npufunc.logit(a)
array([       -inf, -1.09861229,  0.        ,  1.09861229,         inf])


## # 示例具有多个dtypes的NumPy ufunc

#include "Python.h"
#include "math.h"
#include "numpy/ndarraytypes.h"
#include "numpy/ufuncobject.h"
#include "numpy/halffloat.h"

/*
* multi_type_logit.c
* This is the C code for creating your own
* NumPy ufunc for a logit function.
*
* Each function of the form type_logit defines the
* logit function for a different numpy dtype. Each
* of these functions must be modified when you
* create your own ufunc. The computations that must
* be replaced to create a ufunc for
* a different function are marked with BEGIN
* and END.
*
* Details explaining the Python-C API can be found under
* 'Extending and Embedding' and 'Python/C API' at
* docs.python.org .
*
*/

static PyMethodDef LogitMethods[] = {
{NULL, NULL, 0, NULL}
};

/* The loop definitions must precede the PyMODINIT_FUNC. */

static void long_double_logit(char **args, npy_intp *dimensions,
npy_intp* steps, void* data)
{
npy_intp i;
npy_intp n = dimensions[0];
char *in = args[0], *out=args[1];
npy_intp in_step = steps[0], out_step = steps[1];

long double tmp;

for (i = 0; i < n; i++) {
/*BEGIN main ufunc computation*/
tmp = *(long double *)in;
tmp /= 1-tmp;
*((long double *)out) = logl(tmp);
/*END main ufunc computation*/

in += in_step;
out += out_step;
}
}

static void double_logit(char **args, npy_intp *dimensions,
npy_intp* steps, void* data)
{
npy_intp i;
npy_intp n = dimensions[0];
char *in = args[0], *out = args[1];
npy_intp in_step = steps[0], out_step = steps[1];

double tmp;

for (i = 0; i < n; i++) {
/*BEGIN main ufunc computation*/
tmp = *(double *)in;
tmp /= 1-tmp;
*((double *)out) = log(tmp);
/*END main ufunc computation*/

in += in_step;
out += out_step;
}
}

static void float_logit(char **args, npy_intp *dimensions,
npy_intp* steps, void* data)
{
npy_intp i;
npy_intp n = dimensions[0];
char *in=args[0], *out = args[1];
npy_intp in_step = steps[0], out_step = steps[1];

float tmp;

for (i = 0; i < n; i++) {
/*BEGIN main ufunc computation*/
tmp = *(float *)in;
tmp /= 1-tmp;
*((float *)out) = logf(tmp);
/*END main ufunc computation*/

in += in_step;
out += out_step;
}
}

static void half_float_logit(char **args, npy_intp *dimensions,
npy_intp* steps, void* data)
{
npy_intp i;
npy_intp n = dimensions[0];
char *in = args[0], *out = args[1];
npy_intp in_step = steps[0], out_step = steps[1];

float tmp;

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

/*BEGIN main ufunc computation*/
tmp = *(npy_half *)in;
tmp = npy_half_to_float(tmp);
tmp /= 1-tmp;
tmp = logf(tmp);
*((npy_half *)out) = npy_float_to_half(tmp);
/*END main ufunc computation*/

in += in_step;
out += out_step;
}
}

/*This gives pointers to the above functions*/
PyUFuncGenericFunction funcs[4] = {&half_float_logit,
&float_logit,
&double_logit,
&long_double_logit};

static char types[8] = {NPY_HALF, NPY_HALF,
NPY_FLOAT, NPY_FLOAT,
NPY_DOUBLE,NPY_DOUBLE,
NPY_LONGDOUBLE, NPY_LONGDOUBLE};
static void *data[4] = {NULL, NULL, NULL, NULL};

#if PY_VERSION_HEX >= 0x03000000
static struct PyModuleDef moduledef = {
"npufunc",
NULL,
-1,
LogitMethods,
NULL,
NULL,
NULL,
NULL
};

PyMODINIT_FUNC PyInit_npufunc(void)
{
PyObject *m, *logit, *d;
m = PyModule_Create(&moduledef);
if (!m) {
return NULL;
}

import_array();
import_umath();

logit = PyUFunc_FromFuncAndData(funcs, data, types, 4, 1, 1,
PyUFunc_None, "logit",
"logit_docstring", 0);

d = PyModule_GetDict(m);

PyDict_SetItemString(d, "logit", logit);
Py_DECREF(logit);

return m;
}
#else
PyMODINIT_FUNC initnpufunc(void)
{
PyObject *m, *logit, *d;

m = Py_InitModule("npufunc", LogitMethods);
if (m == NULL) {
return;
}

import_array();
import_umath();

logit = PyUFunc_FromFuncAndData(funcs, data, types, 4, 1, 1,
PyUFunc_None, "logit",
"logit_docstring", 0);

d = PyModule_GetDict(m);

PyDict_SetItemString(d, "logit", logit);
Py_DECREF(logit);
}
#endif


'''
setup.py file for logit.c
Note that since this is a numpy extension
distutils from the python standard library.

Calling
$python setup.py build_ext --inplace will build the extension library in the current file. Calling$python setup.py build
will build a file that looks like ./build/lib*, where
lib* is a file that begins with lib. The library will
be in this file and end with a C library extension,
such as .so

Calling
\$python setup.py install
will install the module in your site-packages file.

See the distutils section of
'Extending and Embedding the Python Interpreter'
at docs.python.org  and the documentation
'''

def configuration(parent_package='', top_path=None):
import numpy
from numpy.distutils.misc_util import Configuration
from numpy.distutils.misc_util import get_info

#Necessary for the half-float d-type.
info = get_info('npymath')

config = Configuration('npufunc_directory',
parent_package,
top_path)
['multi_type_logit.c'],
extra_info=info)

return config

if __name__ == "__main__":
from numpy.distutils.core import setup
setup(configuration=configuration)


>>> import numpy as np
>>> import npufunc
>>> npufunc.logit(0.5)
0.0
>>> a = np.linspace(0,1,5)
>>> npufunc.logit(a)
array([       -inf, -1.09861229,  0.        ,  1.09861229,         inf])


## # 示例具有多个参数/返回值的NumPy ufunc

config.add_extension('npufunc', ['single_type_logit.c'])


config.add_extension('npufunc', ['multi_arg_logit.c'])


C文件如下。生成的ufunc接受两个参数A和B.它返回一个元组，其第一个元素是A * B，第二个元素是logit（A * B）。请注意，它会自动支持广播以及ufunc的所有其他属性。

#include "Python.h"
#include "math.h"
#include "numpy/ndarraytypes.h"
#include "numpy/ufuncobject.h"
#include "numpy/halffloat.h"

/*
* multi_arg_logit.c
* This is the C code for creating your own
* NumPy ufunc for a multiple argument, multiple
* return value ufunc. The places where the
* ufunc computation is carried out are marked
*
* Details explaining the Python-C API can be found under
* 'Extending and Embedding' and 'Python/C API' at
* docs.python.org .
*
*/

static PyMethodDef LogitMethods[] = {
{NULL, NULL, 0, NULL}
};

/* The loop definition must precede the PyMODINIT_FUNC. */

static void double_logitprod(char **args, npy_intp *dimensions,
npy_intp* steps, void* data)
{
npy_intp i;
npy_intp n = dimensions[0];
char *in1 = args[0], *in2 = args[1];
char *out1 = args[2], *out2 = args[3];
npy_intp in1_step = steps[0], in2_step = steps[1];
npy_intp out1_step = steps[2], out2_step = steps[3];

double tmp;

for (i = 0; i < n; i++) {
/*BEGIN main ufunc computation*/
tmp = *(double *)in1;
tmp *= *(double *)in2;
*((double *)out1) = tmp;
*((double *)out2) = log(tmp/(1-tmp));
/*END main ufunc computation*/

in1 += in1_step;
in2 += in2_step;
out1 += out1_step;
out2 += out2_step;
}
}

/*This a pointer to the above function*/
PyUFuncGenericFunction funcs[1] = {&double_logitprod};

/* These are the input and return dtypes of logit.*/

static char types[4] = {NPY_DOUBLE, NPY_DOUBLE,
NPY_DOUBLE, NPY_DOUBLE};

static void *data[1] = {NULL};

#if PY_VERSION_HEX >= 0x03000000
static struct PyModuleDef moduledef = {
"npufunc",
NULL,
-1,
LogitMethods,
NULL,
NULL,
NULL,
NULL
};

PyMODINIT_FUNC PyInit_npufunc(void)
{
PyObject *m, *logit, *d;
m = PyModule_Create(&moduledef);
if (!m) {
return NULL;
}

import_array();
import_umath();

logit = PyUFunc_FromFuncAndData(funcs, data, types, 1, 2, 2,
PyUFunc_None, "logit",
"logit_docstring", 0);

d = PyModule_GetDict(m);

PyDict_SetItemString(d, "logit", logit);
Py_DECREF(logit);

return m;
}
#else
PyMODINIT_FUNC initnpufunc(void)
{
PyObject *m, *logit, *d;

m = Py_InitModule("npufunc", LogitMethods);
if (m == NULL) {
return;
}

import_array();
import_umath();

logit = PyUFunc_FromFuncAndData(funcs, data, types, 1, 2, 2,
PyUFunc_None, "logit",
"logit_docstring", 0);

d = PyModule_GetDict(m);

PyDict_SetItemString(d, "logit", logit);
Py_DECREF(logit);
}
#endif


## # 示例带有结构化数组dtype参数的NumPy ufunc

config.add_extension('npufunc', ['single_type_logit.c'])


config.add_extension('npufunc', ['add_triplet.c'])


C文件如下。

#include "Python.h"
#include "math.h"
#include "numpy/ndarraytypes.h"
#include "numpy/ufuncobject.h"
#include "numpy/npy_3kcompat.h"

/*
* 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;
}
}

/* This a pointer to the above function */

/* These are the input and return dtypes of add_uint64_triplet. */
static char types[3] = {NPY_UINT64, NPY_UINT64, NPY_UINT64};

static void *data[1] = {NULL};

#if defined(NPY_PY3K)
static struct PyModuleDef moduledef = {
"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 *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();

/* Create a new ufunc object */
add_triplet = PyUFunc_FromFuncAndData(NULL, NULL, NULL, 0, 2, 1,

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;

/* Register ufunc for structured dtype */
dtype,
dtypes,
NULL);

d = PyModule_GetDict(m);

#if defined(NPY_PY3K)
return m;
#endif
}


static PyUFuncGenericFunction atan2_functions[] = {
PyUFunc_ff_f, PyUFunc_dd_d,
PyUFunc_gg_g, PyUFunc_OO_O_method};
static void* atan2_data[] = {
(void *)atan2f,(void *) atan2,
(void *)atan2l,(void *)"arctan2"};
static char atan2_signatures[] = {
NPY_FLOAT, NPY_FLOAT, NPY_FLOAT,
NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE,
NPY_LONGDOUBLE, NPY_LONGDOUBLE, NPY_LONGDOUBLE
NPY_OBJECT, NPY_OBJECT, NPY_OBJECT};
...
/* in the module initialization code */
PyObject *f, *dict, *module;
...
dict = PyModule_GetDict(module);
...
f = PyUFunc_FromFuncAndData(atan2_functions,
atan2_data, atan2_signatures, 4, 2, 1,
PyUFunc_None, "arctan2",
"a safe and correct arctan(x1/x2)", 0);
PyDict_SetItemString(dict, "arctan2", f);
Py_DECREF(f);
...