import matplotlib.pyplot as plt
import numpy as np
import time
    
# A class to solve the transport equations (partial differential equations) for the concentration of a given reactive substance (A->B) by finite difference methods.
    
    
class Reactor:
    def __init__(self, nfolds=5, alpha=1.0, dt_prod=1e-3, rr=1e4):
        self.nfolds = nfolds  # Parameters that determine the spatial resolution of the simulation domain
        self.alpha = alpha  # Molecular diffusion coefficient
        self.dt_prod = dt_prod  # Predetermined production time
        self.rr = rr  # Coefficients related to the reaction rate
    
    # Function to compute the second-order derivative of the distribution f by second-order central differencing (assuming periodic boundary conditions)
    
    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
    
    # Determine the initial conditions
    
    def __init_field(self, x):
        self.nx = self.nfolds * len(
            x
        )  # Number of mesh points used in the spatial discretization
        self.dx = 1.0 / (self.nx - 1)  # Mesh point spacing
        self.concn_A = np.zeros(self.nx)  # Concentration distribution of A
        self.concn_B = np.zeros(self.nx)  # Concentration distribution of B
        self.x_cord = np.array(
            [i / self.nx - 0.5 for i in range(self.nx)]
        )  # Coordinate of discrete points
        self.concn_A = np.array(
            [float(x[i]) for i in range(len(x)) for j in range(self.nfolds)]
        )  # Generate the initial field for A
    
    # Function to evolve the transport equation in time by dt_prod physical time according to the initial distribution init_A of A and return the total amount of B produced
    
    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  # Time step width in Eulerian method
        lts = int(self.dt_prod / dt)
        if fig:  # Plot of reaction progress
            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)]
                ]
            ):  # Plot of reaction progress
                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)
            )  # Reaction rate
            self.concn_A = (
                self.concn_A
                + (self.alpha * self.__dfdx2(self.concn_A, dfdx2) - omega) * dt
            )  # Time advancement for the concentration of A
            self.concn_B = (
                self.concn_B
                + (self.alpha * self.__dfdx2(self.concn_B, dfdx2) + omega) * dt
            )  # Time advancement for the concentration of B
            self.iter += 1
        if fig:  # Plot of reaction progress
            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  # Measure the computation time
        return (
            np.sum(self.concn_B) * self.dx
        )  # Simplified spatial integration of the concentration of B

# Generate 1D binary array representing C_{A,0} by random numbers
# Discretize 1D space into 100 finite regions to represent the initial distribution of C_A
# Random seed is not fixed, so the contents of c_a0 change with each run
c_a0 = np.random.randint(0, 2, 100)
    
# Reaction simulation is performed by the `integrate`` function of the Reactor class (option to output result image)
amount_B = Reactor().integrate(c_a0, fig=True)
    
# Output the total amount of B produced
print(f"Total amount of B: {amount_B:.2f}")

# Objective function (returns the negative value of the total amount of product B produced)
def my_obj_func(init_A, fig=False):
    my_reactor = Reactor()
    minus_total = -my_reactor.integrate(init_A, fig)
    if fig:
        # (Optional) Displays objective function value, number of integrations in integrate(), and CPU time required for the simulation
        print(f"{minus_total=:.2e}, {my_reactor.iter=}, {my_reactor.cpu_time=:.1e}s")
    return minus_total
    
    
# Example 1: objective function value when a certain binary vector c_a0 (defined in the cell above) is fed
amount_B = my_obj_func(c_a0)
print(f"{amount_B=:.2f}")
    
# Example 2: Objective function value when a certain binary vector c_a0 (defined in the cell above) is fed (log and image as well)
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.client import FixstarsClient
    
client = FixstarsClient()
client.parameters.timeout = 1000  # Timeout 1s
# client.token = "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"  # If you use Amplify in a local environment or Google Colaboratory, enter the Amplify API token.

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
        )  # The first and second terms on the right-hand side are fully connected network
    
    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
    
