2017-12-15 8 views
2

私はCythonでPythonのパーティクルトラッキングコードのパフォーマンスを向上させるのに苦労しています。コードのCython最適化

明らか
from scipy.integrate import odeint 
import numpy as np 
from numpy import sqrt, pi, sin, cos 
from time import time as Time 
import multiprocessing as mp 
from functools import partial 

cLight = 299792458. 
Dim = 6 

class Integrator: 
    def __init__(self, ring): 
     self.ring = ring 

    def equations(self, X, s): 
     dXds = np.zeros(Dim) 

     E, B = self.ring.getEMField([X[0], X[2], s], X[4]) 

     h = 1 + X[0]/self.ring.ringRadius 
     p_s = np.sqrt(X[5]**2 - self.ring.particle.mass**2 - X[1]**2 - X[3]**2) 
     dtds = h*X[5]/p_s 
     gamma = X[5]/self.ring.particle.mass 
     beta = np.array([X[1], X[3], p_s])/X[5] 

     dXds[0] = dtds*beta[0] 
     dXds[2] = dtds*beta[1] 
     dXds[1] = p_s/self.ring.ringRadius + self.ring.particle.charge*(dtds*E[0] + dXds[2]*B[2] - h*B[1]) 
     dXds[3] = self.ring.particle.charge*(dtds*E[1] + h*B[0] - dXds[0]*B[2]) 
     dXds[4] = dtds 
     dXds[5] = self.ring.particle.charge*(dXds[0]*E[0] + dXds[2]*E[1] + h*E[2]) 
     return dXds 

    def odeSolve(self, X0, sRange): 
     sol = odeint(self.equations, X0, sRange) 
     return sol 

class Ring: 
    def __init__(self, particle): 
     self.particle = particle 
     self.ringRadius = 7.112 
     self.magicB0 = self.particle.magicMomentum/self.ringRadius 

    def getEMField(self, pos, time): 
     x, y, s = pos 
     theta = (s/self.ringRadius*180/pi) % 360 
     r = sqrt(x**2 + y**2) 
     arg = 0 if r == 0 else np.angle(complex(x/r, y/r)) 
     rn = r/0.045 

     k2 = 37*24e3 
     k10 = -4*24e3 

     E = np.zeros(3) 
     B = np.array([ 0, self.magicB0, 0 ]) 

     for i in range(4): 
      if ((21.9+90*i < theta < 34.9+90*i or 38.9+90*i < theta < 64.9+90*i) and (-0.05 < x < 0.05 and -0.05 < y < 0.05)): 
       E = np.array([ k2*x/0.045 + k10*rn**9*cos(9*arg), -k2*y/0.045 -k10*rn**9*sin(9*arg), 0]) 
       break 
     return E, B 

class Particle: 
    def __init__(self): 
     self.mass = 105.65837e6 
     self.charge = 1. 
     self.gm2 = 0.001165921 

     self.magicMomentum = self.mass/sqrt(self.gm2) 
     self.magicEnergy = sqrt(self.magicMomentum**2 + self.mass**2) 
     self.magicGamma = self.magicEnergy/self.mass 
     self.magicBeta = self.magicMomentum/(self.magicGamma*self.mass) 


def runSimulation(nParticles, tEnd): 
    particle = Particle() 
    ring = Ring(particle) 
    integrator = Integrator(ring) 

    Xs = np.array([ np.array([45e-3*(np.random.rand()-0.5)*2, 0, 0, 0, 0, particle.magicEnergy]) for i in range(nParticles) ]) 
    sRange = np.arange(0, tEnd, 1e-9)*particle.magicBeta*cLight 

    ode = partial(integrator.odeSolve, sRange=sRange) 

    t1 = Time() 

    pool = mp.Pool() 
    sol = np.array(pool.map(ode, Xs)) 

    t2 = Time() 
    print ("%.3f sec" %(t2-t1)) 

    return t2-t1 

、最も時間のかかるプロセスはodeSolve()とクラスインテグレータで方程式()のように定義ODEを、統合されています

は、ここに私の純粋なPythonのコードです。また、RingクラスのgetEMField()メソッドは、解決プロセス中にequations()メソッドと同じくらい呼び出されます。 私はCythonを使用して(〜20倍、少なくとも10倍)スピードアップ、かなりの量を取得しようとしましたが、私は唯一の次Cythonスクリプトによってスピードアップの〜1.5倍のレベルを得た:

import cython 
import numpy as np 
cimport numpy as np 
from libc.math cimport sqrt, pi, sin, cos 

