CODEWITHSUNDEEP
pythonarraysnumpymatrixscipy
Matthias Nickel
Matthias Nickel

Reputation: 23

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

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!

Upvotes: 2

Views: 778

Answers (2)

panadestein
panadestein

Reputation: 1291

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.

Upvotes: 2

V. Ayrat
V. Ayrat

Reputation: 2719

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

Upvotes: 2

Related Questions