Vectorization of a nested for-loop

问题

I have a function, roughness that is called quite often in a larger piece of code. I need some help with replacing this double for-loop with a simpler vectorized version. Here is the code below:

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)

After the main calculations, the various stats are calculated on them. The values, of c,d,e,f, ndim,half_tile are all single integer values, and the variable z is an array with size ndim x ndim I have tried to vectorize this before, but the values do not come out correctly, although the code does run.

Here is my attempt:

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)

Aside from getting the correct values, I would really like to know if I did the vectorization of the nested for loop correctly, which i'm assuming I didn't.

来源：https://stackoverflow.com/questions/34817500/vectorization-of-a-nested-for-loop

标签

python

performance

for-loop

vectorization