2013-02-22 72 views
2

我有一個3xNxM的numpy數組a,我想遍歷最後兩個軸:a [:,x,y]。不雅的做法是:在多維的numpy數組中迭代向量

import numpy as np 
a = np.arange(60).reshape((3,4,5)) 
M = np. array([[1,0,0], 
       [0,0,0], 
       [0,0,-1]]) 

for x in arange(a.shape[1]): 
    for y in arange(a.shape[2]): 
     a[:,x,y] = M.dot(a[:,x,y]) 

這可以用nditer完成嗎?這樣做的目標是對每個條目執行矩陣乘法,例如, a [:,x,y] = M [:,:,x,y] .dot(a [:,x,y])。另一種方法是用MATLAB的方法重塑一個(3,N * M)和M(3,3 * N * M)並且取一個點積,但這往往會消耗大量的內存。

回答

5

雖然與形狀鬼混可能使你所要完成更清楚,處理這類問題沒有想得太多了最簡單的方法是用np.einsum

In [5]: np.einsum('ij, jkl', M, a) 
Out[5]: 
array([[[ 0, 1, 2, 3, 4], 
     [ 5, 6, 7, 8, 9], 
     [ 10, 11, 12, 13, 14], 
     [ 15, 16, 17, 18, 19]], 

     [[ 0, 0, 0, 0, 0], 
     [ 0, 0, 0, 0, 0], 
     [ 0, 0, 0, 0, 0], 
     [ 0, 0, 0, 0, 0]], 

     [[-40, -41, -42, -43, -44], 
     [-45, -46, -47, -48, -49], 
     [-50, -51, -52, -53, -54], 
     [-55, -56, -57, -58, -59]]]) 

加上它常來與業績獎金:

In [17]: a = np.random.randint(256, size=(3, 1000, 2000)) 

In [18]: %timeit np.dot(M, a.swapaxes(0,1)) 
10 loops, best of 3: 116 ms per loop 

In [19]: %timeit np.einsum('ij, jkl', M, a) 
10 loops, best of 3: 60.7 ms per loop 

編輯einsum爲v強大的巫術力量。您還可以按照以下注釋做OP:

>>> a = np.arange(60).reshape((3,4,5)) 
>>> M = np.array([[1,0,0], [0,0,0], [0,0,-1]]) 
>>> M = M.reshape((3,3,1,1)).repeat(4,axis=2).repeat(5,axis=3) 
>>> np.einsum('ijkl,jkl->ikl', M, b) 
array([[[ 0, 1, 2, 3, 4], 
     [ 5, 6, 7, 8, 9], 
     [ 10, 11, 12, 13, 14], 
     [ 15, 16, 17, 18, 19]], 

     [[ 0, 0, 0, 0, 0], 
     [ 0, 0, 0, 0, 0], 
     [ 0, 0, 0, 0, 0], 
     [ 0, 0, 0, 0, 0]], 

     [[-40, -41, -42, -43, -44], 
     [-45, -46, -47, -48, -49], 
     [-50, -51, -52, -53, -54], 
     [-55, -56, -57, -58, -59]]]) 
+0

謝謝!有沒有一種方法可以將元素乘法包含在Einsum中?例如,我有一個大小爲3xMxN的矢量陣列(如上所示),我想對每個矢量執行矩陣旋轉,大小爲3x3xMxN的張量M. ((3,4,5)) M =數組([[1,0,0], [0,0,0], [0,0,-1 ](a.shape [1])) M = M.reshape((3,3,1,1))。repeat(4,axis = 2).repeat(5,axis = 3) ): for arange(a.shape [2]): a [:,x,y] = M [:,:,x,y] .dot(a [:,x,y]) 什麼是正確的消解方法?我得到了: b = einsum('ijkl,jxy',M,a) c = b [:,0,0,:,] – emarti 2013-02-22 08:09:08

+0

@emarti我編輯了我的答案來掩蓋您的問題評論。 – Jaime 2013-02-22 16:23:39

+1

+1在答案中使用伏都教 – evan54 2014-11-15 04:51:06

2
for x in np.arange(a.shape[1]): 
    for y in np.arange(a.shape[2]): 
     a[:,x,y] = M.dot(a[:,x,y]) 

相當於

a = np.dot(M,a.swapaxes(0,1)) 

In [73]: np.dot(M,a.swapaxes(0,1)) 
Out[73]: 
array([[[ 0, 1, 2, 3, 4], 
     [ 5, 6, 7, 8, 9], 
     [ 10, 11, 12, 13, 14], 
     [ 15, 16, 17, 18, 19]], 

     [[ 0, 0, 0, 0, 0], 
     [ 0, 0, 0, 0, 0], 
     [ 0, 0, 0, 0, 0], 
     [ 0, 0, 0, 0, 0]], 

     [[-40, -41, -42, -43, -44], 
     [-45, -46, -47, -48, -49], 
     [-50, -51, -52, -53, -54], 
     [-55, -56, -57, -58, -59]]]) 

說明:

對於多維陣列,np.dot(M,a)performs a sum product over the last axis of M and the second-to-last axis of a.

a具有形狀(3,4,5),但我們希望對形狀3的軸進行求和。由於倒數第二個軸將被求和,因此我們需要a.swapaxis(0,1) - 它具有形狀( 4,3,5) - 將3移動到倒數第二個軸。

M已成形(3,3),a.swapaxis(0,1)已成形(4,3,5)。刪除M的最後一個軸和a.swapaxis(0,1)的倒數第二軸會使(3,)和(4,5)離開,所以np.dot返回的結果是一個形狀數組(3,4,5) - 正是我們想要的。