python - Numpy interconversion between multidimensional and linear indexing

Question

I'm looking for a fast way to interconvert between linear and multidimensional indexing in Numpy.

To make my usage concrete, I have a large collection of N particles, each assigned 5 float values (dimensions) giving an Nx5 array. I then bin each dimension using numpy.digitize with an appropriate choice of bin boundaries to assign each particle a bin in the 5 dimensional space.

N = 10
ndims = 5
p = numpy.random.normal(size=(N,ndims))
for idim in xrange(ndims):
    bbnds[idim] = numpy.array([-float('inf')]+[-2.,-1.,0.,1.,2.]+[float('inf')])

binassign = ndims*[None]
for idim in xrange(ndims):
    binassign[idim] = numpy.digitize(p[:,idim],bbnds[idim]) - 1

binassign then contains rows that correspond to the multidimensional index. If I then want to convert the multidimensional index to a linear index, I think I would want to do something like:

linind = numpy.arange(6**5).reshape(6,6,6,6,6)

This would give a look-up for each multidimensional index to map it to a linear index. You could then go back using:

mindx = numpy.unravel_index(x,linind.shape)

Where I'm having difficulties is figuring out how to take binassign (the Nx5 array) containing the multidimensional index in each row, and coverting that to an 1d linear index, by using it to slice the linear indexing array linind.

If anyone has a one (or several) line indexing trick to go back and forth between the multidimensional index and the linear index in a way that vectorizes the operation for all N particles, I would appreciate your insight.

Answer 1

You can simply calculate the index of each bin:

box_indices = numpy.dot(ndims**numpy.arange(ndims), binassign)

The scalar product simply does 1*x0 + 5*x1 + 5*5*x2 +… This is done very efficiently through NumPy's dot().