Fastest way to expand the values of a numpy matrix in diagonal blocks

Question

I'm searching for a fast way for resize the matrix in a special way, without using for-loops: I have a squared Matrix:

matrix = [[ 1, 2, 3, 4, 5],
          [ 6, 7, 8, 9,10],
          [11,12,13,14,15],
          [16,17,18,19,20],
          [21,22,23,24,25]]

and my purpose is to resize it 3 (or n) times, where the values are diagonal blocks in the matrix and other values are zeros:

goal_matrix = [[ 1, 0, 0, 2, 0, 0, 3, 0, 0, 4, 0, 0, 5, 0, 0],
               [ 0, 1, 0, 0, 2, 0, 0, 3, 0, 0, 4, 0, 0, 5, 0],
               [ 0, 0, 1, 0, 0, 2, 0, 0, 3, 0, 0, 4, 0, 0, 5],
               [ 6, 0, 0, 7, 0, 0, 8, 0, 0, 9, 0, 0,10, 0, 0],
               [ 0, 6, 0, 0, 7, 0, 0, 8, 0, 0, 9, 0, 0,10, 0],
               [ 0, 0, 6, 0, 0, 7, 0, 0, 8, 0, 0, 9, 0, 0,10],
               [11, 0, 0,12, 0, 0,13, 0, 0,14, 0, 0,15, 0, 0],
               [ 0,11, 0, 0,12, 0, 0,13, 0, 0,14, 0, 0,15, 0],
               [ 0, 0,11, 0, 0,12, 0, 0,13, 0, 0,14, 0, 0,15],
               [16, 0, 0,17, 0, 0,18, 0, 0,19, 0, 0,20, 0, 0],
               [ 0,16, 0, 0,17, 0, 0,18, 0, 0,19, 0, 0,20, 0],
               [ 0, 0,16, 0, 0,17, 0, 0,18, 0, 0,19, 0, 0,20],
               [21, 0, 0,22, 0, 0,23, 0, 0,24, 0, 0,25, 0, 0],
               [ 0,21, 0, 0,22, 0, 0,23, 0, 0,24, 0, 0,25, 0],
               [ 0, 0,21, 0, 0,22, 0, 0,23, 0, 0,24, 0, 0,25]]

It should do something like this question, but without unnecessary zero padding.
Is there any mapping, padding or resizing function for doing this in a fast way?

Answer 1

IMO, it is inappropriate to reject the for loop blindly. Here I provide a solution without the for loop. When n is small, its performance is better than that of @MichaelSzczesny and @SalvatoreDanieleBianco solutions:

def mechanic(mat, n):
    ar = np.zeros((*mat.shape, n * n), mat.dtype)
    ar[..., ::n + 1] = mat[..., None]
    return ar.reshape(
        *mat.shape,
        n,
        n
    ).transpose(0, 3, 1, 2).reshape([s * n for s in mat.shape])

This solution obtains the expected output through a slice assignment, then transpose and reshape, but copies will occur in the last step of reshaping, making it inefficient when n is large.

After a simple test, I found that the solution that simply uses the for loop has the best performance:

def mechanic_for_loop(mat, n):
    ar = np.zeros([s * n for s in mat.shape], mat.dtype)
    for i in range(n):
        ar[i::n, i::n] = mat
    return ar

---

Next is a benchmark test using perfplot. The test functions are as follows:

import numpy as np


def mechanic(mat, n):
    ar = np.zeros((*mat.shape, n * n), mat.dtype)
    ar[..., ::n + 1] = mat[..., None]
    return ar.reshape(
        *mat.shape,
        n,
        n
    ).transpose(0, 3, 1, 2).reshape([s * n for s in mat.shape])


def mechanic_for_loop(mat, n):
    ar = np.zeros([s * n for s in mat.shape], mat.dtype)
    for i in range(n):
        ar[i::n, i::n] = mat
    return ar


def michael_szczesny(mat, n):
    return np.einsum(
        'ij,kl->ikjl',
        mat,
        np.eye(n, dtype=mat.dtype)
    ).reshape([s * n for s in mat.shape])


def salvatore_daniele_bianco(mat, n):
    repeated_matrix = mat.repeat(n, axis=0).repeat(n, axis=1)
    col_ids, row_ids = np.meshgrid(
        np.arange(repeated_matrix.shape[0]),
        np.arange(repeated_matrix.shape[1])
    )
    repeated_matrix[(col_ids % n) - (row_ids % n) != 0] = 0
    return repeated_matrix


functions = [
    mechanic,
    mechanic_for_loop,
    michael_szczesny,
    salvatore_daniele_bianco
]

Resize times unchanged, array size changes:

if __name__ == '__main__':
    from itertools import accumulate, repeat
    from operator import mul
    from perfplot import bench

    bench(
        functions,
        list(accumulate(repeat(2, 11), mul)),
        lambda n: (np.arange(n * n).reshape(n, n), 5),
        xlabel='ar.shape[0]'
    ).show()

Output:

Resize times changes, array size unchanged:

if __name__ == '__main__':
    from itertools import accumulate, repeat
    from operator import mul
    from perfplot import bench

    ar = np.arange(25).reshape(5, 5)
    bench(
        functions,
        list(accumulate(repeat(2, 11), mul)),
        lambda n: (ar, n),
        xlabel='resize times'
    ).show()

Output:

---

---