Tensorflow2/KerasのRNNレイヤーを使ってカオス力学系の計算をする．

TensorFlow2のKerasによって提供されているリカレントニューラルネットワーク(RNN)クラスを使って，力学系の計算をしてみます．

RNNクラスを使ってまったくニューラルでないこともできるというのはある面で面白いのですが，実用性の評価は皆様に委ねます．

環境: Python 3.8, Tensorflow 2.4

## 参考
[TensorFlow公式ガイド Recurrent Neural Networks (RNN) with Keras](https://www.tensorflow.org/guide/keras/rnn?hl=en)
[tf.keras.layers.RNN のレファレンス](https://www.tensorflow.org/api_docs/python/tf/keras/layers/RNN)
[TensorFlow公式ガイド Advanced Automatic differentiation](https://www.tensorflow.org/guide/advanced_autodiff?hl=en)
## 準備

```python3
import numpy as np
import tensorflow as tf

from tensorflow import keras
from tensorflow.keras import layers

tf.keras.backend.set_floatx('float64')
mydtype = tf.float64
```
倍精度で計算します．


## RNN cell としてダイナミクスを定義

カスタマイズしたRNNモデルを作るには，まず1ステップ分を計算するためのクラス（RNNセルと呼ばれる）を定義しインスタンスを作ります．
今回このセルを力学系の1ステップ分を計算するものとして扱います．

cellは`state_size`, `output_size`の属性と`call`メソッドを持っておく必要があります．今回はRNNの内部状態stateをそのまま力学系の状態変数として使うので，`state_size`は力学系の次元の値に設定します． 出力は，内部状態の値そのままを出力するので，`output_size`も同様に力学系の次元の値にします．

`call`メソッドで1ステップ分の値を記述します．外部入力inputsは引数として必要ですが，今回は使いません． 
レファレンスには'A call(input_at_t, states_at_t) method, returning (output_at_t, states_at_t_plus_1). 'と書いてあるとおり，ある時刻での入力と内部状態を引数で受け取り，外部出力と次の時刻の内部状態を返すメソッドとして作成します．それぞれはバッチ処理を前提にしているので，二次元配列となり，1次元めはバッチ内のサンプル数になります．

`call`の中身では，エノン写像$f(x,y) = (1-ax^2+y, bx) $ を計算しています．バッチ処理を前提にしています．

```python3
class HenonCell(keras.layers.Layer):
  '''
  This class calculates one step of Henon-map,
  used as a RNN cell by Keras RNN layer.
  '''
  def __init__(self, a=1.2, b=0.4, **kwargs):

    # parameter
    self.a = tf.constant(a, dtype=mydtype)
    self.b = tf.constant(b, dtype=mydtype)
    # RNN custom cell must have state_size and output_size
    self.state_size = 2 # 2-dimension
    self.output_size = 2
      
    super(HenonCell, self).__init__(**kwargs)

  def build(self, input_shape):
    self.batch_size = input_shape[0]

  def call(self, inputs, states):
    ''' states is a tuple, which only has a (batch_size, 2)-array.
        inputs are ignored.
    '''
    # print(type(states)) # tuple
    # print(states[0].shape) # (2,)
    x =states[0]
    x_new = 1.0 - self.a *x[:,0]*x[:,0] + x[:,1]
    y_new = self.b * x[:,0]
    X_new = tf.stack([x_new,y_new],axis=1)
    return X_new,[X_new] # return output and states. 
```

## 1ステップの計算

cellのインスタンスに値を入れれば，写像を計算することができます．hcの返り値は出力と内部状態のタプルです．出力をx_newに入れています．

```python
if __name__ == '__main__':

  t_length = 10000 #軌道の長さ
  batch_size = 1

  hc = HenonCell(a=1.4, b=0.3)# cell
  x0 = tf.constant([[0.1, 0.1]], dtype=mydtype)  # shape is (1,2)
  x_new, _ = hc(None, [x0])  # one-step
  print(x_new)
```

## RNNレイヤーのクラスを作る

cellインスタンスを利用してRNNレイヤーのクラスを作ります．これが，軌道を一度に計算するために使えます．

```python
  henon_rnn = layers.RNN(hc, return_sequences=True)
```

## 軌道の計算
軌道を計算します．ダミーでinputをいれています．この配列の1次元めのサイズ（時間方向の長さ）で，何ステップ計算するかが決まります．
初期値は RNNをcallするときにinitial_stateで渡すことができます．


```python

 #making dummy input, which determines sequence length.
  dummy_input = tf.zeros(shape=[1, t_length, 2], dtype=mydtype)

  #initial state
  x0 = tf.constant(0.1*np.random.randn(1, 2), dtype=mydtype)
  # run
  X = henon_rnn(dummy_input, initial_state=x0)

  print(f'output sequence shape:{X.shape}')# (1,10000,2)
```

## 結果の確認

```python
  import matplotlib.pyplot as plt

  plt.plot(X[0, :, 0], X[0, :, 1], '.', markersize=0.1)
```

![henon.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/311036/b39eb1dc-3fe5-9a00-9722-b6357c632ccf.png)

## Jacobian

tf.GradientTapeを使えば，勾配を計算する場合と同じように，Jacobianを計算することができます．軌道の各点でのJacobianを一度に計算するには `batch_jacobian` を使います．([公式ガイド参照](https://www.tensorflow.org/guide/advanced_autodiff?hl=en#batch_jacobian))


```python

  X = henon_rnn(dummy_input, initial_state=x0)
  
  #reshape (bs, t_length,2) -> (bs*t_length,2)
  X = tf.reshape(X, [-1, 2])  # (batch_size*t_length, 2)

  with tf.GradientTape() as tape:
    tape.watch(X)
    X_next, _ = hc(None, [X])
  X_next.shape

  Js = tape.batch_jacobian(X_next, X)
  print(Js.shape)  # (1000,2,2)
  print(Js.numpy()[0,:,:])
```

### これはできない
RNNを動かして出てくる軌道Xから直接Jacobianを計算することはできません．

```python
  with tf.GradientTape(persistent=True) as tape:
    tape.watch(x0)
    X = henon_rnn(dummy_input, initial_state=x0)

  Jx = tape.batch_jacobian(X[:,1,:], X[:,0,:]) # 答えはNoneになってしまう．
```