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

My answer builds of @TheCorwoodRep's answer. I am just posting it because I made a few changes to make it more modular so that it would work for different orders of matrices and also changing the values of k1,k2,k3 i.e which decide where the diagonal appears, will take care of the overflow automatically. While calling the function you can specify what values should appear on the diagonals.

import numpy as np
def tridiag(T,x,y,z,k1=-1, k2=0, k3=1):
    a = [x]*(T-abs(k1)); b = [y]*(T-abs(k2)); c = [z]*(T-abs(k3))
    return np.diag(a, k1) + np.diag(b, k2) + np.diag(c, k3)

D=tridiag(10,-1,2,-1)