Skip to content
Advertisement

in numpy what is the multi dimensional equivalent of take

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.

Advertisement

Answer

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.)

User contributions licensed under: CC BY-SA
7 People found this is helpful
Advertisement