CODEWITHSUNDEEP
pythonnumpy
Authman Apatira
Authman Apatira

Reputation: 4054

Numpy Group Reshaping / Indexing

The situation is I'd like to take the following Python / NumPy code:

# Procure some data:
z = np.zeros((32,32))

chunks = []
for i in range(0,32,step):
    for j in range(0,32,step):
        chunks.append( z[i:i+step,j:j+step] )
chunks = np.array(chunks)
chunks.shape # (256, 2, 2)

And vectorize it / remove the for loops. Is this possible? I don't mind much about ordering of the final array, e.g. 256,2,2 vs 2,2,256, as long as the spatial structure remains the same. That is, blocks of 2x2 from the original array.

Perhaps some magic using :: in addition to regular indexing can do this? Any NumPy masters here?

Upvotes: 0

Views: 290

Answers (1)

Quang Hoang
Quang Hoang

Reputation: 150785

You may need transpose:

a = np.arange(1024).reshape(32,32)

a.reshape(16,2,16,2).transpose((0,2,1,3)).reshape(-1,2,2)

Output:

array([[[   0,    1],
        [  32,   33]],

       [[   2,    3],
        [  34,   35]],

       [[   4,    5],
        [  36,   37]],

       ...,

       [[ 986,  987],
        [1018, 1019]],

       [[ 988,  989],
        [1020, 1021]],

       [[ 990,  991],
        [1022, 1023]]])

Upvotes: 2

Related Questions