如何在 C 扩展中访问 numpy 多维数组?

How does one acces numpy multidimensionnal array in c extensions?

为了理解 C 扩展中对 numpy 数组的访问,我已经苦苦挣扎了几天,但我很难理解文档。

编辑:这是我想移植到 c 的代码(grav 函数)

import numpy as np

def grav(p, M):
    G = 6.67408*10**-2     # m³/s²T
    l = len(p[0])
    a = np.empty(shape=(2, l))
    a[:, 0] = 0
    for b in range(1, l):
        # computing the distance between body #b and all previous
        d = p[:, b:b+1] - p[:, :b]
        d2 = (d*d).sum(axis=0)
        d2[d2==0] = 1
        XXX = G * d * d2**(-1.5)

        # computing Newton formula : 
        # acceleration undergone by b from all previous
        a[:, b] = -(M[None, :b] * XXX).sum(axis=1)

        # computing Newton formula : adding for each previous,
        # acceleration undergone by from b
        a[:, :b] += M[b] * XXX
    return a

system_p = np.array([[1., 2., 3., 4., 5., 6., 7., 9., 4., 0.],
                     [3., 2., 5., 6., 3., 5., 6., 3., 5., 8.]])
system_M = np.array( [3., 5., 1., 2., 4., 5., 4., 5., 6., 8.])

system_a = grav(system_p, system_M)
for i in range(len(system_p[0])):
    print('body {:2} mass = {}(ton), position = {}(m), '
          'acceleration = [{:8.4f} {:8.4f}](m/s²)'.format(i,
              system_M[i], system_p[:, i], system_a[0, i], system_a[1, i]))

我找到了 here 一个使用简单迭代器的简单示例。它工作得很好,但它不会超出一维数组范围,并且不会提供有关当数组具有多个维度并且您希望迭代其中的一个子集或者您希望指定顺序(例如 row-wise/column-wise) 你想迭代它们。也就是说,使用这种方法,您可以遍历多维数组,但只能以一种方式进行。

编辑:NpyIter_MultiNew 似乎与多维迭代无关,但与一次迭代多个数组有关。
the documentation中查找,我发现了这个函数:

NpyIter* NpyIter_MultiNew(npy_intp nop,
                          PyArrayObject** op,
                          npy_uint32 flags,
                          NPY_ORDER order,
                          NPY_CASTING casting,
                          npy_uint32* op_flags,
                          PyArray_Descr** op_dtypes)

这可能是我需要的,但我连描述的第一句都看不懂:

Creates an iterator for broadcasting the nop array objects provided in op, [...]

这是什么 «nop 数组对象»?它与op参数有什么关系?我知道我不是以英语为母语的人,但我仍然觉得这个文档可以比它更清晰。

然后我发现了其他资源,如 this,它们似乎采用了完全不同的方法(没有迭代器 - 所以,我想是手动迭代),但它甚至在不更正它的情况下无法编译(仍在工作虽然在上面)。

所以,请问这里有没有人有这方面的经验,谁能提供简单的例子说明如何做到这一点?

简介

可能最好的解决方法是在 Python 中创建迭代器并在那里进行试验。这会很慢,但它会确认你所做的是正确的。然后使用 NpyIter_AdvancedNew,尽可能使用默认参数。

恐怕我自己还没有真正将它翻译成 C 代码 - 这对我来说花了太长时间。因此,我建议您不要接受这个答案,因为它实际上只是给出了一个起点。

我的猜测是,考虑到编写 C 代码的工作量,任何性能改进都会令人失望(特别是因为我猜想编写快速代码需要更深层次的理解)。在答案的最后,我提出了一些我推荐的更简单的替代方案,而不是使用 C API.

示例

我已经从您的代码中翻译了几行作为示例:


d = p[:, b:b+1] - p[:, :b]

变成

with np.nditer([p[:,b],p[:,:b],None],
               op_axes=[[0,-1],[0,1],[0,1]]) as it:
    for x,y,z in it:
        z[...] = x - y
    d = it.operands[2]

