※コメントいただいたので追記します。諸事情によりpython2からアップデートできない状況でして、コードもpython2になってます。大きく結論が変わるわけではないと思いますが、文法が異なるところは適宜ご自身の環境に合わせて読み替えてください。
numpy.arrayの要素をfor文で回すと実行速度がかなり落ちる。
これをもうちょっと早くするために試す対処として、
- list内包表記にする
- 条件が複雑ならnp.whereを使う
- np.frompyfuncを使う(ただし使用に注意が必要）
というのを学んだのでその備忘録。
list内包表記
これはもうどこにでも書いてあるけど、
import numpy as np
a = np.array(range(10))
a2 = []
for x in a:
a2.append(x*2)
とかやるくらいなら、
a2 = [x*2 for x in a]
にしなさいよ、という話。体感でもかなり早くなる。
条件が複雑ならnp.whereを使う
たとえば、いじりたいのは配列aだけど、条件は配列bの要素で決めたい、みたいな場合。
例として、配列bの要素が偶数なら配列aの要素を２倍、それ以外だったら配列aの要素を３倍にするケースを考えてみる。
ついC++的にインデックス使ってforで回したくなってしまうのをぐっとこらえて、np.whereを使う。
import numpy as np
a = np.array(range(10))
print "a = ", a
b = a + 100
print "b = ", b
# インデックス使っちゃうとこんな感じ
result1 = []
for i in range(10) :
answer = a[i]*2 if b[i]%2 == 0 else a[i]*3
result1.append(answer)
print np.array(result1)
# np.whereと関数を使えば１行で書けて速い
def func_double(x) :
return x*2
def func_triple(x) :
return x*3
result2 = np.where(b%2 == 0, func_double(a), func_triple(a))
print result2
なお、コメントにいただいた通り、この程度の関数であればリスト内包表記にそのまま関数を埋め込んでも良いと思います。
（ただ、個人的にはリスト内包表記に関数を直書きで埋め込むのはあまり好きではないです…あとから色々変更しにくいし、変更忘れたりするし、ファーストランゲージがpython世代の若い学生さんがリスト内包表記内にクソ長いコード書いてきたりすると、「すっげえ読みにくいィィッ！！！」とキレそうになるから（笑））
np.frompyfuncを使う（ただし使用に注意が必要）
もうちょっと速くならないか、と探していたら以下のページを見つけたので、有難く使わせていただいた。
Pythonのリストの全要素に任意の関数をapplyする最速の方法
Python高速化実験～map関数とか～
というわけで、frompyfuncを使ってみる。
import numpy as np
# prepare input arrays
a = np.array(range(10))
print "array a is", a
b = a + 100
print "array b is", b
def addition(x, y):
return x + y
np_addition = np.frompyfunc(addition, 2, 1)
print "print a + b using frompyfunc"
print np_addition(a, b)
print "print a + 1 using frompyfunc"
print np_addition(a, 1)
print "print 1 + b using frompyfunc"
print np_addition(1, b)
def subtruction(x, y):
return x - y
np_subtruction = np.frompyfunc(subtruction, 2, 1)
print "using np.where and frompyfuncs"
result2 = np.where(b%2 == 0, np_addition(a, b), np_subtruction(a, b))
print result2
frompyfuncで作ったユニバーサル関数の引数は、１つ目が関数オブジェクト、２つ目が引数の数、３つ目が戻り値の数。これを使うと、ユニバーサル関数の引数に配列を突っ込めば、それぞれの要素に関数を適用した結果を配列として返してくれる。
なんだ、こんな便利なものがあるなら、さっさと使えばよかった。
体感でも、list内包表記より３割くらい速い感じがする。
で、なにを確かめたかったかというと、このfrompyfuncで作ったユニバーサル関数の引数に、一部だけ配列を渡し、残りはただのfloatにする、とかいうことが可能なのか？ という話。 結果は、全然問題なし！ （訂正）そういう問題ではなくて、単純にaddition, subtruct関数の中で使用しているnumpyの＋演算子と-演算子が、引数配列でもスカラーでも受け付けてくれる、という機能が働いただけでした（涙） というわけで、以前書いていた部分は削除。（訂正おわり）以下、記述が間違っていたので畳みます（将来の自分がまた同じミスをしないように畳んで残しておく）
つまり、配列を渡した部分だけを要素ごとに変更しながら、計算した結果を返してくれるのである。
しかも、どの引数を配列にしようが、勝手に判断してくれる。スバラシイ！！
なので、引数のa,bにnp.arrayでなくただのリストを突っ込むと、上の例題死にます。
ただ、frompyfuncには気をつけなければならない側面もあるので、以下それについて記述。
frompyfuncの戻り値（配列）のdtypeはObject型になってしまう！
frompyfuncで作ったユニバーサル関数の戻り値（numpy.array）は、dtypeがObject型になってしまう。（中身の要素のtypeは保持される）
普通に配列の中身を参照するだけなら良いけれど、たとえばnumpy.histogramaddなどで使おうとすると、以下のようにObject型から他の型へのキャストができないと言って怒られる。
...numpy/lib/function_base.py", line 1014, in histogramdd
flatcount = bincount(xy, weights)
TypeError: Cannot cast array data from dtype('O') to dtype('float64') according to the rule 'safe'
これを防ぐには、戻り値のnumpy.arrayをastypeでキャストしてやれば良いわけだが、、、
result2 = np.where(b%2 == 0, np_addition(a, b).astype(np.float64), np_subtruction(a, b).astype(np.float64))
問題は、astypeは新しい配列を作って返す関数なので、ここで配列のコピーが生じてしまうことだ。
条件によっては、このコピー時間のおかげで、結局リスト内包表記でまわしても変わらないような実行速度になるかもしれない。
np.whereとfrompyfuncで作ったユニバーサル関数を併用について
例題の最後の３行で、np.whereとfrompyfuncで作ったユニバーサル関数が併用できることが確認できた。
しかし、どうやら、これをやるとかなり遅くなる。一時１０倍遅くなる、と思ったのはプログラムのミスだったが、２倍くらいは遅くなっている模様。
というわけで、速度を重視する場面では使いにくい感じ。
結局、どうするのが一番速いのか？
ここに挙げたようなシンプルな例題の場合は、ほぼリスト内包表記でかたがつくと思う。
しかし、内包表記内が関数やら条件やらでごちゃごちゃしてとにかく可読性が落ちるのが嫌だ、という場合や、そもそも１行ではとても書き切れない複雑な作業をやらせたい場合というのはある。
そのような場合に、調べるのが面倒でつい昔勉強した言語の常識でコードを書くと、予想外に実行時間がかかってしまう、ということを実感した。
すでに書いてしまったコードを少しでも早くしようと色々画策してはみたが、結局、素材となる関数の引数を、最初からスカラーでも配列でも受け付けるように、かつ速く動くように書き直すのが一番良いのではないか、という気がしてきた。
それができれば、np.whereと既存の関数の合わせ技が一番見た目もスッキリして速いし、frompyfuncでユニバーサル関数を作る必要もない。
私は最初FortranやC++でプログラムを学んだ人間なので、ついデータは1件ずつ処理する関数を書き、それをforでまわして処理する思考から抜け切れていないプログラムを書いてしまった。
numpyをゴリゴリ使うことが最初からわかっている場合、numpyには配列を速く扱う関数がたくさんあるので、最初からデータを複数レコードまとめてマトリクスとして処理するデザインにしておくのが正解だったかな、と思う。