Python - fast way to sum outer products?

Question

I'm looking for a fast way to calculate a sum of n outer products.

Essentially, I start with two matrices generated from normal distributions - there are n vectors with v elements:

A = np.random.normal(size = (n, v))
B = np.random.normal(size = (n, v))

What I'd like is to calculate the outer products of each vector of size v in A and B and sum them together.

Note that A * B.T doesn't work - A is of size n x v whereas B is of size v x n.

The best I can do is create a loop where the outer products are constructed, then summed later. I have it like so:

outers = np.array([A[i] * B[i].T])

This creates an n x v x v array (the loop is within the list comprehension, which is subsequently converted into an array), which I can then sum together by using np.sum(outers, axis = 0). However, this is quite slow, and I was wondering if there's a vectorized function I could use to speed this up.

If anybody has any advice, I would really appreciate it!

Answer 1

It seems to me all you need to do is change the order of the transpositions, and do A.T * B instead of A * B.T.

If that's not quite what you are after, take a look at np.einsum, which can do some very powerful voodoo. For the above example, you would do:

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

Answer 2

Also consider np.outer.

np.array([np.outer(A, B) for i in xrange(n)]).sum(0)

although np.einsum suggested by @Jamie is the clear winner.

In [63]: %timeit np.einsum('ij,ik->jk', A, B)
100000 loops, best of 3: 4.61 us per loop

In [64]: %timeit np.array([np.outer(A[i], B[i]) for i in xrange(n)]).sum(0)
10000 loops, best of 3: 169 us per loop

and, to be sure, their results are identical:

In [65]: np.testing.assert_allclose(method_outer, method_einsum)

But, as an aside, I do not find that A.T * B or A * B.T broadcast successfully.