Optimize a double loop with mesh grids involved

user_7301760 0
Viewed 45

I am doing a double loop to sum a function that has mesh grids as an input. The problem is that it runs very slow... I want to optimize the code with an alternative procedure, maybe using vectorize function of numpy, but I don't see how can be implemented. I show you the code that I have:

import numpy as np
import time

Lxx = 2.
Lyy = 1.0
dxx = dyy = 0.01
nxx = 100
nyy = 100
XX, YY = np.meshgrid(np.arange(0, Lxx+dxx, dxx), np.arange(0, Lyy+dyy, dyy)) #mesh grid


def solution(xx,yy,nnmax,mmmax):
    sol = 0.
    for m in range(nnmax):
        for n in range(mmmax):
           sol = sol+np.sin(XX*0.356*n)+np.cos(YY*2.3*m)
    return sol


start = time.time()  
solution(XX,YY,nxx,nyy)
end = time.time()   
print ("TIME", end-start)

What I want is to make the sum for large values in nxx, nyy. But of course then it takes a lot of time...This is the reason why I want optimize the code.

numpy optimization python
1 Answers

If you notice, the terms of the sum are completely separable: they don't share any loop variables. You can therefore create independent (smaller) arrays for the sum over XX, n and YY, m, and take the trig functions and sum of those. The final grid can be accumulated by broadcasting.

To begin, don't bother making the grid:

x = np.arange(0, Lxx+dxx, dxx)
y = np.arange(0, Lyy+dyy, dyy)

Compute a single sum using broadcasting:

n = np.arange(nyy)[:, None]
m = np.arange(nxx)[:, None]
sumx = np.sin(x + 0.356 * n).sum(0)
sumy = np.cos(y + 2.3 * m).sum(0)

You can use the same broadcasting trick to get the final sum in a grid:

result = sumx[:, None] + sumy
user_29887300
Related