From fe5d33b401c46cfa4bb75cbf9b09675c2b0db320 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Wed, 9 Aug 2023 15:14:15 +0100 Subject: [PATCH 01/16] Add multi TR notebook --- docs/notebooks/multi_trust_region.pct.py | 164 +++++++++++++++++++++++ 1 file changed, 164 insertions(+) create mode 100644 docs/notebooks/multi_trust_region.pct.py diff --git a/docs/notebooks/multi_trust_region.pct.py b/docs/notebooks/multi_trust_region.pct.py new file mode 100644 index 0000000000..45f373f34c --- /dev/null +++ b/docs/notebooks/multi_trust_region.pct.py @@ -0,0 +1,164 @@ +# --- +# jupyter: +# jupytext: +# cell_metadata_filter: -all +# custom_cell_magics: kql +# text_representation: +# extension: .py +# format_name: percent +# format_version: '1.3' +# jupytext_version: 1.11.2 +# kernelspec: +# display_name: .venv_310 +# language: python +# name: python3 +# --- + +# %% [markdown] +# # Batch trust region Bayesian optimization + +# %% +import numpy as np +import tensorflow as tf + +np.random.seed(1793) +tf.random.set_seed(1793) + +# %% [markdown] +# ## Define the problem and model +# +# You can use trust regions for Bayesian optimization in much the same way as we used EGO and EI in the [introduction notebook](expected_improvement.ipynb). Since the setup is much the same as in that tutorial, we'll skip over most of the detail. + +# %% +import trieste +from trieste.objectives import Branin + +branin = Branin.objective +search_space = Branin.search_space + +num_initial_data_points = 10 +initial_query_points = search_space.sample(num_initial_data_points) +observer = trieste.objectives.utils.mk_observer(branin) +initial_data = observer(initial_query_points) + +# %% [markdown] +# As usual, we'll use Gaussian process regression to model the function. Note that we set the likelihood variance to a small number because we are dealing with a noise-free problem. + +# %% +from trieste.models.gpflow import GaussianProcessRegression, build_gpr + +gpflow_model = build_gpr(initial_data, search_space, likelihood_variance=1e-7) +model = GaussianProcessRegression(gpflow_model) + + +# %% [markdown] +# ## Create the batch trust region acquisition rule +# +# We achieve Bayesian optimization with trust region by specifying `MultiTrustRegionBox` as the acquisition rule. +# +# This rule requires an initial number `num_query_points` of sub-spaces (or trust regions) to be provided and performs optimization in parallel across all these sub-spaces. Each region contributes one query point, resulting in each acquisition step collecting `num_query_points` points overall. As the optimization process continues, the bounds of these sub-spaces are dynamically updated. +# +# In addition, this rule requires the specification of a batch aquisition base-rule for performing optimization; for our example we use `EfficientGlobalOptimization` coupled with `ParallelContinuousThompsonSampling`. +# +# Note: the number of sub-spaces/regions must match the number of batch query points. + +# %% +num_query_points = 5 + +init_subspaces = [ + trieste.acquisition.rule.TrustRegionBox(search_space) + for _ in range(num_query_points) +] +base_rule = trieste.acquisition.rule.EfficientGlobalOptimization( # type: ignore[var-annotated] + builder=trieste.acquisition.ParallelContinuousThompsonSampling(), + num_query_points=num_query_points, +) +acq_rule = trieste.acquisition.rule.MultiTrustRegionBox( + init_subspaces, base_rule +) + +# %% [markdown] +# ## Run the optimization loop +# +# We can now run the Bayesian optimization loop by defining a `BayesianOptimizer` and calling its `optimize` method with the trust region rule. Once the optimization loop is complete, the optimizer will return `num_query_points` new query points for every step in the loop. With 5 steps, that's 25 points in total. + +# %% +bo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space) + +num_steps = 5 +result = bo.optimize( + num_steps, initial_data, model, acq_rule, track_state=False +) +dataset = result.try_get_final_dataset() + +# %% [markdown] +# ## Visualizing the result +# +# We can take a look at where we queried the observer, the original query points (crosses), new query points (dots) and the optimum point found (purple dot), and where they lie with respect to the contours of the Branin. + +# %% +from trieste.experimental.plotting import plot_bo_points, plot_function_2d + +arg_min_idx = tf.squeeze(tf.argmin(dataset.observations, axis=0)) +query_points = dataset.query_points.numpy() +observations = dataset.observations.numpy() +_, ax = plot_function_2d( + branin, + search_space.lower, + search_space.upper, + grid_density=40, + contour=True, +) + +plot_bo_points(query_points, ax[0, 0], num_initial_data_points, arg_min_idx) + +# %% [markdown] +# Here we visualize the observations on a three-dimensional plot of the Branin. We'll add the contours of the mean and variance of the model's predictive distribution as translucent surfaces. + +# %% +from trieste.experimental.plotting import plot_model_predictions_plotly + +fig = plot_model_predictions_plotly( + result.try_get_final_model(), + search_space.lower, + search_space.upper, +) + +from trieste.experimental.plotting import add_bo_points_plotly + +fig = add_bo_points_plotly( + x=query_points[:, 0], + y=query_points[:, 1], + z=observations[:, 0], + num_init=num_initial_data_points, + idx_best=arg_min_idx, + fig=fig, + figrow=1, + figcol=1, +) +fig.show() + +# %% [markdown] +# We can also visualize how each successive point compares with the current best by plotting regret. +# This plot shows the observations (crosses and dots), the current best (orange line), and the start of the optimization loop (blue line). + +# %% +import matplotlib.pyplot as plt + +from trieste.experimental.plotting import plot_regret + +suboptimality = observations - Branin.minimum.numpy() + +fig, ax = plt.subplots() +plot_regret( + suboptimality, ax, num_init=num_initial_data_points, idx_best=arg_min_idx +) + +ax.set_yscale("log") +ax.set_ylabel("Regret") +ax.set_xlabel("# evaluations") + +# %% [markdown] +# ## LICENSE +# +# [Apache License 2.0](https://github.com/secondmind-labs/trieste/blob/develop/LICENSE) From 686375977b92ca78fe95420bc6c450b1282b168c Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Tue, 22 Aug 2023 17:20:27 +0100 Subject: [PATCH 02/16] Rename to match upstream changes --- ...{multi_trust_region.pct.py => batch_trust_region.pct.py} | 6 +++--- docs/tutorials.rst | 1 + 2 files changed, 4 insertions(+), 3 deletions(-) rename docs/notebooks/{multi_trust_region.pct.py => batch_trust_region.pct.py} (96%) diff --git a/docs/notebooks/multi_trust_region.pct.py b/docs/notebooks/batch_trust_region.pct.py similarity index 96% rename from docs/notebooks/multi_trust_region.pct.py rename to docs/notebooks/batch_trust_region.pct.py index 45f373f34c..3d6b23514f 100644 --- a/docs/notebooks/multi_trust_region.pct.py +++ b/docs/notebooks/batch_trust_region.pct.py @@ -54,7 +54,7 @@ # %% [markdown] # ## Create the batch trust region acquisition rule # -# We achieve Bayesian optimization with trust region by specifying `MultiTrustRegionBox` as the acquisition rule. +# We achieve Bayesian optimization with trust region by specifying `BatchTrustRegionBox` as the acquisition rule. # # This rule requires an initial number `num_query_points` of sub-spaces (or trust regions) to be provided and performs optimization in parallel across all these sub-spaces. Each region contributes one query point, resulting in each acquisition step collecting `num_query_points` points overall. As the optimization process continues, the bounds of these sub-spaces are dynamically updated. # @@ -66,14 +66,14 @@ num_query_points = 5 init_subspaces = [ - trieste.acquisition.rule.TrustRegionBox(search_space) + trieste.acquisition.rule.SingleObjectiveTRBox(search_space) for _ in range(num_query_points) ] base_rule = trieste.acquisition.rule.EfficientGlobalOptimization( # type: ignore[var-annotated] builder=trieste.acquisition.ParallelContinuousThompsonSampling(), num_query_points=num_query_points, ) -acq_rule = trieste.acquisition.rule.MultiTrustRegionBox( +acq_rule = trieste.acquisition.rule.BatchTrustRegionBox( init_subspaces, base_rule ) diff --git a/docs/tutorials.rst b/docs/tutorials.rst index 72cfee53be..01784004d4 100644 --- a/docs/tutorials.rst +++ b/docs/tutorials.rst @@ -40,6 +40,7 @@ The following tutorials illustrate solving different types of optimization probl notebooks/qhsri-tutorial notebooks/multifidelity_modelling notebooks/rembo + notebooks/batch_trust_region Frequently asked questions -------------------------- From 277fd251da4c414af57be94d25706c7fc3f6455f Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Thu, 24 Aug 2023 11:48:23 +0100 Subject: [PATCH 03/16] Change box class name --- docs/notebooks/batch_trust_region.pct.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/notebooks/batch_trust_region.pct.py b/docs/notebooks/batch_trust_region.pct.py index 3d6b23514f..78f4f3c050 100644 --- a/docs/notebooks/batch_trust_region.pct.py +++ b/docs/notebooks/batch_trust_region.pct.py @@ -66,7 +66,7 @@ num_query_points = 5 init_subspaces = [ - trieste.acquisition.rule.SingleObjectiveTRBox(search_space) + trieste.acquisition.rule.SingleObjectiveTrustRegionBox(search_space) for _ in range(num_query_points) ] base_rule = trieste.acquisition.rule.EfficientGlobalOptimization( # type: ignore[var-annotated] From d0b85d4f82d05bb136b70d577f93bffaa33b1efe Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Thu, 24 Aug 2023 16:42:56 +0100 Subject: [PATCH 04/16] Add gif of trust region optimization --- docs/notebooks/batch_trust_region.pct.py | 86 ++++++++++++++++++++++- docs/notebooks/constraints.txt | 1 + docs/notebooks/requirements.txt | 1 + trieste/experimental/plotting/plotting.py | 6 +- 4 files changed, 88 insertions(+), 6 deletions(-) diff --git a/docs/notebooks/batch_trust_region.pct.py b/docs/notebooks/batch_trust_region.pct.py index 78f4f3c050..ab46471fe6 100644 --- a/docs/notebooks/batch_trust_region.pct.py +++ b/docs/notebooks/batch_trust_region.pct.py @@ -86,9 +86,7 @@ bo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space) num_steps = 5 -result = bo.optimize( - num_steps, initial_data, model, acq_rule, track_state=False -) +result = bo.optimize(num_steps, initial_data, model, acq_rule, track_state=True) dataset = result.try_get_final_dataset() # %% [markdown] @@ -158,6 +156,88 @@ ax.set_ylabel("Regret") ax.set_xlabel("# evaluations") +# %% [markdown] +# Next we visualize the progress of the optimization by plotting the trust regions at each step. The trust regions are shown as translucent boxes, with the current optimum point in each region shown in matching color. + +# %% +import base64 +import io + +import imageio +import IPython +from matplotlib.colors import rgb2hex +from matplotlib.patches import Rectangle +from matplotlib.pyplot import cm + +colors = [ + rgb2hex(color) for color in cm.rainbow(np.linspace(0, 1, num_query_points)) +] +frames = [] + +for step, hist in enumerate(result.history + [result.final_result.unwrap()]): + state = hist.acquisition_state + if state is None: + continue + + # Plot branin contour. + fig, ax = plot_function_2d( + branin, + search_space.lower, + search_space.upper, + grid_density=40, + contour=True, + ) + + query_points = hist.dataset.query_points + new_points_mask = np.zeros(query_points.shape[0], dtype=bool) + new_points_mask[-num_query_points:] = True + + assert isinstance(state, trieste.acquisition.rule.BatchTrustRegionBox.State) + acquisition_space = state.acquisition_space + + # Plot trust regions. + for i, tag in enumerate(acquisition_space.subspace_tags): + lb = acquisition_space.get_subspace(tag).lower + ub = acquisition_space.get_subspace(tag).upper + ax[0, 0].add_patch( + Rectangle( + (lb[0], lb[1]), + ub[0] - lb[0], + ub[1] - lb[1], + facecolor=colors[i], + edgecolor=colors[i], + alpha=0.3, + ) + ) + + # Plot new query points, using failure mask to color them. + plot_bo_points( + query_points, + ax[0, 0], + num_initial_data_points, + mask_fail=new_points_mask, + c_pass="black", + c_fail=colors, + ) + + fig.suptitle(f"step number {step}") + fig.canvas.draw() + size_pix = fig.get_size_inches() * fig.dpi + image = np.frombuffer(fig.canvas.tostring_rgb(), dtype="uint8") + frames.append(image.reshape(list(size_pix[::-1].astype(int)) + [3])) + plt.close(fig) + + +# Create and show the GIF. +gif_file = io.BytesIO() +imageio.mimsave(gif_file, frames, format="gif", loop=0, duration=5000) # type: ignore +gif = IPython.display.HTML( + ''.format( + base64.b64encode(gif_file.getvalue()).decode() + ) +) +IPython.display.display(gif) + # %% [markdown] # ## LICENSE # diff --git a/docs/notebooks/constraints.txt b/docs/notebooks/constraints.txt index 47b2c39cfd..724356c348 100644 --- a/docs/notebooks/constraints.txt +++ b/docs/notebooks/constraints.txt @@ -68,6 +68,7 @@ gym==0.26.2 gym-notices==0.0.8 h5py==3.8.0 idna==3.4 +imageio==2.31.1 ipykernel==6.23.2 ipython==8.14.0 isoduration==20.11.0 diff --git a/docs/notebooks/requirements.txt b/docs/notebooks/requirements.txt index f3ffb86896..cc0cdaed07 100644 --- a/docs/notebooks/requirements.txt +++ b/docs/notebooks/requirements.txt @@ -21,3 +21,4 @@ jupytext gym[box2d] box2d box2d-kengz +imageio diff --git a/trieste/experimental/plotting/plotting.py b/trieste/experimental/plotting/plotting.py index f5e26b1aa8..f4a8d8722a 100644 --- a/trieste/experimental/plotting/plotting.py +++ b/trieste/experimental/plotting/plotting.py @@ -14,7 +14,7 @@ from __future__ import annotations -from typing import Callable, Optional, Sequence +from typing import Callable, List, Optional, Sequence, Union import matplotlib.pyplot as plt import numpy as np @@ -234,7 +234,7 @@ def format_point_markers( m_init: str = "x", m_add: str = "o", c_pass: str = "tab:green", - c_fail: str = "tab:red", + c_fail: Union[str, List[str]] = "tab:red", c_best: str = "tab:purple", ) -> tuple[TensorType, TensorType]: """ @@ -275,7 +275,7 @@ def plot_bo_points( m_init: str = "x", m_add: str = "o", c_pass: str = "tab:green", - c_fail: str = "tab:red", + c_fail: Union[str, List[str]] = "tab:red", c_best: str = "tab:purple", ) -> None: """ From 964871c3ffefc962e2eecb4f01a1b303c6613634 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Fri, 25 Aug 2023 14:47:28 +0100 Subject: [PATCH 05/16] Remove unused type ignore --- docs/notebooks/batch_trust_region.pct.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/notebooks/batch_trust_region.pct.py b/docs/notebooks/batch_trust_region.pct.py index ab46471fe6..95b8ba13f1 100644 --- a/docs/notebooks/batch_trust_region.pct.py +++ b/docs/notebooks/batch_trust_region.pct.py @@ -230,7 +230,7 @@ # Create and show the GIF. gif_file = io.BytesIO() -imageio.mimsave(gif_file, frames, format="gif", loop=0, duration=5000) # type: ignore +imageio.mimsave(gif_file, frames, format="gif", loop=0, duration=5000) gif = IPython.display.HTML( ''.format( base64.b64encode(gif_file.getvalue()).decode() From 3f68541fd63db6c6a6e7b561e624a1f1512a61a8 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Wed, 6 Sep 2023 15:05:30 +0100 Subject: [PATCH 06/16] Move some TR plotting code to exp package --- docs/notebooks/batch_trust_region.pct.py | 103 ++++++++-------------- trieste/experimental/plotting/__init__.py | 1 + trieste/experimental/plotting/plotting.py | 80 +++++++++++++++++ 3 files changed, 120 insertions(+), 64 deletions(-) diff --git a/docs/notebooks/batch_trust_region.pct.py b/docs/notebooks/batch_trust_region.pct.py index 95b8ba13f1..46c467112f 100644 --- a/docs/notebooks/batch_trust_region.pct.py +++ b/docs/notebooks/batch_trust_region.pct.py @@ -137,7 +137,7 @@ fig.show() # %% [markdown] -# We can also visualize how each successive point compares with the current best by plotting regret. +# Next we visualize how each successive point compares with the current best by plotting regret. # This plot shows the observations (crosses and dots), the current best (orange line), and the start of the optimization loop (blue line). # %% @@ -157,7 +157,9 @@ ax.set_xlabel("# evaluations") # %% [markdown] -# Next we visualize the progress of the optimization by plotting the trust regions at each step. The trust regions are shown as translucent boxes, with the current optimum point in each region shown in matching color. +# We can also visualize the progress of the optimization by plotting the trust regions at each step. +# +# However, first we define some helper functions for plotting the trust regions. # %% import base64 @@ -165,78 +167,51 @@ import imageio import IPython -from matplotlib.colors import rgb2hex -from matplotlib.patches import Rectangle -from matplotlib.pyplot import cm -colors = [ - rgb2hex(color) for color in cm.rainbow(np.linspace(0, 1, num_query_points)) -] -frames = [] +from trieste.experimental.plotting import plot_trust_region_history_2d -for step, hist in enumerate(result.history + [result.final_result.unwrap()]): - state = hist.acquisition_state - if state is None: - continue - # Plot branin contour. - fig, ax = plot_function_2d( - branin, - search_space.lower, - search_space.upper, - grid_density=40, - contour=True, - ) +def _fig_to_frame(fig: plt.Figure) -> np.ndarray: + fig.canvas.draw() + size_pix = fig.get_size_inches() * fig.dpi + image = np.frombuffer(fig.canvas.tostring_rgb(), dtype="uint8") + return image.reshape(list(size_pix[::-1].astype(int)) + [3]) - query_points = hist.dataset.query_points - new_points_mask = np.zeros(query_points.shape[0], dtype=bool) - new_points_mask[-num_query_points:] = True - - assert isinstance(state, trieste.acquisition.rule.BatchTrustRegionBox.State) - acquisition_space = state.acquisition_space - - # Plot trust regions. - for i, tag in enumerate(acquisition_space.subspace_tags): - lb = acquisition_space.get_subspace(tag).lower - ub = acquisition_space.get_subspace(tag).upper - ax[0, 0].add_patch( - Rectangle( - (lb[0], lb[1]), - ub[0] - lb[0], - ub[1] - lb[1], - facecolor=colors[i], - edgecolor=colors[i], - alpha=0.3, - ) - ) - # Plot new query points, using failure mask to color them. - plot_bo_points( - query_points, - ax[0, 0], - num_initial_data_points, - mask_fail=new_points_mask, - c_pass="black", - c_fail=colors, +def _frames_to_gif( + frames: list[np.ndarray], duration=5000 +) -> IPython.display.HTML: + gif_file = io.BytesIO() + imageio.mimsave(gif_file, frames, format="gif", loop=0, duration=duration) + gif = IPython.display.HTML( + ''.format( + base64.b64encode(gif_file.getvalue()).decode() + ) ) + return gif - fig.suptitle(f"step number {step}") - fig.canvas.draw() - size_pix = fig.get_size_inches() * fig.dpi - image = np.frombuffer(fig.canvas.tostring_rgb(), dtype="uint8") - frames.append(image.reshape(list(size_pix[::-1].astype(int)) + [3])) - plt.close(fig) +# %% [markdown] +# The trust regions are shown as translucent boxes, with the current optimum point in each region shown in matching color. -# Create and show the GIF. -gif_file = io.BytesIO() -imageio.mimsave(gif_file, frames, format="gif", loop=0, duration=5000) -gif = IPython.display.HTML( - ''.format( - base64.b64encode(gif_file.getvalue()).decode() +# %% +frames = [] +for step, hist in enumerate(result.history + [result.final_result.unwrap()]): + fig, _ = plot_trust_region_history_2d( + branin, + search_space.lower, + search_space.upper, + hist, + num_query_points, + num_initial_data_points, ) -) -IPython.display.display(gif) + + if fig is not None: + fig.suptitle(f"step number {step}") + frames.append(_fig_to_frame(fig)) + plt.close(fig) + +IPython.display.display(_frames_to_gif(frames)) # %% [markdown] # ## LICENSE diff --git a/trieste/experimental/plotting/__init__.py b/trieste/experimental/plotting/__init__.py index 8e6a8345be..7125d30aff 100644 --- a/trieste/experimental/plotting/__init__.py +++ b/trieste/experimental/plotting/__init__.py @@ -30,6 +30,7 @@ plot_mobo_history, plot_mobo_points_in_obj_space, plot_regret, + plot_trust_region_history_2d, ) from .plotting_plotly import ( add_bo_points_plotly, diff --git a/trieste/experimental/plotting/plotting.py b/trieste/experimental/plotting/plotting.py index f4a8d8722a..fda808573c 100644 --- a/trieste/experimental/plotting/plotting.py +++ b/trieste/experimental/plotting/plotting.py @@ -23,11 +23,15 @@ from matplotlib import cm from matplotlib.axes import Axes from matplotlib.collections import Collection +from matplotlib.colors import rgb2hex from matplotlib.contour import ContourSet from matplotlib.figure import Figure +from matplotlib.patches import Rectangle from trieste.acquisition import AcquisitionFunction from trieste.acquisition.multi_objective.dominance import non_dominated +from trieste.bayesian_optimizer import FrozenRecord, Record, StateType +from trieste.space import TaggedMultiSearchSpace from trieste.types import TensorType from trieste.utils import to_numpy @@ -534,3 +538,79 @@ def plot_gp_2d( axx.set_ylim(mins[1], maxs[1]) return fig, ax + + +def plot_trust_region_history_2d( + obj_func: Callable[[TensorType], TensorType], + mins: TensorType, + maxs: TensorType, + history: Record[StateType] | FrozenRecord[StateType], + num_query_points: int, + num_init: Optional[int] = None, +) -> tuple[Optional[Figure], Optional[Axes]]: + """ + Plot the contour of the objective function, query points and the trust regions for a particular + step of the optimization process. + + :param obj_func: the objective function that returns a n-array given a [n, d] array + :param mins: search space 2D lower bounds + :param maxs: search space 2D upper bounds + :param history: the optimization history for a particular step of the optimization process + :param num_query_points: number of query points in this step + :param num_init: initial number of BO points + :return: figure and axes + """ + + state = history.acquisition_state + if state is None: + # If state is None, then there is no trust region state to plot. + return None, None + + # Plot objective contour. + fig, ax = plot_function_2d( + obj_func, + mins, + maxs, + grid_density=40, + contour=True, + ) + + query_points = history.dataset.query_points + new_points_mask = np.zeros(query_points.shape[0], dtype=bool) + new_points_mask[-num_query_points:] = True + + assert hasattr(state, "acquisition_space") + acquisition_space = state.acquisition_space + + colors = [rgb2hex(color) for color in cm.rainbow(np.linspace(0, 1, num_query_points))] + + # Plot trust regions. + if isinstance(acquisition_space, TaggedMultiSearchSpace): + spaces = [acquisition_space.get_subspace(tag) for tag in acquisition_space.subspace_tags] + else: + spaces = [acquisition_space] + for i, space in enumerate(spaces): + lb = space.lower + ub = space.upper + ax[0, 0].add_patch( + Rectangle( + (lb[0], lb[1]), + ub[0] - lb[0], + ub[1] - lb[1], + facecolor=colors[i], + edgecolor=colors[i], + alpha=0.3, + ) + ) + + # Plot new query points, using failure mask to color them. + plot_bo_points( + query_points, + ax[0, 0], + num_init, + mask_fail=new_points_mask, + c_pass="black", + c_fail=colors, + ) + + return fig, ax From c762ab42343c4d92fc1585e916a899e3115db5af Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Wed, 6 Sep 2023 17:10:13 +0100 Subject: [PATCH 07/16] Add TREGO and TurBO to notebook --- docs/notebooks/batch_trust_region.pct.py | 219 ----------------- docs/notebooks/trust_region.pct.py | 283 ++++++++++++++++++++++ trieste/experimental/plotting/plotting.py | 19 +- 3 files changed, 294 insertions(+), 227 deletions(-) delete mode 100644 docs/notebooks/batch_trust_region.pct.py create mode 100644 docs/notebooks/trust_region.pct.py diff --git a/docs/notebooks/batch_trust_region.pct.py b/docs/notebooks/batch_trust_region.pct.py deleted file mode 100644 index 46c467112f..0000000000 --- a/docs/notebooks/batch_trust_region.pct.py +++ /dev/null @@ -1,219 +0,0 @@ -# --- -# jupyter: -# jupytext: -# cell_metadata_filter: -all -# custom_cell_magics: kql -# text_representation: -# extension: .py -# format_name: percent -# format_version: '1.3' -# jupytext_version: 1.11.2 -# kernelspec: -# display_name: .venv_310 -# language: python -# name: python3 -# --- - -# %% [markdown] -# # Batch trust region Bayesian optimization - -# %% -import numpy as np -import tensorflow as tf - -np.random.seed(1793) -tf.random.set_seed(1793) - -# %% [markdown] -# ## Define the problem and model -# -# You can use trust regions for Bayesian optimization in much the same way as we used EGO and EI in the [introduction notebook](expected_improvement.ipynb). Since the setup is much the same as in that tutorial, we'll skip over most of the detail. - -# %% -import trieste -from trieste.objectives import Branin - -branin = Branin.objective -search_space = Branin.search_space - -num_initial_data_points = 10 -initial_query_points = search_space.sample(num_initial_data_points) -observer = trieste.objectives.utils.mk_observer(branin) -initial_data = observer(initial_query_points) - -# %% [markdown] -# As usual, we'll use Gaussian process regression to model the function. Note that we set the likelihood variance to a small number because we are dealing with a noise-free problem. - -# %% -from trieste.models.gpflow import GaussianProcessRegression, build_gpr - -gpflow_model = build_gpr(initial_data, search_space, likelihood_variance=1e-7) -model = GaussianProcessRegression(gpflow_model) - - -# %% [markdown] -# ## Create the batch trust region acquisition rule -# -# We achieve Bayesian optimization with trust region by specifying `BatchTrustRegionBox` as the acquisition rule. -# -# This rule requires an initial number `num_query_points` of sub-spaces (or trust regions) to be provided and performs optimization in parallel across all these sub-spaces. Each region contributes one query point, resulting in each acquisition step collecting `num_query_points` points overall. As the optimization process continues, the bounds of these sub-spaces are dynamically updated. -# -# In addition, this rule requires the specification of a batch aquisition base-rule for performing optimization; for our example we use `EfficientGlobalOptimization` coupled with `ParallelContinuousThompsonSampling`. -# -# Note: the number of sub-spaces/regions must match the number of batch query points. - -# %% -num_query_points = 5 - -init_subspaces = [ - trieste.acquisition.rule.SingleObjectiveTrustRegionBox(search_space) - for _ in range(num_query_points) -] -base_rule = trieste.acquisition.rule.EfficientGlobalOptimization( # type: ignore[var-annotated] - builder=trieste.acquisition.ParallelContinuousThompsonSampling(), - num_query_points=num_query_points, -) -acq_rule = trieste.acquisition.rule.BatchTrustRegionBox( - init_subspaces, base_rule -) - -# %% [markdown] -# ## Run the optimization loop -# -# We can now run the Bayesian optimization loop by defining a `BayesianOptimizer` and calling its `optimize` method with the trust region rule. Once the optimization loop is complete, the optimizer will return `num_query_points` new query points for every step in the loop. With 5 steps, that's 25 points in total. - -# %% -bo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space) - -num_steps = 5 -result = bo.optimize(num_steps, initial_data, model, acq_rule, track_state=True) -dataset = result.try_get_final_dataset() - -# %% [markdown] -# ## Visualizing the result -# -# We can take a look at where we queried the observer, the original query points (crosses), new query points (dots) and the optimum point found (purple dot), and where they lie with respect to the contours of the Branin. - -# %% -from trieste.experimental.plotting import plot_bo_points, plot_function_2d - -arg_min_idx = tf.squeeze(tf.argmin(dataset.observations, axis=0)) -query_points = dataset.query_points.numpy() -observations = dataset.observations.numpy() -_, ax = plot_function_2d( - branin, - search_space.lower, - search_space.upper, - grid_density=40, - contour=True, -) - -plot_bo_points(query_points, ax[0, 0], num_initial_data_points, arg_min_idx) - -# %% [markdown] -# Here we visualize the observations on a three-dimensional plot of the Branin. We'll add the contours of the mean and variance of the model's predictive distribution as translucent surfaces. - -# %% -from trieste.experimental.plotting import plot_model_predictions_plotly - -fig = plot_model_predictions_plotly( - result.try_get_final_model(), - search_space.lower, - search_space.upper, -) - -from trieste.experimental.plotting import add_bo_points_plotly - -fig = add_bo_points_plotly( - x=query_points[:, 0], - y=query_points[:, 1], - z=observations[:, 0], - num_init=num_initial_data_points, - idx_best=arg_min_idx, - fig=fig, - figrow=1, - figcol=1, -) -fig.show() - -# %% [markdown] -# Next we visualize how each successive point compares with the current best by plotting regret. -# This plot shows the observations (crosses and dots), the current best (orange line), and the start of the optimization loop (blue line). - -# %% -import matplotlib.pyplot as plt - -from trieste.experimental.plotting import plot_regret - -suboptimality = observations - Branin.minimum.numpy() - -fig, ax = plt.subplots() -plot_regret( - suboptimality, ax, num_init=num_initial_data_points, idx_best=arg_min_idx -) - -ax.set_yscale("log") -ax.set_ylabel("Regret") -ax.set_xlabel("# evaluations") - -# %% [markdown] -# We can also visualize the progress of the optimization by plotting the trust regions at each step. -# -# However, first we define some helper functions for plotting the trust regions. - -# %% -import base64 -import io - -import imageio -import IPython - -from trieste.experimental.plotting import plot_trust_region_history_2d - - -def _fig_to_frame(fig: plt.Figure) -> np.ndarray: - fig.canvas.draw() - size_pix = fig.get_size_inches() * fig.dpi - image = np.frombuffer(fig.canvas.tostring_rgb(), dtype="uint8") - return image.reshape(list(size_pix[::-1].astype(int)) + [3]) - - -def _frames_to_gif( - frames: list[np.ndarray], duration=5000 -) -> IPython.display.HTML: - gif_file = io.BytesIO() - imageio.mimsave(gif_file, frames, format="gif", loop=0, duration=duration) - gif = IPython.display.HTML( - ''.format( - base64.b64encode(gif_file.getvalue()).decode() - ) - ) - return gif - - -# %% [markdown] -# The trust regions are shown as translucent boxes, with the current optimum point in each region shown in matching color. - -# %% -frames = [] -for step, hist in enumerate(result.history + [result.final_result.unwrap()]): - fig, _ = plot_trust_region_history_2d( - branin, - search_space.lower, - search_space.upper, - hist, - num_query_points, - num_initial_data_points, - ) - - if fig is not None: - fig.suptitle(f"step number {step}") - frames.append(_fig_to_frame(fig)) - plt.close(fig) - -IPython.display.display(_frames_to_gif(frames)) - -# %% [markdown] -# ## LICENSE -# -# [Apache License 2.0](https://github.com/secondmind-labs/trieste/blob/develop/LICENSE) diff --git a/docs/notebooks/trust_region.pct.py b/docs/notebooks/trust_region.pct.py new file mode 100644 index 0000000000..e5bc4d2cba --- /dev/null +++ b/docs/notebooks/trust_region.pct.py @@ -0,0 +1,283 @@ +# --- +# jupyter: +# jupytext: +# cell_metadata_filter: -all +# custom_cell_magics: kql +# text_representation: +# extension: .py +# format_name: percent +# format_version: '1.3' +# jupytext_version: 1.11.2 +# kernelspec: +# display_name: .venv_310 +# language: python +# name: python3 +# --- + +# %% [markdown] +# # Trust region Bayesian optimization +# +# We will demonstrate three trust region Bayesian optimization algorithms in this tutorial. + +# %% +import numpy as np +import tensorflow as tf + +np.random.seed(1793) +tf.random.set_seed(1793) + +# %% [markdown] +# ## Define the problem and model +# +# We can use trust regions for Bayesian optimization in much the same way as we used EGO and EI in +# the [introduction notebook](expected_improvement.ipynb). Since the setup is very similar to +# that tutorial, we'll skip over most of the detail. + +# %% +import trieste +from trieste.objectives import Branin + +branin = Branin.objective +search_space = Branin.search_space + +num_initial_data_points = 10 +initial_query_points = search_space.sample(num_initial_data_points) +observer = trieste.objectives.utils.mk_observer(branin) +initial_data = observer(initial_query_points) + +# %% [markdown] +# As usual, we'll use Gaussian process regression to model the function. Note that we set the +# likelihood variance to a small number because we are dealing with a noise-free problem. + +# %% +from trieste.models.gpflow import GaussianProcessRegression, build_gpr + + +def build_model(): + gpflow_model = build_gpr( + initial_data, search_space, likelihood_variance=1e-7 + ) + return GaussianProcessRegression(gpflow_model) + + +# %% [markdown] +# ## Trust region `TREGO` acquisition rule +# +# First we show how to run Bayesian optimization with the `TREGO` algorithm. This is a trust region +# algorithm that alternates between regular EGO steps and local steps within a trust region. +# +# ### Create the rule and run the optimization loop +# +# We can run the Bayesian optimization loop by defining a `BayesianOptimizer` and calling its +# `optimize` method with the trust region rule. Once the optimization loop is complete, the +# optimizer will return one new query point for every step in the loop. + +# %% +acq_rule = trieste.acquisition.rule.TrustRegion() +bo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space) + +num_steps = 5 +result = bo.optimize( + num_steps, initial_data, build_model(), acq_rule, track_state=True +) +dataset = result.try_get_final_dataset() + +# %% [markdown] +# ### Visualizing the result +# +# Let's take a look at where we queried the observer, the original query points (crosses), new +# query points (dots) and the optimum point found (purple dot), and where they lie with respect to +# the contours of the Branin. + +# %% +from trieste.experimental.plotting import plot_bo_points, plot_function_2d + + +def plot_final_result(_dataset: trieste.data.Dataset) -> None: + arg_min_idx = tf.squeeze(tf.argmin(_dataset.observations, axis=0)) + query_points = _dataset.query_points.numpy() + _, ax = plot_function_2d( + branin, + search_space.lower, + search_space.upper, + grid_density=40, + contour=True, + ) + + plot_bo_points(query_points, ax[0, 0], num_initial_data_points, arg_min_idx) + + +plot_final_result(dataset) + +# %% [markdown] +# We can also visualize the progress of the optimization by plotting the trust regions at each step. +# The trust regions are shown as translucent boxes, with the current optimum point in each region +# shown in matching color. +# +# Note there is only one trust region in this plot, but the rule in the next section will show multiple trust +# regions. + +# %% +import base64 +import io +from typing import List + +import imageio +import IPython +import matplotlib.pyplot as plt + +from trieste.experimental.plotting import plot_trust_region_history_2d + + +def fig_to_frame(fig: plt.Figure) -> np.ndarray: + fig.canvas.draw() + size_pix = fig.get_size_inches() * fig.dpi + image = np.frombuffer(fig.canvas.tostring_rgb(), dtype="uint8") + return image.reshape(list(size_pix[::-1].astype(int)) + [3]) + + +def frames_to_gif( + frames: List[np.ndarray], duration=5000 +) -> IPython.display.HTML: + gif_file = io.BytesIO() + imageio.mimsave(gif_file, frames, format="gif", loop=0, duration=duration) + gif = IPython.display.HTML( + ''.format( + base64.b64encode(gif_file.getvalue()).decode() + ) + ) + return gif + + +def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: + frames = [] + for step, hist in enumerate( + result.history + [result.final_result.unwrap()] + ): + fig, _ = plot_trust_region_history_2d( + branin, + search_space.lower, + search_space.upper, + hist, + num_init=num_initial_data_points, + ) + + if fig is not None: + fig.suptitle(f"step number {step}") + frames.append(fig_to_frame(fig)) + plt.close(fig) + + IPython.display.display(frames_to_gif(frames)) + + +plot_history(result) + +# %% [markdown] +# ## Batch trust region rule +# +# Next we demonstrate how to run Bayesian optimization with the batch trust region rule. +# +# ### Create the batch trust region acquisition rule +# +# We achieve Bayesian optimization with trust region by specifying `BatchTrustRegionBox` as the +# acquisition rule. +# +# This rule needs an initial number `num_query_points` of sub-spaces (or trust regions) to be +# provided and performs optimization in parallel across all these sub-spaces. Each region +# contributes one query point, resulting in each acquisition step collecting `num_query_points` +# points overall. As the optimization process continues, the bounds of these sub-spaces are +# dynamically updated. +# +# In addition, this rule requires the specification of a batch aquisition base-rule for performing +# optimization; for our example we use `EfficientGlobalOptimization` coupled with +# `ParallelContinuousThompsonSampling`. +# +# Note: the number of sub-spaces/regions must match the number of batch query points. + +# %% +num_query_points = 5 + +init_subspaces = [ + trieste.acquisition.rule.SingleObjectiveTrustRegionBox(search_space) + for _ in range(num_query_points) +] +base_rule = trieste.acquisition.rule.EfficientGlobalOptimization( # type: ignore[var-annotated] + builder=trieste.acquisition.ParallelContinuousThompsonSampling(), + num_query_points=num_query_points, +) +acq_rule = trieste.acquisition.rule.BatchTrustRegionBox( # type: ignore[assignment] + init_subspaces, base_rule +) + +# %% [markdown] +# ### Run the optimization loop +# +# We run the Bayesian optimization loop as before by defining a `BayesianOptimizer` and calling its +# `optimize` method with the trust region rule. Once the optimization loop is complete, the +# optimizer will return `num_query_points` new query points for every step in the loop. With +# 5 steps, that's 25 points in total. + +# %% +bo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space) + +num_steps = 5 +result = bo.optimize( + num_steps, initial_data, build_model(), acq_rule, track_state=True +) +dataset = result.try_get_final_dataset() + +# %% [markdown] +# ### Visualizing the result +# +# Next we visualize the results as before. + +# %% +plot_final_result(dataset) + +# %% +plot_history(result) + +# %% [markdown] +# ## Trust region `TurBO` acquisition rule +# +# Finally we show how to run Bayesian optimization with the `TurBO` algorithm. This is a +# trust region algorithm that uses local models and datasets to approximate the objective function +# within the trust region. +# +# ### Create the rule and run the optimization loop +# +# This rule requires the specification of a aquisition base-rule for performing +# optimization within the trust region; for our example we use `DiscreteThompsonSampling`. + +# %% +acq_rule = trieste.acquisition.rule.TURBO( # type: ignore[assignment] + search_space, rule=trieste.acquisition.rule.DiscreteThompsonSampling(500, 3) +) +bo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space) + +num_steps = 5 +result = bo.optimize( + num_steps, + initial_data, + build_model(), + acq_rule, + track_state=True, + fit_model=False, +) +dataset = result.try_get_final_dataset() + +# %% [markdown] +# ### Visualizing the result +# +# Here are the plots `TurBO`. + +# %% +plot_final_result(dataset) + +# %% +plot_history(result) + +# %% [markdown] +# ## LICENSE +# +# [Apache License 2.0](https://github.com/secondmind-labs/trieste/blob/develop/LICENSE) diff --git a/trieste/experimental/plotting/plotting.py b/trieste/experimental/plotting/plotting.py index fda808573c..ed17d820c5 100644 --- a/trieste/experimental/plotting/plotting.py +++ b/trieste/experimental/plotting/plotting.py @@ -545,7 +545,7 @@ def plot_trust_region_history_2d( mins: TensorType, maxs: TensorType, history: Record[StateType] | FrozenRecord[StateType], - num_query_points: int, + num_query_points: Optional[int] = None, num_init: Optional[int] = None, ) -> tuple[Optional[Figure], Optional[Axes]]: """ @@ -575,20 +575,23 @@ def plot_trust_region_history_2d( contour=True, ) - query_points = history.dataset.query_points - new_points_mask = np.zeros(query_points.shape[0], dtype=bool) - new_points_mask[-num_query_points:] = True - assert hasattr(state, "acquisition_space") acquisition_space = state.acquisition_space - colors = [rgb2hex(color) for color in cm.rainbow(np.linspace(0, 1, num_query_points))] - - # Plot trust regions. if isinstance(acquisition_space, TaggedMultiSearchSpace): spaces = [acquisition_space.get_subspace(tag) for tag in acquisition_space.subspace_tags] else: spaces = [acquisition_space] + + if num_query_points is None: + num_query_points = len(spaces) + + query_points = history.dataset.query_points + new_points_mask = np.zeros(query_points.shape[0], dtype=bool) + new_points_mask[-num_query_points:] = True + + # Plot trust regions. + colors = [rgb2hex(color) for color in cm.rainbow(np.linspace(0, 1, num_query_points))] for i, space in enumerate(spaces): lb = space.lower ub = space.upper From 8ce699c366d08687289b24ae03e60b54a28c7a68 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Wed, 6 Sep 2023 17:15:09 +0100 Subject: [PATCH 08/16] Fix ref name --- docs/tutorials.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/tutorials.rst b/docs/tutorials.rst index 01784004d4..749a2eca19 100644 --- a/docs/tutorials.rst +++ b/docs/tutorials.rst @@ -40,7 +40,7 @@ The following tutorials illustrate solving different types of optimization probl notebooks/qhsri-tutorial notebooks/multifidelity_modelling notebooks/rembo - notebooks/batch_trust_region + notebooks/trust_region Frequently asked questions -------------------------- From 53f3b2673cafd32f3c0cf16b32fc9696ffa03dd3 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Thu, 7 Sep 2023 10:03:17 +0100 Subject: [PATCH 09/16] Fix duplicate labels --- docs/notebooks/trust_region.pct.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/notebooks/trust_region.pct.py b/docs/notebooks/trust_region.pct.py index e5bc4d2cba..5493d72065 100644 --- a/docs/notebooks/trust_region.pct.py +++ b/docs/notebooks/trust_region.pct.py @@ -83,7 +83,7 @@ def build_model(): dataset = result.try_get_final_dataset() # %% [markdown] -# ### Visualizing the result +# ### Visualizing `TREGO` results # # Let's take a look at where we queried the observer, the original query points (crosses), new # query points (dots) and the optimum point found (purple dot), and where they lie with respect to @@ -227,7 +227,7 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: dataset = result.try_get_final_dataset() # %% [markdown] -# ### Visualizing the result +# ### Visualizing batch trust region results # # Next we visualize the results as before. @@ -267,7 +267,7 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: dataset = result.try_get_final_dataset() # %% [markdown] -# ### Visualizing the result +# ### Visualizing `TurBO` results # # Here are the plots `TurBO`. From 085868eb80ccee8f803f2ef0833f25d94af5a9dd Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Thu, 7 Sep 2023 10:27:24 +0100 Subject: [PATCH 10/16] Improve text a bit --- docs/notebooks/trust_region.pct.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/notebooks/trust_region.pct.py b/docs/notebooks/trust_region.pct.py index 5493d72065..ef2c7a29f1 100644 --- a/docs/notebooks/trust_region.pct.py +++ b/docs/notebooks/trust_region.pct.py @@ -64,13 +64,13 @@ def build_model(): # ## Trust region `TREGO` acquisition rule # # First we show how to run Bayesian optimization with the `TREGO` algorithm. This is a trust region -# algorithm that alternates between regular EGO steps and local steps within a trust region. +# algorithm that alternates between regular EGO steps and local steps within one trust region. # # ### Create the rule and run the optimization loop # # We can run the Bayesian optimization loop by defining a `BayesianOptimizer` and calling its # `optimize` method with the trust region rule. Once the optimization loop is complete, the -# optimizer will return one new query point for every step in the loop. +# optimizer will return one new query point for every step in the loop; that's 5 points in total. # %% acq_rule = trieste.acquisition.rule.TrustRegion() @@ -229,7 +229,7 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: # %% [markdown] # ### Visualizing batch trust region results # -# Next we visualize the results as before. +# We visualize the results as before. # %% plot_final_result(dataset) @@ -240,13 +240,13 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: # %% [markdown] # ## Trust region `TurBO` acquisition rule # -# Finally we show how to run Bayesian optimization with the `TurBO` algorithm. This is a +# Finally, we show how to run Bayesian optimization with the `TurBO` algorithm. This is a # trust region algorithm that uses local models and datasets to approximate the objective function -# within the trust region. +# within one trust region. # # ### Create the rule and run the optimization loop # -# This rule requires the specification of a aquisition base-rule for performing +# This rule requires the specification of an aquisition base-rule for performing # optimization within the trust region; for our example we use `DiscreteThompsonSampling`. # %% @@ -269,7 +269,7 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: # %% [markdown] # ### Visualizing `TurBO` results # -# Here are the plots `TurBO`. +# We display the results as earlier. # %% plot_final_result(dataset) From 66de1402d3eb9dfa6ca20e74050b6ff51502d934 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Thu, 7 Sep 2023 11:39:27 +0100 Subject: [PATCH 11/16] Fix label duplicate --- docs/notebooks/trust_region.pct.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/notebooks/trust_region.pct.py b/docs/notebooks/trust_region.pct.py index ef2c7a29f1..02a274da7a 100644 --- a/docs/notebooks/trust_region.pct.py +++ b/docs/notebooks/trust_region.pct.py @@ -66,7 +66,7 @@ def build_model(): # First we show how to run Bayesian optimization with the `TREGO` algorithm. This is a trust region # algorithm that alternates between regular EGO steps and local steps within one trust region. # -# ### Create the rule and run the optimization loop +# ### Create `TREGO` rule and run optimization loop # # We can run the Bayesian optimization loop by defining a `BayesianOptimizer` and calling its # `optimize` method with the trust region rule. Once the optimization loop is complete, the @@ -244,7 +244,7 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: # trust region algorithm that uses local models and datasets to approximate the objective function # within one trust region. # -# ### Create the rule and run the optimization loop +# ### Create `TurBO` rule and run optimization loop # # This rule requires the specification of an aquisition base-rule for performing # optimization within the trust region; for our example we use `DiscreteThompsonSampling`. From 86545cccfb187a7a2caeb5f2284749d88a130921 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Fri, 8 Sep 2023 14:27:26 +0100 Subject: [PATCH 12/16] Add note about fit_model=False --- docs/notebooks/trust_region.pct.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/notebooks/trust_region.pct.py b/docs/notebooks/trust_region.pct.py index 02a274da7a..e9d44a6df7 100644 --- a/docs/notebooks/trust_region.pct.py +++ b/docs/notebooks/trust_region.pct.py @@ -248,6 +248,9 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: # # This rule requires the specification of an aquisition base-rule for performing # optimization within the trust region; for our example we use `DiscreteThompsonSampling`. +# +# Note that we switch off global model fitting by setting `fit_model=False`. This is because +# `TurBO` uses a local model and fitting the global model would be redundant and wasteful. # %% acq_rule = trieste.acquisition.rule.TURBO( # type: ignore[assignment] From bfe37523b2b293eaa933b0247b8b71e22b21fe51 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Mon, 11 Sep 2023 13:27:14 +0100 Subject: [PATCH 13/16] Incorporate feedback --- docs/notebooks/trust_region.pct.py | 68 +++++++++++------------ docs/refs.bib | 8 +++ setup.py | 2 +- trieste/experimental/plotting/__init__.py | 2 + trieste/experimental/plotting/plotting.py | 28 ++++++++++ 5 files changed, 70 insertions(+), 38 deletions(-) diff --git a/docs/notebooks/trust_region.pct.py b/docs/notebooks/trust_region.pct.py index e9d44a6df7..9b43d2a483 100644 --- a/docs/notebooks/trust_region.pct.py +++ b/docs/notebooks/trust_region.pct.py @@ -64,16 +64,22 @@ def build_model(): # ## Trust region `TREGO` acquisition rule # # First we show how to run Bayesian optimization with the `TREGO` algorithm. This is a trust region -# algorithm that alternates between regular EGO steps and local steps within one trust region. +# algorithm that alternates between regular EGO steps and local steps within one trust region +# (see ). # # ### Create `TREGO` rule and run optimization loop # # We can run the Bayesian optimization loop by defining a `BayesianOptimizer` and calling its # `optimize` method with the trust region rule. Once the optimization loop is complete, the # optimizer will return one new query point for every step in the loop; that's 5 points in total. +# +# `TREGO` is a "meta" rule that applies a base-rule, either inside a trust region or the whole +# space. The default base-rule is `EfficientGlobalOptimization`, but a different base-rule can be +# provided as an argument to `TREGO`. Here we explicitly set it to make usage clear. # %% -acq_rule = trieste.acquisition.rule.TrustRegion() +base_rule = trieste.acquisition.rule.EfficientGlobalOptimization() # type: ignore[var-annotated] +acq_rule = trieste.acquisition.rule.TrustRegion(base_rule) bo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space) num_steps = 5 @@ -110,43 +116,25 @@ def plot_final_result(_dataset: trieste.data.Dataset) -> None: plot_final_result(dataset) # %% [markdown] -# We can also visualize the progress of the optimization by plotting the trust regions at each step. -# The trust regions are shown as translucent boxes, with the current optimum point in each region -# shown in matching color. +# We can also visualize the progress of the optimization by plotting the acquisition space at each +# step. This space is either the full search space or the trust region, depending on the step, and +# is shown as a translucent box; with the current optimum point in a region shown in matching +# color. # -# Note there is only one trust region in this plot, but the rule in the next section will show multiple trust -# regions. +# Note there is only one trust region in this plot, however the rule in the next section will show +# multiple trust regions. # %% import base64 -import io -from typing import List -import imageio import IPython import matplotlib.pyplot as plt -from trieste.experimental.plotting import plot_trust_region_history_2d - - -def fig_to_frame(fig: plt.Figure) -> np.ndarray: - fig.canvas.draw() - size_pix = fig.get_size_inches() * fig.dpi - image = np.frombuffer(fig.canvas.tostring_rgb(), dtype="uint8") - return image.reshape(list(size_pix[::-1].astype(int)) + [3]) - - -def frames_to_gif( - frames: List[np.ndarray], duration=5000 -) -> IPython.display.HTML: - gif_file = io.BytesIO() - imageio.mimsave(gif_file, frames, format="gif", loop=0, duration=duration) - gif = IPython.display.HTML( - ''.format( - base64.b64encode(gif_file.getvalue()).decode() - ) - ) - return gif +from trieste.experimental.plotting import ( + convert_figure_to_frame, + convert_frames_to_gif, + plot_trust_region_history_2d, +) def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: @@ -164,10 +152,16 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: if fig is not None: fig.suptitle(f"step number {step}") - frames.append(fig_to_frame(fig)) + frames.append(convert_figure_to_frame(fig)) plt.close(fig) - IPython.display.display(frames_to_gif(frames)) + gif_file = convert_frames_to_gif(frames) + gif = IPython.display.HTML( + ''.format( + base64.b64encode(gif_file.getvalue()).decode() + ) + ) + IPython.display.display(gif) plot_history(result) @@ -188,9 +182,9 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: # points overall. As the optimization process continues, the bounds of these sub-spaces are # dynamically updated. # -# In addition, this rule requires the specification of a batch aquisition base-rule for performing -# optimization; for our example we use `EfficientGlobalOptimization` coupled with -# `ParallelContinuousThompsonSampling`. +# In addition, this is a "meta" rule that requires the specification of a batch aquisition +# base-rule for performing optimization; for our example we use `EfficientGlobalOptimization` +# coupled with the `ParallelContinuousThompsonSampling` acquisition function. # # Note: the number of sub-spaces/regions must match the number of batch query points. @@ -246,7 +240,7 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: # # ### Create `TurBO` rule and run optimization loop # -# This rule requires the specification of an aquisition base-rule for performing +# As before, this meta-rule requires the specification of an aquisition base-rule for performing # optimization within the trust region; for our example we use `DiscreteThompsonSampling`. # # Note that we switch off global model fitting by setting `fit_model=False`. This is because diff --git a/docs/refs.bib b/docs/refs.bib index bd82c0e10a..dabcde3dfb 100644 --- a/docs/refs.bib +++ b/docs/refs.bib @@ -521,3 +521,11 @@ @inproceedings{wang2013bayesian year={2013} } +@misc{diouane2022trego, + title={TREGO: a Trust-Region Framework for Efficient Global Optimization}, + author={Youssef Diouane and Victor Picheny and Rodolphe Le Riche and Alexandre Scotto Di Perrotolo}, + year={2022}, + eprint={2101.06808}, + archivePrefix={arXiv}, + primaryClass={math.OC} +} diff --git a/setup.py b/setup.py index 0342b335c1..f767df6ae9 100644 --- a/setup.py +++ b/setup.py @@ -49,7 +49,7 @@ "greenlet>=1.1.0", ], extras_require={ - "plotting": ["seaborn", "plotly"], + "plotting": ["seaborn", "plotly", "imageio"], "qhsri": ["pymoo", "cvxpy"], }, ) diff --git a/trieste/experimental/plotting/__init__.py b/trieste/experimental/plotting/__init__.py index 7125d30aff..6c2ac50503 100644 --- a/trieste/experimental/plotting/__init__.py +++ b/trieste/experimental/plotting/__init__.py @@ -23,6 +23,8 @@ plot_objective_and_constraints, ) from .plotting import ( + convert_figure_to_frame, + convert_frames_to_gif, plot_acq_function_2d, plot_bo_points, plot_function_2d, diff --git a/trieste/experimental/plotting/plotting.py b/trieste/experimental/plotting/plotting.py index ed17d820c5..8faf81e865 100644 --- a/trieste/experimental/plotting/plotting.py +++ b/trieste/experimental/plotting/plotting.py @@ -14,8 +14,10 @@ from __future__ import annotations +import io from typing import Callable, List, Optional, Sequence, Union +import imageio import matplotlib.pyplot as plt import numpy as np import tensorflow as tf @@ -617,3 +619,29 @@ def plot_trust_region_history_2d( ) return fig, ax + + +def convert_figure_to_frame(fig: plt.Figure) -> TensorType: + """ + Converts a matplotlib figure to an array of pixels. + + :param fig: a matplotlib figure + :return: an array of pixels - a frame + """ + fig.canvas.draw() + size_pix = fig.get_size_inches() * fig.dpi + image = np.frombuffer(fig.canvas.tostring_rgb(), dtype="uint8") + return image.reshape(list(size_pix[::-1].astype(int)) + [3]) + + +def convert_frames_to_gif(frames: Sequence[TensorType], duration: int = 5000) -> io.BytesIO: + """ + Converts a sequence of frames (arrays of pixels) to a gif. + + :param frames: sequence of frames + :param duration: duration of each frame in milliseconds + :return: gif file + """ + gif_file = io.BytesIO() + imageio.mimsave(gif_file, frames, format="gif", loop=0, duration=duration) + return gif_file From 6956fbe02c7cd3f88daff289fa26ff19f56e0397 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Mon, 11 Sep 2023 13:39:53 +0100 Subject: [PATCH 14/16] Fix rule type ignores --- docs/notebooks/trust_region.pct.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/docs/notebooks/trust_region.pct.py b/docs/notebooks/trust_region.pct.py index 9b43d2a483..28649731a5 100644 --- a/docs/notebooks/trust_region.pct.py +++ b/docs/notebooks/trust_region.pct.py @@ -78,13 +78,14 @@ def build_model(): # provided as an argument to `TREGO`. Here we explicitly set it to make usage clear. # %% -base_rule = trieste.acquisition.rule.EfficientGlobalOptimization() # type: ignore[var-annotated] -acq_rule = trieste.acquisition.rule.TrustRegion(base_rule) +trego_acq_rule = trieste.acquisition.rule.TrustRegion( + rule=trieste.acquisition.rule.EfficientGlobalOptimization() +) bo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space) num_steps = 5 result = bo.optimize( - num_steps, initial_data, build_model(), acq_rule, track_state=True + num_steps, initial_data, build_model(), trego_acq_rule, track_state=True ) dataset = result.try_get_final_dataset() @@ -199,7 +200,7 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: builder=trieste.acquisition.ParallelContinuousThompsonSampling(), num_query_points=num_query_points, ) -acq_rule = trieste.acquisition.rule.BatchTrustRegionBox( # type: ignore[assignment] +batch_acq_rule = trieste.acquisition.rule.BatchTrustRegionBox( init_subspaces, base_rule ) @@ -216,7 +217,7 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: num_steps = 5 result = bo.optimize( - num_steps, initial_data, build_model(), acq_rule, track_state=True + num_steps, initial_data, build_model(), batch_acq_rule, track_state=True ) dataset = result.try_get_final_dataset() @@ -247,7 +248,7 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: # `TurBO` uses a local model and fitting the global model would be redundant and wasteful. # %% -acq_rule = trieste.acquisition.rule.TURBO( # type: ignore[assignment] +turbo_acq_rule = trieste.acquisition.rule.TURBO( search_space, rule=trieste.acquisition.rule.DiscreteThompsonSampling(500, 3) ) bo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space) @@ -257,7 +258,7 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: num_steps, initial_data, build_model(), - acq_rule, + turbo_acq_rule, track_state=True, fit_model=False, ) From 14c33871d4e6d05c00c99777ac3681ffdd5c1297 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Mon, 11 Sep 2023 16:45:43 +0100 Subject: [PATCH 15/16] Add more explanation --- docs/notebooks/trust_region.pct.py | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/docs/notebooks/trust_region.pct.py b/docs/notebooks/trust_region.pct.py index 28649731a5..f59bc48db8 100644 --- a/docs/notebooks/trust_region.pct.py +++ b/docs/notebooks/trust_region.pct.py @@ -174,18 +174,21 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: # # ### Create the batch trust region acquisition rule # -# We achieve Bayesian optimization with trust region by specifying `BatchTrustRegionBox` as the +# We achieve Bayesian optimization with trust regions by specifying `BatchTrustRegionBox` as the # acquisition rule. -# +# # This rule needs an initial number `num_query_points` of sub-spaces (or trust regions) to be # provided and performs optimization in parallel across all these sub-spaces. Each region # contributes one query point, resulting in each acquisition step collecting `num_query_points` # points overall. As the optimization process continues, the bounds of these sub-spaces are -# dynamically updated. +# dynamically updated. In this example, we create 5 `SingleObjectiveTrustRegionBox` regions. This +# class encapsulates the behavior of a trust region in a single sub-space; being responsible for +# maintaining its own state, initializing it, and updating it after each step. # -# In addition, this is a "meta" rule that requires the specification of a batch aquisition -# base-rule for performing optimization; for our example we use `EfficientGlobalOptimization` -# coupled with the `ParallelContinuousThompsonSampling` acquisition function. +# In addition, `BatchTrustRegionBox` is a "meta" rule that requires the specification of a +# batch aquisition base-rule for performing optimization; for our example we use +# `EfficientGlobalOptimization` coupled with the `ParallelContinuousThompsonSampling` acquisition +# function. # # Note: the number of sub-spaces/regions must match the number of batch query points. @@ -242,10 +245,12 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: # ### Create `TurBO` rule and run optimization loop # # As before, this meta-rule requires the specification of an aquisition base-rule for performing -# optimization within the trust region; for our example we use `DiscreteThompsonSampling`. +# optimization within the trust region; for our example we use the `DiscreteThompsonSampling` rule. # -# Note that we switch off global model fitting by setting `fit_model=False`. This is because -# `TurBO` uses a local model and fitting the global model would be redundant and wasteful. +# Note that trieste maintains a global model that is, by default, automatically trained on each +# iteration. However, this global model is unused for `TurBO`; which uses a local model instead. +# As fitting the global model would be redundant and wasteful, we switch its training off by +# setting `fit_model=False` in the `optimize` method. # %% turbo_acq_rule = trieste.acquisition.rule.TURBO( From 3b22611edd1f26a881c5d5bb30f328cb95693bf0 Mon Sep 17 00:00:00 2001 From: Khurram Ghani Date: Mon, 11 Sep 2023 16:50:19 +0100 Subject: [PATCH 16/16] Fix typo --- docs/notebooks/trust_region.pct.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/notebooks/trust_region.pct.py b/docs/notebooks/trust_region.pct.py index f59bc48db8..5222146110 100644 --- a/docs/notebooks/trust_region.pct.py +++ b/docs/notebooks/trust_region.pct.py @@ -176,7 +176,7 @@ def plot_history(result: trieste.bayesian_optimizer.OptimizationResult) -> None: # # We achieve Bayesian optimization with trust regions by specifying `BatchTrustRegionBox` as the # acquisition rule. -# +# # This rule needs an initial number `num_query_points` of sub-spaces (or trust regions) to be # provided and performs optimization in parallel across all these sub-spaces. Each region # contributes one query point, resulting in each acquisition step collecting `num_query_points`