如何加速泊松 pmf 函数？答案

【问题标题】：How to speed up Poisson pmf function?如何加速泊松 pmf 函数？
【发布时间】：2014-03-31 18:32:53
【问题描述】：

我的用例是在所有小于 10 的点上评估 Poisson pmf，我会使用不同的 lambda 多次调用此类函数。无法提前知道 lambda，因此我无法对 lambda 进行矢量化。

我从某个地方听说了一个秘密技巧，那就是使用_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)

我现在对 Scipy 和 Python 不太信任了。库函数没有我想象的那么先进。

【问题讨论】：

标签： python scipy distribution

【解决方案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

【讨论】：

【解决方案2】：

尝试在 cython 中实现 pmf。如果你的 scipy 是像 Anaconda 或 Enthought 这样的包的一部分，你可能已经安装了 cython。 http://cython.org/
尝试使用 pypy 运行它。 http://pypy.org/
在大型 AWS 服务器（或类似服务器）上租用时间。

【讨论】：

【解决方案3】：

我发现scipy.stats.poisson 类与简单的 python 实现相比速度非常慢。

没有 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

如果您检查scipy.stats.poisson 的source code（甚至是下划线前缀的版本），那么原因就很清楚了！

上面的实现现在只慢 10 倍比 C 中的完全等效（使用 gcc -O3 v9.3 编译）。 scipy 版本至少 再慢 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;
}

【讨论】：