ブラックボックス最適化による化学プラントにおける生産量最大化¶

import matplotlib.pyplot as plt
import numpy as np
import time
    
# ある反応性物質 (A->B) の濃度 に関する次の輸送方程式（偏微分方程式）を有限差分法で解くクラス。
    
    
class Reactor:
    def __init__(self, nfolds=5, alpha=1.0, dt_prod=1e-3, rr=1e4):
        self.nfolds = nfolds  # シミュレーション領域の空間解像度を決定するパラメータ
        self.alpha = alpha  # 分子拡散係数
        self.dt_prod = dt_prod  # 予め決められた生産時間
        self.rr = rr  # 反応速度に関する係数
    
    # 2次精度中心差分法により、分布 f の2階微分を求める関数（周期境界条件を仮定）
    
    def __dfdx2(self, f, dfdx2):
        dfdx2[0] = (f[1] - 2 * f[0] + f[self.nx - 1]) / self.dx / self.dx
        dfdx2[self.nx - 1] = (
            (f[0] - 2 * f[self.nx - 1] + f[self.nx - 2]) / self.dx / self.dx
        )
        dfdx2[1 : self.nx - 1] = (
            np.array([f[i + 1] - 2 * f[i] + f[i - 1] for i in range(1, self.nx - 1)])
            / self.dx
            / self.dx
        )
        return dfdx2
    
    # 場の初期条件を決定
    
    def __init_field(self, x):
        self.nx = self.nfolds * len(x)  # 空間離散化で用いる格子点数
        self.dx = 1.0 / (self.nx - 1)  # 格子点間隔
        self.concn_A = np.zeros(self.nx)  # Aの濃度分布
        self.concn_B = np.zeros(self.nx)  # Bの濃度分布
        self.x_cord = np.array(
            [i / self.nx - 0.5 for i in range(self.nx)]
        )  # 離散点の座標系
        self.concn_A = np.array(
            [float(x[i]) for i in range(len(x)) for j in range(self.nfolds)]
        )  # Aの初期場を作成
    
    # A の初期分布 init_A に応じて、輸送方程式を dt_prod の物理時間だけ時間発展させ、B の総生成量を返す関数
    
    def integrate(self, init_A, fig=False):
        self.__init_field(init_A)
        start = time.perf_counter()
        omega = np.zeros(self.nx)
        dfdx2 = np.zeros(self.nx)
        dt = 0.25 * self.dx * self.dx / self.alpha  # オイラー法での時間ステップ幅
        lts = int(self.dt_prod / dt)
        if fig:  # 反応経過のプロット
            fig = plt.figure(figsize=(6, 4))
            plt.tick_params(labelsize=16)
            plt.xlabel("x", fontsize=16)
            plt.ylabel("Concentration", fontsize=18)
            plt.plot(
                self.x_cord,
                self.concn_A,
                linestyle="-",
                linewidth=1,
                color=[0.6, 0.6, 0.6],
                label="$C_{A,0}$",
            )
        self.iter = 0
        while self.iter * dt < self.dt_prod:
            if fig and any(
                [
                    self.iter == i
                    for i in [0, int(0.1 * lts), int(0.2 * lts), int(0.4 * lts)]
                ]
            ):  # 反応経過のプロット
                plt.plot(
                    self.x_cord, self.concn_B, linestyle="-", linewidth=2, color="r"
                )
            omega = (
                self.rr * np.exp(-self.concn_B) * self.concn_A * (1.0 - self.concn_A)
            )  # 反応速度
            self.concn_A = (
                self.concn_A
                + (self.alpha * self.__dfdx2(self.concn_A, dfdx2) - omega) * dt
            )  # Aの濃度に関する時間発展
            self.concn_B = (
                self.concn_B
                + (self.alpha * self.__dfdx2(self.concn_B, dfdx2) + omega) * dt
            )  # Bの濃度に関する時間発展
            self.iter += 1
        if fig:  # 反応経過のプロット
            plt.plot(
                self.x_cord,
                self.concn_B,
                linestyle="-",
                linewidth=4,
                color="r",
                label="$C_B$",
            )
            plt.legend(fontsize=16)
        self.cpu_time = time.perf_counter() - start  # シミュレーションに要した時間計測
        return np.sum(self.concn_B) * self.dx  # 濃度の簡易的な空間積分

# C_A,0 を表現する1次元バイナリ配列を乱数により生成
# 1次元空間を100個の微小領域に離散化し、C_Aの初期分布を表現する
# 乱数シードを固定していないので、c_a0 の中身が実行毎に変わる
c_a0 = np.random.randint(0, 2, 100)
    
# Reactor クラスの integrate 関数により、反応シミュレーションを実施（オプションで結果画像の出力を選択）
amount_B = Reactor().integrate(c_a0, fig=True)
    
# 取得された B の総生成量を表示
print(f"Total amount of B: {amount_B:.2f}")

# 目的関数（生成された物質 B の総量の負値を返す）
def my_obj_func(init_A, fig=False):
    my_reactor = Reactor()
    minus_total = -my_reactor.integrate(init_A, fig)
    if fig:
        # （オプション）目的関数値、integrate()の積分回数、シミュレーションに要したCPU時間を表示
        print(f"{minus_total=:.2e}, {my_reactor.iter=}, {my_reactor.cpu_time=:.1e}s")
    return minus_total
    
    
