数学とか語学とか楽しいよね

フランス語、ドイツ語、ロシア語、アラビア語、オランダ語、英語、スペイン語、ラテン語とか数学とか数値計算(有限要素法、有限体積法、差分法、格子ボルツマン法、数理最適化、C++コード付き)とか勉強したことをまとめます。右のカテゴリーから興味のある記事を探してください。最近はクラシックの名演も紹介しています。noteにも書いています。https://note.mu/kumag_min

【連立一次方程式】TDMA(三重対角行列アルゴリズム、Tri-Diagonal Matrix Algorithm)C++コード付き

今回は連立一次方程式の係数が三重対角行列のときにとても効率の良い、TDMA(三重対角行列アルゴリズム、Tri-Diagonal Matrix Algorithm)のC++コードを公開します。TDMAは直接法に分類される手法で、流体計算の際のソルバーとして用いられています。理論は

http://penguinitis.g1.xrea.com/study/note/tdma.pdf

がわかりやすく参考になります。方程式が3つの例がわざわざ示されており、とてもわかりやすいです。アルゴリズムの流れを書きます。連立一次方程式を

 
\left( \begin{array}{ccc} 
a_1 & b_1 &   &   &   & 0  \\
c_2 & a_2 & b_2 &   &   &   &   \\
  & c_3 & a_3 & b_3 &   &   &   \\
  &   & \ddots & \ddots & \ddots &  \\
  &   &   & c_{n-1}   &  a_{n-1} &  b_{n-1} \\
0 &   &   &   &  a_{n} &  b_{n} \\
\end{array} \right)

\left( \begin{array}{c} x_1 \\ x_2 \\ x_3 \\ \vdots \\ x_{n-1} \\ x_n \end{array} \right)
 = 
\left( \begin{array}{c} d_1 \\ d_2 \\ d_3 \\ \vdots \\ d_{n-1} \\ d_n \end{array} \right)

とします。このときTDMAのアルゴリズム


1.  P_1=-\frac{b_1}{a_1} Q_1=\frac{d_1}{a_1} と定める
2.  P_i=\frac{-b_i}{a_i+c_i P_{i-1}} Q_i=\frac{d_i-c_i Q_{i-1}}{a_i+c_i P_{i-1}} を使って  i=2 から  i=n までの  P_i Q_i を求める
3.  x_n=Q_n とする
4.  x_i=P_ix_{i+1}+Q_i を使って  i=n-1 から  i=1 までの  x_i を求める(後ろ向きであることに注意!)


となります。

今回解く連立一次方程式は

 Ax=d

とあらわされます。ここで、  A 4 \times 4 の行列、 x d 4 次元のベクトルです。 A d が既知で、 x が未知です。具体的には

 \left( \begin{array}{ccc} 2 & 1 & 0 & 0 \\ 1 & 2 & 1 & 0 \\ 0 & 1 & 2 & 1 \\ 0 & 0 & 1 & 2 \end{array} \right) \left( \begin{array}{c} x_1 \\ x_2 \\ x_3 \\ x_4 \end{array} \right) = \left( \begin{array}{c} 4 \\ 8 \\ 12 \\ 11 \end{array} \right)

であり、これを解くと  x_1=1, x_2=2, x_3=3, x_4=4 となります。

コード上ではそれぞれ対角要素を配列 t[i] に、対角右の要素を配列 r[i] に、対角左の要素を配列 l[i] に収納しています。ではコードを以下に示します。

#include <iostream>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <vector>

using namespace std;

//TDMA//
inline void tdma(double a[], double c[], double b[], double d[], double x[], int size)
{
	int i;
	double *P=new double[size];
	double *Q=new double[size];
	
	//first step//
	P[0]=-b[0]/a[0];
	Q[0]=d[0]/a[0];
	
	//second step//
	for(i=1;i<size;i++)
	{
		P[i]=-b[i]/(a[i]+c[i]*P[i-1]);
		Q[i]=(d[i]-c[i]*Q[i-1])/(a[i]+c[i]*P[i-1]);
	}
	
	//third step//
	x[size-1]=Q[size-1];

	//fourth step, backward//
	for(i=size-2;i>-1;i=i-1)
	{
		x[i]=P[i]*x[i+1]+Q[i];
	}
	
	delete [] P,Q;
}

int main()
{
	int i;
	int size=4;
	
	double *t=new double[size];
	double *l=new double[size];
	double *r=new double[size];
	double *b=new double[size];
	double *x=new double[size];
	
	t[0]=2.0;t[1]=2.0;t[2]=2.0;t[3]=2.0;
	l[0]=0.0;l[1]=1.0;l[2]=1.0;l[3]=1.0;
	r[0]=1.0;r[1]=1.0;r[2]=1.0;r[3]=0.0;
	b[0]=4.0;b[1]=8.0;b[2]=12.0;b[3]=11.0;
	
	
	tdma(t,l,r,b,x,size);
	
	for(i=0;i<size;i++) cout<<x[i]<<endl;
	
	delete [] t,l,r,b,x;

    return 0;
}