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

我正在尝试删除double 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

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

Answer 2

您将无法删除两个for循环，因为计算列i依赖于第i-1列（第二位代码中除第一列之外的零）。

你能做的是：

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以外的值。你想尝试模拟哪种情况？！

如何在Numpy中向量化这个双循环？

2 个答案: