Numba和随机数从泊松分布

Question

我发现我的模拟瓶颈之一是从泊松分布产生的随机数。我原来的代码是这样的

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 

一些问题：为什么是这样的错误发生，如何解决它。

有更好的优化吗？

Answer 1

Numba不支持数组作为lam参数np.random.poisson，所以你要自己做循环：

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一次，并不多（并且与创建50万个随机数相比，完全可以忽略不计）。如果你想加快循环，你不能用纯NumPy的做，但你不应该指望numba（或其他东西）可以以相当于NumPy的功能提供了重要的加速

Numba是极快的。如果很容易让它们变得更快 - 那么NumPy的开发者也会更快。 :)