From 4a654f2e2b770e626a4be65cc4fa5eb599c85da4 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Tue, 5 Sep 2023 14:22:16 -0400 Subject: [PATCH] choosing acquisition functions --- bloptools/bayesian/__init__.py | 104 +++++++++++------------------- bloptools/bayesian/acquisition.py | 62 +++++++++++++++--- 2 files changed, 90 insertions(+), 76 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index f8d5359..7ef6abe 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -23,7 +23,7 @@ from .. import utils from . import acquisition, models -from .acquisition import default_acquisition_plan +from .acquisition import ACQ_FUNC_CONFIG, default_acquisition_plan from .digestion import default_digestion_function os.environ["KMP_DUPLICATE_LIB_OK"] = "True" @@ -42,46 +42,6 @@ TASK_TRANSFORMS = {"log": lambda x: np.log(x)} -ACQ_FUNC_CONFIG = { - "quasi-random": { - "identifiers": ["qr", "quasi-random"], - "pretty_name": "Quasi-random", - "description": "Sobol-sampled quasi-random points.", - }, - "expected_mean": { - "identifiers": ["em", "expected_mean"], - "pretty_name": "Expected mean", - "description": "The expected value at each input.", - }, - "expected_improvement": { - "identifiers": ["ei", "expected_improvement"], - "pretty_name": "Expected improvement", - "description": r"The expected value of max(f(x) - \nu, 0), where \nu is the current maximum.", - }, - "noisy_expected_hypervolume_improvement": { - "identifiers": ["nehvi", "noisy_expected_hypervolume_improvement"], - "pretty_name": "Noisy expected hypervolume improvement", - "description": r"It's like a big box. How big is the box?", - }, - "lower_bound_max_value_entropy": { - "identifiers": ["lbmve", "lbmes", "gibbon", "lower_bound_max_value_entropy"], - "pretty_name": "Lower bound max value entropy", - "description": r"Max entropy search, basically", - }, - "probability_of_improvement": { - "identifiers": ["pi", "probability_of_improvement"], - "pretty_name": "Probability of improvement", - "description": "The probability that this input improves on the current maximum.", - }, - "upper_confidence_bound": { - "identifiers": ["ucb", "upper_confidence_bound"], - "default_args": {"beta": 4}, - "pretty_name": "Upper confidence bound", - "description": r"The expected value, plus some multiple of the uncertainty (typically \mu + 2\sigma).", - }, -} - - def _validate_and_prepare_dofs(dofs): for dof in dofs: if not isinstance(dof, Mapping): @@ -184,7 +144,7 @@ def __init__( self.db = db self.verbose = kwargs.get("verbose", False) - self.allow_acquisition_errors = kwargs.get("allow_acquisition_errors", False) + self.allow_acquisition_errors = kwargs.get("allow_acquisition_errors", True) self.initialization = kwargs.get("initialization", None) self.acquisition_plan = kwargs.get("acquisition_plan", default_acquisition_plan) self.digestion = kwargs.get("digestion", default_digestion_function) @@ -295,7 +255,7 @@ def tell(self, new_table=None, append=True, train=True, **kwargs): # # this ensures that we have equal weight between task fitness and feasibility fitness # self.task_scalarization = botorch.acquisition.objective.ScalarizedPosteriorTransform( - # weights=torch.tensor([*torch.ones(self.n_tasks), self.fitness_variance.sum().sqrt()]).double(), + # weights=torch.tensor([*torch.ones(self.num_tasks), self.fitness_variance.sum().sqrt()]).double(), # offset=0, # ) @@ -328,7 +288,7 @@ def tell(self, new_table=None, append=True, train=True, **kwargs): @property def model(self): - return ModelListGP(*[task["model"] for task in self.tasks]) if self.n_tasks > 1 else self.tasks[0]["model"] + return ModelListGP(*[task["model"] for task in self.tasks]) if self.num_tasks > 1 else self.tasks[0]["model"] @property def target_names(self): @@ -438,7 +398,7 @@ def _acq_func_bounds(self): ).T @property - def n_tasks(self): + def num_tasks(self): return len(self.tasks) @property @@ -548,38 +508,53 @@ def acq_func_info(self): print("\n\n".join(entries)) def get_acquisition_function(self, acq_func_identifier="ei", return_metadata=False, **acq_func_kwargs): + """ + Generates an acquisition function from a supplied identifier. + """ + if not self._initialized: raise RuntimeError( f'Can\'t construct acquisition function "{acq_func_identifier}" (the agent is not initialized!)' ) - if acq_func_identifier.lower() in ACQ_FUNC_CONFIG["expected_improvement"]["identifiers"]: + acq_func_name = None + for _acq_func_name in ACQ_FUNC_CONFIG.keys(): + if acq_func_identifier.lower() in ACQ_FUNC_CONFIG[_acq_func_name]["identifiers"]: + acq_func_name = _acq_func_name + + if acq_func_name is None: + raise ValueError(f'Unrecognized acquisition function "{acq_func_identifier}".') + + if ACQ_FUNC_CONFIG[acq_func_name]["multitask_only"] and (self.num_tasks == 1): + raise ValueError(f'Acquisition function "{acq_func_name}" is only for multi-task optimization problems!') + + if acq_func_name == "expected_improvement": acq_func = acquisition.ConstrainedLogExpectedImprovement( constraint=self.constraint, model=self.model, best_f=(self.train_targets * self.task_weights).sum(axis=-1).max(), posterior_transform=ScalarizedPosteriorTransform(weights=self.task_weights, offset=0), ) - acq_func_meta = {"name": "expected_improvement", "args": {}} + acq_func_meta = {"name": acq_func_name, "args": {}} - elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["probability_of_improvement"]["identifiers"]: + elif acq_func_name == "probability_of_improvement": acq_func = acquisition.ConstrainedLogProbabilityOfImprovement( constraint=self.constraint, model=self.model, best_f=(self.train_targets * self.task_weights).sum(axis=-1).max(), posterior_transform=ScalarizedPosteriorTransform(weights=self.task_weights, offset=0), ) - acq_func_meta = {"name": "probability_of_improvement", "args": {}} + acq_func_meta = {"name": acq_func_name, "args": {}} - elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["lower_bound_max_value_entropy"]["identifiers"]: + elif acq_func_name == "lower_bound_max_value_entropy": acq_func = acquisition.qConstrainedLowerBoundMaxValueEntropy( constraint=self.constraint, model=self.model, candidate_set=self.test_inputs(n=1024).squeeze(1), ) - acq_func_meta = {"name": "lower_bound_max_value_entropy", "args": {}} + acq_func_meta = {"name": acq_func_name, "args": {}} - elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["noisy_expected_hypervolume_improvement"]["identifiers"]: + elif acq_func_name == "noisy_expected_hypervolume_improvement": acq_func = acquisition.qConstrainedNoisyExpectedHypervolumeImprovement( constraint=self.constraint, model=self.model, @@ -587,9 +562,9 @@ def get_acquisition_function(self, acq_func_identifier="ei", return_metadata=Fal X_baseline=self.train_inputs, prune_baseline=True, ) - acq_func_meta = {"name": "noisy_expected_hypervolume_improvement", "args": {}} + acq_func_meta = {"name": acq_func_name, "args": {}} - elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["upper_confidence_bound"]["identifiers"]: + elif acq_func_name == "upper_confidence_bound": config = ACQ_FUNC_CONFIG["upper_confidence_bound"] beta = acq_func_kwargs.get("beta", config["default_args"]["beta"]) @@ -599,14 +574,11 @@ def get_acquisition_function(self, acq_func_identifier="ei", return_metadata=Fal beta=beta, posterior_transform=ScalarizedPosteriorTransform(weights=self.task_weights, offset=0), ) - acq_func_meta = {"name": "upper_confidence_bound", "args": {"beta": beta}} + acq_func_meta = {"name": acq_func_name, "args": {"beta": beta}} - elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["expected_mean"]["identifiers"]: + elif acq_func_name == "expected_mean": acq_func = self.get_acquisition_function(acq_func_identifier="ucb", beta=0, return_metadata=False) - acq_func_meta = {"name": "expected_mean"} - - else: - raise ValueError(f'Unrecognized acquisition acq_func_identifier "{acq_func_identifier}".') + acq_func_meta = {"name": acq_func_name, "args": {}} return (acq_func, acq_func_meta) if return_metadata else acq_func @@ -777,9 +749,9 @@ def plot_tasks(self, **kwargs): def _plot_tasks_one_dof(self, size=16, lw=1e0): self.task_fig, self.task_axes = plt.subplots( - self.n_tasks, + self.num_tasks, 1, - figsize=(6, 4 * self.n_tasks), + figsize=(6, 4 * self.num_tasks), sharex=True, constrained_layout=True, ) @@ -822,9 +794,9 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL gridded = self._len_subset_dofs(kind="active", mode="on") == 2 self.task_fig, self.task_axes = plt.subplots( - self.n_tasks, + self.num_tasks, 3, - figsize=(10, 4 * self.n_tasks), + figsize=(10, 4 * self.num_tasks), sharex=True, sharey=True, constrained_layout=True, @@ -1087,9 +1059,9 @@ def plot_history(self, x_key="index", show_all_tasks=False): num_task_plots = 1 if show_all_tasks: - num_task_plots = self.n_tasks + 1 + num_task_plots = self.num_tasks + 1 - self.n_tasks + 1 if self.n_tasks > 1 else 1 + self.num_tasks + 1 if self.num_tasks > 1 else 1 hist_fig, hist_axes = plt.subplots( num_task_plots, 1, figsize=(6, 4 * num_task_plots), sharex=True, constrained_layout=True, dpi=200 diff --git a/bloptools/bayesian/acquisition.py b/bloptools/bayesian/acquisition.py index 253cce6..73ae31e 100644 --- a/bloptools/bayesian/acquisition.py +++ b/bloptools/bayesian/acquisition.py @@ -1,18 +1,64 @@ +import math + import bluesky.plans as bp import numpy as np - -import math import torch - from botorch.acquisition.analytic import LogExpectedImprovement, LogProbabilityOfImprovement, UpperConfidenceBound -from botorch.acquisition.multi_objective.monte_carlo import qNoisyExpectedHypervolumeImprovement from botorch.acquisition.max_value_entropy_search import qLowerBoundMaxValueEntropy +from botorch.acquisition.multi_objective.monte_carlo import qNoisyExpectedHypervolumeImprovement + def default_acquisition_plan(dofs, inputs, dets): uid = yield from bp.list_scan(dets, *[_ for items in zip(dofs, np.atleast_2d(inputs).T) for _ in items]) return uid +ACQ_FUNC_CONFIG = { + "quasi-random": { + "identifiers": ["qr", "quasi-random"], + "pretty_name": "Quasi-random", + "description": "Sobol-sampled quasi-random points.", + "multitask_only": False, + }, + "expected_mean": { + "identifiers": ["em", "expected_mean"], + "pretty_name": "Expected mean", + "multitask_only": False, + "description": "The expected value at each input.", + }, + "expected_improvement": { + "identifiers": ["ei", "expected_improvement"], + "pretty_name": "Expected improvement", + "multitask_only": False, + "description": r"The expected value of max(f(x) - \nu, 0), where \nu is the current maximum.", + }, + "noisy_expected_hypervolume_improvement": { + "identifiers": ["nehvi", "noisy_expected_hypervolume_improvement"], + "pretty_name": "Noisy expected hypervolume improvement", + "multitask_only": True, + "description": r"It's like a big box. How big is the box?", + }, + "lower_bound_max_value_entropy": { + "identifiers": ["lbmve", "lbmes", "gibbon", "lower_bound_max_value_entropy"], + "pretty_name": "Lower bound max value entropy", + "multitask_only": False, + "description": r"Max entropy search, basically", + }, + "probability_of_improvement": { + "identifiers": ["pi", "probability_of_improvement"], + "pretty_name": "Probability of improvement", + "multitask_only": False, + "description": "The probability that this input improves on the current maximum.", + }, + "upper_confidence_bound": { + "identifiers": ["ucb", "upper_confidence_bound"], + "default_args": {"beta": 4}, + "pretty_name": "Upper confidence bound", + "multitask_only": False, + "description": r"The expected value, plus some multiple of the uncertainty (typically \mu + 2\sigma).", + }, +} + class ConstrainedUpperConfidenceBound(UpperConfidenceBound): def __init__(self, constraint, *args, **kwargs): @@ -20,13 +66,11 @@ def __init__(self, constraint, *args, **kwargs): self.constraint = constraint def forward(self, x): - mean, sigma = self._mean_and_sigma(x) - p_eff = 0.5 * (1 + torch.special.erf(self.beta.sqrt()/math.sqrt(2))) * torch.clamp(self.constraint(x), min=1e-6) - - return (mean if self.maximize else -mean) + sigma * np.sqrt(2) * torch.special.erfinv(2*p_eff-1) + p_eff = 0.5 * (1 + torch.special.erf(self.beta.sqrt() / math.sqrt(2))) * torch.clamp(self.constraint(x), min=1e-6) + return (mean if self.maximize else -mean) + sigma * np.sqrt(2) * torch.special.erfinv(2 * p_eff - 1) class ConstrainedLogExpectedImprovement(LogExpectedImprovement): @@ -63,5 +107,3 @@ def __init__(self, constraint, *args, **kwargs): def forward(self, x): return super().forward(x) * self.constraint(x) - -