from scipy.integrate import odeint 
from time import time as Time 
import multiprocessing as mp 
from functools import partial 

cdef double cLight = 299792458. 
cdef int Dim = 6 

@cython.boundscheck(False) 
cdef class Integrator: 
    cdef Ring ring 

    def __init__(self, ring): 
     self.ring = ring 

    cpdef np.ndarray[np.double_t, ndim=1, negative_indices=False, mode="c"] equations(self, 
        np.ndarray[np.double_t, ndim=1, negative_indices=False, mode="c"] X, 
        double s): 
     cdef np.ndarray[np.double_t, ndim=1, negative_indices=False, mode="c"] dXds = np.zeros(Dim) 
     cdef double h, p_s, dtds, gamma 
     cdef np.ndarray[np.double_t, ndim=1, negative_indices=False, mode="c"] beta, E, B 

     E, B = self.ring.getEMField([X[0], X[2], s], X[4]) 

     h = 1 + X[0]/self.ring.ringRadius 
     p_s = np.sqrt(X[5]*X[5] - self.ring.particle.mass*self.ring.particle.mass - X[1]*X[1] - X[3]*X[3]) 
     dtds = h*X[5]/p_s 
     gamma = X[5]/self.ring.particle.mass 
     beta = np.array([X[1], X[3], p_s])/X[5] 

     dXds[0] = dtds*beta[0] 
     dXds[2] = dtds*beta[1] 
     dXds[1] = p_s/self.ring.ringRadius + self.ring.particle.charge*(dtds*E[0] + dXds[2]*B[2] - h*B[1]) 
     dXds[3] = self.ring.particle.charge*(dtds*E[1] + h*B[0] - dXds[0]*B[2]) 
     dXds[4] = dtds 
     dXds[5] = self.ring.particle.charge*(dXds[0]*E[0] + dXds[2]*E[1] + h*E[2]) 
     return dXds 

    cpdef np.ndarray[np.double_t, ndim=1, negative_indices=False, mode="c"] odeSolve(self, 
       np.ndarray[np.double_t, ndim=1, negative_indices=False, mode="c"] X0, 
       np.ndarray[np.double_t, ndim=1, negative_indices=False, mode="c"] sRange): 
     sol = odeint(self.equations, X0, sRange) 
     return sol 

@cython.boundscheck(False) 
cdef class Ring: 
    cdef Particle particle 
    cdef double ringRadius 
    cdef double magicB0 

    def __init__(self, particle): 
     self.particle = particle 
     self.ringRadius = 7.112 
     self.magicB0 = self.particle.magicMomentum/self.ringRadius 

    cpdef tuple getEMField(self, 
        list pos, 
        double time): 
     cdef double x, y, s 
     cdef double theta, r, rn, arg, k2, k10 
     cdef np.ndarray[np.double_t, ndim=1, negative_indices=False, mode="c"] E, B 

     x, y, s = pos 
     theta = (s/self.ringRadius*180/pi) % 360 
     r = sqrt(x*x + y*y) 
     arg = 0 if r == 0 else np.angle(complex(x/r, y/r)) 
     rn = r/0.045 

     k2 = 37*24e3 
     k10 = -4*24e3 

     E = np.zeros(3) 
     B = np.array([ 0, self.magicB0, 0 ]) 

     for i in range(4): 
      if ((21.9+90*i < theta < 34.9+90*i or 38.9+90*i < theta < 64.9+90*i) and (-0.05 < x < 0.05 and -0.05 < y < 0.05)): 
       E = np.array([ k2*x/0.045 + k10*rn**9*cos(9*arg), -k2*y/0.045 -k10*rn**9*sin(9*arg), 0]) 
       #E = np.array([ k2*x/0.045, -k2*y/0.045, 0]) 
       break 
     return E, B 

cdef class Particle: 
    cdef double mass 
    cdef double charge 
    cdef double gm2 

    cdef double magicMomentum 
    cdef double magicEnergy 
    cdef double magicGamma 
    cdef double magicBeta 

    def __init__(self): 
     self.mass = 105.65837e6 
     self.charge = 1. 
     self.gm2 = 0.001165921 

     self.magicMomentum = self.mass/sqrt(self.gm2) 
     self.magicEnergy = sqrt(self.magicMomentum**2 + self.mass**2) 
     self.magicGamma = self.magicEnergy/self.mass 
     self.magicBeta = self.magicMomentum/(self.magicGamma*self.mass) 

