趣味で計算流砂水理

Computational Sediment Hydraulics for Fun

numba.jit解説1 : 単純な反復計算の高速化

このブロクの人気記事であるnumba.jit関連について、もう少し詳しく書いてみました。

computational-sediment-hyd.hatenablog.jp

私はnumba.jitのヘビーユーザーです。おかげで数日かかるような科学技術計算もpythonで書くようになりました。

numba.jitでは単純な処理の反復が高速化されます。そのため、幾何演算や科学技術計算で効果的です。

自分の勉強のためにも、numba.jitの特徴をまとめておきます。

環境

  • OS:windows10 pro, CPU Core i7-8550U
  • jupyter上で動作確認、時間計測は%%timeitを使用
  • numba 0.47

簡単な事例で動作確認:ポリゴンの面積の計算

ソースコード

ポリゴンの面積を計算するコードでnumba.jitを解説します。

面積を計算する式はこんな感じです。

 S = \frac{1}{2}  \left| \sum^n_{i=2}{(x_{i-1}y_{i} - x_{i}y_{i-1})} \right|

shapelyで半径1の円を簡素化して多角形を作成。64角形になりました。

from shapely.geometry import Point

p = Point(0.0, 0.0)
x = p.buffer(1.0)
s = x.simplify(0.001, preserve_topology=False)
polygon = np.array( [ p for p in s.exterior.coords])
len(polygon)
# 65

f:id:SedimentHydraulics:20200112134907p:plain

上のポリゴンの面積を求めるプログラムはこんな感じでしょうか。当然面積はほぼ3.14になります。

import numpy as np

def poly(p):
    s = np.array( [ p[i-1][0]*p[i][1] - p[i][0]*p[i-1][1] for i in range(1, len(p)) ] )
    return 0.5*np.abs(np.sum(s))

poly(np.array(polygon)) 

numba.jitを使う場合は、defの前に@jit(nopython=True)を付けて、

import numpy as np
from numba import jit

@jit(nopython=True)
def poly_numba(p):
    s = np.array( [ p[i-1][0]*p[i][1] - p[i][0]*p[i-1][1] for i in range(1, len(p)) ] )
    return 0.5*np.abs(np.sum(s))

です。ほとんど何も変わりません。

速度計測

さっそく上の2つのコードの速度比較です。

%%timeit
poly(np.array(polygon)) 
105 µs ± 1.04 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%%timeit
poly_numba(np.array(polygon)) 
1.1 µs ± 34.8 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

なんとこれだけで約100倍の速度差が出ます。

注意点:いつでも使えば良いわけではない。

numba.jitは関数の定義時に通常の場合と比較して若干時間がかかります。関数の定義を含めて時間計測を行うと、

%%timeit 

def poly(p):
    s = np.array( [ p[i-1][0]*p[i][1] - p[i][0]*p[i-1][1] for i in range(1, len(p)) ] )
    return 0.5*np.abs(np.sum(s))

poly(np.array(polygon)) 
105 µs ± 931 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%%timeit

@jit(nopython=True)
def poly_numba(p):
    s = np.array( [ p[i-1][0]*p[i][1] - p[i][0]*p[i-1][1] for i in range(1, len(p)) ] )
    return 0.5*np.abs(np.sum(s))

poly_numba(np.array(polygon)) 
151 ms ± 3.31 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

のとおり、通常の関数の方が早くなります。そのため、一回だけ関数を呼び出す場合は不利になります。 この例では10万回呼び出してようやくnumba.jitのほうが早くなります。

デメリット:pythonらしい書き方が使えない。

今回の例をpythonらしく書くと例えばzip関数を使って、

def poly(p):
    s = np.array( [ p0[0]*p1[1] - p1[0]*p0[1] for p1, p0 in zip(p[1:], p[:-1]) ] ) 
    return 0.5*np.abs(np.sum(s))

と書けますがnumba.jitではエラーがでます。

numbaのupdateで徐々に使える関数が増えてますが 他にもnumpyやscipyの関数も使えないものが多いです。

まとめ

  • numba.jitは大規模演算用
  • 大規模演算では必須
  • 次回はクラス編を書きます。