Numba 和泊松分布的随机数
Numba and random numbers from poisson distribution
我发现我的模拟中的瓶颈之一是从泊松分布生成随机数。我的原始代码是这样的
import numpy as np
#Generating some data. In the actual code this comes from the previous
#steps in the simulation. But this gives an example of the type of data
n = 5000000
pop_n = np.array([range(500000)])
pop_n[:] = np.random.poisson(lam=n*pop_n/np.sum(pop_n))
现在,我看到numba可以非常简单地提高速度。我定义了函数
from numba import jit
@jit()
def poisson(n, pop_n, np=np):
return np.random.poisson(lam=n*pop_n/np.sum(pop_n))
这个确实运行比原来的快。但是,我尝试更进一步 :) 当我写
@jit(nopython=True)
def poisson(n, pop_n, np=np):
return np.random.poisson(lam=n*pop_n/np.sum(pop_n))
我得到了
Failed at nopython (nopython frontend)
Invalid usage of Function(np.random.poisson) with parameters (array(float64, 1d, C))
Known signatures:
* (float64,) -> int64
* () -> int64
* parameterized
一些问题为什么会出现这个错误以及如何解决它。
有没有更好的优化?
Numba 不支持数组作为 np.random.poisson
的 lam
参数,因此您必须自己执行循环:
import numba as nb
import numpy as np
@nb.njit
def poisson(n, pop_n):
res = np.empty_like(pop_n)
pop_n_sum = np.sum(pop_n)
for idx, item in enumerate(range(pop_n.shape[0])):
res[idx] = np.random.poisson(n*pop_n[idx] / pop_n_sum)
return res
n = 5000000
pop_n = np.array(list(range(1, 500000)), dtype=float)
poisson(n, pop_n)
但根据我的时间安排,这与使用纯 NumPy 一样快:
%timeit poisson(n, pop_n)
# 203 ms ± 1.79 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit np.random.poisson(lam=n*pop_n/np.sum(pop_n))
# 203 ms ± 3.97 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
那是因为尽管 Numba 支持 np.random.poisson
和 np.sum
等功能,但这些支持只是为了方便,实际上并没有加速代码(很多)。它可能可以在一定程度上避免函数调用开销,但考虑到它只会在纯 Python 中调用一次 np.random.poisson
,这并不多(与创建半百万个随机数相比完全可以忽略不计)。
如果你想加速你不能用纯 NumPy 做的循环,Numba 是非常快的,但你不应该期望 numba(或其他任何东西)可以为等效的 NumPy 函数提供主要加速。如果可以很容易地使它们更快——NumPy 开发人员也会让它更快。 :)
我发现我的模拟中的瓶颈之一是从泊松分布生成随机数。我的原始代码是这样的
import numpy as np
#Generating some data. In the actual code this comes from the previous
#steps in the simulation. But this gives an example of the type of data
n = 5000000
pop_n = np.array([range(500000)])
pop_n[:] = np.random.poisson(lam=n*pop_n/np.sum(pop_n))
现在,我看到numba可以非常简单地提高速度。我定义了函数
from numba import jit
@jit()
def poisson(n, pop_n, np=np):
return np.random.poisson(lam=n*pop_n/np.sum(pop_n))
这个确实运行比原来的快。但是,我尝试更进一步 :) 当我写
@jit(nopython=True)
def poisson(n, pop_n, np=np):
return np.random.poisson(lam=n*pop_n/np.sum(pop_n))
我得到了
Failed at nopython (nopython frontend)
Invalid usage of Function(np.random.poisson) with parameters (array(float64, 1d, C))
Known signatures:
* (float64,) -> int64
* () -> int64
* parameterized
一些问题为什么会出现这个错误以及如何解决它。
有没有更好的优化?
Numba 不支持数组作为 np.random.poisson
的 lam
参数,因此您必须自己执行循环:
import numba as nb
import numpy as np
@nb.njit
def poisson(n, pop_n):
res = np.empty_like(pop_n)
pop_n_sum = np.sum(pop_n)
for idx, item in enumerate(range(pop_n.shape[0])):
res[idx] = np.random.poisson(n*pop_n[idx] / pop_n_sum)
return res
n = 5000000
pop_n = np.array(list(range(1, 500000)), dtype=float)
poisson(n, pop_n)
但根据我的时间安排,这与使用纯 NumPy 一样快:
%timeit poisson(n, pop_n)
# 203 ms ± 1.79 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit np.random.poisson(lam=n*pop_n/np.sum(pop_n))
# 203 ms ± 3.97 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
那是因为尽管 Numba 支持 np.random.poisson
和 np.sum
等功能,但这些支持只是为了方便,实际上并没有加速代码(很多)。它可能可以在一定程度上避免函数调用开销,但考虑到它只会在纯 Python 中调用一次 np.random.poisson
,这并不多(与创建半百万个随机数相比完全可以忽略不计)。
如果你想加速你不能用纯 NumPy 做的循环,Numba 是非常快的,但你不应该期望 numba(或其他任何东西)可以为等效的 NumPy 函数提供主要加速。如果可以很容易地使它们更快——NumPy 开发人员也会让它更快。 :)