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

score 4 · Accepted Answer · answered May 31 '12 at 12:12

4

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

answered May 31 '12 at 12:12

Chris

44,602
16
137
156

This looks very good to me. I think with a `numpy.array()` call wrapped around the list comprehension the output will be a (a.shape[0]/M,M,M) array. And the index arithmetic can be generalised to off-diagonal blocks. Great! – berian May 31 '12 at 13:02
Glad I could help. Note however, that I haven't really thought about what happens when `M` does not neatly divide into the matrix shape. You'll most likely need some input checking if you wrap this up in a `extract_block_diag` function. – Chris May 31 '12 at 13:10

score 4 · Answer 2 · answered Jun 02 '12 at 13:23

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

score 0 · Answer 3 · answered May 31 '12 at 12:02

Any particular reason you don't want to use a straightforward approach?

>>> A=np.array([[1,1,1,1],[2,2,2,2],[3,3,3,3],[4,4,4,4]])
>>> M1s,M2s=0,2 # start from A[M1s,M2s]
>>> M=2  # extract an MxM block
>>> for a in A[M1s:M1s+M]:
...   print a[M2s:M2s+M]
... 
[1 1]
[2 2]
>>>

score 0 · Answer 4 · edited May 23 '17 at 11:57

Based on unutbu's answer on https://stackoverflow.com/a/8070716/219229, I get the following:(BTW what's wrong with operating at the byte level?)

from numpy.lib.stride_tricks import as_strided

...

def extract_block_diag(A, M=3, k=0):
   ny, nx = A.shape
   ndiags = min(map(lambda x: x//M, A.shape))
   offsets = (nx*M+M, nx, 1)
   strides = map(lambda x:x*A.itemsize, offsets)
   if k > 0:
       B = A[:,k*M]
       ndiags = min(nx//M-k, ny//M)
   else:
       k = -k
       B = A[k*M]
       ndiags = min(nx//M, ny//M-k)
   return as_strided(B, shape=(ndiags,M,M),
                     strides=((nx*M+M)*A.itemsize, nx*A.itemsize, A.itemsize))

Example usage:

# a1, a2, a3 from your example
>>> bigA = scipy.linalg.block_diag(a1, a2, a3)
>>> extract_block_diag ( bigA, 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]]])
>>> extract_block_diag ( bigA, 3 )[2]
array([[3, 3, 3],
       [3, 3, 3],
       [3, 3, 3]])
>>> extract_block_diag(np.arange(1,9*9+1).reshape(9,9),3,1)
[[[ 4  5  6]
 [13 14 15]
 [22 23 24]]

[[34 35 36]
 [43 44 45]
 [52 53 54]]]

Note that the above function returns a view, which if you change anything inside the returned array of arrays, the original is also affected. Make a copy if necessary.

Rasoul Narimani · Answer 5 · 2015-02-17T19:52:01.313

import numpy as np

def extract_blocks(array):

prev = -1
for i in xrange(len(array)-1):
    if array[i+1][i] == 0 and array[i][i+1] == 0:
        yield array[prev + 1: i + 1, prev + 1: i + 1]
        print prev + 1, i
        prev = i
yield array[prev + 1: len(array), prev + 1: len(array)]

d = np.array([[4, 4, 0, 0, 0, 0, 0, 0, 0], [4, 4, 0, 0, 0, 0, 0, 0, 0], [0, 0, 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]])

for h in extract_blocks(d):

print h

[[4 4][4 4]], [[1]], [[2 2 2][2 2 2][2 2 2]], [[3 3 3][3 3 3][3 3 3]]

use extract_blocks(array) to spit all diagonal blocks of any size. — Rasoul Narimani, Feb 16 '15 at 19:11

Extracting diagonal blocks from a numpy array

5 Answers5

Linked