Extract blocks from block diagonal matrix

extract_block_diagonal_blocks.py

import numpy as np

def extract_block_diagonal_blocks(A):
    assert A.shape[0] == A.shape[1], "matrix A should be square"

    N = A.shape[0]

    A_mirrored = A + A.T   # make the matrix symmetric so we only have to check one triangle

    blocks = []
    start = 0
    blocksize = 0

    while start < N:
        blocksize += 1

        if np.all(A_mirrored[start:start + blocksize, start + blocksize:N] == 0):
            block = A[start:start + blocksize, start:start + blocksize]
            blocks.append(block)
            start += blocksize
            blocksize = 0

    return blocks


# Example usage
A = np.array([[1, 0, 0, 0],
              [0, 2, 3, 0],
              [0, 4, 5, 0],
              [0, 0, 0, 6]])

block_diagonal_blocks = extract_block_diagonal_blocks(A)

# Printing the blocks
for i, block in enumerate(block_diagonal_blocks):
    print(f"A{i+1}:")
    print(block)
    print()

"""
Output:

A1:
[[1]]

A2:
[[2 3]
 [4 5]]

A3:
[[6]]
"""