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

Question

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?

Answer 1

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.

Answer 2

You can use numpy.einsum:

np.einsum('ijk,ijl->kl', matrices, matrices)

Example:

np.random.seed(0)
N = 5
matrices = np.random.random((N, 3, 3))

np.einsum('ijk,ijl->kl', matrices, matrices)

array([[5.1367768 , 4.33942585, 5.02664464],
       [4.33942585, 5.73518514, 4.72942748],
       [5.02664464, 4.72942748, 6.56237687]])

Output with your loop:

array([[5.1367768 , 4.33942585, 5.02664464],
       [4.33942585, 5.73518514, 4.72942748],
       [5.02664464, 4.72942748, 6.56237687]])