CODEWITHSUNDEEP
pythonnumpy
MikeL
MikeL

Reputation: 2469

Vectorized syntax for creating a sequence of block matrices in NumPy

I have two 3D arrays A and B with shapes (k, n, n) and (k, m, m) respectively. I would like to create a matrix C of shape (k, n+m, n+m) such that for each 0 <= i < k, the 2D matrix C[i,:,:] is the block diagonal matrix obtained by putting A[i, :, :] at the upper left n x n part and B[i, :, :] at the lower right m x m part.

Currently I am using the following to achieve this is NumPy:

C = np.empty((k, n+m, n+m))
for i in range(k):
    C[i, ...] = np.block([[A[i,...], np.zeros((n,m))],
                          [np.zeros((m,n)), B[i,...]]])

I was wondering if there is a way to do this without the for loop. I think if k is large my solution is not very efficient.

Upvotes: 2

Views: 74

Answers (1)

Divakar
Divakar

Reputation: 221614

IIUC You can simply slice and assign -

C = np.zeros((k, n+m, n+m),dtype=np.result_type(A,B))
C[:,:n,:n] = A
C[:,n:,n:] = B

Upvotes: 1

Related Questions