General functions to iterate over n-D arrays

Question

With Cython, is there a way to write fast general functions, which work for arrays with different dimensions? For example for this simple case of dealiasing functions:

import numpy as np
cimport numpy as np

ctypedef np.uint8_t DTYPEb_t
ctypedef np.complex128_t DTYPEc_t


def dealiasing1D(DTYPEc_t[:, :] data, 
                 DTYPEb_t[:] where_dealiased):
    """Dealiasing data for 1D solvers."""
    cdef Py_ssize_t ik, i0, nk, n0

    nk = data.shape[0]
    n0 = data.shape[1]

    for ik in range(nk):
        for i0 in range(n0):
            if where_dealiased[i0]:
                data[ik, i0] = 0.


def dealiasing2D(DTYPEc_t[:, :, :] data, 
                 DTYPEb_t[:, :] where_dealiased):
    """Dealiasing data for 2D solvers."""
    cdef Py_ssize_t ik, i0, i1, nk, n0, n1

    nk = data.shape[0]
    n0 = data.shape[1]
    n1 = data.shape[2]

    for ik in range(nk):
        for i0 in range(n0):
            for i1 in range(n1):
                if where_dealiased[i0, i1]:
                    data[ik, i0, i1] = 0.


def dealiasing3D(DTYPEc_t[:, :, :, :] data, 
                 DTYPEb_t[:, :, :] where_dealiased):
    """Dealiasing data for 3D solvers."""
    cdef Py_ssize_t ik, i0, i1, i2, nk, n0, n1, n2

    nk = data.shape[0]
    n0 = data.shape[1]
    n1 = data.shape[2]
    n2 = data.shape[3]

    for ik in range(nk):
        for i0 in range(n0):
            for i1 in range(n1):
                for i2 in range(n2):
                    if where_dealiased[i0, i1, i2]:
                        data[ik, i0, i1, i2] = 0.

Here, I need three functions for the one-dimensional, two-dimensional and three-dimensional cases. Is there a good method to write a function that do the job for all (reasonable) dimensions?

PS: Here, I've tried to use memoryviews, but I am not sure it's the right method to do this. I am surprised that the lines if where_dealiased[i0]: data[ik, i0] = 0. are not white in the annotated html produced by the cython -a command. Is there something wrong?

Blake Walsh · Accepted Answer · 2014-10-06T02:00:40.420

The first thing I would say is that there are reasons for wanting to keep the 3 functions, with a more generic function you will probably miss out on optimizations both from the cython compiler and the c compiler.

Making one function which wraps those 3 functions is very doable, it just takes the two arrays as python objects, checks the shape, and calls the relevant other function.

But if was going to attempt this then what I would try is just writing the function for the highest dimension, and then with arrays of lower dimensions recast them as an array of higher dimensions by using the new axis notation:

cdef np.uint8_t [:] a1d = np.zeros((256, ), np.uint8) # 1d
cdef np.uint8_t [:, :] a2d = a1d[None, :]             # 2d
cdef np.uint8_t [:, :, :] a3d = a1d[None, None, :]    # 3d
a2d[0, 100] = 42
a3d[0, 0, 200] = 108
print(a1d[100], a1d[200])
# (42, 108)

cdef np.uint8_t [:, :] data2d = np.zeros((128, 256), np.uint8) #2d
cdef np.uint8_t [:, :, :, :] data4d = data2d[None, None, :, :] #4d
data4d[0, 0, 42, 108] = 64
print(data2d[42, 108])
# 64

As you can see, a memory view can be cast to higher dimensions and that can be used to modify the original data. You would probably still want to write a wrapper function which performs these tricks before passing the new view to the highest-dimensional function. I suspect this trick will work quite well in your case but you'd have to play around to know if it will do what you want with your data.

With your PS:, there is a very simple explanation. The 'extra code' is the code which generates index error, type error and which allows you to use [-1] to index from the end of the array instead of the start (wraparound). You can disable these extra python features and reduce it to c array functionality through the use of compiler directives, for example to eliminate this extra code from the whole file you can include a comment at the start of the file:

# cython: boundscheck=False, wraparound=False, nonecheck=False

Compiler directives can also be applied at a function level using decorators. The doc explains.

Saullo G. P. Castro · Answer 2 · 2014-10-06T11:15:59.127

You can read in a general way the flattened array using numpy.ndindex() with the strided attribute of the np.ndarray object, such that the position is determined by:

indices[0]*strides[0] + indices[1]*strides[1] + ... + indices[n]*strides[n]

which is easily accomplished doing (strides*indices).sum(), when strides is a 1-D array. The code below shows how to construct a working example:

#cython profile=True
#blacython wraparound=False
#blacython boundscheck=False
#blacython nonecheck=False
#blacython cdivision=True
cimport numpy as np
import numpy as np

def readNDArray(x):
    if not isinstance(x, np.ndarray):
        raise ValueError('x must be a valid np.ndarray object')
    if x.itemsize != 8:
        raise ValueError('x.dtype must be float64')
    cdef np.ndarray[double, ndim=1] v # view of x
    cdef np.ndarray[int, ndim=1] strides
    cdef int pos

    shape = list(x.shape)
    strides = np.array([s//x.itemsize for s in x.strides], dtype=np.int32)
    v = x.ravel()
    for indices in np.ndindex(*shape):
        pos = (strides*indices).sum()
        v[pos] = 2.
    return np.reshape(v, newshape=shape)

This algorithm will not copy the original array if it is C contiguous:

def main():
    # case 1
    x = np.array(np.random.random((3,4,5,6)), order='F')
    y = readNDArray(x)
    print(np.may_share_memory(x, y))
    # case 2
    x = np.array(np.random.random((3,4,5,6)), order='C')
    y = readNDArray(x)
    print np.may_share_memory(x, y)
    return 0

Results in:

False
True

General functions to iterate over n-D arrays

2 Answers2