请注意,您需要事先对 p 数组进行切片。我已将其中一个数组作为 None 传递。这转化为 C 中的 NULL 指针,意味着创建的数组具有适当的大小(使用标准广播规则)。

op_axes而言,第一个数组只有一维,所以我说"iterate over axis 0 first; there isn't an axis 1"。第二个和第三个数组是两个 D 所以我说 "iterate over axis 0 then axis 1".

在Python中它会自动推断出op_flags。我不知道它是否会在 C 中执行此操作。如果不是,那么它们应该是:

npy_uint32 op_flags[] = { NPY_ITER_READONLY,
               NPY_ITER_READONLY,
               NPY_ITER_WRITEONLY | NPY_ITER_ALLOCATE };

最重要的是分配了第三个轴。

我的观点是您希望在 C 中将 op_dtypes 指定为

{ PyArray_DescrFromType(NPY_DOUBLE), 
  PyArray_DescrFromType(NPY_DOUBLE), 
  NULL }

强制数组为正确的类型(第三个分配数组的类型可以从两个输入中算出)。这样你就可以在 C.

中将你的数据指针转换为 double*

行:

d2 = (d*d).sum(axis=0)

转换为

with np.nditer([d,None],
                   flags=['reduce_ok'],
                   op_flags=[['readonly'],['readwrite','allocate']],
                   op_axes=[[1,0],[0,-1]]) as it:
    it.operands[1][...] = 0
    for x,z in it:
        z[...] += x*x
    d2 = it.operands[1]

最重要的区别是它是一个缩减(第二个输出数组小于输入,因为其中一个轴是一个总和)。因此我们将 'reduce_ok' 作为标志传递。

第二个数组只有一个轴,所以op_axes[0, -1]。轴是第二个数组匹配第一个数组的轴 1,因此第一个数组的 op_axes 设置为 [1, 0].

当翻译成 C 时,行 it.operands[1][...] = 0 变得更复杂:

Note that if you want to do a reduction on an automatically allocated output, you must use NpyIter_GetOperandArray to get its reference, then set every value to the reduction unit before doing the iteration loop.

在 C 中,我可能会先将 d2 分配为零数组,然后将其传递给迭代器。

备选方案

为此编写 C API 代码涉及大量代码、错误检查、引用计数等。虽然它应该是 "simple" 翻译(nditer API 在 C 和 Python 中基本相同)这并不容易。

如果我是你,你会考虑使用一些标准工具来加速 Python,例如 Numba、NumExpr 或 Cython。 Numba 和 NumExpr 是 just-in-time 编译器,它们可以做一些事情,比如避免分配中间数组。 Cython 是一种 "Python-like" 语言,您可以在其中指定类型。显示翻译成 Cython 的前几部分:

def grav3(double[:,:] p, M):
    G = 6.67408e-2     # m³/s²T
    cdef int l = p.shape[1]
    a = np.empty(shape=(2, l))
    a[:, 0] = 0
    cdef double[:,::1] d
    cdef double[::1] d2
    cdef Py_ssize_t b, i, j
    for b in range(1, l):
        # computing the distance between body #b and all previous
        d = np.empty((p.shape[0],b))
        for i in range(d.shape[0]):
            for j in range(b):
                d[i,j] = p[i,b] - p[i,j]

        d2 = np.zeros((d.shape[1]))
        for j in range(d.shape[1]):                
            for i in range(d.shape[0]):
                d2[j] += d[i,j]*d[i,j]
            if d2[j] == 0:
                d2[j] = 1

这里我指定了一些数组为一维或二维双数组 double[:]double[:,:]。然后我明确地写出了循环,这避免了创建中间体。


Cython 在获取 PyArray_DATA 的地方生成 C 代码,然后使用 PyArray_STRIDES 计算出在二维数组中访问的位置。您可能会发现这比使用迭代器更容易。您可以检查 Cython 生成的代码,看看它是如何做到的。 Numpy 中也有 PyArray_GetPtr functions 可以执行这种访问,您可能会发现它比使用迭代器更容易。

