对数伽马函数的快速算法

Fast algorithm for log gamma function

我正在尝试编写一个快速算法来计算 log gamma function。目前我的实现看起来很幼稚,只是迭代了 1000 万次来计算伽玛函数的对数(我也在使用 numba 来优化代码)。

import numpy as np
from numba import njit
EULER_MAS = 0.577215664901532 # euler mascheroni constant
HARMONC_10MIL = 16.695311365860007 # sum of 1/k from 1 to 10,000,000

@njit(fastmath=True)
def gammaln(z):
"""Compute log of gamma function for some real positive float z"""
    out = -EULER_MAS*z - np.log(z) + z*HARMONC_10MIL
    n = 10000000 # number of iters
    for k in range(1,n+1,4):
        # loop unrolling
        v1 = np.log(1 + z/k)
        v2 = np.log(1 + z/(k+1))
        v3 = np.log(1 + z/(k+2))
        v4 = np.log(1 + z/(k+3))
        out -= v1 + v2 + v3 + v4

    return out

我根据 scipy.special.gammaln 实现对我的代码计时,我的代码实际上慢了 100,000 倍。所以我正在做一些非常错误或非常幼稚的事情(可能两者都有)。尽管与 scipy.

相比,我的答案至少在小数点后 4 位以内是正确的

我试图阅读实现 scipy 的 gammaln 函数的 _ufunc 代码,但是我不理解编写 _gammaln 函数的 cython 代码。

有没有更快更优化的方法来计算对数伽马函数?我怎样才能理解 scipy 的实现,以便将其与我的实现结合起来?

您的函数的运行时间将随着迭代次数线性扩展(达到一些恒定的开销)。因此,减少迭代次数是加快算法速度的关键。虽然预先计算 HARMONIC_10MIL 是一个聪明的主意,但实际上在截断序列时会导致更差的准确性;结果证明只计算系列的一部分可以提供更高的准确性。

下面的代码是上面发布的代码的修改版本(尽管使用 cython 而不是 numba)。

from libc.math cimport log, log1p
cimport cython
cdef:
    float EULER_MAS = 0.577215664901532 # euler mascheroni constant

@cython.cdivision(True)
def gammaln(float z, int n=1000):
    """Compute log of gamma function for some real positive float z"""
    cdef:
        float out = -EULER_MAS*z - log(z)
        int k
        float t
    for k in range(1, n):
        t = z / k
        out += t - log1p(t)

    return out

如下图所示,经过100次逼近也能得到近似值。

在 100 次迭代时,其运行时间与 scipy.special.gammaln:

处于同一数量级
%timeit special.gammaln(5)
# 932 ns ± 19 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
%timeit gammaln(5, 100)
# 1.25 µs ± 20.3 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

剩下的问题当然是使用多少次迭代。函数 log1p(t) 可以展开为小 t 的泰勒级数(这与大 k 的极限有关)。特别是,

log1p(t) = t - t ** 2 / 2 + ...

这样,对于大 k,总和的参数变为

t - log1p(t) = t ** 2 / 2 + ...

因此,在 t 中,和的自变量在二阶之前为零,如果 t 足够小,则可以忽略不计。也就是说,迭代次数至少要和z一样大,最好至少大一个数量级。

但是,如果可能的话,我会坚持使用 scipy 经过充分测试的实现。

通过尝试 numba 的并行模式并主要使用矢量化函数,我设法将性能提高了大约 3 倍(遗憾的是,numba 无法理解 numpy.substract.reduce

from functools import reduce
import numpy as np
from numba import njit

@njit(fastmath=True, parallel=True)
def gammaln_vec(z):
    out = -EULER_MAS*z - np.log(z) + z*HARMONC_10MIL
    n = 10000000

    v = np.log(1 + z/np.arange(1, n+1))

    return out-reduce(lambda x1, x2: x1-x2, v, 0)

次数:

#Your function:
%timeit gammaln(1.5)
48.6 ms ± 1.23 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

#My function:
%timeit gammaln_vec(1.5)
15 ms ± 340 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

#scpiy's function
%timeit gammaln_sp(1.5)
1.07 µs ± 18.7 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

所以,使用 scipy 的功能,你会过得更好。没有 C 代码,我不知道如何进一步分解它

关于您之前的问题,我想将 scipy.special 函数包装到 Numba 的示例也很有用。

例子

只要只涉及简单的数据类型(int、double、double*,...),包装 Cython cdef 函数就非常容易和可移植。有关如何调用 scipy.special 函数 have a look at this 的文档。您实际需要包装函数的函数名称在 scipy.special.cython_special.__pyx_capi__ 中。可以用不同数据类型调用的函数名称被破坏,但确定正确的名称非常容易(只需查看数据类型)

#slightly modified version of https://github.com/numba/numba/issues/3086
from numba.extending import get_cython_function_address
from numba import vectorize, njit
import ctypes
import numpy as np

_PTR = ctypes.POINTER
_dble = ctypes.c_double
_ptr_dble = _PTR(_dble)

addr = get_cython_function_address("scipy.special.cython_special", "gammaln")
functype = ctypes.CFUNCTYPE(_dble, _dble)
gammaln_float64 = functype(addr)

@njit
def numba_gammaln(x):
  return gammaln_float64(x)

Numba 中的用法

#Numba example with loops
import numba as nb
import numpy as np
@nb.njit()
def Test_func(A):
  out=np.empty(A.shape[0])
  for i in range(A.shape[0]):
    out[i]=numba_gammaln(A[i])
  return out

计时

data=np.random.rand(1_000_000)
Test_func(A): 39.1ms
gammaln(A):   39.1ms

当然,您可以轻松并行化此函数并超越 scipy 中的单线程 gammaln 实现,并且您可以在任何 Numba 编译函数中高效地调用此函数。