TensorFlowを(中途半端に)使って常微分方程式 (ODE)の数値計算をする

TensorFlowを使って常微分方程式(ODE)の初期値問題の数値計算をやってみます。速度の面では実用性はありませんが、TensorFlowに備わっている強力な自動微分を使いこなせれば色々解析に便利な面があると思われ、たとえばNeural Networkなどで複雑なモデルを作っていたとしてもヤコビ行列の計算などを自動的にしてくれることなどが利点と思います。

<!-- また最近話題の*neuralODE*を実装するための準備にもなります。 -->

今回はグラフを使って微分方程式を定義してみますが、実際の数値計算はscipyに入っている`solve_ivp`を使用します。 実は拡張ライブラリの[TensorFlow Probability](https://www.tensorflow.org/probability/) においてTensorFlow上で動くODEソルバーが実装されているようですが、理解不足もあり今回は使いません。

# 環境
- TensorFlow 1.14
- Python 3.6.8
- scipy 1.2.1


# やりたいこと

一階の常微分方程式 

```math
\frac{d \boldsymbol{x}}{dt} = f(\boldsymbol{x})
```
の初期値問題の近似解を数値計算により求めます。 $\boldsymbol{x} \in \mathbb{R}^{d}$は$d$次元ベクトル、$f$は$\mathbb{R}^{d}$ から $\mathbb{R}^{d}$への写像（ベクトル場と呼ぶ）とします。

例として次の[ローレンツ方程式](https://en.wikipedia.org/wiki/Lorenz_system)を以後使います。３次元の微分方程式であり、E.ローレンツがカオスを発見したことで有名です。 $\boldsymbol{x}=(x,y,z)^T$ として

```math
\begin{equation}
\frac{d \boldsymbol{x}}{dt} = 
\left( \begin{matrix}
f_x(\boldsymbol{x})  \\
f_y(\boldsymbol{x})  \\
f_z(\boldsymbol{x})
\end{matrix}\right) = 
\left( \begin{matrix}
\sigma(y-x)  \\
x(\rho -z) -y  \\
 xy - \beta z
\end{matrix}\right)  \tag{1}
\end{equation}
```
$\sigma, \rho, \beta $ は定数のパラメータです。

この方程式をTensorFlowのグラフ上で定義し、ODEソルバーを使って数値解を求めてみたいと思います。尚，今回は倍精度で計算を行います．

#グラフの定義

##準備
tensorflow, scipyのodeソルバー呼び出しに使う```solve_ivp```，その他をインポートします．

```
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 unused import
import time 

import tensorflow as tf
from scipy.integrate import solve_ivp
```

##ベクトル場の定義

まず変数 $\boldsymbol{x}$ のテンソルですが外(ODEソルバー）から値を入れて使うので、今回はこんな感じでplaceholder型にします。

```
x = tf.placeholder(dtype=tf.float64, shape=(3), name='x')
```

入力$x$からベクトル場を計算するグラフを作るクラスを作ります。インスタンスを作っておき，パラメータはメンバ変数に持たせるようにしておくと便利です。（参考: https://github.com/titu1994/tfdiffeq/blob/master/examples/lorenz_attractor.py )

```
class Lorenz(object):

    def __init__(self, sigma=10., beta=8 / 3., rho=28., **kwargs):
        self.sigma = float(sigma)
        self.beta = float(beta)
        self.rho = float(rho)

    def __call__(self, t, x):
        dx_dt = self.sigma * (x[1] - x[0])
        dy_dt = x[0] * (self.rho - x[2]) - x[1]
        dz_dt = x[0] * x[1] - self.beta * x[2]

        dX_dt = tf.stack([dx_dt, dy_dt, dz_dt])
        return dX_dt
```

## グラフの作成

```
f_lorenz = Lorenz()
fx = f_lorenz(None, x)  # r.h.s. of ODE
```

# Sessionを動かして計算

とりあえず初期値を設定しておきます．

```
x0 = np.array([1, 10, 10], dtype=np.float64)   # initial value
```

## sessionの開始，動作の確認

```
sess = tf.Session()
sess.run(tf.initializers.global_variables())
```
$\boldsymbol{x}$に初期値$(1,10,10)$を入れたときに$f(\boldsymbol{x})$を正しく計算してるか見てみます．

```
print('###value of f(x)###')
print(sess.run(fx, feed_dict={x:x0})) 
```
結果

```
###value of f(x)###
[ 90.           8.         -16.66666667]
```
手計算だと$(90,8,-50/3)$ですので，ちゃんと計算していることがわかります．

# ODEソルバーによる数値計算

では準備ができたので，これをODEソルバーに渡して数値解を計算してみます．
scipyの```solve_ivp``` のドキュメントは[ここ](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.solve_ivp.html)です．デフォルトでは(適応的時間刻み制御付）4次のルンゲ=クッタ法が実際に使われるようです．

## 時間の設定

積分の開始，終了時刻を設定しておきます．今回きれいに図を描くため時間刻み0.01ごとの値を得たいため，配列tsを作っておきます．

```
dt=0.01
tstart=0.0
tend=100.0
ts=np.arange(tstart, tend+dt, dt)  # 表示のため、0.01ステップで値を出力させる
```

## ODEソルバーに渡すラッパー関数の作成

solve_ode に渡す関数の形式にするためのラッパー関数です．引数のxtをplaceholderに渡し，fxを評価するrunを走らせるだけです． lambdaを使えばより簡潔にすることもできます．

```
def f_lorenz_tf(t,xt):
    return sess.run(fx, feed_dict={x: xt})
```

## ソルバーによる数値解の計算

 ソルバー```solve_ivp```に上で作ったベクトル場を計算する関数，初期値```x0```,積分の開始，終了時刻，値を出力する時間のリスト，を渡して数値計算を行わせ，結果などを```sol_lorenz```に受け取ります．返り値のなかの```['t']```に時間ステップ，```['y']```に解が入っているので，それぞれを取り出します．


```
start_time =time.time() # measurement of time
sol_lorenz = solve_ivp(fun=f_lorenz_tf,
                       t_span=[tstart, tend], y0=x0, t_eval=ts)
integration_time_tf = time.time() - start_time
t_lo = sol_lorenz['t']   # 各ステップの時刻を取得
x_lo = sol_lorenz['y']  # 各ステップのx(t)の値を取得

print("processing time (tf) %.5f" % integration_time_tf)
```

## 結果の表示

解曲線を表示してみます．

```
fig=plt.figure(figsize=(8,6))
ax=fig.add_subplot(111)

ax.plot(t_lo,x_lo[0],'-')
ax.set_xlabel('time')
ax.set_ylabel('x')
# 3dim phase space
fig3 = plt.figure(figsize=(8,6))
ax3 = fig3.add_subplot(111,projection='3d')

ax3.plot(x_lo[0], x_lo[1], x_lo[2], '-', lw=0.5)
ax3.set_xlabel("X Axis")
ax3.set_ylabel("Y Axis")
ax3.set_zlabel("Z Axis")
ax3.set_title("Lorenz Attractor")
```
![lorenz_x.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/311036/03a82c00-de03-f0c5-caae-52a1be4208d1.png)
![lorenz_xyz.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/311036/cb5c4cfb-3ca1-1a05-53e9-44d97a34adc6.png)

## numpyとの比較

numpyのみを使って微分方程式を計算した場合と比較してみます。
まずtensorflowを使わずにローレンツ方程式のfを定義します。扱っているものはnumpy配列かtensorかで違うのですが，コード上では最後のreturn文以外は同じ処理になります。

```
class Lorenz_np(object):

    def __init__(self, sigma=10., beta=8 / 3., rho=28., **kwargs):
        self.sigma = float(sigma)
        self.beta = float(beta)
        self.rho = float(rho)

    def __call__(self, t, x):
        """ x here is [x, y, z] """
        dxdt=  self.sigma * ( x[1] - x[0])
        dydt= x[0] * (self.rho- x[2] ) - x[1]
        dzdt= x[0]*x[1]-self.beta*x[2]
        return np.array([dxdt,dydt,dzdt   ])
```

そしてこの関数をodeソルバーに渡して計算します．

```
f_lorenz_np = Lorenz_np()
start_time = time.time()  # measurement of time
sol_lorenz_np = solve_ivp(fun=f_lorenz_np,
                          t_span=[tstart, tend], y0=x0, t_eval=ts)
integration_time_np= time.time() - start_time
print("processing time (numpy) %.5f" % integration_time_np)

t_lo_np = sol_lorenz_np['t']
x_lo_np = sol_lorenz_np['y']
```
解軌道を比較してみます．

```
fig3=plt.figure(figsize=(8,6))
ax3=fig3.add_subplot(111)

ax3.plot(t_lo,x_lo[0],'-')
ax3.plot(t_lo_np,x_lo_np[0],'-')
```
![lorenz_compare.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/311036/63101af4-8f4a-562f-d2f3-d9ee9bee31b9.png)

２つの解曲線がまったく同一なので，重なってひとつにみえています．どうやら内部的には完全に同じ計算を
しているようですね．　この方程式にはカオスの初期値鋭敏姓があるので，数学的には同一の式でも，計算の手順や精度を少しでも変えるだけでも数値計算の結果は変わってしまいます．たとえば
たとえばLorenz_npの```__call___```の１行目を

```
   dxdt=  sigma *  x[1] - sigma*x[0]
```
と少し変えるだけで，丸め誤差は発生し，その誤差が時間とともに指数的に拡大するため，数値計算の結果は目に見えて変わります．

## 速度について
コードの中には積分にかかった時間を計算していました．手元のマシンでは

```
processing time (TF): 2.97551
processing time (numpy) 0.22292
```
という結果になりました． 3秒とは...遅いですね。numpyと比べてもTensorFlowのほうが10倍遅いです。sess.run をsolve_ivpの中で10000回近く呼んでいることのオーバーヘッドでしょうか。まあnumpyもC++などに比べると100倍くらい遅いのではないかと思いますが．
尚これはCPUのみでの結果ですが、**GPUを使うとさらに遅くなります**。

#終わりに

ひとまずTensorFlowを使ってベクトル場を計算することで微分方程式の計算が可能であるということは確かめられました．

# 次回
今回の範囲ではTensorFlowを使う意義がまったくない感じられないと思いですので，次回は勾配を計算する機能を使ってヤコビ行列を計算することにより、不動点の安定性解析やリアプノフ指数の計算といった力学系の解析をやってみたいと思います．

# TODO
- バッチ処理に対応できるようにする。
- Tensorflow上でODEsolverを動かす。

# code
githubにjupyterNotebookを公開しました．
https://github.com/yymgch/ode_tf/blob/master/ode_by_tf.ipynb