好的,我终于成功了。由于最大的困难是找到好的介绍material,我留下一个例子作为示例代码。以下是我使用或考虑使用的 API 的函数⁽¹⁾:
(1): 描述 in the documentation

PyArray_Descr *PyArray_DESCR(PyArrayObject* arr)¶

是一个宏,它 "returns" C PyArrayObject 结构的 PyArray_Descr *descr 字段,它是指向数组的 dtype 属性 的指针。

int PyArray_NDIM(PyArrayObject *arr)

是一个宏,其中 "returns" C PyArrayObject 结构的 int nd 字段,其中包含数组的维数。

npy_intp *PyArray_DIMS(PyArrayObject *arr)
npy_intp *PyArray_SHAPE(PyArrayObject *arr)

是 "return" C PyArrayObject 结构的 npy_intp *dimensions 字段的同义词宏,它指向包含数组所有维度大小的 C 数组,或

npy_intp PyArray_DIM(PyArrayObject* arr, int n)

前一个数组中的第 n 个条目"returns"(即第 n 维的大小。

npy_intp *PyArray_STRIDES(PyArrayObject* arr)

npy_intp PyArray_STRIDE(PyArrayObject* arr, int n)

是宏,分别 "return" C PyArrayObject 结构的 npy_intp *strides 字段指向数组的(数组)步幅,或第 nth 这个数组中的条目。步幅是在 "lines" 之间跳过的字节数,用于数组的所有维度。由于数组是连续的,所以不需要这样做,但可以避免程序必须自己乘以单元格的数量乘以它们的大小。

PyObject* PyArray_NewLikeArray(PyArrayObject* prototype, NPY_ORDER order, PyArray_Descr* descr, int subok)

是一个创建新 numpy 数组的函数,其形状与作为参数传递的原型相同。此数组未初始化。

PyArray_FILLWBYTE(PyObject* obj, int val)

是一个函数,它调用 memset 来初始化给定的 numpy 数组。

void *PyArray_DATA(PyArrayObject *arr)

是一个宏"returns" C PyArrayObject结构的char *data字段,它指向数组的真实数据space,它的形状与C相同数组。

这里是PyArrayObject结构体的声明,as described in the documentation:

typedef struct PyArrayObject {
    PyObject_HEAD
    char *data;
    int nd;
    npy_intp *dimensions;
    npy_intp *strides;
    PyObject *base;
    PyArray_Descr *descr;
    int flags;
    PyObject *weakreflist;
} PyArrayObject;

这里是示例代码:

#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include <numpy/arrayobject.h>


#define G 6.67408E-8L 

void * failure(PyObject *type, const char *message) {
    PyErr_SetString(type, message);
    return NULL;
}

void * success(PyObject *var){
    Py_INCREF(var);
    return var;
}

static PyObject *
Py_grav_c(PyObject *self, PyObject *args)
{
    PyArrayObject *p, *M;
    PyObject *a;
    int i, j, k;
    double *pq0, *pq1, *Mq0, *Mq1, *aq0, *aq1, *p0, *p1, *a0, *a1;


    if (!PyArg_ParseTuple(args, "O!O!", &PyArray_Type, &p, &PyArray_Type, &M))
        return failure(PyExc_RuntimeError, "Failed to parse parameters.");

    if (PyArray_DESCR(p)->type_num != NPY_DOUBLE)
        return failure(PyExc_TypeError, "Type np.float64 expected for p array.");

    if (PyArray_DESCR(M)->type_num != NPY_DOUBLE)
        return failure(PyExc_TypeError, "Type np.float64 expected for M array.");

    if (PyArray_NDIM(p)!=2)
        return failure(PyExc_TypeError, "p must be a 2 dimensionnal array.");

    if (PyArray_NDIM(M)!=1)
        return failure(PyExc_TypeError, "M must be a 1 dimensionnal array.");

    int K = PyArray_DIM(p, 0);     // Number of dimensions you want
    int L = PyArray_DIM(p, 1);     // Number of bodies in the system
    int S0 = PyArray_STRIDE(p, 0); // Normally, the arrays should be contiguous
    int S1 = PyArray_STRIDE(p, 1); // But since they provide this Stride info
    int SM = PyArray_STRIDE(M, 0); // I supposed they might not be (alignment)

    if (PyArray_DIM(M, 0) != L)
        return failure(PyExc_TypeError, 
                       "P and M must have the same number of bodies.");

    a = PyArray_NewLikeArray(p, NPY_ANYORDER, NULL, 0);
    if (a == NULL)
        return failure(PyExc_RuntimeError, "Failed to create output array.");
    PyArray_FILLWBYTE(a, 0);

    // For all bodies except first which has no previous body
    for (i = 1,
         pq0 = (double *)(PyArray_DATA(p)+S1),
         Mq0 = (double *)(PyArray_DATA(M)+SM),
         aq0 = (double *)(PyArray_DATA(a)+S1);
         i < L;
         i++,
         *(void **)&pq0 += S1,
         *(void **)&Mq0 += SM,
         *(void **)&aq0 += S1
         ) {
        // For all previous bodies
        for (j = 0,
            pq1 = (double *)PyArray_DATA(p),
            Mq1 = (double *)PyArray_DATA(M),
            aq1 = (double *)PyArray_DATA(a);
            j < i;
            j++,
            *(void **)&pq1 += S1,
            *(void **)&Mq1 += SM,
            *(void **)&aq1 += S1
             ) {
            // For all dimensions calculate deltas
            long double d[K], d2 = 0, VVV, M0xVVV, M1xVVV;
            for (k = 0,
                 p0 = pq0,
                 p1 = pq1;
                 k<K;
                 k++,
                 *(void **)&p0 += S0,
                 *(void **)&p1 += S0) {
                d[k] = *p1 - *p0;
            }
            // calculate Hypotenuse squared
            for (k = 0, d2 = 0; k<K; k++) {
                d2 += d[k]*d[k];
            }
            // calculate interm. results once for each bodies pair (optimization)
            VVV = G * (d2>0 ? pow(d2, -1.5) : 1); // anonymous intermediate result 
#define LIM = 1
//            VVV = G * pow(max(d2, LIM), -1.5);  // Variation on collision case
            M0xVVV = *Mq0 * VVV;                  // anonymous intermediate result
            M1xVVV = *Mq1 * VVV;                  // anonymous intermediate result
            // For all dimensions calculate component of acceleration
            for (k = 0,
                 a0 = aq0,
                 a1 = aq1;
                 k<K;
                 k++,
                 *(void **)&a0 += S0,
                 *(void **)&a1 += S0) {
                *a0 += M1xVVV*d[k];
                *a1 -= M0xVVV*d[k];
            }
        }
    }

    /*  clean up and return the result */
    return success(a);
}



// exported functions list

static PyMethodDef grav_c_Methods[] = {
    {"grav_c", Py_grav_c, METH_VARARGS, "grav_c(p, M)\n"
"\n"
"grav_c takes the positions and masses of m bodies in Newtonian"
" attraction in a n dimensional universe,\n"
"and returns the accelerations each body undergoes.\n"
"input data take the for of a row of fload64 for each dimension of the"
" position (in p) and one row for the masses.\n"
"It returns and array of the same shape as p for the accelerations."},
    {NULL, NULL, 0, NULL} // pour terminer la liste.
};


static char grav_c_doc[] = "Compute attractions between n bodies.";



static struct PyModuleDef grav_c_module = {
    PyModuleDef_HEAD_INIT,
    "grav_c",   /* name of module */
    grav_c_doc, /* module documentation, may be NULL */
    -1,         /* size of per-interpreter state of the module,
                 or -1 if the module keeps state in global variables. */
    grav_c_Methods
};



PyMODINIT_FUNC
PyInit_grav_c(void)
{
    // I don't understand why yet, but the program segfaults without this.
    import_array();

    return PyModule_Create(&grav_c_module);
}