DSKim DSKim - 2 months ago 16
Python Question

how to efficiently compute the column-wise outer vector in numpy

Suppose that the matrix A has the size m x r and that B has n x r. In numpy, how can I do column-wise outer-product between A and B?

That is, I want a matrix, C (size: m x n x r), where C[:, :, i] is the outer product of A[:, i] and B[:, i].

Of course, I can use a for-loop over the columns but I am wondering if there is a vectorized way to make the computation fast.


I usually call on einsum for easy vectorization in cases like this:

>>> m,n,r = 30,50,70
>>> A = np.random.random((m,r))
>>> B = np.random.random((n,r))
>>> C = np.einsum('ik,jk->ijk', A, B)
>>> C2 = np.dstack([np.outer(A[:,i], B[:,i]) for i in range(r)])
>>> C.shape
(30, 50, 70)
>>> np.allclose(C, C2)

This way we can be explicit about what we want to happen to each coordinate, while still being faster than a loop:

>>> %timeit C = np.einsum('ik,jk->ijk', A, B)
The slowest run took 4.83 times longer than the fastest. This could mean that an intermediate result is being cached.
1000 loops, best of 3: 216 ┬Ás per loop

>>> %timeit C2 = np.dstack([np.outer(A[:,i], B[:,i]) for i in range(r)])
100 loops, best of 3: 2.08 ms per loop