如何加快泊松pmf函数的速度？

Question

我的用例是在小于10的所有点上评估Poisson pmf，我会用不同的lambdas多次调用这样的函数。兰巴达是不知道提前，所以我无法向量化兰巴达。

我从某个地方听说了一个秘密技巧，就是使用_pmf。这样做有什么坏处？但是，它还是有点慢，有没有办法改进它而不用从头开始用C重写pmf呢？

%timeit scipy.stats.poisson.pmf(np.arange(0,10),3.3)
%timeit scipy.stats.poisson._pmf(np.arange(0,10),3.3)
a = np.arange(0,10)
%timeit scipy.stats.poisson._pmf(a,3.3)

10000 loops, best of 3: 94.5 µs per loop
100000 loops, best of 3: 15.2 µs per loop
100000 loops, best of 3: 13.7 µs per loop

更新

好吧，我只是太懒了，不想用cython写东西。我原以为对所有离散分布都有一个更快的解决方案，可以对连续的x按顺序(迭代)进行评估。例如P(X=3) = P(X=2) * lambda / 3 if X ~ Pois(lambda)

相关：比用户提供的慢？

我现在对西西和Python不那么有信心了。图书馆的功能没有我所期望的那么先进。

Answer 1

大多数scipy.stats发行版都支持矢量化评估：

>>> poisson.pmf(1, [5, 6, 7, 8])
array([ 0.03368973,  0.01487251,  0.00638317,  0.0026837 ])

这可能足够快，也可能不够快，但您可以尝试将pmf调用从循环中删除。

pmf和_pmf之间的区别:真正的工作是在突出的函数(_pmf、_cdf等)中完成，而公共函数(pmf、cdf)则确保只有有效的参数才能到达_pmf (如果参数无效，_pmf的输出不能保证是有意义的，所以请自己承担风险)。

>>> poisson.pmf(1, -1)
nan
>>> poisson._pmf(1, -1)
/home/br/virtualenvs/scipy-dev/local/lib/python2.7/site-packages/scipy/stats/_discrete_distns.py:432: RuntimeWarning: invalid value encountered in log
  Pk = k*log(mu)-gamln(k+1) - mu
nan

详细信息：infrastructure.py#L2721

Answer 2

1. 尝试在cython中实现pmf。如果您的scipy是一个包的一部分，如Anaconda或Enthought，您可能已经安装了cython。http://cython.org/
2. 试着用pypy运行它。http://pypy.org/
3. 租用大型AWS服务器(或类似服务器)上的时间。

Answer 3

我发现与简单的python实现相比，scipy.stats.poisson类速度慢得可怜。

没有cython或载体什么的。

import math


def poisson_pmf(x, mu):
    return mu**x / math.factorial(x) * math.exp(-mu)


def poisson_cdf(k, mu):
    p_total = 0.0
    for x in range(k + 1):
        p_total += poisson_pmf(x, mu)
    return p_total

如果您检查源代码 of scipy.stats.poisson (甚至是下划线前缀版本)，很明显原因是什么！

上面的实现现在只比C(用gcc -O3 v9.3编译)慢10倍。The版本至少是慢10倍的。

#include <math.h>

unsigned long factorial(unsigned n) {
  unsigned long fact = 1;
  for (unsigned k = 2; k <= n; ++k)
    fact *= k;
  return fact;
}

double poisson_pmf(unsigned x, double mu) {
  return pow(mu, x) / factorial(x) * exp(-mu);
}

double poisson_cdf(unsigned k, double mu) {
  double p_total = 0.0;
  for (unsigned x = 0; x <= k; ++x)
    p_total += poisson_pmf(x, mu);
  return p_total;
}