如何有效地将元素分配到 3d 矩阵的对角线上？

Question

a=np.zeros((3,3,3))
b=np.arange(3)
c=np.arange(9).reshape(3,3)

我想将数组的元素b或c沿着 3d 矩阵（张量）的对角线（或对角线上方/下方）a相对于特定轴放置。

我累了numpy.diagflat，但它只适用于二维矩阵。

例如，如何制作以下矩阵？

array([[[ 0.,  0.,  0.],
        [ 0.,  1.,  0.],
        [ 0.,  0.,  2.]],

       [[ 0.,  0.,  0.],
        [ 0.,  1.,  0.],
        [ 0.,  0.,  2.]],

       [[ 0.,  0.,  0.],
        [ 0.,  1.,  0.],
        [ 0.,  0.,  2.]]])

Answer 1

对于主对角线，您可以使用np.einsum. 例如：

>> np.einsum('iii->i', a)[...] = b
>>> a
array([[[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  0.]],

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

       [[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  2.]]])

或者：

>>> np.einsum('iji->ji', a)[...] = c
>>> a
array([[[ 0.,  0.,  0.],
        [ 3.,  0.,  0.],
        [ 6.,  0.,  0.]],

       [[ 0.,  1.,  0.],
        [ 0.,  4.,  0.],
        [ 0.,  7.,  0.]],

       [[ 0.,  0.,  2.],
        [ 0.,  0.,  5.],
        [ 0.,  0.,  8.]]])

编辑：广播工作正常：

>>> np.einsum('ijj->ij', a)[...] = b
>>> a
array([[[ 0.,  0.,  0.],
        [ 0.,  1.,  0.],
        [ 0.,  0.,  2.]],

       [[ 0.,  0.,  0.],
        [ 0.,  1.,  0.],
        [ 0.,  0.,  2.]],

       [[ 0.,  0.,  0.],
        [ 0.,  1.,  0.],
        [ 0.,  0.,  2.]]])

子对角线也可以使用，但由于需要一些手动切片而更加棘手。例如：

>>> a=np.zeros((3,3,3))
>>> np.einsum('iij->ij', a[:2,1:])[...] = c[1:]
>>> a
array([[[ 0.,  0.,  0.],
        [ 3.,  4.,  5.],
        [ 0.,  0.,  0.]],

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

       [[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  0.]]])