逆誤差伝播法による微分のコーディング・ノート#2

## 2体間ポテンシャル実装のテスト
#1に続き，ここではより分子動力学法での応用に近づけるため，2体間ポテンシャルについて逆誤差伝播法（back propagation）を応用した微分のコーディングを紹介する．Morse型ポテンシャルを用い，Pythonでの実装を行う．

## ポテンシャル関数
Morseポテンシャル
```math
\phi_{ij} := \phi(r) := \varepsilon\{\exp[-2\alpha(r-r_0)] - 2\exp[-\alpha(r-r_0)]\}
```
を用いる．なお$r$に関する微分は
```math
\phi'(r) = -2\varepsilon\alpha\{\exp[-2\alpha(r-r_0)] - \exp[-\alpha(r-r_0)]\}
```

原子$i$のエネルギーは
```math
E_i = \frac{1}{2}\sum_{j\ne i}\phi_{ij}
```
で，系全体のエネルギーは
```math
E = \sum_i E_i
```
で与えられる．

## 計算の流れ
エネルギーを計算するためには
1. $r_{ij} = |\mathbf{r}_i-\mathbf{r}_j|$
1. $\phi_{ij} = \phi(r_{ij})$
1. $E_i = \frac{1}{2}\sum_{j\ne i}\phi_{ij}$
1. $E=\sum_i E_i$

の順に計算する．すなわち，

```math
\mathbf{r}_i, \mathbf{r}_j \to r_{ij} \to \phi_{ij} \to E_i \to E
```

のように計算していく（forward処理）．それぞれの計算に対応する関数を用意し，forward処理を実装する．

原子に働く力，すなわち
```math
\mathbf{F}_i = -\frac{\partial E}{\partial \mathbf{r}_i}
```

の計算をするため，原子座標に関するエネルギーの微分を考える．#1でやったように，各変数での微分を*をつけて表すことにする．すなわち，

```math
E_i^* := \frac{\partial E}{\partial E_i}
```
```math
\begin{align}
\phi_{ij}^* &:= \frac{\partial E}{\partial \phi_{ij}} = \sum_k \frac{\partial E}{\partial E_k}\frac{\partial E_k}{\partial \phi_{ij}} \\
&= E_i^*\frac{\partial E_i}{\partial \phi_{ij}}+E_j^*\frac{\partial E_j}{\partial \phi_{ij}}
\end{align}
```
```math
\begin{align}
r_{ij}^* &:= \frac{\partial E}{\partial r_{ij}} = \frac{\partial E}{\partial E_i}\frac{\partial E_i}{\partial r_{ij}} +\frac{\partial E}{\partial E_j}\frac{\partial E_j}{\partial r_{ij}} \\
& = \left(E_i^*\frac{\partial E_i}{\partial \phi_{ij}}+E_j^*\frac{\partial E_j}{\partial \phi_{ij}}\right)
\frac{\partial \phi_{ij}}{\partial r_{ij}} \\
&= \phi_{ij}^*\frac{\partial \phi_{ij}}{\partial r_{ij}}
\end{align}
```
```math
\mathbf{R}_l^* := \frac{\partial E}{\partial \mathbf{R}_l} = \sum_i \sum_j r_{ij}^*\frac{\partial r_{ij}}{\partial \mathbf{R}_l}
```

なお$\mathbf{F}_l = -\mathbf{R}_l^*$であることに注意する．

## Forward/Backward関数の実装

$r_{ij}, \phi_{ij}, E_i, E$ を計算する関数を用意する．Pythonコードでの関数名はそれぞれBeadR, BeadPHI, BeadEI, BeadEとする．

たとえば，BeadRのforward処理は原子座標$\mathbf{r}_i, \mathbf{r}_j$ から
```math
r_{ij} = |\mathbf{r}_i-\mathbf{r}_j|
```
を計算する．

Backward処理については，
- $E_i^*$をBeadEのbackward処理で
- $\phi_{ij}^*$をBeadEIのbackward処理で
- $r_{ij}^*$をBeadPHIのbackward処理で
- $\mathbf{R}_l^*$をBeadRのbackward処理で

求めるようにする．

### Morseポテンシャル関数

```python
import numpy as np

ep = 0.2703
al = 1.1646
ro = 3.253
ev = 1.6021892e-19

def phi(r): # Morse potential
    rr = r*1.0e10
    ex = np.exp(-al*(rr-ro))
    p = ep * (ex*ex - 2.0*ex) * ev # J
    return p

def dphi(r): # Derivative of Morse potential
    rr = r*1.0e10
    ex = np.exp(-al*(rr-ro))
    dp = -2.0 * ep * al * (ex*ex - ex) * ev *1.0e10 # J/m = N
    return dp
```

