Cythonのスピードアップが期待通りではありません

Question

大規模な数値（N〜10^3）の粒子間の対の電磁気的相互作用を計算し、その結果をNxN complex128 ndarrayに保存するPython関数を書いています。それは実行されますが、N = 900の場合には約40秒かかって大きなプログラムの中で最も遅いものです[修正]。元のコードは次のようになります。

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のの本当の力はベクトル化された方法で要素の膨大な数の全体に操作を実行する代わりに、ループにまたがっチャンクでその操作を使用しています。あなたのケースでは、2つのネストされたループと1つの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) 

ランタイムテストや出力検証 - - だから、そのような前提をフォローするベクトル化の実装は次のようになり

ケース＃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