优化涉及网格的双循环

Question

我正在做一个双循环来对一个以网格为输入的函数求和。问题是它运行速度很慢......我想用另一种程序优化代码，也许使用numpy的vectorize函数，但我不知道如何实现。我向您展示我拥有的代码：

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)

我想要的是对 nxx, nyy 中的大值求和。但是当然这需要很多时间......这就是我要优化代码的原因。

Answer 1

如果您注意到，总和的项是完全可分离的：它们不共享任何循环变量。因此，您可以为 XX、n 和 YY、m 的总和创建独立（较小的）数组，并获取三角函数和它们的总和。最终的网格可以通过广播来累积。

首先，不要费心制作网格：

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

使用广播计算单个总和：

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)

您可以使用相同的广播技巧来获得网格中的最终总和：

result = sumx[:, None] + sumy