Merge pull request #2 from NeptuneProjects/qropt

Qropt

demo/hartmann6_demo.py

@@ -7,18 +7,14 @@
 
 from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
 from ax.modelbridge.registry import Models
-from ax.models.torch.botorch_modular.surrogate import Surrogate
 from ax.service.ax_client import AxClient
 from ax.utils.measurement.synthetic_functions import hartmann6
 from ax.utils.notebook.plotting import render
-from botorch.acquisition import qExpectedImprovement
-from botorch.models.gp_regression import SingleTaskGP
-from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood
 import numpy as np
 
 sys.path.insert(0, str(Path(__file__).parents[1]))
 from oao.objective import NoiselessFormattedObjective
-from oao.optimizer import BayesianOptimization, GridSearch
+from oao.optimizer import BayesianOptimization, GridSearch, QuasiRandom
 from oao.results import get_results
 from oao.space import SearchParameter, SearchSpace
 from oao.strategy import GridStrategy
@@ -39,93 +35,99 @@ def main():
     h6 = Hartmann6Objective()
     objective = NoiselessFormattedObjective(h6, "hartmann6", {"minimize": True})
 
+    # Define the search space.
+    search_space = [
+        {"name": f"x{i + 1}", "type": "range", "bounds": [-2.0, 2.0]} for i in range(6)
+    ]
+    space = SearchSpace([SearchParameter(**d) for d in search_space])
+
+    ## Bayesian Optimization ===========================================
     # Define the generation strategy.
     gs = GenerationStrategy(
         [
             GenerationStep(
                 model=Models.SOBOL,
                 num_trials=64,
-                max_parallelism=64,
+                max_parallelism=16,
                 model_kwargs={"seed": 0},
             ),
             GenerationStep(
                 model=Models.GPEI,
                 num_trials=16,
                 max_parallelism=4,
-            )
-            # GenerationStep(
-            #     model=Models.BOTORCH_MODULAR,
-            #     num_trials=16,
-            #     max_parallelism=4,
-            #     model_kwargs={
-            #         "surrogate": Surrogate(
-            #             botorch_model_class=SingleTaskGP,
-            #             mll_class=ExactMarginalLogLikelihood,
-            #         ),
-            #         "botorch_acqf_class": qExpectedImprovement,
-            #     },
-            #     model_gen_kwargs={
-            #         "model_gen_options": {
-            #             "optimizer_kwargs": {
-            #                 "num_restarts": 40,
-            #                 "raw_samples": 1024,
-            #             }
-            #         }
-            #     },
-            # ),
+            ),
         ]
     )
-
-    # Define the search space.
-    search_space = [
-        {"name": f"x{i + 1}", "type": "range", "bounds": [0.0, 1.0]} for i in range(6)
-    ]
-
-    # search_space = [
-    #     {"name": f"x{i + 1}", "type": "choice", "values": [0.0, 1.0]} for i in range(6)
-    # ]
-
-    space = SearchSpace([SearchParameter(**d) for d in search_space])
-
     # Instantiate and run the optimizers.
     opt_bo = BayesianOptimization(
         objective=objective,
         search_space=space,
         strategy=gs,
     )
     opt_bo.run(name="demo_bo")
