Numpy 对称 4D 矩阵构造

Numpy symmetric 4D matrix construction

我想构造一个具有以下结构的数组:

A[i,j,i,j,] = B[i,j] 与所有其他条目 0:A[i,j,l,k]=0 # (i,j) =\= (l,k)

即如果我构建了 B 矩阵,我该如何创建矩阵 A,最好以矢量化方式创建。

明确地,让B = [[1,2],[3,4]]

然后:

A[1,1,:,:]=[[1,0],[0,0]]
A[1,2,:,:]=[[0,2],[0,0]]
A[2,1,:,:]=[[0,0],[3,0]]
A[2,2,:,:]=[[0,0],[0,4]]

我们可以使用 open grid 分配给 A 在轴上广播索引数组:

B = np.array([[1,2],[3,4]])
i,j = B.shape
A = np.zeros([i,j,i,j])
i, j = np.ogrid[:i, :j]
A[i,j,i,j] = B

print(A)

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

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


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

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

这是我的索引解决方案:

x,y = np.meshgrid(np.arange(B.shape[1]),np.arange(B.shape[0]))
x,y = x.ravel(), y.ravel()

A = np.zeros(B.shape + B.shape)

A[y.ravel(), x.ravel(), y.ravel(), x.ravel()] = B.ravel()

# checking
for i in range(2):
    for j in range(2):
        print(f'A[{i},{j}]:',A[i,j])

输出:

A[0,0]: [[1. 0.]
 [0. 0.]]
A[0,1]: [[0. 2.]
 [0. 0.]]
A[1,0]: [[0. 0.]
 [3. 0.]]
A[1,1]: [[0. 0.]
 [0. 4.]]