Python/Cythonループを高速化します。

Question

私はできるだけ早く実行するためにPythonでループを取得しようとしています。だから私はNumPyとCythonに飛び込んだ。これは、k個のインデックスを超えるだけの合計です

def calculate_bsf_u_loop(uvel,dy,dz): 
    """ 
    Calculate barotropic stream function from zonal velocity 

    uvel (t,z,y,x) 
    dy (y,x) 
    dz (t,z,y,x) 

    bsf (t,y,x) 
    """ 

    nt = uvel.shape[0] 
    nz = uvel.shape[1] 
    ny = uvel.shape[2] 
    nx = uvel.shape[3] 

    bsf = np.zeros((nt,ny,nx)) 

    for jn in range(0,nt): 
     for jk in range(0,nz): 
     for jj in range(0,ny): 
      for ji in range(0,nx): 
       bsf[jn,jj,ji] = bsf[jn,jj,ji] + uvel[jn,jk,jj,ji] * dz[jn,jk,jj,ji] * dy[jj,ji] 

    return bsf 

： はここで、元のPythonコードです。配列のサイズは、nt = 12、nz = 75、ny = 559、nx = 1442、〜725万要素です。 それは68秒かかりました。さて、私はcythonで

import numpy as np 
cimport numpy as np 
cimport cython 

@cython.boundscheck(False) # turn off bounds-checking for entire function 
@cython.wraparound(False) # turn off negative index wrapping for entire function 

## Use cpdef instead of def 
## Define types for arrays 
cpdef calculate_bsf_u_loop(np.ndarray[np.float64_t, ndim=4] uvel, np.ndarray[np.float64_t, ndim=2] dy, np.ndarray[np.float64_t, ndim=4] dz): 
    """ 
    Calculate barotropic stream function from zonal velocity 

    uvel (t,z,y,x) 
    dy (y,x) 
    dz (t,z,y,x) 

    bsf (t,y,x) 
    """ 

    ## cdef the constants 
    cdef int nt = uvel.shape[0] 
    cdef int nz = uvel.shape[1] 
    cdef int ny = uvel.shape[2] 
    cdef int nx = uvel.shape[3] 

    ## cdef loop indices 
    cdef ji,jj,jk,jn 

    ## cdef. Note that the cdef is followed by cython type 
    ## but the np.zeros function as python (numpy) type 
    cdef np.ndarray[np.float64_t, ndim=3] bsf = np.zeros([nt,ny,nx], dtype=np.float64) 

    for jn in xrange(0,nt): 
     for jk in xrange(0,nz): 
     for jj in xrange(0,ny): 
      for ji in xrange(0,nx): 
       bsf[jn,jj,ji] += uvel[jn,jk,jj,ji] * dz[jn,jk,jj,ji] * dy[jj,ji] 

    return bsf 

とそれをして49分を要しました。 しかし、

for jn in range(0,nt): 
     for jk in range(0,nz): 
     bsf[jn,:,:] = bsf[jn,:,:] + uvel[jn,jk,:,:] * dz[jn,jk,:,:] * dy[:,:] 

のためのループを交換することのみ0.29秒かかります！残念ながら、私は完全なコードでこれを行うことはできません。

NumPyがCythonループよりもはるかに高速にスライスするのはなぜですか？ 私はNumPyが速かったと思っていました。それで、彼らは似たようなスピードではないでしょうか？

ご覧のとおり、私はcythonで境界チェックを無効にしました。また、「高速数学」を使用してコンパイルしました。しかし、これはわずかなスピードアップしか与えません。 NumPyスライシングと同様の速度のループを作成するにはどうしてですか？またはスライスより常に遅くループしていますか？

ご協力いただきありがとうございます。 /ヨアキム

Answer 1

すなわちコードは、numpy.einsumは非常に効率的な方法で行い 味方essenti 4D積アレイの第2の軸にelementwise-multiplication、次いでsum-reductionを行っていることを考えると、numpy.einsum'sの介入のために叫んされます。あなたのケースを解決するために、次の2つの方法でnumpy.einsumを使用することができます -

bsf = np.einsum('ijkl,ijkl,kl->ikl',uvel,dz,dy) 

bsf = np.einsum('ijkl,ijkl->ikl',uvel,dz)*dy 

ランタイムは&は、出力を確認してテストする -

In [100]: # Take a (1/5)th of original input shapes 
    ...: original_shape = [12,75, 559,1442] 
    ...: m,n,p,q = (np.array(original_shape)/5).astype(int) 
    ...: 
    ...: # Generate random arrays with given shapes 
    ...: uvel = np.random.rand(m,n,p,q) 
    ...: dy = np.random.rand(p,q) 
    ...: dz = np.random.rand(m,n,p,q) 
    ...: 

In [101]: bsf = calculate_bsf_u_loop(uvel,dy,dz) 

In [102]: print(np.allclose(bsf,np.einsum('ijkl,ijkl,kl->ikl',uvel,dz,dy))) 
True 

In [103]: print(np.allclose(bsf,np.einsum('ijkl,ijkl->ikl',uvel,dz)*dy)) 
True 

In [104]: %timeit calculate_bsf_u_loop(uvel,dy,dz) 
1 loops, best of 3: 2.16 s per loop 

In [105]: %timeit np.einsum('ijkl,ijkl,kl->ikl',uvel,dz,dy) 
100 loops, best of 3: 3.94 ms per loop 

In [106]: %timeit np.einsum('ijkl,ijkl->ikl',uvel,dz)*dy 
100 loops, best of 3: 3.96 ms per loo