Merge pull request #230 from ChristopherMayes/custom_objectives

Custom objectives

mkdocs.yml

@@ -31,6 +31,7 @@ nav:
                 - Time dependent upper confidence bound: examples/single_objective_bayes_opt/time_dependent_bo.ipynb
                 - Bayesian Algorithm Execution: examples/single_objective_bayes_opt/bax_tutorial.ipynb
                 - Bayesian Optimization with fixed features: examples/single_objective_bayes_opt/fixed_features.ipynb
+                - Custom objectives: examples/single_objective_bayes_opt/custom_objective.ipynb
                 - Heteroskedastic modeling: examples/single_objective_bayes_opt/heteroskedastic_noise_tutorial.ipynb
 
             - Multi-Objective Bayesian Optimization:

tests/generators/bayesian/test_expected_improvement.py

@@ -7,6 +7,7 @@
 from xopt.base import Xopt
 from xopt.evaluator import Evaluator
 from xopt.generators.bayesian.expected_improvement import ExpectedImprovementGenerator
+from xopt.generators.bayesian.objectives import CustomXoptObjective
 from xopt.resources.testing import TEST_VOCS_BASE, TEST_VOCS_DATA, xtest_callable
 from xopt.vocs import ObjectiveEnum, VOCS
 
@@ -61,6 +62,28 @@ def test_in_xopt(self):
         for _ in range(1):
             xopt.step()
 
+    def test_custom_objectives(self):
+        train_x = torch.tensor([0.01, 0.3, 0.6, 0.99]).double()
+        train_y = torch.sin(2 * torch.pi * train_x)
+        train_c = torch.cos(2 * torch.pi * train_x)
+        train_data = pd.DataFrame(
+            {"x1": train_x.numpy(), "y1": train_y.numpy(), "c1": train_c.numpy()}
+        )
+        vocs = VOCS(**{"variables": {"x1": [0.0, 1.0]}, "observables": ["y1", "c1"]})
+
+        class MyObjective(CustomXoptObjective):
+            def forward(self, samples, X=None):
+                return samples[..., self.vocs.output_names.index("y1")] ** 2
+
+        generator = ExpectedImprovementGenerator(
+            vocs=vocs, custom_objective=MyObjective(vocs)
+        )
+        generator.add_data(train_data)
+        best_f = generator._get_best_f(generator.data, generator.custom_objective)
+        assert float(best_f) == float(torch.max(train_y**2))
+
+        generator.generate(1)
+
     def test_acquisition_accuracy(self):
         train_x = torch.tensor([0.01, 0.3, 0.6, 0.99]).double()
         train_y = torch.sin(2 * torch.pi * train_x)

xopt/generator.py

@@ -42,7 +42,7 @@ class Generator(XoptBaseModel, ABC):
 
     @field_validator("vocs", mode="after")
     def validate_vocs(cls, v, info: ValidationInfo):
-        if v.n_objectives != 1 and not info.data["supports_multi_objective"]:
+        if v.n_objectives > 1 and not info.data["supports_multi_objective"]:
             raise ValueError("this generator only supports single objective")
         return v
 

xopt/generators/__init__.py

@@ -76,35 +76,31 @@ def get_generator_dynamic(name: str):
                 "WARNING: `scipy` not found, NelderMeadGenerator and LatinHypercubeGenerator are not available"
             )
     elif name in all_generator_names["bo"]:
-        try:
-            from xopt.generators.bayesian.bayesian_exploration import (
-                BayesianExplorationGenerator,
-            )
-            from xopt.generators.bayesian.expected_improvement import (
-                ExpectedImprovementGenerator,
-            )
-            from xopt.generators.bayesian.mobo import MOBOGenerator
-            from xopt.generators.bayesian.multi_fidelity import MultiFidelityGenerator
-            from xopt.generators.bayesian.upper_confidence_bound import (
-                TDUpperConfidenceBoundGenerator,
-                UpperConfidenceBoundGenerator,
-            )
+        from xopt.generators.bayesian.bayesian_exploration import (
+            BayesianExplorationGenerator,
+        )
+        from xopt.generators.bayesian.expected_improvement import (
+            ExpectedImprovementGenerator,
+        )
+        from xopt.generators.bayesian.mobo import MOBOGenerator
+        from xopt.generators.bayesian.multi_fidelity import MultiFidelityGenerator
+        from xopt.generators.bayesian.upper_confidence_bound import (
+            TDUpperConfidenceBoundGenerator,
+            UpperConfidenceBoundGenerator,
+        )
+
+        registered_generators = [
+            UpperConfidenceBoundGenerator,
+            MOBOGenerator,
+            BayesianExplorationGenerator,
+            TDUpperConfidenceBoundGenerator,
+            ExpectedImprovementGenerator,
+            MultiFidelityGenerator,
+        ]
+        for gen in registered_generators:
+            generators[gen.name] = gen
+        return generators[name]
 