# 例①：あるバイナリベクトル c_a0（上のセルで定義）を入力したときの目的関数値
amount_B = my_obj_func(c_a0)
print(f"{amount_B=:.2f}")
    
# 例②：あるバイナリベクトル c_a0（上のセルで定義）を入力したときの目的関数値（ログ及び画像表示）
amount_B = my_obj_func(c_a0, fig=True)
print(f"{amount_B=:.2f}")

import os
import torch
import numpy as np
    
    
def seed_everything(seed=0):
    os.environ["PYTHONHASHSEED"] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

from amplify import FixstarsClient
from datetime import timedelta
    
client = FixstarsClient()
timedelta(milliseconds=1000)  # 1000 ミリ秒
# client.token = "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"  # ローカル環境等で使用する場合は、Amplify AEのアクセストークンを入力してください。

import torch.nn as nn
    
    
class TorchFM(nn.Module):
    def __init__(self, d: int, k: int):
        super().__init__()
        self.V = nn.Parameter(torch.randn(d, k), requires_grad=True)
        self.lin = nn.Linear(d, 1)  # 右辺第1項及び2項は全結合ネットワーク
    
    def forward(self, x):
        out_1 = torch.matmul(x, self.V).pow(2).sum(1, keepdim=True)
        out_2 = torch.matmul(x.pow(2), self.V.pow(2)).sum(1, keepdim=True)
        out_inter = 0.5 * (out_1 - out_2)
        out_lin = self.lin(x)
        out = out_inter + out_lin
        return out

from torch.utils.data import TensorDataset, DataLoader
from sklearn.model_selection import train_test_split
from typing import Type
    
import copy
    
    
def train(
    X: np.ndarray,
    y: np.ndarray,
    model_class: Type[nn.Module],
    model_params: dict[str, int | float],
    batch_size=1024,
    epochs=3000,
    criterion=nn.MSELoss(),
    optimizer_class=torch.optim.AdamW,
    opt_params={"lr": 1},
    lr_sche_class=None,
    lr_sche_params=None,
):
    X_tensor, y_tensor = (
        torch.from_numpy(X).float(),
        torch.from_numpy(y).float(),
    )
    indices = np.array(range(X.shape[0]))
    indices_train, indices_valid = train_test_split(
        indices, test_size=0.2, random_state=42
    )
    
    train_set = TensorDataset(X_tensor[indices_train], y_tensor[indices_train])
    valid_set = TensorDataset(X_tensor[indices_valid], y_tensor[indices_valid])
    loaders = {
        "train": DataLoader(train_set, batch_size=batch_size, shuffle=True),
        "valid": DataLoader(valid_set, batch_size=batch_size, shuffle=False),
    }
    
    model = model_class(**model_params)
    best_model_wts = copy.deepcopy(model.state_dict())
    optimizer = optimizer_class(model.parameters(), **opt_params)
    scheduler = None
    if lr_sche_class is not None:
        scheduler = lr_sche_class(optimizer, **lr_sche_params)
    best_score = 1e18
    for _ in range(epochs):
        losses = {"train": 0.0, "valid": 0.0}
    
        for phase in ["train", "valid"]:
            if phase == "train":
                model.train()
            else:
                model.eval()
    
            for batch_x, batch_y in loaders[phase]:
                optimizer.zero_grad()
                out = model(batch_x).T[0]
                loss = criterion(out, batch_y)
                losses[phase] += loss.item() * batch_x.size(0)
    
                with torch.set_grad_enabled(phase == "train"):
                    if phase == "train":
                        loss.backward()
                        optimizer.step()
    
            losses[phase] /= len(loaders[phase].dataset)
    
        with torch.no_grad():
            model.eval()
            if best_score > losses["valid"]:
                best_model_wts = copy.deepcopy(model.state_dict())
                best_score = losses["valid"]
        if scheduler is not None:
            scheduler.step()
    
    with torch.no_grad():
        model.load_state_dict(best_model_wts)
        model.eval()
    return model

def gen_training_data(D: int, N0: int, true_func):
    assert N0 < 2**D
    # N0個の入力値を乱数を用いて取得
    X = np.random.randint(0, 2, 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)
    y = np.zeros(N0)
    # N0個の入力値に対応する出力値を目的関数を評価して取得
    for i in range(N0):
        if i % 10 == 0:
            print(f"Generating {i}-th training data set.")
        y[i] = true_func(X[i])
    return X, y

from amplify import VariableGenerator, solve
    
