Is there any way to calculate the diagonal value of matrix C which is equal to numpy.matmul(A, B) without the need to calculation off-diagonal elements? Because I am working with two large matrices A and B and I only need to have diagonal elements of their multiplication in the rest of my code. And the run time is considerably large if I calculate the whole matrix C. Thanks in advance.
You can achieve it as :
arr1 = np.array([[1,2,3],[4,5,6],[7,8,9]])
arr2 = np.array([[10,11,12],[13,14,15],[16,17,18]])
diag_elem = [sum(arr1[i,]*arr2[:,i]) for i in range(arr1.shape[0])]
print(diag_elem)
output:
[84, 216, 366]
which are the same diagonal elements you get using the np.matmul
print(np.matmul(arr1,arr2))
results in
array([[ 84, 90, 96],
[201, 216, 231],
[318, 342, 366]])