I am searching for a neat way to extract the diagonal blocks of size 2x2 that lie along the main diagonal of a (2N)x(2N) numpy array (that is, there will be N such blocks). This generalises numpy.diag, which returns elements along the main diagonal, that one might think of as 1x1 blocks (though of course numpy doesn't represent them this way).
To phrase this a little more broadly, one might wish to extract N MxM blocks from a (MN)x(MN) array. This seems the complement of scipy.linalg.block_diag, ably discussed in How can I transform blocks into a blockdiagonal matrix (NumPy), pulling blocks out of the places that block_diag will have put them.
Borrowing code from the solution to that question, here's how this could be set up:
>>> a1 = np.array([[1,1,1],[1,1,1],[1,1,1]])
>>> a2 = np.array([[2,2,2],[2,2,2],[2,2,2]])
>>> a3 = np.array([[3,3,3],[3,3,3],[3,3,3]])
>>> import scipy.linalg
>>> scipy.linalg.block_diag(a1, a2, a3)
array([[1, 1, 1, 0, 0, 0, 0, 0, 0],[1, 1, 1, 0, 0, 0, 0, 0, 0],[1, 1, 1, 0, 0, 0, 0, 0, 0],[0, 0, 0, 2, 2, 2, 0, 0, 0],[0, 0, 0, 2, 2, 2, 0, 0, 0],[0, 0, 0, 2, 2, 2, 0, 0, 0],[0, 0, 0, 0, 0, 0, 3, 3, 3],[0, 0, 0, 0, 0, 0, 3, 3, 3],[0, 0, 0, 0, 0, 0, 3, 3, 3]])
Then, we would wish to have a function like
>>> A = scipy.linalg.block_diag(a1, a2, a3)
>>> extract_block_diag(A, M=3)
array([[[1, 1, 1],[1, 1, 1],[1, 1, 1]],[[2, 2, 2],[2, 2, 2],[2, 2, 2]],[[3, 3, 3],[3, 3, 3],[3, 3, 3]]])
To continue the analogy with numpy.diag, one might also wish to extract the off-diagonal blocks: N - k of them on the kth block diagonal. (And in passing, an extension of block_diag allowing block to be placed off the principal diagonal would certainly be useful, but that is not the scope of this question.) In the case of the array above, this could yield:
>>> extract_block_diag(A, M=3, k=1)
array([[[0, 0, 0],[0, 0, 0],[0, 0, 0]],[[0, 0, 0],[0, 0, 0],[0, 0, 0]]])
I see that the use of stride_tricks covered in this question aims to produce something similar to this functionality, but I understand that striding operates at the byte level, which doesn't sound quite appropriate.
By way of motivation, this arises from the situation where I wish to extract the diagonal elements of a covariance matrix (that is, the variances), where the elements themselves are not scalar but 2x2 matrices.
Edit: Based on the suggestion from Chris, I have made the following attempt:
def extract_block_diag(A,M,k=0):"""Extracts blocks of size M from the kth diagonalof square matrix A, whose size must be a multiple of M."""# Check that the matrix can be block dividedif A.shape[0] != A.shape[1] or A.shape[0] % M != 0:raise StandardError('Matrix must be square and a multiple of block size')# Assign indices for offset from main diagonalif abs(k) > M - 1:raise StandardError('kth diagonal does not exist in matrix')elif k > 0:ro = 0co = abs(k)*M elif k < 0:ro = abs(k)*Mco = 0else:ro = 0co = 0blocks = np.array([A[i+ro:i+ro+M,i+co:i+co+M] for i in range(0,len(A)-abs(k)*M,M)])return blocks
where the the following results will be returned for the data above:
D = extract_block_diag(A,3)
[[[1 1 1][1 1 1][1 1 1]][[2 2 2][2 2 2][2 2 2]][[3 3 3][3 3 3][3 3 3]]]D = extract_block_diag(A,3,-1)
[[[0 0 0][0 0 0][0 0 0]][[0 0 0][0 0 0][0 0 0]]]