Implementing tridiagonal matrix algorithm (TDMA) with NumPy

Question

Hi.

I'm implementing TDMA in Python using NumPy. The tridiagonal matrix is stored in three arrays:

a = array([...])
b = array([...])
c = array([...])

I'd like to calculate alpha-coefficients efficiently. The algorithm is as follows:

# n = size of the given matrix - 1
alpha = zeros(n)
alpha[0] = b[0] / c[0]
for i in range(n-1):
    alpha[i+1] = b[i] / (c[i] - a[i] * alpha[i])

However, this is not efficient because of Python's for loop. Want I want is something like this approach:

# n = size of the given matrix - 1
alpha = zeros(n)
alpha[0] = b[0] / c[0]
alpha[1:] = b[1:] / (c[1:] - a[1:] * alpha[:-1])

In this latter case the result is incorrect because NumPy stores the right part of the last expression in a temprorary array and then assigns references to its elements to alpha[1:]. Therefore a[1:] * alpha[:-1] is just an array of zeros.

Is there a way to tell NumPy to use values of alpha calculated on previous steps within its internal loop?

Thanks.

Answer 1

If its tridiagonal systems you want to solve there is solve_banded() in numpy.linalg. Not sure if that's what you're looking for.

Answer 2

Apparently, there is no way to do this in Python without using C or its pythonic variations.