ENH: SMBO: Add weighted random sampling to further improve benchmarks

Author: kernc

sambo/_smbo.py

@@ -11,7 +11,7 @@
     _initialize_population,
     _sample_population,
     _check_bounds,
-    _check_random_state, _sanitize_constraints, lru_cache,
+    _check_random_state, _sanitize_constraints, lru_cache, recompute_kde, weighted_uniform_sampling,
 )
 
 
@@ -117,7 +117,7 @@ def __init__(
             max_iter: int = INT32_MAX,
             n_init: Optional[int] = None,
             n_candidates: Optional[int] = None,
-            n_iter_no_change: int = 10,
+            n_iter_no_change: int = 5,
             n_models: int = 1,
             tol: float = FLOAT32_PRECISION,
             estimator: Literal['gp', 'et', 'gb'] | _SklearnLikeRegressor = None,
@@ -147,7 +147,9 @@ def __init__(
         rng = _check_random_state(rng)
 
         if n_init is None:
-            n_init = 0 if not callable(fun) else min(max(1, max_iter - 20), 150 * len(bounds))
+            n_init = (0 if not callable(fun) else
+                      min(max(1, max_iter - 20),
+                          int(40 * len(bounds) * max(1, np.log2(len(bounds))))))
         assert max_iter >= n_init, (max_iter, n_init)
 
         if n_candidates is None:
@@ -201,6 +203,9 @@ def __init__(
         self._y = y
         assert len(X) == len(y), (X, y)
 
+        self._kde = None
+        self._prev_y_min = np.inf
+
         # Cache methods on the _instance_
         self._init_once = lru_cache(1)(self._init_once)
         self.top_k = lru_cache(1)(self.top_k)
@@ -334,8 +339,18 @@ def ask(
         assert isinstance(kappa, (Real, Iterable)), kappa
         self._init_once()
 
-        n_points = max(10_000, 1000 * int(len(self.bounds)**1.2))
-        X = _sample_population(self.bounds, n_points, self.constraints, self.rng)
+        n_points = min(80_000, 20_000 * int(len(self.bounds)**2))  # TODO: Make this a param?
+        nfev = len(self._X)
+        if nfev < 10 * len(self.bounds)**2:
+            X = _sample_population(self.bounds, n_points, self.constraints, self.rng)
+        else:
+            y_min = np.min(self._y)
+            if self._kde is None or (nfev < 200 or nfev % 5 == 0 or y_min < self._prev_y_min):
+                self._prev_y_min = y_min
+                self._kde = recompute_kde(np.array(self._X), np.array(self._y))
+            X = weighted_uniform_sampling(
+                self._kde, self.bounds, n_points, self.constraints, self.rng)
+
         X, mean, std = self._predict(X)
         criterion = acq_func(mean=mean, std=std, kappa=kappa)
         n_candidates = min(n_candidates, criterion.shape[1])
@@ -552,7 +567,7 @@ def smbo(
         max_iter: int = INT32_MAX,
         n_init: Optional[int] = None,
         n_candidates: Optional[int] = None,
-        n_iter_no_change: int = 10,
+        n_iter_no_change: int = 5,
         n_models: int = 1,
         tol: float = FLOAT32_PRECISION,
         estimator: Optional[str | _SklearnLikeRegressor] = None,

sambo/_test.py

@@ -13,15 +13,15 @@
     Bounds, NonlinearConstraint, rosen, minimize as scipy_minimize,
     shgo as scipy_shgo,
 )
+from scipy.stats import gaussian_kde
 
 from sambo import minimize, Optimizer, SamboSearchCV
 from sambo._space import Space
 from sambo._sceua import sceua
 from sambo._shgo import shgo
 from sambo._smbo import smbo
 from sambo.plot import plot_convergence, plot_evaluations, plot_objective, plot_regret
-from sambo._util import OptimizeResult
-
+from sambo._util import OptimizeResult, recompute_kde, weighted_uniform_sampling
 
 minimize = partial(minimize, rng=0)
 sceua = partial(sceua, rng=0)
@@ -214,7 +214,7 @@ def test_sceua(self):
 
     def test_smbo(self):
         res = minimize(**ROSEN_TEST_PARAMS, method='smbo', max_iter=20, estimator='gp')
-        check_result(res, 0, atol=55)
+        check_result(res, 0, atol=5)
 
     def test_args(self):
         def f(x, a):
@@ -240,7 +240,7 @@ def f(x):
 
         counter = 0
         _ = _minimize(f, method='smbo')
-        self.assertEqual(counter, MAX_ITER)
+        self.assertLessEqual(counter, MAX_ITER)
 
     def test_constraints(self):
         def f(x):
@@ -313,16 +313,11 @@ def test_our_params_match_scipy_optimize_params(self):
 
 class TestSklearnEstimators(unittest.TestCase):
     def test_estimator_factory(self):
-        DEFAULT_KWARGS = {'max_iter': 50, 'n_iter_no_change': 10, 'rng': 0}
-        ESTIMATOR_KWARGS = {
-            'gp': {},
-            'et': {'n_iter_no_change': 40},
-            'gb': {'n_iter_no_change': 20, 'rng': 2},
-        }
+        DEFAULT_KWARGS = {'max_iter': 20, 'n_iter_no_change': 5, 'rng': 0}
         for estimator in BUILTIN_ESTIMATORS:
             with self.subTest(estimator=estimator):
                 res = smbo(lambda x: sum((x-2)**2), bounds=[(-100, 100)], estimator=estimator,
-                           **dict(DEFAULT_KWARGS, **ESTIMATOR_KWARGS[estimator]))
+                           **dict(DEFAULT_KWARGS))
                 self.assertLess(res.fun, 1, msg=res)
 
     def test_SamboSearchCV_large_param_grid(self):
@@ -353,11 +348,17 @@ def test_SamboSearchCV_large_param_grid(self):
 
 class TestDocs(unittest.TestCase):
     def test_make_doc_plots(self):
+        KWARGS = {
+            'shgo': dict(n_init=30),
+            'smbo': dict(n_init=30),
+            'sceua': dict(n_complexes=3),
+        }
         results = [
             minimize(
-                rosen, bounds=[(-2., 2.), (-2., 2.)],
-                constraints=lambda x: sum(x**2) <= 2*len(x),
-                max_iter=120, method=method, rng=2,
+                rosen, bounds=[(-2., 2.)]*2,
+                constraints=lambda x: sum(x**2) <= 2**len(x),
+                max_iter=100, method=method, rng=2,
+                **KWARGS.get(method, {}),
             )
             for method in BUILTIN_METHODS
         ]
@@ -392,7 +393,7 @@ def test_make_doc_plots(self):
     def test_website_example1(self):
         res = minimize(
             rosen, bounds=[(-2., 2.), ] * 2,
-            constraints=lambda x: sum(x**2) <= len(x),
+            constraints=lambda x: sum(x**2) <= 2**len(x),
             n_init=7, method='shgo', rng=0,
         )
         print(type(res), res, sep='\n\n')
@@ -478,5 +479,35 @@ def test_annotations(self):
                         annot[arg], annot_ref[arg], msg=f'{fun.__qualname__} / {arg}')
 
 
+class TestUtil(unittest.TestCase):
+    def test_weighted_uniform_sampling(self):
+        rng = np.random.default_rng(2)
+        X = rng.uniform(-10, 10, (100, 2))
+        y = rng.uniform(1, 10, 100)
+        bounds = [(-10, 10), (-10, 10)]
+        n_samples = 10000
+        kde = recompute_kde(X, y)
+        sampled_points = weighted_uniform_sampling(kde, bounds, n_samples, None, 0)
+
+        # Verify results
+        hist, xedges, yedges = np.histogram2d(*sampled_points.T, range=bounds)
+        # Compare histogram density with weight distribution
+        kde = gaussian_kde(X.T, weights=(np.max(y) - y) / np.sum(np.max(y) - y))
+        test_grid = np.array(np.meshgrid(xedges[:-1], yedges[:-1])).T.reshape(-1, 2)
+        pdf_values = kde(test_grid.T).reshape(hist.shape)
+        # Normalize for direct comparison
+        hist_normalized = hist / np.sum(hist)
+        pdf_normalized = pdf_values / np.sum(pdf_values)
+        # Plot results
+        fig, axes = plt.subplots(1, 2, figsize=(8, 4))
+        extent = np.array(bounds).flatten()
+        axes[0].imshow(hist_normalized, extent=extent, origin='lower', cmap='Blues')
+        axes[1].imshow(pdf_normalized, extent=extent, origin='lower', cmap='Reds')
+        plt.show()
+
+        diff = np.abs(hist_normalized - pdf_normalized).mean()
+        self.assertLess(diff, .05)
+
+
 if __name__ == '__main__':
     unittest.main()

sambo/_util.py

@@ -1,11 +1,13 @@
 import heapq
 from functools import lru_cache as _lru_cache, wraps
+from itertools import islice
 from numbers import Integral, Real
 from threading import Lock
 from typing import Any, Callable, Optional, Protocol, runtime_checkable
 
 import numpy as np
 from scipy.optimize import Bounds, NonlinearConstraint, OptimizeResult as _OptimizeResult
+from scipy.stats import gaussian_kde
 from scipy.stats.qmc import LatinHypercube
 
 FLOAT32_PRECISION = 10**-np.finfo(np.float32).precision
@@ -276,3 +278,32 @@ def constraints(x, *, _c=constraints):
                 def constraints(x, *, _c=constraints['fun']):
                     return _c(x) == 0
     return constraints
+
+
+def weighted_uniform_sampling(kde, bounds, size, constraints, rng):
+    """Sample points from a weighted density within given bounds"""
+    rng = _check_random_state(rng)
+    lb, ub = np.array(bounds).T
+    if constraints is None:
+        def constraints(_):
+            return True
+    try:
+        points = np.array(
+            list(islice(
+                (x for _ in range(1000)
+                 for x in kde.resample(size, seed=rng).T
+                 if constraints(x) and np.all((lb <= x) & (x <= ub))),
+                size)))
+    except ValueError as ex:
+        if 'too short' in str(ex):
+            raise RuntimeError('Constraints seemingly cannot be satisfied.')
+        raise
+    return points
+
+
+def recompute_kde(X, y):
+    w = np.max(y) - y
+    w /= np.sum(w)
+    w **= 3
+    kde = gaussian_kde(X.T, bw_method='silverman', weights=w)
+    return kde