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

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 强制该 View 可写，并写入 View :

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