编写您自己的ufunc

I have the Power!

— He-Man

创建一个新的ufunc

在阅读本文之前，通过阅读/略读扩展和嵌入Python解释器的第1部分中的教程以及如何扩展NumPy，可以帮助您熟悉Python的C扩展基础知识。

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

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

创建自己的ufunc并不困难。所需要的只是您想要支持的每种数据类型的1-d循环。每个1-d循环必须具有特定签名，并且只能使用固定大小数据类型的ufunc。下面给出了用于创建新的ufunc以处理内置数据类型的函数调用。使用不同的机制为用户定义的数据类型注册ufunc。

在接下来的几节中，我们提供了可以轻松修改的示例代码，以创建自己的ufunc。这些示例是logit函数的连续更完整或复杂版本，这是统计建模中的常见功能。Logit也很有趣，因为由于IEEE标准（特别是IEEE 754）的神奇之处，下面创建的所有logit函数都自动具有以下行为。

  1. >>> logit(0)
  2. -inf
  3. >>> logit(1)
  4. inf
  5. >>> logit(2)
  6. nan
  7. >>> logit(-2)
  8. nan

这很好，因为函数编写器不必手动传播infs或nans。

示例非ufunc扩展名

为了比较和阅读器的一般启发，我们提供了一个简单的logit C扩展实现，它没有使用numpy。

为此，我们需要两个文件。第一个是包含实际代码的C文件，第二个是用于创建模块的setup.py文件。

  1. #include <Python.h>
  2. #include <math.h>
  4. /*
  5.  * spammodule.c
  6.  * This is the C code for a non-numpy Python extension to
  7.  * define the logit function, where logit(p) = log(p/(1-p)).
  8.  * This function will not work on numpy arrays automatically.
  9.  * numpy.vectorize must be called in python to generate
  10.  * a numpy-friendly function.
  11.  *
  12.  * Details explaining the Python-C API can be found under
  13.  * 'Extending and Embedding' and 'Python/C API' at
  14.  * docs.python.org .
  15.  */
  18. /* This declares the logit function */
  19. static PyObject* spam_logit(PyObject *self, PyObject *args);
  22. /*
  23.  * This tells Python what methods this module has.
  24.  * See the Python-C API for more information.
  25.  */
  26. static PyMethodDef SpamMethods[] = {
  27.     {"logit",
  28.         spam_logit,
  29.         METH_VARARGS, "compute logit"},
  30.     {NULL, NULL, 0, NULL}
  31. };
  34. /*
  35.  * This actually defines the logit function for
  36.  * input args from Python.
  37.  */
  39. static PyObject* spam_logit(PyObject *self, PyObject *args)
  40. {
  41.     double p;
  43.     /* This parses the Python argument into a double */
  44.     if(!PyArg_ParseTuple(args, "d", &p)) {
  45.         return NULL;
  46.     }
  48.     /* THE ACTUAL LOGIT FUNCTION */
  49.     p = p/(1-p);
  50.     p = log(p);
  52.     /*This builds the answer back into a python object */
  53.     return Py_BuildValue("d", p);
  54. }
  57. /* This initiates the module using the above definitions. */
  58. #if PY_VERSION_HEX >= 0x03000000
  59. static struct PyModuleDef moduledef = {
  60.     PyModuleDef_HEAD_INIT,
  61.     "spam",
  62.     NULL,
  63.     -1,
  64.     SpamMethods,
  65.     NULL,
  66.     NULL,
  67.     NULL,
  68.     NULL
  69. };
  71. PyMODINIT_FUNC PyInit_spam(void)
  72. {
  73.     PyObject *m;
  74.     m = PyModule_Create(&moduledef);
  75.     if (!m) {
  76.         return NULL;
  77.     }
  78.     return m;
  79. }
  80. #else
  81. PyMODINIT_FUNC initspam(void)
  82. {
  83.     PyObject *m;
  85.     m = Py_InitModule("spam", SpamMethods);
  86.     if (m == NULL) {
  87.         return;
  88.     }
  89. }
  90. #endif

