Tridiagonal Matrix Algorithm - Python solver

tdma.py

# Solves the matrix equation A x = d and returns the solution x
def tridiagonal_thomas (A, d):
    A = np.copy(A)
    d = np.copy(d)
    n = len(d)
    
    # Gaussian elimination of lower diagonal
    for i in range(1, n):
        s = A[i, i-1] / A[i-1, i-1]
        A[i, i] = A[i, i] - s*A[i-1, i]
        d[i] = d[i] - s*d[i-1]
    
    # Backward substitution
    x = np.zeros(n, np.float64)
    x[n-1] = d[n-1] / A[n-1, n-1]
    
    for i in range(n-2, -1, -1):
        x[i] = (d[i] - A[i, i+1]*x[i+1]) / A[i, i]
        
    return x

# Example (as given in task 4)
A = np.array([
    [3., -1, 0., 0., 0., 0.],
    [-1, 3., -1, 0., 0., 0.],
    [0., -1, 3., -1, 0., 0.],
    [0., 0., -1, 3., -1, 0.],
    [0., 0., 0., -1, 3., -1],
    [0., 0., 0., 0., -1, 3.]
], np.float64)
d = np.array([0.2, 0.2, 0.2, 0.2, 0.2, 0.2]);

# Solve using TDMA
x = tridiagonal_thomas(A, d)
np.dot(A, x)