将矩阵的每一行与其共轭转置的numpy相乘

Question

I have a numpy.ndarray variable A of size MxN . I wish to take each row and multiply with it's conjugate transposed. For the first row we will get:

np.matmul(np.expand_dims(A[0,:],axis=1),np.expand_dims(A[0,:].conj(),axis=0))

we get an NxN sized result. I want the final result for the total operation to be of size MxNxN .

I can fo this with a simple loop which iterates over the rows of A and concatenates the results. I wish to avoid a for loop for a faster run time with SIMD operations. Is there a way to do this in a single code line with broadcasting?

Otherwise, can I do something else and somehow reshape the results into my requierment?

Answer 1

The next code does what the same as your code snippet but without for-loop. On the other hand, it uses np.repeat twice, so you will need to benchmark both versions and compare them to test their memory/time performance.

import numpy as np

m, n = A.shape

x, y = A.conj().repeat(n, axis=0), A.reshape([-1, 1]).repeat(n, axis=1)
B = (x * y).reshape([m, n, n])

How it works

Basically x holds the conjugate values of the array A in a single column and then is repeated n times on the column axis (it has a shape m*n by n ).

y repeats each row in the conjugate matrix of A , n consecutive times (its final shape is m*n by n also)

x and y are multiplied element-wise and the result is unwrapped to a matrix of shape m by n by n stored in B

Answer 2

列表理解可以解决这个问题：

result = np.array([np.matmul(np.expand_dims(A[i,:],axis=1), np.expand_dims(A[i,:].conj(),axis=0)) for i in range(A.shape[0])])