如何在 Numpy 中矢量化这个 double for 循环？

Question

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

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

我正在尝试删除双 for 循环并将 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 的某个地方，我搞砸了。你能帮我找出我的错误吗？

Answer 1

您无法在问题的时间维度中向量化扩散计算，这仍然需要循环。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

对于不平凡的空间大小，您应该会看到相当大的速度提升。