scipy/numpyのufunc備忘録

scipy/numpyの行列演算は詳しいことを知らなくても何となく使えて便利なのですが,慣れていないとついついループや内包表記を使ってしまいます.pythonのループで書くと遅くなるのは分かっているのに…….最近「ufuncのbroadcastingが分かっていれば大分マシになるのでは無いか」と考えるようになったので,備忘録も兼ねて適当なことを書いてみます.


ちなみに私の数値計算の知識はサッパリなので注意.またnumpy/scipyに関して何も知らない場合は他の記事を当たった方が良いと思います.
計算の速度やメモリのことは良く分からないので,難しいことを考えるのを止めて,目標を ループを使わずにとにかく1行で処理を書く にします.おいおい.

broadcasting

universal function(以下ufunc)はarray*1に作用する関数のことで,これの"array broadcasting"のおかげで行列とスカラーの足し算等が簡単に出来たりします. 詳しくは公式の解説を読め!で終わる話なのですが;

broadcastingのルールは以下の4つ.

  1. 各入力のarrayに対してndimが最大の入力より小さければ,1がshapeの先頭に追加される.
  2. 出力の各次元の大きさは,入力の対応する次元の最大値となる.
  3. 入力の各次元の大きさは,対応する出力の次元の大きさと等しいか1で無ければならない.
  4. 入力のある次元の大きさが1の時,その要素がその次元の計算で繰り返し利用される.

何のこっちゃさっぱり分かりません.ということで簡単な例で試してみます.
ちなみにndimは次元数,shapeは各次元の大きさのタプルです.array([ [1,2,3], [4,5,6] ]) の ndimは2,shapeは(2, 3) といった感じ.

例1

行列とベクトルの足し算を考えます.

import scipy as sp
a = sp.array([[1,2,3],[4,5,6],[7,8,9]])
b = sp.array([10,20,30])
print(a+b)
[[1, 2, 3],
 [4, 5, 6],  + [10, 20, 30]
 [7, 8, 9]]

a,bのndimが一致していないので,ルール1よりbのshapeの先頭に1が追加される(b.shape -> (1,3))

[[1, 2, 3],                  
 [4, 5, 6],  + [[10, 20, 30]]
 [7, 8, 9]]                  

a,bのshapeはそれぞれ (3,3), (1,3) なのでルール2より出力は(3,3)となる

[[1, 2, 3],                     [[?, ?, ?],
 [4, 5, 6],  + [[10, 20, 30]] =  [?, ?, ?],
 [7, 8, 9]]                      [?, ?, ?]]

ルール3は満たしているのでエラーは起きない.ルール4よりbの[10,20,30]が繰り返し使われる

[[1, 2, 3],    [[10, 20, 30],   [[?, ?, ?],
 [4, 5, 6],  +  [10, 20, 30], =  [?, ?, ?],
 [7, 8, 9]]     [10, 20, 30]]    [?, ?, ?]]

これより結果は

[[11, 22, 33], 
 [14, 15, 16], 
 [17, 28, 39]]

となります.

例2

a = [a0, a1, ..., aM] と b = [b0, b1, ..., bN] から

[[a0*b0  a0*b1 ... a0*bN ],
 [a1*b0  a1*b1 ... a1*bN ],
 ...
 [aM*b0            aM*bN ]]

を求めたい.共分散行列とかそんな時に.ちなみにnumpy.outer使えば終わる話なのですが;*2
aとbのshapeはそれぞれ(M,)と(N,)なので結果の(M,N)を求めるには,aのshapeを(M,1)にすれば上手く行きそうです.このような操作はnumpy.newaxisを利用してやればOK*3.具体的には

In [49]: import scipy as sp
In [50]: a = sp.array([1,2,3,4,5])
In [51]: b = sp.array([2,2,2,3,3])
In [52]: a[:, sp.newaxis]
Out[52]: 
array([[1],
       [2],
       [3],
       [4],
       [5]])
In [53]: a[:, sp.newaxis]*b
Out[53]: 
array([[ 2,  2,  2,  3,  3],
       [ 4,  4,  4,  6,  6],
       [ 6,  6,  6,  9,  9],
       [ 8,  8,  8, 12, 12],
       [10, 10, 10, 15, 15]])

具体的な添字の代わりに":"を指定すると「その次元の要素を全て」と認識しているが,実際どうなのかはよく知りません…….
ちなみに":"の個数が分からない場合は(入力の次元数が変化する),"..."とかを利用すると良い感じ.()

例3

EMアルゴリズムで混合ガウス分布のパラメータを求める.PRMLの9章に載ってるやつですね.
以前書いた記事の使い回しですが…….

