Implementing tridiagonal matrix algorithm (TDMA) with NumPy

Question

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.

this isn't possible. see stackoverflow.com/questions/1587367/… for a similar issue. if you really need the speed increase, try cython. — Autoplectic
– Autoplectic, Commented Dec 18, 2009 at 20:17

dwf · Accepted Answer · 2009-12-18 22:28:14Z

2

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.

answered Dec 18, 2009 at 22:28

dwf

3,5732 gold badges22 silver badges25 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

fedvasu Over a year ago

i don't think there is any provision for banded matrix solving , in numpy or even in scipy , there is no sp.sparse ,as far as i know.

fedvasu Over a year ago

I think it's my bad , scipy.sparse exists(it doesn't exist on my system , it existsin docs so it must actually exist) , but not solve_banded() in numpy.linalg , may have been when present in 2009 , not now.

Alexei Sholik · Accepted Answer · 2010-01-10 21:03:19Z

2

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

answered Jan 10, 2010 at 21:03

community wiki

Alexei Sholik

Collectives™ on Stack Overflow

Implementing tridiagonal matrix algorithm (TDMA) with NumPy

2 Answers 2

2 Comments

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

2 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related