-            registered_generators = [
-                UpperConfidenceBoundGenerator,
-                MOBOGenerator,
-                BayesianExplorationGenerator,
-                TDUpperConfidenceBoundGenerator,
-                ExpectedImprovementGenerator,
-                MultiFidelityGenerator,
-            ]
-            for gen in registered_generators:
-                generators[gen.name] = gen
-            return generators[name]
-        except ModuleNotFoundError:
-            warnings.warn(
-                "WARNING: `botorch` not found, Bayesian generators are not available"
-            )
     elif name in all_generator_names["ga"]:
         try:
             from xopt.generators.ga import CNSGAGenerator

xopt/generators/bayesian/bayesian_generator.py

@@ -30,6 +30,7 @@
 from xopt.generators.bayesian.objectives import (
     create_constraint_callables,
     create_mc_objective,
+    CustomXoptObjective,
 )
 from xopt.generators.bayesian.turbo import (
     OptimizeTurboController,
@@ -160,6 +161,10 @@ class BayesianGenerator(Generator, ABC):
         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
 
     n_candidates: int = 1
@@ -550,17 +555,21 @@ def _get_acquisition(self, model):
         pass
 
     def _get_objective(self):
-        """
-        return default objective (scalar objective) determined by vocs
-
-        If objectives are specified the returned function will weight model
-        by +/- 1.0 according to MAXIMIZE/MINIMIZE keys in vocs.
-
-        If no objectives are specified, the returned function will weight observable
-        models by +1.0. This is used in Bayesian exploration.
+        """return default objective (scalar objective) determined by vocs or if
+        defined in custom_objective"""
+        # check to make sure that if we specify a custom objective that no objectives
+        # are specified in vocs
+        if self.custom_objective is not None:
+            if self.vocs.n_objectives:
+                raise RuntimeError(
+                    "cannot specify objectives in VOCS "
+                    "and a custom objective for the generator at the "
+                    "same time"
+                )
 
-        """
-        return create_mc_objective(self.vocs, self._tkwargs)
+            return self.custom_objective
+        else:
+            return create_mc_objective(self.vocs, self._tkwargs)
 
     def _get_constraint_callables(self):
         """return constratint callable determined by vocs"""

xopt/generators/bayesian/expected_improvement.py

@@ -9,6 +9,7 @@
     BayesianGenerator,
     formatted_base_docstring,
 )
+from xopt.generators.bayesian.objectives import CustomXoptObjective
 from xopt.generators.bayesian.utils import set_botorch_weights
 
 
@@ -22,24 +23,42 @@ class ExpectedImprovementGenerator(BayesianGenerator):
     )
 
     def _get_acquisition(self, model):
-        objective_data = self.vocs.objective_data(self.data, "").dropna()
-        best_f = -torch.tensor(objective_data.min().values, **self._tkwargs)
+        objective = self._get_objective()
+        best_f = self._get_best_f(self.data, objective)
 
-        if self.n_candidates > 1:
+        if self.n_candidates > 1 or isinstance(objective, CustomXoptObjective):
             # MC sampling for generating multiple candidate points
             sampler = self._get_sampler(model)
             acq = qExpectedImprovement(
                 model,
                 best_f=best_f,
                 sampler=sampler,
-                objective=self._get_objective(),
+                objective=objective,
             )
         else:
-            # analytic acquisition function for single candidate generation
+            # analytic acquisition function for single candidate generation with
+            # basic objective
+            # note that the analytic version cannot handle custom objectives
             weights = set_botorch_weights(self.vocs).to(**self._tkwargs)
             posterior_transform = ScalarizedPosteriorTransform(weights)
             acq = ExpectedImprovement(
                 model, best_f=best_f, posterior_transform=posterior_transform
             )
 
         return acq
+
+    def _get_best_f(self, data, objective):
+        """get best function value for EI based on the objective"""
+        if isinstance(objective, CustomXoptObjective):
+            best_f = objective(
+                torch.tensor(
+                    self.vocs.observable_data(data).to_numpy(), **self._tkwargs
+                )
+            ).max()
+        else:
+            # analytic acquisition function for single candidate generation
+            best_f = -torch.tensor(
+                self.vocs.objective_data(data).min().values, **self._tkwargs
+            )
+
+        return best_f

xopt/generators/bayesian/objectives.py

@@ -1,16 +1,55 @@
+from abc import ABC, abstractmethod
 from functools import partial
+from typing import Optional
 
 import torch
-from botorch.acquisition import GenericMCObjective
+from botorch.acquisition import GenericMCObjective, MCAcquisitionObjective
 from botorch.acquisition.multi_objective import WeightedMCMultiOutputObjective
 from botorch.sampling import get_sampler
+from torch import Tensor
+
+from xopt import VOCS
 
 from xopt.generators.bayesian.custom_botorch.constrained_acquisition import (
     FeasibilityObjective,
 )
 from xopt.generators.bayesian.utils import set_botorch_weights
 
 
