ブラックボックス最適化と流体シミュレーションによる翼形状の最適化¶

from utils import joukowski
    
# 翼型生成クラス、joukowski.WingGenerator のインスタンス化
wing_generator = joukowski.WingGenerator()
    
# 翼型生成パラメータの整数インデックステーブルの表示
wing_generator.show_index_table()

# 翼型生成（写像変換）の実施。xi0, eta0, alpha は整数インデックスで指定。インデックスのテーブルは wing_generator.show_index_table() の出力参照
wing = wing_generator.generate_wing(10, 15, 20)
    
# 構築された翼型を表示
wing.draw()
    
# 変換の際の写像パラメータを表示
wing_generator.show_parameters("My wing")

from utils import lbm
    
# シミュレーション格子サイズ（nx：横、ny：縦）
nx, ny = 400, 160
    
# lbm.Solver クラスのインスタンス化。
solver = lbm.Solver(nx, ny)
    
# 上で構築した翼型 wing を境界条件（形状モデル）として追加
solver.add_model(wing)
    
# 時間ステップ 3000 回分のシミュレーション実施
solver.integrate(steps=3000)
    
# シミュレーション結果から、目的関数である揚抗比の負値を計算（今回の FMQA では、この値を最小化する）
cost = -solver.fy_average / solver.fx_average
    
# 目的関数の評価値を出力
print(f"{cost=:.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()
client.parameters.timeout = 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_flow_onehot(D: int, N0: int, true_func, input_generator):
    X = np.zeros((N0, D))
    y = np.zeros(N0)
    print("")
    print("###################################################")
    print(" Making initial training data")
    print("###################################################")
    for i in range(N0):
        print(f"Generating training data set #{i}.")
        # (0 or 1)の要素からなる入力ベクトルをランダムに決定
        x = input_generator()
        # 既に全く同じ入力ベクトルが与えられている場合、入力値を再生成
        is_identical = True
        while is_identical:
            is_identical = False
            for j in range(i + 1):
                if np.all(x == X[j, :]):
                    x = input_generator()
                    is_identical = True
                    break
        # 目的関数 f(x) で入力値 x を評価し、入出力ペア (x,f(x)) を教師データにコピー
        X[i, :] = x
        y[i] = true_func(x)
        print(f"------------------------")
    return X, y

from amplify import VariableGenerator, solve, one_hot, sum
import matplotlib.pyplot as plt
import sys
    
# FMQA for non-binary inputs
    
    
class FMQA_NB:
    def __init__(self, D, N, N0, k, true_func, client, encoder_decoder) -> 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.nbins_array = encoder_decoder.nbins_array
        self.encoder = encoder_decoder.encoder
        self.decoder = encoder_decoder.decoder
        self.y = None
    
    # 教師データに基づいて N-N0 回の FMQA を教師データを追加しながら繰り返し実施するメンバー関数
    def cycle(self, X, y, log=False) -> np.ndarray:
        # one-hot 制約条件の重み係数
        constraint_weight = max([1, int(np.abs(y).max() + 0.5)])
        print("")
        print("###################################################")
        print(f" Starting FMQA cycles... {constraint_weight=}")
        print("###################################################")
    
        pred_x = X[0]
        pred_y = 1e18
        for i in range(self.N - self.N0):
            print(f"FMQA Cycle #{i} ")
            try:
                x_hat = self.step(X, y, constraint_weight)
            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, :]):
                        # バイナリベクトルから整数ベクトルへデコードし、inputs へコピー
                        inputs = self.decoder(x_hat)
                        # 周辺の値を取得（inputs は整数ベクトル）
                        inputs += np.random.randint(-1, 2, len(inputs))
                        for i_inp in range(len(inputs)):
                            if inputs[i_inp] < 0:
                                inputs[i_inp] = 0
                            elif inputs[i_inp] > self.nbins_array[i_inp] - 1:
                                inputs[i_inp] = self.nbins_array[i_inp] - 1
                        # 整数ベクトルからバイナリベクトルへエンコードし、x_hat へコピー
                        x_hat = self.encoder(inputs)
                        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("")
            print(f"------------------------")
        self.y = y
        return pred_x
    
    # 1回のFMQAを実施するメンバー関数
    def step(self, X, y, constraint_weight) -> 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": 0.5},
            lr_sche_class=torch.optim.lr_scheduler.StepLR,
            lr_sche_params={"step_size": 50, "gamma": 0.9},
        )
        # 学習済みモデルから、FM パラメータの抽出
        v, w, w0 = list(model.parameters())
        v = v.detach().numpy()
        w = w.detach().numpy()[0]
        w0 = w0.detach().numpy()[0]
        # ここから量子アニーリング (QA) を実施
        gen = VariableGenerator()  # 変数ジェネレータを宣言
        q = gen.array("Binary", self.D)  # 決定変数の作成
        cost = self.__FM_as_QUBO(q, w0, w, v)  # FM パラメータから QUBO として FM を定義
        # 一つの変数が一つの値を持つように one-hot 制約
        constraint_list = []
        ista = 0
        iend = 0
        for i in range(len(self.nbins_array)):
            iend += self.nbins_array[i]
            constraint_list.append(one_hot(q[ista:iend]))
            ista = iend
        constraints = sum(constraint_list)
        # 目的関数と制約条件を足し合わし、Amplify のソルバーに受け渡し
        model = cost + constraint_weight * constraints
        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 サイクル内で実施した N 回の目的関数評価値の履歴をプロットする関数
    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 サイクル時の目的関数評価値
        plt.xlabel("i-th evaluation of f(x)", fontsize=18)
        plt.ylabel("f(x)", fontsize=18)
        plt.tick_params(labelsize=18)
        return fig

