L'accelerazione di Cython non è grande come previsto

Question

Ho scritto una funzione Python che calcola interazioni elettromagnetiche a coppie tra un numero maggiore (N ~ 10^3) di particelle e memorizza i risultati in un complesso NxN128 ndarray. Funziona, ma è la parte più lenta di un programma più grande, impiegando circa 40 secondi quando N = 900 [corretto]. Il codice originale è simile al seguente:

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))) 

non avevo mai usato in precedenza Cython, ma che sembrava un buon punto di partenza nel mio sforzo di velocizzare le cose, così ho praticamente alla cieca adattato le tecniche che ho trovato in online esercitazioni. Ho avuto un po 'di accelerazione (30 secondi contro 40 secondi), ma non così drammatico come mi aspettavo, quindi mi chiedo se sto facendo qualcosa di sbagliato o mi manca un passaggio critico. Quanto segue è il mio miglior tentativo di cythonizing la routine di cui sopra:

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

Il vero potere di NumPy è in esecuzione di un'operazione in un enorme numero di elementi in modo vettorizzati invece di usare quella operazione in blocchi sparsi in loop. Nel tuo caso, stai usando due cicli annidati e una dichiarazione condizionale IF. Proporrei di estendere le dimensioni degli array intermedi, che porterebbero in gioco NumPy's powerful broadcasting capability e quindi le stesse operazioni potrebbero essere utilizzate su tutti gli elementi in una volta sola invece di piccoli pezzi di dati all'interno dei loop.

Per estendere le dimensioni, è possibile utilizzare None/np.newaxis. Così, l'attuazione vettorializzare di seguire una tale premessa sarebbe simile a questa -

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) 

test Runtime e la verifica di uscita -

Caso # 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 

Caso # 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 

Caso n. 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