CODEWITHSUNDEEP
pythonarraysmultidimensional-arraymatrixnumpy
dolphin
dolphin

Reputation: 1233

Selecting some dimensions from a multi-dimensional array

I have a 4-D array, and I need to process all 1-D vectors from this array along a given dimension. This works well:

    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

...but I believe there is some way to avoid the nested "if" part, and maybe also be more efficient? Like creating other views of this array before the loops and then iterating through these views?

Upvotes: 1

Views: 194

Answers (1)

Bi Rico
Bi Rico

Reputation: 25813

One common way to handle this is to use 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

Note that np.rollaxis should return a view so it doesn't actually copy the array.

Upvotes: 3

Related Questions