Search code examples
pythonperformancenumpymatrixscientific-computing

Calculating the outer product for a sequence of numpy ndarrays


I have a list of 3D points p stored in an ndarray with shape (N, 3). I want to compute the outer product for each 3d point with itself:

N = int(1e4)
p = np.random.random((N, 3))
result = np.zeros((N, 3, 3))
for i in range(N):
    result[i, :, :] = np.outer(p[i, :], p[i, :])

Is there a way to compute this outer product without any python-level loops? The problem is that np.outer does not support anything like an axis argument.


Solution

  • You can use broadcasting:

    p[..., None] * p[:, None, :]
    

    This syntax inserts an axis at the end of the first term (making it Nx3x1) and the middle of the second term (making it Nx1x3). These are then broadcast and yield an Nx3x3 result.