如何使用 numba 加速这个 python 函数？

Question

我正在尝试加快此 python 函数的速度：

def twoFreq_orig(z, source_z, num, den, matrix, e):
    Z1, Z2 = np.meshgrid(source_z, np.conj(z))
    Z1 **= num
    Z2 **= den - 1
    M = (e ** ((num + den - 2) / 2.0)) * Z1 * Z2
    return np.sum(matrix * M, 1)

其中z和source_z是np.ndarray（1d，dtype=np.complex128），num和den是np.ndarray (2d, dtype=np.float64), matrix 是一个 np.ndarray (2d, dtype=np.complex128) 而 e 是一个 np.float64.

我对Numba没有太多经验，但是在阅读了一些教程之后，我想出了这个实现：

@nb.jit(nb.f8[:](nb.c16[:], nb.c16[:], nb.f8[:, :], nb.f8[:, :], nb.c16[:, :], nb.f8))
def twoFreq(z, source_z, num, den, matrix, e):
    N1, N2 = len(z), len(source_z)
    out = np.zeros(N1)
    for r in xrange(N1):
        tmp = 0
        for c in xrange(N2):
            n, d = num[r, c], den[r, c] - 1
            z1 = source_z[c] ** n
            z2 = z[r] ** d
            tmp += matrix[r, c] * e ** ((n + d - 1) / 2.0) * z1 * z2
        out[r] = tmp
    return out

不幸的是，Numba 实现没有加速，反而比原来慢了几倍。我不知道如何正确使用 Numba。那里有任何 Numba 大师可以帮助我吗？

Answer 1

实际上，我认为如果不对数组的属性有更多的了解（是否有一些数学技巧可以更快地完成一些计算），您可以做很多事情来加速您的 numba 函数。

但我注意到一个错误：例如，您没有在 numba 版本中结合数组，我编辑了一些行以使其更加精简（其中一些可能只是味道）。我在适当的地方添加了评论：

@nb.njit
def twoFreq(z, source_z, num, den, matrix, e):
    #Replace z with conjugate of z (otherwise the result is wrong!)
    z = np.conj(z)
    # Size instead of len() don't know if it actually makes a difference but it's cleaner
    N1, N2 = z.size, source_z.size
    # Must be zeros_like otherwise you create a float array where you want a complex one
    out = np.zeros_like(z)
    # I'm using python 3 so you need to replace this by xrange later
    for r in range(N1):
        for c in range(N2):
            n, d = num[r, c], den[r, c] - 1
            z1 = source_z[c] ** n
            z2 = z[r] ** d
            # Multiply with 0.5 instead of dividing by 2
            # Work on the out array directly instead of a tmp variable
            out[r] += matrix[r, c] * e ** ((n + d - 1) * 0.5) * z1 * z2
    return out

def twoFreq_orig(z, source_z, num, den, matrix, e):
    Z1, Z2 = np.meshgrid(source_z, np.conj(z))
    Z1 **= num
    Z2 **= den - 1
    M = (e ** ((num + den - 2) / 2.0)) * Z1 * Z2
    return np.sum(matrix * M, 1)


numb = 1000
z = np.random.uniform(0,1,numb) + 1j*np.random.uniform(0,1,numb)
source_z = np.random.uniform(0,10,numb) + 1j*np.random.uniform(0,1,numb)
num = np.random.uniform(0,1,(numb,numb))
den = np.random.uniform(0,1,(numb,numb))
matrix = np.random.uniform(0,1,(numb,numb)) + 1j*np.random.uniform(0,1,(numb, numb))
e = 5.5

# This failed for your initial version:
np.testing.assert_array_almost_equal(twoFreq(z, source_z, num, den, matrix, e),
                                     twoFreq_orig(z, source_z, num, den, matrix, e))

我电脑上的运行时间是：

%timeit twoFreq(z, source_z, num, den, matrix, e)

1 loop, best of 3: 246 ms per loop

%timeit twoFreq_orig(z, source_z, num, den, matrix, e)

1 loop, best of 3: 344 ms per loop

它比您的 numpy 解决方案快大约 30%。但我认为通过巧妙地使用广播可以使 numpy 解决方案更快一些。但是，尽管如此，我获得的大部分加速都是由于省略了签名：请注意，您可能使用了 C 连续数组，但您给出了任意顺序（因此 numba 可能会慢一点，具体取决于计算机体系结构）。可能通过定义 c16[::-1] 你会得到相同的速度，但通常只是让 numba 推断类型，它可能会尽可能快。例外：您希望每个变量的精度输入不同（例如，您希望 z 为 complex128 和 complex64）

当您的 numpy 解决方案内存不足时，您将获得惊人的加速（因为您的 numpy 解决方案是矢量化的，它将需要更多的内存！）使用 numb = 5000，numba 版本比 numpy 快大约 3 倍一.

编辑：

聪明的广播是指

np.conj(z[:,None]**(den-1)) * source_z[None, :]**(num)

等于

z1, z2 = np.meshgrid(source_z, np.conj(z))
z1**(num) * z2**(den-1)

但是对于第一个变体，你只能对 numb 元素进行幂运算，而你有一个 (numb, numb) 形状的数组，所以你执行的 "power" 操作比必要的多得多（即使我想对于小型数组，结果可能主要是缓存的并且不是很昂贵）

没有 mgrid 的 numpy 版本（产生相同的结果）如下所示：

def twoFreq_orig2(z, source_z, num, den, matrix, e):
    z1z2 = source_z[None,:]**(num) * np.conj(z)[:, None]**(den-1)
    M = (e ** ((num + den - 2) / 2.0)) * z1z2
    return np.sum(matrix * M, 1)

如何使用 numba 加速这个 python 函数？

How to speed up this python function with numba?

python

python-2.7

numba

编辑：