change expected improvement to log expected improvement

- also removes log option in BO generators
xopt-org · Jan 7, 2025 · e0f545b · e0f545b
1 parent b6bcd6e
commit e0f545b
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 27 deletions.
diff --git a/xopt/generators/bayesian/bayesian_generator.py b/xopt/generators/bayesian/bayesian_generator.py
@@ -97,10 +97,6 @@ class BayesianGenerator(Generator, ABC):
     computation_time : Optional[pd.DataFrame]
         A data frame tracking computation time in seconds.
 
-    log_transform_acquisition_function: Optional[bool]
-        Flag to determine if final acquisition function value should be
-        log-transformed before optimization.
-
     n_interpolate_samples: Optional[PositiveInt]
         Number of interpolation points to generate between last observation and next
         observation, requires n_candidates to be 1.
@@ -159,16 +155,11 @@ class BayesianGenerator(Generator, ABC):
         None,
         description="data frame tracking computation time in seconds",
     )
-    log_transform_acquisition_function: Optional[bool] = Field(
-        False,
-        description="flag to log transform the acquisition function before optimization",
-    )
     custom_objective: Optional[CustomXoptObjective] = Field(
         None,
         description="custom objective for optimization, replaces objective specified by VOCS",
     )
     n_interpolate_points: Optional[PositiveInt] = None
-    memory_length: Optional[PositiveInt] = None
 
     n_candidates: int = 1
 
@@ -507,6 +498,9 @@ def get_acquisition(self, model):
                 model, acq, self._get_constraint_callables(), sampler=sampler
             )
 
+            # log transform the result to handle the constraints
+            acq = LogAcquisitionFunction(acq)
+
         # apply fixed features if specified in the generator
         if self.fixed_features is not None:
             # get input dim
@@ -521,9 +515,6 @@ def get_acquisition(self, model):
                 acq_function=acq, d=dim, columns=columns, values=values
             )
 
-        if self.log_transform_acquisition_function:
-            acq = LogAcquisitionFunction(acq)
-
         return acq
 
     def get_optimum(self):

diff --git a/xopt/generators/bayesian/expected_improvement.py b/xopt/generators/bayesian/expected_improvement.py
@@ -1,8 +1,6 @@
 import torch
 from botorch.acquisition import (
-    ExpectedImprovement,
-    qExpectedImprovement,
-    ScalarizedPosteriorTransform,
+    ScalarizedPosteriorTransform, LogExpectedImprovement, qLogExpectedImprovement,
 )
 
 from xopt.generators.bayesian.bayesian_generator import (
@@ -30,7 +28,7 @@ def _get_acquisition(self, model):
         if self.n_candidates > 1 or isinstance(objective, CustomXoptObjective):
             # MC sampling for generating multiple candidate points
             sampler = self._get_sampler(model)
-            acq = qExpectedImprovement(
+            acq = qLogExpectedImprovement(
                 model,
                 best_f=best_f,
                 sampler=sampler,
@@ -42,7 +40,7 @@ def _get_acquisition(self, model):
             # note that the analytic version cannot handle custom objectives
             weights = set_botorch_weights(self.vocs).to(**self.tkwargs)
             posterior_transform = ScalarizedPosteriorTransform(weights)
-            acq = ExpectedImprovement(
+            acq = LogExpectedImprovement(
                 model, best_f=best_f, posterior_transform=posterior_transform
             )
 

diff --git a/xopt/generators/bayesian/mggpo.py b/xopt/generators/bayesian/mggpo.py
@@ -2,7 +2,7 @@
 
 import pandas as pd
 import torch
-from botorch.acquisition.multi_objective import qNoisyExpectedHypervolumeImprovement
+from botorch.acquisition.multi_objective import qLogNoisyExpectedHypervolumeImprovement
 from pydantic import Field
 
 from xopt.generators.bayesian.objectives import create_mobo_objective
@@ -68,7 +68,7 @@ def _get_acquisition(self, model):
         inputs = self.get_input_data(self.data)
         sampler = self._get_sampler(model)
 
-        acq = qNoisyExpectedHypervolumeImprovement(
+        acq = qLogNoisyExpectedHypervolumeImprovement(
             model,
             X_baseline=inputs,
             prune_baseline=True,

diff --git a/xopt/generators/bayesian/mobo.py b/xopt/generators/bayesian/mobo.py
@@ -2,7 +2,6 @@
 
 import torch
 from botorch.acquisition import FixedFeatureAcquisitionFunction
-from botorch.acquisition.multi_objective import qNoisyExpectedHypervolumeImprovement
 from botorch.acquisition.multi_objective.logei import (
     qLogNoisyExpectedHypervolumeImprovement,
 )
@@ -24,7 +23,7 @@ class MOBOGenerator(MultiObjectiveBayesianGenerator):
         description="flag to specify if pareto front points are to be used during "
         "optimization of the acquisition function",
     )
-    __doc__ = """Implements Multi-Objective Bayesian Optimization using the Expected
+    __doc__ = """Implements Multi-Objective Bayesian Optimization using the Log Expected
             Hypervolume Improvement acquisition function"""
 
     def _get_objective(self):
@@ -60,12 +59,7 @@ def _get_acquisition(self, model):
         inputs = self.get_input_data(self.data)
         sampler = self._get_sampler(model)
 
-        if self.log_transform_acquisition_function:
-            acqclass = qLogNoisyExpectedHypervolumeImprovement
-        else:
-            acqclass = qNoisyExpectedHypervolumeImprovement
-
-        acq = acqclass(
+        acq = qLogNoisyExpectedHypervolumeImprovement(
             model,
             X_baseline=inputs,
             constraints=self._get_constraint_callables(),