Say I have a set of $n \times n$ matrices $A_1, ..., A_m$ as numpy arrays. I'd like to create the block matrix defined below.

enter image description here

I'm looking for a clean, elegant, and easy-to-interpret way of doing this in numpy. I tried this with np.block:

a1, a2 = np.full((2, 2), 1), np.full((2, 2), 2)
out = np.block([[a1, (a1+a2)/2],
                [(a1+a2)/2, a2]])

but that approach doesn't generalize to an arbitrary number of $m$ matrices.

An approach that I found which is general is the following:

A = np.array([a1, a2])
out = (A[:, :, None, :] + A.transpose(1, 0, 2)[None, :, :, :]).reshape(n * m, -1)

but that one, while being efficient, is fairly hard-to-read (this code will be read much more often than written).

scipy.linalg.block_diag gets me halfway there, but I don't get the off-diagonals.

Can anyone think of a good alternative solution? I was thinking of looking into numpy's array-generation-from-function routines, but haven't found a good way of going about that yet.

mats_row=np.array([[A1,A2,A3,...,A_n]]) #Create array of matrices with shape (1,M,N,N)
mats_column=np.transpose(mats_row,(1,0,2,3)) #Make a copy with shape (M,1,N,N)
block=.5*(mats_row+mats_column) #Add, broadcasting to a (M,M,N,N) array

This works by utilizing Numpy's broadcasting capabilities. The basic idea is similar to if you added a matrix where each row was $[A_1,A_2,..A_n]$ to a matrix where each column was $[A_1,A_2,..A_n]$, but broadcasting allows you to do this without ever explicitly forming these intermediates matrices.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.