Using numpy einsum to compute inner product of column-vectors of a matrix

Question

Suppose I have a numpy matrix like this:

[[   1    2    3]
 [  10  100 1000]]

I would like to compute the inner product of each column with itself, so the result would be:

[1*1 + 10*10    2*2 + 100*100    3*3 + 1000*1000] == [101, 10004, 1000009]

I would like to know if this is possible using the einsum function (and to better understand it).

So far, the closest result I could have is:

import numpy as np

arr = np.array([[1, 2, 3], [10, 100, 1000]])

res = np.einsum('ij,ik->jk', arr, arr)

# [[    101    1002   10003]
#  [   1002   10004  100006]
#  [  10003  100006 1000009]]

The diagonal contains the expected result, but I would like to know if I can avoid edge calculations.

Answer 1

Use np.einsum, like so -

np.einsum('ij,ij->j',arr,arr)

Sample run -

In [243]: np.einsum('ij,ij->j',arr,arr)
Out[243]: array([    101,   10004, 1000009])

Or with np.sum -

In [244]: (arr**2).sum(0)
Out[244]: array([    101,   10004, 1000009])

Or with numexpr module -

In [248]: import numexpr as ne

In [249]: ne.evaluate('sum(arr**2,0)')
Out[249]: array([    101,   10004, 1000009])

Answer 2

What you're expecting here can be understood intuitively, with one intermediate step from Divakar's einsum answer.

In [19]: arr
Out[19]: 
array([[   1,    2,    3],
       [  10,  100, 1000]])

# simply take element-wise product with the array itself
In [20]: np.einsum('ij, ij -> ij', arr, arr)
Out[20]: 
array([[      1,       4,       9],
       [    100,   10000, 1000000]])

But, this doesn't give the result that you expected. So, if you observe the above result, we just need to sum along the first dimension (i.e. axis 0). So, we omit the subscript i after -> in the einsum result, which means we ask it to sum along that axis, and that yields the expected result which is:

In [21]: np.einsum('ij, ij -> j', arr, arr)
Out[21]: array([    101,   10004, 1000009])

---

P.S. Also, for a general understanding of np.einsum, see the detailed discussion here: understanding-numpy-einsum