Scipy.WeaveでPythonにC++埋め込み 5 Numpyの利用

前回の続き
前回はweave.inlineメソッドへのNumpy受け渡しの基本事項について書きました。今回は実際に計算をしていきます。
まず、前回にもお見せした、行と列のインデックスを足して3の倍数となる要素を0にする計算のpure Pythonバージョンをお見せします。

def numpy_mult3_0_pure(npArr):
    for r in xrange(npArr.shape[0]):
        for c in xrange(npArr.shape[1]):
            if (r + c) % 3 == 0:
                npArr[r, c] = 0

これを次のように呼び出す。

>>> side = 1000
>>> npArr = np.ones((side, side))
>>> numpy_mult3_0_pure(npArr)
>>> print npArr

結果は、以下。

[[ 0.  1.  1. ...,  1.  1.  0.]
 [ 1.  1.  0. ...,  1.  0.  1.]
 [ 1.  0.  1. ...,  0.  1.  1.]
 ..., 
 [ 1.  1.  0. ...,  1.  0.  1.]
 [ 1.  0.  1. ...,  0.  1.  1.]
 [ 0.  1.  1. ...,  1.  1.  0.]]

まったく前回のコピペです。
これのweaveバージョンは以下のようになります。

def numpy_mult3_0_weave(npArr):
    code = """
    for(int r = 0; r < NnpArr[0]; r++) { 
        for(int c = 0; c < NnpArr[1]; c++) {
            if((r + c) % 3 == 0) {
                NPARR2(r, c) = 0;
            }
        }
    }
    """
    weave.inline(code, ['npArr'])

前回やった、N+（変数名）, （変数名大文字）+（次元数）のアクセス用マクロ関数を使えば、上のように書けるのです。どのくらい早くなるのか時間を計測してみます。

side = 1000
npArr = np.ones((side, side))
from context_timer import Timer
with Timer("Pure"):
    numpy_mult3_0_pure(npArr)
with Timer("Weave"):
    numpy_mult3_0_weave(npArr)

このTimerについてはこちらの記事をご覧ください。出力は以下です。

[Pure] Elapsed: 0.272620201111
[Weave] Elapsed: 0.00578308105469

Weaveの方が50倍近く早くなっていることがわかります。こんな感じにPythonコードにC++のコードを書くことで高速化することができます。