Ускорение Cython не так велико, как ожидалось

Question

Я написал функцию Python, которая вычисляет попарно электромагнитные взаимодействия между большим числом (N ~ 10^3) частиц и сохраняет результаты в комплексе NxN128 ndarray. Он работает, но это самая медленная часть более крупной программы, занимающая около 40 секунд, когда N = 900 [исправлено]. Исходный код выглядит следующим образом:

import numpy as np 
def interaction(s,alpha,kprop): # s is an Nx3 real array 
           # alpha is complex 
           # kprop is float 

    ndipoles = s.shape[0] 

    Amat = np.zeros((ndipoles,3, ndipoles, 3), dtype=np.complex128) 
    I = np.array([[1,0,0],[0,1,0],[0,0,1]]) 
    im = complex(0,1) 

    k2 = kprop*kprop 

    for i in range(ndipoles): 
     xi = s[i,:] 
     for j in range(ndipoles): 
      if i != j: 
       xj = s[j,:] 
       dx = xi-xj 
       R = np.sqrt(dx.dot(dx)) 
       n = dx/R 
       kR = kprop*R 
       kR2 = kR*kR 
       A = ((1./kR2) - im/kR) 
       nxn = np.outer(n, n) 
       nxn = (3*A-1)*nxn + (1-A)*I 
       nxn *= -alpha*(k2*np.exp(im*kR))/R 
      else: 
       nxn = I 

      Amat[i,:,j,:] = nxn 

    return(Amat.reshape((3*ndipoles,3*ndipoles))) 

Я никогда раньше не использовал Cython, но это казалось как хорошее место, чтобы начать в моих усилиях, чтобы ускорить вещи, так что я в значительной степени вслепую адаптировал методы, которые я нашел в Интернете учебные пособия. Я получил ускорение (30 секунд против 40 секунд), но не так драматично, как я ожидал, поэтому мне интересно, что я делаю что-то неправильно или не вижу критического шага. Ниже моя лучшая попытка cythonizing вышеуказанную процедуру:

import numpy as np 
cimport numpy as np 

DTYPE = np.complex128 
ctypedef np.complex128_t DTYPE_t 

def interaction(np.ndarray s, DTYPE_t alpha, float kprop): 

    cdef float k2 = kprop*kprop 
    cdef int i,j 
    cdef np.ndarray xi, xj, dx, n, nxn 
    cdef float R, kR, kR2 
    cdef DTYPE_t A 

    cdef int ndipoles = s.shape[0] 
    cdef np.ndarray Amat = np.zeros((ndipoles,3, ndipoles, 3), dtype=DTYPE) 
    cdef np.ndarray I = np.array([[1,0,0],[0,1,0],[0,0,1]]) 
    cdef DTYPE_t im = complex(0,1) 

    for i in range(ndipoles): 
     xi = s[i,:] 
     for j in range(ndipoles): 
      if i != j: 
       xj = s[j,:] 
       dx = xi-xj 
       R = np.sqrt(dx.dot(dx)) 
       n = dx/R 
       kR = kprop*R 
       kR2 = kR*kR 
       A = ((1./kR2) - im/kR) 
       nxn = np.outer(n, n) 
       nxn = (3*A-1)*nxn + (1-A)*I 
       nxn *= -alpha*(k2*np.exp(im*kR))/R 
      else: 
       nxn = I 

      Amat[i,:,j,:] = nxn 

    return(Amat.reshape((3*ndipoles,3*ndipoles))) 

Answer 1

Реальная сила NumPy в выполнении операции по огромному количеству элементов в векторизованного образом, вместо того, чтобы использовать эту операцию на куски, разбросанных по петлями. В вашем случае вы используете два вложенных цикла и один условный оператор IF. Я бы предложил расширить размеры промежуточных массивов, которые приведут к включению в NumPy's powerful broadcasting capability, и, таким образом, одни и те же операции могут использоваться для всех элементов за один раз вместо небольших фрагментов данных в петлях.