-    opt = opt_bo
-
-    # opt_gs = GridSearch(
-    #     objective=objective,
-    #     search_space=space,
-    #     strategy=GridStrategy(num_trials=2, max_parallelism=4),
-    # )
-    # opt_gs.run(name="demo_gs")
-    # opt = opt_gs
-
-    # # Save the results to CSV files.
-    # get_results(
-    #     opt_bo.client,
-    #     times=opt.batch_execution_times,
-    #     minimize=objective.properties.minimize,
-    # ).to_csv("demo/results_bo.csv")
-
-    # get_results(
-    #     opt_gs.client,
-    #     times=opt.batch_execution_times,
-    #     minimize=objective.properties.minimize,
-    # ).to_csv("demo/results_gs.csv")
-
-    # # Save the clients to JSON files.
-    # opt_bo.client.save_to_json_file("demo/experiment_bo.json")
-    # opt_gs.client.save_to_json_file("demo/experiment_gs.json")
-
-    # # Load the results from the JSON file and render the optimization trace.
-    # restored_client_bo = AxClient.load_from_json_file("demo/experiment_bo.json")
-    # restored_client_gs = AxClient.load_from_json_file("demo/experiment_gs.json")
-    # render(restored_client_bo.get_optimization_trace(objective_optimum=hartmann6.fmin))
-    # render(restored_client_gs.get_optimization_trace(objective_optimum=hartmann6.fmin))
+
+    ## Grid Search =====================================================
+    # Define the generation strategy.
+    gs = GridStrategy(num_trials=2, max_parallelism=4)
+    # Instantiate and run the optimizers.
+    opt_gs = GridSearch(
+        objective=objective,
+        search_space=space,
+        strategy=gs,
+    )
+    opt_gs.run(name="demo_gs")
+
+    ## Quasi-Random Search =============================================
+    # Define the generation strategy.
+    gs = GenerationStrategy(
+        [
+            GenerationStep(
+                model=Models.SOBOL,
+                num_trials=64,
+                max_parallelism=16,
+                model_kwargs={"seed": 0},
+            ),
+        ]
+    )
+    # Instantiate and run the optimizers.
+    opt_qr = QuasiRandom(
+        objective=objective,
+        search_space=space,
+        strategy=gs,
+    )
+    opt_qr.run(name="demo_bo")
+
+    ## Save the results to CSV files. ==================================
+    get_results(
+        opt_bo.client,
+        times=opt_bo.batch_execution_times,
+        minimize=objective.properties.minimize,
+    ).to_csv("demo/results_bo.csv")
+
+    get_results(
+        opt_gs.client,
+        times=opt_gs.batch_execution_times,
+        minimize=objective.properties.minimize,
+    ).to_csv("demo/results_gs.csv")
+
+    get_results(
+        opt_qr.client,
+        times=opt_qr.batch_execution_times,
+        minimize=objective.properties.minimize,
+    ).to_csv("demo/results_qr.csv")
+
+    # Save the clients to JSON files.
+    opt_bo.client.save_to_json_file("demo/experiment_bo.json")
+    opt_gs.client.save_to_json_file("demo/experiment_gs.json")
+    opt_qr.client.save_to_json_file("demo/experiment_qr.json")
+
+    # Load the results from the JSON file and render the optimization trace.
+    restored_client_bo = AxClient.load_from_json_file("demo/experiment_bo.json")
+    restored_client_gs = AxClient.load_from_json_file("demo/experiment_gs.json")
+    restored_client_qr = AxClient.load_from_json_file("demo/experiment_qr.json")
+    render(restored_client_bo.get_optimization_trace(objective_optimum=hartmann6.fmin))
+    render(restored_client_gs.get_optimization_trace(objective_optimum=hartmann6.fmin))
+    render(restored_client_qr.get_optimization_trace(objective_optimum=hartmann6.fmin))
 
 
 if __name__ == "__main__":

oao/__init__.py

