如何扩展 NumPy#

静态且重复的东西是乏味的.动态的东西
而随机是令人困惑的.在这之间存在着艺术.
John A. Locke
科学是一个微分方程.宗教是一个边界条件.
Alan Turing

编写一个扩展模块#

虽然 ndarray 对象旨在允许在 Python 中进行快速计算,但它也被设计为通用且满足各种计算需求.因此,如果绝对速度至关重要,那么精心编写的、针对您的应用程序和硬件的编译循环是无法替代的.这也是 numpy 包含 f2py 的原因之一,以便为将(简单)C/C++ 和(任意)Fortran 代码直接链接到 Python 提供易于使用的机制.鼓励您使用并改进此机制.本节的目的不是记录此工具,而是记录此工具所依赖的编写扩展模块的更基本步骤.

当一个扩展模块被编写、编译并安装到 Python 路径(sys.path)中的某个位置时,该代码可以像标准 python 文件一样导入到 Python 中.它将包含在 C 代码中定义和编译的对象和方法.在 Python 中执行此操作的基本步骤有详细的文档记录,您可以在 www.python.org 上在线提供的 Python 文档中找到更多信息.

除了 Python C-API 之外,NumPy 还提供了一个完整且丰富的 C-API,允许在 C 级别进行复杂的操作.然而,对于大多数应用程序,通常只会使用少数几个 API 调用.例如,如果您只需要提取一个指向内存的指针以及一些形状信息传递给另一个计算例程,那么您将使用与创建新的数组类型或为 ndarrays 添加新数据类型非常不同的调用.本章节记录了最常用的 API 调用和宏.

必需的子程序#

在你的C代码中,必须定义一个函数,以便Python将其用作扩展模块.该函数必须命名为init{name},其中{name}是Python中模块的名称.此函数必须声明为对外部代码可见.除了添加你想要的方法和常量外,此子程序还必须包含类似``import_array()``和/或``import_ufunc()``的调用,具体取决于所需的C-API.忘记放置这些命令将表现为在实际调用任何C-API子程序时出现丑陋的分段错误(崩溃).实际上,在一个文件中可以有多个init{name}函数,在这种情况下,该文件将定义多个模块.然而,要正确实现这一点有一些技巧,这里不涉及.

一个最小的 init{name} 方法看起来像:

PyMODINIT_FUNC
init{name}(void)
{
   (void)Py_InitModule({name}, mymethods);
   import_array();
}

mymethods 必须是一个数组(通常是静态声明的),包含 PyMethodDef 结构,这些结构包含方法名、实际的 C 函数、一个变量指示方法是否使用关键字参数以及文档字符串.这些将在下一节中解释.如果你想向模块添加常量,那么你可以存储从 Py_InitModule 返回的值,这是一个模块对象.向模块添加项目最通用的方法是使用 PyModule_GetDict(module) 获取模块字典.有了模块字典,你可以手动向模块添加任何你喜欢的内容.向模块添加对象的一个更简单的方法是使用三个额外的 Python C-API 调用之一,这些调用不需要单独提取模块字典.这些在 Python 文档中有记录,但为了方便,这里重复一下:

int PyModule_AddObject(PyObject *module, char *name, PyObject *value)#
int PyModule_AddIntConstant(PyObject *module, char *name, long value)#
int PyModule_AddStringConstant(PyObject *module, char *name, char *value)#

这三个函数都需要 module 对象(Py_InitModule 的返回值).*name* 是一个字符串,用于标记模块中的值.根据调用的函数,*value* 参数可以是通用对象(PyModule_AddObject 会窃取对其的引用)、整数常量或字符串常量.

定义函数#

传递给 Py_InitModule 函数的第二个参数是一个结构体,它使得在模块中定义函数变得容易.在上面的例子中,mymethods 结构体应该在文件的较早部分定义(通常在 init{name} 子程序之前),如下:

static PyMethodDef mymethods[] = {
    { nokeywordfunc,nokeyword_cfunc,
      METH_VARARGS,
      Doc string},
    { keywordfunc, keyword_cfunc,
      METH_VARARGS|METH_KEYWORDS,
      Doc string},
    {NULL, NULL, 0, NULL} /* Sentinel */
}

