はじめに
ある日のこと、業務でめちゃくちゃ処理が遅いなぁと思いコードをのぞいてみると2重のfor文を発見
何とかできないかと探したときに出会ったのが私とapply_along_axisとの出会いです！
apply_along_axisで何ができるかといいますと、、、ずばり「関数をnumpy配列の列or行ごとに適用できる」です！
ということで今回はnumpyのapply_along_axisを紹介します！
環境
- Google Colaboratory
- Python 3.7.13
- numpy 2.21.6
問題設定
今回は公式ドキュメントを参考にかなりシンプルな処理で試してみます。
まず乱数を使って100000×3のサイズのnumpy配列を用意しました。
import numpy as np
arr = np.random.randint(1, 10, (10**5, 3))
print(arr[:5])
>>
array([[6, 2, 9],
[2, 2, 7],
[8, 1, 8],
[9, 3, 8],
[3, 8, 9]])
出力は100000行の内先頭5行を表示しています。
そして適用させたい関数が以下です。
def func(a):
return (a[0] + a[-1]) * 0.5
この関数は1次元配列の先頭の値と末尾の値を足して0.5を掛けるシンプルな関数となっています。
今回の目標は「関数funcを配列arrの行ごとに適用させること」です！
※ちなみに
今回の例でいうと関数など使わずに以下コードで済んでしまいます。
result = (arr[:, 0] + arr[:, -1]) * 0.5
result[:5]
>> array([7.5, 4.5, 8. , 8.5, 6. ])
より複雑な関数を適用させたいとき、apply_along_axisの真のありがたみに気づくでしょう
for文を使った方法
for文を使って1行づつ処理をした場合が以下コードです。
result = np.array([])
for a in arr:
calc = func(a)
result = np.append(result, calc)
result[:5]
>> array([7.5, 4.5, 8. , 8.5, 6. ])
実行してみると7秒ほどかかりました。
apply_along_axisを使った方法
ではapply_along_axisを使ってみます。
result = np.apply_along_axis(func, 1, arr)
result[:5]
>> array([7.5, 4.5, 8. , 8.5, 6. ])
こちらは1秒もかからずに計算が終わりました。
引数は適用させたい関数、軸、関数を適用する配列の順です。
発展
funcの0.5だった部分を引数bで指定できるとします。
def func(a, b):
return (a[0] + a[-1]) * b
この場合はfunctoolのpartialを使うとよいです。
from functools import partial
_f = partial(func, b=0.5)
result = np.apply_along_axis(_f, 1, arr)
result[:5]
>> array([7.5, 4.5, 8. , 8.5, 6. ])
partialでfuncの引数bを予め固定することで同様にapply_along_axisを使うことができます。
最後に
読んでいただきありがとうございます。
今回はnumpyのapply_along_axisについてかなりざっくりとした説明となりましたが、本記事が誰かのためになれば幸いです
参考