CODEWITHSUNDEEP
pythonnumpy
bzm3r
bzm3r

Reputation: 4596

Faster alternative to numpy.einsum for taking the "element-wise" dot product of two lists of vectors?

Let's say you're given two arrays of vectors:

v1 = np.array([ [1, 2], [3, 4] ]) v2 = np.array([ [10, 20], [30, 40]])

We would like to generate an array that is equivalent to:

v3 = np.array([ np.dot(v1[0], v2[0]), np.dot(v1[1], v2[1]) ])

Currently I use:

v3 = np.einsum('ij,ij->i', v1, v2)

However, I do this a lot in my code, so speed ups here would be very helpful for me.

How could we speed it up? np.einsum is already quite efficient, but I wonder if for this particular use-case, there is a faster way?

Upvotes: 3

Views: 2043

Answers (2)

Fnord
Fnord

Reputation: 5895

Try inner1d

import numpy as np
import cProfile
from numpy.core.umath_tests import inner1d

v1 = np.random.random((10**7,2,))  # 10 million random vectors
v2 = np.random.random((10**7,2,))  # 10 million random vectors
v3 = np.einsum('ij,ij->i', v1, v2) # einsum
v4 = inner1d(v1,v2)                # inner1d (~2x faster than einsum)

cProfile.run("np.einsum('ij,ij->i', v1, v2)") # cProfile: 3 function calls in 0.065 seconds
cProfile.run("inner1d(v1,v2)") # cProfile: 2 function calls in 0.033 seconds

print np.allclose(v3,v4) # Returns True

Upvotes: 1

hpaulj
hpaulj

Reputation: 231385

einsum does the best of 3 options that I can think of:

In [73]: timeit v3=np.einsum('ij,ij->i',v1,v2)
100000 loops, best of 3: 5.14 us per loop

In [74]: timeit np.diag(np.dot(v1,v2.T))
100000 loops, best of 3: 7.43 us per loop

In [75]: timeit np.sum(v1*v2,axis=1)
100000 loops, best of 3: 16.8 us per loop

Several questions to ask:

  • is this calculation really that expensive ?
  • if it is relatively expensive do you have to doing so often ?
  • can you consolidate the einsum calls - concatenate arrays ?

Upvotes: 2

Related Questions