Blockwise reshape in numpy?

Question

I have a numpy array with columns that are in blocks. I want to transpose the blocks. It's conceptually simple, and I guess one can do it simply, but I dont know how.

Given a numpy array on block form np.hstack(list_of_blocks), I want to get np.vstack(list_of_blocks).

To make it more precis, I want to go from array a to array b in the snippet below.

import numpy as np
a = np.zeros((3,6))
b = np.zeros((9,2))
t_max = 3
for col in range(1,7):
    for time in range(1,t_max+1):    
        val = ((1+col)//2)*100+((col+1) % 2)*10+time
        a[time-1,col-1]= val
        b[time+t_max*(((1+col)//2)-1)-1,((col+1) % 2)] = val

and the matrixes look like:

>>> print(a)
[[101. 111. 201. 211. 301. 311.]
 [102. 112. 202. 212. 302. 312.]
 [103. 113. 203. 213. 303. 313.]]

>>> print(b)
[[101. 111.]
 [102. 112.]
 [103. 113.]
 [201. 211.]
 [202. 212.]
 [203. 213.]
 [301. 311.]
 [302. 312.]
 [303. 313.]]

Of course the matrixes are not 3 x (2*3) but rather n x (k*m)

What is an efficient way to reshape like this in numpy?

Answer 1

Reshape, permute axes and reshape -

N = a.shape[1]//t_max
b_out = a.reshape(a.shape[0],-1,N).swapaxes(0,1).reshape(-1,N)

More info on intuition behind nd-to-nd array transformation.

Answer 2

np.vstack(np.hsplit(a,3)) does exactly what was asked for, is readable, but is less performant than the answer of Divkar.

Answer 3

This post is old but wants to share the code I wrote for my FYP. This code is not very clean and some cases are not handle. But handles the blocks perfectly.

def getblocks(image: np.ndarray, blockshape: tuple, moveAxis: bool = True, info: bool = False, addChannel: bool = True) -> np.ndarray:
    '''
    takes the array of image in grey= 2D and in RGB = 3D
    takes the numpy array and converts it the the blocks in the fastest way
    '''
    if(info):
        print("Image Shape:", image.shape)
        print("Block Shape:", blockshape)

    oldshape = list(image.shape)
    if addChannel and len(image.shape) == 2:
        mode = "grey"
        image = image.reshape((*image.shape, 1))
    else:
        mode = "color"

    if addChannel:
        img_height, img_width, channels = image.shape
    else:
        img_height, img_width = image.shape

    tile_height, tile_width = blockshape

    if addChannel:
        shp = img_height//tile_height, tile_height, img_width//tile_width, tile_width, channels
    else:
        shp = img_height//tile_height, tile_height, img_width//tile_width, tile_width

    def printinfo():
        print("Old Shape:", oldshape)
        print("Image Shape:", image.shape)
        print("Block Shape:", blockshape)
        print("New Shape Initial:", shp)
        print("img_height % tile_height != 0 :", img_height % tile_height != 0)
        print("img_width % tile_width != 0 :", img_width % tile_width != 0)

    if img_height % tile_height != 0 or img_width % tile_width != 0:
        print("warning: Block size is not fit for the image!")
        printinfo()

    if(info):
        printinfo()

    tiled_array = image.reshape(shp)
    tiled_array = tiled_array.swapaxes(1, 2)

    if moveAxis:
        if(addChannel):
            tiled_array = tiled_array.reshape(-1,
                                              *(tile_height, tile_width, channels))
            tiled_array = np.moveaxis(tiled_array, source=len(
                tiled_array.shape)-1, destination=1)
        else:
            tiled_array = tiled_array.reshape(-1, *(tile_height, tile_width))

    return tiled_array

To reverse this process. Make the original image from blocks.

def combineBlocks(tiled_array: np.ndarray, imageshape: tuple, blockshape: tuple, movedAxis: bool = True, channel: bool = True) -> np.ndarray:

    if channel:
        if len(imageshape) == 2:
            mode = "grey"
            imageshape = *imageshape, 1
        else:
            mode = "color"

    if channel:
        img_height, img_width, channels = imageshape
    else:
        img_height, img_width = imageshape

    tile_height, tile_width = blockshape

    if movedAxis:
        image = tiled_array.copy()
        if(channel):
            image = image.reshape(img_height//tile_height, tile_height,
                                  img_width//tile_width, tile_width, channels)
            swapaxisShape = list(image.shape)
            swapaxisShape[1], swapaxisShape[2] = swapaxisShape[2], swapaxisShape[1]
            image = image.reshape(swapaxisShape)
            image = image.swapaxes(1, 2)
        else:
            f = image.reshape(img_height//tile_height, tile_height,
                              img_width//tile_width, tile_width)
            swapaxisShape = list(f.shape)
            swapaxisShape[1], swapaxisShape[2] = swapaxisShape[2], swapaxisShape[1]
            tmp = f.reshape(swapaxisShape)
            image = tmp.swapaxes(1, 2)
    else:
        image = tiled_array
        # I haven't completed this else case. Btw we aren't using this case lol :)

    return image.reshape(imageshape)

This codes has unhandled cases, like RGB image (3D) but time complexity is very good.