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



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

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函数可以快速执行。

文档称之为函数式编程例程的最不为人所知的numpy函数之一是np.frompyfunc。这将从Python函数创建一个numpy ufunc。不是其他一些物体能很好地模拟一个愚蠢的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函数调用开销)。

cdist速度很快,因为它是用高度优化的C代码编写的(正如您已经指出的),它只支持一小部分预定义的metric s。

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

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

  1. 使用xrange而不是range(如果在python2.x中,在python3中,range已经是类似生成器的)
  2. 使用enumerate,而不是范围+显式索引
  3. 使用python速度的"魔术",如cythonnumba,来加快循环过程

如果您可以对foo做进一步的假设,那么可能会进一步加快它的速度。

正如@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的例子故意过于简单:实际上,你可能想添加一些形状/步幅检查,在函数内分配内存等。

最新更新