Morseポテンシャル内部のパラメータ等はeV, Åであることに注意．forwardの出力が[J]，backwardの出力が[J/m]=[N]となるように単位換算している． 

### 親クラスの準備
コンストラクタとリセット機能を書いておく．ninsがforwad処理の入力変数の次元，
noutが出力変数のそれを表す．従って，backward処理についてはこれらが逆転する．

```python
class Bead(object):
    def __init__(self, nins, nout):
        self.nins = nins
        self.nout = nout
        if (nins == 1):
            self.ar_in = 0
            self.ar_gout = 0
        else:
            self.ar_gout = np.zeros(nins)
            self.ar_in = np.zeros(nins)
        if (nout == 1):
            self.ar_out = 0
            self.ar_gin = 0
        else:
            self.ar_out = np.zeros(nout)
            self.ar_gin = np.zeros(nout)
    def reset(self):
        if (self.nout == 1):
            self.ar_out = 0
        else:
            self.ar_out[:] = 0
    def resetg(self):
        if (self.nins == 1):
            self.ar_gout = 0
        else:
            self.ar_gout[:] = 0

```

### 関数の実装(1)
BeadRのforward, backward処理を実装する．forwardのインプットとして$\\{\mathbf{r}_i, \mathbf{r}_j\\}=\\{r_i^x,r_i^y,r_i^z,r_j^x,r_j^y,r_j^z\\}$の６次元ベクトルをとる．

アウトプットは$r_{ij}$とスカラーなので１次元となる．

```math
r_{ij} = |\mathbf{r}_i-\mathbf{r}_j|
```

```math
\mathbf{R}_l^* = \sum_i \sum_{j\ne i} r_{ij}^*\frac{\partial r_{ij}}{\partial \mathbf{R}_i}
```

であったから，
```math
\mathbf{R}_i^* = \sum_k \sum_{l\ne k} r_{kl}^*\frac{\partial r_{kl}}{\partial \mathbf{R}_i}
=\sum_{j\ne i}r_{ij}^*\frac{\partial r_{ij}}{\partial \mathbf{R}_i}
+\sum_{j\ne i}r_{ji}^*\frac{\partial r_{ji}}{\partial \mathbf{R}_i}
```

各原子の$\mathbf{R}^*$を求めるには$i$と$j$の二重ループを回すことになるが，その際に
```math
r_{ij}^*\frac{\partial r_{ij}}{\partial \mathbf{R}_i} \to \mathbf{R}_i^*\quad に
```
```math
r_{ij}^*\frac{\partial r_{ij}}{\partial \mathbf{R}_j} \to \mathbf{R}_j^*\quad に
```

加算していけばよい．したがって
```math
\frac{\partial r_{ij}}{\partial \mathbf{R}_l} = 
\begin{cases}
(\mathbf{r}_i-\mathbf{r}_j)/r_{ij} \quad (l = i)\\
(\mathbf{r}_j-\mathbf{r}_i)/r_{ij} \quad (l = j)\\
0 \quad (l\ne i, l\ne j)
\end{cases}
```

を用いると，以下のように書ける．


```python
class BeadR(Bead):
    def forward(self, arri_in, arrj_in): # takes [i, r_i(x,y,z), j, r_j(x,y,z)] (6dim.) as input
        for k in range(3):
            self.ar_in[k  ] = arri_in[k]
            self.ar_in[k+3] = arrj_in[k]
        dx = arrj_in[0] - arri_in[0]
        dy = arrj_in[1] - arri_in[1]
        dz = arrj_in[2] - arri_in[2]        
        self.ar_out = np.sqrt(dx*dx + dy*dy + dz*dz)
        return self.ar_out
    def backward(self, ar_gin):
        self.ar_gin = ar_gin
        xx = ar_gin / self.ar_out 
        self.ar_gout[0] = xx * (self.ar_in[0] - self.ar_in[3]) 
        self.ar_gout[1] = xx * (self.ar_in[1] - self.ar_in[4]) 
        self.ar_gout[2] = xx * (self.ar_in[2] - self.ar_in[5]) 
        self.ar_gout[3] = -self.ar_gout[0]
        self.ar_gout[4] = -self.ar_gout[1]
        self.ar_gout[5] = -self.ar_gout[2]
        return self.ar_gout
```

### 関数の実装(2)
BeadPHIのforward, backward処理を実装する．
```math
\phi_{ij} = \phi(r_{ij})
```

