使用 numpy 对多个向量进行对角化

Question

Say I have a matrix of shape (2,3), I need to diagonalize the 3-elements vector into matrix of shape (3,3), for all the 2 vectors at once. That is, I need to return matrix with shape (2,3,3). How can I do that with Numpy elegantly?

given data = np.array([[1,2,3],[4,5,6]])

i want the result [[[1,0,0],
                    [0,2,0],
                    [0,0,3]],
                    
                   [[4,0,0],
                    [0,5,0],
                    [0,0,6]]]
             

Thanks

Answer 1

tl;dr, my one-liner: mydiag=np.vectorize(np.diag, signature='(n)->(n,n)')

I suppose here that by "diagonalize" you mean "applying np.diag". Which, as a teacher of linear algebra, tickles me a bit. Since "diagonalizing" has a specific meaning, which is not that (it is computing eigen vectors and values, and from there, writing M=P⁻¹ΛP. Which you cannot do from the inputs you have).

So, I suppose that if input matrix is

[[1, 2, 3],
 [9, 8, 7]]

The output matrix you want is

[[[1, 0, 0],
  [0, 2, 0],
  [0, 0, 3]],
 [[9, 0, 0],
  [0, 8, 0],
  [0, 0, 7]]]

If not, you can ignore this answer [Edit: in the meantime, you explained exactly that. So yo may continue to read].

There are many way to do that. My one liner would be

mydiag=np.vectorize(np.diag, signature='(n)->(n,n)')

Which build a new functions which does what you want (it interprets the input as a list of 1D-array, call np.diag of each of them, to get a 2D-array, and put each 2D-array in a numpy array, thus getting a 3D-array)

Then, you just call mydiag(M)

One advantage of vectorize, is that it uses numpy broadcasting. In other words, the loops are executed in C, not in python. In yet other words, it is faster. Well it is supposed to be (on small matrix, it is in fact slower than Michael's method - in comment; on large matrix, it is has the exact same speed. Which is frustrating, since einsum doc itself specify that it sacrifices broadcasting).

Plus, it is a one-liner, which has no other interest than bragging on forums. But well, here we are.

Answer 2

Here is one way with indexing:

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

x,y = np.indices(data.shape).reshape(2, -1)
out[x,y,y] = data.ravel()

output:

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

       [[4, 0, 0],
        [0, 5, 0],
        [0, 0, 6]]])

Answer 3

We use array indexing to precisely grab those elements that are on the diagonal. Note that array indexing allows broadcasting between the indices, so we have index1 contain the index of the array, and index2 contain the index of the diagonal element.

index1 = np.arange(2)[:, None] # 2 is the number of arrays
index2 = np.arange(3)[None, :] # 3 is the square size of each matrix

result = np.zeros((2, 3, 3))
result[index1, index2, index2] = data