(API C de Numpy) Itération sur un tableau unique : NpyIter vs boucle for (avec PyArray_DATA)

Question

J'écris du code d'extension en C pour un module python. La fonction que je veux écrire est (en python)

output = 1./(1. + input)

donde input est un tableau numpy de n'importe quelle forme.

À l'origine, j'utilisais 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;
}

Cependant, comme il s'agit d'un seul tableau (c'est-à-dire que je n'ai pas besoin de me préoccuper de la diffusion de plusieurs tableaux), Y a-t-il une raison pour que je ne puisse pas itérer sur les données moi-même en utilisant, PyArray_DATA , a for et la taille du tableau ?

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;
}

Cette deuxième version s'exécute plus rapidement et le code est plus court.

Y a-t-il des dangers auxquels je dois faire attention lorsque j'utilise.., PyArray_DATA avec un for au lieu de NpyIter_MultiNew ?

De la PyArray_DATA documentation :

Si vous n'avez pas garanti un tableau contigu et/ou aligné, assurez-vous de comprendre comment accéder aux données dans le tableau pour éviter les problèmes de mémoire et/ou d'alignement.

Mais je crois que cela est pris en charge par PyArray_FROM_OTF avec le NPY_ARRAY_IN_ARRAY drapeau.

Answer 1

Vous devriez vous en sortir dans ce cas. Comme vous l'avez mentionné NPY_ARRAY_IN_ARRAY con PyArray_FROM_OTF s'occupe de ce problème et puisque vous coulez le pointeur de données vers le bon type, tout devrait bien se passer. Cependant, en général, comme vous semblez le savoir, vous ne pouvez pas utiliser cette approche si vous acceptez directement un tableau NumPy à partir du code Python.

Bonne écriture de l'extension C.