CODEWITHSUNDEEP
pythonperformancenumpyvectorizationmatrix-multiplication
M.Gu
M.Gu

Reputation: 65

how to vectorize this python code to make it more efficient? (in speed)

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

Answers (2)

Nayuki
Nayuki

Reputation: 18552

Assuming that both matrices are square and have the same dimensions, this is what we can do:

  1. 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.

  2. 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

Divakar
Divakar

Reputation: 221774

You can use np.einsum as we need to keep the first axis of M aligned with the second axis of N, while reducing/losing the leftover axes from the inputs. Thus, we would have an einsum based solution, like so:

m = np.einsum('ij,ji->i',M[:,:9],N[:9,:])

Upvotes: 6

Related Questions