How to split a matrix into 4 blocks using numpy?

Question

I'm implementing Strassen's Matrix Multiplication using python. In divide step, we divide a larger matrix into smaller sub-matrices. Is there a built-in numpy function to split a matrix?

Answer 1

According to this answer, you might use the swapaxes:

You can create a helper method as:

def split(array, nrows, ncols):
    """Split a matrix into sub-matrices."""

    r, h = array.shape
    return (array.reshape(h//nrows, nrows, -1, ncols)
                 .swapaxes(1, 2)
                 .reshape(-1, nrows, ncols))

Here is an example of using it

import numpy as np

array = np.array([
    [1, 1, 2, 2],
    [3, 3, 4, 4],
    [5, 5, 6, 6],
    [7, 7, 8, 8]])

A, B, C, D =  split(array, 2, 2)
# A = 
# [[1 1]
#  [3 3]]

# B = 
# [[2 2]
#  [4 4]]

# C = 
# [[5 5]
#  [7 7]]

# D =
# [[6 6]
#  [8 8]]
print('A = \n{}\n\n'
      'B = \n{}\n\n'
      'C = \n{}\n\n'
      'D =\n{}'.format(A, B, C, D))