Reputation: 65
I am trying to vectorize my code using numpy modules, the original code is like:
m = [M[i,:9].dot(N[:9,i]) for i in xrange(9)]
and I have improved the code as:
m = np.diagonal(M[:9,:9].dot(N[:9,:9]))
yet this will lead to some unnecessary calculations (especially when the index is much greater than 9). What can I do to improve the efficiency further?
Edit: Basically what I intend to do is to calculate the diagonal elements of the dot product of two matrixes M and N.
Upvotes: 4
Views: 132
Reputation: 18552
Assuming that both matrices are square and have the same dimensions, this is what we can do:
By the definition of matrix multiplication, the result you want to calculate is basically a vector of dot products. The first element of the result vector is the dot product of M's first row with N's first column, followed by the dot product of M's second row with N's second column, etc.
We can express this calculation by transposing N, then multiplying M and NT element-wise, then adding up each row. This gives us a resulting column vector. This is the code:
Nprime = np.transpose(N[:,:9])
product = np.multiply(M[:9,:], Nprime)
result = np.sum(product, axis=1)
By Divakar's suggestion, it can be condensed into one line:
result = (M[:,:9] * (N[:9,:].T)).sum(1)
Example (with size 3 instead of 9):
import numpy as np
>>> M = np.array(
[[1, 2, 3],
[6, 5, 4],
[8, 7, 9]])
>>> N = np.array(
[[0, 9, 7],
[2, 4, 5],
[6, 8, 1]])
>>> M.dot(N)
array([[22, 41, 20],
[34, 106, 71],
[68, 172, 100]])
>>> np.diagonal(M.dot(N))
array([22, 106, 100])
# ^ The reference answer
>>> Nprime = np.transpose(N)
array([[0, 2, 6],
[9, 4, 8],
[7, 5, 1]])
>>> product = np.multiply(M, Nprime)
array([[ 0, 4, 18],
[54, 20, 32],
[56, 35, 9]])
>>> result = np.sum(product, axis=1)
array([22, 106, 100])
Upvotes: 4