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

Question

我有一个N维数组。我想将它扩展为一个(N+1)-dimensional数组，方法是将最终维数的值放在对角线上。

例如，使用显式循环：

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)-diagonal矩阵，而不需要多次调用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。

将矩阵重塑为2维：

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