Extracting diagonal blocks from a numpy array

Question

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 diagonal
    of square matrix A, whose size must be a multiple of M."""

    # Check that the matrix can be block divided
    if 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 diagonal
    if abs(k) > M - 1:
        raise StandardError('kth diagonal does not exist in matrix')
    elif k > 0:
        ro = 0
        co = abs(k)*M 
    elif k < 0:
        ro = abs(k)*M
        co = 0
    else:
        ro = 0
        co = 0

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

Answer 1

As a starting point you could just use something like

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

>>> M = 3
>>> [a[i*M:(i+1)*M,i*M:(i+1)*M] for i in range(a.shape[0]/M)]
[array([[1, 1, 1],
   [1, 1, 1],
   [1, 1, 1]]), array([[2, 2, 2],
   [2, 2, 2],
   [2, 2, 2]]), array([[3, 3, 3],
   [3, 3, 3],
   [3, 3, 3]])]

Answer 2

You can also do it with views. This is probably faster than the indexing approach.

import numpy as np
import scipy.linalg

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

b = scipy.linalg.block_diag(a1, a2, a3)
b[1,4] = 4
b[1,7] = 5
b[4,1] = 6
b[4,7] = 7
b[7,1] = 8
b[7,4] = 9
print b

def extract_block_diag(a, n, k=0):
    a = np.asarray(a)
    if a.ndim != 2:
        raise ValueError("Only 2-D arrays handled")
    if not (n > 0):
        raise ValueError("Must have n >= 0")

    if k > 0:
        a = a[:,n*k:] 
    else:
        a = a[-n*k:]

    n_blocks = min(a.shape[0]//n, a.shape[1]//n)

    new_shape = (n_blocks, n, n)
    new_strides = (n*a.strides[0] + n*a.strides[1],
                   a.strides[0], a.strides[1])

    return np.lib.stride_tricks.as_strided(a, new_shape, new_strides)

print "-- Diagonal blocks:"
c = extract_block_diag(b, 3)
print c

print "-- They're views!"
c[0,1,2] = 9
print b[1,2]

print "-- 1st superdiagonal"
c = extract_block_diag(b, 3, 1)
print c

print "-- 2nd superdiagonal"
c = extract_block_diag(b, 3, 2)
print c

print "-- 3rd superdiagonal (empty)"
c = extract_block_diag(b, 3, 3)
print c

print "-- 1st subdiagonal"
c = extract_block_diag(b, 3, -1)
print c

print "-- 2nd subdiagonal"
c = extract_block_diag(b, 3, -2)
print c