python - Create block diagonal numpy array from a given numpy array

Question

I have a 2-dimensional numpy array with an equal number of columns and rows. I would like to arrange them into a bigger array having the smaller ones on the diagonal. It should be possible to specify how often the starting matrix should be on the diagonal. For example:

a = numpy.array([[5, 7], 
                 [6, 3]])

So if I wanted this array 2 times on the diagonal the desired output would be:

array([[5, 7, 0, 0], 
       [6, 3, 0, 0], 
       [0, 0, 5, 7], 
       [0, 0, 6, 3]])

For 3 times:

array([[5, 7, 0, 0, 0, 0], 
       [6, 3, 0, 0, 0, 0], 
       [0, 0, 5, 7, 0, 0], 
       [0, 0, 6, 3, 0, 0],
       [0, 0, 0, 0, 5, 7],
       [0, 0, 0, 0, 6, 3]])

Is there a fast way to implement this with numpy methods and for arbitrary sizes of the starting array (still considering the starting array to have the same number of rows and columns)?

Answer 1

Approach #1

Classic case of numpy.kron -

np.kron(np.eye(r,dtype=int),a) # r is number of repeats

Sample run -

In [184]: a
Out[184]: 
array([[1, 2, 3],
       [3, 4, 5]])

In [185]: r = 3 # number of repeats

In [186]: np.kron(np.eye(r,dtype=int),a)
Out[186]: 
array([[1, 2, 3, 0, 0, 0, 0, 0, 0],
       [3, 4, 5, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 1, 2, 3, 0, 0, 0],
       [0, 0, 0, 3, 4, 5, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 1, 2, 3],
       [0, 0, 0, 0, 0, 0, 3, 4, 5]])

Approach #2

Another efficient one with diagonal-viewed-array-assignment -

def repeat_along_diag(a, r):
    m,n = a.shape
    out = np.zeros((r,m,r,n), dtype=a.dtype)
    diag = np.einsum('ijik->ijk',out)
    diag[:] = a
    return out.reshape(-1,n*r)

Sample run -

In [188]: repeat_along_diag(a,3)
Out[188]: 
array([[1, 2, 3, 0, 0, 0, 0, 0, 0],
       [3, 4, 5, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 1, 2, 3, 0, 0, 0],
       [0, 0, 0, 3, 4, 5, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 1, 2, 3],
       [0, 0, 0, 0, 0, 0, 3, 4, 5]])