使用 numpy.apply_along_axis 时提高代码速度

Question

我有一个函数，它接受一个二元数组作为输入。现在我有大数据（shape = (360000,2)），并想使用 numpy.apply_along_axis 评估每个点的函数。此线程（numpy np.apply_along_axis function speed up?）中给出的答案之一说 numpy.apply_along_axis 不是为了速度。我的功能是矢量化的。如何在不使用 jit/cython 的情况下改进所有数据的演化时间。

我将包含我正在尝试执行的操作的示例代码

import numpy as np
import random
def sample(x):
    return np.sin(x[0])*np.cos(x[1])
data = np.random.normal(size=600*600*2)
data = data.reshape(600*600,2)
%timeit np.sum(np.apply_along_axis(sample, 1,data)) #using the apply_along_axis

def loop_way():  # using loop
    result = []
    for i in data:
        result += [sample(i)]
    return np.sum(result)
%timeit loop_way()


output when using np. apply_along_axis: 1 loop, best of 3: 4.06 s per loop

output for loop_way function: 1 loop, best of 3: 2.41 s per loop

Answer 1

np.sin 和 * 是矢量化操作，因此，您可以将它们应用于整个数组：

np.sin(data[:, 0]) * np.cos(data[:, 1])

data[:, 0] 是第一列，data[:, 1] 是第二列。

请注意，这应该很快:)

---

这是一个测试每种方法速度的笔记本：notebook。

平均运行时间：

• 方法一（使用numpy.apply_along_axis）：2.08s
• 方法2（循环应用函数到行）：1.14s
• 方法3（这个答案）：17.3ms