从多形状数组的列创建对角矩阵

Question

假设我有一个像这样的数组

mat = np.array([[1,1],[2,2],[3,3]])

其形状为 (3,2)。我想通过获取数组的最后一个轴并从中创建对角矩阵来创建形状为 (3,2,2) 的新数组。我可以用 for 循环来做到这一点

mat2 = np.zeros((3,2,2))
for i in np.arange(0,3):
    mat2[i] = np.diag(mat[i])

产生所需的输出

[[[1. 0.]
  [0. 1.]]

 [[2. 0.]
  [0. 2.]]

 [[3. 0.]
  [0. 3.]]]

但是有没有办法在直接矢量化（更快？！）版本中做到这一点？在我的实际问题中，我有一个形状为 (...,n) 的大型高维数组，需要将最后一个轴转换为最后形状为 (...,n,n) 的对角矩阵。

Answer 1

这里有一个可能的实现：使用np.diagonal获取相关对角线的视图，并强制该视图可写setflags，并写入视图：

expanded = np.zeros(mat.shape + mat.shape[-1:], dtype=mat.dtype)
diagonals = np.diagonal(expanded, axis1=-2, axis2=-1)
diagonals.setflags(write=True)

diagonals[:] = mat

expanded
array([[[1, 0],
        [0, 1]],

       [[2, 0],
        [0, 2]],

       [[3, 0],
        [0, 3]]])