编写自己的 ufunc#
我拥有力量！
— 希曼 (He-Man)
创建新的通用函数 (ufunc)#
在阅读本节之前，建议先阅读/浏览 扩展和嵌入 Python 解释器 第 1 节中的教程以及 如何扩展 NumPy，以熟悉 Python C 扩展的基础知识。
umath 模块是一个计算机生成的 C 模块，它创建了许多 ufunc。它提供了大量关于如何创建通用函数的示例。创建一个利用 ufunc 机制的自定义 ufunc 并不困难。假设你有一个希望对输入进行逐元素操作的函数。通过创建新的 ufunc，你将获得一个能够处理以下特性的函数：
广播 (Broadcasting)
N 维循环 (N-dimensional looping)
自动类型转换（内存占用极小）
可选的输出数组
创建自己的 ufunc 并不难。你只需要为要支持的每种数据类型提供一个一维循环即可。每个一维循环都必须具有特定的签名，且只能使用固定大小数据类型的 ufunc。下面给出了用于创建处理内置数据类型的新 ufunc 的函数调用方法。对于用户自定义数据类型，注册 ufunc 则需要使用不同的机制。
在接下来的几节中，我们将提供示例代码，你可以轻松修改这些代码来创建自己的 ufunc。这些示例是 logit 函数的逐步完善或复杂化版本，logit 是统计建模中常见的函数。Logit 也很有趣，因为得益于 IEEE 标准（特别是 IEEE 754）的魔力，下面创建的所有 logit 函数都会自动具备以下行为。
>>> logit(0)
-inf
>>> logit(1)
inf
>>> logit(2)
nan
>>> logit(-2)
nan
这非常棒，因为函数编写者不必手动处理 inf 或 nan 的传播。
非 ufunc 扩展示例#
为了进行比较并启发读者，我们提供了一个不使用 numpy 的简单 logit C 扩展实现。
为此，我们需要三个文件。第一个是包含实际代码的 C 文件，另外两个是描述如何创建模块的项目文件。
#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; }
创建模块的步骤与 Python 包相同：创建一个定义构建后端的 pyproject.toml 文件，然后为该后端创建另一个描述如何编译代码的文件。对于后端，我们推荐 meson-python，因为 numpy 本身也使用它，但下面我们也展示了如何使用较旧的 setuptools。
示例 pyproject.toml 和 meson.build。
[project]
name = "spam"
version = "0.1"
[build-system]
requires = ["meson-python"]
build-backend = "mesonpy"
project('spam', 'c')
py = import('python').find_installation()
sources = files('spammodule.c')
extension_module = py.extension_module(
'spam',
sources,
install: true,
)
示例 pyproject.toml 和 setup.py。
[project]
name = "spam"
version = "0.1"
[build-system]
requires = ["setuptools"]
build-backend = "setuptools.build_meta"
from setuptools import setup, Extension
spammodule = Extension('spam', sources=['spammodule.c'])
setup(name='spam', version='1.0',
ext_modules=[spammodule])
使用以上任一方式，都可以通过以下命令构建并安装 spam 包：
pip install .
一旦 spam 模块导入到 Python 中，就可以通过 spam.logit 调用 logit。注意，上面使用的函数不能直接应用于 numpy 数组。要做到这一点，必须在其上调用 numpy.vectorize。例如
>>> 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 示例#
为了简单起见，我们提供一个针对单一 dtype（即 'f8' double）的 ufunc。和上一节一样，我们首先给出 .c 文件，然后给出用于创建包含该 ufunc 的 npufunc 模块的文件。
代码中对应 ufunc 实际计算的位置标记为 /* BEGIN main ufunc computation */ 和 /* END main ufunc computation */。这两行之间的代码是创建自己的 ufunc 时主要需要修改的部分。
#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 const 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; }
对于创建模块所需的文件，与我们之前的示例相比，主要区别在于现在需要声明对 numpy 的依赖。
示例 pyproject.toml 和 meson.build。
[project]
name = "npufunc"
dependencies = ["numpy"]
version = "0.1"
[build-system]
requires = ["meson-python", "numpy"]
build-backend = "mesonpy"
project('npufunc', 'c')
py = import('python').find_installation()
np_dep = dependency('numpy')
sources = files('single_type_logit.c')
extension_module = py.extension_module(
'npufunc',
sources,
dependencies: [np_dep],
install: true,
)
示例 pyproject.toml 和 setup.py。
[project]
name = "npufunc"
dependencies = ["numpy"]
version = "0.1"
[build-system]
requires = ["setuptools", "numpy"]
build-backend = "setuptools.build_meta"
from setuptools import setup, Extension
from numpy import get_include
npufunc = Extension('npufunc',
sources=['single_type_logit.c'],
include_dirs=[get_include()])
setup(name='npufunc', version='1.0', ext_modules=[npufunc])
安装上述内容后，可以按如下方式导入并使用：
>>> import numpy as np
>>> import npufunc
>>> npufunc.logit(0.5)
np.float64(0.0)
>>> a = np.linspace(0, 1, 5)
>>> npufunc.logit(a)
array([ -inf, -1.09861229, 0. , 1.09861229, inf])
多个 dtype 的 NumPy ufunc 示例#
我们现在将其扩展为完整的 logit ufunc，其中包含针对 float、double 和 long double 的内部循环。在这里，我们可以使用与上述相同的构建文件，只需将源文件从 single_type_logit.c 更改为 multi_type_logit.c 即可。
代码中对应 ufunc 实际计算的位置标记为 /* BEGIN main ufunc computation */ 和 /* END main ufunc computation */。这两行之间的代码是创建自己的 ufunc 时主要需要修改的部分。
#define PY_SSIZE_T_CLEAN #include <Python.h> #include "numpy/ndarraytypes.h" #include "numpy/ufuncobject.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; } } /*This gives pointers to the above functions*/ PyUFuncGenericFunction funcs[3] = {&float_logit, &double_logit, &long_double_logit}; static const char types[6] = {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; }
安装上述内容后，可以按如下方式导入并使用。
>>> import numpy as np
>>> import npufunc
>>> npufunc.logit(0.5)
np.float64(0.0)
>>> a = np.linspace(0, 1, 5, dtype=np.float32)
>>> npufunc.logit(a)
<python-input-4>:1: RuntimeWarning: divide by zero encountered in logit
array([ -inf, -1.0986123, 0. , 1.0986123, inf],
dtype=float32)
注意
由于 float16（半精度）具有非标准 C 表示形式和转换要求，在自定义 ufunc 中支持它更为复杂。上述代码可以处理 float16 输入，但会通过将其转换为 float32 来实现。结果也将是 float32，但你可以通过传入合适的输出（如 npufunc.logit(a, out=np.empty_like(a))）将其转回 float16。有关实际 float16 循环的示例，请参阅 numpy 源代码。
多参数/多返回值 NumPy ufunc 示例#
创建具有多个参数的 ufunc 并不难。在这里，我们修改了 logit ufunc 的代码，以计算 (A * B, logit(A * B))。为简单起见，我们仅为 double 类型创建一个循环。
我们再次只提供 C 代码，因为创建模块所需的文件与之前相同，只需将源文件名替换为 multi_arg_logit.c。
下面给出了 C 文件。生成的 ufunc 接受两个参数 A 和 B，并返回一个元组，其中第一个元素是 A * B，第二个元素是 logit(A * B)。注意，它自动支持广播以及 ufunc 的所有其他属性。
#define PY_SSIZE_T_CLEAN #include <Python.h> #include "numpy/ndarraytypes.h" #include "numpy/ufuncobject.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 const 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 示例#
此示例展示了如何为结构化数组 dtype 创建 ufunc。在此示例中，我们展示了一个用于对两个 dtype 为 'u8,u8,u8' 的数组进行加法的平凡 ufunc。该过程与其他示例略有不同，因为调用 PyUFunc_FromFuncAndData 无法为自定义 dtype 和结构化数组 dtype 注册 ufunc。我们需要额外调用 PyUFunc_RegisterLoopForDescr 来完成 ufunc 的设置。
我们只提供 C 代码，因为构建模块所需的文件与之前完全相同，只是源文件现在是 add_triplet.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; } } static struct PyModuleDef moduledef = { PyModuleDef_HEAD_INIT, "npufunc", NULL, -1, StructUfuncTestMethods, NULL, NULL, NULL, NULL }; PyMODINIT_FUNC PyInit_npufunc(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((PyUFuncObject *)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; }
使用示例
>>> import npufunc
>>> import numpy as np
>>> a = np.array([(1, 2, 3), (4, 5, 6)], "u8,u8,u8")
>>> npufunc.add_triplet(a, a)
array([(2, 4, 6), (8, 10, 12)],
dtype=[('f0', '<u8'), ('f1', '<u8'), ('f2', '<u8')])