Implementation of Constrained Multi-Objective Optimization for material discovery using qNEHVI

username: Paeyus
cooked: <p>I am working on a material discovery problem where I am using qNEHVI for constrained multiobjective optimization (two objectives and one constrained), with a batch size of 3. I am a beginner with both PyTorch and Botorch, so I hope I can get help in solving my issue.</p>
<p>Here is my code</p>
<p>import pandas as pd<br>
import numpy as np<br>
import torch<br>
import gpytorch<br>
import botorch<br>
import matplotlib.pyplot as plt<br>
plt.style.use(“seaborn-v0_8-dark”)<br>
plt.rcParams[“figure.figsize”] = (8, 6)</p>
<p>from tqdm.notebook import tqdm<br>
from botorch.acquisition.monte_carlo import qExpectedImprovement<br>
from botorch.acquisition.multi_objective.monte_carlo import qNoisyExpectedHypervolumeImprovement<br>
from botorch.acquisition.multi_objective.monte_carlo import qExpectedHypervolumeImprovement</p>
<p>from botorch.acquisition.analytic import ExpectedImprovement</p>
<p>from botorch.utils.multi_objective.pareto import is_non_dominated<br>
from botorch.utils.multi_objective.box_decompositions.dominated import DominatedPartitioning<br>
from botorch.utils.multi_objective.box_decompositions.non_dominated import FastNondominatedPartitioning<br>
from botorch.models.model_list_gp_regression import ModelListGP<br>
from botorch.optim.optimize import optimize_acqf</p>
<p>import warnings</p>
<p>train_x = torch.tensor([[3., 0., 0.], [0., 0., 0.], [1., 0., 0.]]).type(torch.float64)<br>
train_y1 = torch.tensor([17189., 10068., 11139.]).to(train_x.dtype)<br>
train_y2 = torch.tensor([77., 4., 17.]).to(train_x.dtype)</p>
<h1><a name="combining-objectives-1-2-1" class="anchor" href="#combining-objectives-1-2-1"></a>combining objectives 1 &amp; 2</h1>
<p>train_Y = torch.vstack([train_y1.flatten(), train_y2.flatten()]).transpose(-1, -2)</p>
<p>constraint_Y = torch.tensor([2.5, 0., 0.]).to(train_x.dtype).unsqueeze(-1)</p>
<h1><a name="normalize-train_y-and-constraint-2" class="anchor" href="#normalize-train_y-and-constraint-2"></a>Normalize train_Y and Constraint</h1>
<p>train_y = (train_Y - train_Y.mean()) / train_Y.std()<br>
constraint = (constraint_Y - constraint_Y.mean()) / constraint_Y.std()</p>
<p>class GPModel(gpytorch.models.ExactGP, botorch.models.gpytorch.GPyTorchModel):<br>
_num_outputs = 1</p>
<pre><code>def __init__(self, train_x, train_y, likelihood):
    super().__init__(train_x, train_y, likelihood)
    self.mean_module = gpytorch.means.ConstantMean()
    self.covar_module = gpytorch.kernels.ScaleKernel(
        gpytorch.kernels.MaternKernel(nu=2.5, ard_num_dims=train_x.shape[1])
    )

def forward(self, x):
    mean_x = self.mean_module(x)
    covar_x = self.covar_module(x)
    return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
</code></pre>
<p>def fit_gp_model(train_x, train_y, num_train_iters=500):<br>
# declare the GP<br>
noise = 1e-4</p>
<pre><code>likelihood = gpytorch.likelihoods.GaussianLikelihood()
model = GPModel(train_x, train_y, likelihood)
model.likelihood.noise = noise

# train the hyperparameter (the constant)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

model.train()
likelihood.train()

for i in range(num_train_iters):
    optimizer.zero_grad()

    output = model(train_x)
    loss = -mll(output, train_y)

    loss.backward()
    optimizer.step()

model.eval()
likelihood.eval()

return model, likelihood
</code></pre>
<p>ref_point = torch.tensor([train_y[:, 0].min(), train_y[:, 1].min()])<br>
bounds = torch.tensor([[-10., -10., -10.], [10., 10., 10.]])</p>
<p>num_queries = 50<br>
num_repeats = 10<br>
batch_size = 3</p>
<p>hypervolumes =torch.zeros((num_repeats, num_queries))</p>
<p>for trial in range(num_repeats):<br>
print(“trial”, trial)</p>
<pre><code>torch.manual_seed(trial)

