в NumPy, что является многомерным эквивалентом дубля

Question

У меня есть этот бит кода

def build_tree_base(blocks, x, y, z):
   indicies = [
        (x  ,z  ,y  ),
        (x  ,z+1,y  ),
        (x  ,z  ,y+1),
        (x  ,z+1,y+1),
        (x+1,z  ,y  ),
        (x+1,z+1,y  ),
        (x+1,z  ,y+1),
        (x+1,z+1,y+1),
    ]
    children = [blocks[i] for i in indicies]
    return Node(children=children)

Там, где блоки - это 3-мерный массив NumPy.

Что я хотел бы сделать, это заменить понимание списка чем-то вроде NUMPY.take, однако, take, кажется, имеет дело только с единичными индексами измерения.Есть ли что-то вроде take, которое будет работать с многомерными индексами?

Также я знаю, что вы могли бы сделать это с транспонированием, срезом, а затем изменить форму, но это было медленно, поэтому я ищу лучший вариант.

Answer 1

Индексирование Numpy делает это довольно простым ... Вы должны быть в состоянии что-то вроде этого:

def build_tree_base(blocks, x, y, z):
    idx = [x, x, x, x, x+1, x+1, x+1, x+1]
    idz = [z, z+1, z, z+1, z, z+1, z, z+1]
    idy = [y, y, y+1, y+1, y, y, y+1, y+1]
    children = blocks[idx, idz, idy]
    return Node(children=children)

Редактировать: я должен указать, что это (или любое другое "модное" индексирование) вернет копию, а не представление в исходный массив ...

Answer 2

Как насчет взятия среза 2x2x2, тогда flat?

import numpy as np
blocks = np.arange(2*3*4.).reshape((2,3,4))
i,j,k = 0,1,2
print [x for x in blocks[i:i+2, j:j+2, k:k+2].flat]

(flat - итератор; разверните его так или с помощью np.fromiter(), или позвольте Node iter над ним.)