从多维数组中选择一些尺寸

Question

我有一个4-D数组，我需要沿着给定的维度处理这个数组中的所有1-D向量。 这很好用：

def myfun(arr4d,selected_dim):  # selected_dim can only be 2 or 3
        print arr4d.shape   # (2, 5, 10, 10)
        for i in xrange(arr4d.shape[0]):
            for j in xrange(arr4d.shape[1]):
                for k in xrange(arr4d.shape[selected_dim]):
                    if selected_dim==2:
                        arr=arr4d[i,j,k,:]
                    elif selected_dim==3:
                        arr=arr4d[i,j,:,k]
                    do_something(arr)  # arr is 1-D and has 10 items

...但我相信有一些方法可以避免嵌套的“if”部分，也可能更有效率？ 就像在循环之前创建此数组的其他视图然后迭代这些视图一样？

Answer 1

处理此问题的一种常见方法是使用np.rollaxis ：

def myfun(arr4d, selected_dim):  # selected_dim can only be 2 or 3
    arr4d = np.rollaxis(arr4d, selected_dim)
    print arr4d.shape   # (10, 2, 5, 10)
    for i in xrange(arr4d.shape[1]):
        for j in xrange(arr4d.shape[2]):
            for k in xrange(arr4d.shape[0]):
                arr=arr4d[k, i, j, :]
                do_something(arr)  # arr is 1-D and has 10 items

请注意， np.rollaxis应返回一个视图，因此它实际上不会复制该数组。