在多維numpy數組中迭代向量

Question

我有一個3xNxM numpy數組a，我想在最后兩個軸上進行迭代：a [：，x，y]。 優雅的方法是：

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完成嗎？ 這樣做的目的是對每個條目執行矩陣乘法，例如[：，x，y] = M [：，：，x，y] .dot（a [：，x，y]）。 另一種MATLAB風格的方法是將a（3，N * M）和M重塑為（3,3 * N * M）並采用點積，但這往往會占用大量內存。

Answer 1

雖然對形狀進行愚弄可能會使你想要完成的事情變得更加清晰，但是在不考慮太多問題的情況下處理這類問題的最簡單方法是使用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是非常強大的伏都einsum 。 你也可以在下面的評論中做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]]])

Answer 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) 在M的最后一個軸和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)的倒數第二個軸a.swapaxis(0,1)您擁有（3，）和（4,5），因此np.dot返回的結果是一個形狀數組（3,4,5） - 正是我們想要的。