記事書くの疲れてきたので数式は省略します.下巻p154とか見ればいいんじゃないでしょうか.各変数のshapeは

  • gamma: (N,K)
  • Nk: (K,)
  • mu: (K,M)
  • sigma: (K,M,M)
  • pi: (K,)
  • X: (N,M)

とします.Mステップの更新だと
Nkはgamma(N,K)から(K,)を足し算で生成すれば良さそうです.ということで

Nk = sp.sum(gamma, 0)

sumの第二引数に数字を指定すると,その次元で足し算します.

sigmaはgamma(N,K)とX(N,M)とmu(K,M)とNk(K,)から(K,M,M)を作れば良さそうです.
まずは(N,K,M,M)を作り,そのあとsumでΣ相当のことをすると考えて

xmu = X[:, sp.newaxis, :] - mu
sigma = sp.sum(gamma[..., sp.newaxis, sp.newaxis] * xmu[:, :, sp.newaxis, :] * xmu[:, :, :, sp.newaxis], 0) / Nk[..., sp.newaxis, sp.newaxis]

PEP8(笑) *4

例4

以前しましまさんにコメントで尋ねられた例.

ところで numpy/scipy とかpython数値計算系はお詳しいですか?
私は numpy 修行中の身なのですが,演算の broadcasting とかがよく分からず,for文に手をだしてしまいます.

for x in xrange(m.nx):
for y in xrange(m.ny):
q[:, x, y] = m.pz * m.px[x, :] * m.py[y, :]

とかを1行で書こうとして挫折しました.scipy のサイトや,Python Scripting for Computational Science 本以外に良い資料をご存じでしたら教えていただけませんか?

http://d.hatena.ne.jp/se-kichi/20100519/1274275285#c1274496705

nxやnyやpxが何かよく分からないので「m.nxやm.px等はscipy/numpyとは関係無い」として考えます.
コメントのコードを見ると

  • qの次元数は3.
  • m.px と m.py の次元数は2.*5

と考えられる.m.nx, m.ny(それぞれX, Yと置く)でループをしているので,この値が添字で使われている次元の大きさに等しいと仮定すれば

  • q.shape -> (K, X, Y)
  • m.px.shape -> (X, L)
  • m.py.shape -> (Y, M)

となる.

  • q[:,x,y].shape -> (K,)
  • m.px[x, :].shape -> (L,)
  • m.py[y, :].shape -> (M,)

であるから K=L=M と分かる.m.pzは()*6と(K,)の可能性があるが,前者ならばループの外側で乗算を行うだろうから,後者だと考えて

  • q.shape -> (K, X, Y)
  • m.px.shape -> (X, K)
  • m.py.shape -> (Y, K)
  • m.pz.shape -> (K,)

となる.結局 m.pz, m.px, m.py から (K, X, Y) を乗算で作るので

  • m.px.shape -> (X, K) -> (K, X) -> (K, X, 1)
  • m.py.shape -> (Y, K) -> (K, Y) -> (K, 1, Y)
  • m.pz.shape -> (K,) -> (K, 1, 1)

と変換して計算すれば良さそう(というかこれ以外の選択肢無い……).コードにすると

q = m.pz[:, np.newaxis, np.newaxis]*m.px.T[:, :, np.newaxix]*m.py.T[:, np.newaxix, :]

実際に対話環境で試してみると

In [84]: import scipy as sp
In [85]: K = 10
In [86]: X = 13
In [87]: Y = 19
In [88]: px = sp.random.randint(0, 100, (X, K))
In [89]: py = sp.random.randint(0, 100, (Y, K))
In [90]: pz = sp.random.randint(0, 100, (K,))
In [91]: q1 = pz[:, sp.newaxis, sp.newaxis]*px.T[:, :, sp.newaxis]*py.T[:, sp.newaxis, :]
In [97]: q2 = sp.zeros((K, X, Y))
In [98]: for x in xrange(X):
             for y in xrange(Y):
                 q2[:, x, y] = pz * px[x, :] * py[y, :]
In [105]: sp.all(q2 == q1)
Out[105]: True

*1:正確にはndarray

*2:ちなみにufuncにはouterメソッドがあるので,足し算で似たようなことがやりたければ numpy.add.outer(a, b)

*3:ちなみに newaxis は None

*4:Limit all lines to a maximum of 79 characters.

*5:m.px[x, :]から分かるのは「m.pxの次元数 >= 2」だが,q[:, x, y]の次元数は1なのでm.pxの次元数は2.

*6:スカラー