泊松分布中的 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 的速度非常快，但你不应该期望 numba（或其他任何东西）可以提供相当于同等速度的主要加速NumPy 函数。如果可以很容易地使它们更快 - NumPy 开发人员也会使其更快。 :)