CODEWITHSUNDEEP
pythonnumpy
lehoanglittle
lehoanglittle

Reputation: 25

Vectorized computation on Numpy array

Currently, I am using the loops to do the task below. It is quite similar to my previous question: Dot product two 4D Numpy array

Just wondering if anyone can help me to vectorize the computation without the loops.

A1 = np.random.rand(2, 320, 320)
A2 = np.random.rand(15, 2, 320, 320)
B   = np.zeros((2, 320, 320))

for row in range (320):
    for col in range (320):
        C = np.diag(A1[:, row, col]) - (np.dot(A2[:, :, row, col].T, A2[:, :, row, col]) / 15) # the shape of array C is (2, 2)
        C = np.diag(C)                
        B[:, row, col] = C

Any help or suggestions would be greatly appreciated!

Upvotes: 0

Views: 78

Answers (1)

Akshay Sehgal
Akshay Sehgal

Reputation: 19322

Similar to the last solution -

output = A1-(A2*A2).sum(0)/15
np.allclose(output,B)

True

Upvotes: 1

Related Questions