python - 迭代多维numpy数组中的向量

我有一个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完成吗？其目的是对每个条目执行矩阵乘法，例如a[：，x，y]=M[：，：，x，y].dot（a[：，x，y]）。另一种类似于MATLAB的方法是重塑a（3，N*M）和M（3，3*N*M）的形状并取一个点积，但这会消耗大量内存。

最佳答案

虽然玩弄形状可能会让你想要完成的事情更加清晰，但处理这类问题而不想太多的最简单方法是使用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

编辑是非常强大的巫术。您也可以按照操作人员在下面的注释中的要求执行以下操作：

>>> 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]]])

关于python - 迭代多维numpy数组中的向量，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/15015471/