三对角矩阵算法(TDMA)又名Thomas算法,使用Python和NumPy数组

时间:2012-01-04 19:46:20

标签: python arrays algorithm matlab matrix

我在MATLAB中找到了thomas算法或TDMA的实现。

function x = TDMAsolver(a,b,c,d)
    %a, b, c are the column vectors for the compressed tridiagonal matrix, d is the right vector
    n = length(b); % n is the number of rows

    % Modify the first-row coefficients
    c(1) = c(1) / b(1);    % Division by zero risk.
    d(1) = d(1) / b(1);    % Division by zero would imply a singular matrix.

    for i = 2:n-1
        temp = b(i) - a(i) * c(i-1);
        c(i) = c(i) / temp;
        d(i) = (d(i) - a(i) * d(i-1))/temp;
    end

    d(n) = (d(n) - a(n) * d(n-1))/( b(n) - a(n) * c(n-1));

    % Now back substitute.
    x(n) = d(n);
    for i = n-1:-1:1
        x(i) = d(i) - c(i) * x(i + 1);
    end
end

我需要在python中使用numpy数组,这是我在python中首次尝试算法。

import numpy

aa = (0.,8.,9.,3.,4.)
bb = (4.,5.,9.,4.,7.)
cc = (9.,4.,5.,7.,0.)
dd = (8.,4.,5.,9.,6.)

ary = numpy.array

a = ary(aa)
b = ary(bb)
c = ary(cc)
d = ary(dd)

n = len(b)## n is the number of rows

## Modify the first-row coefficients
c[0] = c[0]/ b[0]    ## risk of Division by zero.
d[0] = d[0]/ b[0]

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

d[-1] = (d[-1] - a[-1] * d[-2])/( b[-1] - a[-1] * c[-2])

## Now back substitute.
x = numpy.zeros(5)
x[-1] = d[-1]
for i in range(-2, -n-1, -1):
    x[i] = d[i] - c[i] * x[i + 1]

他们给出了不同的结果,所以我做错了什么?

4 个答案:

答案 0 :(得分:1)

两者之间至少有一个区别:

for i in range(1,n,1):
Python中的

从索引1迭代到最后一个索引n-1,而

for i = 2:n-1

从索引1(从零开始)迭代到(last-1)索引,因为Matlab有一个基于索引的索引。

答案 1 :(得分:1)

我做了这个,因为python的在线实现都没有实际工作。我已经针对内置矩阵求逆进行了测试,结果匹配。

这里a =下诊断,b =主诊断,c =上诊断,d =解决方案载体

import numpy as np

def TDMA(a,b,c,d):
    n = len(d)
    w= np.zeros(n-1,float)
    g= np.zeros(n, float)
    p = np.zeros(n,float)

    w[0] = c[0]/b[0]
    g[0] = d[0]/b[0]

    for i in range(1,n-1):
        w[i] = c[i]/(b[i] - a[i-1]*w[i-1])
    for i in range(1,n):
        g[i] = (d[i] - a[i-1]*g[i-1])/(b[i] - a[i-1]*w[i-1])
    p[n-1] = g[n-1]
    for i in range(n-1,0,-1):
        p[i-1] = g[i-1] - w[i-1]*p[i]
    return p

为了便于大型矩阵的性能提升,请使用numba!这段代码在我的测试中胜过np.linalg.inv():

import numpy as np
from numba import jit    

@jit
def TDMA(a,b,c,d):
    n = len(d)
    w= np.zeros(n-1,float)
    g= np.zeros(n, float)
    p = np.zeros(n,float)

    w[0] = c[0]/b[0]
    g[0] = d[0]/b[0]

    for i in range(1,n-1):
        w[i] = c[i]/(b[i] - a[i-1]*w[i-1])
    for i in range(1,n):
        g[i] = (d[i] - a[i-1]*g[i-1])/(b[i] - a[i-1]*w[i-1])
    p[n-1] = g[n-1]
    for i in range(n-1,0,-1):
        p[i-1] = g[i-1] - w[i-1]*p[i]
    return p

答案 2 :(得分:0)

在你的循环中,Matlab版本迭代第二个到第二个到最后一个元素。要在Python中执行相同的操作,您需要:

for i in range(1,n-1):

(正如voithos的评论中所述,这是因为范围函数排除了最后一个索引,因此除了更改为0索引之外,还需要对此进行更正。)

答案 3 :(得分:0)

用python写这样的东西会很慢。你最好使用 LAPACK 来完成数值繁重的工作,并使用 python 处理它周围的一切。 LAPACK 是经过编译的,因此它的运行速度比 Python 快得多,而且它的优化程度也比我们大多数人可以匹配的要高得多。

SciPY 为 LAPACK 提供了低级包装器,以便您可以非常简单地从 python 中调用它,您可以在这里找到您正在寻找的那个:

https://docs.scipy.org/doc/scipy/reference/generated/scipy.linalg.lapack.dgtsv.html#scipy.linalg.lapack.dgtsv