編寫您自己的ufunc

創(chuàng)建一個(gè)新的ufunc

在閱讀本文之前,通過(guò)閱讀/略讀擴(kuò)展和嵌入Python解釋器的

第1部分中的教程以及如何擴(kuò)展NumPy,可以幫助您熟悉Python的C擴(kuò)展基礎(chǔ)知識(shí)。

umath模塊是一個(gè)計(jì)算機(jī)生成的C模塊,可以創(chuàng)建許多ufunc。它提供了許多如何創(chuàng)建通用函數(shù)的示例。使用ufunc機(jī)制創(chuàng)建自己的ufunc也不困難。假設(shè)您有一個(gè)函數(shù),您想要在其輸入上逐個(gè)元素地操作。通過(guò)創(chuàng)建一個(gè)新的ufunc,您將獲得一個(gè)處理的函數(shù)

  • 廣播
  • N維循環(huán)
  • 自動(dòng)類型轉(zhuǎn)換,內(nèi)存使用量最少
  • 可選的輸出數(shù)組

創(chuàng)建自己的ufunc并不困難。所需要的只是您想要支持的每種數(shù)據(jù)類型的1-d循環(huán)。每個(gè)1-d循環(huán)必須具有特定簽名,并且只能使用固定大小數(shù)據(jù)類型的ufunc。下面給出了用于創(chuàng)建新的ufunc以處理內(nèi)置數(shù)據(jù)類型的函數(shù)調(diào)用。使用不同的機(jī)制為用戶定義的數(shù)據(jù)類型注冊(cè)u(píng)func。

在接下來(lái)的幾節(jié)中,我們提供了可以輕松修改的示例代碼,以創(chuàng)建自己的ufunc。這些示例是logit函數(shù)的連續(xù)更完整或復(fù)雜版本,這是統(tǒng)計(jì)建模中的常見功能。Logit也很有趣,因?yàn)橛捎贗EEE標(biāo)準(zhǔn)(特別是IEEE 754)的神奇之處,下面創(chuàng)建的所有l(wèi)ogit函數(shù)都自動(dòng)具有以下行為。

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

這很好,因?yàn)楹瘮?shù)編寫器不必手動(dòng)傳播infs或nans。

示例非ufunc擴(kuò)展名

為了比較和閱讀器的一般啟發(fā),我們提供了一個(gè)簡(jiǎn)單的logit C擴(kuò)展實(shí)現(xiàn),它沒有使用numpy。

為此,我們需要兩個(gè)文件。第一個(gè)是包含實(shí)際代碼的C文件,第二個(gè)是用于創(chuàng)建模塊的setup.py文件。

#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. */
#if PY_VERSION_HEX >= 0x03000000
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;
}
#else
PyMODINIT_FUNC initspam(void)
{
    PyObject *m;

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

要使用setup.py文件,請(qǐng)將setup.py和spammodule.c放在同一文件夾中。然后python setup.py build將構(gòu)建要導(dǎo)入的模塊,或者setup.py install將模塊安裝到您的site-packages目錄。

'''
    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])

將垃圾郵件模塊導(dǎo)入python后,您可以通過(guò)spam.logit調(diào)用logit。請(qǐng)注意,上面使用的函數(shù)不能按原樣應(yīng)用于numpy數(shù)組。為此,我們必須在其上調(diào)用numpy.vectorize。例如,如果在包含垃圾郵件庫(kù)或垃圾郵件的文件中打開了python解釋器,則可以執(zhí)行以下命令:

>>> 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])

結(jié)果編輯功能并不快!numpy.vectorize只是循環(huán)遍歷spam.logit。循環(huán)在C級(jí)完成,但numpy數(shù)組不斷被解析并重新構(gòu)建。這很貴。當(dāng)作者將numpy.vectorize(spam.logit)與下面構(gòu)造的logit ufuncs進(jìn)行比較時(shí),logit ufuncs幾乎快4倍。當(dāng)然,取決于功能的性質(zhì),可以實(shí)現(xiàn)更大或更小的加速。

一種dtype的NumPy ufunc示例

為簡(jiǎn)單起見,我們?yōu)閱蝹€(gè)dtype提供了一個(gè)ufunc,即'f8'雙精度型。與前一節(jié)一樣,我們首先給出.c文件,然后是用于創(chuàng)建包含ufunc的模塊的setup.py文件。

代碼中與ufunc的實(shí)際計(jì)算相對(duì)應(yīng)的位置標(biāo)有/ * BEGIN main ufunc computation * /和/ * END main ufunc computation * /。這些行之間的代碼是必須更改以創(chuàng)建自己的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文件。和以前一樣,可以通過(guò)在命令提示符下調(diào)用python setup.py build來(lái)構(gòu)建模塊,也可以通過(guò)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
    on numpy.distutils for more information.
'''


def configuration(parent_package='', top_path=None):
    import numpy
    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)

