python - (Numpy C API)迭代单个数组 : NpyIter vs for loop (with PyArray_DATA)

Question

我正在为 python 模块编写一些 C 扩展代码。我要写的函数是(在python中)

output = 1./(1. + input)

其中 input 是任意形状的 numpy 数组。

最初我使用的是 NpyIter_MultiNew :

static PyObject *
helper_calc1(PyObject *self, PyObject *args){

    PyObject * input;
    PyObject * output = NULL;

    if (!PyArg_ParseTuple(args, "O", &input)){
        return NULL;
    }

    // -- input -----------------------------------------------
    PyArrayObject * in_arr;

    in_arr = (PyArrayObject *) PyArray_FROM_OTF(input, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
    if (in_arr == NULL){
        goto fail;
    }

    // -- set up iterator -------------------------------------
    PyArrayObject * op[2];
    npy_uint32 op_flags[2];
    npy_uint32 flags;

    op[0] = in_arr;
    op_flags[0] = NPY_ITER_READONLY;

    op[1] = NULL;
    op_flags[1] = NPY_ITER_WRITEONLY | NPY_ITER_ALLOCATE;

    flags = NPY_ITER_EXTERNAL_LOOP | NPY_ITER_BUFFERED | NPY_ITER_GROWINNER;

    NpyIter * iter = NpyIter_MultiNew(2, op,
                            flags, NPY_KEEPORDER, NPY_NO_CASTING,
                            op_flags, NULL);

    if (iter == NULL){
        goto fail;
    };

    NpyIter_IterNextFunc * iternext = NpyIter_GetIterNext(iter, NULL);
    if (iternext == NULL){
        NpyIter_Deallocate(iter);
        goto fail;
    };

    // -- iterate ---------------------------------------------
    npy_intp count;
    char ** dataptr = NpyIter_GetDataPtrArray(iter);
    npy_intp * strideptr = NpyIter_GetInnerStrideArray(iter);
    npy_intp * innersizeptr = NpyIter_GetInnerLoopSizePtr(iter);

    do {
        count = *innersizeptr;

        while (count--){

            *(double *) dataptr[1] = 1. / (1. + *(double *)dataptr[0]);

            dataptr[0] += strideptr[0];
            dataptr[1] += strideptr[1];
        }

    } while (iternext(iter));

    output = NpyIter_GetOperandArray(iter)[1];

    if (NpyIter_Deallocate(iter) != NPY_SUCCEED){
        goto fail;
    }

    Py_DECREF(in_arr);

    return output;

    fail:
        Py_XDECREF(in_arr);
        Py_XDECREF(output);
        return NULL;
}

但是，由于这只是一个数组(即我不需要担心广播多个数组)，有什么原因我不能自己使用迭代数据，PyArray_DATA ，for 循环和数组大小？

static PyObject *
helper_calc2(PyObject *self, PyObject *args){

    PyObject * input;
    PyObject * output = NULL;

    if (!PyArg_ParseTuple(args, "O", & in)){
        return NULL;
    }

    // -- input -----------------------------------------------
    PyArrayObject * in_arr;

    in_arr = (PyArrayObject *) PyArray_FROM_OTF(input, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
    if (in_arr == NULL){
        Py_XDECREF(in_arr);
        return NULL;
    }

    int ndim = PyArray_NDIM(in_arr);
    npy_intp * shape = PyArray_DIMS(in_arr);
    int size = (int) PyArray_SIZE(in_arr);

    double * in_data = (double *) PyArray_DATA(in_arr);

    output =  PyArray_SimpleNew(ndim, shape, NPY_DOUBLE);
    double * out_data = (double *) PyArray_DATA((PyArrayObject *) output);

    for (int i = 0; i < size; i++){
        out_data[i] = 1. / (1. + in_data[i]);
    }

    Py_DECREF(in_arr);
    return output;

fail:
    Py_XDECREF(in_arr);
    Py_XDECREF(output);
    return NULL;
}

第二个版本运行速度更快，代码更短。

使用时有没有需要注意的危险，PyArray_DATA使用 for 循环代替 NpyIter_MultiNew ？

来自PyArray_DATA文档:

If you have not guaranteed a contiguous and/or aligned array then be sure you understand how to access the data in the array to avoid memory and/or alignment problems.

但我相信这是由 PyArray_FROM_OTF 处理的与 NPY_ARRAY_IN_ARRAY旗帜。

Answer 1

在这种情况下你应该没问题。就像你提到的 NPY_ARRAY_IN_ARRAY 和 PyArray_FROM_OTF 会解决这个问题，因为你将数据指针转换为正确的类型，你应该没问题。但是，一般来说，如您所知，如果您直接从 Python 代码接受 NumPy 数组，则不能使用此方法。

快乐的 C 扩展编写。