Numpy 多维与线性索引的相互转换

Question

我在寻找一种快速的方法，可以在Numpy中进行线性索引和多维索引之间的转换。

为了让我的使用场景更具体，我有一大堆N个粒子，每个粒子都有5个浮点值（维度），这就形成了一个Nx5的数组。接着，我使用numpy.digitize对每个维度进行分箱，选择合适的分箱边界，以便为每个粒子在5维空间中分配一个箱子。

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就包含了对应多维索引的行。如果我想把多维索引转换成线性索引，我想我需要做类似这样的事情：

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

这将为每个多维索引提供一个查找，以将其映射到线性索引。然后你可以通过以下方式再转换回来：

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

我遇到的困难是，如何将包含多维索引的binassign（这个Nx5的数组）转换为一维线性索引，方法是用它来切片线性索引数组linind。

如果有人有一种（或几种）简单的索引技巧，可以在多维索引和线性索引之间来回转换，并且能够对所有N个粒子进行向量化操作，我将非常感激你的见解。

Answer 1

你可以很简单地计算每个箱子的索引：

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

标量乘积就是把每个数字和它的权重相乘，然后加起来，比如说就是 1*x0 + 5*x1 + 5*5*x2 +…… 这个过程通过NumPy的 dot() 函数做得非常高效。

Answer 2

虽然我非常喜欢EOL的回答，但我想把它稍微普遍化一下，以适应每个方向上箱子的数量不均匀的情况，同时也想强调一下C风格和F风格的排序差异。下面是一个示例解决方案：

ndims = 5
N = 10

# Define bin boundaries 
binbnds = ndims*[None]
nbins = []
for idim in xrange(ndims):
    binbnds[idim] = numpy.linspace(-10.0,10.0,numpy.random.randint(2,15))
    binbnds[idim][0] = -float('inf')
    binbnds[idim][-1] = float('inf')
    nbins.append(binbnds[idim].shape[0]-1)

nstates = numpy.cumprod(nbins)[-1]

# Define variable values for N particles in ndims dimensions
p = numpy.random.normal(size=(N,ndims))

# Assign to bins along each dimension
binassign = ndims*[None]
for idim in xrange(ndims):
    binassign[idim] = numpy.digitize(p[:,idim],binbnds[idim]) - 1

binassign = numpy.array(binassign)

# multidimensional array with elements mapping from multidim to linear index
# Two different arrays for C vs F ordering
linind_C = numpy.arange(nstates).reshape(nbins,order='C')
linind_F = numpy.arange(nstates).reshape(nbins,order='F')

现在进行转换

# Fast conversion to linear index
b_F = numpy.cumprod([1] + nbins)[:-1]
b_C = numpy.cumprod([1] + nbins[::-1])[:-1][::-1]

box_index_F = numpy.dot(b_F,binassign)
box_index_C = numpy.dot(b_C,binassign)

并检查结果是否正确：

# Check
print 'Checking correct mapping for each particle F order'
for k in xrange(N):
    ii = box_index_F[k]
    jj = linind_F[tuple(binassign[:,k])]
    print 'particle %d %s (%d %d)' % (k,ii == jj,ii,jj)

print 'Checking correct mapping for each particle C order'
for k in xrange(N):
    ii = box_index_C[k]
    jj = linind_C[tuple(binassign[:,k])]
    print 'particle %d %s (%d %d)' % (k,ii == jj,ii,jj)

为了完整起见，如果你想快速地从一维索引返回到多维索引，可以使用一种向量化的方式：

print 'Convert C-style from linear to multi'
x = box_index_C.reshape(-1,1)
bassign_rev_C = x / b_C % nbins 

print 'Convert F-style from linear to multi'
x = box_index_F.reshape(-1,1)
bassign_rev_F = x / b_F % nbins

再检查一下：

print 'Check C-order'
for k in xrange(N):
    ii = tuple(binassign[:,k])
    jj = tuple(bassign_rev_C[k,:])
    print ii==jj,ii,jj

print 'Check F-order'
for k in xrange(N):
    ii = tuple(binassign[:,k])
    jj = tuple(bassign_rev_F[k,:])
    print ii==jj,ii,jj