根据 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

使用numpy.diagonal 来查看形状正确的N+1 维数组的相关对角线，强制视图可使用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 2

您可以使用无处不在的np.einsum 的几乎不可能猜到的功能。如下使用einsum 将返回广义对角线的可写视图：

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

Answer 3

将N-D数组的最后一维转换为对角矩阵的一般方法：

我们需要降低数组的维度，将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]])