安裝完上述內(nèi)容后,可以按如下方式導(dǎo)入和使用。

>>> 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])

示例具有多個(gè)dtypes的NumPy ufunc

我們最后給出了一個(gè)完整的ufunc示例,內(nèi)部循環(huán)用于半浮點(diǎn)數(shù),浮點(diǎn)數(shù),雙精度數(shù)和長(zhǎng)雙精度數(shù)。與前面的部分一樣,我們首先給出.c文件,然后是相應(yīng)的setup.py文件。

代碼中與ufunc的實(shí)際計(jì)算相對(duì)應(yīng)的位置標(biāo)有/ * BEGIN main ufunc computation * /和/ * END main ufunc computation * /。這些行之間的代碼是必須更改以創(chuàng)建自己的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 = {
    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, 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文件。和以前一樣,可以通過(guò)在命令提示符下調(diào)用python setup.py build來(lái)構(gòu)建模塊,也可以通過(guò)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
    on numpy.distutils for more information.
'''


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)
    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)

安裝完上述內(nèi)容后,可以按如下方式導(dǎo)入和使用。

>>> 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])

示例具有多個(gè)參數(shù)/返回值的NumPy ufunc

我們的最后一個(gè)例子是一個(gè)帶有多個(gè)參數(shù)的ufunc。它是對(duì)具有單個(gè)dtype的數(shù)據(jù)的logit ufunc的代碼的修改。我們計(jì)算 (A*B, logit(A*B))。

我們只給出 C 代碼,因?yàn)閟etup.py文件與一種dtype的NumPy ufunc示例中的setup.py文件完全相同,只是一行

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

被替換為

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






C文件如下。生成的ufunc接受兩個(gè)參數(shù)A和B.它返回一個(gè)元組,其第一個(gè)元素是A * B,第二個(gè)元素是logit(A * B)。請(qǐng)注意,它會(huì)自動(dòng)支持廣播以及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
 * 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, 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 = {
    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, 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

示例帶有結(jié)構(gòu)化數(shù)組dtype參數(shù)的NumPy ufunc

此示例顯示如何為結(jié)構(gòu)化數(shù)組dtype創(chuàng)建ufunc。在這個(gè)例子中,我們展示了一個(gè)簡(jiǎn)單的ufunc,用于添加兩個(gè)帶有dtype'u8,u8,u8'的數(shù)組。該過(guò)程與其他示例略有不同,因?yàn)檎{(diào)用PyUFunc_FromFuncAndData

不會(huì)為自定義dtypes和結(jié)構(gòu)化數(shù)組dtypes完全注冊(cè)u(píng)func。我們還需要調(diào)用
PyUFunc_RegisterLoopForDescr

完成設(shè)置ufunc。

我們只提供C代碼,因?yàn)閟etup.py文件與一種dtype的NumPy ufunc示例的setup.py文件完全相同,只有一行。

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"


/*
 * 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, 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 */
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 void *data[1] = {NULL};

#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();

    /* 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);
#if defined(NPY_PY3K)
    return m;
#endif
}

返回的ufunc對(duì)象是一個(gè)可調(diào)用的Python對(duì)象。它應(yīng)該放在一個(gè)(模塊)字典中,其名稱與ufunc-creation例程的name參數(shù)中使用的字典相同。以下示例是從umath模塊改編而來(lái)的

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);
...

作者:柯廣的網(wǎng)絡(luò)日志 ? 編寫您自己的ufunc

微信公眾號(hào):Java大數(shù)據(jù)與數(shù)據(jù)倉(cāng)庫(kù)