Is there a reverse operation to scipy.linalg.block_diag?

Viewed 28

I know scipy offers a handy function to convert from a list of arrays to a block diagonal matrix, e.g.

>>> from scipy.linalg import block_diag
>>> A = [[1, 0],
         [0, 1]]
>>> B = [[3, 4, 5],
         [6, 7, 8]]
>>> C = [[7]]
>>> D = block_diag(A, B, C)
>>> D
array([[1, 0, 0, 0, 0, 0],
       [0, 1, 0, 0, 0, 0],
       [0, 0, 3, 4, 5, 0],
       [0, 0, 6, 7, 8, 0],
       [0, 0, 0, 0, 0, 7]])

Is there a reverse operation to this? I.e. take a block diagonal matrix and a list of chunk sizes and decompose to a list of arrays.

a, b, c = foo(D, block_sizes=[(2,2), (2,3), (1,1)])

If there's no handy built-in way to accomplish this, is there a better (more performant) implementation than looping over the input array? The naive implementation probably looks something like this:

def foo(matrix, block_sizes):
    result = []
    curr_row, curr_col = 0, 0 
    for nrows, ncols in block_sizes:
        result.append(matrix[curr_row:curr_row + nrows, curr_col:curr_col + ncols])
        curr_row += nrows
        curr_col += ncols
    return result
1 Answers

Look at the code for block_diag:

The core part allocates a out array, and then iteratively assigns the blocks to it:

out = np.zeros(np.sum(shapes, axis=0), dtype=out_dtype)
r, c = 0, 0
for i, (rr, cc) in enumerate(shapes):
    out[r:r + rr, c:c + cc] = arrs[i]
    r += rr
    c += cc

Were you imagining something more "performant"?

What you want is a list of arrays that (may) differ in shape. How are you going to get those without some sort of (python level) iteration?

Look at the code for np.split (and variants). It too iterates, taking multiplie slices to produce a list of arrays.

Related