三重対角行列の解法について(scipy.linalg.solve_banded)

三重対角行列については、scipy.linalg.solve_bandedを用いることで高速に解くことができる。GitHubにサンプルソースあり。

三重対角行列は以下のようのもの

A=
\begin{bmatrix}
a_{11} & a_{12}&0&0&0\\
a_{21} & a_{22}&a_{22}&0&0\\
0 & a_{32}&a_{33}&a_{34}&0\\
0& 0&a_{43}&a_{44}&a_{45}\\
0& 0&0&a_{54}&a_{55}\\
\end{bmatrix}

つまり、対角成分とその上下以外0となる行列。ちなみに、英語ではbanded matrixという。
上の行列で$Ax=b$を解きたい場合、scipy.linalg.solve_bandedの使う前に、次の行列を準備する。

\tilde{A}=
\begin{bmatrix}
0&a_{12}&a_{22}&a_{34}&a_{45}\\
a_{11}&a_{22}&a_{33}&a_{44}&a_{55}\\
a_{21}&a_{32}&a_{43}&a_{54}&0
\end{bmatrix}

つまり、三重対角成分を並べた行列である。 scipy.linalg.solve_bandedを用いて、$Ax=b$を解くには、

scipy.linalg.solve_banded((1,1),$\tilde{A}$,b)

としてやればいい。以下実例を示すと、

from scipy.linalg import solve_banded
import numpy as np
a = np.array([[5,2, 0, 0, 0],
              [1,4, 2, 0, 0],
              [0,1, 3, 2, 0],
              [0,0, 1, 2, 2],
              [0,0, 0, 1, 1]])
ab = np.array([[0,  2, 2, 2, 2],
               [5,  4,  3,  2,  1],
               [1,  1,  1,  1,  0]])
b = np.array([0, 1, 2, 2, 3])
x = solve_banded((1, 1), ab, b)
print(x)
np.dot(a,x)

これの実行結果は、

[-1.    2.5  -4.    5.75 -2.75]
array([0., 1., 2., 2., 3.])

となる。