2

我有一个 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) 并采用点积，但这往往会占用大量内存。

4

2 回答 2

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

EDIT 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]]])
于 2013-02-22T02:41:34.543 回答
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) 轴上执行求和积。Ma

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) 的数组 - 正是我们想。

于 2013-02-22T01:51:21.267 回答