Numpy：对于二维数组中每行的对角线计算最快的方法

Question

给定一个 2d Numpy 数组，我想以最快的方式计算每一行的对角线。目前我正在使用列表推导式，但我想知道是否可以进行向量化处理？

例如，使用以下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]]

我想计算下列数组：

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

- Kel Solaar

---

你需要所有这些对角矩阵可能有很好的理由，但对于大的 n 值，你将浪费大量空间。我不知道你在做什么会使对角矩阵比包含相同信息的向量更可取，但你可能想看看是否有一种使用数据的紧凑表示方法来完成它。 - jme

我正在实现 Von Kries 色彩适应变换的向量化形式：http://en.wikipedia.org/wiki/Von_Kries_Coefficient_Law#von_Kries_Transform - Kel Solaar

如果你只需要将矩阵乘以对角矩阵，那么最终的结果就相当于向量与矩阵的广播乘法。虽然你的矩阵只有三列宽，所以你不会浪费太多空间。这只是一个想法。 - jme

Answer 1

以下是一种使用逐元素乘法的方法：np.eye(3)（3x3单位矩阵）和稍微重新调整形状的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]]])

通过“重新塑形M”，我是指将M的行面向z轴而不是y轴，使M具有形状(5, 1, 3)。

- Alex Riley

---

太棒了，正是我想要的，比我的使用情况快了7倍以上。 - Kel Solaar

Answer 2

尽管@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

时间显示大约快了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

- Saullo G. P. Castro

---

在我的数百万行数据中，差异并不是很大，但仍然稍微快一些，我会将您的答案标记为被接受的。谢谢！ - Kel Solaar

@KelSolaar 你使用的NumPy版本是哪个？在1.9.0版本中，花式索引有显著变化。 - Saullo G. P. Castro

目前是1.8.1版本，很可能是速度提升不如您自己的测试的原因。 - Kel Solaar

很高兴看到有一个更快的方法（+1）！对我来说，这只是在1.8.0上的轻微改进：是时候升级我的NumPy版本了... - Alex Riley

在numpy 1.8中，您可能会更有效地使用切片：b = np.zeros((M.shape[0], M.shape[1]*M.shape[1])); b[:, ::M.shape[1]+1] = M; return b.reshape(M.shape[0], M.shape[1], M.shape[1])。在1.9中，与Saullo的答案相比，差别很小，尽管对于非常大的数组来说，它稍微快一些。 - Jaime