mymethods 数组中的每个条目都是一个 PyMethodDef 结构,包含 1) Python 名称,2) 实现该函数的 C 函数,3) 指示是否接受该函数的关键字的标志,以及 4) 该函数的文档字符串.通过向该表添加更多条目,可以为一个模块定义任意数量的函数.最后一个条目必须全为 NULL,如所示,作为哨兵.Python 查找此条目以知道该模块的所有函数都已定义.

完成扩展模块必须做的最后一件事是实际编写执行所需功能的代码.有两种函数:不接受关键字参数的函数和接受关键字参数的函数.

没有关键字参数的函数#

不接受关键字参数的函数应写为:

static PyObject*
nokeyword_cfunc (PyObject *dummy, PyObject *args)
{
    /* convert Python arguments */
    /* do function */
    /* return something */
}

dummy 参数在此上下文中未被使用,可以安全地忽略.*args* 参数包含作为元组传递给函数的所有参数.你可以在这一点上做任何你想做的事情,但通常管理输入参数最简单的方法是调用 PyArg_ParseTuple (args, format_string, addresses_to_C_variables…) 或 PyArg_UnpackTuple (tuple, “name”, min, max, …).关于如何使用第一个函数的详细描述包含在 Python C-API 参考手册的第 5.5 节(解析参数和构建值)中.你应该特别注意使用转换器函数在 Python 对象和 C 对象之间转换的 “O&” 格式.所有其他格式函数可以(大部分)被认为是这一通用规则的特例.NumPy C-API 中定义了几个转换器函数,可能会有所帮助.特别是,:c:func:PyArray_DescrConverter 函数非常有用,支持任意数据类型规范.此函数将任何有效的数据类型 Python 对象转换为 PyArray_Descr* 对象.记得传入应填充的 C 变量的地址.

在NumPy源代码中有很多如何使用 PyArg_ParseTuple 的例子.标准用法如下:

PyObject *input;
PyArray_Descr *dtype;
if (!PyArg_ParseTuple(args, "OO&", &input,
                      PyArray_DescrConverter,
                      &dtype)) return NULL;

记住在使用 “O” 格式字符串时,你会得到对象的一个 借用 引用,这很重要.然而,转换器函数通常需要某种形式的内存处理.在这个例子中,如果转换成功,*dtype* 将持有对一个 PyArray_Descr* 对象的新引用,而 input 将持有借用引用.因此,如果这种转换与其他转换(比如转换为整数)混合,并且数据类型转换成功但整数转换失败,那么在返回之前你需要释放对数据类型对象的引用计数.一种典型的做法是在调用 PyArg_ParseTuple 之前将 dtype 设置为 NULL,然后在返回之前对 dtype 使用 Py_XDECREF.

在处理输入参数之后,实际执行工作的代码被编写(可能根据需要调用其他函数).C函数的最后一步是返回某些内容.如果遇到错误,则应返回 ``NULL``(确保实际上已设置错误).如果不需要返回任何内容,则递增 Py_None 并返回它.如果应返回单个对象,则返回它(确保首先拥有对该对象的引用).如果应返回多个对象,则需要返回一个元组.:c:func:Py_BuildValue (format_string, c_variables…) 函数使得从C变量构建Python对象的元组变得容易.特别注意格式字符串中 ‘N’ 和 ‘O’ 之间的区别,否则很容易创建内存泄漏.’O’ 格式字符串递增与其对应的 PyObject* C变量的引用计数,而 ‘N’ 格式字符串窃取对相应 PyObject* C变量的引用.如果你已经为对象创建了一个引用,并且只想将该引用提供给元组,则应使用 ‘N’.如果你只有一个借用的对象引用,并且需要创建一个引用以提供给元组,则应使用 ‘O’.

带有关键字参数的函数#

这些函数与没有关键字参数的函数非常相似.唯一的区别是函数签名是:

static PyObject*
keyword_cfunc (PyObject *dummy, PyObject *args, PyObject *kwds)
{
...
}

kwds 参数包含一个 Python 字典,其键是关键字参数的名称,其值是对应的关键字参数值.这个字典可以按照你认为合适的方式进行处理.然而,最简单的方法是用对 PyArg_ParseTupleAndKeywords (args, kwds, format_string, char *kwlist[], addresses…) 函数的调用替换 PyArg_ParseTuple (args, format_string, addresses…) 函数.这个函数的 kwlist 参数是一个以 NULL 结尾的字符串数组,提供预期的关键字参数.格式字符串中的每个条目都应该有一个字符串.使用这个函数将在传递无效的关键字参数时引发 TypeError.