@@ -12,7 +12,7 @@
 acoustic parameters in the ocean using uninformed search methods and 
 Bayesian optimization.
 """
-__version__ = "1.0.0"
+__version__ = "1.0.1"
 
 from ax.storage.botorch_modular_registry import ACQUISITION_FUNCTION_REGISTRY
 from ax.storage.botorch_modular_registry import REVERSE_ACQUISITION_FUNCTION_REGISTRY

oao/optimizer.py

@@ -11,7 +11,7 @@
 import numpy as np
 
 from oao.objective import Objective
-from oao.space import SearchSpace, get_parameterized_grid
+from oao.space import SearchSpace, get_parameterized_grid, get_parameterized_sobol
 from oao.strategy import GridStrategy
 
 
@@ -250,3 +250,98 @@ def _run_loop(self) -> None:
         # Log metrics.
         if self.monitor:
             self.monitor(self.client)
+
+
+class QuasiRandom(Optimizer):
+    def __init__(
+        self,
+        objective: Objective,
+        search_space: SearchSpace,
+        strategy: GenerationStrategy,
+        random_seed: Optional[int] = None,
+        monitor: Optional[callable] = None,
+    ) -> None:
+        """
+        Initialize Bayesian optimization strategy.
+
+        :param objective: Objective function for optimization; must be callable.
+        :type objective: Objective
+        :param search_space: Search space definition.
+        :type search_space: SearchSpace
+        :param strategy: Specify the generation strategy for optimization.
+        :type strategy: GenerationStrategy
+        :param random_seed: Set the random seed, defaults to None
+        :type random_seed: Optional[int], optional
+        :param monitor: Callable to log metrics, defaults to None
+        :type monitor: Optional[callable], optional
+        """
+        self.objective = objective
+        self.search_space = search_space
+        self.strategy = strategy
+        self.random_seed = random_seed
+        self.monitor = monitor
+        self.client = AxClient(generation_strategy=strategy, random_seed=random_seed)
+        self.batch_execution_times = []
+
+    def run(self, name: str = None) -> None:
+        """Runs the optimization using provided configurations."""
+        self._create_experiment(name)
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=FutureWarning)
+            self._run_steps()
+
+    def _create_experiment(self, name: str) -> None:
+        """
+        Create an Ax experiment.
+
+        :param name: Name of the experiment.
+        :type name: str
+        """
+        self.client.create_experiment(
+            name=name,
+            parameters=self.search_space.to_dict(),
+            objectives={self.objective.name: self.objective.properties},
+        )
+
+    def _run_loop(self, step: GenerationStep) -> None:
+        """
+        Run a single step of the generation strategy.
+
+        :param step: Contains optimization loop specification.
+        :type step: GenerationStep
+        """
+        if step.max_parallelism is None:
+            step.max_parallelism = 1
+
+        parameters = get_parameterized_sobol(
+            search_space=self.search_space, num_samples=step.num_trials
+        )
+
+        # Attach trials.
+        param_list, trial_indexes = list(
+            zip(*[self.client.attach_trial(parameters=p) for p in parameters])
+        )
+
+        # Start timer.
+        t0 = time.time()
+
+        # Evaluate trials.
+        with ThreadPoolExecutor(max_workers=step.max_parallelism) as executor:
+            results = executor.map(self.objective, param_list)
+
+        # Mark trials as complete and update model.
+        [
+            self.client.complete_trial(trial_index, raw_data=result)
+            for trial_index, result in zip(trial_indexes, results)
+        ]
+
+        # Log batch execution time.
+        self.batch_execution_times.extend([time.time() - t0] * len(trial_indexes))
+
+        # Log metrics.
+        if self.monitor:
+            self.monitor(self.client)
+
+    def _run_steps(self) -> None:
+        """Run the steps of the generation strategy."""
+        [self._run_loop(step) for step in self.strategy._steps]

oao/space.py

@@ -5,6 +5,7 @@
 from typing import Optional, Union
 
 import numpy as np
+import scipy.stats.qmc
 
 
 @dataclass
@@ -113,3 +114,36 @@ def get_parameterized_grid(
     return [
         {p.name: i for i, p in zip(row, search_space.parameters)} for row in param_array
     ]
+
+
+def get_parameterized_sobol(search_space: SearchSpace, num_samples: int) -> list[dict]:
+    """
+    Given a search space, return a list of parameterized Sobol sequences.
+
+    :param search_space: Search space to parameterize with Sobol sequence.
+    :type search_space: SearchSpace
+    :param num_samples: Number of samples to take along each dimension of the
+        Sobol sequence.
+    :return: A list of parameterized Sobol sequences.
+    :rtype: list[dict]
+    """
+    samples = scipy.stats.qmc.Sobol(d=len(search_space.parameters)).random(num_samples)
+
+    param_array = np.zeros_like(samples)
+    for i in range(len(search_space.parameters)):
+        sample_min = 0.0
+        sample_max = 1.0
+        sample_range = sample_max - sample_min
+
+        data_min = search_space.parameters[i].bounds[0]
+        data_max = search_space.parameters[i].bounds[1]
+        data_range = data_max - data_min
+
+        scaled_data = (
+            data_min + (samples[:, i] - sample_min) * data_range / sample_range
+        )
+        param_array[:, i] = scaled_data
+
+    return [
+        {p.name: i for i, p in zip(row, search_space.parameters)} for row in param_array
+    ]