"reshape" numpy array of (N, 2) shape into (N, 2, 2) where each column (size 2) become a diag (2,2) block?

Question

Is there an efficient way to do this? For example I have

[[1, 2, 3],
 [4, 5, 6]]

I would like to get:

[[[1, 0],
  [0, 4]],

 [[2, 0],
  [0, 5]],

 [[3, 0],
  [0, 6]]]

Answer 1

For large arrays I recommend np.einsum as follows:

>>> data
array([[1, 2, 3],
       [4, 5, 6]])
>>> out = np.zeros((*reversed(data.shape),2),data.dtype)
>>> np.einsum("...ii->...i",out)[...] = data.T
>>> out
array([[[1, 0],
        [0, 4]],

       [[2, 0],
        [0, 5]],

       [[3, 0],
        [0, 6]]])

einsum creates a writable strided view of the memory locations holding the diagonal elements. This is about as efficient as it gets in numpy.