python - 从 3D 创建 4D 上对角线数组

Question

假设我有一个 (x, y, z) 大小的矩阵。现在，我希望创建一个维度为 (x, y, i, i) 的新矩阵，其中 (i, i) 矩阵是上对角线，并由z 维度上的值。在 numpy 中是否有一些简单的方法可以在不使用超过 1 个 for 循环(循环 x)的情况下执行此操作？谢谢。

编辑

original = np.array([
    [
        [0, 1, 3],
        [4, 5, 6]
    ],
    [
        [7, 8, 9],
        [3, 2, 1]
    ],
])

new = np.array([
    [
        [
            [0, 1],
            [0, 3]
        ],
        [
            [4, 5],
            [0, 6]
        ]
    ],
    [
        [
            [7, 8],
            [0, 9]
        ],
        [
            [3, 2],
            [0, 1]
        ]
    ]
])

因此，使用上面的内容我们可以看到

original[0, 0, :] = [0 1 3]

new[0, 0, :, :] = [[0 1]
                  [0 3]]

Answer 1

这是一种使用 bool 索引的方法 -

n = 2 # This would depend on a.shape[-1]
out = np.zeros(a.shape[:2] + (n,n,),dtype=a.dtype)
out[:,:,np.arange(n)[:,None] <= np.arange(n)] = a

示例运行 -

In [247]: a
Out[247]: 
array([[[0, 1, 3],
        [4, 5, 6]],

       [[7, 8, 9],
        [3, 2, 1]]])

In [248]: out
Out[248]: 
array([[[[0, 1],
         [0, 3]],

        [[4, 5],
         [0, 6]]],


       [[[7, 8],
         [0, 9]],

        [[3, 2],
         [0, 1]]]])

可以建议使用下标索引替代最后一步的另一种方法 -

r,c = np.triu_indices(n)
out[:,:,r,c] = a

注意:如前所述，n 将取决于 a.shape[-1]。在这里，我们将 a.shape[-1] 设为 3，因此 n 为 2。如果 a.shape[-1] 为 6，则 n 将为 3，依此类推。关系是:(n*(n+1))//2 == a.shape[-1]。