How to use broadcasting to speed up this code?

Question

I have the following multi-dimensional array. The first axis denotes a 3-dimensional vector. I want to calculate the 3-by-3 matrix x⋅x' for each of those.

My current solution:

arr.shape
# (3, 64, 64, 33, 187)

dm = arr.reshape(3,-1)
dm.shape
# (3, 25276416)

cov = np.empty((3,3,dm.shape[1]))
cov.shape
# (3, 3, 25276416)

this for-loop iterates over all 25,276,416 elements and takes around 1 or 2 min.

for i in range(dm.shape[1]):
    cov[...,i] = dm[:,i].reshape(3,1).dot(dm[:,i].reshape(1,3))

cov = cov.reshape((3,) + arr.shape)
cov.shape
# (3, 3, 64, 64, 33, 187)

Answer 1

Well you are not really reducing any axes with that matrix-multiplication using np.dot and it's just broadcasted elementwise multiplication there. So, you can simply use NumPy broadcasting for the whole thing, like so -

cov = dm[:,None]*dm

Or use it directly on arr to avoid creating dm and all that reshaping, like so -

cov = arr[:,None]*arr