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), содержащий многомерный индекс в каждой строке, и покрыть его до 1d линейного индекса, используя его, чтобы разрезать линейный индексный массив linind.

Если у кого-то есть один (или несколько) трюк для индексации строк, позволяющий переключаться между многомерным индексом и линейным индексом таким образом, чтобы векторизовать операцию для всех N частиц, я был бы признателен за ваше понимание.

Answer 1

Вы можете просто рассчитать индекс каждой корзины:

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

Скалярный продукт просто делает 1 * x0 + 5 * x1 + 5 * 5 * x2 + ... Это делается очень эффективно через dot().

от NumPy

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)

и проверьте на правильность:хотите быстро вернуться из 1-го индекса к многомерному индексу в векторном стиле:

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