使用 numpy 数组对函数进行矢量化

Question

我正在尝试加快我编写的一些代码的速度，但这样做会遇到很多麻烦。我知道能够删除 for 循环并使用 numpy 可以帮助做到这一点，所以这就是我一直在尝试但收效甚微的事情。

没有任何加速的工作函数是

def acf(x, y, z, cutoff=0):

    steps = x.shape[1]

    natoms = x.shape[0]

    z_x = np.zeros((steps,natoms))

    z_y, z_z = np.zeros_like(z_x), np.zeros_like(z_x)

    xmean = np.mean(x, axis=1)

    ymean = np.mean(y, axis=1)

    zmean = np.mean(z, axis=1)

    for k in range(steps-cutoff): # x.shape[1]

        xtemp, ytemp, ztemp = [], [], []

        for i in range(x.shape[0]): # natoms

            xtop, ytop, ztop = 0.0, 0.0, 0.0

            xbot, ybot, zbot = 0.0, 0.0, 0.0

            for j in range(steps-k): # x.shape[1]-k

                xtop += (x[i][j] - xmean[i]) * (x[i][j+k] - xmean[i])

                ytop += (y[i][j] - ymean[i]) * (y[i][j+k] - ymean[i])

                ztop += (z[i][j] - zmean[i]) * (z[i][j+k] - zmean[i])

                xbot += (x[i][j] - xmean[i])**2

                ybot += (y[i][j] - ymean[i])**2

                zbot += (z[i][j] - zmean[i])**2

            xtemp.append(xtop/xbot)

            ytemp.append(ytop/ybot)

            ztemp.append(ztop/zbot)

        z_x[k] = xtemp

        z_y[k] = ytemp

        z_z[k] = ztemp

    z_x = np.mean(np.array(z_x), axis=1)

    z_y = np.mean(np.array(z_y), axis=1)

    z_z = np.mean(np.array(z_z), axis=1)

    return z_x, z_y, z_z

此函数的输入 x、y 和 z 是相同维度的 numpy 数组。x（或 y 或 z）的一个例子是：

x = np.array([[1,2,3],[4,5,6]])

到目前为止，我能够做的是

def acf_quick(x, y, z, cutoff=0):

    steps = x.shape[1]

    natoms = x.shape[0]

    z_x = np.zeros((steps,natoms))

    z_y, z_z = np.zeros_like(z_x), np.zeros_like(z_x)

    x -= np.mean(x, axis=1, keepdims=True)

    y -= np.mean(y, axis=1, keepdims=True)

    z -= np.mean(z, axis=1, keepdims=True)

这将其速度提高了约 33%，但我相信有一种方法可以删除for i in range(natoms)使用类似x[:][j]. 到目前为止，我一直没有成功，任何帮助将不胜感激。

在有人问之前，我知道这是一个自相关函数，并且有几个内置于 numpy、scipy 等，但我需要自己编写。

Answer 1

这是循环的矢量化形式：def acf_quick_new(x, y, z, cutoff=0):    steps = x.shape[1]    natoms = x.shape[0]    lst_inputs = [x.copy(),y.copy(),z.copy()]    lst_outputs = []    for x_ in lst_inputs:        z_x_ = np.zeros((steps,natoms))        x_ -= np.mean(x_, axis=1, keepdims=True)        x_top = np.diag(np.dot(x_,x_.T))        x_bot = np.sum(x_**2, axis=1)        z_x_[0,:] = np.divide(x_top, x_bot)        for k in range(1,steps-cutoff): # x.shape[1]            x_top = np.diag(np.dot(x_[:,:-k],x_.T[k:,:]))            x_bot = np.sum(x_[:,:-k]**2, axis=1)            z_x_[k,:] = np.divide(x_top, x_bot)        z_x_ = np.mean(np.array(z_x_), axis=1)        lst_outputs.append(z_x_)        return lst_outputs请注意，在您的 _quick-function 中有一个小错误：您总是除以 xbot 而不是 xbot、ybot 和 zbot。此外，我的建议可以写得更好一些，但它应该可以解决您的问题并大大加快计算速度:)