分散をnumpyより高速に計算(numba, Python)

ループ中に分散を取る機会があったので調査。

以下より numpy はオーバーヘッドが大きい疑いがある。

numpy.max はヘビーループでは遅い(Python)

numpy では十分なパフォーマンスが得られなかったため、
ググって numba というのを試すことにした。

これには Just-In-Time Compiler が含まれており、
@JIT デコレータの付いた関数を実行時にコンパイルする。
関数の複数回呼び出し時にパフォーマンスの向上が期待できる。

以下は長さ 10 の配列に対する母分散算出を 100万回 実行した結果である。
Y 軸は処理時間(秒)で、対数となっている。

以下は長さ 100万の配列に対する母分散算出を 1回実行する結果である。

両結果から、numba の JIT コンパイルが、これらのケースに対しては、十分な性能(速度向上)を提供することが分かる。

実行したコードは以下。

import time
import numpy
import statistics
from numba import jit

@jit
def varp(arr):
    sm = 0
    for a in arr: sm += a
    ln = len(arr)
    av = sm / ln
    dv = 0
    for a in arr:
        dv += (a - av)**2
    return dv / ln

def measure(arr, times, fn):
    t = time.time()
    for i in range(times): fn(arr)
    print(time.time() - t)

d = numpy.random.rand(10)
MIL = 1000000
measure(d, MIL, varp)
measure(d, MIL, numpy.var)
measure(d, MIL, statistics.pvariance)
e = numpy.random.rand(MIL)
measure(e, 1, varp)
measure(e, 1, numpy.var)
measure(e, 1, statistics.pvariance)

なお、不偏分散の場合は以下のように最後の割り算を -1 する。

from numba import jit
@jit
def var(arr):
    sm = 0
    for a in arr: sm += a
    ln = len(arr)
    av = sm / ln
    dv = 0
    for a in arr:
        dv += (a - av)**2
    return dv / (ln - 1)