class EncDecOneHot:
    def __init__(self, D: int, nbins_array):
        self.D = D
        self.nbins_array = nbins_array
    
    def gen_random_input(self):
        ista = 0
        iend = 0
        x = np.zeros(self.D, int)
        for i in range(len(self.nbins_array)):
            iend += self.nbins_array[i]
            idx = np.random.randint(ista, iend, 1)[0]
            x[idx] = 1
            ista = iend
        return x
    
    def encoder(self, inputs):
        if len(inputs) != len(self.nbins_array):
            raise RuntimeError("inputs should be the same length as nbins_array!")
        x = np.zeros(self.D, int)
        i_offset = 0
        for i in range(len(inputs)):
            x[inputs[i] + i_offset] = 1
            i_offset += self.nbins_array[i]
        return x
    
    def decoder(self, x):
        bit_locs = np.where(x == 1)[0]
        if len(bit_locs) != len(self.nbins_array):
            raise RuntimeError("bit_locs should be the same length as nbins_array!")
        i_offset = 0
        for i in range(1, len(bit_locs)):
            i_offset += self.nbins_array[i - 1]
            bit_locs[i] -= i_offset
        return bit_locs

# 乱数シード値を初期化
seed_everything()
    
# 翼型生成クラス、joukowski.WingGenerator のインスタンス化
wing_generator = joukowski.WingGenerator()
    
# 翼型生成パラメータの整数インデックステーブルの表示
wing_generator.show_index_table()
    
# 各変数の one-hot エンコーディングの bin サイズ
nbins_array = np.array(
    [wing_generator.nbins_xi0, wing_generator.nbins_eta0, wing_generator.nbins_alpha]
)
    
# バイナリ変数の全体のサイズ（＝QUBOが取り扱うバイナリ決定変数のサイズ）
D = wing_generator.nbins_xi0 + wing_generator.nbins_eta0 + wing_generator.nbins_alpha
    
# One-hot エンコーダー・デコーダークラス（EncDecOneHot）のインスタンス化
enc_dec_one_hot = EncDecOneHot(D, nbins_array)
    
# 目的関数評価回数のカウンター
n_eval = 0
    
# 目的関数（翼型の構築 → シミュレーション実施 → 目的関数値の計算）
    
    
def obj_func(x):
    global n_eval
    # バイナリ変数から整数へデコード（写像パラメータインデックス）
    i_xi0, i_eta0, i_alpha = enc_dec_one_hot.decoder(x)
    # 翼型生成（写像変換）の実施。xi0, eta0, alpha を整数インデックスで指定。
    # インデックスのテーブルは wing_generator.show_index_table() の出力参照
    wing = wing_generator.generate_wing(i_xi0, i_eta0, i_alpha)
    # 変換の際の写像パラメータを表示
    wing_generator.show_parameters("My wing")
    # シミュレーション格子サイズ（nx：横、ny：縦）
    nx, ny = 400, 160
    # lbm.Solver クラスのインスタンス化。
    solver = lbm.Solver(nx, ny)
    # 上で構築した翼型 wing を境界条件（形状モデル）として追加
    solver.add_model(wing)
    # 時間ステップ 5000 回分のシミュレーション実施
    solver.integrate(steps=5000)
    # 揚抗比の負値
    cost = -solver.fy_average / solver.fx_average
    # 流体シミュレーションが発散した場合の処理
    if np.isnan(cost):
        cost = 0
    print(f"#{n_eval} evaluation finished. cost:{cost:.2f}")
    n_eval += 1
    return cost
    
    
N = 5  # 関数を評価できる回数。『3.3. 本サンプルコードによる FMQA 実行例』では 40 と設定
N0 = 3  # 初期教師データのサンプル数。『3.3. 本サンプルコードによる FMQA 実行例』では 20 と設定
k = 20  # FMにおけるベクトルの次元（ハイパーパラメータ）
    
# 初期教師データの生成
X, y = gen_training_data_flow_onehot(D, N0, obj_func, enc_dec_one_hot.gen_random_input)
    
# FMQA のインスタンス化
fmqa_solver = FMQA_NB(D, N, N0, k, obj_func, client, enc_dec_one_hot)
    
# FMQA サイクルの実行
pred_x = fmqa_solver.cycle(X, y)
    
# 最適化結果の出力
print("###################################################")
print(" Optimal input values and corresponding output")
print("###################################################")
print("pred x:", pred_x)
print("pred value:", obj_func(pred_x))

fig = fmqa_solver.plot_history()