编写你自己的 ufunc #
创建一个新的通用函数#
在阅读本文之前,阅读/浏览扩展和嵌入 Python 解释器第 1 节以及如何扩展 NumPy中的教程可能有助于熟悉 Python 的 C 扩展的基础知识
umath 模块是计算机生成的 C 模块,可创建许多 ufunc。它提供了大量关于如何创建通用函数的示例。创建您自己的使用 ufunc 机制的 ufunc 也不困难。假设您有一个函数,您想要对其输入逐个元素进行操作。通过创建一个新的 ufunc,您将获得一个处理函数
广播
N 维循环
自动类型转换,占用内存最少
可选输出数组
创建自己的ufunc并不困难。所需要的只是为您想要支持的每种数据类型建立一个一维循环。每个一维循环必须有一个特定的签名,并且只能使用固定大小数据类型的 ufunc。下面给出了用于创建新的 ufunc 以处理内置数据类型的函数调用。使用不同的机制为用户定义的数据类型注册 ufunc。
在接下来的几节中,我们提供了示例代码,可以轻松修改这些代码以创建您自己的 ufunc。这些示例依次是 logit 函数的更完整或更复杂的版本,logit 函数是统计建模中的常见函数。 Logit 也很有趣,因为由于 IEEE 标准(特别是 IEEE 754)的魔力,下面创建的所有 logit 函数都会自动具有以下行为。
>>> logit(0)
-inf
>>> logit(1)
inf
>>> logit(2)
nan
>>> logit(-2)
nan
这非常棒,因为函数编写者不必手动传播 infs 或 nan。
非 ufunc 扩展示例#
为了对读者进行比较和一般启发,我们提供了一个logit
不使用 numpy 的 C 扩展的简单实现。
为此,我们需要两个文件。第一个是包含实际代码的 C 文件,第二个是setup.py
用于创建模块的文件。
#define PY_SSIZE_T_CLEAN #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. * See the Python-C API for more information. */ 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. */ static struct PyModuleDef moduledef = { PyModuleDef_HEAD_INIT, "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; }
要使用,请将和放在
同一文件夹中。然后将构建要导入的模块,或者将模块安装到您的 site-packages 目录。setup.py file
setup.py
spammodule.c
python setup.py build
python setup.py install
''' 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])
将 spam 模块导入 python 后,您可以通过 调用 logit spam.logit
。请注意,上面使用的函数不能按原样应用于 numpy 数组。为此,我们必须呼吁numpy.vectorize
它。例如,如果在包含spam库的文件中打开python解释器或者已经安装了spam,则可以执行以下命令:
>>> 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])
生成的 Logit 函数并不快!numpy.vectorize
只是循环spam.logit
。循环是在 C 级别完成的,但 numpy 数组不断被解析和构建备份。这很贵。当作者numpy.vectorize(spam.logit)
与下面构建的 logit ufunc 进行比较时,logit ufunc 几乎快了 4 倍。当然,根据函数的性质,更大或更小的加速是可能的。
一种 dtype 的 NumPy ufunc 示例#
为简单起见,我们为单个数据类型提供 ufunc,即'f8'
double
.与上一节一样,我们首先给出文件.c
,然后setup.py
给出用于创建包含 ufunc 的模块的文件。
代码中与 ufunc 实际计算相对应的位置用 和标记
。这些行之间的代码是创建您自己的 ufunc 必须更改的主要内容。/\* BEGIN main ufunc computation \*/
/\* END main ufunc computation \*/
#define PY_SSIZE_T_CLEAN #include <Python.h> #include "numpy/ndarraytypes.h" #include "numpy/ufuncobject.h" #include "numpy/npy_3kcompat.h" #include <math.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, const npy_intp *dimensions, const 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 struct PyModuleDef moduledef = { PyModuleDef_HEAD_INIT, "npufunc", NULL, -1, LogitMethods, NULL, NULL, NULL, NULL }; PyMODINIT_FUNC PyInit_npufunc(void) { PyObject *m, *logit, *d; import_array(); import_umath(); m = PyModule_Create(&moduledef); if (!m) { return NULL; } logit = PyUFunc_FromFuncAndData(funcs, NULL, types, 1, 1, 1, PyUFunc_None, "logit", "logit_docstring", 0); d = PyModule_GetDict(m); PyDict_SetItemString(d, "logit", logit); Py_DECREF(logit); return m; }
这是上面代码的一个。和以前一样,可以通过在命令提示符处调用来构建该模块,或者通过.该模块也可以放置到本地文件夹中,例如下面使用.setup.py file
python setup.py build
python setup.py install
npufunc_directory
python setup.py build_ext --inplace
''' setup.py file for single_type_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 npufunc_directory. 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 on numpy.distutils for more information. ''' def configuration(parent_package='', top_path=None): from numpy.distutils.misc_util import Configuration config = Configuration('npufunc_directory', parent_package, top_path) config.add_extension('npufunc', ['single_type_logit.c']) 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 示例#
最后,我们给出一个完整的 ufunc 示例,其中包含半浮点数、浮点数、双精度数和长双精度数的内部循环。和前面几节一样,我们首先给出文件.c
,然后给出相应的
setup.py
文件。
代码中与 ufunc 实际计算相对应的位置用 和标记
。这些行之间的代码是创建您自己的 ufunc 必须更改的主要内容。/\* BEGIN main ufunc computation \*/
/\* END main ufunc computation \*/
#define PY_SSIZE_T_CLEAN #include <Python.h> #include "numpy/ndarraytypes.h" #include "numpy/ufuncobject.h" #include "numpy/halffloat.h" #include <math.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, const npy_intp *dimensions, const 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, const npy_intp *dimensions, const 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, const npy_intp *dimensions, const 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, const npy_intp *dimensions, const 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_to_float(*(npy_half *)in); 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 struct PyModuleDef moduledef = { PyModuleDef_HEAD_INIT, "npufunc", NULL, -1, LogitMethods, NULL, NULL, NULL, NULL }; PyMODINIT_FUNC PyInit_npufunc(void) { PyObject *m, *logit, *d; import_array(); import_umath(); m = PyModule_Create(&moduledef); if (!m) { return NULL; } logit = PyUFunc_FromFuncAndData(funcs, NULL, types, 4, 1, 1, PyUFunc_None, "logit", "logit_docstring", 0); d = PyModule_GetDict(m); PyDict_SetItemString(d, "logit", logit); Py_DECREF(logit); return m; }
这是setup.py
上述代码的文件。和以前一样,可以通过在命令提示符处调用来构建该模块,或者通过.python setup.py build
python setup.py install
''' setup.py file for multi_type_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 on numpy.distutils for more information. ''' def configuration(parent_package='', top_path=None): from numpy.distutils.misc_util import Configuration, get_info #Necessary for the half-float d-type. info = get_info('npymath') config = Configuration('npufunc_directory', parent_package, top_path) config.add_extension('npufunc', ['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 示例#
我们的最后一个示例是具有多个参数的 ufunc。它是对具有单一数据类型的数据的 logit ufunc 代码的修改。我们计算.(A * B, logit(A * B))
我们只提供 C 代码,因为 setup.py 文件与示例 NumPy ufuncsetup.py
中针对一种 dtype 的文件完全相同,除了该行
config.add_extension('npufunc', ['single_type_logit.c'])
被替换为
config.add_extension('npufunc', ['multi_arg_logit.c'])
C 文件如下。生成的 ufunc 有两个参数A
和B
。它返回一个元组,其第一个元素是,第二个元素是。请注意,它自动支持广播以及 ufunc 的所有其他属性。A * B
logit(A * B)
#define PY_SSIZE_T_CLEAN #include <Python.h> #include "numpy/ndarraytypes.h" #include "numpy/ufuncobject.h" #include "numpy/halffloat.h" #include <math.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 * with comments. * * 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, const npy_intp *dimensions, const 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 struct PyModuleDef moduledef = { PyModuleDef_HEAD_INIT, "npufunc", NULL, -1, LogitMethods, NULL, NULL, NULL, NULL }; PyMODINIT_FUNC PyInit_npufunc(void) { PyObject *m, *logit, *d; import_array(); import_umath(); m = PyModule_Create(&moduledef); if (!m) { return NULL; } logit = PyUFunc_FromFuncAndData(funcs, NULL, types, 1, 2, 2, PyUFunc_None, "logit", "logit_docstring", 0); d = PyModule_GetDict(m); PyDict_SetItemString(d, "logit", logit); Py_DECREF(logit); return m; }
带有结构化数组 dtype 参数的 NumPy ufunc 示例#
此示例演示如何为结构化数组数据类型创建 ufunc。在示例中,我们展示了一个简单的 ufunc,用于将两个具有 dtype 的数组相加
'u8,u8,u8'
。该过程与其他示例略有不同,因为调用PyUFunc_FromFuncAndData
不会完全注册自定义数据类型和结构化数组数据类型的 ufunc。我们还需要调用
PyUFunc_RegisterLoopForDescr
来完成 ufunc 的设置。
我们只给出 C 代码,因为该文件与示例 NumPy ufunc中的一种 dtype 的文件setup.py
完全相同,除了该行setup.py
config.add_extension('npufunc', ['single_type_logit.c'])
被替换为
config.add_extension('npufunc', ['add_triplet.c'])
C 文件如下。
#define PY_SSIZE_T_CLEAN #include <Python.h> #include "numpy/ndarraytypes.h" #include "numpy/ufuncobject.h" #include "numpy/npy_3kcompat.h" #include <math.h> /* * add_triplet.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, const npy_intp *dimensions, const 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 */ PyUFuncGenericFunction funcs[1] = {&add_uint64_triplet}; /* These are the input and return dtypes of add_uint64_triplet. */ static char types[3] = {NPY_UINT64, NPY_UINT64, NPY_UINT64}; static struct PyModuleDef moduledef = { PyModuleDef_HEAD_INIT, "struct_ufunc_test", NULL, -1, StructUfuncTestMethods, NULL, NULL, NULL, NULL }; PyMODINIT_FUNC PyInit_struct_ufunc_test(void) { PyObject *m, *add_triplet, *d; PyObject *dtype_dict; PyArray_Descr *dtype; PyArray_Descr *dtypes[3]; import_array(); import_umath(); m = PyModule_Create(&moduledef); if (m == NULL) { return NULL; } /* Create a new ufunc object */ 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; /* Register ufunc for structured 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); return m; }
返回的 ufunc 对象是一个可调用的 Python 对象。它应该放置在(模块)字典中,其名称与 ufunc-creation 例程的名称参数中使用的名称相同。以下示例改编自umath模块
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); ...