5

我有一些运行缓慢的 Python / Numpy 代码,我认为这是因为使用了双 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。这是我的尝试。

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 的某个地方,我搞砸了。你能帮我找出我的错误吗?

4

2 回答 2

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

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

于 2012-10-19T07:55:59.867 回答
1

您将无法删除两个 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 以外的值。您要模拟哪种情况?!

于 2012-10-19T07:59:29.163 回答