优化嵌套 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。其中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]

Answer 2

不是真的。如果你想为矩阵中的每个元素做一些事情（你做了），你将不得不以某种方式对每个元素进行操作（最明显的方法是使用 for 循环。不太明显的方法要么执行相同或更糟）。

如果您因为循环很慢而试图避免循环，请知道有时循环对于解决某种问题是必要的。但是，有很多方法可以提高循环的效率。

通常对于像这样的矩阵问题，您正在查看相邻元素，一个好的解决方案是使用某种动态编程或记忆（节省您的工作，因此您不必经常重复计算）。就像，假设你想对每个元素和它周围的所有东西取平均值（这就是模糊图像的工作原理）。每个像素有 8 个相邻像素，因此平均值将是总和 / 9。好吧，假设您保存列的总和（保存 NW + W + SW、N + me + S、NE + E + SE）。好吧，当您转到右侧的下一个时，只需将上一个中间列的值、上一个最后一列的值和一个新列的值（右边的新列）相加。您刚刚用加 5 代替了加 9。在比加更复杂的运算中，将 9 减少到 5 可能意味着性能的巨大提升。

我查看了您必须做的事情，但我想不出一个像我刚才描述的那样做的好方法。但是看看你能不能想出类似的东西。

另外，请记住乘法比加法更昂贵。因此，如果您有一个循环，例如，您必须将某个数字乘以循环变量，而不是执行 1x、2x、3x...，您可以执行 (value last time + x)。