python - 优化嵌套 for 循环

Question

我尽量避免使用 for 循环来运行我的计算。但我不知道该怎么做。我有一个形状为 (40,100) 的矩阵 w。每条线在 t 时间内保持波浪的位置。例如，第一行 w[0] 是初始条件(也是 w[1]，原因我将展示)。

要计算我使用的下一个线条元素，对于形状范围上的每个 t 和 x:

w[t+1,x] = a * w[t,x] + b * ( w[t,x-1] + w[t,x+1] ) - w[t-1,x]

其中a和b是一些基于方程求解的常量(真的无所谓)，a = 2(1-r)，b=r，r=(c*(dt/dx))**2。其中c是波速，dt，dx与x和t上的增量有关 方向。

有什么办法可以避免这样的for循环:

for t in range(1,nt-1):
    for x in range(1,nx-1):
      w[t+1,x] = a * w[t,x] + b * ( w[t,x-1] + w[t,x+1] ) - w[t-1,x]

nt 和nx 是w 矩阵的形状。

Answer 1

我假设你事先设置了 w[:,0] 和 w[:-1] (一些常量？)因为我没有看到它在循环。如果是这样，您可以消除对这部分代码进行矢量化的 for x 循环:

for t in range(1,nt-1):
    w[t+1,1:-1] = a*w[t,1:-1] + b*(w[t,:-2] + w[t,2:]) - w[t-1,1:-1]