有关此函数的更多帮助,请参见 Python 文档中扩展和嵌入教程的第 1.8 节(扩展函数的键参数).

引用计数#

编写扩展模块时最大的困难是引用计数.这是f2py、weave、Cython、ctypes等工具受欢迎的重要原因.如果你错误处理引用计数,可能会遇到从内存泄漏到段错误的问题.我知道的唯一正确处理引用计数的策略是血、汗和泪.首先,你必须牢记每个Python变量都有一个引用计数.然后,你必须准确理解每个函数对你的对象引用计数的影响,以便在需要时正确使用DECREF和INCREF.引用计数真的可以考验你对编程工作的耐心和勤奋程度.尽管描述得很严峻,但大多数引用计数的情况相当直接,最常见的困难是在由于某些错误过早退出例程之前没有对对象使用DECREF.其次是常见的错误,即没有拥有传递给会窃取引用(例如:c:func:PyTuple_SET_ITEM,以及大多数接受:c:type:`PyArray_Descr`对象的函数)的函数的对象的引用.

通常,当变量被创建或作为某些函数的返回值时,你会得到一个新的变量引用(然而,也有一些显著的例外,比如从元组或字典中获取一个项目).当你拥有这个引用时,你有责任确保在变量不再需要时调用 Py_DECREF (var)(并且没有其他函数”偷走”它的引用).此外,如果你将一个Python对象传递给一个会”偷走”引用的函数,那么你需要确保你拥有它(或者使用 Py_INCREF 来获取你自己的引用).你还会遇到借用引用的概念.借用引用的函数不会改变对象的引用计数,也不期望”持有”引用.它只是暂时使用对象.当你使用 PyArg_ParseTuplePyArg_UnpackTuple 时,你会收到元组中对象的借用引用,并且不应该在你的函数中改变它们的引用计数.通过实践,你可以学会正确地进行引用计数,但一开始可能会感到沮丧.

引用计数错误的一个常见来源是 Py_BuildValue 函数.请仔细注意 ‘N’ 格式字符和 ‘O’ 格式字符之间的区别.如果你在你的子程序中创建了一个新对象(例如一个输出数组),并且你将它作为返回值的元组传递回去,那么你最有可能应该在 Py_BuildValue 中使用 ‘N’ 格式字符.’O’ 字符会增加引用计数一次.这将使调用者对一个全新的数组有两个引用计数.当变量被删除并且引用计数减少一次时,仍然会有那个额外的引用计数,数组将永远不会被释放.你将会有一个由引用计数引起的内存泄漏.使用 ‘N’ 字符将避免这种情况,因为它将返回给调用者一个引用计数为单个的对象(在元组内).

处理数组对象#

大多数为 NumPy 编写的扩展模块都需要访问 ndarray 对象(或其子类之一)的内存.最简单的方法不需要你了解太多关于 NumPy 内部的知识.方法是

  1. 确保你处理的是一个行为良好的数组(对齐、机器字节顺序和单段),具有正确的类型和维数.

    1. 通过使用 PyArray_FromAny 或基于它的宏将其从某个 Python 对象转换.

    2. 通过使用 PyArray_NewFromDescr 或基于它的一个更简单的宏或函数,构造一个你所需形状和类型的新 ndarray.

  2. 获取数组的形状和指向其实际数据的指针.

  3. 将数据和形状信息传递给实际执行计算的子程序或其他代码部分.

  4. 如果你正在编写算法,那么我建议你使用数组中包含的步幅信息来访问数组的元素(PyArray_GetPtr 宏使这变得无痛).然后,你可以放宽你的要求,以免强制要求单段数组和可能导致的数拷贝.

这些子主题中的每一个都在以下小节中进行了介绍.

转换一个任意的序列对象#

从任何可以转换为数组的Python对象获取数组的主程序是 PyArray_FromAny.这个函数非常灵活,有许多输入参数.几个宏使基本函数的使用更加容易.:c:func:PyArray_FROM_OTF 可以说是最常用的宏之一,适用于最常见的用途.它允许你将任意Python对象转换为特定内置数据类型的数组(例如浮点数),同时指定一组特定的要求(例如连续、对齐和可写).语法是

PyArray_FROM_OTF

