python - 如何使用 Cython 重写一些 numpy 代码？

Question

我只是想加速我用 numpy 编写的数值算法。一个关键部分是计算对数似然函数(两个截断的正态 CDF 之间的差异)。我的函数非常慢(每个循环 31.9 毫秒)，每次迭代我都需要运行它 2000 次。

我尝试使用 scipy 的“norm.cdf”函数而不是“ecfc”。但速度较慢。我还尝试了 Numba 包中的“@jit”。但它也比原始代码慢。

我想也许我需要使用 Cython。但我对 C 几乎一无所知。我尝试从 Cython for numpy users 学习 Cython网页，但这对我来说真的太难了。

谁能帮我用 Cython 重写代码？或者建议我如何更快地编写它？

import numpy as np 
from scipy.special import erfc
# The bloody function for calculating the difference between two truncated normal CDFs
def my_loglikelihood2(x,b,c,z):
    log_likelihood=np.zeros(np.shape(z)[0])
    log_likelihood[x==1]=np.log(0.5*erfc(-(c[1]-np.dot(z[x==1,:],b)) / np.sqrt(2.)) - 0.5*erfc(-(c[0]-np.dot(z[x==1],b)) / np.sqrt(2.)))
    log_likelihood[x==2]=np.log(0.5*erfc(-(c[2]-np.dot(z[x==2,:],b)) / np.sqrt(2.)) - 0.5*erfc(-(c[1]-np.dot(z[x==2],b)) / np.sqrt(2.)))
    log_likelihood[x==3]=np.log(0.5*erfc(-(c[3]-np.dot(z[x==3,:],b)) / np.sqrt(2.)) - 0.5*erfc(-(c[2]-np.dot(z[x==3],b)) / np.sqrt(2.)))
    return log_likelihood

# generate random values
x=np.random.randint(low=1, high=4, size=50000)
b=np.random.normal(0,1,70)
c=np.array([-999,-1,1,999],dtype='f')  
z=np.random.multivariate_normal(np.zeros(70), np.eye(70), 50000) 

%timeit my_loglikelihood2(x,b,c,z)

# 10 loops, best of 3: 31.9 ms per loop   :(

更新 1 - 基于@jackvdp 的建议。它已经加速了 4.5 倍。但我仍在寻找更快的代码:

def up_cutoff(x,c):
    x[x==1]=c[1]
    x[x==2]=c[2]
    x[x==3]=c[3] 
    return x
def low_cutoff(x,c):
    x[x==1]=c[0]
    x[x==2]=c[1]
    x[x==3]=c[2] 
def my_loglikelihood2(x,b,low_c,up_c,z):
    up_c=up_cutoff(x,c)
    low_c=low_cutoff(x,c)
    return np.log(0.5*erfc(-(up_c-np.dot(z,b)) / np.sqrt(2.)) - 0.5*erfc(-(low_c-np.dot(z,b)) / np.sqrt(2.))) 

%timeit my_loglikelihood2(x,b,low_c,up_c,z)
100 loops, best of 3: 6.58 ms per loop

更新 2 - 基于@DSM 的建议。将 np.dot(z,b) 替换为 zdotb = z.dot(b) 。提高了 1.5 毫秒

def my_loglikelihood2(x,b,low_c,up_c,z):
    up_c=up_cutoff(x,c)
    low_c=low_cutoff(x,c)
    zdotb = z.dot(b)
    return np.log(0.5*erfc(-(up_c-zdotb) / np.sqrt(2.)) - 0.5*erfc(-(low_c-zdotb) / np.sqrt(2.))) 

%timeit my_loglikelihood2(x,b,low_c,up_c,z)
100 loops, best of 3: 5.02 ms per loop

Answer 1

如果您的代码由于 Python 中的循环而变慢，那么将其移植到 Cython 可以看到很大的改进。但是您的示例只调用了现有的 numpy/scipy 函数六次。

主要是调用 np.log、erfc、np.dot、np.sqrt。我不确定 erfc 但其他人已经使用编译代码。 Cython 不会触及那些。

我们可以检查 erfc。

但最好的办法是用更大的数组调用这段代码。