Для расширения размеров можно использовать None/np.newaxis. Таким образом, векторизация реализация следовать такой предпосылке будет выглядеть следующим образом -

def vectorized_interaction(s,alpha,kprop): 

    im = complex(0,1) 
    I = np.array([[1,0,0],[0,1,0],[0,0,1]]) 
    k2 = kprop*kprop 

    # Vectorized calculations for dx, R, n, kR, A 
    sd = s[:,None] - s 
    Rv = np.sqrt((sd**2).sum(2)) 
    nv = sd/Rv[:,:,None] 
    kRv = Rv*kprop 
    Av = (1./(kRv*kRv)) - im/kRv 

    # Vectorized calculation for: "nxn = np.outer(n, n)" 
    nxnv = nv[:,:,:,None]*nv[:,:,None,:] 

    # Vectorized calculation for: "(3*A-1)*nxn + (1-A)*I" 
    P = (3*Av[:,:,None,None]-1)*nxnv + (1-Av[:,:,None,None])*I 

    # Vectorized calculation for: "-alpha*(k2*np.exp(im*kR))/R"  
    multv = -alpha*(k2*np.exp(im*kRv))/Rv 

    # Vectorized calculation for: "nxn *= -alpha*(k2*np.exp(im*kR))/R" 
    outv = P*multv[:,:,None,None] 


    # Simulate ELSE part of the conditional statement"if i != j:" 
    # with masked setting to I on the last two dimensions 
    outv[np.eye((N),dtype=bool)] = I 

    return outv.transpose(0,2,1,3).reshape(N*3,-1) 

испытаний Время воспроизведения и вывода проверки -

Case # 1:

In [703]: N = 10 
    ...: s = np.random.rand(N,3) + complex(0,1)*np.random.rand(N,3) 
    ...: alpha = 3j 
    ...: kprop = 5.4 
    ...: 

In [704]: out_org = interaction(s,alpha,kprop) 
    ...: out_vect = vectorized_interaction(s,alpha,kprop) 
    ...: print np.allclose(np.real(out_org),np.real(out_vect)) 
    ...: print np.allclose(np.imag(out_org),np.imag(out_vect)) 
    ...: 
True 
True 

In [705]: %timeit interaction(s,alpha,kprop) 
100 loops, best of 3: 7.6 ms per loop 

In [706]: %timeit vectorized_interaction(s,alpha,kprop) 
1000 loops, best of 3: 304 µs per loop 

Дело № 2:

In [707]: N = 100 
    ...: s = np.random.rand(N,3) + complex(0,1)*np.random.rand(N,3) 
    ...: alpha = 3j 
    ...: kprop = 5.4 
    ...: 

In [708]: out_org = interaction(s,alpha,kprop) 
    ...: out_vect = vectorized_interaction(s,alpha,kprop) 
    ...: print np.allclose(np.real(out_org),np.real(out_vect)) 
    ...: print np.allclose(np.imag(out_org),np.imag(out_vect)) 
    ...: 
True 
True 

In [709]: %timeit interaction(s,alpha,kprop) 
1 loops, best of 3: 826 ms per loop 

In [710]: %timeit vectorized_interaction(s,alpha,kprop) 
100 loops, best of 3: 14 ms per loop 

Дело № 3:

In [711]: N = 900 
    ...: s = np.random.rand(N,3) + complex(0,1)*np.random.rand(N,3) 
    ...: alpha = 3j 
    ...: kprop = 5.4 
    ...: 

In [712]: out_org = interaction(s,alpha,kprop) 
    ...: out_vect = vectorized_interaction(s,alpha,kprop) 
    ...: print np.allclose(np.real(out_org),np.real(out_vect)) 
    ...: print np.allclose(np.imag(out_org),np.imag(out_vect)) 
    ...: 
True 
True 

In [713]: %timeit interaction(s,alpha,kprop) 
1 loops, best of 3: 1min 7s per loop 

In [714]: %timeit vectorized_interaction(s,alpha,kprop) 
1 loops, best of 3: 1.59 s per loop