从任何可以转换为数组的 Python 对象 obj 返回一个 ndarray.返回数组的维数由对象决定.返回数组的期望数据类型在 typenum 中提供,该类型应为枚举类型之一.返回数组的要求可以是任何标准数组标志的组合.这些参数在下面有更详细的解释.成功时,您会收到对数组的新引用.失败时,返回 NULL 并设置异常.

obj

对象可以是任何可以转换为 ndarray 的 Python 对象.如果对象已经是满足要求的 ndarray(或其子类),则返回一个新的引用.否则,构造一个新数组.除非使用数组接口,否则 obj 的内容会被复制到新数组中,这样数据就不必复制.可以转换为数组的对象包括:1) 任何嵌套的序列对象,2) 任何暴露数组接口的对象,3) 任何具有 __array__ 方法的对象(该方法应返回一个 ndarray),以及 4) 任何标量对象(变成一个零维数组).符合要求的 ndarray 子类将被传递.如果你想确保一个基类 ndarray,那么在要求标志中使用 NPY_ARRAY_ENSUREARRAY.只有在必要时才会进行复制.如果你想保证复制,那么在要求标志中传递 NPY_ARRAY_ENSURECOPY.

typenum

枚举类型之一或 NPY_NOTYPE 如果数据类型应从对象本身确定.可以使用基于C的名称:

或者,可以使用平台支持的位宽名称.例如:

对象只有在不会丢失精度的情况下才会被转换为所需类型.否则将返回 NULL 并引发错误.在需求标志中使用 NPY_ARRAY_FORCECAST 以覆盖此行为.

requirements

ndarray 的内存模型允许在每个维度中任意跨步以前进到数组的下一个元素.然而,通常情况下,你需要与期望 C 连续或 Fortran 连续内存布局的代码进行接口.此外,ndarray 可能会未对齐(元素的地址不是元素大小的整数倍),这可能会导致你的程序崩溃(或至少工作得更慢),如果你尝试解引用指向数组数据的指针.这两个问题都可以通过将 Python 对象转换为对你特定使用更”行为良好”的数组来解决.

requirements 标志允许指定可以接受哪种数组.如果传入的对象不满足这些要求,则会创建一个副本,以使返回的对象将满足这些要求.这些 ndarray 可以使用非常通用的内存指针.此标志允许指定返回数组对象所需的属性.所有标志在详细的 API 章节中都有解释.最常用的标志是 NPY_ARRAY_IN_ARRAYNPY_ARRAY_OUT_ARRAYNPY_ARRAY_INOUT_ARRAY:

NPY_ARRAY_IN_ARRAY

此标志对于必须按 C 连续顺序和对齐的数组非常有用.这些类型的数组通常是某些算法的输入数组.

NPY_ARRAY_OUT_ARRAY

此标志用于指定一个按 C 连续顺序排列、对齐且可以写入的数组.这样的数组通常作为输出返回(尽管通常这样的输出数组是从头创建的).

NPY_ARRAY_INOUT_ARRAY

此标志用于指定一个将同时用于输入和输出的数组.在接口例程结束时调用 Py_DECREF 之前,必须调用 PyArray_ResolveWritebackIfCopy 以将临时数据写回传入的原始数组.使用 NPY_ARRAY_WRITEBACKIFCOPY 标志要求输入对象已经是一个数组(因为其他对象不能以这种方式自动更新).如果发生错误,请对设置了这些标志的数组使用 PyArray_DiscardWritebackIfCopy (obj).这将使底层基础数组可写,而不会导致内容被复制回原始数组.

其他有用的标志可以作为附加要求进行或操作的有:

NPY_ARRAY_FORCECAST

转换为所需类型,即使这样做会丢失信息.

NPY_ARRAY_ENSURECOPY

确保生成的数组是原始数组的副本.

NPY_ARRAY_ENSUREARRAY

确保生成的对象是一个实际的 ndarray,而不是子类.

备注

一个数组是否字节交换取决于数组的数据类型.本地字节顺序的数组总是由 PyArray_FROM_OTF 请求,因此在需求参数中不需要 NPY_ARRAY_NOTSWAPPED 标志.也没有办法从这个例程中获取字节交换的数组.

创建一个新的 ndarray#

通常,必须从扩展模块代码中创建新数组.可能需要一个输出数组,并且你不希望调用者必须提供它.可能只需要一个临时数组来保存中间计算.无论需要什么,都有简单的方法来获取所需数据类型的 ndarray 对象.最通用的函数是 PyArray_NewFromDescr.所有数组创建函数都通过这个被大量重用的代码.由于其灵活性,使用起来可能有些混乱.因此,存在更简单的形式,更容易使用.这些形式是 PyArray_SimpleNew 系列函数的一部分,通过为常见用例提供默认值来简化接口.

