用两个 numpy 向量中的元素对上的函数填充矩阵的最快方法？

Question

我有两个一维 numpy 向量va，vb它们用于通过将所有对组合传递给函数来填充矩阵。

na = len(va)
nb = len(vb)
D = np.zeros((na, nb))
for i in range(na):
    for j in range(nb):
        D[i, j] = foo(va[i], vb[j])

就目前情况而言，由于 va 和 vb 相对较大（4626 和 737），这段代码需要很长时间才能运行。然而，我希望这可以得到改进，因为使用cdistscipy 中的方法执行类似的过程具有非常好的性能。

D = cdist(va, vb, metric)

我显然知道 scipy 具有在 C 中而不是在 python 中运行这段代码的好处 - 但我希望有一些我不知道的 numpy 函数可以快速执行它。

Answer 1

文档中称为函数式编程例程的最不为人所知的 numpy 函数之一是np.frompyfunc. 这会从 Python 函数创建一个 numpy ufunc。不是其他一些紧密模拟 numpy ufunc 的对象，而是一个具有所有附加功能的适当的 ufunc。虽然该行为在许多方面与 非常相似np.vectorize，但它具有一些明显的优点，希望以下代码应该突出显示：

In [2]: def f(a, b):
   ...:     return a + b
   ...:

In [3]: f_vec = np.vectorize(f)

In [4]: f_ufunc = np.frompyfunc(f, 2, 1)  # 2 inputs, 1 output

In [5]: a = np.random.rand(1000)

In [6]: b = np.random.rand(2000)

In [7]: %timeit np.add.outer(a, b)  # a baseline for comparison
100 loops, best of 3: 9.89 ms per loop

In [8]: %timeit f_vec(a[:, None], b)  # 50x slower than np.add
1 loops, best of 3: 488 ms per loop

In [9]: %timeit f_ufunc(a[:, None], b)  # ~20% faster than np.vectorize...
1 loops, best of 3: 425 ms per loop

In [10]: %timeit f_ufunc.outer(a, b)  # ...and you get to use ufunc methods
1 loops, best of 3: 427 ms per loop

因此，虽然它仍然明显不如正确的矢量化实现，但它更快一点（循环是用 C 语言编写的，但仍然有 Python 函数调用开销）。