python - 加速 numpy 小函数

Question

这个小函数在应用程序中被多次调用，与 Matlab 中的相同代码实现相比运行速度非常慢。在那里它的工作速度可能快 10-100 倍，所以我想知道在 python/numpy 中可以在哪里改进它。

def ahamming(n,mid):
    data = np.zeros(n)
    wid1 = mid - 1
    wid2 = n - mid
    wid = max(wid1,wid2)

    for i in range(n):
        arg = (i+1) - mid;
        data[i] = np.cos((np.pi*arg)/wid);
    return 0.54 + 0.46*data

Answer 1

这是一个简单的向量化方法。使用 numpy 数组的好处是您可以避免循环并利用 numpy 的速度。

您可以按如下方式替换循环:

def ahamming(n,mid):
    data = np.zeros(n)
    wid1 = mid - 1
    wid2 = n - mid
    wid = max(wid1,wid2)

    i = np.arange(n)
    arg = (i+1) - mid
    data = 0.54 + 0.46*np.cos((np.pi*arg)/wid)

    return data

效率稍微高一些，但可能不太直观

i = np.arange(1, n+1)
arg = i - mid

编辑:速度结果出来了。n = 500 的循环版本需要 3.97 秒进行 10000 次计算。 numpy 版本工具 0.10 秒，速度提高了 40 倍。