Function NpyIter_CreateCompatibleStrides

numpy/core/src/multiarray/nditer_api.c:1038–1080

NUMPY_API * Builds a set of strides which are the same as the strides of an * output array created using the NPY_ITER_ALLOCATE flag, where NULL * was passed for op_axes. This is for data packed contiguously, * but not necessarily in C or Fortran order. This should be used * together with NpyIter_GetShape and NpyIter_GetNDim. * * A use case for this function is to match the shape and layout

 * Returns NPY_SUCCEED or NPY_FAIL.
 */
NPY_NO_EXPORT int
NpyIter_CreateCompatibleStrides(NpyIter *iter,
                            npy_intp itemsize, npy_intp *outstrides)
{
    npy_uint32 itflags = NIT_ITFLAGS(iter);
    int idim, ndim = NIT_NDIM(iter);
    int nop = NIT_NOP(iter);

    npy_intp sizeof_axisdata;
    NpyIter_AxisData *axisdata;
    npy_int8 *perm;

    if (!(itflags&NPY_ITFLAG_HASMULTIINDEX)) {
        PyErr_SetString(PyExc_RuntimeError,
                "Iterator CreateCompatibleStrides may only be called "
                "if a multi-index is being tracked");
        return NPY_FAIL;
    }

    axisdata = NIT_AXISDATA(iter);
    sizeof_axisdata = NIT_AXISDATA_SIZEOF(itflags, ndim, nop);

    perm = NIT_PERM(iter);
    for(idim = 0; idim < ndim; ++idim) {
        npy_bool flipped;
        npy_int8 axis = npyiter_undo_iter_axis_perm(idim, ndim, perm, &flipped);
        if (flipped) {
            PyErr_SetString(PyExc_RuntimeError,
                    "Iterator CreateCompatibleStrides may only be called "
                    "if DONT_NEGATE_STRIDES was used to prevent reverse "
                    "iteration of an axis");
            return NPY_FAIL;
        }
        else {
            outstrides[axis] = itemsize;
        }

        itemsize *= NAD_SHAPE(axisdata);
        NIT_ADVANCE_AXISDATA(axisdata, 1);
    }

    return NPY_SUCCEED;
}

/*NUMPY_API
 * Get the array of data pointers (1 per object being iterated)