从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。