+class CustomXoptObjective(MCAcquisitionObjective, ABC):
+    """
+    Custom objective function wrapper for use in Bayesian generators
+
+    """
+
+    def __init__(self, vocs: VOCS, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.vocs = vocs
+
+    @abstractmethod
+    def forward(self, samples: Tensor, X: Optional[Tensor] = None) -> Tensor:
+        r"""Evaluate the objective on the samples.
+
+        Args:
+            samples: A `sample_shape x batch_shape x q x m`-dim Tensors of
+                samples from a model posterior.
+            X: A `batch_shape x q x d`-dim tensor of inputs. Relevant only if
+                the objective depends on the inputs explicitly.
+
+        Returns:
+            Tensor: A `sample_shape x batch_shape x q`-dim Tensor of objective
+            values (assuming maximization).
+
+        This method is usually not called directly, but via the objectives.
+
+        Example:
+            >>> # `__call__` method:
+            >>> samples = sampler(posterior)
+            >>> outcome = mc_obj(samples)
+        """
+        pass
+
+
 def feasibility(X, model, vocs, posterior_transform=None):
     constraints = create_constraint_callables(vocs)
     posterior = model.posterior(X=X, posterior_transform=posterior_transform)

xopt/generators/bayesian/upper_confidence_bound.py

@@ -11,6 +11,7 @@
     BayesianGenerator,
     formatted_base_docstring,
 )
+from xopt.generators.bayesian.objectives import CustomXoptObjective
 from xopt.generators.bayesian.time_dependent import TimeDependentBayesianGenerator
 from xopt.generators.bayesian.utils import set_botorch_weights
 
@@ -50,7 +51,8 @@ def validate_log_transform_acquisition_function(cls, v):
             )
 
     def _get_acquisition(self, model):
-        if self.n_candidates > 1:
+        objective = self._get_objective()
+        if self.n_candidates > 1 or isinstance(objective, CustomXoptObjective):
             # MC sampling for generating multiple candidate points
             sampler = self._get_sampler(model)
             acq = qUpperConfidenceBound(

xopt/generators/bayesian/visualize.py

@@ -688,8 +688,12 @@ def _plot2d_prediction(
     kwargs = locals()
     axis.locator_params(axis="both", nbins=5)
     pcm = axis.pcolormesh(
-        input_mesh[:, vocs.variable_names.index(variable_names[0])].reshape(n_grid, n_grid).numpy(),
-        input_mesh[:, vocs.variable_names.index(variable_names[1])].reshape(n_grid, n_grid).numpy(),
+        input_mesh[:, vocs.variable_names.index(variable_names[0])]
+        .reshape(n_grid, n_grid)
+        .numpy(),
+        input_mesh[:, vocs.variable_names.index(variable_names[1])]
+        .reshape(n_grid, n_grid)
+        .numpy(),
         prediction.reshape(n_grid, n_grid),
         rasterized=True,
     )

xopt/resources/test_functions/haverly_pooling.py

@@ -0,0 +1,59 @@
+import numpy as np
+
+from xopt import VOCS
+
+"""
+C.A. Floudas, P.M. Pardalos
+A Collection of Test Problems for Constrained
+Global Optimization Algorithms, vol. 455,
+Springer Science & Business Media (1990)
+"""
+
+variables = {
+    "x1": [0, 100],
+    "x2": [0, 200],
+    "x3": [0, 100],
+    "x4": [0, 100],
+    "x5": [0, 100],
+    "x6": [0, 100],
+    "x7": [0, 200],
+    "x8": [0, 100],
+    "x9": [0, 200],
+}
+objectives = {"f": "MAXIMIZE"}
+tol = 0.5
+constraints = {
+    "h1": ["LESS_THAN", tol],
+    "h2": ["LESS_THAN", tol],
+    "h3": ["LESS_THAN", tol],
+    "h4": ["LESS_THAN", tol],
+    # "g1": ["LESS_THAN", 0.0],
+    # "g2": ["LESS_THAN", 0.0]
+}
+
+vocs_haverly = VOCS(variables=variables, objectives=objectives, constraints=constraints)
+
+
+def evaluate_haverly(input_dict):
+    x1 = input_dict["x1"]
+    x2 = input_dict["x2"]
+    x3 = input_dict["x3"]
+    x4 = input_dict["x4"]
+    x5 = input_dict["x5"]
+    x6 = input_dict["x6"]
+    x7 = input_dict["x7"]
+    x8 = input_dict["x8"]
+    x9 = input_dict["x9"]
+
+    result = {
+        "f": 9 * x1 + 15 * x2 - 6 * x3 - 16 * x4 - 10 * (x5 + x6),
+        "h1": np.abs(x7 + x8 - x4 - x3) / 100,
+        "h2": np.abs(x1 - x5 - x7) / 100,
+        "h3": np.abs(x2 - x6 - x8) / 100,
+        "h4": np.abs(x9 * x7 + x9 * x8 - 3 * x3 - x4) / 10000,
+        "g1": (x9 * x7 + 2 * x5 - 2.5 * x1) / 1000,
+        "g2": (x9 * x8 + 2 * x6 - 1.5 * x2) / 1000,
+    }
+    result["f"] = result["f"] / 1000.0
+
+    return result