Why does this implementation of the Block Tridiagonal Thomas algorithm give such large errors?

Question

See below my attempt at implementing the block tridiagonal thomas algorithm. If you run this however, you get a relatively large (10^-2) error in the block TMDA compared to the np direct solve (10^-15), even for this very simple case. More complicated test cases give larger errors - I think the errors start growing on the back substitution. Any help as to why would be much appreciated!

import numpy as np
import torch

def solve_block_tridiagonal(a, b, c, d):

    N = len(b)
    x = np.zeros_like(d)
    
    # Forward elimination with explicit C* and d* storage
    C_star = np.zeros_like(c)
    d_star = np.zeros_like(d)

    # Initial calculations for C_0* and d_0*
    C_star[0] = np.linalg.solve(b[0], c[0])
    d_star[0] = np.linalg.solve(b[0], d[0])

    # Forward elimination
    for i in range(1, N - 1):
        C_star[i] = np.linalg.solve(b[i] - a[i-1] @ C_star[i-1], c[i])
        d_star[i] = np.linalg.solve(b[i] - a[i-1] @ C_star[i-1], d[i] - a[i-1] @ d_star[i-1])

    # Last d_star update for the last block
    d_star[-1] = np.linalg.solve(b[-1] - a[-2] @ C_star[-2], d[-1] - a[-2] @ d_star[-2])

    # Backward substitution
    x[-1] = d_star[-1]
    for i in range(N-2, -1, -1):
        x[i] = d_star[i] - C_star[i] @ x[i+1]

    return x


def test_block_tridiagonal_solver():

    N = 4

    a = np.array([
        [[1, 0.5], [0.5, 1]],  
        [[1, 0.5], [0.5, 1]],
        [[1, 0.5], [0.5, 1]]
    ], dtype=np.float64)
    
    b = np.array([
        [[5, 0.5], [0.5, 5]],  
        [[5, 0.5], [0.5, 5]],
        [[5, 0.5], [0.5, 5]],
        [[5, 0.5], [0.5, 5]]
    ], dtype=np.float64)
    
    c = np.array([
        [[1, 0.5], [0.5, 1]],  
        [[1, 0.5], [0.5, 1]],
        [[1, 0.5], [0.5, 1]]
    ], dtype=np.float64)

    d = np.array([
        [1, 2], 
        [2, 3], 
        [3, 4], 
        [4, 5]
    ], dtype=np.float64)
    
    x = solve_block_tridiagonal(a, b, c, d)

    # Construct the equivalent full matrix A_full and right-hand side d_full
    A_full = np.block([
        [b[0], c[0], np.zeros((2, 2)), np.zeros((2, 2))],
        [a[0], b[1], c[1], np.zeros((2, 2))],
        [np.zeros((2, 2)), a[1], b[2], c[2]],
        [np.zeros((2, 2)), np.zeros((2, 2)), a[2], b[3]]
    ])
    
    d_full = d.flatten()  # Flatten d for compatibility with the full system

    # Solve using numpy's direct solve for comparison
    x_np = np.linalg.solve(A_full, d_full).reshape((N, 2))
    # Print the solutions for comparison
    print("Solution x from block tridiagonal solver (TMDA):\n", x, "\nResidual:", torch.sum(torch.abs(torch.tensor(A_full)@torch.tensor(x).flatten() - torch.tensor(d).flatten())))
    print("Solution x from direct full matrix solver:\n", x_np, "\nResidual np:", torch.sum(torch.abs(torch.tensor(A_full)@torch.tensor(x_np).flatten() - torch.tensor(d).flatten())))
# Run the test function
test_block_tridiagonal_solver()

Answer 1

10^-2 likely means you have a mistake in your formulas, not just an accuracy problem. There is no simple and painless solution: you'll have to grab a cup of coffee, sit down, and figure out where the bug is.

Some bugs can be found simply by staring intently at the code and going through the algorithm in your head (or explaining it to a rubber duck), but often one is not that lucky. A more general strategy is the following:

• Find the simplest, most trivial case in which your code doesn't work. Does it work with 1x1 blocks? With 1x1 blocks and N=1? With N=2? With N=3? In a block upper or lower triangular case? With the right-hand side equal to [1,0,0,...,0]?

• Take that case, and add a lot of print statements (or use a debugger); compute the result of each intermediate value on the side, either in the REPL or by hand, to figure out where your code differs from the theory.

You will probably not find someone here that wishes to help you with this task: debugging someone else's code is not fun, especially if it is not well commented (what are a, b, c and d? Which one of those is the diagonal?). But learning to find these bugs is a useful learning experience.

Answer 2

There are a lot of reasons, why your custom implemented code may not reach the performance of the standard solvers.

Immediate Observations from your code: Repeated calculations and inplace modification of arrays. For Eg:b[i] - a[i-1] @ C_star[i-1] has been calculated multiple times and there are places where instead of making a copy, some of the modifications to arrays are made in place. These things combined with the lot of inverse operations results in accumulation of numerical errors, which becomes worse if the condition number of the problem increases.

One can try to make these smaller changes and improve the accuracy. However, we cannot reach the accuracy of the established solvers. Because, In General, The solvers that libraries like numpy use on the backend are highly sophisticated and use libraries like BLAS/LAPACK in the backend. These routines employ techniques such as scaling, pivoting and different factorisation techniques (depending on the nature of the matrix) to get the better solution. This is why these methods (which have been iteratively improved over decades) perform way better than the naive implementation of solvers.