numba-stats

我们提供了常用概率分布的JIT编译（使用numba）实现。

• 均匀分布
• (截断) 正态分布
• 对数正态分布
• 泊松分布
• 二项分布
• (截断) 指数分布
• 学生t分布
• Voigtian分布
• Crystal Ball分布
• 广义双边Crystal Ball分布
• Tsallis-Hagedorn，一个用于最小偏斜pT分布的模型
• Q-Gaussian分布
• Bernstein密度（未归一化到1，在扩展似然拟合中使用此密度）
• Cruijff密度（未归一化到1，在扩展似然拟合中使用此密度）
• CMS-Shape分布

速度提升非常显著，与scipy相比可高达100倍。基准测试包括在仓库中，并由pytest运行。

这些分布针对在最大似然拟合中使用进行了优化，其中您使用一组参数在多个点查询分布。

用法

每个分布都在一个子模块中实现。导入您需要的子模块，并调用模块中的函数。

from numba_stats import norm
import numpy as np

x = np.linspace(-10, 10)
mu = 2.0
sigma = 3.0

p = norm.pdf(x, mu, sigma)
c = norm.cdf(x, mu, sigma)

函数在变量x上进行了向量化，但不是在分布的形状参数上。理想情况下，为每个分布实现以下函数

• pdf：概率密度函数
• logpdf：概率密度函数的对数（对于某些分布可以更有效地和准确地计算）
• cdf：概率密度函数的积分
• ppf：cdf的逆
• rvs：生成随机变量

对于某些分布（例如 voigt）缺少 cdf 和 ppf，如果目前没有快速实现。只有当与计算 log(dist.pdf(...)) 相比更高效和准确时才实现 logpdf。只有对具有 ppf 的分布实现 rvs，其中 ppf 用于生成随机变量。目前 rvs 的实现尚未针对最高性能进行优化，但在实际中仍然很有用。

numba_stats 中的分布可以在其他 numba-JIT 编译的函数中使用。在 numba_stats 中的函数使用单个线程，但实现方式被编写成可以受益于自动并行化。要启用此功能，请从具有 parallel=True,fastmath=True 参数的 JIT 编译函数中调用它们。您应该始终将 parallel=True 与 fastmath=True 结合使用，因为后者可以增强自动并行化的收益。

from numba_stats import norm
import numba as nb
import numpy as np

@nb.njit(parallel=True, fastmath=True)
def norm_pdf(x, mu, sigma):
  return norm.pdf(x, mu, sigma)

# this must be an array of float
x = np.linspace(-10, 10)

# these must be floats
mu = 2.0
sigma = 3.0

# uses all your CPU cores
p = norm_pdf(x, mu, sigma)

请注意，这只有在 x 有足够的长度（约 1000 个元素或更多）时才更快。否则，并行化开销将使调用变慢，请参见下面的基准测试。

注意事项和解决方案

类型错误

当您在编译函数中使用 numba-stats 分布时，需要传递预期的数据类型。第一个参数必须是浮点数（32 位或 64 位）的 numpy 数组。后续参数必须是浮点数。如果您传递了错误的参数，您将得到类似以下 numba 错误（其中参数被传递为整数而不是浮点数）。

numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<function pdf at 0x7ff7186b7be0>) found for signature:

 >>> pdf(array(float64, 1d, C), int64, int64)

当您在编译函数外部调用 numba-stats PDF 时，您不会遇到这些错误，因为我在其中添加了一些包装器，这些包装器可以自动转换数据类型以方便使用。这就是为什么您可以调用 norm.pdf(1, 2, 3) 但 norm_pdf(1, 2, 3)（如上面所示实现）将失败的原因。

高维数组

为了保持实现简单，PDF 都是在 1D 数组参数上操作的。如果您有一个更高维的数组，您可以将其重塑，将其传递给我们的函数，然后将形状返回。这是一个便宜的运算。

x = ... # some higher dimensional array
# y = norm_pdf(x, 0.0, 1.0) this fails
y = norm_pdf(x.reshape(-1), 0.0, 1.0).reshape(x.shape)  # OK

文档

要获取文档，请使用 Python 解释器的 help()。

具有 scipy.stats 中等效函数的函数遵循 scipy 调用约定，但对于以 trunc... 开头的分布，由于 scipy 的行为非常不实用，因此遵循不同的约定。即使如此，请注意，scipy 的约定有时有些不寻常，特别是指数分布、对数正态分布和均匀分布。有关详细信息，请参阅 scipy 文档。

引用

如果您在科学工作中使用此包，请引用我们。您可以在 Zenodo 网站 生成您首选格式的引用。

基准测试

以下基准测试是在 Intel(R) Core(TM) i7-8569U CPU @ 2.80GHz 上针对 SciPy-1.10.1 生成的。右手图的虚线显示了预期从具有四个物理核心的 CPU 上的并行化获得的加速（4 倍）。

对于几乎所有的分布，我们都看到了与 scipy 相比的大幅加速。由于减少了调用开销，短数组的调用也受益于 numba_stats。由于我们调用的是用 FORTRAN 编写的相应的 scipy 实现，因此 voigt.pdf 和 t.ppf 的函数运行速度不会比 scipy 版本快。在这里，numba_stats 提供的优势是您可以从中调用这些函数的其他 numba-JIT 编译的函数，而这在 scipy 实现中是不可能的，并且 voigt.pdf 仍然受益于自动并行化。

由于内部实现无法轻松自动并行化，bernstein.density 不受益于自动并行化，反而会变得非常慢，因此应该避免这种情况。这是一个已知问题。

贡献

您可以帮助添加更多分布，欢迎提供补丁。实现概率分布很简单。您需要用简单的Python编写，让numba能够理解。在一定的包装后，可以使用scipy.special中的特殊函数，请参考子模块numba_stats._special.py了解具体实现方法。

numba-stats 和 numba-scipy

numba-scipy 是官方的快速numba-accelerated scipy函数包和仓库，我们是否在重复造轮子？

理想情况下，这个包的功能应该在numba-scipy中，我们希望最终能够实现这一点。在这个包中，我们不提供像numba-scipy那样的scipy函数和类的重载。这极大地简化了实现。numba-stats是一个临时解决方案，直到快速统计函数被包含在numba-scipy中。numba-stats目前不依赖于numba-scipy，只依赖于numba和scipy。