import copy
    
    
def train(
    X,
    y,
    model_class=None,
    model_params=None,
    batch_size=1024,
    epochs=3000,
    criterion=None,
    optimizer_class=None,
    opt_params=None,
    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)
    if lr_sche_class is not None:
        scheduler = lr_sche_class(optimizer, **lr_sche_params)
    best_score = 1e18
    for epoch 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 lr_sche_class 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 input values are obtained using random numbers
    X = np.random.randint(0, 2, size=(N0, D))
    # Remove duplicate input values and add new input values using random numbers
    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)
    # Obtain output values corresponding to N0 input values by evaluating the objective function, true_func
    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 (
    Solver,
    BinarySymbolGenerator,
    sum_poly,
    BinaryMatrix,
    BinaryQuadraticModel,
)
import matplotlib.pyplot as plt
import sys
    
    
class FMQA:
    def __init__(self, D: int, N: int, N0: int, k: int, true_func, solver) -> None:
        assert N0 < N
        self.D = D
        self.N = N
        self.N0 = N0
        self.k = k
        self.true_func = true_func
        self.solver = solver
        self.y = None
    
    # A member function that repeatedly performs (N-N0)x FMQA based on the training data with adding new training data
    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}")
            # If an input value identical to the found x_hat already exists in the current training data set, a neighboring value is used as a new 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
            # Evaluate objective function f() with x_hat
            y_hat = self.true_func(x_hat)
            # Add an input-output pair [x_hat, y_hat] to the training data set
            X = np.vstack((X, x_hat))
            y = np.append(y, y_hat)
            # Copy the input-output pair to [pred_x, pred_y] when the evaluated value of the objective function updates the minimum value
            if pred_y > y_hat:
                pred_y = y_hat
                pred_x = x_hat
                print(f"variable updated, {pred_y=}")
            else:
                print("")
            # Exit the "for" statement if all inputs have been fully explored
            if len(y) >= 2**self.D:
                print(f"Fully searched at {i=}. Terminating FMQA cycles.")
                break
        self.y = y
        return pred_x
    
    # Member function to perform one FMQA cycle
    def step(self, X, y) -> np.ndarray:
        # Train 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},
        )
         # Extract FM parameters from the trained FM model
        v, w, w0 = list(model.parameters())
        v = v.detach().numpy()
        w = w.detach().numpy()[0]
        w0 = w0.detach().numpy()[0]
        # Solve a QUBO problem using a quantum annealing or Ising machine
        gen = BinarySymbolGenerator()  # Declare a variable generator, BinaryPoly
        q = gen.array(self.D)  # Generate decision variables using BinaryPoly
        cost = self.__FM_as_QUBO(
            q, w0, w, v
        )  # Define FM as a QUBO equation from FM parameters
        result = self.solver.solve(
            cost
        )  # Pass the objective function to Amplify solver
        if len(result.solutions) == 0:
            raise RuntimeError("No solution was found.")
        values = result.solutions[0].values
        q_values = q.decode(values)
        return q_values
    
    # A function that defines FM as a QUBO equation from FM parameters. As with the previously defined TorchFM class, the formula is written as per the acquisition function form of g(x).
    def __FM_as_QUBO(self, x, w0, w, v):
        lin = w0 + (x.T @ w)
        D = w.shape[0]
        out_1 = sum_poly(self.k, lambda i: sum_poly(D, lambda j: x[j] * v[j, i]) ** 2)
        # Note that x[j] = x[j]^2 since x[j] is a binary variable in the following equation.
        out_2 = sum_poly(
            self.k, lambda i: sum_poly(D, lambda j: x[j] * v[j, i] * v[j, i])
        )
        return lin + (out_1 - out_2) / 2
    
    """The sum_poly used in __FM_as_QUBO above is inefficient in terms of computation speed and 
    memory. In the case of FM, where the interaction terms of the decision variables are generally 
    nonzero, the following implementation using BinaryMatrix is more efficient. Here, the quadratic 
    terms in BinaryMatrix correspond to the non-diagonal terms represented by the upper triangular 
    matrix, so x(1/2) for the quadratic terms in the FM formula is unnecessary. Also, although x is 
    taken as an argument just to match the function signature with __FM_as_QUBO above (implementation 
    using sum_poly), it is not needed in this implementation using BinaryMatrix.
    def __FM_as_QUBO(self, x, w0, w, v):
        out_1_matrix = v @ v.T
        out_2_matrix = np.diag((v * v).sum(axis=1))
        matrix = BinaryMatrix(out_1_matrix - out_2_matrix + np.diag(w))
        # Do not forget to put the constant term w0 in the second argument of BinaryQuadraticModel.
        model = BinaryQuadraticModel(matrix, w0)
        return model
    """
    
    # A function to plot the history of i-th objective function evaluations performed within the initial training data construction (blue) and during FMQA cycles (red).
    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",
        )  # Objective function evaluation values at the time of initial training data generation (random process)
        plt.plot(
            [i for i in range(self.N0, self.N)],
            self.y[self.N0 :],
            marker="o",
            linestyle="-",
            color="r",
        )  # Objective function evaluation values during the FMQA cycles (FMQA cycle process)
        plt.xlabel("i-th evaluation of f(x)", fontsize=18)
        plt.ylabel("f(x)", fontsize=18)
        plt.tick_params(labelsize=18)
        return fig

