从 N 维数组中的值构造 (N+1) 维对角矩阵

Question

我有一个N维数组。我想通过将最终维度的值放在对角线上来将其扩展为 (N+1) 维数组。

例如，使用显式循环：

In [197]: M = arange(5*3).reshape(5, 3)

In [198]: numpy.dstack([numpy.diag(M[i, :]) for i in range(M.shape[0])]).T
Out[198]: 
array([[[ 0,  0,  0],
        [ 0,  1,  0],
        [ 0,  0,  2]],

       [[ 3,  0,  0],
        [ 0,  4,  0],
        [ 0,  0,  5]],

       [[ 6,  0,  0],
        [ 0,  7,  0],
        [ 0,  0,  8]],

       [[ 9,  0,  0],
        [ 0, 10,  0],
        [ 0,  0, 11]],

       [[12,  0,  0],
        [ 0, 13,  0],
        [ 0,  0, 14]]])

这是一个5×3×3的数组。

我的实际数组很大，我想避免显式循环（隐藏 map 中的循环而不是列表理解没有性能提升；它仍然是一个循环）。尽管 numpy.diag 适用于构建常规的二维对角矩阵，但它不会扩展到更高的维度（当给定一个二维数组时，它将提取其对角线）。 numpy.diagflat 返回的数组将所有内容都变成一个大对角线，结果是一个 15×15 的数组，其中包含更多的零，无法重新整形为 5×3×3。

有没有一种方法可以从 N 维数组中的值高效地构造 (N+1) 对角矩阵，而无需多次调用 diag？

Answer 1

将N维数组最后一维转为对角矩阵的通用方法：

我们需要降低数组的维度，对每个向量应用 numpy.diag() 函数，然后将其重建为原始维度 + 1。

将矩阵重塑为二维：

M.reshape(-1, M.shape[-1])

然后使用 map 对其应用 np.diag，并使用以下方法重建具有附加维度的矩阵：

result.reshape([*M.shape, M.shape[-1]])

所有这些结合起来得出以下结果：

result = np.array(list(map(
   np.diag,
   M.reshape(-1, M.shape[-1])
))).reshape([*M.shape, M.shape[-1]])

一个例子：

shape = np.arange(2,8)
M = np.arange(shape.prod()).reshape(shape)
print(M.shape)  # (2, 3, 4, 5, 6, 7)

result = np.array(list(map(np.diag, M.reshape(-1, M.shape[-1])))).reshape([*M.shape, M.shape[-1]])

print(result.shape)  # (2, 3, 4, 5, 6, 7, 7)

和 res[0,0,0,0,2,:] 包含以下内容：

array([[14,  0,  0,  0,  0,  0,  0],
       [ 0, 15,  0,  0,  0,  0,  0],
       [ 0,  0, 16,  0,  0,  0,  0],
       [ 0,  0,  0, 17,  0,  0,  0],
       [ 0,  0,  0,  0, 18,  0,  0],
       [ 0,  0,  0,  0,  0, 19,  0],
       [ 0,  0,  0,  0,  0,  0, 20]])

Answer 2

使用numpy.diagonal to take a view of the relevant diagonals of a properly-shaped N+1-dimensional array, force the view to be writeable with setflags，并写入视图：

expanded = numpy.zeros(M.shape + M.shape[-1:], dtype=M.dtype)

diagonals = numpy.diagonal(expanded, axis1=-2, axis2=-1)
diagonals.setflags(write=True)

diagonals[:] = M

这会生成您想要的数组 expanded。

Answer 3

您可以使用无处不在的 np.einsum 的几乎不可能猜测的功能。当按如下方式使用时 einsum 将 return 广义对角线的可写视图：

>>> import numpy as np
>>> M = np.arange(5*3).reshape(5, 3)
>>> 
>>> out = np.zeros((*M.shape, M.shape[-1]), M.dtype)
>>> np.einsum('...jj->...j', out)[...] = M
>>> out
array([[[ 0,  0,  0],
        [ 0,  1,  0],
        [ 0,  0,  2]],

       [[ 3,  0,  0],
        [ 0,  4,  0],
        [ 0,  0,  5]],

       [[ 6,  0,  0],
        [ 0,  7,  0],
        [ 0,  0,  8]],

       [[ 9,  0,  0],
        [ 0, 10,  0],
        [ 0,  0, 11]],

       [[12,  0,  0],
        [ 0, 13,  0],
        [ 0,  0, 14]]])

从 N 维数组中的值构造 (N+1) 维对角矩阵

Construct (N+1)-dimensional diagonal matrix from values in N-dimensional array

python

arrays

numpy

multidimensional-array

diagonal