Thompson samplingによる離散バッチベイズ最適化

## 離散バッチベイズ最適化の需要と困難

材料科学のベイズ最適化において、1サイクルの実験で1条件のみを試すことは効率が悪く、1サイクルにつき複数条件を同時に実験したいという要望が多い。
このようなベイズ最適化をバッチベイズ最適化と呼び、`BoTroch`では`q～～`という獲得関数を利用することで対応することができる。
predictorが連続変数のときは獲得関数の最適化が容易であり、`optimize_acqf`関数により効率良く最適化できるが、材料科学においてはカテゴリ変数や離散値 (分子記述子など) がpredictorにある場合が多く、獲得関数の最適化が困難となる。
特にバッチベイズ最適化においては、考えられる条件の組合せ$N$通りからバッチ数$q$個を取り出した場合の全ての獲得関数を計算すると、$_NC_q$通りの計算が必要となり、容易に組合せ爆発を起こす。

## 簡易な対応方法

ベイズ最適化ライブラリ[BoTorchのDocument](https://botorch.org/docs/optimization)では、連続変数の最適化において、バッチ数が大きくなると逐次最適化法 (Fantasyモデルを用いた方法) が望ましいとの記載があり、この方法を離散最適化でも適用することで、離散バッチベイズ最適化に対応できると考えられる。
逐次最適化法は、[明治大学 金子先生も紹介](https://datachemeng.com/post-3347/) しており、実装については金子先生のGithubを参照いただきたい。
GPや獲得関数の計算の部分を`GPyTorch` `BoTorch`に変更することで、さらに柔軟な対応が可能となる。

## Thompson samplingによる対応

ガウス過程回帰から関数をバッチ数だけサンプリングし、各関数における最適値を次の実験候補とする方法。
この方法は、全条件の数 ($N$) に依存する $N \times N$ 行列の逆行列を計算するため、探索空間が広いとき計算困難とデメリットはあるものの、対応方法が少ない離散バッチベイズ最適化に対する解法の1つとして紹介する。
ベイズ最適化ライブラリ`Physbo`では、この逆行列計算のコストを低減しており、'BoTorch'で実装するよりも適用範囲が広がる。

## 実装 (全条件組合せ数 N=10000)

### インポート

```python
import sys
import physbo
import itertools
import numpy as np
import torch
from botorch.models import SingleTaskGP
from botorch.fit import fit_gpytorch_mll
from gpytorch.mlls import ExactMarginalLogLikelihood
from botorch.test_functions import Branin
import matplotlib.pyplot as plt
import seaborn as sns
import japanize_matplotlib

sys.path.append("src/")
import visualize

sns.set(rc={
    "savefig.bbox": "tight",
    "legend.frameon": False,
    "figure.figsize": (5, 5)
})
japanize_matplotlib.japanize()
```

```visualize.py
import torch
from botorch.test_functions import Branin
import matplotlib.pyplot as plt

def plot_Branin(
    fig=None, 
    axes=None, 
    func=Branin(negate=True), 
    ncol=1, 
    nrow=1, 
    i=0
):
    
    if (fig is None) or (axes is None):
        fig = plt.figure()
        axes = []
    
    x1 = torch.linspace(-5, 10, 100)
    x2 = torch.linspace(0, 15, 100)
    X1, X2 = torch.meshgrid(x1, x2, indexing="ij")
    X = torch.stack([X1.reshape(-1), X2.reshape(-1)], dim=-1)
    
    Z = func(X.double()).reshape(100, 100).detach().numpy()
    
    optima = torch.tensor([
        [-torch.pi, 12.275],
        [torch.pi, 2.275],
        [9.42478, 2.475]
    ])
    
    axes.append(fig.add_subplot(nrow, ncol, i+1))
    contours = axes[i].contour(X1.numpy(), X2.numpy(), Z, levels=20, cmap="viridis")
    axes[i].clabel(contours, inline=True, fontsize=8)
    axes[i].set_aspect("equal", "box")
    
    return fig, axes
```

### データの生成

```python
func = Branin(negate=True)

bounds = torch.tensor([
    [-5, 0],
    [10, 15],
], dtype=torch.float64)

choices = torch.tensor(list(itertools.product(
    torch.linspace(bounds[0, 0], bounds[1, 0], 100), 
    torch.linspace(bounds[0, 1], bounds[1, 1], 100)
)))
y = func(choices)

train_id = np.random.choice(range(choices.shape[0]), 20, replace=False)
train_X = choices[train_id]
train_y = y[train_id].unsqueeze(-1)
```

最適化対象の関数の形状は以下の通り。
```python
plot_Branin()
plt.show()
```
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3599033/45ce40ab-ee8f-4544-929e-d76bd204141f.png)

### `Physbo`による実装

```python
policy = physbo.search.discrete.policy(
    test_X=choices.numpy(), 
    initial_data=[train_id, y.numpy()[train_id]]
)

actions = policy.bayes_search(
    max_num_probes=5,
    num_search_each_probe=10, 
    simulator=None, 
    score="TS", 
    interval=0, 
    num_rand_basis=5000
)
```

```python
fig = plt.figure()
ax = []

visualize.plot_Branin(fig, ax)
ax[0].scatter(*choices[actions].T, marker="*", s=200, c="r")
ax[0].scatter(*train_X.T, c="k")
plt.show()
```
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3599033/c8ca20f5-ee22-4fa0-aa9a-e9a62d072cdc.png)


### `BoTorch` による実装

GPの学習

```python
model = SingleTaskGP(train_X, train_y)
mll = ExactMarginalLogLikelihood(model.likelihood, model)
fit_gpytorch_mll(mll)
```

全候補条件に対するThompson sampling獲得関数の計算

```python
model.eval()
with torch.no_grad():
    posterior = model.posterior(choices)
    samples = posterior.rsample(sample_shape=torch.Size([10])) 

idx = list(range(choices.shape[0]))
candidates = []

for batch in samples:
    batch = batch[idx, :]
    best_idx = batch.argmax()
    best_choice = choices[best_idx]
    idx.pop(best_idx)
    candidates.append(best_choice)

candidates = np.array(candidates)
```

```
fig = plt.figure()
ax = []

visualize.plot_Branin(fig, ax)
ax[0].scatter(*candidates.T, marker="*", s=200, c="r")
ax[0].scatter(*train_X.T, c="k")
plt.show()
```
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3599033/acb93a6e-f266-4864-8c6c-6ff1281d050e.png)

サンプリングした各関数の予測値の分布

```
ncol = 5
nrow = 2

fig = plt.figure(figsize=4 * np.array([ncol * 4/3, nrow]))
axes = []
for i, sample in enumerate(samples):
    axes.append(fig.add_subplot(nrow, ncol, i+1))
    mappable = axes[i].scatter(*choices.T, c=sample, cmap="viridis")
    plt.colorbar(mappable, ax=axes[i])
    best_idx = sample.argmax()
    axes[i].scatter(*choices[best_idx].T, marker="*", c="r", s=500)
    axes[i].set_aspect("equal", "box")
plt.tight_layout()
plt.show()
```
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3599033/8edfd89e-7011-47a2-9631-fa61f5103ae1.png)

### 2つのライブラリによる計算結果の比較

2つのライブラリによる計算結果を比較すると、`Physbo`は最適条件付近の提案数が多いが、条件の多様性に乏しいような…。
全条件数$N$が100,000になると、`Physbo`では計算できるが、`BoTorch`はメモリオーバーフロー (使用した計算機のメモリは32Gb) により計算できなかった。
$N$が1,000,000になると、`Physbo`でも計算不可。