嵌套for循环的矢量化

Question

我有一个函数roughness，它经常在更大的代码中调用。我需要一些帮助，用更简单的矢量化版本替换这个双重for循环。以下是代码：

def roughness(c,d,e,f,z,ndim,half_tile,dx):
    imin=0-half_tile
    imax=half_tile
    z_calc = np.zeros((ndim,ndim), dtype=float)
    for j in range(ndim):
        y=(j-half_tile)*dx
        for i in range(ndim):
            x=(i-half_tile)*dx
            z_calc[i,j] = c*x*y + d*x + e*y + f - z[i,j]

    z_min=z_calc[z_calc!=0].min()
    z_max=z_calc[z_calc!=0].max()

    # Calculate some statistics for the difference tile

    difference = np.reshape(z_calc,ndim*ndim)

    mean = np.mean(difference)
    var = stats.tvar(difference,limits=None)
    skew = stats.skew(difference,axis=None)
    kurt = stats.kurtosis(difference, axis=None)

    return(z_min,z_max,mean,var,skew,kurt)

在主要计算之后，计算各种统计数据。 c,d,e,f, ndim,half_tile的值都是单个整数值，变量z是一个大小为ndim x ndim的数组我之前试图对其进行向量化，但值不正确，虽然代码确实运行。

这是我的尝试：

def roughness(c,d,e,f,z,ndim,half_tile,dx):

    z_calc = np.zeros((ndim,ndim), dtype=float)
    x = np.zeros((ndim,ndim), dtype=float)
    y = np.zeros((ndim,ndim), dtype=float)

    x,y = np.mgrid[1:ndim+1,1:ndim+1]
    x = (x-half_tile)*dx
    y = (y-half_tile)*dx

    z_calc = c*x*y + d*x + e*y + f - z

    z_min=z_calc[z_calc!=0].min()
    z_max=z_calc[z_calc!=0].max()

    # Calculate some statistics for the difference tile

    difference = np.reshape(z_calc,ndim*ndim)

    mean = np.mean(difference)
    var = stats.tvar(difference,limits=None)
    skew = stats.skew(difference,axis=None)
    kurt = stats.kurtosis(difference, axis=None)

    return(z_min,z_max,mean,var,skew,kurt)

除了获得正确的值之外，我真的想知道我是否正确地进行了嵌套for循环的矢量化，我假设我没有这样做。