diff --git a/MANIFEST.in b/MANIFEST.in index 5ec1548..e0ceedc 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -11,6 +11,7 @@ recursive-include docs *.rst conf.py Makefile make.bat include versioneer.py include bloptools/_version.py +include bloptools/bayesian/acquisition/config.yml # If including data files in the package, add them like: # include path/to/data_file diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index b951ce0..dceb5de 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -16,7 +16,6 @@ import pandas as pd import scipy as sp import torch -from botorch.acquisition.objective import ScalarizedPosteriorTransform from botorch.models.deterministic import GenericDeterministicModel from botorch.models.model_list_gp_regression import ModelListGP from matplotlib import pyplot as plt @@ -24,7 +23,7 @@ from .. import utils from . import acquisition, models -from .acquisition import ACQ_FUNC_CONFIG, default_acquisition_plan +from .acquisition import default_acquisition_plan from .digestion import default_digestion_function os.environ["KMP_DUPLICATE_LIB_OK"] = "True" @@ -41,9 +40,11 @@ TASK_TRANSFORMS = {"log": lambda x: np.log(x)} DEFAULT_COLOR_LIST = ["dodgerblue", "tomato", "mediumseagreen", "goldenrod"] -DEFAULT_COLORMAP = "turbo" +DEFAULT_COLORMAP = "viridis" DEFAULT_SCATTER_SIZE = 16 +DEFAULT_MINIMUM_SNR = 2e1 + def _validate_and_prepare_dofs(dofs): for i_dof, dof in enumerate(dofs): @@ -95,6 +96,8 @@ def _validate_and_prepare_tasks(tasks): task["weight"] = 1 if "limits" not in task.keys(): task["limits"] = (-np.inf, np.inf) + if "min_snr" not in task.keys(): + task["min_snr"] = DEFAULT_MINIMUM_SNR task_keys = [task["key"] for task in tasks] unique_task_keys, counts = np.unique(task_keys, return_counts=True) @@ -162,68 +165,18 @@ def __init__( self.trigger_delay = kwargs.get("trigger_delay", 0) - self.acq_func_config = kwargs.get("acq_func_config", ACQ_FUNC_CONFIG) + self.acq_func_config = kwargs.get("acq_func_config", acquisition.config) self.sample_center_on_init = kwargs.get("sample_center_on_init", False) self.table = pd.DataFrame() - self._initialized = False + self.initialized = False self._train_models = True self.a_priori_hypers = None self.plots = {"tasks": {}} - def reset(self): - """ - Reset the agent. - """ - self.table = pd.DataFrame() - self._initialized = False - - def initialize( - self, - acq_func=None, - n_init=4, - data=None, - hypers=None, - ): - """ - An initialization plan for the self. - This must be run before the agent can learn. - It should be passed to a Bluesky RunEngine. - """ - - # experiment-specific stuff - if self.initialization is not None: - yield from self.initialization() - - if hypers is not None: - self.a_priori_hypers = self.load_hypers(hypers) - - if data is not None: - new_table = yield from self.acquire(self._acq_func_bounds.mean(axis=0)) - new_table.loc[:, "acq_func"] = "sample_center_on_init" - self.tell(new_table=new_table, train=False) - - if data is not None: - if type(data) == str: - self.tell(new_table=pd.read_hdf(data, key="table")) - else: - self.tell(new_table=data) - - # now let's get bayesian - elif acq_func in ["qr"]: - yield from self.learn("qr", n_iter=1, n_per_iter=n_init, route=True) - - else: - raise Exception( - """Could not initialize model! Either load a table, or specify an acq_func from: -['qr'].""" - ) - - self._initialized = True - def tell(self, new_table=None, append=True, train=True, **kwargs): """ Inform the agent about new inputs and targets for the model. @@ -235,10 +188,10 @@ def tell(self, new_table=None, append=True, train=True, **kwargs): skew_dims = self.latent_dim_tuples - if self._initialized: + if self.initialized: cached_hypers = self.hypers - inputs = self.inputs.loc[:, self._subset_dof_names(mode="on")].values + inputs = self.model_inputs.values for i, task in enumerate(self.tasks): self.table.loc[:, f"{task['key']}_fitness"] = targets = self._get_task_fitness(i) @@ -250,11 +203,17 @@ def tell(self, new_table=None, append=True, train=True, **kwargs): train_inputs = torch.tensor(inputs[train_index]).double() train_targets = torch.tensor(targets[train_index]).double().unsqueeze(-1) # .unsqueeze(0) + # for constructing the log normal noise prior + target_snr = 2e2 + scale = 2e0 + loc = np.log(1 / target_snr**2) + scale**2 + likelihood = gpytorch.likelihoods.GaussianLikelihood( noise_constraint=gpytorch.constraints.Interval( - torch.tensor(1e-6).square(), - torch.tensor(1e0).square(), + torch.tensor(1e-4).square(), + torch.tensor(1 / task["min_snr"]).square(), ), + noise_prior=gpytorch.priors.torch_priors.LogNormalPrior(loc=loc, scale=scale), ).double() outcome_transform = botorch.models.transforms.outcome.Standardize(m=1) # , batch_shape=torch.Size((1,))) @@ -288,13 +247,153 @@ def tell(self, new_table=None, append=True, train=True, **kwargs): try: self.train_models() except botorch.exceptions.errors.ModelFittingError: - if self._initialized: + if self.initialized: self._set_hypers(cached_hypers) else: raise RuntimeError("Could not fit model on initialization!") self.constraint = GenericDeterministicModel(f=lambda x: self.classifier.probabilities(x)[..., -1].squeeze(-1)) + def ask(self, acq_func_identifier="qei", n=1, route=True, sequential=True, **acq_func_kwargs): + """ + Ask the agent for the best point to sample, given an acquisition function. + + acq_func_identifier: which acquisition function to use + n: how many points to get + """ + + acq_func_name = acquisition.parse_acq_func(acq_func_identifier) + acq_func_type = acquisition.config[acq_func_name]["type"] + + start_time = ttime.monotonic() + + if acq_func_type in ["analytic", "monte_carlo"]: + if not self.initialized: + raise RuntimeError( + f'Can\'t construct non-trivial acquisition function "{acq_func_identifier}"' + f" (the agent is not initialized!)" + ) + + if acq_func_type == "analytic" and n > 1: + raise ValueError("Can't generate multiple design points for analytic acquisition functions.") + + acq_func, acq_func_meta = acquisition.get_acquisition_function(self, acq_func_identifier=acq_func_identifier) + + NUM_RESTARTS = 8 + RAW_SAMPLES = 1024 + + candidates, _ = botorch.optim.optimize_acqf( + acq_function=acq_func, + bounds=self.acq_func_bounds, + q=n, + sequential=sequential, + num_restarts=NUM_RESTARTS, + raw_samples=RAW_SAMPLES, # used for intialization heuristic + ) + + x = candidates.numpy().astype(float) + + active_X = x[..., [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")]] + passive_X = x[..., [dof["kind"] != "active" for dof in self._subset_dofs(mode="on")]] + acq_func_meta["passive_values"] = passive_X + + else: + if acq_func_name == "quasi-random": + active_X = self._subset_inputs_sampler(n=n, kind="active", mode="on").squeeze(1).numpy() + acq_func_meta = {"name": "quasi-random", "args": {}} + + else: + raise ValueError() + + acq_func_meta["duration"] = duration = ttime.monotonic() - start_time + + if self.verbose: + print(f"found points {active_X} in {duration:.01f} seconds") + + if route and n > 1: + active_X = active_X[utils.route(self._read_subset_devices(kind="active", mode="on"), active_X)] + + return active_X, acq_func_meta + + def acquire(self, active_inputs): + """ + Acquire and digest according to the self's acquisition and digestion plans. + + This should yield a table of sampled tasks with the same length as the sampled inputs. + """ + try: + active_devices = self._subset_devices(kind="active", mode="on") + passive_devices = [*self._subset_devices(kind="passive"), *self._subset_devices(kind="active", mode="off")] + + uid = yield from self.acquisition_plan( + active_devices, + active_inputs.astype(float), + [*self.dets, *passive_devices], + delay=self.trigger_delay, + ) + + products = self.digestion(self.db, uid) + + # compute the fitness for each task + # for index, entry in products.iterrows(): + # for task in self.tasks: + # products.loc[index, task["key"]] = getattr(entry, task["key"]) + + except Exception as error: + if not self.allow_acquisition_errors: + raise error + logging.warning(f"Error in acquisition/digestion: {repr(error)}") + products = pd.DataFrame(active_inputs, columns=self._subset_dof_names(kind="active", mode="on")) + for task in self.tasks: + products.loc[:, task["key"]] = np.nan + + if not len(active_inputs) == len(products): + raise ValueError("The table returned by the digestion must be the same length as the sampled inputs!") + + return products + + def learn( + self, + acq_func, + n=1, + iterations=1, + upsample=1, + train=True, + data=None, + **kwargs, + ): + """ + This iterates the learning algorithm, looping over ask -> acquire -> tell. + It should be passed to a Bluesky RunEngine. + """ + + if data is not None: + if type(data) == str: + self.tell(new_table=pd.read_hdf(data, key="table")) + else: + self.tell(new_table=data) + + if self.sample_center_on_init and not self.initialized: + new_table = yield from self.acquire(self.acq_func_bounds.mean(axis=0)) + new_table.loc[:, "acq_func"] = "sample_center_on_init" + self.tell(new_table=new_table, train=False) + + for i in range(iterations): + x, acq_func_meta = self.ask(n=n, acq_func_identifier=acq_func, **kwargs) + + new_table = yield from self.acquire(x) + new_table.loc[:, "acq_func"] = acq_func_meta["name"] + self.tell(new_table=new_table, train=train) + + self.initialized = True + + def reset(self): + """ + Reset the agent. + """ + self.table = pd.DataFrame() + self.initialized = False + @property def model(self): """ @@ -365,7 +464,7 @@ def test_inputs_grid(self): ).swapaxes(0, -1) @property - def _acq_func_bounds(self): + def acq_func_bounds(self): return torch.tensor( [ dof["limits"] if dof["kind"] == "active" else tuple(2 * [dof["device"].read()[dof["device"].name]["value"]]) @@ -454,7 +553,7 @@ def latent_dim_tuples(self): return [tuple(np.where(uinv == i)[0]) for i in range(len(u))] def test_inputs(self, n=MAX_TEST_INPUTS): - return utils.sobol_sampler(self._acq_func_bounds, n=n) + return utils.sobol_sampler(self.acq_func_bounds, n=n) def _subset_input_transform(self, kind=None, mode=None, tags=[]): limits = self._subset_dof_limits(kind, mode, tags) @@ -536,7 +635,7 @@ def train_models(self, **kwargs): gpytorch.mlls.ExactMarginalLogLikelihood(self.classifier.likelihood, self.classifier), **kwargs ) if self.verbose: - print(f"trained models in {ttime.monotonic() - t0:.02f} seconds") + print(f"trained models in {ttime.monotonic() - t0:.01f} seconds") @property def acq_func_info(self): @@ -549,232 +648,27 @@ 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!)' - ) - - 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.best_scalarized_fitness, - posterior_transform=ScalarizedPosteriorTransform(weights=self.task_weights, offset=0), - ) - acq_func_meta = {"name": acq_func_name, "args": {}} - - elif acq_func_name == "probability_of_improvement": - acq_func = acquisition.ConstrainedLogProbabilityOfImprovement( - constraint=self.constraint, - model=self.model, - best_f=self.best_scalarized_fitness, - posterior_transform=ScalarizedPosteriorTransform(weights=self.task_weights, offset=0), - ) - acq_func_meta = {"name": acq_func_name, "args": {}} - - 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": acq_func_name, "args": {}} - - elif acq_func_name == "noisy_expected_hypervolume_improvement": - acq_func = acquisition.qConstrainedNoisyExpectedHypervolumeImprovement( - constraint=self.constraint, - model=self.model, - ref_point=self.train_targets.min(dim=0).values, - X_baseline=self.train_inputs, - prune_baseline=True, - ) - acq_func_meta = {"name": acq_func_name, "args": {}} - - elif acq_func_name == "upper_confidence_bound": - config = ACQ_FUNC_CONFIG["upper_confidence_bound"] - beta = acq_func_kwargs.get("beta", config["default_args"]["beta"]) - - acq_func = acquisition.ConstrainedUpperConfidenceBound( - constraint=self.constraint, - model=self.model, - beta=beta, - posterior_transform=ScalarizedPosteriorTransform(weights=self.task_weights, offset=0), - ) - acq_func_meta = {"name": acq_func_name, "args": {"beta": beta}} - - 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": acq_func_name, "args": {}} - - return (acq_func, acq_func_meta) if return_metadata else acq_func - - def ask(self, acq_func_identifier="ei", n=1, route=True, return_metadata=False): - if acq_func_identifier.lower() == "qr": - active_X = self._subset_inputs_sampler(n=n, kind="active", mode="on").squeeze(1).numpy() - acq_func_meta = {"name": "quasi-random", "args": {}} - - elif n == 1: - active_X, acq_func_meta = self.ask_single(acq_func_identifier, return_metadata=True) - - elif n > 1: - active_x_list = [] - for i in range(n): - active_x, acq_func_meta = self.ask_single(acq_func_identifier, return_metadata=True) - active_x_list.append(active_x) - - if i < (n - 1): - x = np.c_[active_x, acq_func_meta["passive_values"]] - task_samples = [task["model"].posterior(torch.tensor(x)).sample().item() for task in self.tasks] - fantasy_table = pd.DataFrame( - np.c_[active_x, acq_func_meta["passive_values"], np.atleast_2d(task_samples)], - columns=[ - *self._subset_dof_names(kind="active", mode="on"), - *self._subset_dof_names(kind="passive", mode="on"), - *self.task_keys, - ], - ) - self.tell(fantasy_table, train=False) - - active_X = np.concatenate(active_x_list, axis=0) - self.forget(self.table.index[-(n - 1) :]) - - if route: - active_X = active_X[utils.route(self._read_subset_devices(kind="active", mode="on"), active_X)] - - return (active_X, acq_func_meta) if return_metadata else active_X - - def ask_single( - self, - acq_func_identifier="ei", - return_metadata=False, - ): - """ - The next $n$ points to sample, recommended by the self. Returns - """ - - t0 = ttime.monotonic() - - acq_func, acq_func_meta = self.get_acquisition_function( - acq_func_identifier=acq_func_identifier, return_metadata=True - ) - - BATCH_SIZE = 1 - NUM_RESTARTS = 4 - RAW_SAMPLES = 256 - - candidates, _ = botorch.optim.optimize_acqf( - acq_function=acq_func, - bounds=self._acq_func_bounds, - q=BATCH_SIZE, - num_restarts=NUM_RESTARTS, - raw_samples=RAW_SAMPLES, # used for intialization heuristic - ) - - x = candidates.numpy().astype(float) - - active_x = x[..., [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")]] - passive_x = x[..., [dof["kind"] != "active" for dof in self._subset_dofs(mode="on")]] - - acq_func_meta["passive_values"] = passive_x - - if self.verbose: - print(f"found point {x} in {ttime.monotonic() - t0:.02f} seconds") - - return (active_x, acq_func_meta) if return_metadata else active_x - - def acquire(self, active_inputs): - """ - Acquire and digest according to the self's acquisition and digestion plans. - - This should yield a table of sampled tasks with the same length as the sampled inputs. - """ - try: - active_devices = self._subset_devices(kind="active", mode="on") - passive_devices = [*self._subset_devices(kind="passive"), *self._subset_devices(kind="active", mode="off")] - - uid = yield from self.acquisition_plan( - active_devices, - active_inputs.astype(float), - [*self.dets, *passive_devices], - delay=self.trigger_delay, - ) - - products = self.digestion(self.db, uid) - - # compute the fitness for each task - # for index, entry in products.iterrows(): - # for task in self.tasks: - # products.loc[index, task["key"]] = getattr(entry, task["key"]) - - except Exception as error: - if not self.allow_acquisition_errors: - raise error - logging.warning(f"Error in acquisition/digestion: {repr(error)}") - products = pd.DataFrame(active_inputs, columns=self._subset_dof_names(kind="active", mode="on")) - for task in self.tasks: - products.loc[:, task["key"]] = np.nan - - if not len(active_inputs) == len(products): - raise ValueError("The table returned by the digestion must be the same length as the sampled inputs!") - - return products - - def learn( - self, - acq_func_identifier, - n_iter=1, - n_per_iter=1, - reuse_hypers=True, - train=True, - upsample=1, - verbose=True, - plots=[], - **kwargs, - ): - """ - This iterates the learning algorithm, looping over ask -> acquire -> tell. - It should be passed to a Bluesky RunEngine. - """ - - for i in range(n_iter): - x, acq_func_meta = self.ask( - n=n_per_iter, acq_func_identifier=acq_func_identifier, return_metadata=True, **kwargs - ) - - new_table = yield from self.acquire(x) - new_table.loc[:, "acq_func"] = acq_func_meta["name"] - self.tell(new_table=new_table, train=train) - @property def inputs(self): + return self.table.loc[:, self._subset_dof_names()].astype(float) + + @property + def model_inputs(self): return self.table.loc[:, self._subset_dof_names(mode="on")].astype(float) + @property + def active_inputs(self): + return self.table.loc[:, self._subset_dof_names(kind="active", mode="on")].astype(float) + @property def best_inputs(self): - return self.inputs.values[np.nanargmax(self.scalarized_fitness)] + return self.active_inputs.values[np.nanargmax(self.scalarized_fitness)] - def go_to(self, inputs): + def go_to(self, active_inputs): args = [] - for dof, value in zip(self._subset_dofs(mode="on"), np.atleast_1d(inputs).T): - if dof["kind"] == "active": - args.append(dof["device"]) - args.append(value) + for dof, value in zip(self._subset_dofs(kind="active"), np.atleast_1d(active_inputs).T): + args.append(dof["device"]) + args.append(value) yield from bps.mv(*args) def go_to_best(self): @@ -797,24 +691,27 @@ def _plot_tasks_one_dof(self, size=16, lw=1e0): self.task_axes = np.atleast_1d(self.task_axes) - for itask, task in enumerate(self.tasks): + for task_index, task in enumerate(self.tasks): task_plots = self.plots["tasks"][task["name"]] = {} - color = DEFAULT_COLOR_LIST[itask] + task_fitness = self._get_task_fitness(task_index=task_index) + + color = DEFAULT_COLOR_LIST[task_index] - self.task_axes[itask].set_ylabel(task["key"]) + self.task_axes[task_index].set_ylabel(task["key"]) x = self.test_inputs_grid task_posterior = task["model"].posterior(x) task_mean = task_posterior.mean.detach().numpy() task_sigma = task_posterior.variance.sqrt().detach().numpy() - task_plots["sampled"] = self.task_axes[itask].scatter([], [], s=size, color=color) + sampled_inputs = self.active_inputs.values + task_plots["sampled"] = self.task_axes[task_index].scatter(sampled_inputs, task_fitness, s=size, color=color) on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] for z in [0, 1, 2]: - self.task_axes[itask].fill_between( + self.task_axes[task_index].fill_between( x[..., on_dofs_are_active_mask].squeeze(), (task_mean - z * task_sigma).squeeze(), (task_mean + z * task_sigma).squeeze(), @@ -823,7 +720,7 @@ def _plot_tasks_one_dof(self, size=16, lw=1e0): alpha=0.5**z, ) - self.task_axes[itask].set_xlim(self._subset_dofs(kind="active", mode="on")[0]["limits"]) + self.task_axes[task_index].set_xlim(self._subset_dofs(kind="active", mode="on")[0]["limits"]) def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=16): if gridded is None: @@ -841,23 +738,24 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL self.task_axes = np.atleast_2d(self.task_axes) # self.task_fig.suptitle(f"(x,y)=({self.dofs[axes[0]].name},{self.dofs[axes[1]].name})") - for itask, task in enumerate(self.tasks): - sampled_fitness = np.where(self.all_tasks_valid, self.table.loc[:, f'{task["key"]}_fitness'].values, np.nan) - task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) + for task_index, task in enumerate(self.tasks): + task_fitness = self._get_task_fitness(task_index=task_index) + + task_vmin, task_vmax = np.nanpercentile(task_fitness, q=[1, 99]) task_norm = mpl.colors.Normalize(task_vmin, task_vmax) # if task["transform"] == "log": # task_norm = mpl.colors.LogNorm(task_vmin, task_vmax) # else: - self.task_axes[itask, 0].set_ylabel(f'{task["key"]}_fitness') + self.task_axes[task_index, 0].set_ylabel(f'{task["key"]}_fitness') - self.task_axes[itask, 0].set_title("samples") - self.task_axes[itask, 1].set_title("posterior mean") - self.task_axes[itask, 2].set_title("posterior std. dev.") + self.task_axes[task_index, 0].set_title("samples") + self.task_axes[task_index, 1].set_title("posterior mean") + self.task_axes[task_index, 2].set_title("posterior std. dev.") - data_ax = self.task_axes[itask, 0].scatter( - *self.inputs.values.T[axes], s=size, c=sampled_fitness, norm=task_norm, cmap=cmap + data_ax = self.task_axes[task_index, 0].scatter( + *self.active_inputs.values.T[axes], s=size, c=task_fitness, norm=task_norm, cmap=cmap ) x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) @@ -869,7 +767,7 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL if gridded: if not x.ndim == 3: raise ValueError() - self.task_axes[itask, 1].pcolormesh( + self.task_axes[task_index, 1].pcolormesh( x[..., 0].detach().numpy(), x[..., 1].detach().numpy(), task_mean[..., 0].detach().numpy(), @@ -877,7 +775,7 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL cmap=cmap, norm=task_norm, ) - sigma_ax = self.task_axes[itask, 2].pcolormesh( + sigma_ax = self.task_axes[task_index, 2].pcolormesh( x[..., 0].detach().numpy(), x[..., 1].detach().numpy(), task_sigma[..., 0].detach().numpy(), @@ -886,7 +784,7 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL ) else: - self.task_axes[itask, 1].scatter( + self.task_axes[task_index, 1].scatter( x.detach().numpy()[..., axes[0]], x.detach().numpy()[..., axes[1]], c=task_mean[..., 0].detach().numpy(), @@ -894,7 +792,7 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL norm=task_norm, cmap=cmap, ) - sigma_ax = self.task_axes[itask, 2].scatter( + sigma_ax = self.task_axes[task_index, 2].scatter( x.detach().numpy()[..., axes[0]], x.detach().numpy()[..., axes[1]], c=task_sigma[..., 0].detach().numpy(), @@ -902,8 +800,8 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL cmap=cmap, ) - self.task_fig.colorbar(data_ax, ax=self.task_axes[itask, :2], location="bottom", aspect=32, shrink=0.8) - self.task_fig.colorbar(sigma_ax, ax=self.task_axes[itask, 2], location="bottom", aspect=32, shrink=0.8) + self.task_fig.colorbar(data_ax, ax=self.task_axes[task_index, :2], location="bottom", aspect=32, shrink=0.8) + self.task_fig.colorbar(sigma_ax, ax=self.task_axes[task_index, 2], location="bottom", aspect=32, shrink=0.8) for ax in self.task_axes.ravel(): ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) @@ -930,7 +828,7 @@ def _plot_acq_one_dof(self, acq_funcs, lw=1e0, **kwargs): for iacq_func, acq_func_identifier in enumerate(acq_funcs): color = DEFAULT_COLOR_LIST[iacq_func] - acq_func, acq_func_meta = self.get_acquisition_function(acq_func_identifier, return_metadata=True) + acq_func, acq_func_meta = acquisition.get_acquisition_function(self, acq_func_identifier) x = self.test_inputs_grid *input_shape, input_dim = x.shape @@ -970,7 +868,7 @@ def _plot_acq_many_dofs( *input_shape, input_dim = x.shape for iacq_func, acq_func_identifier in enumerate(acq_funcs): - acq_func, acq_func_meta = self.get_acquisition_function(acq_func_identifier, return_metadata=True) + acq_func, acq_func_meta = acquisition.get_acquisition_function(self, acq_func_identifier) obj = acq_func.forward(x.reshape(-1, 1, input_dim)).reshape(input_shape) if acq_func_identifier in ["ei", "pi"]: @@ -1016,7 +914,7 @@ def _plot_valid_one_dof(self, size=16, lw=1e0): *input_shape, input_dim = x.shape constraint = self.classifier.probabilities(x.reshape(-1, 1, input_dim))[..., -1].reshape(input_shape) - self.valid_ax.scatter(self.inputs.values, self.all_tasks_valid, s=size) + self.valid_ax.scatter(self.active_inputs.values, self.all_tasks_valid, s=size) on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] @@ -1033,7 +931,7 @@ def _plot_valid_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL gridded = self._len_subset_dofs(kind="active", mode="on") == 2 data_ax = self.valid_axes[0].scatter( - *self.inputs.values.T[:2], + *self.active_inputs.values.T[:2], c=self.all_tasks_valid, s=size, vmin=0, @@ -1111,11 +1009,11 @@ def plot_history(self, x_key="index", show_all_tasks=False): sample_colors = np.array(DEFAULT_COLOR_LIST)[acq_func_inverse] if show_all_tasks: - for itask, task in enumerate(self.tasks): + for task_index, task in enumerate(self.tasks): y = self.table.loc[:, f'{task["key"]}_fitness'].values - hist_axes[itask].scatter(x, y, c=sample_colors) - hist_axes[itask].plot(x, y, lw=5e-1, c="k") - hist_axes[itask].set_ylabel(task["key"]) + hist_axes[task_index].scatter(x, y, c=sample_colors) + hist_axes[task_index].plot(x, y, lw=5e-1, c="k") + hist_axes[task_index].set_ylabel(task["key"]) y = self.scalarized_fitness diff --git a/bloptools/bayesian/acquisition.py b/bloptools/bayesian/acquisition.py deleted file mode 100644 index e6d3835..0000000 --- a/bloptools/bayesian/acquisition.py +++ /dev/null @@ -1,151 +0,0 @@ -import math - -import bluesky.plan_stubs as bps -import bluesky.plans as bp -import numpy as np -import torch -from botorch.acquisition.analytic import LogExpectedImprovement, LogProbabilityOfImprovement, UpperConfidenceBound -from botorch.acquisition.max_value_entropy_search import qLowerBoundMaxValueEntropy -from botorch.acquisition.multi_objective.monte_carlo import qNoisyExpectedHypervolumeImprovement - - -def list_scan_with_delay(*args, delay=0, **kwargs): - "Accepts all the normal 'scan' parameters, plus an optional delay." - - def one_nd_step_with_delay(detectors, step, pos_cache): - "This is a copy of bluesky.plan_stubs.one_nd_step with a sleep added." - motors = step.keys() - yield from bps.move_per_step(step, pos_cache) - yield from bps.sleep(delay) - yield from bps.trigger_and_read(list(detectors) + list(motors)) - - kwargs.setdefault("per_step", one_nd_step_with_delay) - uid = yield from bp.list_scan(*args, **kwargs) - return uid - - -def default_acquisition_plan(dofs, inputs, dets, **kwargs): - delay = kwargs.get("delay", 0) - args = [] - for dof, points in zip(dofs, np.atleast_2d(inputs).T): - args.append(dof) - args.append(list(points)) - - uid = yield from list_scan_with_delay(dets, *args, delay=delay) - return uid - - -# def sleepy_acquisition_plan(dofs, inputs, dets): - -# args = [] -# for dof, points in zip(dofs, np.atleast_2d(inputs).T): -# args.append(dof) -# args.append(list(points)) - -# for point in inputs: -# args = [] -# for dof, value in zip(dofs, point): -# args.append(dof) -# args.append(value) - -# yield from bps.mv(*args) -# yield from bps.count([*dets, *dofs]) -# yield from bps.sleep(1) - -# 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): - super().__init__(*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) - - -class ConstrainedLogExpectedImprovement(LogExpectedImprovement): - def __init__(self, constraint, *args, **kwargs): - super().__init__(*args, **kwargs) - self.constraint = constraint - - def forward(self, x): - return super().forward(x) + self.constraint(x).log() - - -class ConstrainedLogProbabilityOfImprovement(LogProbabilityOfImprovement): - def __init__(self, constraint, *args, **kwargs): - super().__init__(*args, **kwargs) - self.constraint = constraint - - def forward(self, x): - return super().forward(x) + self.constraint(x).log() - - -class qConstrainedNoisyExpectedHypervolumeImprovement(qNoisyExpectedHypervolumeImprovement): - def __init__(self, constraint, *args, **kwargs): - super().__init__(*args, **kwargs) - self.constraint = constraint - - def forward(self, x): - return super().forward(x) * self.constraint(x) - - -class qConstrainedLowerBoundMaxValueEntropy(qLowerBoundMaxValueEntropy): - def __init__(self, constraint, *args, **kwargs): - super().__init__(*args, **kwargs) - self.constraint = constraint - - def forward(self, x): - return super().forward(x) * self.constraint(x) diff --git a/bloptools/bayesian/acquisition/__init__.py b/bloptools/bayesian/acquisition/__init__.py new file mode 100644 index 0000000..7931974 --- /dev/null +++ b/bloptools/bayesian/acquisition/__init__.py @@ -0,0 +1,129 @@ +import os + +import yaml +from botorch.acquisition.objective import ScalarizedPosteriorTransform + +from . import analytic, monte_carlo +from .analytic import * # noqa F401 +from .monte_carlo import * # noqa F401 + +here, this_filename = os.path.split(__file__) + +with open(f"{here}/config.yml", "r") as f: + config = yaml.safe_load(f) + + +# supplying the full name is also a valid identifier +for acq_func_name in config.keys(): + config[acq_func_name]["identifiers"].append(acq_func_name) + + +def parse_acq_func(acq_func_identifier): + acq_func_name = None + for _acq_func_name in config.keys(): + if acq_func_identifier.lower() in config[_acq_func_name]["identifiers"]: + acq_func_name = _acq_func_name + + if acq_func_name is None: + raise ValueError(f'Unrecognized acquisition function identifier "{acq_func_identifier}".') + + return acq_func_name + + +def get_acquisition_function(agent, acq_func_identifier="qei", **acq_func_kwargs): + """ + Generates an acquisition function from a supplied identifier. + """ + + acq_func_name = parse_acq_func(acq_func_identifier) + acq_func_config = agent.acq_func_config["upper_confidence_bound"] + + if agent.acq_func_config[acq_func_name]["multitask_only"] and (agent.num_tasks == 1): + raise ValueError(f'Acquisition function "{acq_func_name}" is only for multi-task optimization problems!') + + # there is probably a better way to structure this + if acq_func_name == "expected_improvement": + acq_func = analytic.ConstrainedLogExpectedImprovement( + constraint=agent.constraint, + model=agent.model, + best_f=agent.best_scalarized_fitness, + posterior_transform=ScalarizedPosteriorTransform(weights=agent.task_weights, offset=0), + ) + acq_func_meta = {"name": acq_func_name, "args": {}} + + elif acq_func_name == "monte_carlo_expected_improvement": + acq_func = monte_carlo.qConstrainedExpectedImprovement( + constraint=agent.constraint, + model=agent.model, + best_f=agent.best_scalarized_fitness, + posterior_transform=ScalarizedPosteriorTransform(weights=agent.task_weights, offset=0), + ) + acq_func_meta = {"name": acq_func_name, "args": {}} + + elif acq_func_name == "probability_of_improvement": + acq_func = analytic.ConstrainedLogProbabilityOfImprovement( + constraint=agent.constraint, + model=agent.model, + best_f=agent.best_scalarized_fitness, + posterior_transform=ScalarizedPosteriorTransform(weights=agent.task_weights, offset=0), + ) + acq_func_meta = {"name": acq_func_name, "args": {}} + + elif acq_func_name == "monte_carlo_probability_of_improvement": + acq_func = monte_carlo.qConstrainedProbabilityOfImprovement( + constraint=agent.constraint, + model=agent.model, + best_f=agent.best_scalarized_fitness, + posterior_transform=ScalarizedPosteriorTransform(weights=agent.task_weights, offset=0), + ) + acq_func_meta = {"name": acq_func_name, "args": {}} + + elif acq_func_name == "lower_bound_max_value_entropy": + acq_func = monte_carlo.qConstrainedLowerBoundMaxValueEntropy( + constraint=agent.constraint, + model=agent.model, + candidate_set=agent.test_inputs(n=1024).squeeze(1), + ) + acq_func_meta = {"name": acq_func_name, "args": {}} + + elif acq_func_name == "noisy_expected_hypervolume_improvement": + acq_func = monte_carlo.qConstrainedNoisyExpectedHypervolumeImprovement( + constraint=agent.constraint, + model=agent.model, + ref_point=agent.train_targets.min(dim=0).values, + X_baseline=agent.train_inputs, + prune_baseline=True, + ) + acq_func_meta = {"name": acq_func_name, "args": {}} + + elif acq_func_name == "upper_confidence_bound": + beta = acq_func_kwargs.get("beta", acq_func_config["default_args"]["beta"]) + + acq_func = analytic.ConstrainedUpperConfidenceBound( + constraint=agent.constraint, + model=agent.model, + beta=beta, + posterior_transform=ScalarizedPosteriorTransform(weights=agent.task_weights, offset=0), + ) + acq_func_meta = {"name": acq_func_name, "args": {"beta": beta}} + + elif acq_func_name == "monte_carlo_upper_confidence_bound": + beta = acq_func_kwargs.get("beta", acq_func_config["default_args"]["beta"]) + + acq_func = monte_carlo.qConstrainedUpperConfidenceBound( + constraint=agent.constraint, + model=agent.model, + beta=beta, + posterior_transform=ScalarizedPosteriorTransform(weights=agent.task_weights, offset=0), + ) + acq_func_meta = {"name": acq_func_name, "args": {"beta": beta}} + + elif acq_func_name == "expected_mean": + acq_func = get_acquisition_function(agent, acq_func_identifier="ucb", beta=0, return_metadata=False) + acq_func_meta = {"name": acq_func_name, "args": {}} + + elif acq_func_name == "monte_carlo_expected_mean": + acq_func = get_acquisition_function(agent, acq_func_identifier="qucb", beta=0, return_metadata=False) + acq_func_meta = {"name": acq_func_name, "args": {}} + + return acq_func, acq_func_meta diff --git a/bloptools/bayesian/acquisition/analytic.py b/bloptools/bayesian/acquisition/analytic.py new file mode 100644 index 0000000..6c5cc0e --- /dev/null +++ b/bloptools/bayesian/acquisition/analytic.py @@ -0,0 +1,84 @@ +import math + +import bluesky.plan_stubs as bps +import bluesky.plans as bp +import numpy as np +import torch +from botorch.acquisition.analytic import LogExpectedImprovement, LogProbabilityOfImprovement, UpperConfidenceBound + + +def list_scan_with_delay(*args, delay=0, **kwargs): + "Accepts all the normal 'scan' parameters, plus an optional delay." + + def one_nd_step_with_delay(detectors, step, pos_cache): + "This is a copy of bluesky.plan_stubs.one_nd_step with a sleep added." + motors = step.keys() + yield from bps.move_per_step(step, pos_cache) + yield from bps.sleep(delay) + yield from bps.trigger_and_read(list(detectors) + list(motors)) + + kwargs.setdefault("per_step", one_nd_step_with_delay) + uid = yield from bp.list_scan(*args, **kwargs) + return uid + + +def default_acquisition_plan(dofs, inputs, dets, **kwargs): + delay = kwargs.get("delay", 0) + args = [] + for dof, points in zip(dofs, np.atleast_2d(inputs).T): + args.append(dof) + args.append(list(points)) + + uid = yield from list_scan_with_delay(dets, *args, delay=delay) + return uid + + +# def sleepy_acquisition_plan(dofs, inputs, dets): + +# args = [] +# for dof, points in zip(dofs, np.atleast_2d(inputs).T): +# args.append(dof) +# args.append(list(points)) + +# for point in inputs: +# args = [] +# for dof, value in zip(dofs, point): +# args.append(dof) +# args.append(value) + +# yield from bps.mv(*args) +# yield from bps.count([*dets, *dofs]) +# yield from bps.sleep(1) + +# return uid + + +class ConstrainedUpperConfidenceBound(UpperConfidenceBound): + def __init__(self, constraint, *args, **kwargs): + super().__init__(*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 + sigma * np.sqrt(2) * torch.special.erfinv(2 * p_eff - 1) + + +class ConstrainedLogExpectedImprovement(LogExpectedImprovement): + def __init__(self, constraint, *args, **kwargs): + super().__init__(*args, **kwargs) + self.constraint = constraint + + def forward(self, x): + return super().forward(x) + self.constraint(x).log() + + +class ConstrainedLogProbabilityOfImprovement(LogProbabilityOfImprovement): + def __init__(self, constraint, *args, **kwargs): + super().__init__(*args, **kwargs) + self.constraint = constraint + + def forward(self, x): + return super().forward(x) + self.constraint(x).log() diff --git a/bloptools/bayesian/acquisition/config.yml b/bloptools/bayesian/acquisition/config.yml new file mode 100644 index 0000000..5eb38ac --- /dev/null +++ b/bloptools/bayesian/acquisition/config.yml @@ -0,0 +1,93 @@ +expected_improvement: + pretty_name: Expected improvement + description: The expected value of max(f(x) - \nu, 0), where \nu is the current maximum. + identifiers: + - ei + multitask_only: false + type: analytic + +monte_carlo_expected_improvement: + description: The expected value of max(f(x) - \nu, 0), where \nu is the current maximum. + identifiers: + - qei + multitask_only: false + pretty_name: Monte Carlo Expected improvement + type: monte_carlo + +expected_mean: + description: The expected value at each input. + identifiers: + - em + multitask_only: false + pretty_name: Expected mean + type: analytic + +monte_carlo_expected_mean: + description: The expected value at each input. + identifiers: + - qem + multitask_only: false + pretty_name: Monte Carlo expected mean + type: monte_carlo + +lower_bound_max_value_entropy: + description: Max entropy search, basically + identifiers: + - lbmve + - lbmes + - gibbon + multitask_only: false + pretty_name: Lower bound max value entropy + type: monte_carlo + +noisy_expected_hypervolume_improvement: + description: It's like a big box. How big is the box? + identifiers: + - nehvi + multitask_only: true + pretty_name: Noisy expected hypervolume improvement + type: analytic + +probability_of_improvement: + description: The probability that this input improves on the current maximum. + identifiers: + - pi + multitask_only: false + pretty_name: Probability of improvement + type: analytic + +monte_carlo_probability_of_improvement: + description: The probability that this input improves on the current maximum. + identifiers: + - qpi + multitask_only: false + pretty_name: Monte Carlo probability of improvement + type: monte_carlo + +quasi-random: + description: Sobol-sampled quasi-random points. + identifiers: + - qr + multitask_only: false + pretty_name: Quasi-random + type: none + +upper_confidence_bound: + default_args: + beta: 4 + description: The expected value, plus some multiple of the uncertainty (typically \mu + 2\sigma). + identifiers: + - ucb + multitask_only: false + pretty_name: Upper confidence bound + type: analytic + +monte_carlo_upper_confidence_bound: + default_args: + beta: 4 + description: The expected value, plus some multiple of the uncertainty (typically \mu + 2\sigma). + identifiers: + - qucb + multitask_only: false + pretty_name: Monte Carlo upper confidence bound + type: monte_carlo diff --git a/bloptools/bayesian/acquisition/monte_carlo.py b/bloptools/bayesian/acquisition/monte_carlo.py new file mode 100644 index 0000000..ce9c681 --- /dev/null +++ b/bloptools/bayesian/acquisition/monte_carlo.py @@ -0,0 +1,58 @@ +import math + +import numpy as np +import torch +from botorch.acquisition.max_value_entropy_search import qLowerBoundMaxValueEntropy +from botorch.acquisition.monte_carlo import qExpectedImprovement, qProbabilityOfImprovement, qUpperConfidenceBound +from botorch.acquisition.multi_objective.monte_carlo import qNoisyExpectedHypervolumeImprovement + + +class qConstrainedUpperConfidenceBound(qUpperConfidenceBound): + def __init__(self, constraint, beta=4, *args, **kwargs): + super().__init__(beta=beta, *args, **kwargs) + self.constraint = constraint + self.beta = torch.tensor(beta) + + def forward(self, x): + posterior = self.model.posterior(x) + mean, sigma = posterior.mean, posterior.variance.sqrt() + + p_eff = 0.5 * (1 + torch.special.erf(self.beta.sqrt() / math.sqrt(2))) * torch.clamp(self.constraint(x), min=1e-6) + + return mean + sigma * np.sqrt(2) * torch.special.erfinv(2 * p_eff - 1) + + +class qConstrainedExpectedImprovement(qExpectedImprovement): + def __init__(self, constraint, *args, **kwargs): + super().__init__(*args, **kwargs) + self.constraint = constraint + + def forward(self, x): + return super().forward(x) * self.constraint(x) + + +class qConstrainedProbabilityOfImprovement(qProbabilityOfImprovement): + def __init__(self, constraint, *args, **kwargs): + super().__init__(*args, **kwargs) + self.constraint = constraint + + def forward(self, x): + return super().forward(x) * self.constraint(x) + + +class qConstrainedNoisyExpectedHypervolumeImprovement(qNoisyExpectedHypervolumeImprovement): + def __init__(self, constraint, *args, **kwargs): + super().__init__(*args, **kwargs) + self.constraint = constraint + + def forward(self, x): + return super().forward(x) * self.constraint(x) + + +class qConstrainedLowerBoundMaxValueEntropy(qLowerBoundMaxValueEntropy): + def __init__(self, constraint, *args, **kwargs): + super().__init__(*args, **kwargs) + self.constraint = constraint + + def forward(self, x): + return super().forward(x) * self.constraint(x) diff --git a/bloptools/bayesian/digestion.py b/bloptools/bayesian/digestion.py index eed3199..05805a0 100644 --- a/bloptools/bayesian/digestion.py +++ b/bloptools/bayesian/digestion.py @@ -1,4 +1,3 @@ def default_digestion_function(db, uid): products = db[uid].table(fill=True) - print(products) return products diff --git a/bloptools/bayesian/models.py b/bloptools/bayesian/models.py index 8733ac8..eff9a64 100644 --- a/bloptools/bayesian/models.py +++ b/bloptools/bayesian/models.py @@ -25,7 +25,7 @@ class LatentDirichletClassifier(LatentGP): def __init__(self, train_inputs, train_targets, skew_dims=True, *args, **kwargs): super().__init__(train_inputs, train_targets, skew_dims, *args, **kwargs) - def probabilities(self, x, n_samples=256): + def probabilities(self, x, n_samples=1024): """ Takes in a (..., m) dimension tensor and returns a (..., n_classes) tensor """ diff --git a/bloptools/tests/test_acq_funcs.py b/bloptools/tests/test_acq_funcs.py index 0dcead6..701d3d1 100644 --- a/bloptools/tests/test_acq_funcs.py +++ b/bloptools/tests/test_acq_funcs.py @@ -22,6 +22,7 @@ def test_acq_funcs_single_task(RE, db): db=db, ) - RE(agent.initialize("qr", n_init=64)) - RE(agent.learn("ei", n_iter=2)) - RE(agent.learn("pi", n_iter=2)) + RE(agent.learn("qr", n=64)) + + RE(agent.learn("qei", n=2)) + RE(agent.learn("qpi", n=2)) diff --git a/bloptools/tests/test_agent.py b/bloptools/tests/test_agent.py deleted file mode 100644 index 66657ce..0000000 --- a/bloptools/tests/test_agent.py +++ /dev/null @@ -1,25 +0,0 @@ -import os - -import pytest - - -@pytest.mark.test_func -def test_writing_hypers(RE, agent): - RE(agent.initialize("qr", n_init=32)) - - agent.save_hypers("hypers.h5") - - RE(agent.initialize("qr", n_init=8, hypers="hypers.h5")) - - os.remove("hypers.h5") - - -@pytest.mark.test_func -def test_writing_hypers_multitask(RE, multitask_agent): - RE(multitask_agent.initialize("qr", n_init=32)) - - multitask_agent.save_hypers("hypers.h5") - - RE(multitask_agent.initialize("qr", n_init=8, hypers="hypers.h5")) - - os.remove("hypers.h5") diff --git a/bloptools/tests/test_passive_dofs.py b/bloptools/tests/test_passive_dofs.py index d6aa6b7..cf81020 100644 --- a/bloptools/tests/test_passive_dofs.py +++ b/bloptools/tests/test_passive_dofs.py @@ -26,7 +26,7 @@ def test_passive_dofs(RE, db): tolerate_acquisition_errors=False, ) - RE(agent.initialize("qr", n_init=32)) + RE(agent.learn("qr", n=32)) agent.plot_tasks() agent.plot_acquisition() diff --git a/bloptools/tests/test_plots.py b/bloptools/tests/test_plots.py index aa1aa42..1d0d864 100644 --- a/bloptools/tests/test_plots.py +++ b/bloptools/tests/test_plots.py @@ -3,7 +3,7 @@ @pytest.mark.test_func def test_plots(RE, agent): - RE(agent.initialize("qr", n_init=32)) + RE(agent.learn("qr", n=32)) agent.plot_tasks() agent.plot_acquisition() diff --git a/docs/source/tutorials.rst b/docs/source/tutorials.rst index 297b303..3506da7 100644 --- a/docs/source/tutorials.rst +++ b/docs/source/tutorials.rst @@ -4,7 +4,6 @@ Tutorials .. toctree:: :maxdepth: 2 - tutorials/introduction.ipynb - tutorials/constrained-himmelblau.ipynb + tutorials/himmelblau.ipynb tutorials/hyperparameters.ipynb tutorials/passive-dofs.ipynb diff --git a/docs/source/tutorials/himmelblau.ipynb b/docs/source/tutorials/himmelblau.ipynb new file mode 100644 index 0000000..1a93aa7 --- /dev/null +++ b/docs/source/tutorials/himmelblau.ipynb @@ -0,0 +1,310 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "id": "e7b5e13a-c059-441d-8d4f-fff080d52054", + "metadata": {}, + "source": [ + "# Introduction (Himmelblau's function)\n", + "\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "c18ef717", + "metadata": {}, + "source": [ + "Let's use ``bloptools`` to minimize Himmelblau's function, which has four global minima:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "22438de8", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import matplotlib as mpl\n", + "from matplotlib import pyplot as plt\n", + "from bloptools import test_functions\n", + "\n", + "x1 = x2 = np.linspace(-6, 6, 1024)\n", + "X1, X2 = np.meshgrid(x1, x2)\n", + "from bloptools.tasks import Task\n", + "\n", + "task = Task(key=\"himmelblau\", kind=\"min\")\n", + "F = test_functions.himmelblau(X1, X2)\n", + "\n", + "plt.pcolormesh(x1, x2, F, norm=mpl.colors.LogNorm())\n", + "plt.colorbar()\n", + "plt.xlabel(\"x1\")\n", + "plt.ylabel(\"x2\")" + ] + }, + { + "cell_type": "markdown", + "id": "2500c410", + "metadata": {}, + "source": [ + "There are several things that our agent will need. The first ingredient is some degrees of freedom (these are always `ophyd` devices) which the agent will move around to different inputs within each DOF's bounds (the second ingredient). We define these here:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5d6df7a4", + "metadata": {}, + "outputs": [], + "source": [ + "from bloptools import devices\n", + "\n", + "dofs = [\n", + " {\"device\": devices.DOF(name=\"x1\"), \"limits\": (-6, 6), \"kind\": \"active\"},\n", + " {\"device\": devices.DOF(name=\"x2\"), \"limits\": (-6, 6), \"kind\": \"active\"},\n", + "]" + ] + }, + { + "cell_type": "markdown", + "id": "54b6f23e", + "metadata": {}, + "source": [ + "We also need to give the agent something to do. We want our agent to look in the feedback for a variable called \"himmelblau\", and try to minimize it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c8556bc9", + "metadata": {}, + "outputs": [], + "source": [ + "tasks = [\n", + " {\"key\": \"himmelblau\", \"kind\": \"minimize\"},\n", + "]" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "7a88c7bd", + "metadata": {}, + "source": [ + "In our digestion function, we define our objective as a deterministic function of the inputs:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e6bfcf73", + "metadata": {}, + "outputs": [], + "source": [ + "def digestion(db, uid):\n", + " products = db[uid].table()\n", + "\n", + " for index, entry in products.iterrows():\n", + " products.loc[index, \"himmelblau\"] = test_functions.himmelblau(entry.x1, entry.x2)\n", + "\n", + " return products" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "0d3d91c3", + "metadata": {}, + "source": [ + "We then combine these ingredients into an agent, giving it an instance of ``databroker`` so that it can see the output of the plans it runs." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "071a829f-a390-40dc-9d5b-ae75702e119e", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from bloptools.utils import prepare_re_env\n", + "\n", + "%run -i $prepare_re_env.__file__ --db-type=temp\n", + "from bloptools.bayesian import Agent\n", + "\n", + "agent = Agent(\n", + " dofs=dofs,\n", + " tasks=tasks,\n", + " digestion=digestion,\n", + " db=db,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "b7a608d6", + "metadata": {}, + "source": [ + "To decide which points to sample, the agent needs an acquisition function. The available acquisition function are here:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fb06739b", + "metadata": {}, + "outputs": [], + "source": [ + "agent.acq_func_info" + ] + }, + { + "cell_type": "markdown", + "id": "27685849", + "metadata": {}, + "source": [ + "Without any data, we can't make any inferences about what the function looks like, and so we can't use any non-trivial acquisition functions. Let's start by quasi-randomly sampling the parameter space, and plotting our model of the function:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "996da937", + "metadata": {}, + "outputs": [], + "source": [ + "RE(agent.learn(\"quasi-random\", n=32))\n", + "agent.plot_tasks()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "ab608930", + "metadata": {}, + "source": [ + "Now we can start to learn intelligently. Using the shorthand acquisition functions shown above, we can see the output of a few different ones:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "43b55f4f", + "metadata": {}, + "outputs": [], + "source": [ + "agent.plot_acquisition(acq_funcs=[\"qei\", \"qpi\", \"ucb\"])" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "18210f81-0e23-42b7-8589-77dc260e3131", + "metadata": {}, + "source": [ + "To decide where to go, the agent will find the inputs that maximize a given acquisition function:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b902172e-e89c-4346-89f3-bf9571cba6b3", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "agent.ask(\"qei\", n=1)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "9a888385-4e09-4fea-9282-cd6a6fe2c3df", + "metadata": {}, + "source": [ + "We can also ask the agent for multiple points to sample and it will jointly maximize the acquisition function over all sets of inputs, and find the most efficient route between them:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "28c5c0df", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "X, _ = agent.ask(\"qei\", n=8)\n", + "agent.plot_acquisition(acq_funcs=[\"qei\"])\n", + "plt.plot(*X.T, lw=5e-1, c=\"r\", marker=\"x\")" + ] + }, + { + "cell_type": "markdown", + "id": "23f3f7ef-c024-4ac1-9144-d0b6fb8a3944", + "metadata": {}, + "source": [ + "All of this is automated inside the ``learn`` method, which will find a point (or points) to sample, sample them, and retrain the model and its hyperparameters with the new data. To do 4 learning iterations of 8 points each, we can run" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ff1c5f1c", + "metadata": {}, + "outputs": [], + "source": [ + "RE(agent.learn(\"qei\", n=8, iterations=4))" + ] + }, + { + "cell_type": "markdown", + "id": "b52f3352-3b67-431c-b5af-057e02def5ba", + "metadata": {}, + "source": [ + "Our agent has found all the global minima of Himmelblau's function using Bayesian optimization, and we can ask it for the best point: " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0d5cc0c8-33cf-4fb1-b91c-81828e249f6a", + "metadata": {}, + "outputs": [], + "source": [ + "agent.plot_tasks()\n", + "print(agent.best_inputs)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.4" + }, + "vscode": { + "interpreter": { + "hash": "eee21ccc240bdddd7cf04478199e20f7257541e2f592ca1a4d34ebdc0225d742" + } + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/source/tutorials/hyperparameters.ipynb b/docs/source/tutorials/hyperparameters.ipynb index e3908f5..e4026e1 100644 --- a/docs/source/tutorials/hyperparameters.ipynb +++ b/docs/source/tutorials/hyperparameters.ipynb @@ -91,7 +91,7 @@ " db=db,\n", ")\n", "\n", - "RE(agent.initialize(acq_func=\"qr\", n_init=16))\n", + "RE(agent.learn(acq_func=\"qr\", n=16))\n", "\n", "agent.plot_tasks()" ] @@ -114,7 +114,7 @@ }, "outputs": [], "source": [ - "agent.plot_acquisition(strategy=[\"ei\", \"pi\", \"ucb\"])" + "agent.plot_acquisition(acq_funcs=[\"ei\", \"pi\", \"ucb\"])" ] }, { @@ -124,7 +124,7 @@ "metadata": {}, "outputs": [], "source": [ - "RE(agent.learn(\"ei\", n_iter=8))\n", + "RE(agent.learn(\"qei\", n=4, iterations=4))\n", "agent.plot_tasks()" ] } diff --git a/docs/source/tutorials/passive-dofs.ipynb b/docs/source/tutorials/passive-dofs.ipynb index 2ce4245..363a2a0 100644 --- a/docs/source/tutorials/passive-dofs.ipynb +++ b/docs/source/tutorials/passive-dofs.ipynb @@ -59,7 +59,7 @@ " db=db,\n", ")\n", "\n", - "RE(agent.initialize(\"qr\", n_init=32))\n", + "RE(agent.learn(\"qr\", n=32))\n", "\n", "agent.plot_tasks()" ] @@ -67,7 +67,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3.11.4 64-bit", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, @@ -81,7 +81,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.5" + "version": "3.11.4" }, "vscode": { "interpreter": { diff --git a/docs/source/tutorials/constrained-himmelblau.ipynb b/docs/wip/constrained-himmelblau copy.ipynb similarity index 59% rename from docs/source/tutorials/constrained-himmelblau.ipynb rename to docs/wip/constrained-himmelblau copy.ipynb index 95f863d..b317aa1 100644 --- a/docs/source/tutorials/constrained-himmelblau.ipynb +++ b/docs/wip/constrained-himmelblau copy.ipynb @@ -6,7 +6,7 @@ "id": "e7b5e13a-c059-441d-8d4f-fff080d52054", "metadata": {}, "source": [ - "# Feasibility modeling\n", + "# Introduction (Himmelblau's function)\n", "\n" ] }, @@ -16,7 +16,7 @@ "id": "c18ef717", "metadata": {}, "source": [ - "A more complicated example involves minimizing in two dimensions, where some parts of the parameter space are off-limits. Let's minimize Himmelblau's function, subject to the constraint that $x_1^2 + x_2^2 < 50$" + "Let's use ``bloptools`` to minimize Himmelblau's function, subject to the constraint that $x_1^2 + x_2^2 < 50$. Our function looks like this:" ] }, { @@ -38,19 +38,62 @@ "task = Task(key=\"himmelblau\", kind=\"min\")\n", "F = test_functions.constrained_himmelblau(X1, X2)\n", "\n", - "plt.pcolormesh(x1, x2, F, norm=mpl.colors.LogNorm(), shading=\"auto\")\n", + "plt.pcolormesh(x1, x2, F, norm=mpl.colors.LogNorm(), cmap=\"gnuplot\")\n", "plt.colorbar()\n", "plt.xlabel(\"x1\")\n", "plt.ylabel(\"x2\")" ] }, + { + "cell_type": "markdown", + "id": "2500c410", + "metadata": {}, + "source": [ + "There are several things that our agent will need. The first ingredient is some degrees of freedom (these are always `ophyd` devices) which the agent will move around to different inputs within each DOF's bounds (the second ingredient). We define these here:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5d6df7a4", + "metadata": {}, + "outputs": [], + "source": [ + "from bloptools import devices\n", + "\n", + "dofs = [\n", + " {\"device\": devices.DOF(name=\"x1\"), \"limits\": (-8, 8), \"kind\": \"active\"},\n", + " {\"device\": devices.DOF(name=\"x2\"), \"limits\": (-8, 8), \"kind\": \"active\"},\n", + "]" + ] + }, + { + "cell_type": "markdown", + "id": "54b6f23e", + "metadata": {}, + "source": [ + "We also need to give the agent something to do. We want our agent to look in the feedback for a variable called \"himmelblau\", and try to minimize it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c8556bc9", + "metadata": {}, + "outputs": [], + "source": [ + "tasks = [\n", + " {\"key\": \"himmelblau\", \"kind\": \"minimize\"},\n", + "]" + ] + }, { "attachments": {}, "cell_type": "markdown", "id": "7a88c7bd", "metadata": {}, "source": [ - "where everything outside our constraint is undefined. In our digestion function, we return a `NaN` when we violate the constraint:" + "In our digestion function, we define our objective as a deterministic function of the inputs, returning a `NaN` when we violate the constraint:" ] }, { @@ -75,7 +118,7 @@ "id": "0d3d91c3", "metadata": {}, "source": [ - "and create the agent in the usual way:" + "We then combine these ingredients into an agent, giving it an instance of ``databroker`` so that it can see the output of the plans it runs." ] }, { @@ -90,31 +133,63 @@ "from bloptools.utils import prepare_re_env\n", "\n", "%run -i $prepare_re_env.__file__ --db-type=temp\n", - "\n", - "from bloptools import devices\n", "from bloptools.bayesian import Agent\n", "\n", - "dofs = [\n", - " {\"device\": devices.DOF(name=\"x1\"), \"limits\": (-8, 8), \"kind\": \"active\"},\n", - " {\"device\": devices.DOF(name=\"x2\"), \"limits\": (-8, 8), \"kind\": \"active\"},\n", - "]\n", - "\n", - "tasks = [\n", - " {\"key\": \"himmelblau\", \"kind\": \"minimize\"},\n", - "]\n", - "\n", "agent = Agent(\n", " dofs=dofs,\n", " tasks=tasks,\n", " digestion=digestion,\n", " db=db,\n", - ")\n", - "\n", - "RE(agent.initialize(\"qr\", n_init=64))\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c4ec72a5", + "metadata": {}, + "outputs": [], + "source": [ + "import bloptools\n", "\n", + "bloptools.bayesian.acquisition.parse_acq_func(acq_func_identifier=\"quasi-random\")" + ] + }, + { + "cell_type": "markdown", + "id": "27685849", + "metadata": {}, + "source": [ + "Without any data, we can't make any inferences about what the function looks like, and so we can't use any non-trivial acquisition functions. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "996da937", + "metadata": {}, + "outputs": [], + "source": [ + "RE(agent.learn(\"quasi-random\", n=64))\n", "agent.plot_tasks()" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "fb2fa8e9", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e1a56c4c", + "metadata": {}, + "outputs": [], + "source": [] + }, { "attachments": {}, "cell_type": "markdown", @@ -133,7 +208,7 @@ }, "outputs": [], "source": [ - "agent.plot_validity()" + "agent.plot_validity(cmap=\"viridis\")" ] }, { @@ -154,7 +229,44 @@ }, "outputs": [], "source": [ - "agent.plot_acquisition(acq_func=[\"ei\", \"pi\", \"ucb\"])" + "X = agent.ask(\"qei\", n=8)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "50d627fe", + "metadata": {}, + "outputs": [], + "source": [ + "import scipy as sp\n", + "\n", + "X = sp.interpolate.interp1d(np.arange(len(X)), X, axis=0, kind=\"cubic\")(np.linspace(0, len(X) - 1, 16))\n", + "plt.plot(*X.T)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "43b55f4f", + "metadata": {}, + "outputs": [], + "source": [ + "agent.plot_acquisition(acq_func=[\"ei\", \"pi\", \"ucb\"], cmap=\"viridis\")\n", + "plt.plot(*X.T, c=\"r\", marker=\"x\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ca6cf39f", + "metadata": {}, + "outputs": [], + "source": [ + "import yaml\n", + "\n", + "with open(\"config.yml\", \"w\") as f:\n", + " yaml.safe_dump(ACQ_FUNC_CONFIG, f)" ] }, { @@ -164,7 +276,7 @@ "metadata": {}, "outputs": [], "source": [ - "RE(agent.learn(\"ei\", n_per_iter=4))" + "RE(agent.learn(\"qei\", n_per_iter=4))" ] }, { diff --git a/docs/source/tutorials/introduction.ipynb b/docs/wip/introduction.ipynb similarity index 100% rename from docs/source/tutorials/introduction.ipynb rename to docs/wip/introduction.ipynb