Replacing nested loops in Python using NumPy

Question

So I have these 6 nested loops, and their purpose is only to multiply and add arrays X and Y over different indices to get array Z.

import numpy as np

dim_a = 5
dim_b = 9
Z = np.zeros((dim_a,dim_b,dim_b,dim_a))
X = np.arange(2025).reshape(dim_a,dim_b,dim_b,dim_a)
Y = np.arange(2025).reshape(dim_a,dim_b,dim_b,dim_a)

for i in range(0, dim_a):
    for j in range(0,dim_a):
        for a in range(0, dim_b):
            for b in range(0, dim_b):
                for m in range(0,dim_a):
                    for e in range(0,dim_b):
                        Z[i,a,b,j] += X[m,e,b,j] * Y[m,e,a,i] * 2

Is there a way to write it using just a few lines of code using NumPy? The computational effort of these nested loops is enormous. I have a feeling that NumPy can optimize it significantly.

Quang Hoang · Accepted Answer · 2021-04-07 17:40:30Z

2

You can certainly do with np.einsum:

Z[i,a,b,j] += X[m,e,b,j] * Y[m,e,a,i] * 2

translates to

Z = np.einsum('mebj,meai->iabj', X,Y) * 2

answered Apr 7, 2021 at 17:40

Quang Hoang

151k11 gold badges64 silver badges86 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Grzegorz Skrzyński Over a year ago

This is amazing! I checked the computational times: 6 nested loops: 0.14539456367492676 np.einsum: 0.000209808349609375

Kevin Over a year ago

Also the same as np.einsum('mebj,meai,...->iabj', X, Y, 2). Although, not sure if this is any faster.

Collectives™ on Stack Overflow

Replacing nested loops in Python using NumPy

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related