优化嵌套的for循环

Question

我正在尝试避免使用for循环来运行我的计算.但我不知道该怎么做.我有一个w形状矩阵(40,100).每条线在一段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在哪里,与增量和方向有关.cdtdxxt

有没有办法避免像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]