Add Gradient Enhanced EGO optimization (SMTorg#340)

* Add Gradient Enhanced EGO optimization

* Fix test_ego_gek result assertion

* Clarify superclass object of ego_gek test problem class

* Simplify y_data update in ego.optimize()

* Change location of GEKPLS name declaration for correct printing purposes

* Change ego_gek test optimization function to exp

* Set n_comp=2 for GEKPLS objects to pass option assertions

* Update author list

Author: Laurentww

AUTHORS.md

@@ -26,3 +26,4 @@ SMT has been developed thanks to contributions from:
 * Ruben Conde
 * Steven Berguin
 * Vincent Drouet
+* Laurent Wilkens

smt/applications/ego.py

@@ -22,7 +22,7 @@
 )
 from smt.utils.misc import compute_rms_error
 
-from smt.surrogate_models import KPLS, KRG, KPLSK, MGP
+from smt.surrogate_models import KPLS, KRG, KPLSK, MGP, GEKPLS
 from smt.sampling_methods import LHS
 
 
@@ -119,7 +119,7 @@ def _initialize(self):
         declare(
             "surrogate",
             KRG(print_global=False),
-            types=(KRG, KPLS, KPLSK, MGP),
+            types=(KRG, KPLS, KPLSK, GEKPLS, MGP),
             desc="SMT kriging-based surrogate model used internaly",
         )
         declare(
@@ -180,7 +180,8 @@ def optimize(self, fun):
                 y_et_k = self._get_virtual_point(np.atleast_2d(x_et_k), y_data)
 
                 # Update y_data with predicted value
-                y_data = np.atleast_2d(np.append(y_data, y_et_k)).T
+                y_data = y_data.reshape(y_data.shape[0], self.gpr.ny)
+                y_data = np.vstack((y_data, y_et_k))
                 x_data = np.atleast_2d(np.append(x_data, x_et_k, axis=0))
 
             # Compute the real values of y_data
@@ -191,7 +192,7 @@ def optimize(self, fun):
             y_data[-n_parallel:] = y
 
         # Find the optimal point
-        ind_best = np.argmin(y_data)
+        ind_best = np.argmin(y_data if y_data.ndim == 1 else y_data[:, 0])
         x_opt = x_data[ind_best]
         y_opt = y_data[ind_best]
 
@@ -335,6 +336,13 @@ def _find_best_point(self, x_data=None, y_data=None, enable_tunneling=False):
 
         """
         self.gpr.set_training_values(x_data, y_data)
+        if self.gpr.supports["training_derivatives"]:
+            for kx in range(self.gpr.nx):
+                self.gpr.set_training_derivatives(
+                    x_data,
+                    y_data[:, 1 + kx].reshape((y_data.shape[0], 1)),
+                    kx
+                )
         self.gpr.train()
 
         criterion = self.options["criterion"]
@@ -364,7 +372,7 @@ def _find_best_point(self, x_data=None, y_data=None, enable_tunneling=False):
                 try:
                     opt_all.append(
                         minimize(
-                            lambda x: float(self.obj_k(x)),
+                            lambda x: float(np.array(self.obj_k(x)).flat[0]),
                             x_start[ii, :],
                             method="SLSQP",
                             bounds=bounds,

smt/applications/tests/test_ego.py

@@ -23,7 +23,7 @@
 from smt.problems import Branin, Rosenbrock
 from smt.sampling_methods import FullFactorial
 from multiprocessing import Pool
-from smt.surrogate_models import KRG, QP
+from smt.surrogate_models import KRG, QP, GEKPLS
 from smt.applications.mixed_integer import (
     MixedIntegerContext,
     MixedIntegerSamplingMethod,
@@ -464,6 +464,53 @@ def test_find_best_point(self):
         x, _ = ego._find_best_point(xdoe, ydoe, enable_tunneling=False)
         self.assertAlmostEqual(6.5, float(x), delta=1)
 
+    def test_ego_gek(self):
+        from smt.problems import TensorProduct
+
+        class TensorProductIndirect(TensorProduct):
+            def __init__(self, **kwargs):
+                super().__init__(**kwargs)
+                self.super = super()
+
+            def _evaluate(self, x, kx):
+                assert kx is None
+                response = self.super._evaluate(x, kx)
+                sens = np.hstack(self.super._evaluate(x, ki) for ki in range(x.shape[1]))
+                return np.hstack((response, sens))
+
+        fun = TensorProductIndirect(ndim=2, func="exp")
+
+        # Construction of the DOE
+        sampling = LHS(xlimits=fun.xlimits, criterion="m")
+        xdoe = sampling(20)
+        ydoe = fun(xdoe)
+
+        # Build the GEKPLS surrogate model
+        n_comp = 2
+        sm = GEKPLS(
+            theta0=[1e-2] * n_comp,
+            xlimits=fun.xlimits,
+            extra_points=1,
+            print_prediction=False,
+            n_comp=n_comp,
+        )
+
+        # Build the EGO optimizer and optimize
+        ego = EGO(
+            xdoe=xdoe,
+            ydoe=ydoe,
+            n_iter=5,
+            criterion="LCB",
+            xlimits=fun.xlimits,
+            surrogate=sm,
+            n_start=30,
+            enable_tunneling=False,
+        )
+        x_opt, _, _, _, _ = ego.optimize(fun=fun)
+
+        self.assertAlmostEqual(-1.0, float(x_opt[0]), delta=1e-4)
+        self.assertAlmostEqual(-1.0, float(x_opt[1]), delta=1e-4)
+
     def test_qei_criterion_default(self):
         fun = TestEGO.function_test_1d
         xlimits = np.array([[0.0, 25.0]])

smt/surrogate_models/gekpls.py

@@ -10,7 +10,6 @@
 
 
 class GEKPLS(KPLS):
-    name = "GEKPLS"
 
     def _initialize(self):
         super(GEKPLS, self)._initialize()
@@ -37,6 +36,8 @@ def _initialize(self):
         )
         self.supports["training_derivatives"] = True
 
+        self.name = "GEKPLS"
+
     def _compute_pls(self, X, y):
         if 0 in self.training_points[None]:
             self.coeff_pls, XX, yy = ge_compute_pls(

smt/surrogate_models/tests/test_surrogate_model_examples.py

@@ -411,8 +411,13 @@ def test_gekpls(self):
             yt = np.concatenate((yt, yd), axis=1)
 
         # Build the GEKPLS model
+        n_comp = 2
         sm = GEKPLS(
-            theta0=[1e-2], xlimits=fun.xlimits, extra_points=1, print_prediction=False
+            theta0=[1e-2] * n_comp,
+            xlimits=fun.xlimits,
+            extra_points=1,
+            print_prediction=False,
+            n_comp=n_comp,
         )
         sm.set_training_values(xt, yt[:, 0])
         for i in range(2):

smt/tests/test_all.py

@@ -73,7 +73,7 @@ def setUp(self):
         sms["KPLS"] = KPLS(theta0=[1e-2] * ncomp, n_comp=ncomp)
         sms["KPLSK"] = KPLSK(theta0=[1] * ncomp, n_comp=ncomp)
         sms["MGP"] = KPLSK(theta0=[1e-2] * ncomp, n_comp=ncomp)
-        sms["GEKPLS"] = GEKPLS(theta0=[1e-2] * ncomp, n_comp=ncomp, delta_x=1e-1)
+        sms["GEKPLS"] = GEKPLS(theta0=[1e-2] * 2, n_comp=2, delta_x=1e-1)
         sms["GENN"] = genn()
         if compiled_available:
             sms["IDW"] = IDW()

smt/utils/kriging_utils.py

@@ -1048,7 +1048,7 @@ def ge_compute_pls(X, y, n_comp, pts, delta_x, xlimits, extra_points):
     X: np.ndarray [n_obs,dim]
             - - The input variables.
 
-    y: np.ndarray [n_obs,1]
+    y: np.ndarray [n_obs,ny]
             - The output variable
 
     n_comp: int
@@ -1081,7 +1081,7 @@ def ge_compute_pls(X, y, n_comp, pts, delta_x, xlimits, extra_points):
     """
     nt, dim = X.shape
     XX = np.empty(shape=(0, dim))
-    yy = np.empty(shape=(0, 1))
+    yy = np.empty(shape=(0, y.shape[1]))
     _pls = pls(n_comp)
 
     coeff_pls = np.zeros((nt, dim, n_comp))
@@ -1168,9 +1168,9 @@ def ge_compute_pls(X, y, n_comp, pts, delta_x, xlimits, extra_points):
             for ii in max_coeff:
                 XX = np.vstack((XX, X[i, :]))
                 XX[-1, ii] += delta_x * (xlimits[ii, 1] - xlimits[ii, 0])
-                yy = np.vstack((yy, y[i, 0]))
-                yy[-1, 0] += (
-                    pts[None][1 + ii][1][i, 0]
+                yy = np.vstack((yy, y[i]))
+                yy[-1] += (
+                    pts[None][1 + ii][1][i]
                     * delta_x
                     * (xlimits[ii, 1] - xlimits[ii, 0])
                 )