# Initialize random seed values
seed_everything()
    
D = 100  # Size of input values (problem size)
N = 30  # Number of times the function can be evaluated
N0 = 20  # Number of samples of initial training data
k = 20  # Dimension of the vector in FM (hyperparameters)
    
# Instantiate Solver class with Amplify client created earlier
solver = Solver(client)
    
# Generate initial training data
X, y = gen_training_data(D, N0, my_obj_func)
    
# Instantiate FMQA class
fmqa_reactor = FMQA(D, N, N0, k, my_obj_func, solver)
    
# Run FMQA cycle
pred_x = fmqa_reactor.cycle(X, y, log=True)
    
# Output optimization results
print("pred x:", pred_x)
print("pred value:", my_obj_func(pred_x, fig=True))

fig = fmqa_reactor.plot_history()

Generating 0-th training data set.
Generating 10-th training data set.
Starting FMQA cycles...
FMQA Cycle #0 variable updated, pred_y=-0.0
FMQA Cycle #1 variable updated, pred_y=-0.7341318540009673
FMQA Cycle #2 variable updated, pred_y=-0.7836727189544249
FMQA Cycle #3 variable updated, pred_y=-0.7862647410081264
FMQA Cycle #4 
FMQA Cycle #5 
FMQA Cycle #6 
FMQA Cycle #7 variable updated, pred_y=-0.8310535978823115
FMQA Cycle #8 
FMQA Cycle #9 
pred x: [1. 1. 1. 1. 1. 1. 1. 0. 1. 1. 1. 1. 1. 1. 0. 1. 1. 0. 1. 1. 1. 1. 1. 1.
1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 0. 1. 1. 1. 1. 1. 1. 1. 1.
2. 1. 1. 1. 1. 1. 1. 0. 1. 1. 0. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.
3. 1. 1. 1. 1. 1. 1. 1. 0. 1. 1. 1. 0. 1. 1. 1. 1. 1. 0. 1. 1. 1. 1. 1.
4. 1. 1. 1.]
minus_total=-8.31e-01, my_reactor.iter=997, my_reactor.cpu_time=3.7e-01s
pred value: -0.8310535978823115

Black-Box Optimization of Operating Condition in a Chemical Reactor¶

1. Problem setting¶

1.1. Reactor model and optimization target¶

1.2. Description of the reactor simulator¶

1.3. Implementation of the reactor simulator¶

1.4. Simulation examples and definition of the objective function¶

2. FMQA program implementation¶

2.1．Random seed initialization¶

2.2．Configuration of Amplify client¶

2.3．Implementing FM with PyTorch¶

2.4．Construction of initial training data¶

2.5．Execution class for FMQA cycle¶

3．Search for optimal operating conditions¶

3.1. FMQA execution example¶

3.2. Transition of objective function values during the optimization process¶

3.3. Example output from this FMQA sample program¶