Merge pull request #137 from optimas-org/feature/adjust_n_init

Reduce number of initialization evaluations if external data is supplied

optimas/generators/ax/service/ax_client.py

@@ -75,6 +75,7 @@ def __init__(
             varying_parameters=varying_parameters,
             objectives=objectives,
             analyzed_parameters=analyzed_parameters,
+            enforce_n_init=True,
             use_cuda=use_cuda,
             gpu_id=gpu_id,
             dedicated_resources=dedicated_resources,

optimas/generators/ax/service/base.py

@@ -5,6 +5,7 @@
 import os
 
 from ax.service.ax_client import AxClient
+from ax.modelbridge.registry import Models
 
 from optimas.utils.other import update_object
 from optimas.core import Objective, Trial, VaryingParameter, Parameter
@@ -26,7 +27,12 @@ class AxServiceGenerator(AxGenerator):
         optimization objectives. By default ``None``.
     n_init : int, optional
         Number of evaluations to perform during the initialization phase using
-        Sobol sampling. By default, ``4``.
+        Sobol sampling. If external data is attached to the exploration, the
+        number of initialization evaluations will be reduced by the same
+        amount, unless `enforce_n_init=True`. By default, ``4``.
+    enforce_n_init : bool, optional
+        Whether to enforce the generation of `n_init` Sobol trials, even if
+        external data is supplied. By default, ``False``.
     use_cuda : bool, optional
         Whether to allow the generator to run on a CUDA GPU. By default
         ``False``.
@@ -54,6 +60,7 @@ def __init__(
         objectives: List[Objective],
         analyzed_parameters: Optional[List[Parameter]] = None,
         n_init: Optional[int] = 4,
+        enforce_n_init: Optional[bool] = False,
         use_cuda: Optional[bool] = False,
         gpu_id: Optional[int] = 0,
         dedicated_resources: Optional[bool] = False,
@@ -73,6 +80,7 @@ def __init__(
             model_history_dir=model_history_dir,
         )
         self._n_init = n_init
+        self._enforce_n_init = enforce_n_init
         self._ax_client = self._create_ax_client()
 
     def _ask(self, trials: List[Trial]) -> List[Trial]:
@@ -104,6 +112,15 @@ def _tell(self, trials: List[Trial]) -> None:
                 _, trial_id = self._ax_client.attach_trial(params)
                 self._ax_client.complete_trial(trial_id, objective_eval)
 
+                # Since data was given externally, reduce number of
+                # initialization trials.
+                if not self._enforce_n_init:
+                    gs = self._ax_client.generation_strategy
+                    ngen, _ = gs._num_trials_to_gen_and_complete_in_curr_step()
+                    # Reduce only if there are still Sobol trials to generate.
+                    if gs.current_step.model == Models.SOBOL and ngen > 0:
+                        gs.current_step.num_trials -= 1
+
     def _create_ax_client(self) -> AxClient:
         """Create Ax client (must be implemented by subclasses)."""
         raise NotImplementedError

optimas/generators/ax/service/multi_fidelity.py

@@ -29,7 +29,12 @@ class AxMultiFidelityGenerator(AxServiceGenerator):
         optimization objectives. By default ``None``.
     n_init : int, optional
         Number of evaluations to perform during the initialization phase using
-        Sobol sampling. By default, ``4``.
+        Sobol sampling. If external data is attached to the exploration, the
+        number of initialization evaluations will be reduced by the same
+        amount, unless `enforce_n_init=True`. By default, ``4``.
+    enforce_n_init : bool, optional
+        Whether to enforce the generation of `n_init` Sobol trials, even if
+        external data is supplied. By default, ``False``.
     fidel_cost_intercept : float, optional
         The cost intercept for the affine cost of the form
         `cost_intercept + n`, where `n` is the number of generated points.
@@ -61,6 +66,7 @@ def __init__(
         objectives: List[Objective],
         analyzed_parameters: Optional[List[Parameter]] = None,
         n_init: Optional[int] = 4,
+        enforce_n_init: Optional[bool] = False,
         fidel_cost_intercept: Optional[float] = 1.0,
         use_cuda: Optional[bool] = False,
         gpu_id: Optional[int] = 0,
@@ -75,6 +81,7 @@ def __init__(
             objectives=objectives,
             analyzed_parameters=analyzed_parameters,
             n_init=n_init,
+            enforce_n_init=enforce_n_init,
             use_cuda=use_cuda,
             gpu_id=gpu_id,
             dedicated_resources=dedicated_resources,

optimas/generators/ax/service/single_fidelity.py

@@ -38,7 +38,12 @@ class AxSingleFidelityGenerator(AxServiceGenerator):
         optimization objectives. By default ``None``.
     n_init : int, optional
         Number of evaluations to perform during the initialization phase using
-        Sobol sampling. By default, ``4``.
+        Sobol sampling. If external data is attached to the exploration, the
+        number of initialization evaluations will be reduced by the same
+        amount, unless `enforce_n_init=True`. By default, ``4``.
+    enforce_n_init : bool, optional
+        Whether to enforce the generation of `n_init` Sobol trials, even if
+        external data is supplied. By default, ``False``.
     fully_bayesian : bool, optional
         Whether to optimize the hyperparameters of the GP with a fully
         Bayesian approach (using SAAS priors) instead of maximizing
@@ -81,6 +86,7 @@ def __init__(
         objectives: List[Objective],
         analyzed_parameters: Optional[List[Parameter]] = None,
         n_init: Optional[int] = 4,
+        enforce_n_init: Optional[bool] = False,
         fully_bayesian: Optional[bool] = False,
         use_cuda: Optional[bool] = False,
         gpu_id: Optional[int] = 0,
@@ -95,6 +101,7 @@ def __init__(
             objectives=objectives,
             analyzed_parameters=analyzed_parameters,
             n_init=n_init,
+            enforce_n_init=enforce_n_init,
             use_cuda=use_cuda,
             gpu_id=gpu_id,
             dedicated_resources=dedicated_resources,

tests/test_ax_generators.py

@@ -427,6 +427,94 @@ def test_ax_multitask_with_history():
     exploration.run()
 
 
+def test_ax_service_init():
+    """
+    Test that an exploration with using an AxServiceGenerator correctly
+    reduces the number of `n_init` SOBOL evaluations if external trials
+    or evaluations are given.
+    """
+
+    var1 = VaryingParameter("x0", -50.0, 5.0)
+    var2 = VaryingParameter("x1", -5.0, 15.0)
+    obj = Objective("f", minimize=False)
+
+    n_init = 4
+    n_external = 6
+
+    for i in range(n_external):
+        gen = AxSingleFidelityGenerator(
+            varying_parameters=[var1, var2], objectives=[obj], n_init=n_init
+        )
+        ev = FunctionEvaluator(function=eval_func_sf)
+        exploration = Exploration(
+            generator=gen,
+            evaluator=ev,
+            max_evals=6,
+            sim_workers=2,
+            exploration_dir_path=f"./tests_output/test_ax_service_init_{i}",
+        )
+
+        # Get reference to AxClient.
+        ax_client = gen._ax_client
+
+        for _ in range(i):
+            exploration.evaluate_trials(
+                {
+                    "x0": [-2.0 + np.random.rand(1)[0]],
+                    "x1": [2.7 + np.random.rand(1)[0]],
+                }
+            )
+        # Run exploration.
+        exploration.run()
+
+        # Check that the number of SOBOL trials is reduced and that they
+        # are replaced by Manual trials.
+        df = ax_client.get_trials_data_frame()
+        for j in range(i):
+            assert df["generation_method"][j] == "Manual"
+        for k in range(i, n_init - 1):
+            assert df["generation_method"][k] == "Sobol"
+        df["generation_method"][min(i, n_init)] == "GPEI"
+
+    # Test single case with `enforce_n_init=True`
+    gen = AxSingleFidelityGenerator(
+        varying_parameters=[var1, var2],
+        objectives=[obj],
+        n_init=n_init,
+        enforce_n_init=True,
+    )
+    ev = FunctionEvaluator(function=eval_func_sf)
+    exploration = Exploration(
+        generator=gen,
+        evaluator=ev,
+        max_evals=15,
+        sim_workers=2,
+        exploration_dir_path="./tests_output/test_ax_service_init_enforce",
+    )
+
+    # Get reference to AxClient.
+    ax_client = gen._ax_client
+
+    for _ in range(n_external):
+        exploration.evaluate_trials(
+            {
+                "x0": [-2.0 + np.random.rand(1)[0]],
+                "x1": [2.7 + np.random.rand(1)[0]],
+            }
+        )
+    # Run exploration.
+    exploration.run()
+
+    # Check that the number of SOBOL trials is still `n_init` after adding
+    # `n_external` Manual trials.
+    df = ax_client.get_trials_data_frame()
+    for j in range(n_external):
+        assert df["generation_method"][j] == "Manual"
+    for k in range(n_external, n_external + n_init):
+        assert df["generation_method"][k] == "Sobol"
+    df["generation_method"][n_external + n_init] == "GPEI"
+
+
 if __name__ == "__main__":
     test_ax_single_fidelity()
     test_ax_single_fidelity_moo()
@@ -438,3 +526,4 @@ def test_ax_multitask_with_history():
     test_ax_single_fidelity_with_history()
     test_ax_multi_fidelity_with_history()
     test_ax_multitask_with_history()
+    test_ax_service_init()