共役勾配法で連立方程式を解く

#はじめに
この記事は、「共役勾配法」を元に実際に共役勾配法を動かしてみよう、というものである。共役勾配法の詳しい原理については「共役勾配法」を参照していただきたい。

#実装の準備
共役勾配法を使うには、$K$ステップ目のある地点$x^{(K)}$における接ベクトル$t^{(K)}$、共役勾配$\boldsymbol{m}^{(K)}$と勾配$\nabla f^{(K)}$、ヘッセ行列$\boldsymbol{H}^{(K)}$を用いて作られる

\begin{align*}
\boldsymbol{m}^{(K)} &= \nabla f^{(K)}-\dfrac{(\boldsymbol{m}^{(K-1)},\boldsymbol{H}^{(K)}\nabla f^{(K)})}{(\boldsymbol{m}^{(K-1)},\boldsymbol{H}^{(K)}\boldsymbol{m}^{(K-1)})}m^{(K-1)}\\
\boldsymbol{m}^{(0)} &= \nabla f^{(0)} \qquad\qquad\qquad\qquad\qquad\qquad(1)\\
t^{(K)}&=-\dfrac{(\boldsymbol{m}^{(K)},\nabla f^{(K)})}{(\boldsymbol{m}^{(K)},\boldsymbol{H}^{(K)}\boldsymbol{m}^{(K)})}\\
\boldsymbol{x}^{(K+1)} &= \boldsymbol{x}^{(K)}+t^{(K)}\boldsymbol{m}^{(K)}
\end{align*}

をfor文等で回せばよいものであった。
今回は、次の問題を共役勾配法で解くことにしよう。

次の連立1次方程式を共役勾配法を用いて計算せよ

\boldsymbol{Ax}=\boldsymbol{b}

$\boldsymbol{Ax}=\boldsymbol{b}$の解を見つけることは、$\boldsymbol{Ax}-\boldsymbol{b}=0$を満たす$\boldsymbol{x}$を見つけることと同値である。すなわち、次の関数が最小となるような$\boldsymbol{x}$を共役勾配法で解けば良いということである。

f(\boldsymbol{x})=||\boldsymbol{Ax}-\boldsymbol{b}||^2

これを変形していくと

\begin{align*}
f(\boldsymbol{x})&=(\boldsymbol{Ax}-\boldsymbol{b},\boldsymbol{Ax}-\boldsymbol{b})\\
&=(\boldsymbol{Ax},\boldsymbol{Ax})-2(\boldsymbol{A}^T\boldsymbol{b},\boldsymbol{x})+(\boldsymbol{b},\boldsymbol{b})\\
&=(x,\boldsymbol{A}^T\boldsymbol{A}\boldsymbol{x})-2(\boldsymbol{A}^T\boldsymbol{b},\boldsymbol{x})+(\boldsymbol{b},\boldsymbol{b})
\end{align*}

ここから、

\begin{align*}
\nabla f &= \boldsymbol{A}^T(\boldsymbol{Ax}-\boldsymbol{b})\\
\boldsymbol{H} &= \boldsymbol{A}^T\boldsymbol{A}
\end{align*}

を得る。ただし、ベクトルの微分の公式

\nabla (\boldsymbol{a},\boldsymbol{x})=\boldsymbol{a}

\nabla (\boldsymbol{x},\boldsymbol{Ax})=2\boldsymbol{Ax}

を使って計算した。では、この勾配とヘッセ行列を式(1)に入れると
[メインの式]

\begin{align*}
\alpha^{(K)}&=\dfrac{(\boldsymbol{m}^{(K-1)},\boldsymbol{A}^T\boldsymbol{A}\boldsymbol{A}^T(\boldsymbol{Ax}-\boldsymbol{b}))}{(\boldsymbol{m}^{(K-1)},\boldsymbol{A}^T\boldsymbol{A}\boldsymbol{m}^{(K-1)})}\\
\boldsymbol{m}^{(K)} &= \boldsymbol{A}^T(\boldsymbol{Ax}-\boldsymbol{b})-\alpha^{(K)}\boldsymbol{m}^{(K-1)}\qquad\qquad(2)\\
t^{(K)}&=-\dfrac{(\boldsymbol{m}^{(K)},\boldsymbol{A}^T(\boldsymbol{Ax}-\boldsymbol{b}))}{(\boldsymbol{m}^{(K)},\boldsymbol{A}^T\boldsymbol{A}\boldsymbol{m}^{(K)})}\\
\boldsymbol{x}^{(K+1)} &= \boldsymbol{x}^{(K)}+t^{(K)}\boldsymbol{m}^{(K)}
\end{align*}

