Python: Dot product of each vector in two lists of vectors

Python: Dot product of each vector in two lists of vectors

The fastest method here is using einsum

out = numpy.einsum(ij, ij->i, A, B)

It outperforms the multiply-add version by a factor of 2; the list comprehension version is a few hundred times slower.

enter


Code to reproduce the plot:

import numpy
import perfplot

perfplot.show(
    setup=lambda n: (numpy.random.rand(n, 3), numpy.random.rand(n, 3)),
    kernels=[
        lambda data: numpy.einsum(ij, ij->i, data[0], data[1]),
        lambda data: numpy.multiply(data[0], data[1]).sum(1),
        lambda data: [
            sum(ai * bi for ai, bi in zip(a, b)) for a, b in zip(data[0], data[1])
        ],
    ],
    labels=[einsum, multiply+sum, sum+zip],
    n_range=[2 ** k for k in range(18)],
    xlabel=len(a),
)

In pure Python, try a nested list/generator comprehension:

>>> [sum(ai * bi for ai, bi in zip(a, b))
...  for a, b in zip(A, B)]
[2, 0]

or with numpy, you can do an element-wise product of the 2-dimensional arrays, followed by a sum along each row:

>>> import numpy as np
>>> np.multiply(A, B).sum(1)
array([2, 0])

Note that the numpy solution will only work if all vectors are the same length – each list of vectors is implicitly converted into a two-dimensional matrix.

Python: Dot product of each vector in two lists of vectors

We can use zipping, summing, and list comprehensions for a fancy one-liner:

A = [(1,1,1), (0,1,1)]
B = [(1,0,1), (1,0,0)]
C = [sum(i*j for i, j in zip(a, b)) for a, b in zip(A, B)]
print(C)  # [2, 0]

Leave a Reply

Your email address will not be published. Required fields are marked *