Numpy：计算2d阵列每行对角线的最快方法

Given a 2d Numpy array, I would like to be able to compute the diagonal for each row in the fastest way possible, I'm right now using a list comprehension but I'm wondering if it can be vectorised somehow?

给定一个2d Numpy数组,我希望能够以最快的方式计算每一行的对角线,我现在正在使用列表理解但是我想知道它是否可以以某种方式进行矢量化?

For example using the following M array:

例如,使用以下M数组:

M = np.random.rand(5, 3)


[[ 0.25891593  0.07299478  0.36586996]
 [ 0.30851087  0.37131459  0.16274825]
 [ 0.71061831  0.67718718  0.09562581]
 [ 0.71588836  0.76772047  0.15476079]
 [ 0.92985142  0.22263399  0.88027331]]

I would like to compute the following array:

我想计算以下数组:

np.array([np.diag(row) for row in M])

array([[[ 0.25891593,  0.        ,  0.        ],
        [ 0.        ,  0.07299478,  0.        ],
        [ 0.        ,  0.        ,  0.36586996]],

       [[ 0.30851087,  0.        ,  0.        ],
        [ 0.        ,  0.37131459,  0.        ],
        [ 0.        ,  0.        ,  0.16274825]],

       [[ 0.71061831,  0.        ,  0.        ],
        [ 0.        ,  0.67718718,  0.        ],
        [ 0.        ,  0.        ,  0.09562581]],

       [[ 0.71588836,  0.        ,  0.        ],
        [ 0.        ,  0.76772047,  0.        ],
        [ 0.        ,  0.        ,  0.15476079]],

       [[ 0.92985142,  0.        ,  0.        ],
        [ 0.        ,  0.22263399,  0.        ],
        [ 0.        ,  0.        ,  0.88027331]]])

2 个解决方案

#1

Despite the good answer of @ajcr, a much faster alternative can be achieved with fancy indexing (tested in NumPy 1.9.0):

尽管@ajcr得到了很好的答案,但通过花式索引可以实现更快的替代方案(在NumPy 1.9.0中测试):

import numpy as np

def sol0(M):
    return np.eye(M.shape[1]) * M[:,np.newaxis,:]

def sol1(M):
    b = np.zeros((M.shape[0], M.shape[1], M.shape[1]))
    diag = np.arange(M.shape[1])
    b[:, diag, diag] = M
    return b

where the timing shows this is approximately 4X faster:

时间显示这大约快4倍:

M = np.random.random((1000, 3))
%timeit sol0(M)
#10000 loops, best of 3: 111 µs per loop
%timeit sol1(M)
#10000 loops, best of 3: 23.8 µs per loop

#2

Here's one way using element-wise multiplication of np.eye(3) (the 3x3 identity array) and a slightly re-shaped M:

这是使用np.eye(3)(3x3身份数组)和略微重新形状M的元素乘法的一种方法:

>>> M = np.random.rand(5, 3)
>>> np.eye(3) * M[:,np.newaxis,:]
array([[[ 0.42527357,  0.        ,  0.        ],
        [ 0.        ,  0.17557419,  0.        ],
        [ 0.        ,  0.        ,  0.61920924]],

       [[ 0.04991268,  0.        ,  0.        ],
        [ 0.        ,  0.74000307,  0.        ],
        [ 0.        ,  0.        ,  0.34541354]],

       [[ 0.71464307,  0.        ,  0.        ],
        [ 0.        ,  0.11878955,  0.        ],
        [ 0.        ,  0.        ,  0.65411844]],

       [[ 0.01699954,  0.        ,  0.        ],
        [ 0.        ,  0.39927673,  0.        ],
        [ 0.        ,  0.        ,  0.14378892]],

       [[ 0.5209439 ,  0.        ,  0.        ],
        [ 0.        ,  0.34520876,  0.        ],
        [ 0.        ,  0.        ,  0.53862677]]])

(By "re-shaped M" I mean that the rows of M are made to face out along the z-axis rather than across the y-axis, giving M the shape (5, 1, 3).)

(通过“重新成形M”我的意思是使M行沿z轴而不是横穿y轴,使M形状(5,1,3)。)

#1