3

Say I have a matrix A of dimension N by M.

I wish to return an N dimensional vector V where the nth element is the double sum of all pairwise product of the entries in the nth row of A.

Formula

In loops, I guess I could do:

V = np.zeros(A.shape[0])
for n in range(A.shape[0]):
    for i in range(A.shape[1]):
        for j in range(A.shape[1]):
            V[n] += A[n,i] * A[n,j]

I want to vectorise this and I guess I could do:

V_temp = np.einsum('ij,ik->ijk', A, A)
V = np.einsum('ijk->i', A)

But I don't think this is very memory efficient way as the intermediate step V_temp is unnecessarily storing the whole outer products when all I need are sums. Is there a better way to do this?

Thanks

Tadhg McDonald-Jensen
  • 20,699
  • 5
  • 35
  • 59
chanyoungs
  • 33
  • 2

3 Answers3

2

You can use 

V=np.einsum("ni,nj->n",A,A)
Roun
  • 1,449
  • 1
  • 18
  • 25
2

You are actually calculating

A.sum(-1)**2

In other words, the sum over an outer product is just the product of the sums of the factors.

Demo:

A = np.random.random((1000,1000))
np.allclose(np.einsum('ij,ik->i', A, A), A.sum(-1)**2)
# True
t = timeit.timeit('np.einsum("ij,ik->i",A,A)', globals=dict(A=A,np=np), number=10)*100; f"{t:8.4f} ms"
# '948.4210 ms'
t = timeit.timeit('A.sum(-1)**2', globals=dict(A=A,np=np), number=10)*100; f"{t:8.4f} ms"
# '  0.7396 ms'
Paul Panzer
  • 51,835
  • 3
  • 54
  • 99
0

Perhaps you can use

np.einsum('ij,ik->i', A, A)

or the equivalent

np.einsum(A, [0,1], A, [0,2], [0])

On a 2015 Macbook, I get

In [35]: A = np.random.rand(100,100)

In [37]: %timeit for_loops(A)
640 ms ± 24.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [38]: %timeit np.einsum('ij,ik->i', A, A)
658 µs ± 7.25 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [39]: %timeit np.einsum(A, [0,1], A, [0,2], [0])
672 µs ± 19.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
James Lim
  • 12,915
  • 4
  • 40
  • 65