Refactor acquisition functions to use functools.partial instead of lambda's. Goal: enable pickling

Moreover:

- Create greedy approach for FarEnoughSampleFilter. More reliable as a fallback
- PartialVariableFunction as a common interface for MinimizeSurrogate and cptv
- Record all points used in the local solver to improve the surrogate in cptv-L

Author: weslleyspereira

examples/vlse_benchmark/vlse_bench.py

@@ -127,15 +127,15 @@ def run_optimizer(
                 maxeval=maxEval - 2 * (nArgs + 1),
                 surrogateModel=modelIter,
                 acquisitionFunc=acquisitionFuncIter,
-                seed=2*i+1,
+                seed=2 * i + 1,
             )
         else:
             res = optimizer(
                 objf,
                 bounds=bounds,
                 maxeval=maxEval - 2 * (nArgs + 1),
                 surrogateModel=modelIter,
-                seed=2*i+1,
+                seed=2 * i + 1,
             )
         optres.append(res)
 

pyproject.toml

@@ -45,7 +45,6 @@ test = [
     "jupyter",
     "ipykernel",
     "pygoblet",
-    "PyNomadBBO",
 ]
 lint = [
     "ruff",

soogo/acquisition/endpoints_pareto_front.py

@@ -22,6 +22,7 @@
 
 import numpy as np
 from typing import Optional
+import functools
 
 from pymoo.optimize import minimize as pymoo_minimize
 
@@ -88,7 +89,7 @@ def optimize(
         endpoints = np.empty((0, dim))
         for i in range(objdim):
             minimumPointProblem = PymooProblem(
-                lambda x: surrogateModel(x, i=i), bounds, iindex
+                functools.partial(surrogateModel, i=i), bounds, iindex
             )
             res = pymoo_minimize(
                 minimumPointProblem,

soogo/acquisition/maximize_distance.py

@@ -19,6 +19,7 @@
 
 import numpy as np
 from typing import Optional
+import functools
 
 from pymoo.optimize import minimize as pymoo_minimize
 
@@ -28,6 +29,16 @@
 from .utils import FarEnoughSampleFilter
 
 
+def _negative_distance(tree, x):
+    """Compute negative distance to nearest neighbor.
+
+    :param tree: KDTree of previously sampled points.
+    :param x: Point to evaluate.
+    :return: Negative distance to nearest neighbor.
+    """
+    return -tree.query(x)[0]
+
+
 class MaximizeDistance(Acquisition):
     """
     Maximizing distance acquisition function as described in [#]_.
@@ -86,7 +97,7 @@ def optimize(
         filter = FarEnoughSampleFilter(exclusion_set, self.tol(bounds))
 
         problem = PymooProblem(
-            lambda x: -filter.tree.query(x)[0],
+            functools.partial(_negative_distance, filter.tree),
             bounds,
             iindex,
         )

soogo/acquisition/maximize_ei.py

@@ -22,6 +22,7 @@
 from scipy.linalg import cholesky, solve_triangular
 from typing import Optional
 import logging
+import functools
 
 from pymoo.optimize import minimize as pymoo_minimize
 
@@ -34,6 +35,17 @@
 logger = logging.getLogger(__name__)
 
 
+def _negative_ei(surrogate, best_value, x):
+    """Negative expected improvement (for minimization).
+
+    :param surrogate: Surrogate model.
+    :param best_value: Best objective value so far.
+    :param x: Point to evaluate.
+    :return: Negative expected improvement at x.
+    """
+    return -surrogate.expected_improvement(x, best_value)
+
+
 class MaximizeEI(Acquisition):
     """Acquisition by maximizing the expected improvement (EI) of a
     :class:`.GaussianProcess`.
@@ -156,7 +168,7 @@ def optimize(
 
         # Use the point that maximizes the EI
         problem = PymooProblem(
-            lambda x: -surrogateModel.expected_improvement(x, ybest),
+            functools.partial(_negative_ei, surrogateModel, ybest),
             bounds,
         )
         res = pymoo_minimize(

soogo/acquisition/minimize_surrogate.py

@@ -25,40 +25,24 @@
 from scipy.optimize import minimize
 from scipy.stats.qmc import LatinHypercube
 from typing import Optional
+import functools
 
 from .base import Acquisition
 from ..model import Surrogate
 from ..sampling import random_sample
 from .utils import FarEnoughSampleFilter
+from ..optimize.utils import PartialVariableFunction
 
 
-class _SurrogateMinimizerWrapper:
-    """Helper class to wrap surrogate model for minimization.
+def _ith_function_value(fun, x, i):
+    """Returns the i-th component of the function value at x.
 
-    :param surrogateModel: Surrogate model to be minimized.
-    :param x: Base point in the space.
-    :param i: Index(s) that will be optimized.
+    :param fun: Function that computes the function value at x.
+    :param x: Point at which to evaluate the function.
+    :param i: Index of the component to return.
+    :return: i-th component of the function value at x.
     """
-
-    def __init__(self, surrogateModel, x, i):
-        self.surrogateModel = surrogateModel
-        self.x = x
-        self.i = i
-
-    def fun(self, xi):
-        """Evaluate surrogate model at modified point x[i] = xi."""
-        _x = self.x.copy()
-        _x[self.i] = xi
-        return self.surrogateModel(_x)
-
-    def jac(self, xi):
-        """Evaluate surrogate model Jacobian at modified point x[i] = xi.
-
-        Only the components corresponding to indices i are returned.
-        """
-        _x = self.x.copy()
-        _x[self.i] = xi
-        return self.surrogateModel.jac(_x)[self.i]
+    return fun(x)[i]
 
 
 class MinimizeSurrogate(Acquisition):
@@ -144,8 +128,15 @@ def optimize(
         critdist = (
             (gamma(1 + (dim / 2)) * volumeBounds * sigma) ** (1 / dim)
         ) / np.sqrt(np.pi)  # critical distance when 2 points are equal
-        has_jac = hasattr(surrogateModel, "jac") and callable(
-            getattr(surrogateModel, "jac")
+        model_jac = (
+            functools.partial(
+                _ith_function_value, surrogateModel.jac, i=cindex
+            )
+            if (
+                hasattr(surrogateModel, "jac")
+                and callable(getattr(surrogateModel, "jac"))
+            )
+            else None
         )
 
         # Local space to store information
@@ -220,14 +211,19 @@ def optimize(
             for i in range(nSelected):
                 xi = candidates[chosenIds[i], :]
 
-                wrapper = _SurrogateMinimizerWrapper(
+                wrapper_fun = PartialVariableFunction(
                     surrogateModel, xi, cindex
                 )
+                wrapper_jac = (
+                    PartialVariableFunction(model_jac, xi, cindex)
+                    if model_jac is not None
+                    else None
+                )
                 res = minimize(
-                    wrapper.fun,
+                    wrapper_fun,
                     xi[cindex],
                     method="L-BFGS-B",
-                    jac=wrapper.jac if has_jac else None,
+                    jac=wrapper_jac,
                     bounds=cbounds,
                     options={
                         "maxfun": remevals,

soogo/acquisition/pareto_front.py

@@ -21,6 +21,7 @@
 import numpy as np
 from scipy.spatial import KDTree
 from typing import Optional
+import functools
 
 from pymoo.core.initialization import Initialization
 from pymoo.optimize import minimize as pymoo_minimize
@@ -33,6 +34,17 @@
 from .utils import FarEnoughSampleFilter
 
 
+def _distance_to_target(surrogate, target_value, x):
+    """Compute distance from surrogate prediction to target.
+
+    :param surrogate: Surrogate model.
+    :param target_value: Target value to compare against.
+    :param x: Point to evaluate.
+    :return: Distance from surrogate prediction at x to target_value.
+    """
+    return np.absolute(surrogate(x) - target_value)
+
+
 class ParetoFront(Acquisition):
     """Obtain sample points that fill gaps in the Pareto front from [#]_.
 
@@ -231,7 +243,7 @@ def optimize(
             # For discontinuous Pareto fronts in the original problem, such set
             # may not exist, or it may be too far from the target value.
             multiobjTVProblem = PymooProblem(
-                lambda x: np.absolute(surrogateModel(x) - tau),
+                functools.partial(_distance_to_target, surrogateModel, tau),
                 bounds,
                 iindex,
                 n_obj=objdim,

soogo/acquisition/target_value_acquisition.py

@@ -20,6 +20,7 @@
 import numpy as np
 import logging
 from typing import Optional
+import functools
 
 from pymoo.optimize import minimize as pymoo_minimize
 
@@ -268,8 +269,11 @@ def optimize(
                     surrogateModel.prepare_mu_measure()
                     mu_measure_is_prepared = True
                 problem = PymooProblem(
-                    lambda x: TargetValueAcquisition.bumpiness_measure(
-                        surrogateModel, x, f_target, target_range
+                    functools.partial(
+                        TargetValueAcquisition.bumpiness_measure,
+                        surrogateModel,
+                        target=f_target,
+                        target_range=target_range,
                     ),
                     bounds,
                     iindex,
@@ -305,8 +309,11 @@ def optimize(
                         surrogateModel.prepare_mu_measure()
                         mu_measure_is_prepared = True
                     problem = PymooProblem(
-                        lambda x: TargetValueAcquisition.bumpiness_measure(
-                            surrogateModel, x, f_target, target_range
+                        functools.partial(
+                            TargetValueAcquisition.bumpiness_measure,
+                            surrogateModel,
+                            target=f_target,
+                            target_range=target_range,
                         ),
                         bounds,
                         iindex,

soogo/acquisition/utils.py

@@ -230,6 +230,10 @@ class FarEnoughSampleFilter:
 
     :param X: Matrix of existing sample points (n x d).
     :param tol: Minimum distance threshold.
+    :param approach: Method to select points when multiple candidates are close
+        to each other. Options are:
+        "independent_set" (default) which uses a maximum independent set
+        algorithm, and "greedy" which uses a simple greedy approach.
 
     .. attribute:: tree
 
@@ -239,11 +243,17 @@ class FarEnoughSampleFilter:
     .. attribute:: tol
 
         Minimum distance threshold. Points closer than this are filtered out.
+
+    .. attribute:: approach
+
+        Method to select points when multiple candidates are close to each
+        other.
     """
 
-    def __init__(self, X, tol):
+    def __init__(self, X, tol, approach="independent_set"):
         self.tree = KDTree(X)
         self.tol = tol
+        self.approach = approach
 
     def is_far_enough(self, x):
         """Check if a point is far enough from existing samples.
@@ -254,6 +264,28 @@ def is_far_enough(self, x):
         dist, _ = self.tree.query(x.reshape(1, -1))
         return dist[0] >= self.tol
 
+    def max_independent_set(self, X):
+        dist = cdist(X, X)
+        np.fill_diagonal(dist, np.inf)
+        g = nx.Graph()
+        g.add_nodes_from(range(len(X)))
+        g.add_edges_from(
+            [
+                (i, j)
+                for i in range(len(X))
+                for j in range(i + 1, len(X))
+                if dist[i, j] < self.tol
+            ]
+        )
+        return maximum_independent_set(g)
+
+    def greedy_independent_set(self, X):
+        idx = []
+        for i in range(len(X)):
+            if all(np.linalg.norm(X[i] - X[j]) >= self.tol for j in idx):
+                idx.append(i)
+        return idx
+
     def indices(self, Xc):
         """Filter candidates based on minimum distance criterion.
 
@@ -266,19 +298,18 @@ def indices(self, Xc):
         Xc0 = Xc[mask0]
 
         # Find the maximum independent set among the remaining points
-        dist = cdist(Xc0, Xc0)
-        np.fill_diagonal(dist, np.inf)
-        g = nx.Graph()
-        g.add_nodes_from(range(len(Xc0)))
-        g.add_edges_from(
-            [
-                (i, j)
-                for i in range(len(Xc0))
-                for j in range(i + 1, len(Xc0))
-                if dist[i, j] < self.tol
-            ]
-        )
-        idx = maximum_independent_set(g)
+        if self.approach == "independent_set":
+            try:
+                idx = self.max_independent_set(Xc0)
+            except Exception as e:
+                logger.warning(
+                    "Error in maximum independent set computation: %s. "
+                    "Falling back to a greedy approach.",
+                    e,
+                )
+                idx = self.greedy_independent_set(Xc0)
+        else:
+            idx = self.greedy_independent_set(Xc0)
 
         # Recover original indices
         original_indices = np.where(mask0)[0]