Multiplication for a 3d array and slicing

Question

I have a matrix of size 5 x 98 x 3. I want to find the transpose of each block of 98 x 3 and multiply it with itself to find the standard deviation. Hence, I want my final answer to be of the size 5 x 3 x 3. What would be an efficient way of doing this using numpy.

I can currently do this using the following code:

MU.shape[0] = 5
rows = 98
SIGMA = []
    for i in np.arange(MU.shape[0]):
        SIGMA.append([])
        SIGMA[i] = np.matmul(np.transpose(diff[i]),diff[i])
    SIGMA = np.array(SIGMA)
    SIGMA = SIGMA/rows

Here diff is of the size 5 x 98 x 3.

Answer 1

Use np.einsum to sum-reduce the last axes off against each other -

SIGMA = np.einsum('ijk,ijl->ikl',diff,diff)
SIGMA = SIGMA/rows

Use optimize flag with True value in np.einsum to leverage BLAS .

We can also use np.matmul to get those sum-reductions -

SIGMA = np.matmul(diff.swapaxes(1,2),diff)

Answer 2

You can use this:

my_result = arr1.swapaxes(1,2) @ arr1

Testing it out:

import numpy as np

NINETY_EIGHT = 10
arr1 = np.arange(5*NINETY_EIGHT*3).reshape(5,NINETY_EIGHT,3)

my_result = arr1.swapaxes(1,2) @ arr1
print (my_result.shape)

Output:

(5, 3, 3)