I am writing a lot of Cython code that needs to generalize to n dimensions (most of the time this will be n = 1, 2, or 3).
Take the following code as an example, this is a simple sum of the elements.
cimport cython
from cython cimport floating
cpdef floating nd_sum(floating [:, :, :] arr):
cdef:
Py_ssize_t [3] N = [arr.shape[0], arr.shape[1], arr.shape[2]]
Py_ssize_t [3] i
floating total = 0
for i[0] in range(N[0]):
for i[1] in range(N[1]):
for i[2] in range(N[2]):
total += arr[i[0], i[1], i[2]]
return total
Conceptually, the generalization to n dimensions is quite obvious. However, I can't seem to be implementing it in code. I would hate having to basically replicate this code for each case of n=1,2,3,...
In general, I also need to be able to access neighbors of each array element in every dimension (think n-dimensional kernel convolutions etc), which makes things like flattening the array infeasible.
My main problems are:
How do I tell Cython to expect an array with arbitrary dimensions as input? I imagine I may have to resort to
np.ndarray
here...How can I get an element from an n-dimensional array given an array of n indices?
How can I generalize the construct of n nested loops? Recursion seems inevitable here, but how do I get access to all n indices that way?