Create a 3D (partially) diagonal array from a 2D array

Question

I'd like to ask how can I efficiently generate a numpy 3D array from a 2D array with each row filling the diagonal part of the new array? For example, the input 2D array is

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

and I want the output to be

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

       [[3, 0],
        [0, 4]],

       [[5, 0],
        [0, 6]],

       [[7, 0],
        [0, 8]]])

Typically, the size of the first dimensional is very large. Thanks in advance.

Answer 1

Assuming a the input and using indexing with unravel_index :

x, y = np.unravel_index(np.arange(a.size), a.shape)

out = np.zeros(a.shape+(a.shape[-1],), dtype=a.dtype)

out[x, y, y] = a.flat

Output:

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

       [[3, 0],
        [0, 4]],

       [[5, 0],
        [0, 6]],

       [[7, 0],
        [0, 8]]])

timings:

Answer 2

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