Numpy symmetric 4D matrix construction

Question

I would like to construct an array with the following structure:

A[i,j,i,j,] = B[i,j] with all other entries 0: A[i,j,l,k]=0 # (i,j) =\= (l,k)

Ie if I have the B matrix constructed how can I create the matrix A , preferably in a vectorized manner.

Explicitly, let B = [[1,2],[3,4]]

Then:

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

Answer 1

We can use an open grid to assign to A broadcasting the indexing arrays across the axes:

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

Answer 2

This is my solution with indexing:

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

Output:

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