Search code examples
pythonarraysalgorithmnumpyindex-error

what is the solution to a tridiagonal matrix in python?


I'm working on a problem right now on tridiagonal matrix, i used the tridiagonal matrix algorithim in wiki https://en.wikipedia.org/wiki/Tridiagonal_matrix_algorithm to implement a solution and i have attempted it but my solution is not complete.

I'm confused and i need help also here is my script using jupyter notebook

import numpy as np 

# The diagonals and the solution vector 

b = np.array((5, 6, 6, 6, 6, 6, 6, 6, 5), dtype=float) # Main Diagonal
a= np.array((1, 1, 1, 1, 1, 1, 1, 1), dtype = float) # Lower Diagonal
c = np.array((1, 1, 1, 1, 1, 1, 1, 1), dtype = float) # Upper Diagonal
d = np.array((3, 2, 1, 3, 1, 3, 1, 2, 3), dtype = float) # Solution Vector

#number of rows 
n = d.size

newC = np.zeros(n, dtype= float)

newD = np.zeros(n, dtype = float)

x = np.zeros(n, dtype = float)

# the first coefficents 
newC[0] = c[0] / b[0]
newD[0] = d[0] / b[0]

for i in range(1, n - 1):
newC[i] = c[i] / (b[i] - a[i] * newC[i - 1])

for i in range(1, n -1):
newD[i] = (d[i] - a[i] * newD[i - 1]) / (b[i] - a[i] * newC[i - 1])

x[n - 1] = newD[n - 1]

for i in reversed(range(0, n - 1)):
    x[i] = newD[i] - newC[i] * x[i + 1]


x

Solution

  • The vectors a and c should be the same length as b and d, so just prepend/append zero the them respectively. Additionally, the range should be range(1,n) otherwise your last solution element is 0 when it shouldn't be. You can see this modified code, as well as a comparison to a known algorithm, showing that it gets the same answer.

    import numpy as np 
    
    # The diagonals and the solution vector 
    
    b = np.array((5, 6, 6, 6, 6, 6, 6, 6, 5), dtype=float) # Main Diagonal
    a= np.array((0, 1, 1, 1, 1, 1, 1, 1, 1), dtype = float) # Lower Diagonal
    c = np.array((1, 1, 1, 1, 1, 1, 1, 1, 0), dtype = float) # Upper Diagonal
    d = np.array((3, 2, 1, 3, 1, 3, 1, 2, 3), dtype = float) # Solution Vector
    
    print(b.size)
    print(a.size)
    
    #number of rows 
    n = d.size
    
    newC = np.zeros(n, dtype= float)
    
    newD = np.zeros(n, dtype = float)
    
    x = np.zeros(n, dtype = float)
    
    # the first coefficents 
    newC[0] = c[0] / b[0]
    newD[0] = d[0] / b[0]
    
    for i in range(1, n):
        newC[i] = c[i] / (b[i] - a[i] * newC[i - 1])
    
    for i in range(1, n):
        newD[i] = (d[i] - a[i] * newD[i - 1]) / (b[i] - a[i] * newC[i - 1])
    
    x[n - 1] = newD[n - 1]
    
    for i in reversed(range(0, n - 1)):
        x[i] = newD[i] - newC[i] * x[i + 1]
    
    # Test with know algorithme
    mat = np.zeros((n, n))
    for i in range(n):
        mat[i,i] = b[i]
        if i < n-1:
            mat[i, i+1] = a[i+1]
            mat[i+1, i] = c[i]
    
    print(mat)
    
    sol = np.linalg.solve(mat, d)
    print(x)
    print(sol)