pythonで二肢選択で次の探索点を求めて行くプログラムを書いていたのですが、多肢選択に変更したい。

前提・実現したいこと

タイトルの通りです。

発生している問題・エラーメッセージ

多肢選択で似たようなことをやっている人がいなく、どうしたらいいかわからない。

該当のソースコード

python
1accuracy_3D.py
2ソースコード
3
4import os
5import numpy as np
6import matplotlib.pyplot as plt
7from scipy.stats import zscore
8
9from botorch.models.pairwise_gp import PairwiseGP, PairwiseLaplaceMarginalLogLikelihood
10from botorch.fit import fit_gpytorch_model
11from botorch.acquisition import qNoisyExpectedImprovement
12from botorch.optim import optimize_acqf
13import torch
14
15from create_function import generate_data_3D, generate_comparisons_3D, utility_3D
16
17IMAGE_DIR = "./accuracy/qNEI1_qNEI2/image_3D"
18if not os.path.exists(IMAGE_DIR):
19    # ディレクトリが存在しない場合、ディレクトリを作成する
20    os.makedirs(IMAGE_DIR)
21def observation_max_points(model, results):
22    observation_point = model.posterior(results).mean.tolist()
23    next_x_index = observation_point.index(max(observation_point))
24    next_x = results[next_x_index]
25
26    return next_x
27
28def plot_3d(x, y, image_name):
29    plt.clf()
30    y_2D_12 = y.mean(axis=2)
31    plt.contourf(x, x, y_2D_12)
32    plt.savefig(image_name + "_12.png")
33    plt.clf()
34    y_2D_13 = y.mean(axis=1)
35    plt.contourf(x, x, y_2D_13)
36    plt.savefig(image_name + "_13.png")
37    plt.clf()
38    y_2D_23 = y.mean(axis=0)
39    plt.contourf(x, x, y_2D_23)
40    plt.savefig(image_name + "_23.png")
41
42def main(n, m, dim, noise, bounds):
43    mse_list = []
44    points = torch.linspace(0, 1, 11)
45    points_0, points_1, points_2 = torch.meshgrid(points, points, points)
46    points_mesh = torch.stack((points_0, points_1, points_2), dim=-1)
47
48    for i in range(10):
49        train_X, train_y = generate_data_3D(n, dim=dim)
50        train_comp = generate_comparisons_3D(train_y, m, noise=noise)
51
52        model = PairwiseGP(train_X, train_comp)
53        mll = PairwiseLaplaceMarginalLogLikelihood(model)
54        mll = fit_gpytorch_model(mll)
55        pred_y_std = zscore(model.posterior(points_mesh).mean.squeeze().detach().numpy())
56        pred_y = model.posterior(points_mesh).mean.squeeze().detach().numpy()
57        val_y = zscore(utility_3D(points_mesh)).reshape(11, 11, 11)
58
59        mse = np.square(val_y - pred_y_std).mean()
60        mse_list.append(mse)
61        print("MSE", mse)
62
63        # 検証用の関数(2次元だけ取り出してあとは平均値)
64        image_name = "./accuracy/qNEI1_qNEI2/image_3D/val_2D"
65        plot_3d(points, val_y, image_name=image_name)
66
67        # 予測平均をplot(2次元だけ取り出してあとは平均値)
68        image_name = "./accuracy/qNEI1_qNEI2/image_3D/pred_y_mean_2D_0"
69        plot_3d(points, pred_y, image_name=image_name)
70
71        for j in range(10):
72            qNEI = qNoisyExpectedImprovement(
73                model=model,
74                X_baseline=train_X.float()
75            )
76            next_X, acq_val = optimize_acqf(
77                acq_function=qNEI,
78                bounds=bounds,
79                q=2,
80                num_restarts=5,
81                raw_samples=256
82            )
83            next_X = torch.cat((next_X[0].reshape(-1, 3), next_X[1].reshape(-1, 3)), dim=0).type(torch.float64)
84            train_X = torch.cat((train_X, next_X), dim=0)
85
86            if utility_3D(next_X[0]) > utility_3D(next_X[1]):
87                preference = torch.Tensor([len(train_X) - 2, len(train_X) - 1]).long().reshape(-1, 2)
88                train_comp = torch.cat((train_comp, preference), dim=0)
89            else:
90                preference = torch.Tensor([len(train_X) - 1, len(train_X) - 2]).long().reshape(-1, 2)
91                train_comp = torch.cat((train_comp, preference), dim=0)
92
93            model = PairwiseGP(train_X, train_comp)
94            mll = PairwiseLaplaceMarginalLogLikelihood(model)
95            mll = fit_gpytorch_model(mll)
96            pred_y_std = zscore(model.posterior(points_mesh).mean.squeeze().detach().numpy())
97            pred_y = model.posterior(points_mesh).mean.squeeze().detach().numpy()
98
99            mse = np.square(val_y - pred_y_std).mean()
100            mse_list.append(mse)
101            print("MSE", i, ":", mse)
102
103            # 予測平均をplot(2次元だけ取り出してあとは平均値)
104            image_name = "./accuracy/qNEI1_qNEI2/image_3D/pred_y_mean_2D_" + str(i + 1)
105            plot_3d(points, pred_y, image_name=image_name)
106
107    number_of_comparisons = [m]
108    for k in range(1, 11):
109        number_of_comparisons.append(m + k)
110
111    mse_list = np.array(mse_list).reshape(-1, 11)
112    rmse_list = np.sqrt(mse_list)
113    mse_list_mean = mse_list.mean(axis=0)
114    rmse_list_mean = rmse_list.mean(axis=0)
115
116    plt.clf()
117    plt.plot(number_of_comparisons, mse_list_mean)
118    plt.savefig("./accuracy/qNEI1_qNEI2/image_3D/mse_mean.png")
119
120    plt.clf()
121    plt.plot(number_of_comparisons, rmse_list_mean)
122    plt.savefig("./accuracy/qNEI1_qNEI2/image_3D/rmse_mean.png")

