MCMCのオーバーラップ問題

## 概要
[ゼロからできるMCMC (花田政範・松浦壮)](https://www.kspub.co.jp/book/detail/5201749.html)で紹介されていたMCMCのオーバーラップ問題への対処法を実際に実装して試してみた記事になっています.

MCMCについて分かりやすくまとまっており非常に良い本だったので, MCMCに入門したい方は購入をお勧めします.
(MCMCの各手法のコードは著者のGitHubで公開されています)

https://github.com/masanorihanada/MCMC-Sample-Codes

オーバーラップ問題への対処については実装がなかったので, 今回実際に試してみました.

## オーバーラップ問題
本と同様、2変数のメトロポリス法で試してみたいと思います.
```Python
def metropolis_2var(action, x_init=0, y_init=0, x_step=0.5, y_step=0.5, iter=100000):
    """
    2変数のメトロポリス法
    
    Arges:
        action:　作用関数
        x_init: xの初期値
        y_init: yの初期値
        x_step: xのステップサイズ 
        y_step: yのステップサイズ
        iter: イテレーション数
    """
    x = x_init
    y = y_init
    naccept = 0

    sample_x = []
    sample_y = []
    
    # MCMC実行
    for iter in range(iter):
        backup_x = x
        backup_y = y
        action_init = action(x, y)  # 更新前のS(x)
        # 一様乱数を生成
        dx = np.random.uniform(-x_step, x_step)
        dy = np.random.uniform(-y_step, y_step)
        x += dx
        y += dy
        action_fin = action(x, y) # 更新後のS(x)
        
        # メトロポリステスト
        metropolis = np.random.uniform(0,1)
        if(np.exp(action_init-action_fin) > metropolis):
            # 採択
            pass
        else:
            #　棄却
            x = backup_x
            y = backup_y
        # サンプル格納
        sample_x.append(x)
        sample_y.append(y)
    return sample_x, sample_y
```
2変数$\boldsymbol{x}=(x, y)$, 分散共分散行列$\Sigma=\begin{bmatrix} 1 & \frac{1}{2} \\\ \frac{1}{2} & 1\end{bmatrix}$の多変量ガウス分布を考えます.

```math
P(\boldsymbol{x}) = \frac{1}{(2\pi)^d \sqrt{\det{\Sigma}}}

\exp \left\{-\frac{1}{2}(\boldsymbol{x}-\boldsymbol{\mu})^{\top}\Sigma^{-1} (\boldsymbol{x}-\boldsymbol{\mu})\right\}
```
より、作用$S(x, y)$は
```math
S(x,y)= -\frac{1}{2}(\boldsymbol{x}-\boldsymbol{\mu})^{\top}\Sigma^{-1} (\boldsymbol{x}-\boldsymbol{\mu}) 

= -\frac{1}{2} 
\begin{bmatrix}x & y 
\end{bmatrix} 

\begin{bmatrix}1 & \frac{1}{2} \\\frac{1}{2} & 1 
\\\end{bmatrix}

\begin{bmatrix}x \\ y 
\end{bmatrix} 

= -\frac{x^2+y^2+xy}{2}
```
となります.

以降で$y$を$\alpha$だけずらした作用を考えるので、先に実装しておきます.
```Python
def action(alpha=0):
    def action_(x, y):
        y = y - alpha
        return 0.5*(x*x+y*y+x*y)
    return action_
```
上記のガウス分布のもとで、以下の関数の期待値を計算します.
```math
f_{\rm target}(y) = \left\{\begin{array}{ll}0 & (y\le2) \\\frac{y^2}{2} & (y>2)\end{array}\right.
```
```Python
def f_target_(x):
    return x**2 / 2 if x > 2 else 0
```
<img width="500" alt="mojikyo45_640-2.gif" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/884891/c2843ac5-87a6-3050-ca4a-137a8427a740.png">

```Python
# MCMCによる期待値計算（オリジナル）
true_e_value = 0.1305

fig = plt.figure()

# 10回実行
for j in range(10):
    # MCMCによるサンプル取得
    sample_x, sample_y = metropolis_2var(action=action(alpha=0), x_init=0, y_init=0, iter=500000)
    # 関数値に変換
    target_values = [f_target(y) for y in sample_y]
    
    # 各iterでの期待値計算(時間かかるので逐次計算で)
    e_value = target_values[0]
    e_values = [e_value]
    for i in range(1, len(target_values)-1):
        e_value = e_value + (1 / (i+1))*(target_values[i] - e_value)
        e_values.append(e_value)
        
    # プロット
    iters = list(range(len(e_values)))
    plt.plot(iters, e_values, linestyle="dashed", alpha=0.6)
    
# 真の期待値プロット
plt.plot(iters, [true_e_value]*len(iters), color='black', label='true')
plt.title('naive')
plt.xlabel('n_iter')
plt.ylabel('$\mathbb{E}_{\\alpha_0}[f]$')
plt.ylim(true_e_value-0.15, true_e_value+0.15)
plt.legend()
plt.show()
```
![naive_mcmc.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/884891/8671889e-8175-32e0-9f20-93cf8e443c90.png)

メトロポリス法で500000サンプル取得しましたが, まだ厳密値である$0.1305\cdots$には収束しきっていません.
これは, 期待値計算において重要な配位は$y>2$のサンプルですが, 確率分布は平均$0$の正規分布であるため, 分布のピークと期待値計算において重要な配位がうまく重なっていないことでサンプル効率が非常に悪くなってしまっていることに起因しています. 

## 解決方法の実装
この問題に対して, 紹介されていた解決手法を試してみました. 
**期待値計算で重要な配位と分布のピークを合わせることを考えます.**

元々の作用から$y$の値を$\alpha$だけずらした作用を以下で定義します.
（上記で実装済みです）
```math
S(x, y;\alpha) = \frac{x^2+(y-\alpha)^2+x(y-\alpha)}{2}
```
ここで, 
```math
Z_{\alpha} = \int\int e^{-S(x, y;\alpha)}dxdy
```
と定義すると, $\mathbb{E}_{\alpha=0}[f\_{\rm target}(y)]$は以下のように式変形できます.
```math
\begin{aligned}
\mathbb{E}_{\alpha=0}[f_{\rm target}(y)]
&=\frac{1}{Z_0}\int\int e^{-S(x, y;0)}dxdy

\\
&=\frac{Z_{\alpha_1}}{Z_0}\cdot\frac{Z_{\alpha_2}}{Z_{\alpha_1}}\cdots
\frac{Z_{\alpha_k}}{Z_{\alpha_{k-1}}}
\cdot\frac{1}{Z_{\alpha_k}}
\int\int e^{-S(x, y;0)}dxdy

\end{aligned}
```
このとき, 
```math
\frac{Z_{\alpha_{i+1}}}{Z_{\alpha_i}}

= \mathbb{E}_{\alpha=\alpha_i}
[
e^{-S(x, y;\alpha_{i+1}) + S(x, y;\alpha_{i})}
]

\equiv
\mathbb{E}_{\alpha=\alpha_i}
[
e^{-\Delta_{i+1,i}}
]
```
として, 以下のように変形できます.
```math
\mathbb{E}_{\alpha=0}[f_{\rm target}(y)]
=
\mathbb{E}_{\alpha=0}
[
e^{-\Delta_{1,0}}
]
\cdot
\mathbb{E}_{\alpha=\alpha_1}
[
e^{-\Delta_{2,1}}
]
\cdots
\mathbb{E}_{\alpha=\alpha_{k-1}}
[
e^{-\Delta_{k,k-1}}
]
\cdot
\mathbb{E}_{\alpha=\alpha_{k}}
[
e^{-\Delta_{0,k}}f_{\rm target}(y)
]
```
:::note info
$f_{\rm target}(y)$に関連する期待値を計算する際に, 元の分布ではなく$y$を$\alpha_k$だけずらした分布で計算することができるということですね.
:::

以下で$e^{-\Delta_{\alpha_1,\alpha_2}}$を求める関数を生成する関数を実装しています. 

```Python
def f_tmp(alpha_1, alpha_2):
    def f_tmp_(x, y):
        action_1 = action(alpha=alpha_1)
        action_2 = action(alpha=alpha_2)
        s = - action_1(x, y) + action_2(x, y)
        return np.exp(s)
    return f_tmp_
```
今回の場合は, 重要な配位と分布のピークが重なるように, $\alpha_0 = 0, \alpha_1 = 1.5, \alpha_2 = 3$として, 3回に分けて計算することで十分と紹介されていました.
(この辺は分布に関する事前知識と経験則から調整する感じですかね...?)


```Python
fig = plt.figure()

for j in range(2):
    # alpha=0の分布からサンプリング
    sample_x, sample_y = metropolis_2var(action=action(alpha=0), x_init=0, y_init=0, iter=500000)
    f = f_tmp(alpha_1=1.5, alpha_2=0)
    f_values = [f(x, y) for x, y in zip(sample_x, sample_y)]
    e_1_0 = np.mean(f_values)  # 期待値計算
    
    # alpha=1.5の分布からサンプリング
    sample_x, sample_y = metropolis_2var(action=action(alpha=1.5), x_init=0, y_init=1.5, iter=500000)
    f = f_tmp(alpha_1=3, alpha_2=1.5)
    f_values = [f(x, y) for x, y in zip(sample_x, sample_y)]
    e_2_1 = np.mean(f_values)  # 期待値計算
    
    # alpha=3の分布からサンプリング
    sample_x, sample_y = metropolis_2var(action=action(alpha=3), x_init=0, y_init=3, iter=500000)
    f = f_tmp(alpha_1=0, alpha_2=3)
    target_values = [f(x, y) * f_target(y) for x, y in zip(sample_x, sample_y)]
    
    # 各iterでの期待値計算(時間かかるので逐次計算で)
    e_value = e_1_0 * e_2_1 * target_values[0]
    e_values = [e_value]
    for i in range(1, len(target_values)-1):
        e = e_1_0 * e_2_1 * target_values[i]
        e_value = e_value + (1 / (i+1))*(target_values[i] - e_value)
        e_values.append(e_value)
        
    # プロット
    iters = list(range(len(e_values)))
    plt.plot(iters, e_values, linestyle="dashed", alpha=0.6)
    
plt.plot(iters, [true_e_value]*len(iters), color='black', label='true')
plt.title('treat_overlap_prob')
plt.xlabel('n_iter')
plt.ylabel('$\mathbb{E}_{\\alpha_0}[f]$')
plt.ylim(true_e_value-0.15, true_e_value+0.15)
plt.legend()
plt.show()
```
![treat_overlap_prob_mcmc.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/884891/847e3a90-59a5-5578-443a-147aca7044e4.png)

並べるとこんな感じでした. 
左がナイーブな期待値計算で右が３回に分けて計算した結果です. 複数回に分けてサンプリングすることで分布のピークと重要な配位が重なり収束が早くなっていることが確認できました. 

![図1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/884891/e279dca8-a38d-079c-bdbc-06a2fb392490.png)
※ 上記の例では, 単純にサンプリング回数だけで言えば, 複数回MCMCを実行しているため右の方が多くなっていることになります.  

**参考文献**
[ゼロからできるMCMC (花田政範・松浦壮)](https://www.kspub.co.jp/book/detail/5201749.html)