How to turn a 2D array into a 3D diagonal matrix with numpy

January 26, 2022

Having a 2D numpy array with shape (i,j), where each column is a vector with i entries, how can I obtain a 3D numpy array with shape (j, i, i) where each j corresponds to a diagonal matrix with each vector on its diagonal without using for loops?

For example, having the 2D array [[1,5],[2,6],[3,7],[4,8]], I wanted to obtain a 3D array where for j=0, the matrix would be

[[1,0,0,0],[0,2,0,0],[0,0,3,0],[0,0,0,4]]

and for j=1 the matrix would be

[[5,0,0,0],[0,6,0,0],[0,0,7,0],[0,0,0,8]].

I tried np.diag but apparently it outputs the matrix diagonal when a 2D matrix is provided.

>Solution :

You can use np.einsum:

np.einsum('ij,ik->jik', a, np.eye(a.shape[0], dtype=a.dtype))

Output:

array([[[1, 0, 0, 0],
        [0, 2, 0, 0],
        [0, 0, 3, 0],
        [0, 0, 0, 4]],

       [[5, 0, 0, 0],
        [0, 6, 0, 0],
        [0, 0, 7, 0],
        [0, 0, 0, 8]]])