python
1create_function.py
2ソースコード
3
4import warnings
5
6warnings.filterwarnings("ignore")
7
8from itertools import combinations
9import numpy as np
10import torch
11
12from botorch.test_functions import Hartmann
13
14def utility_6D(X):
15    """ Given X, output corresponding utility (i.e., the latent function)
16    """
17    f_test = Hartmann(dim=6)
18    X = X.reshape(-1, 6)
19    y = f_test.evaluate_true(X)
20    return y
21
22def generate_data_6D(n, dim=6):
23    """ Generate data X and y """
24    # X is randomly sampled from dim-dimentional unit cube
25    # we recommend using double as opposed to float tensor here for
26    # better numerical stability
27    X = torch.rand(n, dim, dtype=torch.float64)  # 書き換えてる
28
29    y = utility_6D(X)
30
31    return X, y
32
33def generate_comparisons_6D(y, n_comp, replace=False, noise=None):
34    """  Create pairwise comparisons with noise """
35    # generate all possible pairs of elements in y
36    all_pairs = np.array(list(combinations(range(y.shape[0]), 2)))
37    # randomly select n_comp pairs from all_pairs
38    comp_pairs = all_pairs[np.random.choice(range(len(all_pairs)), n_comp, replace=replace)]
39    # add gaussian noise to the latent y values
40    c0 = y[comp_pairs[:, 0]] + np.random.standard_normal(len(comp_pairs)) * noise
41    c1 = y[comp_pairs[:, 1]] + np.random.standard_normal(len(comp_pairs)) * noise
42    reverse_comp = (c0 < c1)
43    comp_pairs[reverse_comp, :] = np.flip(comp_pairs[reverse_comp, :], 1)
44    comp_pairs = torch.tensor(comp_pairs).long()
45
46    return comp_pairs
47
48def utility_3D(X):
49    """ Given X, output corresponding utility (i.e., the latent function)
50    """
51    f_test = Hartmann(dim=3)
52    X = X.reshape(-1, 3)
53    y = f_test.evaluate_true(X)
54    return y
55
56def generate_data_3D(n, dim=3):
57    """ Generate data X and y """
58    # X is randomly sampled from dim-dimentional unit cube
59    # we recommend using double as opposed to float tensor here for
60    # better numerical stability
61    X = torch.rand(n, dim, dtype=torch.float64)  # 書き換えてる
62
63    y = utility_3D(X)
64
65    return X, y
66
67def generate_comparisons_3D(y, n_comp, replace=False, noise=None):
68    """  Create pairwise comparisons with noise """
69    # generate all possible pairs of elements in y
70    all_pairs = np.array(list(combinations(range(y.shape[0]), 2)))
71    # randomly select n_comp pairs from all_pairs
72    comp_pairs = all_pairs[np.random.choice(range(len(all_pairs)), n_comp, replace=replace)]
73    # add gaussian noise to the latent y values
74    c0 = y[comp_pairs[:, 0]] + np.random.standard_normal(len(comp_pairs)) * noise
75    c1 = y[comp_pairs[:, 1]] + np.random.standard_normal(len(comp_pairs)) * noise
76    reverse_comp = (c0 < c1)
77    comp_pairs[reverse_comp, :] = np.flip(comp_pairs[reverse_comp, :], 1)
78    comp_pairs = torch.tensor(comp_pairs).long()
79
80    return comp_pairs
81def generate_train_x(n=2, dim=3):
82    x = torch.rand(n, dim, dtype=torch.float64)
83    return x
84
85def generate_comparisons(x1, x2):
86    if x1 > x2:
87        preference = torch.Tensor([0, 1]).reshape(1, 2)
88    elif x1 < x2:
89        preference = torch.Tensor([1, 0]).reshape(1, 2)
90    return preference

試したこと

qNEIを用いて探索点を決めて行くことにしました。

補足情報（FW/ツールのバージョンなど）

テスト関数にはハートマン3関数を、評価指標にはRMSEを用いています。
バージョンは以下の通りです。
python==3.8.5
botorch==0.3.1
gpytorch==1.2.0
matplotlib==3.3.2
numpy==1.19.2
scipy==1.5.3
torch==1.6.0