M.Gu - 1 year ago 34

Python Question

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.

Answer Source

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,:])
```