import matplotlib.pyplot as plt
import sys
    
    
class FMQA:
    def __init__(self, D: int, N: int, N0: int, k: int, true_func, client) -> None:
        assert N0 < N
        self.D = D
        self.N = N
        self.N0 = N0
        self.k = k
        self.true_func = true_func
        self.client = client
        self.y = None
    
    # 教師データに基づいて N-N0 回のFMQAを教師データを追加しながら繰り返し実施するメンバー関数
    def cycle(self, X, y, log=False) -> np.ndarray:
        print(f"Starting FMQA cycles...")
        pred_x = X[0]
        pred_y = 1e18
        for i in range(self.N - self.N0):
            print(f"FMQA Cycle #{i} ", end="")
            try:
                x_hat = self.step(X, y)
            except RuntimeError:
                sys.exit(f"Unknown error, i = {i}")
            # x_hat として既に全く同じ入力が教師データ内に存在する場合、その周辺の値を x_hat とする。
            is_identical = True
            while is_identical:
                is_identical = False
                for j in range(i + self.N0):
                    if np.all(x_hat == X[j, :]):
                        change_id = np.random.randint(0, self.D, 1)
                        x_hat[change_id.item()] = 1 - x_hat[change_id.item()]
                        if log:
                            print(f"{i=}, Identical x is found, {x_hat=}")
                        is_identical = True
                        break
            # hat{x} で目的関数 f() を評価
            y_hat = self.true_func(x_hat)
            # 最適点近傍における入出力ペア [x_hat, y_hat] を教師データに追加
            X = np.vstack((X, x_hat))
            y = np.append(y, y_hat)
            # 目的関数の評価値が最小値を更新したら、その入出力ペアを [pred_x, pred_y] へコピー
            if pred_y > y_hat:
                pred_y = y_hat
                pred_x = x_hat
                print(f"variable updated, {pred_y=}")
            else:
                print("")
            # 全ての入力を全探索済みの場合は、for文を抜ける
            if len(y) >= 2**self.D:
                print(f"Fully searched at {i=}. Terminating FMQA cycles.")
                break
        self.y = y
        return pred_x
    
    # 1回のFMQAを実施するメンバー関数
    def step(self, X, y) -> np.ndarray:
        # FM を機械学習
        model = train(
            X,
            y,
            model_class=TorchFM,
            model_params={"d": self.D, "k": self.k},
            batch_size=8,
            epochs=2000,
            criterion=nn.MSELoss(),
            optimizer_class=torch.optim.AdamW,
            opt_params={"lr": 1},
        )
        # 学習済みモデルから、FM パラメータの抽出
        v, w, w0 = list(model.parameters())
        v = v.detach().numpy()
        w = w.detach().numpy()[0]
        w0 = w0.detach().numpy()[0]
        # ここから量子アニーリング・イジングマシンによる求解を実施
        gen = VariableGenerator()  # 変数ジェネレータを宣言
        q = gen.array("Binary", self.D)  # 決定変数の作成
        model = self.__FM_as_QUBO(
            q, w0, w, v
        )  # FM パラメータから QUBO として FM を定義
        result = solve(model, self.client)  # 目的関数を Amplify のソルバーに受け渡し
        if len(result.solutions) == 0:
            raise RuntimeError("No solution was found.")
        q_values = q.evaluate(result.best.values)
        return q_values
    
    # FM パラメータから QUBO として FM を定義する関数。前定義の TorchFM クラスと同様に、g(x) の関数形通りに数式を記述。
    def __FM_as_QUBO(self, x, w0, w, v):
        lin = w0 + (x.T @ w)
        out_1 = np.array([(x * v[:, i]).sum() ** 2 for i in range(self.k)]).sum()
        # 次式において、x[j] はバイナリ変数なので、x[j] = x[j]^2 であることに注意。
        out_2 = np.array([(x * v[:, i] * v[:, i]).sum() for i in range(self.k)]).sum()
        return lin + (out_1 - out_2) / 2
    
    # 初期教師データ及び各 FMQA サイクル内で実施した i 回の目的関数評価値の履歴をプロットする関数
    def plot_history(self):
        assert self.y is not None
        fig = plt.figure(figsize=(6, 4))
        plt.plot(
            [i for i in range(self.N0)],
            self.y[: self.N0],
            marker="o",
            linestyle="-",
            color="b",
        )  # 初期教師データ生成時の目的関数評価値（ランダム過程）
        plt.plot(
            [i for i in range(self.N0, self.N)],
            self.y[self.N0 :],
            marker="o",
            linestyle="-",
            color="r",
        )  # FMQA サイクル時の目的関数評価値（FMQA サイクル過程）
        plt.xlabel("i-th evaluation of f(x)", fontsize=18)
        plt.ylabel("f(x)", fontsize=18)
        plt.tick_params(labelsize=18)
        return fig

seed_everything()  # 乱数シードを初期化
    
D = 100  # 決定変数の数（問題サイズ）
N = 30  # 目的関数の最大評価回数
N0 = 20  # 初期教師データのサンプル数
k = 20  # FM ハイパーパラメータ（FM パラメータ数に影響）
    
# 初期教師データの生成
X, y = gen_training_data(D, N0, my_obj_func)
    
# FMQA クラスのインスタンス化
fmqa_reactor = FMQA(D, N, N0, k, my_obj_func, client)
    
# FMQA サイクルの実行
pred_x = fmqa_reactor.cycle(X, y, log=True)
    
# 最適化結果の出力
print("pred x:", pred_x)
print("pred value:", my_obj_func(pred_x, fig=True))

fig = fmqa_reactor.plot_history()