おっぱいそん!

pythonを数値計算に使うときのテクニックとかをまとめていきたい。

要素積を任意の軸について

要素積 A*B or np.multiply(A,B) をndarray A, B の任意の軸について実行する方法。
(基本的な使い方は、以前の記事http://oppython.hatenablog.com/entry/2014/01/05/004454で紹介したので、そちらを参考に。)

np.multiply(A, B)はn次元配列Aと1次元配列Bを掛けると、Aの最後の軸とBの要素積をする。:
C[i,j,k] = A[i,j,k] B[k]
したがって、
C[i,j,k] = A[i,j,k] B[j]
のような、最後の軸以外の軸と要素積をしたい場合には、そのままでは使えない。

そこで、少し工夫して、最後の軸以外の軸との要素積を実装する。

もっとも、メジャーな方法は、1次元配列にnp.newaxis*1を使って軸を追加する方法。
np.newaxisはあるn次元配列B(例えば B.shape=(4,5,3) )に対し、Bnew = B[:,np.newaxis,:,np.newaxis,:]とすることで、Bnew.shape = (4,1,5,1,3)のように軸を追加できる機能。
例えば、上のC[i,j,k] = A[i,j,k] B[j]の例であれば、

import numpy as np

A = np.arange(24).reshape((2,3,4))
B = np.array([0,2,3])

print A.shape
print B.shape
print (A*B[np.newaxis,:,np.newaxis]).shape  # or (A*B[:,np.newaxis]) 

(2, 3, 4)
(3,)
(2, 3, 4)

とすれば良い。
参考:
http://stackoverflow.com/questions/7096371/in-numpy-what-is-the-fastest-way-to-multiply-the-second-dimension-of-a-3-dimens


他にも、配列の次元に依って、いくつか方法がある。
もっとも、単純な場合である2次元配列A(つまり行列)と1次元配列B(つまりベクトル)について、Aの行とBの要素積をやり方を以下に列挙する(どれも同じ結果を返す)。

A = np.arange(6).reshape((2,3))
B = np.array([0,2])

print A
print B
print A*B[:,None] # or A*B[:,np.newaxis] or B[:,None]*A
print A*np.expand_dims(B,1) 
print (A.T * B).T
print np.dot(np.diag(B), A)
print np.einsum('ij,i->ij',A,B)

[[0 1 2]
[3 4 5]]
[0 2]
[[ 0 0 0]
[ 6 8 10]]
....

None、np.newaxis、np.expand_dimsはどれも軸の追加で、同じもの。
転置とdot+diagは配列の次元が上がったり、要素積したい軸が変わると使えなくなる。
np.einsumはアインシュタインの縮約で、これ以外にも色々使い方があるみたい。

速度は方法に依って異なる。
速度の点では、
np.einsum > 転置 = np.newaxis > dot+diag の順で優れているっぽい。
参考:
http://stackoverflow.com/questions/18522216/multiplying-across-in-a-numpy-array


速度を気にしない時には、分かりやすいnp.newaxisを、速度が気になる時にはnp.einsumを使っとけばよさげ。