开发者

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 resha开发者_运维知识库pe, but that was slow so I'm looking for a better option.


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


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

0

上一篇:

下一篇:

精彩评论

暂无评论...
验证码 换一张
取 消

最新问答

问答排行榜