if文を含むPythonコードの高速化をNumbaとCythonでやってみた

#はじめに
PythonコードをNumbaで高速化したときのメモ
で、Numbaがfor文を使ったテクニカル指標関数の高速化に効果があることはわかったのですが、ほかにif文を多用する指標もあります。

そのひとつにパラボリックSARというのがあります。これは特に珍しい指標ではなく、割とポピュラーなものです。ただ、上昇モードと下降モードが切り替わったり、ステップ幅が変化したりするので、どうしてもfor文だけでは記述できません。MetaTraderのテクニカル指標をPythonに移植する際に一番最後になったものです。

今回はこれを高速化したときのメモです。

#パラボリックSARのPythonコード

import numpy as np
import pandas as pd
dataM1 = pd.read_csv('DAT_ASCII_EURUSD_M1_2015.csv', sep=';',
                     names=('Time','Open','High','Low','Close', ''),
                     index_col='Time', parse_dates=True)

def iSAR(df, step, maximum):
    last_period = 0
    dir_long = True
    ACC = step
    SAR = df['Close'].copy()
    for i in range(1,len(df)):
        last_period += 1    
        if dir_long == True:
            Ep1 = df['High'][i-last_period:i].max()
            SAR[i] = SAR[i-1]+ACC*(Ep1-SAR[i-1])
            Ep0 = max([Ep1, df['High'][i]])
            if Ep0 > Ep1 and ACC+step <= maximum: ACC+=step
            if SAR[i] > df['Low'][i]:
                dir_long = False
                SAR[i] = Ep0
                last_period = 0
                ACC = step
        else:
            Ep1 = df['Low'][i-last_period:i].min()
            SAR[i] = SAR[i-1]+ACC*(Ep1-SAR[i-1])
            Ep0 = min([Ep1, df['Low'][i]])
            if Ep0 < Ep1 and ACC+step <= maximum: ACC+=step
            if SAR[i] < df['High'][i]:
                dir_long = True
                SAR[i] = Ep0
                last_period = 0
                ACC = step
    return SAR

%timeit y = iSAR(dataM1, 0.02, 0.2)

for文は１重ですが、それなりに時間がかかります。

1 loop, best of 3: 1min 19s per loop

#Numbaで高速化
まず、Numbaで高速化してみます。pandasの配列をnumpyの配列に変えて@jitを追加するだけです。

from numba import jit
@jit
def iSARjit(df, step, maximum):
    last_period = 0
    dir_long = True
    ACC = step
    SAR = df['Close'].values.copy()
    High = df['High'].values
    Low = df['Low'].values
    for i in range(1,len(SAR)):
        last_period += 1    
        if dir_long == True:
            Ep1 = High[i-last_period:i].max()
            SAR[i] = SAR[i-1]+ACC*(Ep1-SAR[i-1])
            Ep0 = max([Ep1, High[i]])
            if Ep0 > Ep1 and ACC+step <= maximum: ACC+=step
            if SAR[i] > Low[i]:
                dir_long = False
                SAR[i] = Ep0
                last_period = 0
                ACC = step
        else:
            Ep1 = Low[i-last_period:i].min()
            SAR[i] = SAR[i-1]+ACC*(Ep1-SAR[i-1])
            Ep0 = min([Ep1, Low[i]])
            if Ep0 < Ep1 and ACC+step <= maximum: ACC+=step
            if SAR[i] < High[i]:
                dir_long = True
                SAR[i] = Ep0
                last_period = 0
                ACC = step
    return SAR

%timeit y = iSARjit(dataM1, 0.02, 0.2)

1 loop, best of 3: 1.43 s per loop

55倍ほど速くなりました。コードの修正はほとんどないので、まあまあの結果です。

#Cythonで高速化
次にCythonで高速化してみます。Cythonは設定が面倒かと思っていましたが、Jupyter notebookだと、結構簡単に導入することができました。ただ、外部のコンパイラを利用するので、Visual C++をインストールする必要があります。Anacondaをビルドしたバージョンに合わせる必要があったので、今回は以下のコンパイラをインストールしました。

Visual Studio Community 2015

最初はコードは変えずにCythonの設定だけ行った場合です。

%load_ext Cython

%%cython
cimport numpy
cimport cython
def iSAR_c0(df, step, maximum):
    last_period = 0
    dir_long = True
    ACC = step
    SAR = df['Close'].values.copy()
    High = df['High'].values
    Low = df['Low'].values
    for i in range(1,len(SAR)):
        last_period += 1    
        if dir_long == True:
            Ep1 = High[i-last_period:i].max()
            SAR[i] = SAR[i-1]+ACC*(Ep1-SAR[i-1])
            Ep0 = max([Ep1, High[i]])
            if Ep0 > Ep1 and ACC+step <= maximum: ACC+=step
            if SAR[i] > Low[i]:
                dir_long = False
                SAR[i] = Ep0
                last_period = 0
                ACC = step
        else:
            Ep1 = Low[i-last_period:i].min()
            SAR[i] = SAR[i-1]+ACC*(Ep1-SAR[i-1])
            Ep0 = min([Ep1, Low[i]])
            if Ep0 < Ep1 and ACC+step <= maximum: ACC+=step
            if SAR[i] < High[i]:
                dir_long = True
                SAR[i] = Ep0
                last_period = 0
                ACC = step
    return SAR

%timeit y = iSAR_c0(dataM1, 0.02, 0.2)

結果

1 loop, best of 3: 1.07 s per loop

同じコードのままでもCythonの方が少し速くなりました。

次にcdefで変数の型宣言を追加した場合です。

%%cython
cimport numpy
cimport cython
def iSARnew(df, double step, double maximum):
    cdef int last_period = 0
    dir_long = True
    cdef double ACC = step
    cdef numpy.ndarray[numpy.float64_t, ndim=1] SAR = df['Close'].values.copy()
    cdef numpy.ndarray[numpy.float64_t, ndim=1] High = df['High'].values
    cdef numpy.ndarray[numpy.float64_t, ndim=1] Low = df['Low'].values
    cdef double Ep0, Ep1
    cdef int i, N=len(SAR)
    for i in range(1,N):
        last_period += 1    
        if dir_long == True:
            Ep1 = max(High[i-last_period:i])
            SAR[i] = SAR[i-1]+ACC*(Ep1-SAR[i-1])
            Ep0 = max([Ep1, High[i]])
            if Ep0 > Ep1 and ACC+step <= maximum: ACC+=step
            if SAR[i] > Low[i]:
                dir_long = False
                SAR[i] = Ep0
                last_period = 0
                ACC = step
        else:
            Ep1 = min(Low[i-last_period:i])
            SAR[i] = SAR[i-1]+ACC*(Ep1-SAR[i-1])
            Ep0 = min([Ep1, Low[i]])
            if Ep0 < Ep1 and ACC+step <= maximum: ACC+=step
            if SAR[i] < High[i]:
                dir_long = True
                SAR[i] = Ep0
                last_period = 0
                ACC = step
    return SAR

%timeit y = iSARnew(dataM1, 0.02, 0.2)

結果は

1 loop, best of 3: 533 ms per loop

でした。さらに２倍くらい速くなりました。チューニングすればもっと速くなるかもしれませんが、コードの可読性が悪くなるおそれがあるので、ここまでにしておきます。

#まとめ
for文だけの場合、Numbaでも結構高速化の効果があるのですが、if文も入ってくると、効果が下がってきます。もう少し速くしたい場合には、多少コードの修正を伴いますが、Cython使うのもありかもしれません。