如何在Numpy中对这个double for循环进行向量化？

Question

我有一些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

我正在尝试删除双for循环并将Z矢量化。这是我的尝试。在

^{pr2}$

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

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

所以在尝试矢量化Z时，我搞砸了。你能帮我找出我的错误吗？在

Answer 1

由于计算列i依赖于列i-1（在您的第二部分代码中）除了第一列之外，其余的值都是零，所以不能同时删除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):
        Z[1:-1,i] = Z[1:-1,i-1] + D*q*(Z[:-2,i-1] - 2*Z[1:-1,i-1] + Z[2:,i-1])
    return Z

返回代码： 您正在用（仅第一项）Z[1:-1，：-1]填充Z[1,1:-1]。形状上的不匹配在这里很明显。在

不考虑第二个索引（因为无论如何都要循环），矢量化解决方案使用了与非矢量化方法不同的假设：在第一个版本中，您有一侧u0（Z[：，0]）和两侧0（Z[0，：]和Z[-1，：]），而在矢量化解决方案中，您确实尝试通过填充Z[1:，i]将Z[-1，：]设置为0以外的值。你想模拟什么情况？！在

Answer 2

你不能在问题的时间维度上对扩散计算进行矢量化，这仍然需要一个循环。这里唯一明显的优化是用调用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

对于非平凡的空间大小，您应该会看到相当大的加速。在