Efficient way of constructing a 3D stack of block diagonal matrix in numpy/scipy from a 3D stack of matrices

Question

I am trying to construct a stack of block diagonal matrix in the form of nXMXM in numpy/scipy from a given stacks of matrices (nXmXm), where M=k*m with k the number of stacks of matrices. At the moment, I'm using the scipy.linalg.block_diag function in a for loop to perform this task:

import numpy as np
import scipy.linalg as linalg

a = np.ones((5,2,2))
b = np.ones((5,2,2))
c = np.ones((5,2,2))

result = np.zeros((5,6,6))

for k in range(0,5):
    result[k,:,:] = linalg.block_diag(a[k,:,:],b[k,:,:],c[k,:,:])

However, since n is in my case getting quite large, I'm looking for a more efficient way than a for loop. I found 3D numpy array into block diagonal matrix but this does not really solve my problem. Anything I could imagine is transforming each stack of matrices into block diagonals

import numpy as np
import scipy.linalg as linalg

a = np.ones((5,2,2))
b = np.ones((5,2,2))
c = np.ones((5,2,2))

a = linalg.block_diag(*a)
b = linalg.block_diag(*b)
c = linalg.block_diag(*c)

and constructing the resulting matrix from it by reshaping

result = linalg.block_diag(a,b,c)

result = result.reshape((5,6,6))

which does not reshape. I don't even know, if this approach would be more efficient, so I'm asking if I'm on the right track or if somebody knows a better way of constructing this block diagonal 3D matrix or if I have to stick with the for loop solution.

Edit: Since I'm new to this platform, I don't know where to leave this (Edit or Answer?), but I want to share my final solution: The highlightet solution from panadestein worked very nice and easy, but I'm now using higher dimensional arrays, where my matrices reside in the last two dimensions. Additionally my matrices are no longer of the same dimension (mostly a mixture of 1x1, 2x2, 3x3), so I adopted V. Ayrat's solution with minor changes:

def nd_block_diag(arrs):
    shapes = np.array([i.shape for i in arrs])

    out = np.zeros(np.append(np.amax(shapes[:,:-2],axis=0), [shapes[:,-2].sum(), shapes[:,-1].sum()]))
    r, c = 0, 0
    for i, (rr, cc) in enumerate(shapes[:,-2:]):
        out[..., r:r + rr, c:c + cc] = arrs[i]
        r += rr
        c += cc

    return out

which works also with array broadcasting, if the input arrays are shaped properly (i.e. the dimensions, which are to be broadcasted are not added automatically). Thanks to pandestein and V. Ayrat for your kind and fast help, I've learned a lot about the possibilites of list comprehensions and array indexing/slicing!

panadestein · Accepted Answer · 2020-06-15 14:48:39Z

2

I don't think that you can escape all possible loops to solve your problem. One way that I find convenient and perhaps more efficient than your for loop is to use a list comprehension:

import numpy as np
from scipy.linalg import block_diag

# Define input matrices

a = np.ones((5, 2, 2))
b = np.ones((5, 2, 2))
c = np.ones((5, 2, 2))

# Generate block diagonal matrices

mats = np.array([a, b, c]).reshape(5, 3, 2, 2)
result = [block_diag(*bmats) for bmats in mats]

Maybe this can give you some ideas to improve your implementation.

edited Jun 15, 2020 at 14:48

answered Jun 15, 2020 at 9:45

panadestein

1,29110 silver badges22 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Matthias Nickel Over a year ago

Thanks for your quick reply and kind help, I think I'll use this solution for my first draft for the sake of simplicity.

panadestein Over a year ago

@MatthiasNickel no problem, I am happy it helped. Further improvements could perhaps be obtained by using numpy's vectorize method, although it might be not so obvious to implement.

V. Ayrat · Accepted Answer · 2020-06-15 11:09:39Z

2

block_diag also just iterate through shapes. Almost all time spend in copying data so you can do it whatever way your want for example with little change of source code of block_diag

arrs = a, b, c
shapes = np.array([i.shape for i in arrs])
out = np.zeros([shapes[0, 0], shapes[:, 1].sum(), shapes[:, 2].sum()])
r, c = 0, 0

for i, (_, rr, cc) in enumerate(shapes):
    out[:, r:r + rr, c:c + cc] = arrs[i]
    r += rr
    c += cc

print(np.allclose(result, out))
# True

answered Jun 15, 2020 at 11:09

V. Ayrat

2,74912 silver badges10 bronze badges

1 Comment

Matthias Nickel Over a year ago

Thanks also for this suggestion, I think it will be handy to have a seperate function as soon as my code gets more complex and when I'm going to use this functionality more often. I just thought, I've overlooked some builtin function, which solves this task, but maybe It's not as common as I thought...

Collectives™ on Stack Overflow

Efficient way of constructing a 3D stack of block diagonal matrix in numpy/scipy from a 3D stack of matrices

2 Answers 2

2 Comments

1 Comment

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

2 Comments

1 Comment

Your Answer

Sign up or log in

Post as a guest

Linked

Related