Add EI and Branin example in (#715)

* create branch

* multiplying the EI acquisition function since we are minimizing it

* minor fix

---------

Co-authored-by: Tucker Hartland <tucker.hartland@gmail.com>

Author: nychiang

pyproject.toml

@@ -1,3 +1,3 @@
 [build-system]
-requires = ["setuptools", "wheel", "numpy"]
+requires = ["setuptools", "wheel", "numpy", "smt", "cyipopt"]
 build-backend = "setuptools.build_meta"

setup.py

@@ -12,14 +12,14 @@
 
 metadata = dict(
         name="hiopbbpy",
-        version="0.0.1",
+        version="0.0.2",
         description="HiOp black box optimization (hiopbbpy)",
         author="Tucker hartland et al.",
         author_email="hartland1@llnl.gov",
         license="BSD-3",
         packages=find_packages(where="src"),
         package_dir={"": "src"},
-        install_requires=["smt"],
+        install_requires=["smt", "cyipopt"],
         python_requires=">=3.9",
         zip_safe=False,
         url="https://github.com/LLNL/hiop",

src/Drivers/hiopbbpy/BODriver.py

@@ -14,6 +14,7 @@
 from LpNormProblem import LpNormProblem
 from hiopbbpy.surrogate_modeling import smtKRG
 from hiopbbpy.opt import BOAlgorithm
+from hiopbbpy.problems import BraninProblem
 
 
 ### parameters
@@ -23,7 +24,9 @@
 nx = 2 # dimension of the problem
 xlimits = np.array([[-5, 5], [-5, 5]]) # bounds on optimization variable
 
-problem = LpNormProblem(nx, xlimits)
+#problem = LpNormProblem(nx, xlimits)
+problem = BraninProblem()
+
 print(problem.name, " problem")
 
 ### initial training set
@@ -34,9 +37,14 @@
 gp_model = smtKRG(theta, xlimits, nx)
 gp_model.train(x_train, y_train)
 
-# Instantiate and run Bayesian Optimization
-bo = BOAlgorithm(gp_model, x_train, y_train)
-bo.optimize(problem)
-
-# Retrieve optimal point
-x_opt, y_opt = bo.getOptimalPoint()
+acquisition_types = ["LCB", "EI"]
+for acquisition_type in acquisition_types:
+    print("acquisition type: ", acquisition_type)
+    
+    # Instantiate and run Bayesian Optimization
+    bo = BOAlgorithm(gp_model, x_train, y_train, acquisition_type = acquisition_type) #EI or LCB
+    bo.optimize(problem)
+    
+    # Retrieve optimal point
+    x_opt, y_opt = bo.getOptimalPoint()
+    print()

src/hiopbbpy/opt/__init__.py

@@ -1,9 +1,10 @@
 from .boalgorithm import (BOAlgorithmBase, BOAlgorithm)
-from .acquisition import (acquisition, LCBacquisition)
+from .acquisition import (acquisition, LCBacquisition, EIacquisition)
 
 __all__ = [
         "BOAlgorithmBase"
         "BOAlgorithm"
         "acquisition"
         "LCBacquisition"
+        "EIacquisition"
         ]

src/hiopbbpy/opt/acquisition.py

@@ -6,6 +6,7 @@
 """
 
 import numpy as np
+from scipy.stats import norm
 from ..surrogate_modeling.gp import GaussianProcess
 
 # A base class for acquisition functions
@@ -30,3 +31,28 @@ def evaluate(self, x : np.ndarray) -> np.ndarray:
         mu = self.gpsurrogate.mean(x)
         sig = self.gpsurrogate.variance(x)
         return mu - self.beta * np.sqrt(sig)
+
+# A subclass of acquisition, implementing the Expected improvement (EI) acquisition function.
+class EIacquisition(acquisition):
+    def __init__(self, gpsurrogate):
+        super().__init__(gpsurrogate)
+
+    # Method to evaluate the acquisition function at x.
+    def evaluate(self, x : np.ndarray) -> np.ndarray:        
+        y_data = self.gpsurrogate.training_y
+        y_min = y_data[np.argmin(y_data[:, 0])]
+
+        pred = self.gpsurrogate.mean(x)
+        sig = self.gpsurrogate.variance(x)
+
+        retval = []
+        if sig.size == 1 and np.abs(sig) > 1e-12:
+            arg0 = (y_min - pred) / sig
+            retval = (y_min - pred) * norm.cdf(arg0) + sig * norm.pdf(arg0)
+            retval *= -1.
+        elif sig.size == 1 and np.abs(sig) <= 1e-12:
+            retval = 0.0
+        elif sig.size > 1:
+            NotImplementedError("TODO --- Not implemented yet!")
+
+        return retval

src/hiopbbpy/opt/boalgorithm.py

@@ -9,7 +9,7 @@
 from numpy.random import uniform
 from scipy.optimize import minimize
 from ..surrogate_modeling.gp import GaussianProcess
-from .acquisition import LCBacquisition
+from .acquisition import LCBacquisition, EIacquisition
 from ..problems.problem import Problem
 
 # A base class defining a general framework for Bayesian Optimization
@@ -59,7 +59,7 @@ class BOAlgorithm(BOAlgorithmBase):
     def __init__(self, gpsurrogate, xtrain, ytrain, acquisition_type = "LCB"):
         super().__init__()
         assert isinstance(gpsurrogate, GaussianProcess)
-        assert acquisition_type in ["LCB"]
+        assert acquisition_type in ["LCB", "EI"]
         self.setTrainingData(xtrain, ytrain)
         self.setAcquisitionType(acquisition_type)
         self.gpsurrogate = gpsurrogate
@@ -84,6 +84,8 @@ def _find_best_point(self, x_train, y_train, x0 = None):
 
         if self.acquisition_type == "LCB":
             acqf = LCBacquisition(self.gpsurrogate)
+        elif self.acquisition_type == "EI":
+            acqf = EIacquisition(self.gpsurrogate)
         else:
             raise NotImplementedError("No implemented acquisition_type associated to"+self.acquisition_type)
 

src/hiopbbpy/problems/BraninProblem.py

@@ -0,0 +1,56 @@
+"""
+Implementation of the Branin problem
+
+min f(x_1, x_2) = \left( x_2 - \frac{5.1}{4\pi^2} x_1^2 + \frac{5}{\pi} x_1 - 6 \right)^2 + 10 \left(1 - \frac{1}{8\pi}\right) \cos(x_1) + 10.
+s.t x_1 \in [-5, 10], \quad x_2 \in [0, 15].
+
+It has three global minima at:
+(x_1, x_2) = (-\pi, 12.275) 
+             (\pi, 2.275)
+             (9.42478, 2.475)
+and the optimal objective 0.3979
+
+Authors:    Tucker Hartland <hartland1@llnl.gov>
+            Nai-Yuan Chiang <chiang7@llnl.gov>
+"""
+import numpy as np
+from .problem import Problem
+from numpy.random import uniform
+
+# define the Branin problem class
+class BraninProblem(Problem):
+    def __init__(self):
+        ndim = 2
+        xlimits = np.array([[-5.0, 10], [0.0, 15]]) 
+        name = 'Branin'
+        super().__init__(ndim, xlimits, name=name)
+            
+    def _evaluate(self, x: np.ndarray) -> np.ndarray:
+        
+        ne, nx = x.shape
+        assert nx == self.ndim
+        
+        y = np.zeros((ne, 1), complex)
+        b = 5.1 / (4.0 * (np.pi) ** 2)
+        c = 5.0 / np.pi
+        r = 6.0
+        s = 10.0
+        t = 1.0 / (8.0 * np.pi)
+        
+        arg1 = (x[:,1] - b * x[:,0]**2 + c * x[:,0] - r)
+        y[:,0] = arg1**2 + s * (1 - t) * np.cos(x[:,0]) + s
+        
+        '''
+        # compute derivatives
+        dy_dx0 = 2*arg1*(-2*b*x[:,0]+c) - s*(1-t)*np.sin(x[:,0])
+        dy_dx1 = 2*arg1
+        
+        dy_dx = np.array([dy_dx0, dy_dx1])
+        '''
+
+        return y
+
+
+
+
+

src/hiopbbpy/problems/__init__.py

@@ -1,6 +1,8 @@
 from .problem import Problem
+from .BraninProblem import (BraninProblem)
 
 __all__ = [
         "Problem"
+        "BraninProblem"
         ]
 

src/hiopbbpy/problems/problem.py

@@ -6,6 +6,7 @@
 """
 import numpy as np
 from numpy.random import uniform
+from scipy.stats import qmc
 
 # define the general optimization problem class
 class Problem:
@@ -17,6 +18,7 @@ def __init__(self, ndim, xlimits, name=None):
             self.name = " "
         else:
             self.name = name
+        self.sampler = qmc.LatinHypercube(ndim)
 
     def _evaluate(self, x: np.ndarray) -> np.ndarray:
         """
@@ -66,9 +68,16 @@ def sample(self, nsample: int) -> np.ndarray:
         ndarray[nsample, nx]
            Samples from domain defined by xlimits
         """
-        xsample = np.zeros((nsample, self.ndim))
-        for j in range(self.ndim):
-            xsample[:, j] = uniform(self.xlimits[j][0], self.xlimits[j][1], size=nsample)
+
+        # uniform
+        # xsample = np.zeros((nsample, self.ndim))
+        # for j in range(self.ndim):
+        #    xsample[:, j] = uniform(self.xlimits[j][0], self.xlimits[j][1], size=nsample)
+
+        # from predefined sampler
+        xsample = self.sampler.random(nsample)
+        xsample = self.xlimits[:,0] + (self.xlimits[:,1] - self.xlimits[:,0]) * xsample
+
         return xsample
 
 

src/hiopbbpy/surrogate_modeling/gp.py

@@ -12,6 +12,9 @@ class GaussianProcess:
     def __init__(self, ndim, xlimits=None):
         self.ndim = ndim
         self.xlimits = xlimits
+        self.training_x = []
+        self.training_y = []
+        self.trained = False
 
     # Abstract method for computing the mean of the GP at a given input x
     def mean(self, x: np.ndarray) -> np.ndarray:
@@ -72,3 +75,15 @@ def get_bounds(self):
         else:
             return [(self.xlimits[i][0], self.xlimits[i][1]) for i in range(self.ndim)]
 
+    # Abstract method for training the GP
+    def train(self, x: np.ndarray, y: np.ndarray) -> np.ndarray:
+        """
+        train the GP model
+
+        Parameters
+        ---------
+        x : ndarray[n, nx]
+        y : ndarray[n, 1]
+
+        """
+        NotImplementedError("Child class of GaussianProcess should implement method train")