Cython + Numpy variable ndim?

Question

I need to initialize arrays of variable shape (dim,) + (nbins,)*dim, where dim is usually small, but nbins can be large, so that the array has ndims = dim + 1. For example if dim = 1 I need an array of shape (1, nbins), if dim = 2 the shape is (2, nbins, nbins) etc.

Is it possible to type numpy arrays accordingly? I tried something like

 ctypedef uint32_t  uint_t
 ctypedef float     real_t
 ...
     cdef:
         uint_t dim = input_data.shape[1]
         np.ndarray[real_t, ndim=dim+1] my_array = np.zeros((dim,)\
     + (nbins,)*dim, dtype = np.float32)

Yeah, I had a hunch it wouldn't work, but had to try anyway ;)

Will it be possible to do something like this or do I have to use pointers/ memory allocation etc? Or do I have to (gulp!) just use a one dimensional array and just reshape it at the end??

Answer 1

Or do I have to (gulp!) just use a one dimensional array and just reshape it at the end?

I'm afraid you do. Cython can only do its array access optimization magic if the number of dimensions is known at compile time. Don't go mucking with malloc, it's not worth it.

cdef:
    np.npy_intp size = dim * n_bins ** dim
    np.ndarray[float, ndim=1, mode='c'] arr = np.zeros(size, dtype=np.float32)

    # do work, using "manual" index calculations

    return arr.reshape(dim, (n_bins,) * dim)

(Side note: the proper type for shapes is np.npy_intp, not uint32_t.)