となる。ここでは$\alpha,x$初期条件は0とすることによって初期値は次のように計算する。
[初期条件、初期値]

\begin{align*}
\alpha^{(0)}&=0\\
x^{(0)}&=0\\
m^{(0)}&=\boldsymbol{A}^T(\boldsymbol{Ax}-\boldsymbol{b})\\
t^{(K)}&=-\dfrac{(\boldsymbol{m}^{(0)},\boldsymbol{A}^T(\boldsymbol{Ax}-\boldsymbol{b}))}{(\boldsymbol{m}^{(0)},\boldsymbol{A}^T\boldsymbol{A}\boldsymbol{m}^{(0)})}\\
\boldsymbol{x}^{(1)} &= \boldsymbol{x}^{(K)}+t^{(0)}\boldsymbol{m}^{(0)}
\end{align*}

#いざ、実装
共役勾配法で解く綺麗な方法はいくらかあるが、今回は初めて共役勾配法を学習し、それをそのままPython3で実装することに主眼を置いている。したがって、若干見辛い点や改善点はあると思われるが、ご容赦願いたい。使うライブラリはnumpyのみとする。

import numpy as np

まず、内積を求める関数を作ろう。ベクトルの内積の時はnp.dotが使えたが、今回は行列の内積を計算したいので新しく関数を立てることにしよう。

naiseki

def naiseki(a,b):
    ans = np.sum(np.multiply(a,b))
    return ans

さて、上でたてた式(2)をそのまま実装してみよう。

CG

def solve_CG(A,b,k):
    alpha = 0
    x = np.array([[0],[0],[0]])
    m = A.T * (A*x-b)
    t = -(naiseki(m,A.T*(A*x-b)))/(naiseki(m,A.T*A*m))
    x = x + t*m
    
    for i in range(k):
        alpha = - (naiseki(m,A.T * A * A.T*(A*x-b)))/(naiseki(m, A.T*A*m))
        m = A.T * (A*x-b) + alpha*m
        t = -(naiseki(m,A.T*(A*x-b)))/(naiseki(m,A.T*A*m))
        x = x + t*m
        
    return x

#テスト
本当にできてるか、簡単な連立方程式を与えてみよう。

A=\left(\begin{array}{ccc}1&0&0\\0&2&0\\0&0&1\end{array}\right),b=\left(\begin{array}{c}4\\5\\6\end{array}\right)

とするとき、容易に$x=4,y=2.5,z=6$となることがわかるが、ちゃんとできているだろうか。

A = np.matrix([[1,0,0],[0,2,0],[0,0,1]])
b = np.array([[4],[5],[6]])
solve_CG(A,b,2)

さて、その出力は...

matrix([[4. ],
        [2.5],
        [6. ]])

やったぜ！！！！できた！！！

#追記(2018/5/15)
行列の内積について「np.dot」で計算できないため新しい関数を立てましたが、実は「np.tensordot」で計算できるという助言をいただきましたので、補足しておきます。

def naiseki(a,b):
    ans = np.tensordot(a,b)
    return ans

とすればそのまま関数「solve_CG」を活用することができます。もしくは、「solve_CG」中の「naiseki」を「np.tensordot」に置換することによって関数１つで解くこともできます。

def solve_CG(A,b,k):
    alpha = 0
    x = np.array([[0],[0],[0]])
    m = A.T * (A*x-b)
    t = -(np.tensordot(m,A.T*(A*x-b)))/(np.tensordot(m,A.T*A*m))
    x = x + t*m
    
    for i in range(k):
        alpha = - (np.tensordot(m,A.T * A * A.T*(A*x-b)))/(np.tensordot(m, A.T*A*m))
        m = A.T * (A*x-b) + alpha*m
        t = -(np.tensordot(m,A.T*(A*x-b)))/(np.tensordot(m,A.T*A*m))
        x = x + t*m
        
    return x

実はこの記事は学部２年生〜４年生合同の自主ゼミ用に書いた記事であり、Pythonを知らない学部2年生にも分かるようにとしてnaiseki関数を作成しました。また、他の参考資料を見たときに「行列間の内積は直接計算不可」という記述を目にしていた、ということもありnaiseki関数を立てたという背景もあります。初めて「np.tensordot」を知るきっかけにもなったので、追記として残しておきます
#さいごに
この実装は本当に単純なものであり、もう少し色々改良の余地がありそう。ただ、得られた数式をとりあえず実装してみる第１歩として捉えていただければ幸いです。