解三对角线性方程组的追赶法
追赶法在解三对角矩阵时非常实用,效率极高,特别对在隐式求解双曲型或者抛物型偏微分方程的时候。算法具体讲解如图所示
import matplotlib
import math
matplotlib.use('TkAgg')
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator, FormatStrFormatter
def catchup(n, a, b, c, f): #输入三个对角线上的值a,b,c,f是等号右边的向量
m = np.zeros(f.size, dtype=float)#中间变量
x = np.zeros(f.size, dtype=float)
for i in range(1, n): #正序,追
m[i] = a[i]/b[i-1]
b[i] = np.copy(b[i]) -m[i]*c[i-1]
f[i] = np.copy(f[i]) -m[i]*f[i-1]
print(i, m[i], f[i-1], f[i])
x[-1] = f[-1]/b[-1]
for j in range (n-2, -1, -1):#倒序,赶
x[j] = (f[j]-c[j]*x[j+1])/b[j]
return x
def get_tridiagonal(A, F):# 获取三对角矩阵的对角线,其中主对角线是b,左对角线是a,右对角线是c
n=F.size
b = np.zeros(n,dtype=float)
a=np.zeros(n,dtype=float)
c=np.zeros(n,dtype=float)
for i in range(0,n):
b[i] = A[i,i]# 主对角线
if i >0:
a[i] = A[i,i-1]
if i<(n-1):
c[i] = A[i,i+1]
return n,a, b,c
#下面是例子,方程为Ax = F
A = np.array([[2.,-1.,0.,0.,0.],[-1.,2.,-1.,0.,0.],[0.,-1.,2.,-1.,0.],[0.,0.,-1.,2.,-1.],[0.,0.,0.,-1.,2.]] )
F=np.array([1.,0.,0.,0.,0.])
n,a,b,c = get_tridiagonal(A,F)
print(b)
print(a)
print(c)
result = catchup(n,a,b,c,F)
print(result)