numpyで2次元配列上のこのトリプルループをベクトル化するにはどうすればよいですか？

Question

私はこの計算では、すべてのPythonのループを排除することができる：x[i]は、y[j]は、z[k]は長Nとx、yのベクトルで

result[i,j,k] = (x[i] * y[j] * z[k]).sum() 

は、zは長A、B有する第一の寸法を有し、C S。T.出力は形状(A,B,C)であり、各要素は で、3つの積の合計（要素ごと）です。

私は3から1ループ（下のコード）からダウンすることができますが、最後のループを解消するために を試みました。

A=B=C（必要な場合は少量のパディング）。あなたのループのバージョンと同じです

np.einsum('im,jm,km->ijk',x,y,z) 

：

# Example with 3 loops, 2 loops, 1 loop (testing omitted) 

N = 100 # more like 100k in real problem 
A = 2 # more like 20 in real problem 
B = 3 # more like 20 in real problem 
C = 4 # more like 20 in real problem 

import numpy 
x = numpy.random.rand(A, N) 
y = numpy.random.rand(B, N) 
z = numpy.random.rand(C, N) 

# outputs of each variant 
result_slow = numpy.empty((A,B,C)) 
result_vec_C = numpy.empty((A,B,C)) 
result_vec_CB = numpy.empty((A,B,C)) 

# 3 nested loops 
for i in range(A): 
    for j in range(B): 
     for k in range(C): 
      result_slow[i,j,k] = (x[i] * y[j] * z[k]).sum() 

# vectorize loop over C (2 nested loops) 
for i in range(A): 
    for j in range(B): 
     result_vec_C[i,j,:] = (x[i] * y[j] * z).sum(axis=1) 

# vectorize one C and B (one loop) 
for i in range(A): 
    result_vec_CB[i,:,:] = numpy.dot(x[i] * y, z.transpose()) 

numpy.testing.assert_almost_equal(result_slow, result_vec_C) 
numpy.testing.assert_almost_equal(result_slow, result_vec_CB) 

Answer 1

あなたがnumpyの> 1.6を使用している場合には、素晴らしいnp.einsum機能があります。実際の問題で配列のサイズに達すると、これが効率的にどのように公平になるかはわかりません（実際には私のマシン上でsegfaultを取得しています）。私がこの種の問題に好まれる他の解決策は、cythonを使ってメソッドを書き直すことです。

Answer 2

einsumを使用すると、あなたのケースで多くの意味があります。あなたは手でこれをかなり簡単に行うことができます。そのトリックは、配列を相互にブロードキャスト可能にすることです。それは、各配列がそれ自身の軸に沿って独立して変化するようにそれらを再形成することを意味します。次にそれらを掛け合わせて、numpyが放送を世話するようにします。最後（右端）の軸に沿って合計します。

>>> x = numpy.arange(2 * 4).reshape(2, 4) 
>>> y = numpy.arange(3 * 4).reshape(3, 4) 
>>> z = numpy.arange(4 * 4).reshape(4, 4) 
>>> (x.reshape(2, 1, 1, 4) * 
... y.reshape(1, 3, 1, 4) * 
... z.reshape(1, 1, 4, 4)).sum(axis=3) 
array([[[ 36, 92, 148, 204], 
     [ 92, 244, 396, 548], 
     [ 148, 396, 644, 892]], 

     [[ 92, 244, 396, 548], 
     [ 244, 748, 1252, 1756], 
     [ 396, 1252, 2108, 2964]]]) 

あなたはスライス表記を使用して、このもう少し一般化することができ、（以下同様Noneで動作しますので、Noneに等しい）newaxis値、及びsumが負の軸の値を受け入れるという事実（-1は最後を示し、-2は最後のものを示します）。この方法では、配列の元の形状を知る必要はありません。最後の軸に互換性がある限り、最初の3つを一緒にブロードキャストします：

>>> (x[:, numpy.newaxis, numpy.newaxis, :] * 
... y[numpy.newaxis, :, numpy.newaxis, :] * 
... z[numpy.newaxis, numpy.newaxis, :, :]).sum(axis=-1) 
array([[[ 36, 92, 148, 204], 
     [ 92, 244, 396, 548], 
     [ 148, 396, 644, 892]], 

     [[ 92, 244, 396, 548], 
     [ 244, 748, 1252, 1756], 
     [ 396, 1252, 2108, 2964]]])