如何在Numpy中将这个double for循环向量化？(How do I vectorize this double for loop in Numpy?)

如何在Numpy中将这个double for循环向量化？(How do I vectorize this double for loop in Numpy?)

我有一些运行缓慢的Python / Numpy代码，我认为这是因为使用了double for循环。 这是代码。

def heat(D,u0,q,tdim): xdim = np.size(u0) Z = np.zeros([xdim,tdim]) Z[:,0]=u0; for i in range(1,tdim): for j in range (1,xdim-1): Z[j,i]=Z[j,i-1]+D*q*(Z[j-1,i-1]-2*Z[j,i-1]+Z[j+1,i-1]) return Z

我试图删除双循环和向量化Z.这是我的尝试。

def heat(D,u0,q,tdim): xdim = np.size(u0) Z = np.zeros([xdim,tdim]) Z[:,0]=u0; Z[1:,1:-1]=Z[1:-1,:-1]+D*q*(Z[:-2,:-1]-2*Z[1:-1,:-1]+Z[2:,:-1]) return Z

这不起作用 - 我得到以下错误：

operands could not be broadcast together with shapes (24,73) (23,74)

所以在试图向量化Z的某个地方，我搞砸了。 你能帮我发现我的错误吗？

I have some Python / Numpy code that is running slow and I think it is because of the use of a double for loop. Here is the code.

def heat(D,u0,q,tdim): xdim = np.size(u0) Z = np.zeros([xdim,tdim]) Z[:,0]=u0; for i in range(1,tdim): for j in range (1,xdim-1): Z[j,i]=Z[j,i-1]+D*q*(Z[j-1,i-1]-2*Z[j,i-1]+Z[j+1,i-1]) return Z

I am trying to remove the double for loop and vectorize Z. Here is my attempt.

def heat(D,u0,q,tdim): xdim = np.size(u0) Z = np.zeros([xdim,tdim]) Z[:,0]=u0; Z[1:,1:-1]=Z[1:-1,:-1]+D*q*(Z[:-2,:-1]-2*Z[1:-1,:-1]+Z[2:,:-1]) return Z

This doesn't work - I get the following error:

operands could not be broadcast together with shapes (24,73) (23,74)

So somewhere in trying to vectorize Z, I messed up. Can you please help me spot my error?

最满意答案

您不能在问题的时间维度上将扩散计算向量化，这仍然需要循环。 这里唯一明显的优化是用调用numpy.diff函数（这是预编译的C）来替换拉普拉斯计算，所以你的热方程求解器变成：

def heat(D,u0,q,tdim): xdim = np.size(u0) Z = np.zeros([xdim,tdim]) Z[:,0]=u0; for i in range(1,tdim): Z[1:-1,i]=Z[1:-1,i-1] + D*q*np.diff(Z[:,i-1], 2) return Z

对于非平凡的空间大小，你应该看到相当大的速度。

You cannot vectorize that diffusion calculation in the time dimension of the problem, that still requires a loop. The only obvious optimization here is to replace the Laplacian calculation with a call to the numpy.diff function (which is precompiled C), so your heat equation solver becomes:

def heat(D,u0,q,tdim): xdim = np.size(u0) Z = np.zeros([xdim,tdim]) Z[:,0]=u0; for i in range(1,tdim): Z[1:-1,i]=Z[1:-1,i-1] + D*q*np.diff(Z[:,i-1], 2) return Z

For non-trivial spatial sizes you should see considerable speed up.