def runSimulation(nParticles, tEnd): 
    particle = Particle() 
    ring = Ring(particle) 
    integrator = Integrator(ring) 

    #nParticles = 5 
    Xs = np.array([ np.array([45e-3*(np.random.rand()-0.5)*2, 0, 0, 0, 0, particle.magicEnergy]) for i in range(nParticles) ]) 
    sRange = np.arange(0, tEnd, 1e-9)*particle.magicBeta*cLight 

    ode = partial(integrator.odeSolve, sRange=sRange) 

    t1 = Time() 

    pool = mp.Pool() 
    sol = np.array(pool.map(ode, Xs)) 

    t2 = Time() 
    print ("%.3f sec" %(t2-t1)) 

    return t2-t1 

は私が得るために何をすべきCythonからの最大の効果? (私はCythonの代わりにNumbaを試しましたが、実際にはNumbaのパフォーマンスは非常に高かった(約20倍のスピードアップ)しかし、NumbaをPythonクラスのインスタンスで使うのは非常に難しく、Numbaの代わりにCythonを使用することにしました。参考のため、以下ではそのコンパイルにcython注釈

されています: enter image description here

enter image description here

enter image description here

enter image description here

+3

ボトルネックを見つけるためにコードをベンチマークしましたか?速いリードスルーからは、CythonやNumbaはスピードアップを大幅に向上させることができます。ほとんどの操作は既にベクトル化されています。まず、[ラインプロファイラ](https://jakevdp.github.io/PythonDataScienceHandbook/01.07-timing-and-profiling.html#Line-By-Line-Profiling-with-%lprun)を使用して開始しますスロースポットはどこにありますか? – jakevdp

+0

@jakevdpご意見ありがとうございます。私はラインプロファイラを使用するために調べましたが、まずCythonとPython3でそれを使用する方法を学ぶ必要があるようです...時間がかかります。それが助けになる場合は、私は注釈モードでCythonのコンパイルの結果を追加しました。 – hopeso

+1

彼は元の/非cythonコードでラインプロファイラを使用して、どの操作が遅いかを確認することをお勧めします。それらが基本的なnumpyプリミティブ/ベクトル化された部分であるなら、あなたはcythonが助けにならないことを知っています。 – sascha

答えて

1

私はプロファイリングかしていないので、これは非常に不完全な答えです何かをタイムリーにしたり、同じ答えを与えることをチェックしたりしていました。しかし、ここでCythonが生成するPythonコードの量を減らすいくつかの提案は以下のとおりです。

  • @cython.cdivision(True)コンパイルディレクティブを追加します。つまり、ZeroDivisionErrorは浮動小数点除算では生成されず、代わりにNaNの値が得られます。 (エラーを発生させたくない場合にのみこれを実行してください)。

  • 変更p_s = np.sqrt(...)~p_s = sqrt(...)。これにより、単一の値に対してのみ動作するnumpy呼び出しが削除されます。あなたは他の場所でこれをやったようですので、なぜこの行を逃したのかわかりません。代わりにnumpyの配列の

  • 使用の可能性は、固定サイズのC配列:

    cdef double beta[3] 
    # ... 
    beta[0] = X[1]/X[5] 
    beta[1] = X[3]/X[5] 
    beta[2] = p_s/X[5] 
    

    サイズはコンパイル時(とかなり小さい)で知られているときは、この操作を行うことができますし、返すようにしたくないときそれ。これにより、np.zerosへの呼び出しが回避され、型付きnumpy配列が割り当てられます。私はbetaがこれを行うことができる唯一の場所だと思います。

  • np.angle(complex(x/r, y/r))atan2libc.mathから使用すると、atan2(y/r, x/r)で置き換えることができます。あなたはまた、私は疑う

  • (Cythonは、ループ変数の型を拾う時に自動的にしばしば良いのですが、ここで失敗しているようだ)getEMFieldにすばやくforループを作るのに役立ちますr

  • cdef int iによる除算を失うことができますそれは、アレイ全体としてよりもE要素ごとに割り当てることが迅速です:

     E[0] = k2*x/0.045 + k10*rn**9*cos(9*arg) 
         E[1] = -k2*y/0.045 -k10*rn**9*sin(9*arg) 
    
  • 多くの価値がlistや01などの指定のタイプではありませんであり、コードを少し遅くすることがあります(タイプのチェックに時間がかかるため)。

  • 大きな変更は、ポインタではなく、それらnp.zeros割り当て使ってGetEMFieldEBを渡すことであろう。これにより、それらを静的C配列としてequationscdef double E[3])に割り当てることができます。欠点は、GetEMFieldcdefなので、もはやPythonから呼び出すことができなくなります(でも、好きなようにPythonの呼び出し可能なラッパー関数を作ることができます)。