如何在Numpy C扩展中返回可变长度数组？

Question

我之前做过一些Numpy的C扩展，得到了这个网站的大力帮助，但我发现返回的参数都是固定长度的。

有没有办法让Numpy的C扩展返回一个可变长度的numpy数组呢？

Answer 1

我理解你的问题是“我有一个函数，它接收一个长度为 n 的 NumPy 数组，但它会返回一个长度为 m 的不同数组。”如果是这样的话，你需要在扩展中使用 malloc 来申请一个新的 C 数组，比如：

new_array = malloc(m * sizeof(int64)); // or whatever your data type is

然后用这个新数组创建一个新的 NumPy 数组。这个例子假设你处理的是一维数组：

int npy_intp dims[1];
dims[0] = m;
PyArrayObject *out = (PyArrayObject *)PyArray_SimpleNewFromData(1,          // 1D array
                                                                dims,       // dimensions
                                                                NPY_INT64,  // type
                                                                new_array);
PyArray_ENABLEFLAGS(out, NPY_ARRAY_OWNDATA);

然后返回这个新数组。这里重要的一点是要设置 NPY_ARRAY_OWNDATA 标志，这样你申请的内存会在 Python 对象被垃圾回收时自动释放。

Answer 2

你可能会发现，用Cython来制作numpy扩展会更简单，因为它使用了Numpy的C-API，这样可以让你在Python和C对象之间混合使用。这意味着，创建一个可变长度的数组其实并不难，你只需要指定一个任意形状的数组就可以了。

Cython numpy教程可能是这个主题上最好的资料来源。

比如，这里有一个我最近写的函数：

import numpy as np
cimport numpy as np
cimport cython

dtype = np.double
ctypedef double dtype_t

np.import_ufunc()
np.import_array()

def ewma(a, d, axis):
    #Calculates the exponentially weighted moving average of array a along axis using the parameter d.
    cdef void *args[1]

    cdef double weight[1]
    weight[0] = <double>np.exp(-d)


    args[0] = &weight[0]

    return apply_along_axis(&ewma_func, np.array(a, dtype = float), np.double, np.double, False, &(args[0]), <int>axis)

cdef void ewma_func(int n, void* aData,int astride, void* oData, int ostride, void** args):
    #Exponentially weighted moving average calculation function 

    cdef double avg = 0.0
    cdef double weight = (<double*>(args[0]))[0]
    cdef int i = 0

    for i in range(n): 

        avg = (<double*>((<char*>aData) + i * astride))[0]*weight + avg * (1.0 - weight) 


        (<double*>((<char*>oData) + i * ostride))[0] = avg  

ctypedef void (*func_1d)(int, void*, int, void*, int, void **)

cdef apply_along_axis(func_1d function, a, adtype, odtype, reduce,  void** args, int axis):
    #generic function for applying a cython function along a particular dimension

    oshape = list(a.shape)

    if reduce :
        oshape[axis] = 1

    out = np.empty(oshape, odtype)

    cdef np.flatiter ita, ito

    ita = np.PyArray_IterAllButAxis(a,   &axis)
    ito = np.PyArray_IterAllButAxis(out, &axis)

    cdef int axis_length = a.shape[axis]
    cdef int a_axis_stride = a.strides[axis]
    cdef int o_axis_stride = out.strides[axis]

    if reduce: 
        o_axis_stride = 0

    while np.PyArray_ITER_NOTDONE(ita):

        function(axis_length, np.PyArray_ITER_DATA (ita), a_axis_stride, np.PyArray_ITER_DATA (ito), o_axis_stride, args)

        np.PyArray_ITER_NEXT(ita)
        np.PyArray_ITER_NEXT(ito)

    if reduce:  
        oshape.pop(axis)
        out.shape = oshape

    return out  

如果这个不适合你，还有一个函数可以用来创建一个任意形状的新空数组（链接）。