如何加快 Cython 代码以计算 dirichlet 的条件对数似然性？

Question

我有一个函数可以计算狄利克雷分布的条件（在第 k 个 alpha 上）对数似然。我用 Cython 编写并编译了它，但我的代码调用它大约 12M 次，这似乎是瓶颈，所以我希望加快速度。

cimport numpy as np
import numpy as np
import math
DTYPE = np.float64
ctypedef np.float64_t DTYPE_t

def logFullConAlphaK(np.ndarray p,np.ndarray alpha, np.int k):
    assert p.dtype == np.float64 and alpha.dtype == np.float64
    cdef double t1=sum(np.log(p))
    cdef DTYPE_t y=((alpha[k-1]-1)*t1)-np.log(alpha[k-1])+(p.shape[0]*
                   (math.lgamma(sum(alpha))- math.lgamma(alpha[k-1])))
return y

我将 Cython 编译成我在代码中使用的 .pyd 文件。有什么想法可以加快速度吗？

谢谢

Answer 1

1) 通过声明输入数组和p.shape[0] 的数据类型和维度：

def logFullConAlphaK(np.ndarray[DTYPE_t, ndim=1] p,
                     np.ndarray[DTYPE_t, ndim=1] alpha, int k):

    ...
    cdef int tmp
    tmp = p.shape[0]

2) 通过使用 C 函数而不是模块 math 中的 Python 函数：

cdef extern from "math.h":
    double log(double x) nogil

3) 使用 NumPy 的np.ndarray.sum() 方法

4) 使用 Cython 指令来避免一些开销

总共：

#cython: wraparound=False
#cython: boundscheck=False
#cython: cdivision=True
#cython: nonecheck=False

import math

cimport numpy as np
import numpy as np

cdef extern from "math.h":
    double log(double x) nogil    

DTYPE = np.float64
ctypedef np.float64_t DTYPE_t

def logFullConAlphaK(np.ndarray[DTYPE_t, ndim=1] p,
                     np.ndarray[DTYPE_t, ndim=1] alpha, int k):
    assert p.dtype == np.float64 and alpha.dtype == np.float64
    cdef double t1
    cdef int tmp

    t1 = np.log(p).sum()
    tmp = p.shape[0]

    cdef DTYPE_t y=((alpha[k-1]-1)*t1)-log(alpha[k-1])+(tmp*
                   (math.lgamma(alpha.sum()) - math.lgamma(alpha[k-1])))

    return y

OP 的原始解决方案、@cel 的解决方案和我的解决方案之间的一些性能比较：

In [2]: timeit solOP(a, b, 10)
1000 loops, best of 3: 273 µs per loop

In [3]: timeit solcel(a, b, 10)
10000 loops, best of 3: 30.5 µs per loop

In [4]: timeit solS(a, b, 10)
100000 loops, best of 3: 15.8 µs per loop

Answer 2

拿这个（可能完全不切实际的）样本数据：

n = 1000000
p = np.random.rand(n)
alpha = np.random.rand(n)
k = 12

我得到以下时间：

%timeit logFullConAlphaK(p, alpha, k) -> 1 loops, best of 3: 174 ms per loop

%timeit logFullConAlphaK_opt(p, alpha, k) -> 100 loops, best of 3: 13.3 ms per loop

此版本已经为您提供了一个数量级的速度。请注意，几乎所有的加速都来自使用 np.sum 而不是内置的 sum。所有其他更改只是为了更简洁的代码，它们不会影响速度。

cimport numpy as np
import numpy as np
import math

def logFullConAlphaK_opt(double[:] p, double[:] alpha, int k):
    cdef double t1=np.sum(np.log(p))
    cdef double y=((alpha[k-1]-1)*t1)-np.log(alpha[k-1])+(p.shape[0]*
                   (math.lgamma(np.sum(alpha))- math.lgamma(alpha[k-1])))
    return y