Self-Consistent Histogram Methodによる自由エネルギー計算

Frenkel・Smitの教科書[1]で自由エネルギー計算法のひとつであるSelf-Consistent Histogram Method (WHAMの前身？）について勉強し、その中で紹介されている例についてPythonでシミュレーションを行った。

#問題設定
1つの理想気体粒子が以下のポテンシャル

```math
  U(z) = \begin{cases}
    z & (z>0) \\
    \infty & (z\le0)
  \end{cases}
```

にある状況を考える。逆温度$\beta$で粒子を位置$z$に見つける確率は

```math
p_0(z)=C\exp[-\beta U(z)]
```

と表せる。参考文献[1]では$z=0$を基準点としたときのLandau Free Energy[^1]として

```math
F(z) = -k_BT\ln[p_0(z)]=U(z)=z\tag{1}
```

が定義されており、この場合はポテンシャル関数と等しいことがわかる。

ポテンシャル$U(z)$のもとでモンテカルロシミュレーションを行うことで$p_0(z)$を計算することが可能である。サンプリングが十分行われていればシミュレーションによって求めた$p_0(z)$から計算した自由エネルギーは式(1)のようになるはずなので確認してみる。

モンテカルロシミュレーションのコードを以下に示す。

```python
import numpy as np
from numpy.random import *
import matplotlib.pyplot as plt
import pickle

dz = 0.05
temp = 10.0
eq = 50000
prod = 50000000
save = 100

def Monte_Carlo(zini,zmin,zmax,prod,traj):
    def Pot(x):
        return x

    z = zini
    for i in range(eq):
        zn = z + (2.0*rand()-1.0)*dz
        ez = Pot(z)
        ezn = Pot(zn)
        if (rand() <= np.exp(-temp*(ezn-ez)) and zmin < zn <= zmax):
            z = zn

    for i in range(prod):
        zn = z + (2.0*rand()-1.0)*dz
        ez = Pot(z)
        ezn = Pot(zn)
        if (rand() <= np.exp(-temp*(ezn-ez)) and zmin < zn <= zmax):
            z = zn

        if ( i % save == 0):
            traj.append(ez)

run1 = []
Monte_Carlo(0.0,0.0,10000.0,prod,run1)

r1 = np.array(run1)

with open('default.pickle','wb') as f:
    pickle.dump(r1, f)

hist = plt.hist(r1,100,range=(0.0,1.0),density=True)
plt.title(r"$\beta$ = ${{{0}}}$, step = ${{{1}}}$".format(temp,prod))
plt.xlabel(r"$z$")
plt.ylabel(r"$p_0(z)$")
plt.show()

xr = np.delete(hist[1],-1)
plt.plot(xr+0.5*xr[1],-np.ma.log(hist[0])/temp+np.log(hist[0][0])/temp,zorder=2)
plt.plot(xr,xr,c="black",lw=0.7,zorder=1,ls="--")
plt.xlabel(r"$z$")
plt.ylabel(r"$F(z)$")
plt.axes().set_aspect("equal")
plt.show()
```

$\beta=10$で5千万ステップのモンテカルロシミュレーションを行った結果を示す。
区間は$0\le z<1$でビンの数を100としている。さらに確率密度関数として規格化している（density=True）。
![fig1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/408827/ae0bb8b4-1380-0665-7153-1233e5086559.png)
![fig2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/408827/c4e2a972-142e-f071-419a-ce2a76f046ef.png)



$z>0.6$で誤差が大きく、十分サンプリングが行われていないように見える。

同じステップ数で広い範囲を十分にサンプリングをする方法としてバイアスポテンシャルを加える手法がある。例えば次のようなポテンシャル

```math
W_i(z) = \begin{cases}
    \infty & (z<z_i^{\rm min}) \\
    0 & (z_i^{\rm min} \le z < z_i^{\rm max}) \\ \tag{2}
    \infty & (z \ge z_i^{\rm max})
  \end{cases}
```

を加えるとサンプリングが区間$\ z_i^{\rm min} \le z < z_i^{\rm max}$で行われるようになる。

このことから$z$の領域をいくつかに分割し、その結果を統合して$p_0(z)$を決定するというアプローチが考えられる。今回は式(2)のバイアスポテンシャルをかけることによって区間$0\le z<1$を幅0.2ずつの区間$i(i=1,\ldots,5)$に分けてサンプリングする。

それぞれの区間ごとに計算される$p_i^{\rm est}$から$p_0$を推定するための一つの手法がSelf-Consistent Histogram Method (以下SCHM)である。SCHMの詳細については例によって参考文献[1]を見てもらうこととするが、大事なことは

* $p_0^{\rm est}$は$p_i^{\rm est}$の重ね合わせで表現できる
* $p_0^{\rm est}$の分散が最小になるよう重ね合わせの際の重みを決める

である。それぞれの区間ごとのサンプル点を揃えることにすると最終的に以下のような表式が求められる。

```math
p_0^{\rm est}(z) = \frac{\displaystyle\sum_i^n p_i^{\rm est}(z)}{\displaystyle\sum_j^n \exp[-\beta W_j(z)]\frac{Z_0}{Z_j}}\tag{3}
```

分配関数$Z_i$については後述。


# SCHMによる自由エネルギー計算
総シミュレーションステップを揃えるためにそれぞれの区間ごとに1千万ステップのシミュレーションを行った。基本的に前で書いた関数Monte_Carloを使いまわすことができる。

ただしここでは新しくパラメータ${\rm ovp}$を設定している。これは隣り合う区間との重なりの割合を示している。この操作によって区間の幅が0.2より少し広くなっている。

```python
wid = 0.2
step = 0.01
num = wid/step

prod = 10000000
ovp = 0.1 # < 0.5
result = [[] for i in range(5)]
for j in range(5):
    zmin = (j - ovp) * wid
    zmax = (j + 1 + ovp) * wid
    if (j==0):
        zmin = 0.0

    Monte_Carlo(zmin,zmin,zmax,prod,result[j])

ra = np.array(result)
with open('overlap_{}.pickle'.format(ovp),'wb') as f:
    pickle.dump(ra, f)

hist = []
for j in range(5):
    zmin = (j-ovp) * wid
    zmax = (j+1+ovp) * wid
    if (j==0):
        bnum = int(num*(1+ovp))
        zmin = 0.0
        h = plt.hist(ra[j,:],bnum,range=(zmin,zmax),normed=True,alpha=0.6,label=r"$p_1(z)$")
        hist.append(h[0])
    else:
        bnum = int(num*(1+2.0*ovp))
        h = plt.hist(ra[j,:],bnum,range=(zmin,zmax),normed=True,alpha=0.6,label=r"$p_{{{}}}(z)$".format(j+1))
        hist.append(h[0])

plt.title("ovp={}".format(ovp))
plt.xlabel(r"$z$")
plt.ylabel(r"$p(z)$")
plt.legend()
plt.show()
```

シミュレーションで得られた$p_i^{\rm est}$を以下に示す。ヒストグラムに適宜重なりが生じていることが見て取れる。
![fig3.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/408827/45576159-3cd8-621c-24de-7014d21c206f.png)



次にSCHMのコードである。まず分配関数$Z_i$の計算が必要である。式は

```math
Z_i = \int dz\ \exp[-\beta W_i(z)]\frac{\sum_j^n p_j(z)}{\sum_k^n \exp[-\beta W_k(z)]/Z_k}
```

という形になり本来は自己無撞着に解く必要があるのだが、今回は区間外で$\exp[-\beta W_i(z)]=0$になることを使うと式変形を進めることができる。最終的には漸化式の形で表せて

```math
Z_{i+1} = -\frac{Z_iZ_{i-1}A_i+(Z_iZ_{i-1}+Z_i^2)B_i}{Z_{i-1}A_i+(Z_{i-1}+Z_i)(B_i+C_i)}
```

となる。$A_i,B_i,C_i$の定義は

```math
\begin{align}
A_i &= \int_{z_i^{\rm min}}^{z_{i-1}^{\rm max}} dz\ (p_{i-1}(z) + p_i(z))\\
B_i &= -1+\int_{z_{i-1}^{\rm max}}^{z_{i+1}^{\rm min}} dz\ p_i(z) \\
C_i &= \int_{z_{i+1}^{\rm min}}^{z_{i}^{\rm max}} dz\ (p_{i}(z) + p_{i+1}(z))\\
\end{align}
```

である。これは区間$i$について隣の区間と重なりがある部分とそうでない部分に分けて積分を計算することを示している（下図参照）。また、定義から明らかであるが、$C_i=A_{i+1}$である。

![fig5.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/408827/652f0976-e8c3-4420-e7e3-1579c168b5b8.png)


$Z_i$の絶対値を決めることはできないので、$Z_1=1$としたときの相対値を求める。今回は解析的に表式が導けるので素直に計算をするだけでよい。この計算が関数Calc_Zで行われている。

続くCalc_pでは式(3)による$p(z)$の推定値が計算されている。ここで注意したいのが$Z_i$が比として決定されているために式(3)において$Z_0$を無視することができるという点である。

```python
def Calc_Z():
    k = 0
    ba = np.zeros(6)
    aa = np.zeros(6)

    ba[0] = -1.0
    for j in range(int((1.0-ovp)*num)):
        ba[0] += hist[0][j] * step

    m = j + 1

    for n in range(1,6):
        ba[n] = -1.0
        jmax = int(2.0*ovp*num)
        for j in range(jmax):
            aa[n] += (hist[n-1][-jmax+j] + hist[n][j]) * step

        m = j + 1

        for j in range(int((1.0-2.0*ovp)*num)):
            ba[n] += hist[n][m] * step
            m += 1

    za = []
    za.append(1.0)
    za.append(-(ba[0]/(ba[0]+aa[1])))

    def Func(z1,z2,p):
        A = aa[p-1]
        B = ba[p]
        return -z2*(z1*A+(z1+z2)*B)/(z1*A+(z1+z2)*(B+aa[p]))

    for i in range(1,4):
        za.append(Func(za[i-1],za[i],i+1))

    return za


def Calc_p():
    k = 0
    pa = np.zeros(int(num*5))

    for j in range(int((1.0-ovp)*num)):
        pa[k] += hist[0][j] * Z[0]
        k += 1

    for n in range(1,5):
        jmax = int(2.0*ovp*num)
        for j in range(jmax):
            pa[k] += (hist[n-1][-jmax+j] + hist[n][j]) / (1.0/Z[n-1]+1.0/Z[n])
            k += 1

        m = j + 1

        for j in range(int((1.0-2.0*ovp)*num)):
            pa[k] += hist[n][m] * Z[n]
            k += 1
            m += 1

    return  pa

Z = Calc_Z()
res = Calc_p()

xr = np.arange(0,1,0.01)
plt.title(r"Calculated $F(z)$ from $p_0^{\rm est}(z)$")
plt.plot(xr,xr,c="black",lw=0.7,ls="--")
plt.plot(xr,-np.ma.log(res)/temp+np.log(res[0])/temp)
plt.xlabel(r"$z$")
plt.ylabel(r"$F(z)$")
plt.axes().set_aspect("equal")
plt.show()
```

求められた$p_0^{\rm est}$から計算された$F(z)$は次のようになる。単一のシミュレーションでは誤差が大きかった部分が改善され、大部分が$F(z)=z$の直線状にのっていることが確認できる。

![fig4.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/408827/1acce17d-4764-dee8-1a8d-1005eac4f9cc.png)


#最後に
一粒子一次元の非常に簡単なモデルであるが、SCHMの威力を体感できたような気がする。ただ、理論化学の分野でどういった使い方がなされているかのイメージがまだつかめていないので実際の分子系のシミュレーションで何かしらの計算をやってみたい。


#参考文献
[1] D. Frenkel and B. Smit, Understanding Molecular Simulation: From Algorithms to Applications: ACADEMIC PRESS.


[^1]: Landau Free Energyという用語が厳密に何を指すのかわからなかったが、$z=0$を基準とした自由エネルギーということなら

 ```math
 -\beta\ln[C\exp[-\beta U(z)]] =-\beta \ln[C] +U(z)
     =-\beta \ln[p_0(0)] + U(z)
 ```

 に対して

 ```math
 F(z)\equiv-\beta(\ln[p_0(z)]- \ln[p_0(0)]) = U(z)
 ```

 のように考えることはできるかななどと考えていた。ちなみにプログラム中ではこの式で$F(z)$を計算している。