また
```math
r_{ij}^* = \phi_{ij}^*\frac{\partial \phi_{ij}}{\partial r_{ij}}
```
```math
\frac{\partial \phi_{ij}}{\partial r_{ij}} = \phi'(r_{ij})
```

であるから，

```python
class BeadPHI(Bead):
    def forward(self, ar_in): # takes r_ij as input
        self.ar_in = ar_in
        self.ar_out = phi(self.ar_in)
        return self.ar_out
    def backward(self, ar_gin):
        self.ar_gin = ar_gin
        self.ar_gout = ar_gin * dphi(self.ar_in)
        return self.ar_gout
```

### 関数の実装(3)
BeadEIのforward, backward処理を実装する．

```math
E_i = \frac{1}{2}\sum_{j\ne i}\phi_{ij}
```

また
```math
\begin{align}
\phi_{ij}^* &:= \frac{\partial E}{\partial \phi_{ij}} = \sum_k \frac{\partial E}{\partial E_k}\frac{\partial E_k}{\partial \phi_{ij}} \\
&= E_i^*\frac{\partial E_i}{\partial \phi_{ij}}+E_j^*\frac{\partial E_j}{\partial \phi_{ij}}
\end{align}
```

であったが，
```math
\frac{\partial E_i}{\partial \phi_{ij}} = \frac{\partial}{\partial \phi_{ij}}
\left(\frac{1}{2} \sum_{k\ne i}\phi_{ik}\right) = \frac{1}{2}
```
```math
\frac{\partial E_j}{\partial \phi_{ij}} = \frac{\partial}{\partial \phi_{ij}}
\left(\frac{1}{2} \sum_{k\ne j}\phi_{jk}\right) = 0
```

より
```math
\phi_{ij}^* =  E_i^*\frac{\partial E_i}{\partial \phi_{ij}} = \frac{1}{2}E_i^*
```

となる．

```python
class BeadEI(Bead):
    def forward(self, ar_in): # takes phi_ij as input
        self.ar_in = ar_in
        self.ar_out += 0.5 * self.ar_in
        return self.ar_out
    def backward(self, ar_gin):
        self.ar_gin = ar_gin
        dedp = 0.5
        self.ar_gout = ar_gin * dedp
        return self.ar_gout
```

※ここで，$\phi_{ij}$の変化が$E_i$には効くが$E_j$には効かない，という点に「？」となるかもしれない．これは，
```math
E := \sum_i \sum_{j>i} \phi_{ij} \quad\quad ({\rm A})
```

ではなく
```math
E := \frac{1}{2}\sum_i \sum_{j\ne i} \phi_{ij} \quad\quad ({\rm B})
```

としていることに起因する．(B)の方法では，原子$i$と$j$の組合せが$i-j$, $j-i$と２回現れるが，$i-j$すなわち$\phi_{ij}$を求めたときに，そのエネルギーの半分を原子$i$のエネルギーのみに足し込んでいる．これは$i$と$j$のペアが作り出すエネルギーなので，残りの半分は原子$j$のエネルギーとして足されなければならないが，それは$j-i$が現れたとき（$\phi_{ji}$を求めたとき）に行われる．従って，$E_i$に寄与するのは$\phi_{ij}, \phi_{ji}$のうち$\phi_{ij}$のみである．

この方法は(A)に比べてループが約２倍になり，計算効率という点では劣るため，実際の分子動力学計算では(A)を採用するべきであるが，ここでは分かりやすさ（少なくとも筆者にとっての）を重視して(B)を採用した．

### 関数の実装(4)
BeadEのforward, backward処理を実装する．

```math
E = \sum_i E_i
```

また
```math
E_i^* = \frac{\partial E}{\partial E_i} = 1
```

であるから，
```python
class BeadE(Bead): # takes E_i as input
    def forward(self, ar_in):
        self.ar_in = ar_in
        self.ar_out += self.ar_in
        return self.ar_out
    def backward(self, ar_gin):
        self.ar_gin = ar_gin
        self.ar_gout = 1.0
        return self.ar_gout
```

### 実装した関数を用いてエネルギーと力を計算

```python
r  = BeadR(6,1)
p  = BeadPHI(1,1)
ei = BeadEI(1,1)
e  = BeadE(1,1)

e.reset()
for i in range(n):
    e.backward(1.0)
    ei.reset()
    for j in range(n):
        if (j != i):
            ri = [rx[i],ry[i],rz[i]]
            rj = [rx[j],ry[j],rz[j]]
            rij = r.forward(ri,rj)
            pij = p.forward(r.ar_out)
            eei = ei.forward(p.ar_out)
            ei.backward(e.ar_gout)
            p.backward(ei.ar_gout)
            xxx = r.backward(p.ar_gout)
            fx[i] -= xxx[0]
            fy[i] -= xxx[1]
            fz[i] -= xxx[2]
            fx[j] -= xxx[3]  
            fy[j] -= xxx[4]
            fz[j] -= xxx[5]
    ee = e.forward(ei.ar_out)
```


