（M×N×N）行列の対角線を高速に設定する方法はありますか？ Einsum/n次元のfill_diagonal？

Question

私は行列に基づいて高速で最適化されたコードを書こうとしていますが、最近、著しいスピードアップを達成するためのツールとしてeinsumを発見しました。

これを使用して多次元配列の対角を効率的に設定できますか、それともデータのみを返すことはできますか？

私の問題では、各正方形（N×N）の行列の列を合計することによって正方行列の配列（形状：M×N×N）の対角線を設定しようとしています。

私の現在の（遅い、ループベース）ソリューションは、次のとおりです。

# Build dummy array 
dimx = 2 # Dimension x (likely to be < 100) 
dimy = 3 # Dimension y (likely to be between 2 and 10) 
M = np.random.randint(low=1, high=9, size=[dimx, dimy, dimy]) 

# Blank the diagonals so we can see the intended effect 
np.fill_diagonal(M[0], 0) 
np.fill_diagonal(M[1], 0) 

# Compute diagonals based on summing columns 
diags = np.einsum('ijk->ik', M) 

# Set the diagonal for each matrix 
# THIS IS LOW. CAN IT BE IMPROVED? 
for i in range(len(M)): 
    np.fill_diagonal(M[i], diags[i]) 

# Print result 
M 

これは全くしてください改善することができますか？ np.fill_diagonalは非正方行列を受け入れない（したがって、私のループベースの解を強制する）ようだ。おそらくエインサンもここで助けてくれる？

Answer 1

1つの方法は、2Dに変更することです。対角値を使用してncols+1のステップで列を設定します。整形するとビューが作成され、そのように対角線の位置に直接アクセスすることができます。したがって、実装は次のようになります -

s0,s1,s2 = M.shape 
M.reshape(s0,-1)[:,::s2+1] = diags 

Answer 2

あなたは2Dの場合には、それは

if a.ndim == 2: 
     step = a.shape[1] + 1 
     end = a.shape[1] * a.shape[1] 
    a.flat[:end:step] = val 

@Divakar's「ストライド」アプローチを使用するソリューションをすることによって、あなたの3Dの場合にこれを適用することnp.source(np.fill_diagonal)が表示されますない場合2次元で「平らにする」。

M.sum(axis=1)で列を合計できます。私はぼんやりと、実際にはちょっと速いと感じたタイミングを思い出します。 sumはもう少し従来のものです。

誰かがeinsumでサイズを拡張する機能をリクエストしていますが、それは起こるとは思われません。