在两个 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），这段代码需要很长时间才能运行。但是我希望这可以得到改进，因为使用来自 scipy 的 cdist 方法执行了类似的过程，性能非常好。

D = cdist(va, vb, metric)

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

Answer 1

对于文档称为functional programming routines 的最不为人知的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 函数调用开销）。

Answer 2

cdist 速度很快，因为它是用高度优化的 C 代码编写的（正如您已经指出的那样），并且它只支持一小部分预定义的 metrics。

由于您想将操作一般地应用于任何给定的foo 函数，您别无选择，只能调用该函数na-times-nb 次。这部分不太可能进一步优化。

剩下要优化的是循环和索引。一些尝试的建议：

1. 使用xrange而不是range（如果在python2.x中。在python3中，范围已经是类似生成器的）
2. 使用enumerate，而不是范围+显式索引
3. 使用 python 速度“魔法”，例如 cython 或 numba，加快循环过程。

如果您可以对foo 做出进一步的假设，则可能会进一步加快速度。

Answer 3

就像@shx2 所说，这完全取决于foo 是什么。如果你可以用numpy ufuncs来表达，那就用outer方法：

In [11]: N = 400

In [12]: B = np.empty((N, N))

In [13]: x = np.random.random(N)

In [14]: y = np.random.random(N)

In [15]: %%timeit
for i in range(N):
   for j in range(N):
     B[i, j] = x[i] - y[j]
   ....: 
10 loops, best of 3: 87.2 ms per loop

In [16]: %timeit A = np.subtract.outer(x, y)   # <--- np.subtract is a ufunc
1000 loops, best of 3: 294 µs per loop

否则，您可以将循环向下推到 cython 级别。继续上面的一个简单的例子：

In [45]: %%cython
cimport cython
@cython.boundscheck(False)
@cython.wraparound(False)
def foo(double[::1] x, double[::1] y, double[:, ::1] out):
    cdef int i, j
    for i in xrange(x.shape[0]):
        for j in xrange(y.shape[0]):
            out[i, j] = x[i] - y[j]
   ....: 

In [46]: foo(x, y, B)

In [47]: np.allclose(B, np.subtract.outer(x, y))
Out[47]: True

In [48]: %timeit foo(x, y, B)
10000 loops, best of 3: 149 µs per loop

故意将 cython 示例过于简单化：实际上您可能想要添加一些形状/步幅检查，在函数中分配内存等。