[英]How to speed up this python function with numba?
我正在尝试加速这个 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 大师都可以帮我一把吗?
实际上,我认为在不深入了解数组属性的情况下,您无法做很多事情来加速 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 个循环,最好的 3 个:每个循环 246 毫秒
%timeit twoFreq_orig(z, source_z, num, den, matrix, e)
1 个循环,最好的 3 个:每个循环 344 毫秒
它比您的 numpy 解决方案快约 30%。 但是我认为通过巧妙地使用广播可以使 numpy 解决方案更快一点。 但是,尽管如此,我获得的大部分加速都来自省略签名:请注意,您可能使用 C 连续数组,但您已经给出了任意排序(因此 numba 可能会慢一点,具体取决于计算机体系结构)。 可能通过定义c16[::-1]
你会得到相同的速度,但通常只是让 numba 推断类型,它可能会尽可能快。 例外:您希望每个变量具有不同的精度输入(例如,您希望z
为complex128
和complex64
)
当您的 numpy 解决方案内存不足时,您将获得惊人的加速(因为您的 numpy 解决方案是矢量化的,它将需要更多的 RAM!)当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)
形状的数组,因此您执行的“幂”运算比必要的要多得多(即使我猜对于小数组,结果可能主要是缓存的)而且不是很贵)
没有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)
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.