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

Use the function scipy.sparse.diags.

Example:

from scipy.sparse import diags
import numpy as np
#
n = 10
k = np.array([np.ones(n-1),-2*np.ones(n),np.ones(n-1)])
offset = [-1,0,1]
A = diags(k,offset).toarray()

This returns:

array([[-2.,  1.,  0.,  0.,  0.],
       [ 1., -2.,  1.,  0.,  0.],
       [ 0.,  1., -2.,  1.,  0.],
       [ 0.,  0.,  1., -2.,  1.],
       [ 0.,  0.,  0.,  1., -2.]])

Answer 2

Since tridiagonal matrix is a sparse matrix using a sparse package could be a nice option, see http://pysparse.sourceforge.net/spmatrix.html#matlab-implementation, there are some examples and comparisons with MATLAB even...

Answer 3

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.]]

Answer 4

@TheCorwoodRep's answer can actually be done in a single line. No need for a seperate function.

np.eye(3,3,k=-1) + np.eye(3,3)*2 + np.eye(3,3,k=1)*3

This produces:

array([[ 2.,  3.,  0.],
       [ 1.,  2.,  3.],
       [ 0.,  1.,  2.]])

Answer 5

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)

Answer 6

With "regular" numpy arrays, using numpy.diag:

def tridiag(a, b, c, k1=-1, k2=0, k3=1):
    return np.diag(a, k1) + np.diag(b, k2) + np.diag(c, k3)

a = [1, 1]; b = [2, 2, 2]; c = [3, 3]
A = tridiag(a, b, c)