Trying to translate for loops into numpy matrix operation

Question

This is a bit of an odd question, but I am trying to improve the runtime of my code by losing the for loops and relying on numpy operations. I am still a beginner at handling numpy matrix operations, and I cant seem to translate this properly.

for i in range (Ma):

    for j in range (Na):
    
        for h in range(Mt):
        
            for k in range(Nt):

                dx = xrfa[i] - xrft[h]
                dy = yrfa[j] - yrft[k] 

                Wat[i,j,h,k] = at * np.exp(- ((np.square(dx) + np.square(dy))/ (2  * np.square(sigat)

Any help would be appreciated. I specifically struggle with the dx and dy.

Asked By: Grigoras Vlad Albastrel6

||

Source

Answer 1

Rely on the power of broadcasting!

Setup:

import numpy as np

Ma, Na, Mt, Nt = 2, 3, 4, 5

# Assuming these are all just 1D arrays
xrfa = np.arange(Ma)
yrfa = np.arange(Na)
xrft = np.arange(Mt)
yrft = np.arange(Nt)

# and these scalars:
at = 2
sigat = 3

You could (ab)use np.ix_ to introduce singleton dimensions, but I’ll do it explicitly here:

xrfa = xrfa[:, None, None, None]
yrfa = yrfa[None, :, None, None]
xrft = xrft[None, None, :, None]
yrft = yrft[None, None, None, :]

Now your arrays all broadcast to shape (Ma, Na, Mt, Nt), at which point you can use a one liner to all the calculations (hopefully I fixed your parentheses properly):

Wat = at * np.exp(-(np.square(xrfa - xrft) + np.square(yrfa - yrft)) / (2  * np.square(sigat)))

Answered By: Chrysophylaxs

Trying to translate for loops into numpy matrix operation

Question:

Answers: