ずいぶん間が空きましたが前回の記事の続きです。今回は変分法を用いたデータ同化手法について解説します。例のごとく過去記事の内容は仮定するので、準備とETKFの記事を読んでから来てください。過去記事と違ってAtomのJulia開発環境がVSCode に移行してるのでこの記事のコードも全部VSCodeのJulia拡張を導入して動かしています。

3次元変分法は今まで説明したKalman Filterと同じような考え方をするのですが、4次元変分法はちょっと発展的な考え方をするので、すこしそのあたりのお話をします。前回までのKalman Filterの手法は同じ時刻におけるデータから推定値を作っていたのでした。たとえば、$x_{i}^{f}$は予報値、$y_{i},y_{i+1},y_{i+2}$を観測データとすると、$y_{i+1},y_{i+2}$は使わずに推定値$x_{i}^{a}$を作ったのでした。$y_{i+1},y_{i+2}$も$y_{i}$からそうずれてはいないでしょうから、ちょっともったいないですね。お絵描きするとこんな感じです。

もったいないので$y_{i+1},y_{i+2}$も使いたいのですが、このままでは時刻がずれているのでどうしようもありません。どうすればいいでしょうか。モデル$M$の出番ですね。$x_{i}^{a}$をモデルで時間発展させると時刻をそろえることができます。お絵描きするとこんな感じです。

時刻をそろえることができたので、$M(x_{i}^{a})$と$y_{i+1}$のずれ具合、$M(M(x_{i}^{a}))$と$y_{i+2}$のずれ具合まで考慮して$x_{i}^{a}$を決めることができます。今回説明する4次元変分法はこういう考え方に基づきます。Kalman Filterにもこういう考え方をする手法があり、スムーザーと呼ばれています。

変分法もKalman Filterも基本的な考え方は一緒なのですが、推定値$x_{i}^{a}$をどういう原理に基づいて決定するかだけが異なります。この記事ではそのあたりの話もしていきたいと思います。

変分法の手続き

Forecast Stepではモデル$M$により前ステップの解析値$x_{i-1}^{a}$を時間発展させます。

    x_{i}^{f} = M(x_{i-1}^{a})

Analysis Stepでは予報値$x_{i}^{f}$と観測値から定まる目的関数$\mathcal{J}(x)$を最小化して解析値$x_{i}^{a}$を得ます。

    x_{i}^{a} = \mathrm{argmin}_{x\in \mathbb{R}^{N}} \mathcal{J}(x)

あとはこれを交互に繰り返すだけで、見かけはとてもシンプルです。

目的関数$\mathcal{J}$は適当に事前に定めた$T\geq 0$（Time Windowの大きさ）に対して次のようになります。

       \mathcal{J}(x) = \frac{1}{2}(B_{i}^{-1}(x-x_{i}^{f}), x-x_{i}^{f}) + \sum_{j=0}^{T} \frac{1}{2} (R_{i+j}^{-1}(H_{i+j}(M^{j}(x))-y_{i+j}),H_{i+j}(M^{j}(x))-y_{i+j})

ここで$B_{i}$は以前までの記事でいう$P_{i}^{f}$に相当し、$M^{j}$はモデル$M$の$j$回合成を表します。そのほかの記号は以前までの記事で出てきたものと同じです。特に$T=0$のときは

       \mathcal{J}(x) = \frac{1}{2}(B_{i}^{-1}(x-x_{i}^{f}), x-x_{i}^{f}) + \frac{1}{2} (R_{i}^{-1}(H_{i}(x)-y_{i}),H_{i}(x)-y_{i})

となり、時間方向にずれたデータが出てこないので、3次元変分法といいます。3次元変分法と対比して$T>0$のときを4次元変分法といいます。

上で定義したもの以外にも、事前分布と呼ばれる$x$についての関数$p_{x_{i}^{t}}(x)$を導入して次のように定めることもあります。

\mathcal{J}(x) = \frac{1}{2}(B_{i}^{-1}(x-x_{i}^{f}), x-x_{i}^{f}) + \sum_{j=0}^{T} \frac{1}{2} (R_{i+j}^{-1}(H_{i+j}(M^{j}(x))-y_{i+j}),H_{i+j}(M^{j}(x))-y_{i+j})  - \mathrm{log} p_{x_{i}^{t}}(x)

これ以外にも考えている設定に応じて正則化項やさまざまな項を付け加えることもあるのですが、その辺りについてはデータ同化　観測・実験とモデルを融合するイノベーションを参照してください。あるいは機械学習の本（状態空間モデルについての記述があるとよい）に載ってるかもしれません。

目的関数の導出

統計と測度論についての知識を仮定するので、とにかく手法について知りたい場合は飛ばしてください。

いきなり$\mathcal{J}$を定義してなぜそうなったのかの説明をしないのも味気ないので、手法の導出について書くことにします。本節の内容は数理統計学吉田朋広とかに載ってるので省略した証明などが気になる人はそちらを参照してください。

まず一般論として最小分散推定について説明します。$(\Omega, \mathcal{F}, P)$を確率空間とし、$X \colon \Omega \to \mathbb{R}^{N}$を可積分確率変数とします。$Y \colon \Omega \to \mathbb{R}^{d}$に対して$Y$が生成する$\sigma$-代数を$\sigma(Y)$と書きます。このとき、条件付き期待値$E[X|\sigma(Y)]$のことを$Y$の下での$X$の条件付き期待値といい、単に$E[X|Y]$と書きます。この条件付き期待値は$\Omega$上の関数なので数学的には取り扱いがしやすいです。特に二乗可積分なら「$Y$の関数として書ける関数の中で$X$との$L^{2}$ノルム誤差が最小である」という良い性質が成り立つので、$X$が真値、$Y$が観測データであるとしたとき、推定値として$E[X|Y]$を用いる推定手法を最小分散推定と言い、これまで解説してきたカルマンフィルターは最小分散推定です。推定するときにわかっているのは$\omega \in \Omega$ではなく$Y(\omega) \in \mathbb{R}^{d}$であるという事情があってこのままではちょっと使いづらいので、像測度を利用して変形します。$Y$による$\mathbb{R}^{d}$上の像測度を$P^{Y}$と書きます。$\mu(E) = \int_{Y^{-1}(E)} X(\omega) P(d\omega)$として$\mu \colon \mathcal{B}(\mathbb{R}^{d}) \to \mathbb{R}^{N}$を定義すると$P^{Y}(E)=0$ならば$\mu(E)=0$が成り立ちます。したがって、Radon-Nikodymの定理から$\mu(E) = \int_{E} g(y) P^{Y}(dy)$が成り立つような$g \colon \mathbb{R}^{d} \to \mathbb{R}^{N}$が存在します。$g(y)$を$Y=y$が与えられた下での条件付き期待値といい、$E[X|Y=y]$と書きます。すると$P$-零集合を除いて$E[X|Y] = E[X|Y=y]|_{y=Y(\omega)}$が成り立つので、$Y(\omega)$が見えていれば$E[X|Y=y]$に代入することで条件付き期待値を求めることができます。あとは$E[X|y]$を求めればよいです。$A \in \mathcal{F}$に対して$P_{Y=y}(A) =E[\mathbb{1}_{A}|Y=y]$によって$P_{Y=y}$を定めると$P_{Y=y}$は確率測度の条件を満たすので、これを$Y=y$の下での$P$の条件付き確率測度¹と言います。$X$による$P_{Y=y}$の像測度$P_{Y=y}^{X}$が密度関数を持つとき、この密度関数を$Y=y$の下での$X$の条件付き密度関数といい、$p_{X|Y=y}$と書きます。さて、$X,Y,(X,Y)$が密度関数を持つと仮定し、それぞれ$p_{X},p_{Y},p_{X,Y}$と書くとき、$Y=y$の下での条件付き密度関数$p_{X|Y=y}$は次のように書けることが知られています。

    p_{X|Y=y}(x) = \frac{p_{X,Y}(x,y)}{p_{Y}(y)}

$p_{Y}(y)=0$の時は$p_{X|Y=y}(x)=0$とします。$p_{X|Y=y}$を用いると

    E[X|Y=y] = \int_{\mathbb{R}^{d}} x p_{X|Y=y}(x) dx

が成り立つので、$p_{X|Y=y}(x)$を先述の定義に従って計算したうえで、右辺の積分を計算すれば$E[X|Y=y]$を求めることができます。

以上の一般論をデータ同化の枠組みに対して適用します。表記を単純にするために時刻に関する添え字$i$は書かないことにします。真値$X \colon \Omega \to \mathbb{R}^{N}$、予報$X^{f} \colon \Omega \to \mathbb{R}^{N}$、観測$Y_{j}^{o} \colon \Omega \to \mathbb{R}^{d}, j=0,1,\ldots,T$を確率変数とし、その像の空間の元を表す変数をそれぞれ$x\in \mathbb{R}^{N}, x^{f}\in \mathbb{R}^{N}, y_{0}^{o}\in \mathbb{R}^{d}, \ldots, y_{T}^{o} \in \mathbb{R}^{d}$と書くことにします。$X$としてそのまま$X$、$Y$として$(X^{f}, Y_{0}^{o}, Y_{1}^{o},\ldots, Y_{T}^{o})$を用いて以上の一般論を適用すれば、$X,Y,(X,Y)$が密度関数を持つという仮定の下で、$Y=(x^{f}, y_{0}^{o}, y_{1}^{o}, \ldots, y_{T}^{o})$が与えられた下での$X$の条件付き期待値は次のように表すことができます。

    E[X | Y=(x^{f}, y_{0}^{o}, y_{1}^{o}, \ldots, y_{T}^{o}) ]
    = \int_{\mathbb{R}^{N}} x p_{X | Y=(x^{f}, y_{0}^{o}, y_{1}^{o}, \ldots, y_{T}^{o})}(x) dx

$X^{f},Y_{0},\ldots,Y_{T}$が独立な確率変数であり、$X=x$が与えられた下での条件付き密度関数がそれぞれ
$p_{X^{f}|X=x},p_{Y_{0}|X=x},\ldots,p_{Y_{T}|X=x}$によって与えられることを仮定すれば

p_{X^{f},Y_{0}^{o},Y_{1}^{o},\ldots,Y_{T}^{o}|X=x}(x^{f},y_{0}^{o},\ldots,y_{T}^{o}) = 
    p_{X^{f}|X=x}(x^{f}) \Pi_{j=0}^{T}p_{Y_{j}|X=x}(y_{j}^{o})

が成り立つので、

    p_{X | Y=(x^{f}, y_{0}^{o}, y_{1}^{o}, \ldots, y_{T}^{o})}(x)
= \frac{p_{X}(x) p_{X^{f}|X=x}(x^{f}) \Pi_{j=0}^{T}p_{Y_{j}|X=x}(y_{j})}{p_{Y}(x^{f}, y_{0}^{o}, y_{1}^{o}, \ldots, y_{T}^{o})}

として変形でき、

    E[X | Y=(x^{f}, y_{0}^{o}, y_{1}^{o}, \ldots, y_{T}^{o}) ]
    = \int_{\mathbb{R}^{N}} x \frac{p_{X}(x) p_{X^{f}|X=x}(x^{f}) \Pi_{j=0}^{T}p_{Y_{j}|X=x}(y_{j})}{p_{Y}(x^{f}, y_{0}^{o}, y_{1}^{o}, \ldots, y_{T}^{o})} dx

として条件付き期待値を表すことができます。この式の$x^{f}, y_{0}^{o}, y_{1}^{o}, \ldots, y_{T}^{o}$として手元にある予報値や観測値、すなわち確率変数の実現値$X^{f}(\omega),Y_{0}^{o}(\omega),\ldots,Y_{T}^{o}(\omega)$を代入すれば、最小分散推定としての推定値が得られます。登場するすべての確率変数は正規分布に従うとし、モデルや観測演算子は全て線形であるとすれば解析的に積分が求まり、$T=0$の時はカルマンフィルターの式になります。そのような仮定が成り立たない場合でも使える手法として、直接上記の積分を近似計算する手法（粒子法、マルコフ連鎖モンテカルロ、ハミルトニアンモンテカルロなど）もあるのですが、その解説はのちの機会に譲ります。さて、変分法では別の方法（最尤推定）で推定値$x^{a}$を定義します。

まず、ありとあらゆる真値の実現値$x \in \mathbb{R}^{N}$の中から手元にある予報値や観測値$X^{f}(\omega),Y_{0}^{o}(\omega),\ldots,Y_{T}^{o}(\omega)$が得られる確率
$p_{Y|X=x}(X^{f}(\omega),Y_{0}^{o}(\omega),\ldots,Y_{T}^{o}(\omega))$を最大化するものを推定値とする方法を考えます。

    x^{a} = \mathrm{argmax}_{x\in \mathbb{R}^{N}} p_{Y|X=x}(X^{f}(\omega),Y_{0}^{o}(\omega),\ldots,Y_{T}^{o}(\omega))
 = \mathrm{argmax}_{x\in \mathbb{R}^{N}} p_{X^{f}|X=x}(X^{f}(\omega)) \Pi_{j=0}^{T}p_{Y_{j}^{o}|X=x}(Y_{j}^{o}(\omega))

次に、ありとあらゆる真値の実現値$x \in \mathbb{R}^{N}$の中から手元にある予報値や観測値$(X^{f}(\omega),Y_{0}^{o}(\omega),\ldots,Y_{T}^{o}(\omega))$が得られたもとで真値が$x$となる確率
$p_{X|Y=(X^{f}(\omega),Y_{0}^{o}(\omega),\ldots,Y_{T}^{o}(\omega))} (x)$を最大化するものを推定値とする方法を考えます。

    x^{a} = \mathrm{argmax}_{x\in \mathbb{R}^{N}} p_{X|Y=(X^{f}(\omega),Y_{0}^{o}(\omega),\ldots,Y_{T}^{o}(\omega))} (x)
= \mathrm{argmax}_{x\in \mathbb{R}^{N}} \frac{p_{X}(x) p_{X^{f}|X=x}(X^{f}(\omega)) \Pi_{j=0}^{T}p_{Y_{j}^{o}|X=x}(Y_{j}^{o}(\omega))}{p_{Y}(X^{f}(\omega),Y_{0}^{o}(\omega),\ldots,Y_{T}^{o}(\omega))}

あとは分布の形を適切に仮定して計算すればよいです。特に正規分布の形を仮定すれば式が簡単になり、冒頭で定義した目的関数になります。$H_{j} \colon \mathbb{R}^{N} \to \mathbb{R}^{d}(j=0,1,\ldots,T)$を観測演算子（非線形でもよいです）、$M \colon \mathbb{R}^{N} \to \mathbb{R}^{N}$をモデルとし、$M^{j}(j=0,1,\ldots,T)$を$M$の$j$回合成（$j=0$のときは恒等写像）として、
$p_{X^{f}|X=x}(x^{f})$は平均$x$で分散$B$の正規分布に従い、$p_{Y_{j}^{o}|X=x}(y_{j}^{o})$は平均$H_{j}(M^{j}(x))$で分散$R_{j}$の正規分布に従うことを仮定します。すなわち、

    p_{X^{f}|X=x}(x^{f}) \propto \mathrm{exp}\left(-\frac{1}{2}(B^{-1}(x-x^{f}), x-x^{f})\right) \\
p_{Y_{j}^{o}|X=x}(y_{j}^{o}) \propto \mathrm{exp}\left(-\frac{1}{2} (R_{j}^{-1}(H_{j}(M^{j}(x))-y_{j}),H_{j}(M^{j}(x))-y_{j})\right)

という形を仮定すれば、$\mathrm{argmax}$は対数を取った関数の$-1$倍を最小化すれば求められるので、（$\mathrm{argmax}$には影響しないので）対数を取って定数になる部分を無視して、それぞれ

       \mathcal{J}(x) = \frac{1}{2}(B^{-1}(x-X^{f}(\omega)), x-X^{f}(\omega)) + \sum_{j=0}^{T} \frac{1}{2} (R_{j}^{-1}(H_{j}(M^{j}(x))-Y_{j}(\omega)),H(M^{j}(x))-Y_{j}(\omega)) \\
\mathcal{J}(x) = \frac{1}{2}(B^{-1}(x-X^{f}(\omega)), x-X^{f}(\omega)) + \sum_{j=0}^{T} \frac{1}{2} (R_{j}^{-1}(H_{j}(M^{j}(x))-Y_{j}(\omega)),H_{j}(M^{j}(x))-Y_{j}(\omega)) - \mathrm{log}p_{X}(x)

を最小化する元を推定値$x^{a}$とすればよいことになります。後者の$J(x)$に現れる$p_{X}(x)$は真値の分布ですが、実際にはわからないので適当な分布を仮定する必要があり、事前分布と呼ばれています。データ同化のテキストは大抵の場合前者の$\mathcal{J}(x)$を変分法が最小化する汎函数として定義しています。データ同化　観測・実験とモデルを融合するイノベーションでは後者の$\mathcal{J}(x)$を用いる方法を変分法QCと呼んでおり、観測データの品質を管理するように事前分布を設定しています。あとは設定に合わせて$X^{f}(\omega)$や$Y_{j}(\omega)$を改めて$x_{i}^{f},y_{i+j}^{o}$などと置きなおせば、冒頭で述べた目的関数$\mathcal{J}$が得られます。

最適化のための数値微分

前節で定義した関数を最小化すればよいわけですが、当然手計算では不可能です。今回はBFGS法と呼ばれる勾配法の一種を用いて最小化を行います。BFGS法では目的関数の勾配を用いる必要があります。$p_{X}$は自分で与える関数なので、これを除いた部分を考えます。また数値実験の設定に合わせて$H_{j}$は線形であるとします。このとき、$\mathcal{J}$を微分すると

    \nabla \mathcal{J}(x) = B^{-1}(x-x^{f}) + \sum_{j=0}^{T} J(M^{j})^{\top}(x)H_{j}^{T}R_{j}^{-1}(H_{j}M^{j}(x)-y_{j})

となります。ここで、右辺の$J(M^{j})^{\top}(x)$は、$x$における$M^{j}$のヤコビ行列$J(M^{j})(x)$の転置行列という意味です。右辺第二項を変形していきます。合成関数の微分を用いて、$1\leq i,j\leq N$に対して

    \frac{\partial}{\partial x_{j}} (M(M(x)))_{i} = \sum_{k=1}^{N} \frac{\partial M_{i}}{\partial x_{k}}(M(x)) \frac{\partial M_{k}}{\partial x_{j}} (x)

が成り立つので$J(M^{2})(x) = JM(M(x))JM(x)$が成り立ちます。表記を見やすくするために$M^{j}(x)=x_{j}(j=0,1,\ldots,T)$と書くことにする（$M^{0}(x)=x$と約束します）と、$J(M^{2})(x) = JM(x_{1})JM(x_{0})$が成り立ちます。このような計算を繰り返すことで、$j\geq 1$に対して

    J(M^{j})(x) = JM(x_{j-1})JM(x_{j-2})\ldots JM(x_{1})JM(x)

が成り立つことが分かるので、$d_{j}=H_{j}^{T}R_{j}^{-1}(H_{j}x_{j}-y_{j})$とおくと、$J$の微分は

    \nabla J(x) = B^{-1}(x-x^{f}) + \sum_{j=0}^{T} (JM(x_{j-1})JM(x_{j-2})\ldots JM(x_{1})JM(x))^{\top}(x)d_{j} \\
    = B^{-1}(x-x^{f}) + \sum_{j=0}^{T} JM^{\top}(x)JM^{\top}(x_{1}) \ldots JM^{\top} (x_{j-2}) JM^{\top} (x_{j-1}) d_{j}

と書き直せます。最後に第2項の和をもう少し変形します。

    \sum_{j=0}^{T} JM^{\top}(x)JM^{\top}(x_{1}) \ldots JM^{\top} (x_{j-2}) JM^{\top} (x_{j-1}) d_{j} \\
    = JM^{\top}(x)JM^{\top}(x_{1}) \ldots JM^{\top} (x_{T-2}) JM^{\top} (x_{T-1}) d_{T} \\
        +JM^{\top}(x)JM^{\top}(x_{1}) \ldots JM^{\top} (x_{T-3}) JM^{\top}(x_{T-2}) d_{T-1} \\
        +JM^{\top}(x)JM^{\top}(x_{1}) \ldots JM^{\top} (x_{T-4}) JM^{\top} (x_{T-3}) d_{T-2} \\
        + \ldots \\
        +JM^{\top}(x)d_{1} + d_{0}

として$\sum$記号を外して考えると、一行目と二行目はかなりの部分が共通しているので同類項でくくって残りを$e_{1}$とおき、次の行と同類項でくくって残りを$e_{2}$とおき、と繰り返していくととてもすっきりした式になります。具体的には、$e_{0}=d_{T}$としてから、$k=1,2,\ldots,T$に対して$e_{k} = JM^{\top}(x_{T-k})e_{k-1} + d_{T-k}$を計算してゆくと、$e_{T}$が求める和になります。以上を踏まえて

    \nabla \mathcal{J} (x) = B^{-1}(x-x^{f}) + e_{T}

が求める導関数になるので、これを求める関数を書けばよいです。道中で散々出てきた通り、ベクトル$x,y$に対して$M$のJacobi行列の転置との積$JM^{\top}(x)y$を計算する必要があります。Extended Kalman Filterのように差分で近似計算することも可能ですが、今回のように何度も計算する必要がある時にはとても大変です。精度と計算量の両方でより優れた、接線形コードと双対漸化式を利用するアルゴリズム（アジョイント法）が提案されているので、それを使います。以降ではかなり丁寧にその内容について説明するので、結論だけ見たい方は飛ばしてください。

接線形コード

$x,y \in \mathbb{R}^{N}$と$M \colon \mathbb{R}^{N} \to \mathbb{R}^{N}$に対して、$x$における$M$のJacobi行列$JM(x)$と$y$の積$JM(x)y$を返すコードを接線形コードと言います。$x$のことを基準場、$y$のことを摂動変数と呼ぶことが多いそうです。$JM(x)$と$y$の積を真面目に計算するとコストがかかるので、$\frac{d}{dt}M(x+ty)|_{t=0} = JM(x)y$を利用して計算します。実装は例を通じて説明します。

例１

単純な場合から始めます。モデル$M$とJulia言語での実装が次のように与えられているとき、その接線形コードを実装します。

M(x_{1},x_{2}) = (x_{1}^2+x_{2}^2, x_{1}^2-x_{2}^2) \\
JM(x) =  \begin{pmatrix}
            2x_{1} & 2x_{2} \\
            2x_{1} & -2x_{2}
        \end{pmatrix}

    function M(x)
        result = Array{Float64}(undef,2)
        result[1] = x[1]^2 + x[2]^2
        result[2] = x[1]^2 - x[2]^2
        return result
    end

$\frac{d}{dt}M(x+ty)$を計算して$t=0$とします。

    M(x+ty) = ((x_{1}+ty_{1})^{2}+(x_{2}+ty_{2})^{2}, (x_{1}-ty_{1})^{2}-(x_{2}-ty_{2})^{2}) \\
    \frac{d}{dt}M(x+ty) = 2( (x_{1}+ty_{1})y_{1} + (x_{2} + ty_{2})y_{2}, -(x_{1}-ty_{1})y_{1} + (x_{2} - ty_{2})y_{2}) \\
    \frac{d}{dt}M(x+ty)|_{t=0} = 2( x_{1}y_{1} + x_{2}y_{2}, -x_{1}y_{1} + x_{2}y_{2})

あとはこれを関数にして実装すれば完成です。この操作をすることはresult[1] = x[1]^2 + x[2]^2とresult[2] = x[1]^2 - x[2]^2に対して入力変数に関する全微分をとってdresult[1] = 2x[1]dx[1] + 2x[2]dx[2], dresult[2]=2x[1]dx[1] - 2x[2]dx[2]とし、$dx=(dx[1],dx[2])$を$y$に置き換えることに等価です。一行ずつ全微分を取って置き換えるだけでいいので慣れるとこちらの考え方の方が早いですが、プログラムが複雑だと途中で混乱する場合があります。

    #JM(x)dxを返す
    function TLM(x,dx)
        dresult = Array{Float64}(undef,2)
        dresult[1] = 2.0*x[1]*dx[1] + 2.0*x[2]*dx[2]
        dresult[2] = 2.0*x[1]*dx[1] - 2.0*x[2]*dx[2]
        return dresult
    end

例２

この例はデータ同化流体科学から引用しています。Burgers方程式の差分法による時間発展を$M$としています。

    (M(u))_{i} = u_{i}  - c_{1} u_{i}(u_{i+1} - u_{i-1}) + c_{1}|u_{i}|(u_{i+1} - 2u_{i} + u_{i-1}) + c_{2}(u_{i+1} - 2u_{i} + u_{i-1}) \quad i=1,2,\ldots ,N \\
c1 = \frac{\Delta t}{2\Delta x},\,c_{2}= \frac{\nu \Delta t}{ (\Delta x)^{2}},\, u_{0} = u_{left},\, u_{N+1} = u_{right}

    function M(u, dt, dx, nu=0.01, ul = -1.0, ur = 1.0)
        c1 = dt/(2.0*dx)
        c2 = nu*dt/dx^2
        N = length(u)
        result = Array{Float64}(undef,N)
        for i in 2:N-1
            result[i] = u[i] - c1*u[i]*(u[i+1]-u[i-1]) + c1*abs(u[i])*(u[i+1]-2*u[i]+u[i-1])
                        + c2*(u[i+1]-2*u[i]+u[i-1])
        end
        result[1] = u[1] - c1*u[1]*(u[2]-ul) + c1*abs(u[1])*(u[2]-2*u[1]+ul)
                        + c2*(u[2]-2*u[1]+ul)
        result[N] = u[N] - c1*u[N]*(ur-u[N-1]) + c1*abs(u[N])*(ur-2*u[N]+u[N-1])
                        + c2*(ur-2*u[N]+u[N-1])
        return result
    end

一気に複雑になりましたがやることは変わりません。$i=2,...,N-1$に対して計算をしてみます。

\begin{align*}
   (M(u+tv))_{i} = &u_{i} + tv_{i} - c_{1}(u_{i}+tv_{i})(u_{i+1} + tv_{i+1} - (u_{i-1} + tv_{i-1})) \\
                  &+ c_{1} |u_{i}+tv_{i}| (u_{i+1}+tv_{i+1} - 2(u_{i}+tv_{i}) + u_{i-1} + tv_{i-1}) \\
                  &+ c_{2}(u_{i+1} + tv_{i+1} - 2(u_{i}+tv_{i}) + u_{i-1} + tv_{i-1})
\end{align*}

\begin{align*}
   \frac{d}{dt} (M(u+tv))_{i} = & v_{i} - c_{1}v_{i}(u_{i+1} + tv_{i+1} - (u_{i-1} + tv_{i-1}))
                    + c_{1}(u_{i} + tv_{i})(v_{i+1} - v_{i-1}) \\
                  &+ c_{1} \frac{u_{i}+tv_{i}}{|u_{i}+tv_{i}|}v_{i} (u_{i+1}+tv_{i+1} - 2(u_{i}+tv_{i}) + u_{i-1} + tv_{i-1}) \\
&+ c_{1} |u_{i} + tv_{i} | (v_{i+1} - 2v_{i} + v_{i-1} ) \\
                  &+ c_{2}(v_{i+1} - 2v_{i} + v_{i-1})
\end{align*}

\begin{align*}
   \frac{d}{dt} (M(u+tv))_{i} |_{t=0} = & v_{i} - c_{1}v_{i}(u_{i+1} - u_{i-1}) + c_{1}u_{i} (v_{i+1} - v_{i-1}) \\
                  &+ c_{1} \frac{u_{i}}{|u_{i}|}v_{i} (u_{i+1} - 2u_{i} + u_{i-1}) \\
&+ c_{1} |u_{i} | (v_{i+1} - 2v_{i} + v_{i-1} ) \\
                  &+ c_{2}(v_{i+1} - 2v_{i} + v_{i-1})
\end{align*}

最後の式を関数にすれば完成です。$u_{i}/|u_{i}|$は別の関数に置き換えて$u_{i}=0$付近の振る舞いを制御した方が良いでしょう。$i=1,N$のときだけ少し問題で、$u[0]=ul,u[N]=ur$を定数だと思って$i=1,N$だけ別扱いして同じ計算をするか、あるいは$u,u+tv$が同じ境界条件を満たすためには$v$が両端で$0$でなければならないと解釈して$v[0]=v[N+1]=0$を代入するかのどちらかを採用してください。結果は同じで、コードは次のようになります。


    #本当はx=0で0を返すのが符号関数
    function sgn(x)
        return x>=0 ? 1 : -1
    end

    #JM(u)vを返す
    function TLM(u,v)
        c1 = dt/(2.0*dx)
        c2 = nu*dt/dx^2
        N = length(u)
        result = Array{Float64}(undef,N)
        for i in 2:N-1
            result[i] = v[i] - c1*v[i]*(u[i+1]-u[i-1]) 
                        + c1*u[i]*(v[i+1] - v[i-1])
                        + c1*sgn(u[i])*v[i]*(u[i+1]-2*u[i]+u[i-1])
                        + c1*abs(u[i])*(v[i+1] - 2*v[i] + v[i-1])
                        + c2*(v[i+1]-2*v[i]+v[i-1])
        end
        result[1] = v[1] - c1*v[1]*(u[2]-ul) 
                        + c1*u[1]*v[2]
                        + c1*sgn(u[1])*v[1]*(u[2]-2*u[1]+ul)
                        + c1*abs(u[1])*(v[2] - 2*v[1])
                        + c2*(v[2]-2*v[1])
        result[N] = v[N] - c1*v[N]*(ur-u[N-1]) 
                        + c1*u[N]*(-v[N-1])
                        + c1*sgn(u[N])*v[N]*(ur-2*u[N]+u[N-1])
                        + c1*abs(u[N])*(- 2*v[N] + v[N-1])
                        + c2*(-2*v[N]+v[N-1])
        return result
    end

もちろん1行ずつ全微分を取って処理しても構いません。

例３

本題のL96モデルの時間発展です。

##L96modelの右辺
function L96(u;F=8.0,N=40)
    f = fill(0.0, N)
    f[1] = (u[2]-u[N-1])u[N] - u[1] + F
    f[2] = (u[3]-u[N])u[1] - u[2] + F
    for k in 3:N-1
        f[k] = (u[k+1]-u[k-2])u[k-1] - u[k] + F
    end
    f[N] =  (u[1]-u[N-2])u[N-1] - u[N] + F

    return f
end

#4-4Runge-Kutta
function Model(u;dt=0.05)
    N = length(u)
    result = Array{Float64}(undef,N)
    s1 = L96(u)
    s2 = L96(u + s1*dt/2)
    s3 = L96(u + s2*dt/2)
    s4 = L96(u + s3*dt)
    result = u + (s1 + 2.0*s2 + 2.0*s3 + s4)*(dt/6.0)
    return result
end

別の関数に処理を一部投げており、中間変数があるためそこそこ複雑ですが、やることはこれまでと何も変わりません。

\begin{align*}
    &M(u+tv) = u + tv + \frac{dt}{6}(s_{1}(u+tv) + 2s_{2}(u+tv) + 2s_{3}(u+tv) + s_{4}(u+tv)) \\
    &s_{1}(u+tv) = L96(u+tv) \\
    &s_{2}(u+tv) = L96(u + tv + (dt/2)s_{1}(u+tv)) \\
    &s_{3}(u+tv) = L96(u+tv + (dt/2)s_{2}(u+tv)) \\
    &s_{4}(u+tv) = L96(u+tv + dt s_{3}(u+tv)) 
\end{align*}

\begin{align*}
    &\frac{d}{dt} M(u+tv) = v + \frac{dt}{6}\frac{d}{dt} (s_{1}(u+tv) + 2s_{2}(u+tv) + 2s_{3}(u+tv) + s_{4}(u+tv)) \\
    &\frac{d}{dt} s_{1}(u+tv) = JL96(u+tv)v \\
    &\frac{d}{dt} s_{2}(u+tv) = JL96(u + tv + (dt/2)s_{1}(u+tv)) (v + \frac{dt}{2}\frac{d}{dt} s_{1}(u+tv)) \\
    &\frac{d}{dt} s_{3}(u+tv) = JL96(u+tv + (dt/2)s_{2}(u+tv)) ( v + \frac{dt}{2} \frac{d}{dt} s_{2}(u+tv))\\
    &\frac{d}{dt} s_{4}(u+tv) = JL96(u+tv + dt s_{3}(u+tv)) (v + dt \frac{d}{dt} s_{3}(u+tv)) 
\end{align*}

\begin{align*}
    &\frac{d}{dt} M(u+tv) |_{t=0} = v + \frac{dt}{6} \frac{d}{dt}(s_{1}(u+tv) + 2s_{2}(u+tv) + 2s_{3}(u+tv) + s_{4}(u+tv))|_{t=0}  \\
    &\frac{d}{dt} s_{1}(u+tv) |_{t=0} = JL96(u)v \\
    &\frac{d}{dt} s_{2}(u+tv) |_{t=0} = JL96(u + (dt/2)s_{1}(u)) (v + \frac{dt}{2}\frac{d}{dt} s_{1}(u+tv) |_{t=0}
) \\
    &\frac{d}{dt} s_{3}(u+tv) |_{t=0} = JL96(u + (dt/2)s_{2}(u)) ( v + \frac{dt}{2} \frac{d}{dt} s_{2}(u+tv) |_{t=0} )\\
    &\frac{d}{dt} s_{4}(u+tv) |_{t=0} = JL96(u + dt s_{3}(u)) (v + dt \frac{d}{dt} s_{3}(u+tv) |_{t=0} ) 
\end{align*}

$JL96(u)v$を計算する必要があるので、$L96$の接線形コードも必要になります。作り方は今までと同じなので、$i=3,4,...,N-1$の部分だけ書きます。

    (L96(x+ty))_{i} = (x_{i+1} + ty_{i+1} - (x_{i-2} + ty_{i-2}))(x_{i-1} + ty_{i-1}) - (x_{i} + ty_{i}) + F \\
 \frac{d}{dt} (L96(x+ty))_{i} = (y_{i+1} - y_{i-2})(x_{i-1} + ty_{i-1}) + 
(x_{i+1} + ty_{i+1} - (x_{i-2} + ty_{i-2}))y_{i-1} - y_{i} \\
\frac{d}{dt} (L96(x+ty))_{i} |_{t=0} = (y_{i+1} - y_{i-2}) x_{i-1} + 
(x_{i+1} -x_{i-2})y_{i-1} - y_{i}

あとはこれらの式をコードにするだけです。

    #JL96(x)yを返す
    function TLL96(x,y;F=8.0)
        N = length(x)
        result = Array{Float64}(undef,N)
        result[1] = -x[N]*y[N-1] + (x[2]-x[N-1])y[N] - y[1] + x[N]*y[2]
        result[2] = -x[1]*y[N] + (x[3]-x[N])y[1] - y[2] + x[1]*y[3]
        for k in 3:N-1
            result[k] = -x[k-1]*y[k-2] + (x[k+1]-x[k-2])y[k-1] - y[k] + x[k-1]*y[k+1]
        end
        result[N] = -x[N-1]*y[N-2] + (x[1]-x[N-2])y[N-1] - y[N] + x[N-1]*y[1]

        return result
    end

    #JM(x)yを返す
    function TLM(x,y;dt=0.05)
        s1 = L96(x)
        s2 = L96(x + s1*dt/2.0)
        s3 = L96(x + s2*dt/2.0)
        t1 = TLL96(x,y)
        t2 = TLL96(x+s1*dt/2.0,y) + dt * TLL96(x+s1*dt/2.0, t1) / 2.0
        t3 = TLL96(x+s2*dt/2.0,y) + dt * TLL96(x+s2*dt/2.0, t2) / 2.0
        t4 = TLL96(x+s3*dt,y) + dt * TLL96(x+s3*dt, t3)
        return y + (t1 + 2.0*t2 + 2.0*t3 + t4)*(dt/6.0)
    end

TLM(x,y;dt=0.05)の内部に注目すると、s1,s2,s3,s4のところで元のモデルと同じ計算をしていることがわかります。MとTLMを分けて書いてしまうと同じ計算を二度することになって非効率的なので、同時並行で計算することもあるそうです。

接線形コードのデバッグは差分近似との比較で行います。すなわち、$JM(x)e_{i} = (M(x+he_{i})-M(x))/h$が$h$の2次以上を無視する近似で成り立つことを確かめます。

双対漸化式とアジョイントコード

$x,y \in \mathbb{R}^{N}$と$M \colon \mathbb{R}^{N} \to \mathbb{R}^{N}$に対して、$x$における$M$のJacobi行列の転置$JM(x)^{\top}$と$y$の積$JM(x)^{\top}y$を返すコードをアジョイントコードと言います。$JM(x)^{\top}$を計算して$y$との積を真面目に計算するとコストがかかるので、$JM(x)^{\top}y = \nabla_{z}(z,JM(x)^{\top}y) = \nabla_{z}(JM(x)z,y)$と先ほど作った接線形コードを利用するか、もしくはこれから説明する双対漸化式を利用して計算します。

$y\in \mathbb{R}^{N}$を入力とし、中間変数$v_{1},\ldots,v_{r} \in \mathbb{R}^{n}$を経て出力$z \in \mathbb{R}^{N}$を得る線形の漸化式

    v_{k} = A_{(k)}y + \sum_{j=1}^{k-1} L_{(k,j)} v_{j} \quad (k=1,2,\ldots,r),\quad 
    z = A_{(r+1)}y + \sum_{j=1}^{r} L_{(j)} v_{j}

を考えます。ここで、$(A_{(k)})_{k=1}^{r+1} \subset \mathbb{R}^{N\times N}$や$(L_{(k,j)})_{1\leq k\leq r, 1\leq j \leq k-1} \subset \mathbb{R}^{N\times n}$および$(L_{(j)})_{j=1}^{r} \subset \mathbb{R}^{N\times N}$は係数行列です（添え字は成分のことではないです）。$y$と$z$は線形関係にあるので、適当な行列$\mathcal{M}$を用いて、$z = \mathcal{M}y$と書けます。例えば、$z = JM(x)y$を計算する接線形コードは、$x$によって定まる係数行列を用いた線形の漸化式として上のような形に帰着され、$\mathcal{M} = JM(x)$となります。

この漸化式に対して、次のような手続きで$\hat{z}$を入力として中間変数$\hat{v}_{j}(j=1,2,\ldots,r)$を経て出力$\hat{y}$を計算する漸化式を双対漸化式と言います。

    \hat{y} = A_{(r+1)}^{\top} \hat{z},\quad \hat{v}_{j} = L_{(j)}^{\top} \hat{z} \quad (j=1,2,\ldots,r)\\
    \hat{y} \leftarrow \hat{y} + A_{(k)}^{\top} \hat{v}_{k}, \quad \hat{v}_{j} \leftarrow \hat{v}_{j} + L_{(k,j)}^{\top} \hat{v}_{k} \quad(k=r,r-1,\ldots,1)

途中の$\leftarrow$は右辺を左辺に改めて代入するという意味です。計算する順番が末尾からになっていることに注意してください。この計算の後、$\hat{y} = \mathcal{M}^{\top}\hat{z}$が成り立っています。証明は面倒なので$L_{(k,j)}=0$の場合（中間変数がない場合）やあったとしても$r=1,2$ぐらいの単純な場合に確かめてみるのが良いと思います。特に$z=JM(x)y$の場合に用いると$\hat{y}=JM(x)^{\top}\hat{z}$が成り立つので、双対漸化式を用いて$JM(x)^{\top}\hat{z}$を計算することができます。

背景にはいろいろな理論があったり$\hat{y},\hat{z}$にも入力に対する応答という物理的な意味がついたり自動微分の原理の一つになっていたりするのですが、詳細については立ち入らないので杉原・室田数値計算法の数理やその参考文献を参照してください。

例１

先ほどの例１で計算した接線形コードをアジョイントコードに変換します。

M(x_{1},x_{2}) = (x_{1}^2+x_{2}^2, x_{1}^2-x_{2}^2)

    function M(x)
        result = Array{Float64}(undef,2)
        result[1] = x[1]^2 + x[2]^2
        result[2] = x[1]^2 - x[2]^2
        return result
    end

    #JM(x)dxを返す
    function TLM(x,dx)
        dresult = Array{Float64}(undef,2)
        dresult[1] = 2.0*x[1]*dx[1] + 2.0*x[2]*dx[2]
        dresult[2] = 2.0*x[1]*dx[1] - 2.0*x[2]*dx[2]
        return dresult
    end

まずは$(JM(x)z,y)$を計算してみましょう。

    JM(x)z = (2x_{1}z_{1} + 2x_{2}z_{2}, 2x_{1}z_{1} - 2x_{2}z_{2}) \\
    (JM(x)z, y) = 2x_{1}z_{1}y_{1}+2x_{2}z_{2}y_{1} + 2x_{1}z_{1}y_{2} - 2x_{2}z_{2}y_{2} \\
    \nabla_{z}(JM(x)z,y) = (2x_{1}y_{1} + 2x_{1}y_{2}, 2x_{2}y_{1} - 2x_{2}y_{2})

あとは最後の式を関数にすれば出来上がりです。もうひとつ、双対漸化式を用いたやり方も試してみましょう。

    v_{1} = 2x_{1}y_{1} + 2x_{2}y_{2} \\
    v_{2} =  2x_{1}y_{1} - 2x_{2}y_{2}

これを双対漸化式に変換します。最後の行から行うことに注意してください。

    \hat{v}_{2} = \hat{v}_{2} \\
    \hat{x}_{1} = \hat{x}_{1} + 2y_{1}\hat{v}_{2} \\
    \hat{x}_{2} = \hat{x}_{2} - 2y_{2} \hat{v}_{2}, \\
    \hat{v}_{1} = \hat{v}_{1} \\
    \hat{x}_{1} = \hat{x}_{1} + 2y_{1}\hat{v}_{1} \\
    \hat{x}_{2} = \hat{x}_{2} + 2y_{2}\hat{v}_{2}

あとはこれを関数にすると完成です。$\hat{v}$についての式は今回は不要ですね。一行ずつ変換していけばよいので、複雑なコードに対してはこちらの方が書きやすいです。以降はこちらの方法を練習していきます。

例２

例２で計算したコードをアジョイントコードに書き換えます。一気に複雑になるのですが、一つ一つの手続き自体は単純です。


    #本当はx=0で0を返すのが符号関数
    function sgn(x)
        return x>=0 ? 1 : -1
    end

    #JM(u)vを返す
    function TLM(u,v)
        c1 = dt/(2.0*dx)
        c2 = nu*dt/dx^2
        N = length(u)
        result = Array{Float64}(undef,N)
        for i in 2:N-1
            result[i] = v[i] - c1*v[i]*(u[i+1]-u[i-1]) 
                        + c1*u[i]*(v[i+1] - v[i-1])
                        + c1*sgn(u[i])*v[i]*(u[i+1]-2*u[i]+u[i-1])
                        + c1*abs(u[i])*(v[i+1] - 2*v[i] + v[i-1])
                        + c2*(v[i+1]-2*v[i]+v[i-1])
        end
        result[1] = v[1] - c1*v[1]*(u[2]-ul) 
                        + c1*u[1]*v[2]
                        + c1*sgn(u[1])*v[1]*(u[2]-2*u[1]+ul)
                        + c1*abs(u[1])*(v[2] - 2*v[1])
                        + c2*(v[2]-2*v[1])
        result[N] = v[N] - c1*v[N]*(ur-u[N-1]) 
                        + c1*u[N]*(-v[N-1])
                        + c1*sgn(u[N])*v[N]*(ur-2*u[N]+u[N-1])
                        + c1*abs(u[N])*(- 2*v[N] + v[N-1])
                        + c2*(-2*v[N]+v[N-1])
        return result
    end

一番本質的な部分についてのみ書きます。

    for i in 2:N-1
            result[i] = v[i] - c1*v[i]*(u[i+1]-u[i-1]) 
                        + c1*u[i]*(v[i+1] - v[i-1])
                        + c1*sgn(u[i])*v[i]*(u[i+1]-2*u[i]+u[i-1])
                        + c1*abs(u[i])*(v[i+1] - 2*v[i] + v[i-1])
                        + c2*(v[i+1]-2*v[i]+v[i-1])
    end

見た目は複雑ですが中間変数がないため案外単純です。

    for i in N-1:-1:2
        x_dual[i] += v_dual[i]
        x_dual[i] += -c1*v_dual[i]*(u[i+1]-u[i-1])
        x_dual[i+1] += c1*u[i]*v_dual[i]
        x_dual[i-1] += -c1*u[i]*v_dual[i]
        x_dual[i] += c1*sgn(u[i])*v_dual[i]*(u[i+1]-2*u[i]+u[i-1])
        x_dual[i+1] += c1*abs(u[i])*v_dual[i]
        x_dual[i] += c1*abs(u[i])*(-2*v_dual[i])
        x_dual[i-1] += c1*abs(u[i])*v_dual[i]
        x_dual[i+1] += c2*v_dual[i]
        x_dual[i] += -2*c2*v_dual[i]
        x_dual[i-1] += c2*v_dual[i]
    end

もちろん処理をまとめて整理しても構いません。今回は中間変数が入っていないためループの順番を逆順にしなくても正しい結果を返しますが、中間変数がある場合は誤りになるので注意してください。この点については次の例で説明します。

例３

本題のLorenz方程式のアジョイントコードです。中間変数があり、処理をほかの関数に投げているのでなかなか大変ですが、原則通り一つ一つ処理していきます。

まず接線形コードの変形から入ります。コードを再掲します。

    #JM(x)yを返す
    function TLM(x,y;dt=0.05)
        s1 = L96(x)
        s2 = L96(x + s1*dt/2.0)
        s3 = L96(x + s2*dt/2.0)
        t1 = TLL96(x,y)
        t2 = TLL96(x+s1*dt/2.0,y) + dt * TLL96(x+s1*dt/2.0, t1) / 2.0
        t3 = TLL96(x+s2*dt/2.0,y) + dt * TLL96(x+s2*dt/2.0, t2) / 2.0
        t4 = TLL96(x+s3*dt,y) + dt * TLL96(x+s3*dt, t3)
        result = y + (t1 + 2.0*t2 + 2.0*t3 + t4)*(dt/6.0)
        return result
    end

入力変数と中間変数を双対変数に置き換えたy_dual, t1_dual,t2_dual,t3_dual,t4_dualを$0$で初期化し、末尾から双対漸化式に置き換えていきます。

    #result = y + (t1 + 2.0*t2 + 2.0*t3 + t4)*(dt/6.0)
    y_dual += r_dual
    t1_dual += r_dual * dt / 6.0
    t2_dual += 2.0*r_dual *dt / 6.0
    t3_dual += 2.0*r_dual * dt / 6.0
    t4_dual += r_dual * dt / 6.0

全て$0$で初期化しているので+=は単に=と書いても同じです。残りの式も双対漸化式に直します。
Adj_L96はTL_L96のアジョイントコードで、すぐ後で実装します。

    s1 = L96(x)
    s2 = L96(x + s1*dt/2.0)
    s3 = L96(x + s2*dt/2.0)

    #t4 = TLL96(x+s3*dt,y) + dt * TLL96(x+s3*dt, t3)
    y_dual += Adj_L96(x+s3*dt, t4_dual)
    t3_dual += dt * Adj_L96(x+s3*dt, t3_dual)

    #t3 = TLL96(x+s2*dt/2.0,y) + dt * TLL96(x+s2*dt/2.0, t2) / 2.0
    y_dual += Adj_L96(x+s2*dt/2.0, t3_dual)
    t2_dual += dt * Adj_L96(x+s2*dt/2.0, t3_dual) / 2.0

    #t2 = TLL96(x+s1*dt/2.0,y) + dt * TLL96(x+s1*dt/2.0, t1) / 2.0
    y_dual += Adj_L96(x+s1*dt/2.0, t2_dual)
    t1_dual += dt * Adj_L96(x+s1*dt/2.0, t2_dual) / 2.0

    #t1 = TLL96(x,y)
    y_dual += Adj_L96(x,t1_dual)

    return y_dual

次に、途中で必要になるL96のアジョイントコードの実装をします。
接線形コードは次のように書いたのでした。

function TL_L96(x,u;F=8.0,N=40)
    TL = fill(0.0, N)
    TL[1] = -x[N]*u[N-1] + (x[2]-x[N-1])u[N] - u[1] + x[N]*u[2]
    TL[2] = -x[1]*u[N] + (x[3]-x[N])u[1] - u[2] + x[1]*u[3]
    for k in 3:N-1
        TL[k] = -x[k-1]*u[k-2] + (x[k+1]-x[k-2])u[k-1] - u[k] + x[k-1]*u[k+1]
    end
    TL[N] = -x[N-1]*u[N-2] + (x[1]-x[N-2])u[N-1] - u[N] + x[N-1]*u[1]

    return TL
end

Adj = fill(0.0, N)として$0$ベクトルで初期化した後、一行づつ変形します。全部書くと冗長なのでメインとなる

    for k in 3:N-1
        TL[k] = -x[k-1]*u[k-2] + (x[k+1]-x[k-2])u[k-1] - u[k] + x[k-1]*u[k+1]
    end

の部分だけ変形します。ループの変数を逆順にし、基準場の変数$x$を除いたTL,uを双対変数に置き換えて漸化式を作ります。

    for k in N-1:-1:3
        #TL[k] = -x[k-1]*u[k-2] + (x[k+1]-x[k-2])u[k-1] - u[k] + x[k-1]*u[k+1]
        Adj[k-2] += -x[k-1]*u[k]
        Adj[k-1] += (x[k+1]-x[k-2])u[k]
        Adj[k] += -u[k]
        Adj[k+1] += x[k-1]*u[k]
    end

以上の関数をまとめると、完成です。


function Adj_L96(x,u;N=40,F=8.0)
    Adj = fill(0.0, N)
    #TL[1] = -x[N]*u[N-1] + (x[2]-x[N-1])u[N] - u[1] + x[N]*u[2]
    Adj[N-1] += -x[N]*u[1]
    Adj[N] += (x[2]-x[N-1])u[1]
    Adj[1] += -u[1]
    Adj[2] += x[N]*u[1]

    #TL[2] = -x[1]*u[N] + (x[3]-x[N])u[1] - u[2] + x[1]*u[3]
    Adj[N] += -x[1]*u[2]
    Adj[1] += (x[3]-x[N])u[2]
    Adj[2] += -u[2]
    Adj[3] += x[1]*u[2]

    for k in N-1:-1:3
        #TL[k] = -x[k-1]*u[k-2] + (x[k+1]-x[k-2])u[k-1] - u[k] + x[k-1]*u[k+1]

        Adj[k-2] += -x[k-1]*u[k]
        Adj[k-1] += (x[k+1]-x[k-2])u[k]
        Adj[k] += -u[k]
        Adj[k+1] += x[k-1]*u[k]
    end

    #TL[N] = -x[N-1]*u[N-2] + (x[1]-x[N-2])u[N-1] - u[N] + x[N-1]*u[1]
    Adj[N-2] += -x[N-1]*u[N]
    Adj[N-1] += (x[1]-x[N-2])u[N]
    Adj[N] += -u[N]
    Adj[1] += x[N-1]*u[N]

    return Adj

end

#JM(x)^T yを返す
function Adj_Model(x,y)
    #result = y + (t1 + 2.0*t2 + 2.0*t3 + t4)*(dt/6.0)
    y_dual = y
    t1_dual = y*dt / 6.0
    t2_dual = 2.0*y*dt / 6.0
    t3_dual = 2.0*y*dt / 6.0
    t4_dual = y*dt / 6.0

    s1 = L96(x)
    s2 = L96(x + s1*dt/2.0)
    s3 = L96(x + s2*dt/2.0)

    #t4 = TLL96(x+s3*dt,y) + dt * TLL96(x+s3*dt, t3)
    y_dual += Adj_L96(x+s3*dt, t4_dual)
    t3_dual += dt * Adj_L96(x+s3*dt, t3_dual)

    #t3 = TLL96(x+s2*dt/2.0,y) + dt * TLL96(x+s2*dt/2.0, t2) / 2.0
    y_dual += Adj_L96(x+s2*dt/2.0, t3_dual)
    t2_dual += dt * Adj_L96(x+s2*dt/2.0, t3_dual) / 2.0

    #t2 = TLL96(x+s1*dt/2.0,y) + dt * TLL96(x+s1*dt/2.0, t1) / 2.0
    y_dual += Adj_L96(x+s1*dt/2.0, t2_dual)
    t1_dual += dt * Adj_L96(x+s1*dt/2.0, t2_dual) / 2.0

    #t1 = TLL96(x,y)
    y_dual += Adj_L96(x,t1_dual)

    return y_dual
end

大抵の場合作ったアジョイントコードにはバグがあります。デバッグ方法として、$(JM(x)y,z) = (y, JM(x)^{\top}z)$が成り立つことなどを利用するとよいと思います。

見ての通りアジョイントコードを書くのは大変です。直接アジョイントコードを自動生成することもできるらしいのですが、複雑なコードに対しては正しい結果を返さないようです。手っ取り早く$JM(x)^{\top}y$を求めたい場合、$z \mapsto (JM(x)z,y)$の導関数を自動微分することが考えられます。

変分法の実装と計算結果

以前の記事と同様、観測値の誤差共分散$R_{i}$や観測演算子$H_{i}$は$i$によらず一定とします。

3次元変分法

全体を示してから個別に説明していきます。L96_observation.txtから観測データを読み込んでBFGS法による最適化で解析値$x_{i}^{a}$を計算し、L96_truestate.txtに保存された真値との誤差を計算して結果をL96_3DVAR_DESCENT_output.txtに出力するプログラムです。

アジョイントモデルなどのコード（すでに解説済みなので折り畳み）

using LinearAlgebra
using Statistics
using DelimitedFiles
using Random
using Optim
using Plots


##L96modelの右辺
function L96(u;F=8.0,N=40)
    f = fill(0.0, N)
    f[1] = (u[2]-u[N-1])u[N] - u[1] + F
    f[2] = (u[3]-u[N])u[1] - u[2] + F
    for k in 3:N-1
        f[k] = (u[k+1]-u[k-2])u[k-1] - u[k] + F
    end
    f[N] =  (u[1]-u[N-2])u[N-1] - u[N] + F

    return f
end

function TL_L96(x,u;F=8.0,N=40)
    TL = fill(0.0, N)
    TL[1] = -x[N]*u[N-1] + (x[2]-x[N-1])u[N] - u[1] + x[N]*u[2]
    TL[2] = -x[1]*u[N] + (x[3]-x[N])u[1] - u[2] + x[1]*u[3]
    for k in 3:N-1
        TL[k] = -x[k-1]*u[k-2] + (x[k+1]-x[k-2])u[k-1] - u[k] + x[k-1]*u[k+1]
    end
    TL[N] = -x[N-1]*u[N-2] + (x[1]-x[N-2])u[N-1] - u[N] + x[N-1]*u[1]

    return TL
end

function Adj_L96(x,u;N=40,F=8.0)
    Adj = fill(0.0, N)
    #TL[1] = -x[N]*u[N-1] + (x[2]-x[N-1])u[N] - u[1] + x[N]*u[2]
    Adj[N-1] += -x[N]*u[1]
    Adj[N] += (x[2]-x[N-1])u[1]
    Adj[1] += -u[1]
    Adj[2] += x[N]*u[1]

    #TL[2] = -x[1]*u[N] + (x[3]-x[N])u[1] - u[2] + x[1]*u[3]
    Adj[N] += -x[1]*u[2]
    Adj[1] += (x[3]-x[N])u[2]
    Adj[2] += -u[2]
    Adj[3] += x[1]*u[2]

    for k in N-1:-1:3
        #TL[k] = -x[k-1]*u[k-2] + (x[k+1]-x[k-2])u[k-1] - u[k] + x[k-1]*u[k+1]

        Adj[k-2] += -x[k-1]*u[k]
        Adj[k-1] += (x[k+1]-x[k-2])u[k]
        Adj[k] += -u[k]
        Adj[k+1] += x[k-1]*u[k]
    end

    #TL[N] = -x[N-1]*u[N-2] + (x[1]-x[N-2])u[N-1] - u[N] + x[N-1]*u[1]
    Adj[N-2] += -x[N-1]*u[N]
    Adj[N-1] += (x[1]-x[N-2])u[N]
    Adj[N] += -u[N]
    Adj[1] += x[N-1]*u[N]

    return Adj

end

#4-4Runge-Kutta
function Model(u;dt=0.05)
    du = u
    s1 = L96(u)
    s2 = L96(u + s1*dt/2)
    s3 = L96(u + s2*dt/2)
    s4 = L96(u + s3*dt)
    du += (s1 + 2.0*s2 + 2.0*s3 + s4)*(dt/6.0)
    return du
end

function TL_Model(x,u;dt=0.05)
    s1 = L96(x)
    s2 = L96(x + s1*dt/2.0)
    s3 = L96(x + s2*dt/2.0)
    t1 = TL_L96(x,u)
    t2 = TL_L96(x+s1*dt/2.0,u) + dt * TL_L96(x+s1*dt/2.0, t1) / 2.0
    t3 = TL_L96(x+s2*dt/2.0,u) + dt * TL_L96(x+s2*dt/2.0, t2) / 2.0
    t4 = TL_L96(x+s3*dt,u) + dt * TL_L96(x+s3*dt, t3)
    z = u + (t1 + 2.0*t2 + 2.0*t3 + t4)*(dt/6.0)
    return z
end

function Adj_Model(x,u;dt=0.05)
    s1 = L96(x)
    s2 = L96(x + s1*dt/2.0)
    s3 = L96(x + s2*dt/2.0)

    y_dual = u
    t1_dual = dt*u/6.0
    t2_dual = 2.0*dt*u/6.0
    t3_dual = 2.0*dt*u/6.0
    t4_dual = dt*u/6.0

    y_dual += Adj_L96(x+s3*dt,t4_dual)
    t3_dual += dt*Adj_L96(x+s3*dt, t4_dual)
    y_dual += Adj_L96(x+s2*dt/2.0, t3_dual)
    t2_dual += dt*Adj_L96(x+s2*dt/2.0, t3_dual)/2.0
    y_dual += Adj_L96(x+s1*dt/2.0, t2_dual)
    t1_dual += dt*Adj_L96(x+s1*dt/2.0, t2_dual)/2.0
    y_dual += Adj_L96(x, t1_dual)

    return y_dual
end

3DVAR.jl

function VAR_3D()
    Time_Step = 14600
    N = 40
    M = 40
    F = 8.0
    IN = Matrix(1.0I, N, N)
    H = Matrix(1.0I, M, N)
    R = Matrix(1.0I, M, M)
    alpha = 0.4
    B = alpha * IN

    J(x,xf,y,H,B,R) = ((x-xf)'*inv(B)*(x-xf) + (H*x-y)'*inv(R)*(H*x-y))/2.0

    function grad_J!(G,x,xf,y,H,B,R;N=40)
        gJ = Array{Float64}(undef, N)
        gJ = inv(B)*(x-xf) + H'*inv(R)*(H*x-y)
        for i = 1:N
            G[i] = gJ[i]
        end
    end


    truestate = readdlm("L96_truestate.txt")
    obs = readdlm("L96_observation.txt")
    open("L96_3DVAR_DESCENT_output.txt", "w") do output

        ua = rand(N) .+ F
        for i = 1:Time_Step
            ua = Model(ua)
        end


        for i = 1:Time_Step
            #forecast step
            uf = Model(ua)

            #analysis step
            y = H * view(obs,i,2:N+1)
            opt = optimize(x -> J(x,uf,y,H,B,R), (G,x) -> grad_J!(G,x,uf,y,H,B,R), uf, LBFGS())
            ua = Optim.minimizer(opt)
            res = Optim.iterations(opt)
            is_success = Optim.converged(opt)

            writedlm(output,[(i / 40) (norm(truestate[i, 2:N+1] - ua) / sqrt(N)) res is_success])

        end
    end

end

目的関数とその導関数

       \mathcal{J}(x) = \frac{1}{2}(B_{i}^{-1}(x-x_{i}^{f}), x-x_{i}^{f}) + \frac{1}{2} (R_{i}^{-1}(H_{i}x-y_{i}),H_{i}x-y_{i})  \\
    \nabla \mathcal{J}(x) = B_{i}^{-1}(x-x_{i}^{f}) + H_{i}^{\top}R_{i}^{-1}(H_{i}x-y_{i})

    J(x,xf,y,H,B,R) = ((x-xf)'*inv(B)*(x-xf) + (H*x-y)'*inv(R)*(H*x-y))/2.0

    function grad_J!(G,x,xf,y,H,B,R;N=40)
        gJ = Array{Float64}(undef, N)
        gJ = inv(B)*(x-xf) + H'*inv(R)*(H*x-y)
        for i = 1:N
            G[i] = gJ[i]
        end
    end

そのまま書けます。$\nabla \mathcal{J}(x)$は勾配ベクトル$G$を入力して$G$に$\nabla \mathcal{J}(x)$を書き込んで返す関数として実装していますが、これは最適化に使うライブラリがこのような形で導関数を渡すように要求しているためです。

Forecast Step

    x_{i}^{f} = M(x_{i-1}^{a})

    uf = Model(ua)

そのままです。予報誤差共分散の更新はしようがないのでずっと一定のままにしていますが、カルマンフィルターのように適切に更新していった方がもちろん精度は上がります。この部分は考えどころです。例えば（結構な力技ですが）アンサンブルカルマンフィルターを同時に走らせて$P_{i}^{f}$をそちらから持ってくるなどの方法があるらしいです。

Analysis Step

    x_{i}^{a} = \mathrm{argmin} \mathcal{J}(x)

    y = H * view(obs,i,2:N+1)
    opt = optimize(x -> J(x,uf,y,H,B,R), (G,x) -> grad_J!(G,x,uf,y,H,B,R), uf, LBFGS())
    ua = Optim.minimizer(opt)
    res = Optim.iterations(opt)
    is_success = Optim.converged(opt)

Optim.jlを用いて省メモリBFGSで極小値を求めます。反復回数と収束判定をパスしているかどうかも同時に記録することにします。今回の設定ではそう問題にはなりませんが、$H$が非線形で観測値の次元$d$が小さいときには一般に$\mathcal{J}$は多峰になり、適切に初期値を取らないと極小値は最小値に一致しません。確率的勾配降下法などの力技を使うか、初期値を試行錯誤しないといけないのは最適化の一般論と同様です。

$\nabla \mathcal{J}(x)$の零点は解析的に求まり、$K = BH^{\top}(HBH^{\top}+R)^{-1}$に対して$x_{i}^{a} = (I-KH)x_{i}^{f} + Ky_{i}^{o}$となります。これはカルマンフィルターの式と同じです。理由についてはデータ同化の教科書を参照してください。ここでは最適化の練習として勾配法で求めています。

計算結果

$B=\alpha I, \alpha = 0.4$に対して結果をグラフにすると次のようになります。勾配法は全てのステップで一回で収束していました。

今回の設定ではRMSEは$0.45$前後であり、$0.2$前後だったカルマンフィルターと比較すると劣ります。原因は予報誤差共分散をずっと一定値にしていることにあり、適切に更新していけるなら（今回の設定では解析値の更新式は同じになるので）ほとんど推定値の差はありません。どちらが良いのかについては仮定した分布や$M,H$によっても変わってくるので一概には言えません。さらに言えば数値計算の規模によっては一方の手法がメモリの制約などで使えないことすらあります。結局のところ場合に応じて使い分けるしかないと思います。

4次元変分法

こちらも全体を示してから解説していきたいと思います。。L96_observation.txtから観測データを読み込んでBFGS法による最適化で解析値$x_{i}^{a}$を計算し、L96_truestate.txtに保存された真値との誤差を計算して結果をL96_4DVAR_output.txtに出力するプログラムです。3次元変分法に比べるととても大変です。

アジョイントモデルなどのコード（すでに解説済みなので折り畳み）

using LinearAlgebra
using Statistics
using DelimitedFiles
using Random
using Optim
using Plots


##L96modelの右辺
function L96(u;F=8.0,N=40)
    f = fill(0.0, N)
    f[1] = (u[2]-u[N-1])u[N] - u[1] + F
    f[2] = (u[3]-u[N])u[1] - u[2] + F
    for k in 3:N-1
        f[k] = (u[k+1]-u[k-2])u[k-1] - u[k] + F
    end
    f[N] =  (u[1]-u[N-2])u[N-1] - u[N] + F

    return f
end

function TL_L96(x,u;F=8.0,N=40)
    TL = fill(0.0, N)
    TL[1] = -x[N]*u[N-1] + (x[2]-x[N-1])u[N] - u[1] + x[N]*u[2]
    TL[2] = -x[1]*u[N] + (x[3]-x[N])u[1] - u[2] + x[1]*u[3]
    for k in 3:N-1
        TL[k] = -x[k-1]*u[k-2] + (x[k+1]-x[k-2])u[k-1] - u[k] + x[k-1]*u[k+1]
    end
    TL[N] = -x[N-1]*u[N-2] + (x[1]-x[N-2])u[N-1] - u[N] + x[N-1]*u[1]

    return TL
end

function Adj_L96(x,u;N=40,F=8.0)
    Adj = fill(0.0, N)
    #TL[1] = -x[N]*u[N-1] + (x[2]-x[N-1])u[N] - u[1] + x[N]*u[2]
    Adj[N-1] += -x[N]*u[1]
    Adj[N] += (x[2]-x[N-1])u[1]
    Adj[1] += -u[1]
    Adj[2] += x[N]*u[1]

    #TL[2] = -x[1]*u[N] + (x[3]-x[N])u[1] - u[2] + x[1]*u[3]
    Adj[N] += -x[1]*u[2]
    Adj[1] += (x[3]-x[N])u[2]
    Adj[2] += -u[2]
    Adj[3] += x[1]*u[2]

    for k in N-1:-1:3
        #TL[k] = -x[k-1]*u[k-2] + (x[k+1]-x[k-2])u[k-1] - u[k] + x[k-1]*u[k+1]

        Adj[k-2] += -x[k-1]*u[k]
        Adj[k-1] += (x[k+1]-x[k-2])u[k]
        Adj[k] += -u[k]
        Adj[k+1] += x[k-1]*u[k]
    end

    #TL[N] = -x[N-1]*u[N-2] + (x[1]-x[N-2])u[N-1] - u[N] + x[N-1]*u[1]
    Adj[N-2] += -x[N-1]*u[N]
    Adj[N-1] += (x[1]-x[N-2])u[N]
    Adj[N] += -u[N]
    Adj[1] += x[N-1]*u[N]

    return Adj

end

#4-4Runge-Kutta
function Model(u;dt=0.05)
    du = u
    s1 = L96(u)
    s2 = L96(u + s1*dt/2)
    s3 = L96(u + s2*dt/2)
    s4 = L96(u + s3*dt)
    du += (s1 + 2.0*s2 + 2.0*s3 + s4)*(dt/6.0)
    return du
end

function TL_Model(x,u;dt=0.05)
    s1 = L96(x)
    s2 = L96(x + s1*dt/2.0)
    s3 = L96(x + s2*dt/2.0)
    t1 = TL_L96(x,u)
    t2 = TL_L96(x+s1*dt/2.0,u) + dt * TL_L96(x+s1*dt/2.0, t1) / 2.0
    t3 = TL_L96(x+s2*dt/2.0,u) + dt * TL_L96(x+s2*dt/2.0, t2) / 2.0
    t4 = TL_L96(x+s3*dt,u) + dt * TL_L96(x+s3*dt, t3)
    z = u + (t1 + 2.0*t2 + 2.0*t3 + t4)*(dt/6.0)
    return z
end

function Adj_Model(x,u;dt=0.05)
    s1 = L96(x)
    s2 = L96(x + s1*dt/2.0)
    s3 = L96(x + s2*dt/2.0)

    y_dual = u
    t1_dual = dt*u/6.0
    t2_dual = 2.0*dt*u/6.0
    t3_dual = 2.0*dt*u/6.0
    t4_dual = dt*u/6.0

    y_dual += Adj_L96(x+s3*dt,t4_dual)
    t3_dual += dt*Adj_L96(x+s3*dt, t4_dual)
    y_dual += Adj_L96(x+s2*dt/2.0, t3_dual)
    t2_dual += dt*Adj_L96(x+s2*dt/2.0, t3_dual)/2.0
    y_dual += Adj_L96(x+s1*dt/2.0, t2_dual)
    t1_dual += dt*Adj_L96(x+s1*dt/2.0, t2_dual)/2.0
    y_dual += Adj_L96(x, t1_dual)

    return y_dual
end

4DVAR.jl

function J_4DVAR(x::Array{Float64,1},xf::Array{Float64,1},obs,H::Matrix{Float64},B::Matrix{Float64},R::Matrix{Float64},window_size)
    result::Float64 = (x-xf)'*inv(B)*(x-xf)/2.0

    result += (H*x-H*obs[1,:])'*inv(R)*(H*x-H*obs[1,:])/2.0
    for i in 1:window_size
        x = Model(x) ;
        result += (H*x-H*obs[i+1,:])'*inv(R)*(H*x-H*obs[i+1,:])/2.0
    end
    return result
end

function gradJ_4DVAR!(G::Array{Float64,1}, x::Array{Float64,1}, xf::Array{Float64,1}, obs, H::Array{Float64,2}, B::Array{Float64,2}, R::Array{Float64,2}, window_size ; N=40)
    gJ = inv(B) * (x - xf)

    ad = zeros(window_size, N)
    d = zeros(window_size+1, N)
    x_step = zeros(window_size+1,N)

    x_step[1,:] = x
    d[1,:] = H'*inv(R)*(H*x_step[1,:]-H*obs[1,:])
    for i in 1:window_size
        x_step[i+1,:] = Model(x_step[i,:]) 
        d[i+1,:] = H'*inv(R)*(H*x_step[i+1,:]-H*obs[i+1,:])
    end

    ad = Adj_Model(x_step[window_size,:], d[window_size+1,:]) + d[window_size,:]
    for i in 2:window_size
        ad = Adj_Model(x_step[window_size+1-i,:], ad) + d[window_size+1-i,:]
    end
    gJ += ad

    for i = 1:N
        G[i] = gJ[i]
    end
end

function VAR_4D()
    Time_Step = 14600
    N = 40
    M = 40
    F = 8.0
    IN = Matrix(1.0I, N, N)
    H = Matrix(1.0I, M, N)
    R = Matrix(1.0I, M, M)
    alpha = 0.4
    B = alpha * IN

    #3次元変分法のためのJとその導関数
    #Time_Stepの最後の最後でのみ使う
    J(x, xf, y, H, B, R) = ((x - xf)' * inv(B) * (x - xf) + (H * x - y)' * inv(R) * (H * x - y)) / 2.0

    function grad_J!(G, x, xf, y, H, B, R; N=40)
        gJ = Array{Float64}(undef, N)
        gJ = inv(B) * (x - xf) + H' * inv(R) * (H * x - y)
        for i = 1:N
            G[i] = gJ[i]
        end
    end


    #同化する際の時間幅
    #広すぎるとあまり意味がなく、計算資源の無駄使いになる
    window_size = 2
    
    truestate = readdlm("L96_truestate.txt")
    obs = readdlm("L96_observation.txt")
    open("L96_4DVAR_output.txt", "w") do output
        ua = rand(N) .+ F
        for i = 1:Time_Step
            ua = Model(ua)
        end

        #for iter in 0:window_step-1
        for i in 1:Time_Step
            #forecast step
            uf = Model(ua)

            #analysis step

            #Time_Stepを超えないように取り直す（末尾でしか起こらない）
            if i+window_size>Time_Step
                window_size = Time_Step - i
            end

            #もし0になったら3次元変分法でやる（末尾でしかおこらない）
            if window_size == 0
                y = H * view(obs, i, 2:N+1)
                opt = optimize(x -> J(x, uf, y, H, B, R), (G, x) -> grad_J!(G, x, uf, y, H, B, R), uf, LBFGS())
                ua = Optim.minimizer(opt)
                res = Optim.iterations(opt)
                is_success = Optim.converged(opt)

                writedlm(output, [(i / 40) (norm(truestate[i, 2:N+1] - ua) / sqrt(N)) res is_success])
                continue
            end

            #最適化
            y = view(obs, i:i+window_size, 2:N+1)
            opt = optimize(x -> J_4DVAR(x, uf, y, H, B, R, window_size), (G, x) -> gradJ_4DVAR!(G, x, uf, y, H, B, R, window_size), uf, LBFGS())
            ua = Optim.minimizer(opt)
            res = Optim.iterations(opt)
            is_success = Optim.converged(opt)

            writedlm(output, [(i / 40) (norm(truestate[i, 2:N+1] - ua) / sqrt(N)) res is_success])
           
        end
    end

end

目的関数とその導関数

       \mathcal{J}(x) = \frac{1}{2}(B_{i}^{-1}(x-x_{i}^{f}), x-x_{i}^{f}) + \sum_{j=0}^{T} \frac{1}{2} (R_{i+j}^{-1}(H_{i+j}M^{j}(x)-y_{i+j}),H_{i+j}M^{j}(x)-y_{i+j})

function J_4DVAR(x::Array{Float64,1},xf::Array{Float64,1},obs,H::Matrix{Float64},B::Matrix{Float64},R::Matrix{Float64},window_size)
    result::Float64 = (x-xf)'*inv(B)*(x-xf)/2.0

    result += (H*x-H*obs[1,:])'*inv(R)*(H*x-H*obs[1,:])/2.0
    for i in 1:window_size
        x = Model(x) ;
        result += (H*x-H*obs[i+1,:])'*inv(R)*(H*x-H*obs[i+1,:])/2.0
    end
    return result
end

観測obsとして$(y_{j}^{o})_{j=i}^{i+T}$を一斉に入力しないといけないので、ここでは行列にして$k\geq 1$行目に$y_{i+k-1}$が入っているものとして渡しています。$T$という記号の代わりにwindow_sizeという名前で渡しています。

導関数の実装は3次元変分法に比べるととても大変です。

    x_{i+j} = M^{j}(x_{i}), \quad (j=0,1,2,\ldots,T) \\
    d_{j}=H_{i+j}^{T}R_{i+j}^{-1}(H_{i+j}x_{i+j}-y_{i+j}), \quad (j=0,1,2,\ldots,T) \\
    e_{0} = d_{T},\, e_{k} = JM^{\top}(x_{T-k+i})e_{k-1} + d_{T-k}, \quad (k=1,2,\ldots,T) \\
    \nabla \mathcal{J} (x) = B_{i}^{-1}(x-x_{i}^{f}) + e_{T}

function gradJ_4DVAR!(G::Array{Float64,1}, x::Array{Float64,1}, xf::Array{Float64,1}, obs, H::Array{Float64,2}, B::Array{Float64,2}, R::Array{Float64,2}, window_size ; N=40)
    gJ = inv(B) * (x - xf)
    
    ad = zeros(window_size, N)
    d = zeros(window_size+1, N)
    x_step = zeros(window_size+1,N)

    x_step[1,:] = x
    d[1,:] = H'*inv(R)*(H*x_step[1,:]-H*obs[1,:])
    for i in 1:window_size
        x_step[i+1,:] = Model(x_step[i,:]) 
        d[i+1,:] = H'*inv(R)*(H*x_step[i+1,:]-H*obs[i+1,:])
    end

    ad = Adj_Model(x_step[window_size,:], d[window_size+1,:]) + d[window_size,:]
    for i in 2:window_size
        ad = Adj_Model(x_step[window_size+1-i,:], ad) + d[window_size+1-i,:]
    end
    gJ += ad

    for i = 1:N
        G[i] = gJ[i]
    end
end

$d_{j}$は行列dの$j$列目に対応し、$x_{i+j}$は行列x_stepの$j$列目に対応するとして格納します。$e_{k}$はベクトルadに格納することで持つようにしています。$G$を入力して書き込んで返すようにしているのは3次元変分法と同じ理由で、最適化に使うライブラリがそのように要求しているからです。window_sizeまで時間発展する途中の結果をすべて持っておかないと勾配の計算ができないこと、アジョイントコードによる計算をwindow_size分だけ行うことは認識しておいてください。適切なwindow_sizeの選び方については計算結果のところで解説します。

Forecast Step

    x_{i}^{f} = M(x_{i-1}^{a})

    uf = Model(ua)

そのままです。

Analysis Step

    x_{i}^{a} = \mathrm{argmin} \mathcal{J}(x)

    #analysis step
    #Time_Stepを超えないように取り直す（末尾でしか起こらない）
    if i+window_size>Time_Step
         window_size = Time_Step - i
    end
    
    #最適化
    y = view(obs, i:i+window_size, 2:N+1)
    opt = optimize(x -> J_4DVAR(x, uf, y, H, B, R, window_size), (G, x) -> gradJ_4DVAR!(G, x, uf, y, H, B, R, window_size), uf, LBFGS())
    ua = Optim.minimizer(opt)
    res = Optim.iterations(opt)
    is_success = Optim.converged(opt)

ここは3次元変分法とほとんど変わりません。観測をwindow_size分だけ切り出し、J_4DVARやgradJ_4DVAR!に渡して最適化します。計算の最後のところだけすこし処理をする必要がありますが、それだけです。

計算結果

window_sizeによる依存性を見るため、window_sizeを$1$から$5$まで変えながら同化期間の後半のRMSEの平均を取った図、および最適化が目的値に収束するまでの平均反復回数を示します。

見ての通りwindow_sizeを大きくすると結果は少しずつ改善するのですが、平均反復回数が大きくなっていくので計算にはものすごく時間がかかるようになります。これ以上window_sizeを増やすのは大変です。前処理と初期値の設定についてはいろいろな手法があるので、先述のデータ同化の教科書を参照するとよいと思います。

window_sizeをどうとるのが良いかは難しい問題です。$M$がカオス性を持たない写像の場合は計算資源が許す限り取れば良いと思いますが、今回のようにカオス性を持つ場合ではどうでしょうか。$x \in \mathbb{R}^{N}$を正規分布に従って沢山生成し、$M^{j}(x)$との共分散行列のトレースを計算して$j$を横軸にとってグラフにしてみると下のような図になります。

見ての通り時間とともに相関が弱くなり、100回反復したあたりでほとんど$0$になります。相関が$0$だから独立というのは少々乱暴ですが、知りたい真値の分布と独立な分布からいくらサンプリングしても意味がないわけで、$100$を超えるwindow_sizeをとっても結果はほとんど改善しないように思われます。

逆に言うとこの図を参考にするならwindow_size=100ぐらいまではRMSEが改善するようので、なんとか方法を考えて確かめてみたいところです。3変数Lorenzあたりのもっと計算が軽いモデルでやってみた方がいいかもしれません。

あとがき

今回は変分法についての解説をしました。4次元変分法の場合は実装するのが大変で計算も時間がかかります。しかしながらKalman Filter系列の手法と比べてどちらかが優れているというものではないので、両方使えるのに越したことはないのだと思います。次回は粒子法の説明になると思います。

$A$を固定するごとに$P_{Y=y}(A)$はa.s.でしか定まらないので、除外する零集合が$A$によらずにとれることを言うため、正確には正則条件付き確率というものを導入する必要があります。この定義はあまりにも中学校で習う条件付き確率からかけ離れています。伊藤清確率論p147以降に条件付き確率からの延長線で定義する方法が載っているので、興味のある方は読むとよいと思います(この記事に書いてある方法よりも直感的な定義でとても良いのですが、結構大変です)。 ↩