NumPyの行列演算で二重ループをなくす

100 Numpy Exercisesの86番を見て知ったテクニック。これは数値を2進数表記する問題。

86.Convert a vector of ints into a matrix binary representation (★★★)

# Author: Warren Weckesser
I = np.array([0, 1, 2, 3, 15, 16, 32, 64, 128])
B = ((I.reshape(-1,1) & (2**np.arange(8))) != 0).astype(int)
print(B[:,::-1])

# Author: Daniel T. McDonald
I = np.array([0, 1, 2, 3, 15, 16, 32, 64, 128], dtype=np.uint8)
print(np.unpackbits(I[:, np.newaxis], axis=1))

2番目のほうが直感的でわかりやすいかな。内容はまあいいとして、1番目の2行目で使っているテクニックがん？となった。わかりやすいように別の例に変える。

>>> x = np.arange(5)
>>> x
array([0, 1, 2, 3, 4])
>>> x.reshape(-1,1)
array([[0],
       [1],
       [2],
       [3],
       [4]])

xとx.reshape(-1,1)って演算することってできるんだ！ 数学的には正しい演算ではないだろうが(演算子が*の場合、ベクトルの外積になるらしい。np.outer(np.arange(5), np.arange(5))のように計算できる。なのでその拡張ととらえても良さそう)、プログラム上では結構便利な機能かもしれない。上記の例では論理積だったが、四則演算も当然できる。

>>> y = x*10
>>> y
array([ 0, 10, 20, 30, 40])

>>> x.reshape(-1,1)+y
array([[ 0, 10, 20, 30, 40],
       [ 1, 11, 21, 31, 41],
       [ 2, 12, 22, 32, 42],
       [ 3, 13, 23, 33, 43],
       [ 4, 14, 24, 34, 44]])
>>> x+y.reshape(-1,1)
array([[ 0,  1,  2,  3,  4],
       [10, 11, 12, 13, 14],
       [20, 21, 22, 23, 24],
       [30, 31, 32, 33, 34],
       [40, 41, 42, 43, 44]])
>>> y.reshape(-1,1)+x
array([[ 0,  1,  2,  3,  4],
       [10, 11, 12, 13, 14],
       [20, 21, 22, 23, 24],
       [30, 31, 32, 33, 34],
       [40, 41, 42, 43, 44]])
>>> y+x.reshape(-1,1)
array([[ 0, 10, 20, 30, 40],
       [ 1, 11, 21, 31, 41],
       [ 2, 12, 22, 32, 42],
       [ 3, 13, 23, 33, 43],
       [ 4, 14, 24, 34, 44]])

結果の違いは味わいがいがある。行列の和差は交換法則が成り立つので、reshapeするのに気をつければあんま気にしなくてもいいかなと思ったり。この手の処理は手続き脳だと二重ループで書いてしまいたくなるが、このやり方だと結構簡潔に書けそうだ。

Python使いがループを外しの職人芸をしたくなる気持ちがよくわかる。

この手の処理がループのシンタックスシュガーのように見えて気持ち悪いかもしれないが、そもそもNumPy自体がCのネイティブコードを動かして高速化を図ったものでPythonとは似て非なるもの。Cから見ればNumPy自体が広義のシンタックスシュガー（んなこと言ったら高級言語はみんなそうなるけど）、Pythonのforループは遅いということを先入観として叩き込まれている身としては（Rがそうだったからなんとなくしっくりくる）、ループのような処理は低レベルの言語に任せっきりにしてもいいのではないかなと思う。少なくともそれで速くなるのなら、NumPyの独特の記法にならってループを外したほうがいい。

atsuoishimoto.hatenablog.com

追記：この記事によると、Pythonのループ処理が遅い犯人はfor文ではなく、どちらかというとPythonが静的な型宣言を行わないからだそうだ。読んでいてへーっとなった。