質問: 最近数値計算をpythonでやり始めた者です。 例えば ``` A.shape #(3,5,5) B.shape #(3,) ``` の形をもつA,Bについて ``` numpy.array([ A[i,:,:] * B[i] for i in range(3)]).reshape(5,5) ``` を計算したいのですが、この方法では規模が大きくなると遅いような気がしています。 この操作をより速く動くように書く方法はないものでしょうか?
回答: 別のやり方を説明を兼ねて。 基本的に質問者様が行いたいのは要素積だと思います。numpy.arrayの要素積はnp.multiplyもしくは*で行えますが、今の場合、配列の次元が違うのが厄介です。 np.multiplyは高次元の配列Aと一次元配列Bを掛けあわせた場合には、配列Aの最後の軸とBの要素積を行う仕様になっているようです。具体的には、C[i,j,k,l] = A[i,j,k,l] * B[l]のような積を返します。 今の場合はC[i,j,k] = A[i,j,k] * B[i]という積を行いたいので、そのためには以下の様な方法が可能だと思います(これ以外にも方法はあるかもしれません)。 1. 配列Bの次元をAと揃える。具体的には、B[i]に余分な軸を付け足すことでB[i,j,k]のような配列に変形してから要素積をとる。このやり方としては、np.newaxisを使う、Noneを使うなどがあります。前出の回答のnp.array([[B]])も実質的に同じことをしているはずです。(正確には、追加した軸の次元は1なので、AとBの配列は違う次元ですが、その辺はブロードキャストされて計算できるようです) 2. 転置をとることで、無理やりC[i,j,k] = A[i,j,k] * B[k]の形に落とす。 3. np.einsumを使う。この関数を使うと、任意のテンソルを任意の形で積をとることができるので、単純に実装するだけでいいなら万能だと思いますが、場合によっては速度の問題が起こるかもしれません。というのも、一般には配列の積はどの順番でループをとるかで(numpyを使う場合には明示的ループは書きませんが)、速度が変わるはずですが、この関数ではループの順番を明示的に指定できないからです。 4. 内包表記で具体的に書いてしまう。 くらいでしょうか。 ``` import numpy as np A = np.random.rand(3,5,5) B = np.random.rand(3) #Method 1 print A*B[:,np.newaxis,np.newaxis] print A*B[:,None,None] print A * np.array([[B]]).T print A*np.expand_dims(np.expand_dims(B,1),1) #Method 2 print (A.T * B).T #Method 3 print np.einsum('ijk,i->ijk',A,B) ``` 手元で速度を比べてみましたが、A.shape = (10000, 7, 7)では、方法1,2,3はほぼ同じで、方法4だけ10倍程度遅かったです。