for i in tqdm(range(num_queries)):
    dominated_part = DominatedPartitioning(ref_point, train_y)
    hypervolumes[trial, i] = dominated_part.compute_hypervolume().item()
    
    model1, likelihood1 = fit_gp_model(train_x, train_y[:, 0].squeeze(-1))
    model2, likelihood2 = fit_gp_model(train_x, train_y[:, 1].squeeze(-1))
    
    # add the constraint
    constraint_model, constraint_likelihood = fit_gp_model(train_x, constraint_Y.squeeze(-1))
    
    
    #if strategy == qNEHVI:
        
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")

        policy = qNoisyExpectedHypervolumeImprovement(
            model=ModelListGP(model1, model2, constraint_model),
            ref_point=ref_point,
            X_baseline = train_x,
            partitioning=FastNondominatedPartitioning(ref_point, train_y),
            constraints=[lambda samples: samples[..., 2] &lt;= 6.0])

    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", category=RuntimeWarning)

        next_x, acq_val = optimize_acqf(
            acq_function=policy,
            bounds = bounds, 
            q=batch_size,
            num_restarts=100,
            raw_samples=500)
            
next_y = train_y(next_x)
next_constraint = constraint_Y(next_x)

train_x = torch.cat([train_x, next_x])
train_Y = torch.cat([train_y, next_y])
constraint = torch.cat([constraint_Y, next_constraint])
</code></pre>
<p>Here is the stack trace:</p>
<p>RuntimeError                              Traceback (most recent call last)<br>
Cell In[10], line 24<br>
21 with warnings.catch_warnings():<br>
22     warnings.simplefilter(“ignore”)<br>
—&gt; 24     policy = qNoisyExpectedHypervolumeImprovement(<br>
25         model=ModelListGP(model1, model2, constraint_model),<br>
26         ref_point=ref_point,<br>
27         X_baseline = train_x,<br>
28         partitioning=FastNondominatedPartitioning(ref_point, train_y),<br>
29         constraints=[lambda samples: samples[…, 2] &lt;= 6.0])<br>
31 with warnings.catch_warnings():<br>
32     warnings.filterwarnings(“ignore”, category=RuntimeWarning)</p>
<p>File ~\Anaconda3\envs\Project_Battery\lib\site-packages\botorch\acquisition\multi_objective\monte_carlo.py:515, in qNoisyExpectedHypervolumeImprovement.<strong>init</strong>(self, model, ref_point, X_baseline, sampler, objective, constraints, X_pending, eta, prune_baseline, alpha, cache_pending, max_iep, incremental_nehvi, cache_root, **kwargs)<br>
510 # In the case that X_pending is not None, but there are fewer than<br>
511 # max_iep pending points, the box decompositions are not performed in<br>
512 # set_X_pending. Therefore, we need to perform a box decomposition over<br>
513 # f(X_baseline) here.<br>
514 if X_pending is None or X_pending.shape[-2] &lt;= self._max_iep:<br>
 → 515     self._set_cell_bounds(num_new_points=X_baseline.shape[0])<br>
516 # Set q_in=-1 to so that self.sampler is updated at the next forward call.<br>
517 self.q_in = -1</p>
<p>File ~\Anaconda3\envs\Project_Battery\lib\site-packages\botorch\acquisition\multi_objective\monte_carlo.py:625, in qNoisyExpectedHypervolumeImprovement._set_cell_bounds(self, num_new_points)<br>
622     self.partitioning = BoxDecompositionList(*partitionings)<br>
623 else:<br>
624     # use batched partitioning<br>
 → 625     obj = _pad_batch_pareto_frontier(<br>
626         Y=obj,<br>
627         ref_point=self.ref_point.unsqueeze(0).expand(<br>
628             obj.shape[0], self.ref_point.shape[-1]<br>
629         ),<br>
630         feasibility_mask=feas,<br>
631     )<br>
632     self.partitioning = self.p_class(<br>
633         ref_point=self.ref_point, Y=obj, **self.p_kwargs<br>
634     )<br>
635 cell_bounds = self.partitioning.get_hypercell_bounds().to(self.ref_point)</p>
<p>File ~\Anaconda3\envs\Project_Battery\lib\site-packages\botorch\utils\multi_objective\box_decompositions\utils.py:73, in _pad_batch_pareto_frontier(Y, ref_point, is_pareto, feasibility_mask)<br>
66     raise UnsupportedError(<br>
67         "_pad_batch_pareto_frontier only supports a single "<br>
68         f"batch dimension, but got {len(batch_shape)} "<br>
69         “batch dimensions.”<br>
70     )<br>
71 if feasibility_mask is not None:<br>
72     # set infeasible points to be the reference point (corresponding to the batch)<br>
—&gt; 73     Y = torch.where(feasibility_mask.unsqueeze(-1), Y, ref_point)<br>
74 if not is_pareto:<br>
75     pareto_mask = is_non_dominated(Y)</p>
<p>RuntimeError: The size of tensor a (3) must match the size of tensor b (2) at non-singleton dimension 2</p>