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
READ ALSO
how can I make the arrow border right background colorful

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

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

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

Selecting Names with Highest Aggregate in each Class

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

33