Block tridiagonal matrix python

Question

I would like to create a block tridiagonal matrix starting from three numpy.ndarray. Is there any (direct) way to do that in python?

Thank you in advance!

Ch

Answer 1

For building a block-wise tridiagonal matrix from the three individual blocks (and repeat the blocks for N times), one solution can be:

import numpy as np
from scipy.linalg import block_diag

def tridiag(c, u, d, N): 
    # c, u, d are center, upper and lower blocks, repeat N times
    cc = block_diag(*([c]*N)) 
    shift = c.shape[1]
    uu = block_diag(*([u]*N)) 
    uu = np.hstack((np.zeros((uu.shape[0], shift)), uu[:,:-shift]))
    dd = block_diag(*([d]*N)) 
    dd = np.hstack((dd[:,shift:],np.zeros((uu.shape[0], shift))))
    return cc+uu+dd

For example:

c=np.matrix([[1,2],[3,4]])
u = np.matrix([[5,6],[7,8]])
d = -1*u
N=4
H = tridiag(c,u,d,N)
print(H)

gives the answer

[[ 1.  2.  5.  6.  0.  0.  0.  0.]
 [ 3.  4.  7.  8.  0.  0.  0.  0.]
 [-5. -6.  1.  2.  5.  6.  0.  0.]
 [-7. -8.  3.  4.  7.  8.  0.  0.]
 [ 0.  0. -5. -6.  1.  2.  5.  6.]
 [ 0.  0. -7. -8.  3.  4.  7.  8.]
 [ 0.  0.  0.  0. -5. -6.  1.  2.]
 [ 0.  0.  0.  0. -7. -8.  3.  4.]]