in numpy what is the multi dimensional equivalent of take

Question

I have this bit of code

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)

Where blocks is a 3 dimensional numpy array.

What I'd like to do is replace the list comprehension with something like numpy.take, however take seems to only deal with single dimension indices. Is there something like take that will work with multidimensional indices?

Also I know you could do this with a transpose, slice and then reshape, but that was slow so I'm looking for a better option.

Answer 1

Numpy indexing make this quite easy... You should be able to to something like this:

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)

Edit: I should point out that this (or any other "fancy" indexing) will return a copy, rather than a view into the original array...

Answer 2

How about taking a 2x2x2 slice, then 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 is an iterator; expand it like this, or with np.fromiter(), or let Node iter over it.)