## 作成したコードとテスト

<details><summary>コードはこちら．</summary><div>

```python:bptest03.py
'''
Testing back-propagation for derivative calculation
03: pair potential
r_ij := |r_j - r_i|
phi_ij := phi(r_ij)
E_i := 1/2 sum_{j\ne i} phi_ij
E := \sum_i E_i
phi(r) := epsilon{exp[-2*alpha(r-r_0)] - 2*exp[-alpha(r-r_0)]}
'''

import numpy as np

ep = 0.2703
al = 1.1646
ro = 3.253
ev = 1.6021892e-19

def phi(r): # Morse potential
    rr = r*1.0e10
    ex = np.exp(-al*(rr-ro))
    p = ep * (ex*ex - 2.0*ex) * ev # J
    return p

def dphi(r): # Derivative of Morse potential
    rr = r*1.0e10
    ex = np.exp(-al*(rr-ro))
    dp = -2.0 * ep * al * (ex*ex - ex) * ev *1.0e10 # J/m = N
    return dp

class Bead(object):
    def __init__(self, nins, nout):
        self.nins = nins
        self.nout = nout
        if (nins == 1):
            self.ar_in = 0
            self.ar_gout = 0
        else:
            self.ar_gout = np.zeros(nins)
            self.ar_in = np.zeros(nins)
        if (nout == 1):
            self.ar_out = 0
            self.ar_gin = 0
        else:
            self.ar_out = np.zeros(nout)
            self.ar_gin = np.zeros(nout)
    def reset(self):
        if (self.nout == 1):
            self.ar_out = 0
        else:
            self.ar_out[:] = 0
    def resetg(self):
        if (self.nins == 1):
            self.ar_gout = 0
        else:
            self.ar_gout[:] = 0

class BeadR(Bead):
    def forward(self, arri_in, arrj_in): # takes [i, r_i(x,y,z), j, r_j(x,y,z)] (6dim.) as input
        for k in range(3):
            self.ar_in[k  ] = arri_in[k]
            self.ar_in[k+3] = arrj_in[k]
        dx = arrj_in[0] - arri_in[0]
        dy = arrj_in[1] - arri_in[1]
        dz = arrj_in[2] - arri_in[2]        
        self.ar_out = np.sqrt(dx*dx + dy*dy + dz*dz)
        return self.ar_out
    def backward(self, ar_gin):
        self.ar_gin = ar_gin
        xx = ar_gin / self.ar_out 
        self.ar_gout[0] = xx * (self.ar_in[0] - self.ar_in[3]) 
        self.ar_gout[1] = xx * (self.ar_in[1] - self.ar_in[4]) 
        self.ar_gout[2] = xx * (self.ar_in[2] - self.ar_in[5]) 
        self.ar_gout[3] = -self.ar_gout[0]
        self.ar_gout[4] = -self.ar_gout[1]
        self.ar_gout[5] = -self.ar_gout[2]
        return self.ar_gout

class BeadPHI(Bead):
    def forward(self, ar_in): # takes r_ij as input
        self.ar_in = ar_in
        self.ar_out = phi(self.ar_in)
        return self.ar_out
    def backward(self, ar_gin):
        self.ar_gin = ar_gin
        self.ar_gout = ar_gin * dphi(self.ar_in)
        return self.ar_gout

class BeadEI(Bead):
    def forward(self, ar_in): # takes phi_ij as input
        self.ar_in = ar_in
        self.ar_out += 0.5 * self.ar_in
        return self.ar_out
    def backward(self, ar_gin):
        self.ar_gin = ar_gin
        dedp = 0.5
        self.ar_gout = ar_gin * dedp
        return self.ar_gout

class BeadE(Bead): # takes E_i as input
    def forward(self, ar_in):
        self.ar_in = ar_in
        self.ar_out += self.ar_in
        return self.ar_out
    def backward(self, ar_gin):
        self.ar_gin = ar_gin
        self.ar_gout = 1.0
        return self.ar_gout

rx = [] # m
ry = []
rz = []
rx_org = []
ry_org = []
rz_org = []
vx = [] # m/s
vy = []
vz = []
fx = [] # N
fy = []
fz = []
fx_n = [] # for calculation by numerical derivative
fy_n = []
fz_n = []
f = open('initial3d.d', 'r')
n = 0 # to count the number of total atoms
for line in f:
    xy = line.split()
    rx += [float(xy[0])*1e-10]
    ry += [float(xy[1])*1e-10]
    rz += [float(xy[2])*1e-10]
    rx_org += [float(xy[0])*1e-10]
    ry_org += [float(xy[1])*1e-10]
    rz_org += [float(xy[2])*1e-10]
    vx += [float(xy[3])]
    vy += [float(xy[4])]
    vz += [float(xy[5])]
    fx += [0]
    fy += [0]
    fz += [0]
    fx_n += [0]
    fy_n += [0]
    fz_n += [0]
    n += 1
print("number of atoms = ",n)

r  = BeadR(6,1)
p  = BeadPHI(1,1)
ei = BeadEI(1,1)
e  = BeadE(1,1)


e.reset()
for i in range(n):
    e.backward(1.0)
    ei.reset()
    for j in range(n):
        if (j != i):
            ri = [rx[i],ry[i],rz[i]]
            rj = [rx[j],ry[j],rz[j]]
            rij = r.forward(ri,rj)
            pij = p.forward(r.ar_out)
            eei = ei.forward(p.ar_out)
            ei.backward(e.ar_gout)
            p.backward(ei.ar_gout)
            xxx = r.backward(p.ar_gout)
            fx[i] -= xxx[0]
            fy[i] -= xxx[1]
            fz[i] -= xxx[2]
            fx[j] -= xxx[3]  
            fy[j] -= xxx[4]
            fz[j] -= xxx[5]
    ee = e.forward(ei.ar_out)
print('Etot = ',ee)

# numerical derivative
e0 = ee
ddx = 0.01e-10
for k in range(n):
    for l in range(3):
        # plus
        for i in range(n):
            rx[i] = rx_org[i]
            ry[i] = ry_org[i]
            rz[i] = rz_org[i]
        if (l==0):
            rx[k] = rx_org[k] + ddx
        if (l==1):
            ry[k] = ry_org[k] + ddx
        if (l==2):
            rz[k] = rz_org[k] + ddx
        e.reset()
        for i in range(n):
            ei.reset()
            for j in range(n):
                if (j != i):
                    ri = [rx[i],ry[i],rz[i]]
                    rj = [rx[j],ry[j],rz[j]]
                    rij = r.forward(ri,rj)
                    pij = p.forward(r.ar_out)
                    eei = ei.forward(p.ar_out)
            eep = e.forward(ei.ar_out)
        # minus
        for i in range(n):
            rx[i] = rx_org[i]
            ry[i] = ry_org[i]
            rz[i] = rz_org[i]
        if (l==0):
            rx[k] = rx_org[k] - ddx
        if (l==1):
            ry[k] = ry_org[k] - ddx
        if (l==2):
            rz[k] = rz_org[k] - ddx
        e.reset()
        for i in range(n):
            ei.reset()
            for j in range(n):
                if (j != i):
                    ri = [rx[i],ry[i],rz[i]]
                    rj = [rx[j],ry[j],rz[j]]
                    rij = r.forward(ri,rj)
                    pij = p.forward(r.ar_out)
                    eei = ei.forward(p.ar_out)
            eem = e.forward(ei.ar_out)
        numderiv = -(eep-eem)/ddx/2.0
        if (l==0):
            fx_n[k] = numderiv
        if (l==1):
            fy_n[k] = numderiv
        if (l==2):
            fz_n[k] = numderiv
for i in range(n):
    print('Analytical:',i,fx[i],fy[i],fz[i])
    print('Numerical :',i,fx_n[i],fy_n[i],fz_n[i])
```

</div></details>


このコードは，系のエネルギーと原子に働く力を求めるとともに，原子に微小変位を加えることでエネルギーの数値微分を行い，それによって求めた力の値と比べることで，微分がうまく行っていることを確認している．
原子の配置は以下のファイルで与えられる．

```:initial3d.d
15 10 0 0 0 0
20 10 0 0 0 0
20 15 10 0 0 0
12 8 5 0 0 0
```

左からx,y,z座標（単位はÅ）を与えている．4〜6番コラムは初期速度を与えるために用意しているがここでは用いられない．上記の例では原子数は4である．

実行結果

![Screenshot from 2024-02-16 20-49-35.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/893569/7942c29c-8bd4-472f-6f3d-f500c211488c.png)

back propagationによる微分で求めた力（Analytical）と，数値微分により求めたもの（Numerical）がほぼ一致していることがわかる．