Numpy: function to create block matrix from another matrices, each block should contain values from the same indexes from another matrices

Question

I'm looking for some fast function to create blocks in matrix from another matrices like in a picture (bottom matrix is result): matrix

I also wrote some code to do this but for many matrices is working slow.

import numpy as np
im1 = np.array([[1, 2], [3, 4]])
im2 = np.array([[5, 6], [7, 8]])
im3 = np.array([[9, 10], [11, 12]])
im4 = np.array([[13, 14], [15, 15]])
arrays = [im1, im2, im3, im4]


def array_rearrangement(arrays):
    num = int(len(arrays) ** 0.5)
    size = arrays[0].shape[0] * num

    result = np.zeros((size, size))
    k = -1
    p = -1
    for k1 in range(num * num):
        if k1 % num == 0:
            k = 0
        if k1 % num == 0:
            p += 1
        im = arrays[k1]
        for i in range(num):
            for j in range(num):
                x_index = i * num + p
                y_index = j * num + k
                result[x_index, y_index] = im[i, j]

        k += 1

    print(result)


array_rearrangement(arrays)

Answer 1

This answer is always a good starting point for nd to nd transformation. In your case, this is a reshape, swapaxes and reshape again.

---

arr = np.column_stack(arrays)
arr.reshape(-1, 2, 2).swapaxes(2, 1).reshape(4, 4)

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