「n要素ごとの平均を計算する」の爆速化

はじめに

今回の患者です。

avg = []
for i in range(1000):
    sum = 0
    for j in range(10):
        sum += data[i + j*1000]
    avg.append(sum / 10)

何をしているかの確認

初めに見たときは「え？これiとj逆じゃね？」と思いましたが前に聞いていた話も考えるとあってるようです。
dataには計測データが入っていて10回試行（1回の試行は1000データ）、10回試行の平均を取るという処理をしているようです。

つまり、

>>> data
array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16,
       17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29])

みたいにデータが入っていて（10試行1000データは多すぎるので3試行10データにしています）
試行ごとにデータを分けると、

>>> data.reshape(3, 10)
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14, 15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24, 25, 26, 27, 28, 29]])

となる。これを(0 + 10 + 20) / 3、(1 + 11 + 21) / 3みたいな感じに計算しているようです。

NumPy使うように書き換えよう

「ださーい、自力で平均計算していいのは小学生までだよねー」ということで書き換えましょう。n個飛ばしのデータってどう拾えばいいんだっけ？って少し考えましたがスライスの3要素目を使えば瞬殺でした。

>>> data[::10]
array([ 0, 10, 20])

後は平均、あれ？averageないと思ったらmeanでした。averageとmeanは微妙に違うみたいですね。

>>> data[::10].mean()
10.0

data[1], data[11], data[21]の平均を計算したい場合はスライスの1要素目を指定すればOKです。

>>> data[1::10].mean()
11.0

これで元プログラムの内側forはOK。外側もNumPyだけでやる方法が思いつかなかったので内包表記を使いました。

def calc_avg_improve(data, size, count):
    return [data[i::size].mean() for i in range(size)]

もっとNumPyっぽく

というわけでスライスと内包表記を使って改善完了、としたのですが帰り道に考えてたら以下のようにすればいけるなと思いました。

def calc_avg_improve2(data, size, count):
    return data.reshape(count, -1).T.mean(axis=1)

っていきなり示されてわかる人もあまりいなそうなので段階的に示すと、

>>> data.reshape(3, -1)
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14, 15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24, 25, 26, 27, 28, 29]])
>>> data.reshape(3, -1).T
array([[ 0, 10, 20],
       [ 1, 11, 21],
       [ 2, 12, 22],
       [ 3, 13, 23],
       [ 4, 14, 24],
       [ 5, 15, 25],
       [ 6, 16, 26],
       [ 7, 17, 27],
       [ 8, 18, 28],
       [ 9, 19, 29]])
>>> data.reshape(3, -1).T.mean(axis=1)
array([10., 11., 12., 13., 14., 15., 16., 17., 18., 19.])

Tを使うことで行列の縦横入れ替える、内側の配列それぞれについて平均を取る、となります。

計算結果の確認

念のため計算結果を確認しましょう。
なお、calc_avgは「はじめに」で紹介した元プログラムを関数化して固定値をsizeとcountを使うようにしたものです。

SIZE  = 10
COUNT = 3
d = np.arange(0, (SIZE*COUNT)*0.1, 0.1)

print(d)
print(calc_avg(d, SIZE, COUNT))
print(calc_avg_improve(d, SIZE, COUNT))
print(calc_avg_improve2(d, SIZE, COUNT))

結果。calc_avg_improveとcalc_avg_improve2は同じになるはず、と思いましたがcalc_avg_improveはPythonのリストでcalc_avg_improve2はNumPy配列なので、前者では非常に小さな値の切り捨てが行われてないのだと思います。

[0.  0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.  1.1 1.2 1.3 1.4 1.5 1.6 1.7
 1.8 1.9 2.  2.1 2.2 2.3 2.4 2.5 2.6 2.7 2.8 2.9]
[1.0, 1.1, 1.2000000000000002, 1.3, 1.4000000000000004, 1.5, 1.6000000000000003, 1.7000000000000002, 1.8, 1.9000000000000004]
[1.0, 1.1, 1.2000000000000002, 1.3, 1.4000000000000004, 1.5, 1.6000000000000003, 1.7000000000000002, 1.8, 1.9000000000000004]
[1.  1.1 1.2 1.3 1.4 1.5 1.6 1.7 1.8 1.9]

どれが速いのか

ではお待ちかねの実行時間計測。

import time

LOOP = 100
SIZE = 10000
COUNT = 10
d = np.arange(0, (SIZE*COUNT)*0.1, 0.1)

def perf(f):
    elapsed = []
    for l in range(LOOP):
        start = time.time()
        f(d, SIZE, COUNT)
        end = time.time()
        elapsed.append(end - start)
    print('{:<20} {}'.format(f.__name__, np.mean(elapsed)))

perf(calc_avg)
perf(calc_avg_improve)
perf(calc_avg_improve2)

結果

calc_avg             0.20858944416046143
calc_avg_improve     0.5223674201965331
calc_avg_improve2    0.0005125141143798828

calc_avg_improve改悪しとるやないかい！んーなんだろ、内包表記の部分が関数として実行されるからそのオーバーヘッド？

まとめ

というわけで今回はn要素ごとの平均を取るという処理を爆速化してみました。calc_improve2はややトリッキーなのでできればcalc_improveのように書けば速いぜ！と言いたかったのですが結果は残念なことになってしまいました。ともかくやはりNumPy内で計算が完結できると非常に爆速化できることがわかりますね。