Creating Tridiagonal Block Matrices

[dutta-alankar]

Tridiagonal block matrices are quite frequent in Physics and Engineering problems. For example Hamiltonian Matrices in Quantum Mechanics are often Tridiagonal blocks. So it is really handy to have a pre-built method to deal with it. You may consider adding the following piece to your repository:

def block_tridiag(diagonals): # 'I' or 0 as members are allowed
    #Tridiagonal block matrices are quite frequent in Physics and Engineering problems
    #Ordering from left to right
    #diagonals[0]: diagonal
    #diagonals[1]: upper diag
    #diagnals[2]: lower diag
    if ((len(diagonals[0])!=len(diagonals[1])+1) or (len(diagonals[0])!=len(diagonals[2])+1)):
        print('Off diagonals must be of smaller sizes! Please check')
        return None
    shape = len(diagonals[0])
    mat = ()
    for i in range(shape):
        tup = ()
        for j in range(shape):
            if (i==j):   tup = (*tup,diagonals[0][i])
            elif (i==j-1): tup = (*tup,diagonals[1][-i])
            elif (i==j+1): tup = (*tup,diagonals[2][i-1])
            else: tup = (*tup,0)
        mat = (*mat,tup)        
    return block(mat)

Example usage is as follows:

import numpy as np
d1 = np.array([[1,0],[0,2]])
d2 = 2*d1
d3 = 2*d2
dp = np.array([[0,1j],[1j,0]])
dm = -dp
H = block.block_tridiag([[d1,d2,d3],[0,dp],[dm,'I']])
print(H)
matrix([[ 1.+0.j,  0.+0.j,  0.+0.j,  0.+0.j,  0.+0.j,  0.+0.j],
            [ 0.+0.j,  2.+0.j,  0.+0.j,  0.+0.j,  0.+0.j,  0.+0.j],
            [-0.-0.j, -0.-1.j,  2.+0.j,  0.+0.j,  0.+0.j,  0.+1.j],
            [-0.-1.j, -0.-0.j,  0.+0.j,  4.+0.j,  0.+1.j,  0.+0.j],
            [ 0.+0.j,  0.+0.j,  1.+0.j,  0.+0.j,  4.+0.j,  0.+0.j],
            [ 0.+0.j,  0.+0.j,  0.+0.j,  1.+0.j,  0.+0.j,  8.+0.j]])

I have modified block.py here: https://github.com/dutta-alankar/block

[bamos]

Thanks! Merged in.