使用函数快速填充矩阵，基于两个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

cdist 运行得很快，因为它是用高度优化的 C 语言写的（正如你已经提到的），而且它只支持一小部分预定义的 metric（度量标准）。

由于你想要对任何给定的 foo 函数进行通用操作，你别无选择，只能调用这个函数 na 次乘以 nb 次。这个部分可能不会再有进一步的优化空间。

剩下可以优化的就是循环和索引。这里有一些建议可以尝试：

1. 在 Python 2.x 中使用 xrange 代替 range（在 Python 3 中，range 已经像生成器一样工作了）。
2. 使用 enumerate，而不是 range 加上手动索引。
3. 使用一些 Python 加速的“魔法”，比如 cython 或 numba，来加速循环过程。

如果你能对 foo 做出更多假设，可能会进一步加快速度。

Answer 2

就像@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的例子故意做得过于简单：实际上你可能还想添加一些形状/步幅的检查，或者在你的函数内部分配内存等等。

Answer 3

在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函数的开销还是存在的）。