python - 使用cython将numpy数组列表传递给C

Question

我有一个 3D numpy 数组的列表 list_of_arrays，我想使用模板将其传递给 C 函数

int my_func_c(double **data, int **shape, int n_arrays)

这样

data[i]  : pointer to the numpy array values in list_of_arrays[i]
shape[i] : pointer to the shape of the array in list_of_arrays[i] e.g. [2,3,4]

如何使用 cython 接口(interface)函数调用 my_func_c？

我的第一个想法是做类似下面的事情(可行)，但我觉得有一种更好的方法，只使用 numpy 数组而不进行 mallocing 和释放。

# my_func_c.pyx

import numpy as np
cimport numpy as np
cimport cython
from libc.stdlib cimport malloc, free

cdef extern from "my_func.c":
    double my_func_c(double **data, int **shape, int n_arrays)

def my_func(list list_of_arrays):
    cdef int n_arrays  = len(list_of_arrays)
    cdef double **data = <double **> malloc(n_arrays*sizeof(double *))
    cdef int **shape   = <int **> malloc(n_arrays*sizeof(int *))
    cdef double x;

    cdef np.ndarray[double, ndim=3, mode="c"] temp

    for i in range(n_arrays):
        temp = list_of_arrays[i]
        data[i]  = &temp[0,0,0]
        shape[i] = <int *> malloc(3*sizeof(int))
        for j in range(3):
            shape[i][j] = list_of_arrays[i].shape[j]

    x = my_func_c(data, shape, n_arrays)

    # Free memory
    for i in range(n_arrays):
        free(shape[i])
    free(data)
    free(shape)

    return x

注意

要查看工作示例，我们可以使用一个非常简单的函数来计算列表中所有数组的乘积。

# my_func.c

double my_func_c(double **data, int **shape, int n_arrays) {
    int array_idx, i0, i1, i2;

    double prod = 1.0;

    // Loop over all arrays
    for (array_idx=0; array_idx<n_arrays; array_idx++) {
        for (i0=0; i0<shape[array_idx][0]; i0++) {
            for (i1=0; i1<shape[array_idx][1]; i1++) {
                for (i2=0; i2<shape[array_idx][2]; i2++) {
                    prod = prod*data[array_idx][i0*shape[array_idx][1]*shape[array_idx][2] + i1*shape[array_idx][2] + i2];
                }
            }
        }
    }

    return prod;
}

创建setup.py文件，

# setup.py

from distutils.core import setup
from Cython.Build import cythonize
import numpy as np

setup(
    name='my_func',
    ext_modules = cythonize("my_func_c.pyx"),
    include_dirs=[np.get_include()]
    )

编译

python3 setup.py build_ext --inplace

最后我们可以运行一个简单的测试

# test.py

import numpy as np
from my_func_c import my_func

a = [1+np.random.rand(3,1,2), 1+np.random.rand(4,5,2), 1+np.random.rand(1,2,3)]

print('Numpy product: {}'.format(np.prod([i.prod() for i in a])))
print('my_func product: {}'.format(my_func(a)))

使用

python3 test.py

Answer 1

另一种选择是让 numpy 为您管理内存。您可以使用 np.uintp 的 numpy 数组来做到这一点，它是一个与任何指针大小相同的无符号整数。

不幸的是，这确实需要一些类型转换(在“指针大小的 int”和指针之间)，这是隐藏逻辑错误的好方法，所以我不是 100% 满意。

def my_func(list list_of_arrays):
    cdef int n_arrays  = len(list_of_arrays)
    cdef np.uintp_t[::1] data = np.array((n_arrays,),dtype=np.uintp)
    cdef np.uintp_t[::1] shape = np.array((n_arrays,),dtype=np.uintp)
    cdef double x;

    cdef np.ndarray[double, ndim=3, mode="c"] temp

    for i in range(n_arrays):
        temp = list_of_arrays[i]
        data[i]  = <np.uintp_t>&temp[0,0,0]
        shape[i] = <np.uintp_t>&(temp.shape[0])

    x = my_func_c(<double**>(&data[0]), <np.intp_t**>&shape[0], n_arrays)

(我应该指出，我只是确认它可以编译，并没有进一步测试它，但基本思路应该没问题)

---

您完成它的方式可能是一种非常明智的方式。对应该有效的原始代码稍作简化

shape[i] = <np.uintp_t>&(temp.shape[0])

而不是 malloc 和复制。我还建议将 free 放在 finally block 中以确保它们运行。

---

编辑: @ead 很有帮助地指出 numpy shape is stored as as np.intp_t - 即一个足够大的有符号整数来容纳一个指针，它主要是 64 位的 - 而 int 通常是 32 位。因此，要在不复制的情况下传递形状，您需要更改 C api。转换帮助使错误更难被发现(“隐藏逻辑错误的好方法”)