Numpy - Reshape and partition 2D array to 3D

Question

Is there a way to partition and reshape a 2D array to a 3D one. Like the following example:

Basically, I have a 4x4 matrix on the left and I want a 2x2x4 matrix in the way shown, so that I can apply numpy.mean on the 3rd axis. In reality the matrix I have is really huge, so that's why looping over the blocks is not an option.

Any help is greatly appreciated.

Answer 1

For your example, you can use numpy.lib.stride_tricks.as_strided.

In [1]: A = np.arange(16).reshape(4, 4)

In [2]: A
Out[2]:
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15]])

In [3]: strides = A.itemsize * np.array([8, 2, 4, 1])

In [4]: x = p.lib.stride_tricks.as_strided(A, shape = (2, 2, 2, 2), strides = strides)

In [4]: x
Out[4]:
array([[[[ 0,  1],
         [ 4,  5]],

        [[ 2,  3],
         [ 6,  7]]],


       [[[ 4,  5],
         [ 8,  9]],

        [[ 6,  7],
         [10, 11]]]])

In [5]: x.reshape(4, 2, 2)
Out[5]:
array([[[ 0,  1],
        [ 4,  5]],

       [[ 2,  3],
        [ 6,  7]],

       [[ 8,  9],
        [12, 13]],

       [[10, 11],
        [14, 15]]])

The strides define offsets in bytes to use when traversing the array, and the as_strided function enables you to build a new array with user-define strides.

However, I do not know how efficient it is and how well it will scale for your use.