如何使用 numpy 将 2D 数组转换为 3D 对角矩阵

Question

有一个形状为 (i,j) 的 2D numpy 数组，其中每一列都是一个包含 i 个条目的向量，我如何获得一个形状为 (j, i, i) 的 3D numpy 数组，其中每个 j 对应于每个向量都在其对角线上而不使用 for 循环的对角矩阵？

例如，拥有二维数组 [[1,5],[2,6],[3,7],[4,8]]，我想获得一个 3D 数组，其中对于 j=0，矩阵为

[[1,0,0,0],[0,2,0,0],[0,0,3,0],[0,0,0,4]]

对于 j=1，矩阵将是

[[5,0,0,0],[0,6,0,0],[0,0,7,0],[0,0,0,8]].

我试过 np.diag 但显然它在提供 2D 矩阵时输出矩阵对角线。

Answer 1

您可以使用 np.einsum:

np.einsum('ij,ik->jik', a, np.eye(a.shape[0], dtype=a.dtype))

输出：

array([[[1, 0, 0, 0],
        [0, 2, 0, 0],
        [0, 0, 3, 0],
        [0, 0, 0, 4]],

       [[5, 0, 0, 0],
        [0, 6, 0, 0],
        [0, 0, 7, 0],
        [0, 0, 0, 8]]])

如何使用 numpy 将 2D 数组转换为 3D 对角矩阵

How to turn a 2D array into a 3D diagonal matrix with numpy

python

numpy

multidimensional-array

diagonal