[SOLVED] Numpy multiply each slice of a 3D array for its transpose and sum them

Numpy multiply each slice of a 3D array for its transpose and sum them

I have a 3D numpy array with shape (N, 3, 3).

I want to multiply each 3x3 matrix by its transpose and then sum all of them.

This is my code:

m = np.zeros([3, 3])


for matrix in matrices:
    m += np.dot(matrix.T, matrix)

but I think it could be done just using numpy.

I was thinking to use np.apply_along_axis, but it works only with 1D arrays.

Is there a way to avoid the for loop and make the product and sum in one go?

Solution

With a sample array:

In [22]: arr=np.arange(4*9).reshape(4,3,3)
In [23]: m = np.zeros([3, 3])
    ...: for matrix in arr:
    ...:     m += np.dot(matrix.T, matrix)
    ...:     
In [24]: m
Out[24]: 
array([[4554., 4752., 4950.],
       [4752., 4962., 5172.],
       [4950., 5172., 5394.]])

A matmul solution, transposing the trailing (3,3) arrays, and summing on the lead dimension:

In [26]: (arr.transpose(0,2,1)@arr).sum(axis=0)
Out[26]: 
array([[4554, 4752, 4950],
       [4752, 4962, 5172],
       [4950, 5172, 5394]])

And the einsum approach as given in the other answer.

In [31]: np.einsum('ikj,ikm->jm', arr, arr)
Out[31]: 
array([[4554, 4752, 4950],
       [4752, 4962, 5172],
       [4950, 5172, 5394]])

I had to try several einsum indices to get the matching one. Summing on the shared i dimension was easy, but get the 'transpose' right required some trial-n-error or systematic thinking.

Seeing the repeated ik in the einsum, suggests an alternative - reshape to 2d, and do the normal dot:

In [35]: arr1 = arr.reshape(-1,3); arr1.T@arr1
Out[35]: 
array([[4554, 4752, 4950],
       [4752, 4962, 5172],
       [4950, 5172, 5394]])

For this small sample case this is 2x faster.