Tốc độ Cython không lớn như mong đợi

Question

Tôi đã viết hàm Python tính các tương tác điện từ theo cặp giữa một số lượng lớn (N ~ 10^3) hạt và lưu trữ kết quả trong một NxN complex128 ndarray. Nó chạy, nhưng nó là phần chậm nhất của một chương trình lớn hơn, mất khoảng 40 giây khi N = 900 [sửa]. Mã gốc trông như thế này:

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

tôi chưa bao giờ sử dụng trước đó Cython, nhưng điều đó dường như là một nơi tốt để bắt đầu trong nỗ lực của tôi để tăng tốc, vì vậy tôi khá nhiều một cách mù quáng thích nghi các kỹ thuật tôi tìm thấy trong trực tuyến hướng dẫn. Tôi đã có một số tăng tốc (30 giây so với 40 giây), nhưng không gần như là kịch tính như tôi mong đợi, vì vậy tôi tự hỏi liệu tôi đang làm điều gì đó sai hoặc đang thiếu một bước quan trọng. Sau đây là nỗ lực hết sức mình tại cythonizing thói quen trên:

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

Sức mạnh thực sự của NumPy là trong việc thực hiện một hoạt động trên một số lượng lớn các yếu tố một cách vectorized thay vì sử dụng hoạt động mà trong khối lây lan qua vòng. Trong trường hợp của bạn, bạn đang sử dụng hai vòng lặp lồng nhau và một câu lệnh IF có điều kiện. Tôi sẽ đề xuất mở rộng kích thước của các mảng trung gian, sẽ mang lại cho NumPy's powerful broadcasting capability để đi vào chơi và do đó các hoạt động tương tự có thể được sử dụng trên tất cả các yếu tố trong một lần thay vì các khối dữ liệu nhỏ trong vòng lặp.

Để mở rộng thứ nguyên, None/np.newaxis có thể được sử dụng. Vì vậy, việc thực hiện vectorized theo một tiền đề như vậy sẽ giống như thế này -

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) 

kiểm tra Runtime và xác minh đầu ra -

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 

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

Trường hợp # 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