ブラックボックス最適化によるモデル超電導材料の探索¶

from amplify import FixstarsClient
from datetime import timedelta
    
# ソルバークライアントを Amplify AE に設定
client = FixstarsClient()
# ローカル環境等で実行する場合はコメントを外して Amplify AEのアクセストークンを入力してください
# client.token = "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
# 最適化の実行時間を 2 秒に設定
client.parameters.timeout = timedelta(milliseconds=2000)

import numpy as np
import math
from typing import Callable, Any
    
# 乱数シードの固定
seed = 1234
rng = np.random.default_rng(seed)
    
    
def make_blackbox_func(d: int) -> Callable[[np.ndarray], float]:
    """入力が長さ d のバリナリ値のベクトルで出力が float であるような関数を返却する"""
    
    def set_properties(size: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """（ランダムに決定した）材料物性係数を返却する"""
        mu, sigma, ratio = 0.0, 1.0, 0.2
        table1 = rng.random(size) * 1e5 * (0.1 * math.log(size) - 0.23)
        table2 = rng.lognormal(mu, sigma, size) * ratio
        table3 = rng.lognormal(mu, sigma, size) * ratio
        return table1, table2, table3
    
    def supercon_temperature(
        x: np.ndarray,
        debye_table: np.ndarray,
        state_table: np.ndarray,
        interaction_table: np.ndarray,
    ) -> float:
        """与えられた材料の組合せ（長さ d のバリナリ値のベクトル）に対し、臨界温度を計算し返却する（シミュレーションや実験の代わり）"""
        debye_temperature = np.sum(x * debye_table) / np.sum(x)
        state_density = np.sum(x * state_table) / np.sum(x)
        interaction = np.sum(x * interaction_table) / np.sum(x)
        crit_temp = debye_temperature * math.exp(-1.0 / state_density / interaction)
        return crit_temp
    
    # 係数テーブルの準備
    debye_temperature_table, state_density_table, interaction_table = set_properties(d)
    
    # ブラックボックス関数の定義
    def blackbox(x: np.ndarray) -> float:
        """与えられた材料の組合せ（長さ d のバリナリ値のベクトル）に対し、超電導臨界温度の負値を返却する"""
        assert x.shape == (d,)  # x は要素数 d の一次元配列
        t_c = supercon_temperature(
            x, debye_temperature_table, state_density_table, interaction_table
        )
        return -t_c
    
    return blackbox

num_materials = 100  # 決定変数のサイズ（材料選択肢の数）
    
    
blackbox_func = make_blackbox_func(num_materials)
    
# ランダムな入力 x で ブラックボックス関数を n_cycle 回評価し、得られた最小目的関数値と平均目的関数値を出力。
n_cycle = 100
obj_min = 0.0  # 臨界温度の負値の最小値を格納する変数
obj_mean = 0.0  # 臨界温度の負値の平均値を計算する変数
for i in range(n_cycle):
    x = rng.integers(0, 2, num_materials)
    if np.sum(x) == 0:
        continue
    obj = blackbox_func(x)
    if obj_min > obj:
        obj_min = obj
    obj_mean += obj
obj_mean /= n_cycle
    
print(f"Minimum objective function value: {obj_min:.2f} K")
print(f"Mean objective function value: {obj_mean:.2f} K")

import torch
import torch.nn as nn
    
# 乱数シードの固定
torch.manual_seed(seed)
    
    
class TorchFM(nn.Module):
    def __init__(self, d: int, k: int):
        """モデルを構築する
    
        Args:
            d (int): 入力ベクトルのサイズ
            k (int): パラメータ k
        """
        super().__init__()
        self.d = d
        self.v = torch.randn((d, k), requires_grad=True)
        self.w = torch.randn((d,), requires_grad=True)
        self.w0 = torch.randn((), requires_grad=True)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """入力 x を受け取って y の推定値を出力する
    
        Args:
            x (torch.Tensor): (データ数 × d) の 2 次元 tensor
    
        Returns:
            torch.Tensor: y の推定値 の 1次元 tensor (サイズはデータ数)
        """
        out_linear = torch.matmul(x, self.w) + self.w0
    
        out_1 = torch.matmul(x, self.v).pow(2).sum(1)
        out_2 = torch.matmul(x.pow(2), self.v.pow(2)).sum(1)
        out_quadratic = 0.5 * (out_1 - out_2)
    
        out = out_linear + out_quadratic
        return out
    
    def get_parameters(self) -> tuple[np.ndarray, np.ndarray, float]:
        """パラメータ v, w, w0 を出力する"""
        np_v = self.v.detach().numpy().copy()
        np_w = self.w.detach().numpy().copy()
        np_w0 = self.w0.detach().numpy().copy()
        return np_v, np_w, float(np_w0)

from torch.utils.data import TensorDataset, DataLoader, random_split
from tqdm.auto import tqdm, trange
import copy
    
    
def train(
    x: np.ndarray,
    y: np.ndarray,
    model: TorchFM,
    epochs: int = 2000,
    lr: float = 0.1,
) -> None:
    """FM モデルの学習を行う
    
    Args:
        x (np.ndarray): 学習データ (入力ベクトル)
        y (np.ndarray): 学習データ (出力値)
        model (TorchFM): TorchFM モデル
        epochs (int): イテレーション数
        lr (float): 学習率
    """
    
    # モデルの最適化関数
    optimizer = torch.optim.AdamW([model.v, model.w, model.w0], lr=lr)
    # 損失関数
    loss_func = nn.MSELoss()
    
    # データセットの用意
    x_tensor, y_tensor = (
        torch.from_numpy(x).float(),
        torch.from_numpy(y).float(),
    )
    dataset = TensorDataset(x_tensor, y_tensor)
    train_set, valid_set = random_split(dataset, [0.8, 0.2])
    train_loader = DataLoader(train_set, batch_size=8, shuffle=True)
    valid_loader = DataLoader(valid_set, batch_size=8, shuffle=True)
    
    # 学習の実行
    min_loss = 1e18  # 損失関数の最小値を保存
    best_state = model.state_dict()  # モデルの最も良いパラメータを保存
    
    # `range` の代わりに `tqdm` モジュールを用いて進捗を表示
    for _ in trange(epochs, leave=False):
        # 学習フェイズ
        for x_train, y_train in train_loader:
            optimizer.zero_grad()
            pred_y = model(x_train)
            loss = loss_func(pred_y, y_train)
            loss.backward()
            optimizer.step()
    
        # 検証フェイズ
        with torch.no_grad():
            loss = 0
            for x_valid, y_valid in valid_loader:
                out_valid = model(x_valid)
                loss += loss_func(out_valid, y_valid)
            if loss < min_loss:
                # 損失関数の値が更新されたらパラメータを保存
                best_state = copy.deepcopy(model.state_dict())
                min_loss = loss
    
    # モデルを学習済みパラメータで更新
    model.load_state_dict(best_state)

from amplify import VariableGenerator, Model, solve, Poly
from datetime import timedelta
    
    
def anneal(torch_model: TorchFM) -> np.ndarray:
    """FM モデルを受け取り、それらのパラメータにより記述される FM モデルの最小値を与える x を求める"""
    
    # 長さ d のバイナリ変数の配列を作成
    gen = VariableGenerator()
    x = gen.array("Binary", torch_model.d)
    
    # TorchFM からパラメータ v, w, w0 を取得
    v, w, w0 = torch_model.get_parameters()
    
    # 目的関数を作成
    out_linear = w0 + (x * w).sum()
    out_1 = ((x[:, np.newaxis] * v).sum(axis=0) ** 2).sum()  # type: ignore
    out_2 = ((x[:, np.newaxis] * v) ** 2).sum()
    objective: Poly = out_linear + (out_1 - out_2) / 2
    
    # 組合せ最適化モデルを構築
    amplify_model = Model(objective)
    
    # 最小化を実行
    result = solve(amplify_model, client)
    if len(result.solutions) == 0:
        raise RuntimeError("No solution was found.")
    
    # モデルを最小化する入力ベクトルを返却
    return x.evaluate(result.best.values).astype(int)

def init_training_data(d: int, n0: int):
    """n0 組の初期教師データを作成する"""
    assert n0 < 2**d
    
    # n0 個の 長さ d の入力値を乱数を用いて作成
    x = rng.choice(np.array([0, 1]), size=(n0, d))
    
    # 入力値の重複が発生していたらランダムに値を変更して回避する
    x = np.unique(x, axis=0)
    while x.shape[0] != n0:
        x = np.vstack((x, np.random.randint(0, 2, size=(n0 - x.shape[0], d))))
        x = np.unique(x, axis=0)
    
    # blackbox 関数を評価して入力値に対応する n0 個の出力を得る
    y = np.zeros(n0)
    for i in range(n0):
        y[i] = blackbox_func(x[i])
    
    return x, y
    
    
N0 = 10  # 初期教師データの数
x_init, y_init = init_training_data(num_materials, N0)

print(x_init.shape, y_init.shape)

# FMQA サイクルの実行回数
N = 40
    
# 教師データの初期化
x, y = x_init, y_init
    
# N - N0 回のイテレーションを実行
# `range` の代わりに `tqdm` モジュールを用いて進捗を表示
for i in trange(N):
    # 機械学習モデルの作成
    model = TorchFM(d=num_materials, k=10)
    
    # モデルの学習の実行
    train(x, y, model)
    
    # 学習済みの最小値を与える入力ベクトルの値を取得
    x_hat = anneal(model)
    
    # x_hat が重複しないようにする
    while (x_hat == x).all(axis=1).any():
        flip_idx = rng.choice(np.arange(num_materials))
        x_hat[flip_idx] = 1 - x_hat[flip_idx]
    
    # 推定された入力ベクトルを用いてブラックボックス関数を評価
    y_hat = blackbox_func(x_hat)
    
    # 評価した値をデータセットに追加
    x = np.vstack((x, x_hat))
    y = np.append(y, y_hat)
    
    tqdm.write(f"FMQA cycle {i}: found y = {y_hat:.1f}; current best = {np.min(y):.1f}")

print(x.shape, y.shape)

min_idx = np.argmin(y)
print(f"best x = {x[min_idx]}")
print(f"best y = {y[min_idx]}")

import matplotlib.pyplot as plt
    
fig = plt.figure(figsize=(6, 4))
ax = fig.add_subplot()
# 初期教師データ生成のブラックボックス関数の評価値
ax.plot(
    range(N0),
    y[:N0],
    marker="o",
    linestyle="-",
    color="b",
)
# FMQA サイクルのブラックボックス関数の評価値
ax.plot(
    range(N0, N0 + N),
    y[N0:],
    marker="o",
    linestyle="-",
    color="r",
)
# 目的関数最小値の更新履歴
ax.plot(
    range(N0 + N),
    [y[:i].min() for i in range(1, N0 + N + 1)],
    linestyle="-",
    color="k",
)
ax.set_xlabel("number of iterations", fontsize=18)
ax.set_ylabel("f(x)", fontsize=18)
ax.tick_params(labelsize=18)
plt.show()