获取 ndarray 内存和访问 ndarray 的元素#

如果 obj 是一个 ndarray (PyArrayObject*),那么 ndarray 的数据区域由 void* 指针 PyArray_DATA (obj) 或 char* 指针 PyArray_BYTES (obj) 指向.请记住,(通常情况下)这个数据区域可能没有根据数据类型对齐,它可能表示字节交换的数据,和/或它可能不是可写的.如果数据区域是对齐的并且是本地字节顺序,那么如何获取数组中的特定元素仅由 npy_intp 变量数组 PyArray_STRIDES (obj) 决定.特别是,这个 c 整数数组显示了必须添加到当前元素指针的 字节 数,以在每个维度中获取下一个元素.对于小于 4 维的数组,有 PyArray_GETPTR{k} (obj, …) 宏,其中 {k} 是整数 1, 2, 3 或 4,这使得使用数组步幅更容易.参数 …. 表示 {k} 个非负整数索引到数组中.例如,假设 E 是一个 3 维 ndarray.指向元素 E[i,j,k] 的 (void*) 指针通过 PyArray_GETPTR3 (E, i, j, k) 获得.

如前所述,C 风格的连续数组和 Fortran 风格的连续数组具有特定的步幅模式.两个数组标志 (NPY_ARRAY_C_CONTIGUOUSNPY_ARRAY_F_CONTIGUOUS) 指示特定数组的步幅模式是否匹配 C 风格的连续或 Fortran 风格的连续或两者都不匹配.是否匹配标准的 C 或 Fortran 步幅模式可以使用 PyArray_IS_C_CONTIGUOUS (obj) 和 PyArray_ISFORTRAN (obj) 分别进行测试.大多数第三方库期望连续数组.但是,通常支持通用步幅并不困难.我鼓励您在可能的情况下在您自己的代码中使用步幅信息,并将单段要求保留用于包装第三方代码.使用 ndarray 提供的步幅信息而不是要求连续步幅减少了必须进行的复制.

示例#

以下示例展示了如何编写一个接受两个输入参数(将转换为数组)和一个输出参数(必须是数组)的包装器.该函数返回 None 并更新输出数组.注意 NumPy v1.14 及以上版本中对 WRITEBACKIFCOPY 语义的更新使用.

static PyObject *
example_wrapper(PyObject *dummy, PyObject *args)
{
    PyObject *arg1=NULL, *arg2=NULL, *out=NULL;
    PyObject *arr1=NULL, *arr2=NULL, *oarr=NULL;

    if (!PyArg_ParseTuple(args, "OOO!", &arg1, &arg2,
        &PyArray_Type, &out)) return NULL;

    arr1 = PyArray_FROM_OTF(arg1, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
    if (arr1 == NULL) return NULL;
    arr2 = PyArray_FROM_OTF(arg2, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
    if (arr2 == NULL) goto fail;
#if NPY_API_VERSION >= 0x0000000c
    oarr = PyArray_FROM_OTF(out, NPY_DOUBLE, NPY_ARRAY_INOUT_ARRAY2);
#else
    oarr = PyArray_FROM_OTF(out, NPY_DOUBLE, NPY_ARRAY_INOUT_ARRAY);
#endif
    if (oarr == NULL) goto fail;

    /* code that makes use of arguments */
    /* You will probably need at least
       nd = PyArray_NDIM(<..>)    -- number of dimensions
       dims = PyArray_DIMS(<..>)  -- npy_intp array of length nd
                                     showing length in each dim.
       dptr = (double *)PyArray_DATA(<..>) -- pointer to data.

       If an error occurs goto fail.
     */

    Py_DECREF(arr1);
    Py_DECREF(arr2);
#if NPY_API_VERSION >= 0x0000000c
    PyArray_ResolveWritebackIfCopy(oarr);
#endif
    Py_DECREF(oarr);
    Py_INCREF(Py_None);
    return Py_None;

 fail:
    Py_XDECREF(arr1);
    Py_XDECREF(arr2);
#if NPY_API_VERSION >= 0x0000000c
    PyArray_DiscardWritebackIfCopy(oarr);
#endif
    Py_XDECREF(oarr);
    return NULL;
}