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的開發者也會更快。 :)