使用scipy.sparse的三对角块矩阵

Question

I need help in building a block tridiagonal matrix using scipy sparse.

What I mean by that is for a square matrix B,

I need to create

[[B I 0 0 0]
 [I B I 0 0]
 [0 I B I 0]
 [0 0 I B I]
 [0 0 0 I B]]

Now, I want this to be programmatically done since the size of the matrix may vary.

Thanks!

Answer 1

Solved it!

I just used scipy.sparse.bmat in conjunction with list comprehensions.

A = sparse.bmat([[B if i == j else np.eye(n) if abs(i-j)==1
                else None for i in range(n)]
                for j in range(n)], format='bsr')

Where B is an nxn matrix.