# Replacing nested loops in Python using NumPy

36
April 08, 2021, at 00:40 AM

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.

Answer 1

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
``````
POPULAR ONLINE

#### No responsive page when redirected from facebook

READ ALSO ### how can I make the arrow border right background colorful

I need to make the arrow of the right side colorful #2b6a83

22 ### can't download wordpress posts using javascript

hello i have a list of posts and i use javascript to download a post as pdf but i have a problem because the download button works just on the first post and not in the others and when i inspect the button in the browser i can see the description below but the button...

22 ### How to add string into a RealmList Kotlin

it`s my first question in stackOverflow, and i need help with how to insert data (Strings) into a RealmList

35 ### Selecting Names with Highest Aggregate in each Class

I'm new to SQL and I'm working with the following tables:

33