Numpy Symmetric 4d Matrix Construction

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) I.e. if I have the B matrix constru

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

Solution 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()

# checkingfor i inrange(2):
    for j inrange(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.]]