Как переписать несколько numpy-кодов с помощью Cython?

Question

Я просто хочу ускорить мой численный алгоритм, написанный в numpy. Критическая часть состоит в том, чтобы вычислить функцию логарифмического правдоподобия (разность между двумя усеченными нормальными CDF). Моя функция очень медленная (31,9 мс за цикл), и мне нужно запустить ее за 2000 раз за каждую итерацию.

Я попытался использовать функцию «norm.cdf» scipy вместо «ecfc». Но это медленнее. Я также попробовал «@jit» из пакета Numba. Но он был медленнее исходного кода.

Возможно, мне нужно использовать Cython. Но я почти ничего не знаю о C. Я пытался научиться Cython с веб-страницы Cython for numpy users, но мне было очень сложно.

Может ли кто-нибудь помочь мне переписать код в 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 :(

Update 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 

Update 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, но другие уже используют скомпилированный код. Китон не трогает их.

Мы могли бы изучить erfc.

Но лучше всего назвать этот код более крупными массивами.