要使用setup.py文件，请将setup.py和spammodule.c放在同一文件夹中。然后python setup.py build将构建要导入的模块，或者setup.py install将模块安装到您的site-packages目录。

  1. '''
  2.     setup.py file for spammodule.c
  4.     Calling
  5.     $python setup.py build_ext --inplace
  6.     will build the extension library in the current file.
  8.     Calling
  9.     $python setup.py build
  10.     will build a file that looks like ./build/lib*, where
  11.     lib* is a file that begins with lib. The library will
  12.     be in this file and end with a C library extension,
  13.     such as .so
  15.     Calling
  16.     $python setup.py install
  17.     will install the module in your site-packages file.
  19.     See the distutils section of
  20.     'Extending and Embedding the Python Interpreter'
  21.     at docs.python.org for more information.
  22. '''
  25. from distutils.core import setup, Extension
  27. module1 = Extension('spam', sources=['spammodule.c'],
  28.                         include_dirs=['/usr/local/lib'])
  30. setup(name = 'spam',
  31.         version='1.0',
  32.         description='This is my spam package',
  33.         ext_modules = [module1])

将垃圾邮件模块导入python后，您可以通过spam.logit调用logit。请注意，上面使用的函数不能按原样应用于numpy数组。为此，我们必须在其上调用numpy.vectorize。例如，如果在包含垃圾邮件库或垃圾邮件的文件中打开了python解释器，则可以执行以下命令：

  1. >>> import numpy as np
  2. >>> import spam
  3. >>> spam.logit(0)
  4. -inf
  5. >>> spam.logit(1)
  6. inf
  7. >>> spam.logit(0.5)
  8. 0.0
  9. >>> x = np.linspace(0,1,10)
  10. >>> spam.logit(x)
  11. TypeError: only length-1 arrays can be converted to Python scalars
  12. >>> f = np.vectorize(spam.logit)
  13. >>> f(x)
  14. array([       -inf, -2.07944154, -1.25276297, -0.69314718, -0.22314355,
  15.     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的主要事物。

  1. #include "Python.h"
  2. #include "math.h"
  3. #include "numpy/ndarraytypes.h"
  4. #include "numpy/ufuncobject.h"
  5. #include "numpy/npy_3kcompat.h"
  7. /*
  8.  * single_type_logit.c
  9.  * This is the C code for creating your own
  10.  * NumPy ufunc for a logit function.
  11.  *
  12.  * In this code we only define the ufunc for
  13.  * a single dtype. The computations that must
  14.  * be replaced to create a ufunc for
  15.  * a different function are marked with BEGIN
  16.  * and END.
  17.  *
  18.  * Details explaining the Python-C API can be found under
  19.  * 'Extending and Embedding' and 'Python/C API' at
  20.  * docs.python.org .
  21.  */
  23. static PyMethodDef LogitMethods[] = {
  24.         {NULL, NULL, 0, NULL}
  25. };
  27. /* The loop definition must precede the PyMODINIT_FUNC. */
  29. static void double_logit(char **args, npy_intp *dimensions,
  30.                             npy_intp* steps, void* data)
  31. {
  32.     npy_intp i;
  33.     npy_intp n = dimensions[0];
  34.     char *in = args[0], *out = args[1];
  35.     npy_intp in_step = steps[0], out_step = steps[1];
  37.     double tmp;
  39.     for (i = 0; i < n; i++) {
  40.         /*BEGIN main ufunc computation*/
  41.         tmp = *(double *)in;
  42.         tmp /= 1-tmp;
  43.         *((double *)out) = log(tmp);
  44.         /*END main ufunc computation*/
  46.         in += in_step;
  47.         out += out_step;
  48.     }
  49. }
  51. /*This a pointer to the above function*/
  52. PyUFuncGenericFunction funcs[1] = {&double_logit};
  54. /* These are the input and return dtypes of logit.*/
  55. static char types[2] = {NPY_DOUBLE, NPY_DOUBLE};
  57. static void *data[1] = {NULL};
  59. #if PY_VERSION_HEX >= 0x03000000
  60. static struct PyModuleDef moduledef = {
  61.     PyModuleDef_HEAD_INIT,
  62.     "npufunc",
  63.     NULL,
  64.     -1,
  65.     LogitMethods,
  66.     NULL,
  67.     NULL,
  68.     NULL,
  69.     NULL
  70. };
  72. PyMODINIT_FUNC PyInit_npufunc(void)
  73. {
  74.     PyObject *m, *logit, *d;
  75.     m = PyModule_Create(&moduledef);
  76.     if (!m) {
  77.         return NULL;
  78.     }
  80.     import_array();
  81.     import_umath();
  83.     logit = PyUFunc_FromFuncAndData(funcs, data, types, 1, 1, 1,
  84.                                     PyUFunc_None, "logit",
  85.                                     "logit_docstring", 0);
  87.     d = PyModule_GetDict(m);
  89.     PyDict_SetItemString(d, "logit", logit);
  90.     Py_DECREF(logit);
  92.     return m;
  93. }
  94. #else
  95. PyMODINIT_FUNC initnpufunc(void)
  96. {
  97.     PyObject *m, *logit, *d;
  100.     m = Py_InitModule("npufunc", LogitMethods);
  101.     if (m == NULL) {
  102.         return;
  103.     }
  105.     import_array();
  106.     import_umath();
  108.     logit = PyUFunc_FromFuncAndData(funcs, data, types, 1, 1, 1,
  109.                                     PyUFunc_None, "logit",
  110.                                     "logit_docstring", 0);
  112.     d = PyModule_GetDict(m);
  114.     PyDict_SetItemString(d, "logit", logit);
  115.     Py_DECREF(logit);
  116. }
  117. #endif

这是上面代码的setup.py文件。和以前一样，可以通过在命令提示符下调用python setup.py build来构建模块，也可以通过python setup.py install将其安装到site-packages。

  1. '''
  2.     setup.py file for logit.c
  3.     Note that since this is a numpy extension
  4.     we use numpy.distutils instead of
  5.     distutils from the python standard library.
  7.     Calling
  8.     $python setup.py build_ext --inplace
  9.     will build the extension library in the current file.
  11.     Calling
  12.     $python setup.py build
  13.     will build a file that looks like ./build/lib*, where
  14.     lib* is a file that begins with lib. The library will
  15.     be in this file and end with a C library extension,
  16.     such as .so
  18.     Calling
  19.     $python setup.py install
  20.     will install the module in your site-packages file.
  22.     See the distutils section of
  23.     'Extending and Embedding the Python Interpreter'
  24.     at docs.python.org  and the documentation
  25.     on numpy.distutils for more information.
  26. '''
  29. def configuration(parent_package='', top_path=None):
  30.     import numpy
  31.     from numpy.distutils.misc_util import Configuration
  33.     config = Configuration('npufunc_directory',
  34.                            parent_package,
  35.                            top_path)
  36.     config.add_extension('npufunc', ['single_type_logit.c'])
  38.     return config
  40. if __name__ == "__main__":
  41.     from numpy.distutils.core import setup
  42.     setup(configuration=configuration)

安装完上述内容后，可以按如下方式导入和使用。

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

示例具有多个dtypes的NumPy ufunc

我们最后给出了一个完整的ufunc示例，内部循环用于半浮点数，浮点数，双精度数和长双精度数。与前面的部分一样，我们首先给出.c文件，然后是相应的setup.py文件。

代码中与ufunc的实际计算相对应的位置标有/ BEGIN main ufunc computation /和/ END main ufunc computation /。这些行之间的代码是必须更改以创建自己的ufunc的主要事物。

  1. #include "Python.h"
  2. #include "math.h"
  3. #include "numpy/ndarraytypes.h"
  4. #include "numpy/ufuncobject.h"
  5. #include "numpy/halffloat.h"
  7. /*
  8.  * multi_type_logit.c
  9.  * This is the C code for creating your own
  10.  * NumPy ufunc for a logit function.
  11.  *
  12.  * Each function of the form type_logit defines the
  13.  * logit function for a different numpy dtype. Each
  14.  * of these functions must be modified when you
  15.  * create your own ufunc. The computations that must
  16.  * be replaced to create a ufunc for
  17.  * a different function are marked with BEGIN
  18.  * and END.
  19.  *
  20.  * Details explaining the Python-C API can be found under
  21.  * 'Extending and Embedding' and 'Python/C API' at
  22.  * docs.python.org .
  23.  *
  24.  */
  27. static PyMethodDef LogitMethods[] = {
  28.         {NULL, NULL, 0, NULL}
  29. };
  31. /* The loop definitions must precede the PyMODINIT_FUNC. */
  33. static void long_double_logit(char **args, npy_intp *dimensions,
  34.                               npy_intp* steps, void* data)
  35. {
  36.     npy_intp i;
  37.     npy_intp n = dimensions[0];
  38.     char *in = args[0], *out=args[1];
  39.     npy_intp in_step = steps[0], out_step = steps[1];
  41.     long double tmp;
  43.     for (i = 0; i < n; i++) {
  44.         /*BEGIN main ufunc computation*/
  45.         tmp = *(long double *)in;
  46.         tmp /= 1-tmp;
  47.         *((long double *)out) = logl(tmp);
  48.         /*END main ufunc computation*/
  50.         in += in_step;
  51.         out += out_step;
  52.     }
  53. }
  55. static void double_logit(char **args, npy_intp *dimensions,
  56.                          npy_intp* steps, void* data)
  57. {
  58.     npy_intp i;
  59.     npy_intp n = dimensions[0];
  60.     char *in = args[0], *out = args[1];
  61.     npy_intp in_step = steps[0], out_step = steps[1];
  63.     double tmp;
  65.     for (i = 0; i < n; i++) {
  66.         /*BEGIN main ufunc computation*/
  67.         tmp = *(double *)in;
  68.         tmp /= 1-tmp;
  69.         *((double *)out) = log(tmp);
  70.         /*END main ufunc computation*/
  72.         in += in_step;
  73.         out += out_step;
  74.     }
  75. }
  77. static void float_logit(char **args, npy_intp *dimensions,
  78.                         npy_intp* steps, void* data)
  79. {
  80.     npy_intp i;
  81.     npy_intp n = dimensions[0];
  82.     char *in=args[0], *out = args[1];
  83.     npy_intp in_step = steps[0], out_step = steps[1];
  85.     float tmp;
  87.     for (i = 0; i < n; i++) {
  88.         /*BEGIN main ufunc computation*/
  89.         tmp = *(float *)in;
  90.         tmp /= 1-tmp;
  91.         *((float *)out) = logf(tmp);
  92.         /*END main ufunc computation*/
  94.         in += in_step;
  95.         out += out_step;
  96.     }
  97. }
  100. static void half_float_logit(char **args, npy_intp *dimensions,
  101.                              npy_intp* steps, void* data)
  102. {
  103.     npy_intp i;
  104.     npy_intp n = dimensions[0];
  105.     char *in = args[0], *out = args[1];
  106.     npy_intp in_step = steps[0], out_step = steps[1];
  108.     float tmp;
  110.     for (i = 0; i < n; i++) {
  112.         /*BEGIN main ufunc computation*/
  113.         tmp = *(npy_half *)in;
  114.         tmp = npy_half_to_float(tmp);
  115.         tmp /= 1-tmp;
  116.         tmp = logf(tmp);
  117.         *((npy_half *)out) = npy_float_to_half(tmp);
  118.         /*END main ufunc computation*/
  120.         in += in_step;
  121.         out += out_step;
  122.     }
  123. }
  126. /*This gives pointers to the above functions*/
  127. PyUFuncGenericFunction funcs[4] = {&half_float_logit,
  128.                                    &float_logit,
  129.                                    &double_logit,
  130.                                    &long_double_logit};
  132. static char types[8] = {NPY_HALF, NPY_HALF,
  133.                 NPY_FLOAT, NPY_FLOAT,
  134.                 NPY_DOUBLE,NPY_DOUBLE,
  135.                 NPY_LONGDOUBLE, NPY_LONGDOUBLE};
  136. static void *data[4] = {NULL, NULL, NULL, NULL};
  138. #if PY_VERSION_HEX >= 0x03000000
  139. static struct PyModuleDef moduledef = {
  140.     PyModuleDef_HEAD_INIT,
  141.     "npufunc",
  142.     NULL,
  143.     -1,
  144.     LogitMethods,
  145.     NULL,
  146.     NULL,
  147.     NULL,
  148.     NULL
  149. };
  151. PyMODINIT_FUNC PyInit_npufunc(void)
  152. {
  153.     PyObject *m, *logit, *d;
  154.     m = PyModule_Create(&moduledef);
  155.     if (!m) {
  156.         return NULL;
  157.     }
  159.     import_array();
  160.     import_umath();
  162.     logit = PyUFunc_FromFuncAndData(funcs, data, types, 4, 1, 1,
  163.                                     PyUFunc_None, "logit",
  164.                                     "logit_docstring", 0);
  166.     d = PyModule_GetDict(m);
  168.     PyDict_SetItemString(d, "logit", logit);
  169.     Py_DECREF(logit);
  171.     return m;
  172. }
  173. #else
  174. PyMODINIT_FUNC initnpufunc(void)
  175. {
  176.     PyObject *m, *logit, *d;
  179.     m = Py_InitModule("npufunc", LogitMethods);
  180.     if (m == NULL) {
  181.         return;
  182.     }
  184.     import_array();
  185.     import_umath();
  187.     logit = PyUFunc_FromFuncAndData(funcs, data, types, 4, 1, 1,
  188.                                     PyUFunc_None, "logit",
  189.                                     "logit_docstring", 0);
  191.     d = PyModule_GetDict(m);
  193.     PyDict_SetItemString(d, "logit", logit);
  194.     Py_DECREF(logit);
  195. }
  196. #endif

这是上面代码的setup.py文件。和以前一样，可以通过在命令提示符下调用python setup.py build来构建模块，也可以通过python setup.py install将其安装到site-packages。

  1. '''
  2.     setup.py file for logit.c
  3.     Note that since this is a numpy extension
  4.     we use numpy.distutils instead of
  5.     distutils from the python standard library.
  7.     Calling
  8.     $python setup.py build_ext --inplace
  9.     will build the extension library in the current file.
  11.     Calling
  12.     $python setup.py build
  13.     will build a file that looks like ./build/lib*, where
  14.     lib* is a file that begins with lib. The library will
  15.     be in this file and end with a C library extension,
  16.     such as .so
  18.     Calling
  19.     $python setup.py install
  20.     will install the module in your site-packages file.
  22.     See the distutils section of
  23.     'Extending and Embedding the Python Interpreter'
  24.     at docs.python.org  and the documentation
  25.     on numpy.distutils for more information.
  26. '''
  29. def configuration(parent_package='', top_path=None):
  30.     import numpy
  31.     from numpy.distutils.misc_util import Configuration
  32.     from numpy.distutils.misc_util import get_info
  34.     #Necessary for the half-float d-type.
  35.     info = get_info('npymath')
  37.     config = Configuration('npufunc_directory',
  38.                             parent_package,
  39.                             top_path)
  40.     config.add_extension('npufunc',
  41.                             ['multi_type_logit.c'],
  42.                             extra_info=info)
  44.     return config
  46. if __name__ == "__main__":
  47.     from numpy.distutils.core import setup
  48.     setup(configuration=configuration)

安装完上述内容后，可以按如下方式导入和使用。

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

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

我们的最后一个例子是一个带有多个参数的ufunc。它是对具有单个dtype的数据的logit ufunc的代码的修改。我们计算 (A*B, logit(A*B))。

我们只给出 C 代码，因为setup.py文件与一种dtype的NumPy ufunc示例中的setup.py文件完全相同，只是一行

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

被替换为

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

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

  1. #include "Python.h"
  2. #include "math.h"
  3. #include "numpy/ndarraytypes.h"
  4. #include "numpy/ufuncobject.h"
  5. #include "numpy/halffloat.h"
  7. /*
  8.  * multi_arg_logit.c
  9.  * This is the C code for creating your own
  10.  * NumPy ufunc for a multiple argument, multiple
  11.  * return value ufunc. The places where the
  12.  * ufunc computation is carried out are marked
  13.  * with comments.
  14.  *
  15.  * Details explaining the Python-C API can be found under
  16.  * 'Extending and Embedding' and 'Python/C API' at
  17.  * docs.python.org .
  18.  *
  19.  */
  22. static PyMethodDef LogitMethods[] = {
  23.         {NULL, NULL, 0, NULL}
  24. };
  26. /* The loop definition must precede the PyMODINIT_FUNC. */
  28. static void double_logitprod(char **args, npy_intp *dimensions,
  29.                             npy_intp* steps, void* data)
  30. {
  31.     npy_intp i;
  32.     npy_intp n = dimensions[0];
  33.     char *in1 = args[0], *in2 = args[1];
  34.     char *out1 = args[2], *out2 = args[3];
  35.     npy_intp in1_step = steps[0], in2_step = steps[1];
  36.     npy_intp out1_step = steps[2], out2_step = steps[3];
  38.     double tmp;
  40.     for (i = 0; i < n; i++) {
  41.         /*BEGIN main ufunc computation*/
  42.         tmp = *(double *)in1;
  43.         tmp *= *(double *)in2;
  44.         *((double *)out1) = tmp;
  45.         *((double *)out2) = log(tmp/(1-tmp));
  46.         /*END main ufunc computation*/
  48.         in1 += in1_step;
  49.         in2 += in2_step;
  50.         out1 += out1_step;
  51.         out2 += out2_step;
  52.     }
  53. }
  56. /*This a pointer to the above function*/
  57. PyUFuncGenericFunction funcs[1] = {&double_logitprod};
  59. /* These are the input and return dtypes of logit.*/
  61. static char types[4] = {NPY_DOUBLE, NPY_DOUBLE,
  62.                         NPY_DOUBLE, NPY_DOUBLE};
  65. static void *data[1] = {NULL};
  67. #if PY_VERSION_HEX >= 0x03000000
  68. static struct PyModuleDef moduledef = {
  69.     PyModuleDef_HEAD_INIT,
  70.     "npufunc",
  71.     NULL,
  72.     -1,
  73.     LogitMethods,
  74.     NULL,
  75.     NULL,
  76.     NULL,
  77.     NULL
  78. };
  80. PyMODINIT_FUNC PyInit_npufunc(void)
  81. {
  82.     PyObject *m, *logit, *d;
  83.     m = PyModule_Create(&moduledef);
  84.     if (!m) {
  85.         return NULL;
  86.     }
  88.     import_array();
  89.     import_umath();
  91.     logit = PyUFunc_FromFuncAndData(funcs, data, types, 1, 2, 2,
  92.                                     PyUFunc_None, "logit",
  93.                                     "logit_docstring", 0);
  95.     d = PyModule_GetDict(m);
  97.     PyDict_SetItemString(d, "logit", logit);
  98.     Py_DECREF(logit);
  100.     return m;
  101. }
  102. #else
  103. PyMODINIT_FUNC initnpufunc(void)
  104. {
  105.     PyObject *m, *logit, *d;
  108.     m = Py_InitModule("npufunc", LogitMethods);
  109.     if (m == NULL) {
  110.         return;
  111.     }
  113.     import_array();
  114.     import_umath();
  116.     logit = PyUFunc_FromFuncAndData(funcs, data, types, 1, 2, 2,
  117.                                     PyUFunc_None, "logit",
  118.                                     "logit_docstring", 0);
  120.     d = PyModule_GetDict(m);
  122.     PyDict_SetItemString(d, "logit", logit);
  123.     Py_DECREF(logit);
  124. }
  125. #endif

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

此示例显示如何为结构化数组dtype创建ufunc。在这个例子中，我们展示了一个简单的ufunc，用于添加两个带有dtype’u8，u8，u8’的数组。该过程与其他示例略有不同，因为调用PyUFunc_FromFuncAndData不会为自定义dtypes和结构化数组dtypes完全注册ufunc。我们还需要调用 PyUFunc_RegisterLoopForDescr完成设置ufunc。

我们只提供C代码，因为setup.py文件与一种dtype的NumPy ufunc示例的setup.py文件完全相同，只有一行。

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

被替换为

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

C文件如下。

  1. #include "Python.h"
  2. #include "math.h"
  3. #include "numpy/ndarraytypes.h"
  4. #include "numpy/ufuncobject.h"
  5. #include "numpy/npy_3kcompat.h"
  8. /*
  9.  * add_triplet.c
  10.  * This is the C code for creating your own
  11.  * NumPy ufunc for a structured array dtype.
  12.  *
  13.  * Details explaining the Python-C API can be found under
  14.  * 'Extending and Embedding' and 'Python/C API' at
  15.  * docs.python.org .
  16.  */
  18. static PyMethodDef StructUfuncTestMethods[] = {
  19.     {NULL, NULL, 0, NULL}
  20. };
  22. /* The loop definition must precede the PyMODINIT_FUNC. */
  24. static void add_uint64_triplet(char **args, npy_intp *dimensions,
  25.                             npy_intp* steps, void* data)
  26. {
  27.     npy_intp i;
  28.     npy_intp is1=steps[0];
  29.     npy_intp is2=steps[1];
  30.     npy_intp os=steps[2];
  31.     npy_intp n=dimensions[0];
  32.     uint64_t *x, *y, *z;
  34.     char *i1=args[0];
  35.     char *i2=args[1];
  36.     char *op=args[2];
  38.     for (i = 0; i < n; i++) {
  40.         x = (uint64_t*)i1;
  41.         y = (uint64_t*)i2;
  42.         z = (uint64_t*)op;
  44.         z[0] = x[0] + y[0];
  45.         z[1] = x[1] + y[1];
  46.         z[2] = x[2] + y[2];
  48.         i1 += is1;
  49.         i2 += is2;
  50.         op += os;
  51.     }
  52. }
  54. /* This a pointer to the above function */
  55. PyUFuncGenericFunction funcs[1] = {&add_uint64_triplet};
  57. /* These are the input and return dtypes of add_uint64_triplet. */
  58. static char types[3] = {NPY_UINT64, NPY_UINT64, NPY_UINT64};
  60. static void *data[1] = {NULL};
  62. #if defined(NPY_PY3K)
  63. static struct PyModuleDef moduledef = {
  64.     PyModuleDef_HEAD_INIT,
  65.     "struct_ufunc_test",
  66.     NULL,
  67.     -1,
  68.     StructUfuncTestMethods,
  69.     NULL,
  70.     NULL,
  71.     NULL,
  72.     NULL
  73. };
  74. #endif
  76. #if defined(NPY_PY3K)
  77. PyMODINIT_FUNC PyInit_struct_ufunc_test(void)
  78. #else
  79. PyMODINIT_FUNC initstruct_ufunc_test(void)
  80. #endif
  81. {
  82.     PyObject *m, *add_triplet, *d;
  83.     PyObject *dtype_dict;
  84.     PyArray_Descr *dtype;
  85.     PyArray_Descr *dtypes[3];
  87. #if defined(NPY_PY3K)
  88.     m = PyModule_Create(&moduledef);
  89. #else
  90.     m = Py_InitModule("struct_ufunc_test", StructUfuncTestMethods);
  91. #endif
  93.     if (m == NULL) {
  94. #if defined(NPY_PY3K)
  95.         return NULL;
  96. #else
  97.         return;
  98. #endif
  99.     }
  101.     import_array();
  102.     import_umath();
  104.     /* Create a new ufunc object */
  105.     add_triplet = PyUFunc_FromFuncAndData(NULL, NULL, NULL, 0, 2, 1,
  106.                                     PyUFunc_None, "add_triplet",
  107.                                     "add_triplet_docstring", 0);
  109.     dtype_dict = Py_BuildValue("[(s, s), (s, s), (s, s)]",
  110.         "f0", "u8", "f1", "u8", "f2", "u8");
  111.     PyArray_DescrConverter(dtype_dict, &dtype);
  112.     Py_DECREF(dtype_dict);
  114.     dtypes[0] = dtype;
  115.     dtypes[1] = dtype;
  116.     dtypes[2] = dtype;
  118.     /* Register ufunc for structured dtype */
  119.     PyUFunc_RegisterLoopForDescr(add_triplet,
  120.                                 dtype,
  121.                                 &add_uint64_triplet,
  122.                                 dtypes,
  123.                                 NULL);
  125.     d = PyModule_GetDict(m);
  127.     PyDict_SetItemString(d, "add_triplet", add_triplet);
  128.     Py_DECREF(add_triplet);
  129. #if defined(NPY_PY3K)
  130.     return m;
  131. #endif
  132. }

返回的ufunc对象是一个可调用的Python对象。它应该放在一个（模块）字典中，其名称与ufunc-creation例程的name参数中使用的字典相同。以下示例是从umath模块改编而来的

  1. static PyUFuncGenericFunction atan2_functions[] = {
  2.                       PyUFunc_ff_f, PyUFunc_dd_d,
  3.                       PyUFunc_gg_g, PyUFunc_OO_O_method};
  4. static void* atan2_data[] = {
  5.                       (void *)atan2f,(void *) atan2,
  6.                       (void *)atan2l,(void *)"arctan2"};
  7. static char atan2_signatures[] = {
  8.               NPY_FLOAT, NPY_FLOAT, NPY_FLOAT,
  9.               NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE,
  10.               NPY_LONGDOUBLE, NPY_LONGDOUBLE, NPY_LONGDOUBLE
  11.               NPY_OBJECT, NPY_OBJECT, NPY_OBJECT};
  12. ...
  13. /* in the module initialization code */
  14. PyObject *f, *dict, *module;
  15. ...
  16. dict = PyModule_GetDict(module);
  17. ...
  18. f = PyUFunc_FromFuncAndData(atan2_functions,
  19.     atan2_data, atan2_signatures, 4, 2, 1,
  20.     PyUFunc_None, "arctan2",
  21.     "a safe and correct arctan(x1/x2)", 0);
  22. PyDict_SetItemString(dict, "arctan2", f);
  23. Py_DECREF(f);
  24. ...