From 167768ecabbd8a880ed325aec43d6b4ffa73c2af Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Wed, 21 Jun 2023 16:25:45 -0400 Subject: [PATCH 01/11] pulled from main --- docs/source/tutorials/introduction copy.ipynb | 302 ++++++++++++++++++ 1 file changed, 302 insertions(+) create mode 100644 docs/source/tutorials/introduction copy.ipynb diff --git a/docs/source/tutorials/introduction copy.ipynb b/docs/source/tutorials/introduction copy.ipynb new file mode 100644 index 0000000..90061dc --- /dev/null +++ b/docs/source/tutorials/introduction copy.ipynb @@ -0,0 +1,302 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "id": "e7b5e13a-c059-441d-8d4f-fff080d52054", + "metadata": {}, + "source": [ + "# Kernels and hyperparameters\n", + "\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "c18ef717", + "metadata": {}, + "source": [ + "This tutorial is an introduction to the syntax used by the optimizer, as well as the principles of Bayesian optimization in general.\n", + "\n", + "We'll start by minimizing Booth's function, which looks like this:" + ] + }, + { + "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", + "\n", + "import bloptools\n", + "\n", + "x1 = x2 = np.linspace(-10, 10, 256)\n", + "X1, X2 = np.meshgrid(x1, x2)\n", + "\n", + "plt.pcolormesh(x1, x2, bloptools.experiments.tests.himmelblau(X1, X2), norm=mpl.colors.LogNorm(), shading=\"auto\")\n", + "plt.colorbar()\n", + "plt.xlabel(\"x1\")\n", + "plt.ylabel(\"x2\")\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "ecef8da5", + "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": "4c870567", + "metadata": {}, + "outputs": [], + "source": [ + "import bloptools\n", + "\n", + "dofs = bloptools.experiments.tests.get_dummy_dofs(2)\n", + "bounds = [(-10, 10), (-10, 10)]" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "7a88c7bd", + "metadata": {}, + "source": [ + "The agent automatically samples at different inputs, but we often need some post-processing after data collection. In this case, we need to give the agent a way to compute Himmelblau's function. We accomplish this with a digestion function:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e6bfcf73", + "metadata": {}, + "outputs": [], + "source": [ + "def digestion(db, uid):\n", + "\n", + " table = db[uid].table()\n", + " products = {\"himmelblau\": []}\n", + "\n", + " for index, entry in table.iterrows():\n", + "\n", + " products[\"himmelblau\"].append(bloptools.experiments.tests.himmelblau(entry.x1, entry.x2))\n", + "\n", + " return products" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "dad64303", + "metadata": {}, + "source": [ + "The next ingredient is a task, which gives the agent something to do. We want it to minimize Himmelblau's function, so we make a task that will try to minimize the output of the digestion function called \"himmelblau\". We also include a transform function, which will make it easier to regress over the outputs of the function." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4c14d162", + "metadata": {}, + "outputs": [], + "source": [ + "from bloptools.tasks import Task\n", + "\n", + "task = Task(key=\"himmelblau\", kind=\"min\", transform=lambda x: np.log(1 + 1e-2 * x))" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "0d3d91c3", + "metadata": {}, + "source": [ + "Combining all of these with a databroker instance, we can make an agent:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "071a829f-a390-40dc-9d5b-ae75702e119e", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "%run -i ../../../examples/prepare_bluesky.py # prepare the bluesky environment\n", + "\n", + "boa = bloptools.bayesian.Agent(\n", + " dofs=dofs, # things which we move around\n", + " bounds=bounds, # where we're allowed to move them\n", + " tasks=task, # tasks for the optimizer\n", + " digestion=digestion, # how to process the acquisition into task data\n", + " db=db, # a databroker instance\n", + " )\n", + "\n", + "RE(boa.initialize(init_scheme='quasi-random', n_init=16))" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "d8f2da43", + "metadata": {}, + "source": [ + "We initialized the GP with the \"quasi-random\" strategy, as it doesn't require any prior data. We can view the state of the optimizer's posterior of the tasks over the input parameters:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "996c3c01-f91d-4a25-9b8d-eba5fa964504", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "boa.plot_tasks()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "4274930c", + "metadata": {}, + "source": [ + "We can also the agent's posterior about the probability of goodness over the parameters:" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "d2eb855c", + "metadata": {}, + "source": [ + "We want to learn a bit more, so we can ask the agent where to sample based off of some strategy. Here we use the \"esti\" strategy, which maximizes the expected sum of tasks improvement.\n", + "\n", + "We can ask the agent to \"route\" them using ``ortools``, so that we can sample them more quickly if it requires us to e.g. move motors." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "334f3c66", + "metadata": {}, + "outputs": [], + "source": [ + "X_to_sample = boa.ask(strategy='esti', n=16, optimize=True, route=True)\n", + "plt.scatter(*X_to_sample.T)\n", + "plt.plot(*X_to_sample.T)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "296d9fd2", + "metadata": {}, + "source": [ + "Let's tell the agent to learn a bit more (it will direct itself):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f0e74651", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "RE(boa.learn(strategy='esti', n_iter=1, n_per_iter=4))\n", + "boa.plot_tasks()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "aeab7a9b", + "metadata": {}, + "source": [ + "The agent has updated its model of the tasks, including refitting the hyperparameters. Continuing:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d6b39b54", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "RE(boa.learn(strategy='esti', n_iter=4, n_per_iter=4))\n", + "boa.plot_tasks()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "e955233f", + "metadata": {}, + "source": [ + "Eventually, we reach a point of saturation where no more improvement takes place:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d73e4fd5", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "RE(boa.learn(strategy='esti', n_iter=4, n_per_iter=4))\n", + "boa.plot_tasks()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5ad4b1e7", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "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.9.16" + }, + "vscode": { + "interpreter": { + "hash": "9aced674e98d511b4f654e147532c84d38dc986fe042b1e92785fb9d8df41f75" + } + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From 831fa85b50d81d154f86cfa0d970872e6dc8e880 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Wed, 28 Jun 2023 22:04:21 -0400 Subject: [PATCH 02/11] refactored botorch models --- .flake8 | 2 +- bloptools/__init__.py | 2 +- bloptools/bayesian/__init__.py | 738 +++++++++++------- bloptools/bayesian/acquisition.py | 44 +- bloptools/bayesian/kernels.py | 190 ++++- bloptools/bayesian/models.py | 66 +- bloptools/devices.py | 29 + bloptools/experiments/sirepo/tes.py | 39 +- bloptools/experiments/tests/__init__.py | 60 -- bloptools/{tasks/__init__.py => tasks.py} | 3 +- bloptools/test_functions.py | 39 + bloptools/tests/test_bayesian_shadow.py | 19 +- bloptools/tests/test_bayesian_test_funcs.py | 34 +- .../source/tutorials/custom-acquisition.ipynb | 114 +++ ...n copy.ipynb => introduction copy 2.ipynb} | 0 docs/source/tutorials/multi-task-sirepo.ipynb | 11 +- docs/source/tutorials/passive-dofs.ipynb | 335 ++++++++ 17 files changed, 1277 insertions(+), 448 deletions(-) create mode 100644 bloptools/devices.py delete mode 100644 bloptools/experiments/tests/__init__.py rename bloptools/{tasks/__init__.py => tasks.py} (95%) create mode 100644 bloptools/test_functions.py create mode 100644 docs/source/tutorials/custom-acquisition.ipynb rename docs/source/tutorials/{introduction copy.ipynb => introduction copy 2.ipynb} (100%) create mode 100644 docs/source/tutorials/passive-dofs.ipynb diff --git a/.flake8 b/.flake8 index a82760d..83b1308 100644 --- a/.flake8 +++ b/.flake8 @@ -9,6 +9,6 @@ exclude = examples/*.py, bloptools/tests/test_bayesian_shadow.py, versioneer.py, -max-line-length = 115 +max-line-length = 125 # Ignore some style 'errors' produced while formatting by 'black' ignore = E203, W503 diff --git a/bloptools/__init__.py b/bloptools/__init__.py index 16275c3..5591be5 100644 --- a/bloptools/__init__.py +++ b/bloptools/__init__.py @@ -1,4 +1,4 @@ -from . import bayesian, experiments, tasks # noqa F401 +from . import bayesian, devices, experiments, tasks # noqa F401 from ._version import get_versions __version__ = get_versions()["version"] diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index 1eeb815..793d7fc 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -1,6 +1,3 @@ -import copy -import logging - import bluesky.plan_stubs as bps import bluesky.plans as bp # noqa F401 import botorch @@ -11,6 +8,7 @@ import pandas as pd import scipy as sp import torch +from botorch.acquisition.analytic import LogExpectedImprovement, LogProbabilityOfImprovement from matplotlib import pyplot as plt from .. import utils @@ -18,7 +16,8 @@ mpl.rc("image", cmap="coolwarm") -COLOR_LIST = ["dodgerblue", "tomato", "mediumseagreen"] +DEFAULT_COLOR_LIST = ["dodgerblue", "tomato", "mediumseagreen"] +DEFAULT_COLORMAP = "inferno" def default_acquisition_plan(dofs, inputs, dets): @@ -26,19 +25,28 @@ def default_acquisition_plan(dofs, inputs, dets): return uid +def default_digestion_plan(db, uid): + return db[uid].table() + + +# let's be specific about our terminology. +# +# dofs: degrees of freedom are things that can change +# inputs: these are the values of the dofs, which may be transformed/normalized +# targets: these are what our model tries to predict from the inputs +# tasks: these are quantities that our agent will try to optimize over + + class Agent: + MAX_TEST_POINTS = 2**10 + def __init__( self, - dofs, - bounds, + active_dofs, + active_dof_bounds, tasks, - digestion, db, - acquisition=None, - detectors=None, - initialization=None, - training_iter=256, - verbose=True, + **kwargs, ): """ A Bayesian optimization agent. @@ -58,119 +66,59 @@ def __init__( db : A databroker instance. """ - self.dofs = np.atleast_1d(dofs) - self.bounds = np.atleast_2d(bounds) + self.active_dofs = np.atleast_1d(active_dofs) + self.active_dof_bounds = np.atleast_2d(active_dof_bounds) self.tasks = np.atleast_1d(tasks) - self.initialization = initialization - self.acquisition = acquisition if acquisition is not None else default_acquisition_plan - self.digestion = digestion self.db = db - for dof in self.dofs: - dof.kind = "hinted" - - for i, task in enumerate(self.tasks): - task.index = i - - self.dets = detectors if detectors is not None else np.array([]) - - self.n_dof = len(self.dofs) - self.target_names = [f"{task.name}_fitness" for task in self.tasks] - self.n_tasks = len(self.tasks) - - self.training_iter = training_iter - self.verbose = verbose - - MAX_TEST_POINTS = 2**11 + self.verbose = kwargs.get("verbose", False) + self.ignore_acquisition_errors = kwargs.get("ignore_acquisition_errors", False) - self.test_X = self.sampler(n=MAX_TEST_POINTS) + self.initialization = kwargs.get("initialization", None) + self.acquisition_plan = kwargs.get("acquisition_plan", default_acquisition_plan) + self.digestion = kwargs.get("digestion", default_digestion_plan) - self.test_inputs = self.unnormalize_inputs(self.test_X) + self.acquisition = acquisition.Acquisition() - n_per_dim = int(np.power(MAX_TEST_POINTS, 1 / self.n_dof)) - self.X_samples = np.linspace(0, 1, n_per_dim) - self.test_X_grid = np.swapaxes(np.r_[np.meshgrid(*self.n_dof * [self.X_samples], indexing="ij")], 0, -1) - self.test_inputs_grid = self.unnormalize_inputs(self.test_X_grid) - - self.inputs = np.empty((0, self.n_dof)) - self.targets = np.empty((0, self.n_tasks)) - self.table = pd.DataFrame() + self.dets = np.atleast_1d(kwargs.get("detectors", [])) + self.passive_dofs = np.atleast_1d(kwargs.get("passive_dofs", [])) - self._initialized = False - - def normalize_inputs(self, inputs): - return (inputs - self.bounds.min(axis=1)) / self.bounds.ptp(axis=1) - - def unnormalize_inputs(self, X): - return X * self.bounds.ptp(axis=1) + self.bounds.min(axis=1) - - def normalize_targets(self, targets): - return (targets - self.targets_mean) / (1e-20 + self.targets_scale) + for i, task in enumerate(self.tasks): + task.index = i - def unnormalize_targets(self, targets): - return targets * self.targets_scale + self.targets_mean + self.n_active_dof = self.active_dofs.size + self.n_passive_dof = self.passive_dofs.size - @property - def normalized_bounds(self): - return (self.bounds - self.bounds.min(axis=1)[:, None]) / self.bounds.ptp(axis=1)[:, None] + self.dofs = np.r_[self.active_dofs, self.passive_dofs] - @property - def targets_mean(self): - return np.nanmean(self.targets, axis=0) + self.n_dof = self.n_active_dof + self.n_passive_dof - @property - def targets_scale(self): - return np.nanstd(self.targets, axis=0) + self.n_tasks = len(self.tasks) - @property - def normalized_targets(self): - return self.normalize_targets(self.targets) + self.training_iter = kwargs.get("training_iter", 256) - @property - def current_inputs(self): - return np.array([dof.read()[dof.name]["value"] for dof in self.dofs]) + # self.test_normalized_active_inputs = self.sampler(n=self.MAX_TEST_POINTS) + self.test_active_inputs = self.unnormalize_active_inputs(self.sampler(n=self.MAX_TEST_POINTS)) - @property - def dof_names(self): - return [dof.name for dof in self.dofs] + n_per_active_dim = int(np.power(self.MAX_TEST_POINTS, 1 / self.n_active_dof)) - @property - def det_names(self): - return [det.name for det in self.dets] + test_normalized_active_inputs_grid = np.swapaxes( + np.r_[np.meshgrid(*self.n_active_dof * [np.linspace(0, 1, n_per_active_dim)])], 0, -1 + ) - @property - def best_sum_of_tasks(self): - return np.nanmax(self.targets.sum(axis=1)) + self.test_active_inputs_grid = self.unnormalize_active_inputs(test_normalized_active_inputs_grid) - @property - def best_sum_of_tasks_inputs(self): - return self.inputs[np.nanargmax(self.targets.sum(axis=1))] + self.table = pd.DataFrame() - @property - def go_to(self, inputs): - yield from bps.mv(*[_ for items in zip(self.dofs, np.atleast_1d(inputs).T) for _ in items]) + self._initialized = False @property - def go_to_best_sum_of_tasks(self): - yield from self.go_to(self.best_sum_of_tasks_inputs) - - def inspect_beam(self, index, border=None): - im = self.images[index] - - x_min, x_max, y_min, y_max, width_x, width_y = self.table.loc[ - index, ["x_min", "x_max", "y_min", "y_max", "width_x", "width_y"] - ] - - bbx = np.array([x_min, x_max])[[0, 0, 1, 1, 0]] - bby = np.array([y_min, y_max])[[0, 1, 1, 0, 0]] - - plt.figure() - plt.imshow(im, cmap="gray_r") - plt.plot(bbx, bby, lw=4e0, c="r") - - if border is not None: - plt.xlim(x_min - border * width_x, x_min + border * width_x) - plt.ylim(y_min - border * width_y, y_min + border * width_y) + def input_transform(self): + coefficient = torch.tensor(self.inputs.values.ptp(axis=0)) + offset = torch.tensor(self.inputs.values.min(axis=0)) + return botorch.models.transforms.input.AffineInputTransform( + d=self.n_dof, coefficient=coefficient, offset=offset + ) def save(self, filepath="./agent_data.h5"): """ @@ -181,18 +129,19 @@ def save(self, filepath="./agent_data.h5"): f.create_dataset("inputs", data=self.inputs) self.table.to_hdf(filepath, key="table") - def forget(self, n): - if n >= len(self.inputs): - raise ValueError(f"Cannot forget last {n} points (the agent only has {len(self.inputs)} points).") - self.tell(new_inputs=self.inputs[:-n], new_targets=self.targets[:-n], append=False) + def forget(self, index): + self.tell(new_table=self.table.drop(index=index), append=False) def sampler(self, n): """ - Returns $n$ quasi-randomly sampled points on the [0,1] ^ n_dof hypercube using Sobol sampling. + Returns $n$ quasi-randomly sampled points on the [0,1] ^ n_active_dof hypercube using Sobol sampling. """ - power_of_two = 2 ** int(np.ceil(np.log(n) / np.log(2))) - subset = np.random.choice(power_of_two, size=n, replace=False) - return sp.stats.qmc.Sobol(d=self.n_dof, scramble=True).random(n=power_of_two)[subset] + min_power_of_two = 2 ** int(np.ceil(np.log(n) / np.log(2))) + subset = np.random.choice(min_power_of_two, size=n, replace=False) + return sp.stats.qmc.Sobol(d=self.n_active_dof, scramble=True).random(n=min_power_of_two)[subset] + + def active_dof_sampler(self, n, q=1): + return botorch.utils.sampling.draw_sobol_samples(torch.tensor(self.active_dof_bounds.T), n=n, q=q) def initialize( self, @@ -206,11 +155,10 @@ def initialize( It should be passed to a Bluesky RunEngine. """ + init_table = None + if filepath is not None: - with h5py.File(filepath, "r") as f: - init_inputs = f["inputs"][:] - self.table = pd.read_hdf(filepath, key="table") - init_targets = self.table.loc[:, self.target_names].values + init_table = pd.read_hdf(filepath, key="table") # experiment-specific stuff if self.initialization is not None: @@ -218,11 +166,8 @@ def initialize( # now let's get bayesian if init_scheme == "quasi-random": - init_inputs = self.ask(strategy="quasi-random", n=n_init, route=True) - init_table = yield from self.acquire(self.dofs, init_inputs, self.dets) - self.table = pd.concat([self.table, init_table]) - self.table.index = np.arange(len(self.table)) - init_targets = init_table.loc[:, self.target_names].values + init_dof_inputs = self.ask(strategy="quasi-random", n=n_init, route=True) + init_table = yield from self.acquire(dof_inputs=init_dof_inputs) else: raise Exception( @@ -230,48 +175,56 @@ def initialize( "['quasi-random']." ) - if (init_inputs is None) or (init_targets is None): + if init_table is None: raise RuntimeError("Unhandled initialization error.") - no_good_samples_tasks = np.isnan(init_targets).all(axis=0) + no_good_samples_tasks = np.isnan(init_table.loc[:, self.target_names]).all(axis=0) if no_good_samples_tasks.any(): raise ValueError( f"The tasks {[self.tasks[i].name for i in np.where(no_good_samples_tasks)[0]]} " f"don't have any good samples." ) - self.tell(new_inputs=init_inputs, new_targets=init_targets, reuse_hypers=True, verbose=self.verbose) - + self.tell(new_table=init_table, verbose=self.verbose) self._initialized = True - def tell(self, new_inputs, new_targets, append=True, **kwargs): + def tell(self, new_table=None, append=True, **kwargs): """ Inform the agent about new inputs and targets for the model. """ - if append: - self.inputs = np.r_[self.inputs, np.atleast_2d(new_inputs)] - self.targets = np.r_[self.targets, np.atleast_2d(new_targets)] - else: - self.inputs = new_inputs - self.targets = new_targets - self.X = self.normalize_inputs(self.inputs) + new_table = pd.DataFrame() if new_table is None else new_table - # if hasattr(self.experiment, "IMAGE_NAME"): - # self.images = np.array([im for im in self.table[self.experiment.IMAGE_NAME].values]) + self.table = pd.concat([self.table, new_table]) if append else new_table + + self.table.loc[:, "total_fitness"] = self.table.loc[:, self.task_names].fillna(-np.inf).sum(axis=1) + self.table.index = np.arange(len(self.table)) + + # self.normalized_inputs = self.normalize_inputs(self.inputs) self.all_targets_valid = ~np.isnan(self.targets).any(axis=1) for task in self.tasks: - task.targets = self.targets[:, task.index] - task.normalized_targets = self.normalized_targets[:, task.index] + task.targets = self.targets.loc[:, task.name] + # + # task.targets_mean = np.nanmean(task.targets, axis=0) + # task.targets_scale = np.nanstd(task.targets, axis=0) - task.targets_mean = np.nanmean(task.targets) - task.targets_scale = np.nanstd(task.targets) + # task.normalized_targets = self.normalized_targets.loc[:, task.name] - task.classes = self.all_targets_valid.astype(int) + task.feasibility = self.all_targets_valid.astype(int) + + train_inputs = torch.tensor(self.inputs.loc[task.feasibility == 1].values).double().unsqueeze(0) + train_targets = ( + torch.tensor(task.targets.loc[task.feasibility == 1].values).double().unsqueeze(0).unsqueeze(-1) + ) - regressor_likelihood = gpytorch.likelihoods.GaussianLikelihood( + if train_inputs.ndim == 1: + train_inputs = train_inputs.unsqueeze(-1) + if train_targets.ndim == 1: + train_targets = train_targets.unsqueeze(-1) + + likelihood = gpytorch.likelihoods.GaussianLikelihood( noise_constraint=gpytorch.constraints.Interval( torch.tensor(task.MIN_NOISE_LEVEL).square(), torch.tensor(task.MAX_NOISE_LEVEL).square(), @@ -279,9 +232,11 @@ def tell(self, new_inputs, new_targets, append=True, **kwargs): ).double() task.regressor = models.BoTorchSingleTaskGP( - train_inputs=torch.tensor(self.X[task.classes == 1]).double(), - train_targets=torch.tensor(self.normalized_targets[:, task.index][task.classes == 1]).double(), - likelihood=regressor_likelihood, + train_inputs=train_inputs, + train_targets=train_targets, + likelihood=likelihood, + input_transform=self.input_transform, + outcome_transform=botorch.models.transforms.outcome.Standardize(m=1, batch_shape=torch.Size((1,))), ).double() task.regressor_mll = gpytorch.mlls.ExactMarginalLogLikelihood( @@ -289,19 +244,20 @@ def tell(self, new_inputs, new_targets, append=True, **kwargs): ) botorch.fit.fit_gpytorch_mll(task.regressor_mll, **kwargs) - self.scalarization = botorch.acquisition.objective.ScalarizedPosteriorTransform( - weights=torch.tensor(self.targets_scale).double(), - offset=self.targets_mean.sum(), + self.task_scalarization = botorch.acquisition.objective.ScalarizedPosteriorTransform( + weights=torch.tensor([*[task.weight for task in self.tasks], 1]).double(), + offset=0, ) - dirichlet_classifier_likelihood = gpytorch.likelihoods.DirichletClassificationLikelihood( + dirichlet_likelihood = gpytorch.likelihoods.DirichletClassificationLikelihood( torch.as_tensor(self.all_targets_valid).long(), learn_additional_noise=True ).double() self.dirichlet_classifier = models.BoTorchDirichletClassifier( - train_inputs=torch.tensor(self.X).double(), - train_targets=dirichlet_classifier_likelihood.transformed_targets, - likelihood=dirichlet_classifier_likelihood, + train_inputs=torch.tensor(self.inputs.values).double(), + train_targets=dirichlet_likelihood.transformed_targets.transpose(-1, -2).double(), + likelihood=dirichlet_likelihood, + input_transform=self.input_transform, ).double() self.dirichlet_classifier_mll = gpytorch.mlls.ExactMarginalLogLikelihood( @@ -309,47 +265,141 @@ def tell(self, new_inputs, new_targets, append=True, **kwargs): ) botorch.fit.fit_gpytorch_mll(self.dirichlet_classifier_mll, **kwargs) - self.classifier = botorch.models.deterministic.GenericDeterministicModel( + self.feas_model = botorch.models.deterministic.GenericDeterministicModel( f=lambda X: self.dirichlet_classifier.log_prob(X) ) - self.multimodel = botorch.models.model.ModelList(*[task.regressor for task in self.tasks]) + self.targets_model = botorch.models.model_list_gp_regression.ModelListGP( + *[task.regressor for task in self.tasks] + ) - def sample_acqf(self, acqf, n_test=256, optimize=False): - """ - Return the point that optimizes a given acqusition function. - """ + self.task_model = botorch.models.model.ModelList(*[task.regressor for task in self.tasks], self.feas_model) - def acq_loss(x, *args): - return -acqf(x, *args).detach().numpy() + # def sample_acqf(self, acqf, n_test=256, optimize=False): + # """ + # Return the point that optimizes a given acqusition function. + # """ - acq_args = (self,) + # def acq_loss(x, *args): + # return -acqf(x, *args) - test_X = self.sampler(n=n_test) - init_X = test_X[acq_loss(test_X, *acq_args).argmin()] + # acq_args = (self,) - if optimize: - res = sp.optimize.minimize( - fun=acq_loss, - args=acq_args, - x0=init_X, - bounds=self.normalized_bounds, - method="SLSQP", - options={"maxiter": 256}, - ) - X = res.x - else: - X = init_X + # test_inputs = self.unnormalize_active_inputs(self.sampler(n=n_test)) + # init_inputs = test_inputs[acq_loss(test_inputs, *acq_args).argmin()] + + # if optimize: + # res = sp.optimize.minimize( + # fun=acq_loss, + # args=acq_args, + # x0=init_inputs, + # bounds=self.active_dof_bounds, + # method="SLSQP", + # options={"maxiter": 256}, + # ) + # acq_inputs = res.x + # else: + # acq_inputs = init_inputs - # print(res) + # return acq_inputs, acq_loss(acq_inputs, *acq_args) - return X, acq_loss(X, *acq_args) + # def ask( + # self, + # tasks=None, + # classifier=None, + # strategy=None, + # greedy=True, + # n=1, + # disappointment=0, + # route=True, + # cost_model=None, + # n_test=1024, + # optimize=True, + # ): + # """ + # The next $n$ points to sample, recommended by the agent. + # """ + + # if route: + # unrouted_points = self.ask( + # tasks=tasks, + # classifier=classifier, + # strategy=strategy, + # greedy=greedy, + # n=n, + # disappointment=disappointment, + # route=False, + # cost_model=cost_model, + # n_test=n_test, + # ) + + # routing_index, _ = utils.get_routing(self.read_active_dofs, unrouted_points) + # return unrouted_points[routing_index] + + # if strategy.lower() == "quasi-random": + # return self.unnormalize_active_inputs(self.sampler(n=n)) + + # if not self._initialized: + # raise RuntimeError("The agent is not initialized!") + + # if tasks is None: + # tasks = self.tasks + + # if classifier is None: + # classifier = self.feas_model + + # if (not greedy) or (n == 1): + # acqf = None + + # if strategy.lower() == "ei": # maximize the expected improvement + + # acq_func = LogExpectedImprovement(agent.task_model, + # best_f=agent.best_sum_of_tasks, + # posterior_transform=agent.scalarization) + + # acqf = self.acquisition.log_expected_improvement + + # if strategy.lower() == "pi": # maximize the probability improvement + # acqf = self.acquisition.log_probability_of_improvement + + # if acqf is None: + # raise ValueError(f'Unrecognized strategy "{strategy}".') + + # inputs_to_sample, loss = self.sample_acqf(acqf, optimize=optimize) + + # return np.atleast_2d(inputs_to_sample) + + # if greedy and (n > 1): + # dummy_tasks = [copy.deepcopy(task) for task in self.tasks] + + # inputs_to_sample = np.zeros((0, self.n_dof)) + + # for i in range(n): + + # new_input = self.ask(strategy=strategy, tasks=dummy_tasks, classifier=classifier, greedy=True, n=1) + + # inputs_to_sample = np.r_[inputs_to_sample, new_input] + + # if not (i + 1 == n): # no point if we're on the last iteration + # fantasy_model = botorch.models.model.ModelList(*[task.regressor for task in self.tasks]) + # fantasy_posterior = fantasy_model.posterior( + # torch.tensor(self.normalize_inputs(new_input)).double() + # ) + + # fantasy_targets = fantasy_posterior.mean - disappointment * fantasy_posterior.variance.sqrt() + # new_targets = self.unnormalize_targets(fantasy_targets.detach().numpy()) + + # self.tell(new_table=new_table) + + # self.forget(index=self.table.index[-n:]) # forget what you saw here + + # return np.atleast_2d(inputs_to_sample) def ask( self, tasks=None, classifier=None, - strategy=None, + strategy="ei", greedy=True, n=1, disappointment=0, @@ -375,87 +425,71 @@ def ask( n_test=n_test, ) - routing_index, _ = utils.get_routing(self.current_inputs, unrouted_points) + routing_index, _ = utils.get_routing(self.read_active_dofs, unrouted_points) return unrouted_points[routing_index] if strategy.lower() == "quasi-random": - return self.unnormalize_inputs(self.sampler(n=n)) + return self.unnormalize_active_inputs(self.sampler(n=n)) if not self._initialized: raise RuntimeError("The agent is not initialized!") - if tasks is None: - tasks = self.tasks - - if classifier is None: - classifier = self.classifier - - if (not greedy) or (n == 1): - acqf = None - - if strategy.lower() == "esti": # maximize the expected improvement - acqf = acquisition.log_expected_sum_of_tasks_improvement - - if acqf is None: - raise ValueError(f'Unrecognized strategy "{strategy}".') - - X, loss = self.sample_acqf(acqf, optimize=optimize) - - return self.unnormalize_inputs(np.atleast_2d(X)) + self.acqf = None - if greedy and (n > 1): - dummy_tasks = [copy.deepcopy(task) for task in self.tasks] - - inputs_to_sample = np.zeros((0, self.n_dof)) - - for i in range(n): - new_input = self.ask(strategy=strategy, tasks=dummy_tasks, classifier=classifier, greedy=True, n=1) - - inputs_to_sample = np.r_[inputs_to_sample, new_input] + if strategy.lower() == "ei": # maximize the expected improvement + self.acqf = LogExpectedImprovement( + self.task_model, best_f=self.best_sum_of_tasks, posterior_transform=self.task_scalarization + ) - if not (i + 1 == n): # no point if we're on the last iteration - fantasy_multimodel = botorch.models.model.ModelList(*[task.regressor for task in self.tasks]) - fantasy_posterior = fantasy_multimodel.posterior( - torch.tensor(self.normalize_inputs(new_input)).double() - ) + if strategy.lower() == "pi": # maximize the expected improvement + self.acqf = LogProbabilityOfImprovement( + self.task_model, best_f=self.best_sum_of_tasks, posterior_transform=self.task_scalarization + ) - fantasy_targets = fantasy_posterior.mean - disappointment * fantasy_posterior.variance.sqrt() - new_targets = self.unnormalize_targets(fantasy_targets.detach().numpy()) + if self.acqf is None: + raise ValueError(f'Unrecognized strategy "{strategy}".') - self.tell(new_inputs=new_input, new_targets=new_targets) + BATCH_SIZE = 1 + NUM_RESTARTS = 7 + RAW_SAMPLES = 512 - self.forget(n=n) # forget what you saw here + candidates, _ = botorch.optim.optimize_acqf( + acq_function=self.acqf, + bounds=torch.tensor(self.active_dof_bounds).T, + q=BATCH_SIZE, + num_restarts=NUM_RESTARTS, + raw_samples=RAW_SAMPLES, # used for intialization heuristic + options={"batch_limit": 5, "maxiter": 200}, + ) - return inputs_to_sample + return candidates.detach().numpy() - def acquire(self, dofs, inputs, dets): + def acquire(self, dof_inputs): """ Acquire and digest according to the agent's acquisition and digestion plans. This should yield a table of sampled tasks with the same length as the sampled inputs. """ try: - uid = yield from self.acquisition(dofs, inputs, dets) + uid = yield from self.acquisition_plan( + self.dofs, dof_inputs, [*self.dets, *self.dofs, *self.passive_dofs] + ) products = self.digestion(self.db, uid) - acq_table = pd.DataFrame(inputs, columns=[dof.name for dof in dofs]) - acq_table.insert(0, "timestamp", pd.Timestamp.now()) - - for key, values in products.items(): - acq_table.loc[:, key] = values + if "rejected" not in products.keys(): + products["rejected"] = False # compute the fitness for each task - for index, entry in acq_table.iterrows(): + for index, entry in products.iterrows(): for task in self.tasks: - acq_table.loc[index, f"{task.name}_fitness"] = task.get_fitness(entry) + products.loc[index, task.name] = task.get_fitness(entry) except Exception as err: - acq_table = pd.DataFrame() - logging.warning(repr(err)) + raise err - if not len(inputs) == len(acq_table): + if not len(dof_inputs) == len(products): raise ValueError("The resulting table must be the same length as the sampled inputs!") - return acq_table + return products def learn( self, strategy, n_iter=1, n_per_iter=1, reuse_hypers=True, upsample=1, verbose=True, plots=[], **kwargs @@ -470,13 +504,9 @@ def learn( for i in range(n_iter): inputs_to_sample = np.atleast_2d(self.ask(n=n_per_iter, strategy=strategy, **kwargs)) - new_table = yield from self.acquire(self.dofs, inputs_to_sample, self.dets) + new_table = yield from self.acquire(inputs_to_sample) - self.table = pd.concat([self.table, new_table]) - self.table.index = np.arange(len(self.table)) - new_targets = new_table.loc[:, self.target_names].values - - self.tell(new_inputs=inputs_to_sample, new_targets=new_targets, reuse_hypers=reuse_hypers) + self.tell(new_table=new_table, reuse_hypers=reuse_hypers) def plot_tasks(self, **kwargs): if self.n_dof == 1: @@ -485,26 +515,28 @@ def plot_tasks(self, **kwargs): else: self._plot_tasks_many_dofs(**kwargs) - def plot_constraints(self, **kwargs): + def plot_feasibility(self, **kwargs): if self.n_dof == 1: - self._plot_constraints_one_dof(**kwargs) + self._plot_feasibility_one_dof(**kwargs) else: - self._plot_constraints_many_dofs(**kwargs) + self._plot_feasibility_many_dofs(**kwargs) - def _plot_constraints_one_dof(self, size=32): + def _plot_feasibility_one_dof(self, size=32): self.class_fig, self.class_ax = plt.subplots(1, 1, figsize=(4, 4), sharex=True, constrained_layout=True) - self.class_ax.scatter(self.inputs, self.all_targets_valid.astype(int), s=size) + self.class_ax.scatter(self.inputs.values, self.all_targets_valid.astype(int), s=size) - x = torch.tensor(self.test_X_grid.reshape(-1, self.n_dof)).double() - log_prob = self.dirichlet_classifier.log_prob(x).detach().numpy().reshape(self.test_X_grid.shape[:-1]) + x = torch.tensor(self.test_inputs_grid.reshape(-1, self.n_dof)).double() + log_prob = self.dirichlet_classifier.log_prob(x).detach().numpy().reshape(self.test_inputs_grid.shape[:-1]) self.class_ax.plot(self.test_inputs_grid.ravel(), np.exp(log_prob)) - self.class_ax.set_xlim(*self.bounds[0]) + self.class_ax.set_xlim(*self.active_dof_bounds[0]) - def _plot_constraints_many_dofs(self, axes=[0, 1], shading="nearest", cmap="inferno", size=32, gridded=None): + def _plot_feasibility_many_dofs( + self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, size=32, gridded=None + ): if gridded is None: gridded = self.n_dof == 2 @@ -517,16 +549,18 @@ def _plot_constraints_many_dofs(self, axes=[0, 1], shading="nearest", cmap="infe ax.set_ylabel(self.dofs[axes[1]].name) data_ax = self.class_axes[0].scatter( - *self.inputs.T[:2], s=size, c=self.all_targets_valid.astype(int), vmin=0, vmax=1, cmap=cmap + *self.inputs.values.T[:2], s=size, c=self.all_targets_valid.astype(int), vmin=0, vmax=1, cmap=cmap ) if gridded: - x = torch.tensor(self.test_X_grid.reshape(-1, self.n_dof)).double() - log_prob = self.dirichlet_classifier.log_prob(x).detach().numpy().reshape(self.test_X_grid.shape[:-1]) + x = torch.tensor(self.test_inputs_grid.reshape(-1, self.n_dof)).double() + log_prob = ( + self.dirichlet_classifier.log_prob(x).detach().numpy().reshape(self.test_inputs_grid.shape[:-1]) + ) self.class_axes[1].pcolormesh( - *(self.bounds[axes].ptp(axis=1) * self.X_samples[:, None] + self.bounds[axes].min(axis=1)).T, - np.exp(log_prob), + *np.swapaxes(self.test_inputs_grid, 0, -1), + np.exp(log_prob).T, shading=shading, cmap=cmap, vmin=0, @@ -534,16 +568,18 @@ def _plot_constraints_many_dofs(self, axes=[0, 1], shading="nearest", cmap="infe ) else: - x = torch.tensor(self.test_X).double() + x = torch.tensor(self.normalize_inputs(self.test_inputs)).double() log_prob = self.dirichlet_classifier.log_prob(x).detach().numpy() - self.class_axes[1].scatter(*self.test_X.T[axes], s=size, c=np.exp(log_prob), vmin=0, vmax=1, cmap=cmap) + self.class_axes[1].scatter( + *self.test_inputs.T[axes], s=size, c=np.exp(log_prob), vmin=0, vmax=1, cmap=cmap + ) self.class_fig.colorbar(data_ax, ax=self.class_axes[:2], location="bottom", aspect=32, shrink=0.8) - for ax in self.task_axes.ravel(): - ax.set_xlim(*self.bounds[axes[0]]) - ax.set_ylim(*self.bounds[axes[1]]) + for ax in self.class_axes.ravel(): + ax.set_xlim(*self.active_dof_bounds[axes[0]]) + ax.set_ylim(*self.active_dof_bounds[axes[1]]) def _plot_tasks_one_dof(self, size=32, lw=1e0): self.task_fig, self.task_axes = plt.subplots( @@ -557,17 +593,18 @@ def _plot_tasks_one_dof(self, size=32, lw=1e0): self.task_axes = np.atleast_1d(self.task_axes) for itask, task in enumerate(self.tasks): - color = COLOR_LIST[itask] + color = DEFAULT_COLOR_LIST[itask] self.task_axes[itask].set_ylabel(task.name) - x = torch.tensor(self.test_X_grid).double() - task_posterior = task.regressor.posterior(x) + task_posterior = task.regressor.posterior( + torch.tensor(self.normalize_inputs(self.test_inputs_grid)).double() + ) task_mean = task_posterior.mean.detach().numpy() * task.targets_scale + task.targets_mean task_sigma = task_posterior.variance.sqrt().detach().numpy() * task.targets_scale - self.task_axes[itask].scatter(self.inputs, task.targets, s=size, color=color) - self.task_axes[itask].plot(self.test_inputs_grid.ravel(), task_mean, lw=lw, color=color) + self.task_axes[itask].scatter(self.inputs.values, task.targets, s=size, color=color) + self.task_axes[itask].plot(self.test_active_inputs_grid.ravel(), task_mean, lw=lw, color=color) for z in [1, 2]: self.task_axes[itask].fill_between( @@ -579,9 +616,9 @@ def _plot_tasks_one_dof(self, size=32, lw=1e0): alpha=0.5**z, ) - self.task_axes[itask].set_xlim(*self.bounds[0]) + self.task_axes[itask].set_xlim(*self.active_dof_bounds[0]) - def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap="inferno", gridded=None, size=32): + def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=32): if gridded is None: gridded = self.n_dof == 2 @@ -607,29 +644,30 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap="inferno", self.task_axes[itask, 2].set_title("posterior std. dev.") data_ax = self.task_axes[itask, 0].scatter( - *self.inputs.T[axes], s=size, c=task.targets, norm=task_norm, cmap=cmap + *self.inputs.values.T[axes], s=size, c=task.targets, norm=task_norm, cmap=cmap ) - if gridded: - x = torch.tensor(self.test_X_grid).double() - else: - x = torch.tensor(self.test_X).double() + x = ( + torch.tensor(self.test_inputs_grid).double() + if gridded + else torch.tensor(self.test_inputs).double() + ) task_posterior = task.regressor.posterior(x) - task_mean = task_posterior.mean.detach().numpy() * task.targets_scale + task.targets_mean - task_sigma = task_posterior.variance.sqrt().detach().numpy() * task.targets_scale + task_mean = task_posterior.mean.detach().numpy() # * task.targets_scale + task.targets_mean + task_sigma = task_posterior.variance.sqrt().detach().numpy() # * task.targets_scale if gridded: self.task_axes[itask, 1].pcolormesh( - *(self.bounds[axes].ptp(axis=1) * self.X_samples[:, None] + self.bounds[axes].min(axis=1)).T, - task_mean.reshape(self.test_X_grid.shape[:-1]), + *np.swapaxes(self.test_inputs_grid, 0, -1), + task_mean.reshape(self.test_active_inputs_grid.shape[:-1]).T, shading=shading, cmap=cmap, norm=task_norm, ) sigma_ax = self.task_axes[itask, 2].pcolormesh( - *(self.bounds[axes].ptp(axis=1) * self.X_samples[:, None] + self.bounds[axes].min(axis=1)).T, - task_sigma.reshape(self.test_X_grid.shape[:-1]), + *np.swapaxes(self.test_inputs_grid, 0, -1), + task_sigma.reshape(self.normalize_inputs(self.test_inputs_grid).shape[:-1]).T, shading=shading, cmap=cmap, ) @@ -646,5 +684,155 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap="inferno", self.task_fig.colorbar(sigma_ax, ax=self.task_axes[itask, 2], location="bottom", aspect=32, shrink=0.8) for ax in self.task_axes.ravel(): - ax.set_xlim(*self.bounds[axes[0]]) - ax.set_ylim(*self.bounds[axes[1]]) + ax.set_xlim(*self.active_dof_bounds[axes[0]]) + ax.set_ylim(*self.active_dof_bounds[axes[1]]) + + def normalize_active_inputs(self, inputs): + return (inputs - self.active_dof_bounds.min(axis=1)) / self.active_dof_bounds.ptp(axis=1) + + def unnormalize_active_inputs(self, X): + return X * self.active_dof_bounds.ptp(axis=1) + self.active_dof_bounds.min(axis=1) + + def normalize_inputs(self, inputs): + return (inputs - self.input_bounds.min(axis=1)) / self.input_bounds.ptp(axis=1) + + def unnormalize_inputs(self, X): + return X * self.input_bounds.ptp(axis=1) + self.input_bounds.min(axis=1) + + def normalize_targets(self, targets): + return (targets - self.targets_mean) / (1e-20 + self.targets_scale) + + def unnormalize_targets(self, targets): + return targets * self.targets_scale + self.targets_mean + + @property + def test_inputs(self): + test_passive_inputs = ( + self.passive_inputs.values[-1][None] * np.ones(len(self.test_active_inputs))[..., None] + ) + return np.concatenate([self.test_active_inputs, test_passive_inputs], axis=-1) + + @property + def test_inputs_grid(self): + test_passive_inputs_grid = self.passive_inputs.values[-1] * np.ones( + (*self.test_active_inputs_grid.shape[:-1], self.n_passive_dof) + ) + return np.concatenate([self.test_active_inputs_grid, test_passive_inputs_grid], axis=-1) + + @property + def inputs(self): + return self.table.loc[:, self.dof_names].astype(float) + + @property + def active_inputs(self): + return self.inputs.loc[:, self.active_dof_names] + + @property + def passive_inputs(self): + return self.inputs.loc[:, self.passive_dof_names] + + @property + def targets(self): + return self.table.loc[:, self.task_names].astype(float) + + @property + def feasible(self): + with pd.option_context("mode.use_inf_as_null", True): + feasible = ~self.targets.isna() + return feasible + + @property + def input_bounds(self): + lower_bound = np.r_[ + self.active_dof_bounds[:, 0], np.nanmin(self.passive_inputs.astype(float).values, axis=0) + ] + upper_bound = np.r_[ + self.active_dof_bounds[:, 1], np.nanmax(self.passive_inputs.astype(float).values, axis=0) + ] + return np.c_[lower_bound, upper_bound] + + @property + def targets_mean(self): + return np.nanmean(self.targets, axis=0) + + @property + def targets_scale(self): + return np.nanstd(self.targets, axis=0) + + @property + def normalized_targets(self): + return self.normalize_targets(self.targets) + + @property + def read_active_dofs(self): + return np.array([dof.read()[dof.name]["value"] for dof in self.active_dofs]) + + @property + def read_passive_dofs(self): + return np.array([dof.read()[dof.name]["value"] for dof in self.passive_dofs]) + + @property + def read_dofs(self): + return np.r_[self.read_active_dofs, self.read_passive_dofs] + + @property + def active_dof_names(self): + return [dof.name for dof in self.active_dofs] + + @property + def passive_dof_names(self): + return [dof.name for dof in self.passive_dofs] + + @property + def dof_names(self): + return [dof.name for dof in self.dofs] + + @property + def det_names(self): + return [det.name for det in self.dets] + + @property + def target_names(self): + return [task.name for task in self.tasks] + + @property + def task_names(self): + return [task.name for task in self.tasks] + + @property + def task_weights(self): + return np.array([task.weight for task in self.tasks]) + + @property + def best_sum_of_tasks(self): + return self.targets.fillna(-np.inf).sum(axis=1).max() + + @property + def best_sum_of_tasks_inputs(self): + return self.inputs[np.nanargmax(self.targets.sum(axis=1))] + + @property + def go_to(self, inputs): + yield from bps.mv(*[_ for items in zip(self.dofs, np.atleast_1d(inputs).T) for _ in items]) + + @property + def go_to_best_sum_of_tasks(self): + yield from self.go_to(self.best_sum_of_tasks_inputs) + + def inspect_beam(self, index, border=None): + im = self.images[index] + + x_min, x_max, y_min, y_max, width_x, width_y = self.table.loc[ + index, ["x_min", "x_max", "y_min", "y_max", "width_x", "width_y"] + ] + + bbx = np.array([x_min, x_max])[[0, 0, 1, 1, 0]] + bby = np.array([y_min, y_max])[[0, 1, 1, 0, 0]] + + plt.figure() + plt.imshow(im, cmap="gray_r") + plt.plot(bbx, bby, lw=4e0, c="r") + + if border is not None: + plt.xlim(x_min - border * width_x, x_min + border * width_x) + plt.ylim(y_min - border * width_y, y_min + border * width_y) diff --git a/bloptools/bayesian/acquisition.py b/bloptools/bayesian/acquisition.py index 5586302..0a041dc 100644 --- a/bloptools/bayesian/acquisition.py +++ b/bloptools/bayesian/acquisition.py @@ -1,19 +1,37 @@ import torch -from botorch.acquisition.analytic import LogExpectedImprovement +from botorch.acquisition.analytic import LogExpectedImprovement, LogProbabilityOfImprovement -def log_expected_sum_of_tasks_improvement(candidates, agent): - """ - Return the expected improvement in the sum of tasks. - """ +class Acquisition: + def __init__(self, *args, **kwargs): + ... - *input_shape, n_dim = candidates.shape - x = torch.as_tensor(candidates.reshape(-1, 1, n_dim)).double() + @staticmethod + def log_expected_improvement(candidates, agent): + *input_shape, n_dim = candidates.shape - LEI = LogExpectedImprovement( - model=agent.multimodel, best_f=agent.best_sum_of_tasks, posterior_transform=agent.scalarization - ) - lei = LEI.forward(x).reshape(input_shape) - log_prob = agent.dirichlet_classifier.log_prob(x).reshape(input_shape) + x = torch.as_tensor(candidates.reshape(-1, 1, n_dim)).double() - return torch.clamp(lei + log_prob, min=-16) + LEI = LogExpectedImprovement( + model=agent.task_model, best_f=agent.best_sum_of_tasks, posterior_transform=agent.scalarization + ) + + lei = LEI.forward(x) + feas_log_prob = agent.feas_model(x) + + return (lei.reshape(input_shape) + feas_log_prob.reshape(input_shape)).detach().numpy() + + @staticmethod + def log_probability_of_improvement(candidates, agent): + *input_shape, n_dim = candidates.shape + + x = torch.as_tensor(candidates.reshape(-1, 1, n_dim)).double() + + LPI = LogProbabilityOfImprovement( + model=agent.task_model, best_f=agent.best_sum_of_tasks, posterior_transform=agent.scalarization + ) + + lpi = LPI.forward(x) + feas_log_prob = agent.feas_model(x) + + return (lpi.reshape(input_shape) + feas_log_prob.reshape(input_shape)).detach().numpy() diff --git a/bloptools/bayesian/kernels.py b/bloptools/bayesian/kernels.py index e678667..96481a2 100644 --- a/bloptools/bayesian/kernels.py +++ b/bloptools/bayesian/kernels.py @@ -1,8 +1,134 @@ +import math + import gpytorch import numpy as np import torch +class MultiOutputLatentKernel(gpytorch.kernels.Kernel): + is_stationary = True + + def __init__( + self, + num_inputs=1, + num_outputs=1, + off_diag=False, + diag_prior=False, + **kwargs, + ): + super(MultiOutputLatentKernel, self).__init__() + + self.num_inputs = num_inputs + self.num_outputs = num_outputs + self.n_off_diag = int(num_inputs * (num_inputs - 1) / 2) + self.off_diag = off_diag + + self.nu = kwargs.get("nu", 1.5) + + # self.batch_shape = torch.Size([num_outputs]) + + # output_scale_constraint = gpytorch.constraints.Positive() + diag_params_constraint = gpytorch.constraints.Interval(1e0, 1e2) + skew_params_constraint = gpytorch.constraints.Interval(-1e0, 1e0) + + diag_params_initial = np.sqrt(diag_params_constraint.lower_bound * diag_params_constraint.upper_bound) + raw_diag_params_initial = diag_params_constraint.inverse_transform(diag_params_initial) + + self.register_parameter( + name="raw_diag_params", + parameter=torch.nn.Parameter( + raw_diag_params_initial * torch.ones(self.num_outputs, self.num_inputs).double() + ), + ) + + # self.register_constraint("raw_output_scale", output_scale_constraint) + self.register_constraint("raw_diag_params", diag_params_constraint) + + if diag_prior: + self.register_prior( + name="diag_params_prior", + prior=gpytorch.priors.GammaPrior(concentration=0.5, rate=0.2), + param_or_closure=lambda m: m.diag_params, + setting_closure=lambda m, v: m._set_diag_params(v), + ) + + if self.off_diag: + self.register_parameter( + name="raw_skew_params", + parameter=torch.nn.Parameter(torch.zeros(self.num_outputs, self.n_off_diag).double()), + ) + self.register_constraint("raw_skew_params", skew_params_constraint) + + @property + def diag_params(self): + return self.raw_diag_params_constraint.transform(self.raw_diag_params) + + @property + def skew_params(self): + return self.raw_skew_params_constraint.transform(self.raw_skew_params) + + @diag_params.setter + def diag_params(self, value): + self._set_diag_params(value) + + @skew_params.setter + def skew_params(self, value): + self._set_skew_params(value) + + def _set_skew_params(self, value): + if not torch.is_tensor(value): + value = torch.as_tensor(value).to(self.raw_skew_params) + self.initialize(raw_skew_params=self.raw_skew_params_constraint.inverse_transform(value)) + + def _set_diag_params(self, value): + if not torch.is_tensor(value): + value = torch.as_tensor(value).to(self.raw_diag_params) + self.initialize(raw_diag_params=self.raw_diag_params_constraint.inverse_transform(value)) + + @property + def dimension_transform(self): + # no rotations + if not self.off_diag: + T = torch.eye(self.num_inputs, dtype=torch.float64) + + # construct an orthogonal matrix. fun fact: exp(skew(N)) is the generator of SO(N) + else: + A = torch.zeros((self.num_outputs, self.num_inputs, self.num_inputs), dtype=torch.float64) + upper_indices = np.triu_indices(self.num_inputs, k=1) + for output_index in range(self.num_outputs): + A[(output_index, *upper_indices)] = self.skew_params[output_index] + A += -A.transpose(-1, -2) + T = torch.linalg.matrix_exp(A) + + diagonal_transform = torch.cat([torch.diag(_values).unsqueeze(0) for _values in self.diag_params], dim=0) + T = torch.matmul(diagonal_transform, T) + + return T + + def forward(self, x1, x2, diag=False, **params): + # adapted from the Matern kernel + mean = x1.reshape(-1, x1.size(-1)).mean(0)[(None,) * (x1.dim() - 1)] + + trans_x1 = torch.matmul(self.dimension_transform.unsqueeze(1), (x1 - mean).unsqueeze(-1)).squeeze(-1) + trans_x2 = torch.matmul(self.dimension_transform.unsqueeze(1), (x2 - mean).unsqueeze(-1)).squeeze(-1) + + distance = self.covar_dist(trans_x1, trans_x2, diag=diag, **params) + + # if distance.shape[0] == 1: + # distance = distance.squeeze(0) # this is extremely necessary + + exp_component = torch.exp(-math.sqrt(self.nu * 2) * distance) + + if self.nu == 0.5: + constant_component = 1 + elif self.nu == 1.5: + constant_component = (math.sqrt(3) * distance).add(1) + elif self.nu == 2.5: + constant_component = (math.sqrt(5) * distance).add(1).add(5.0 / 3.0 * distance**2) + + return constant_component * exp_component + + class LatentMaternKernel(gpytorch.kernels.Kernel): def __init__( self, @@ -18,80 +144,78 @@ def __init__( self.off_diag = off_diag # output_scale_constraint = gpytorch.constraints.Positive() - trans_diagonal_constraint = gpytorch.constraints.Interval(1e0, 1e2) - trans_off_diag_constraint = gpytorch.constraints.Interval(-1e0, 1e0) + diag_params_constraint = gpytorch.constraints.Interval(1e-1, 1e2) + skew_params_constraint = gpytorch.constraints.Interval(-1e0, 1e0) - trans_diagonal_initial = np.sqrt( - trans_diagonal_constraint.lower_bound * trans_diagonal_constraint.upper_bound - ) - raw_trans_diagonal_initial = trans_diagonal_constraint.inverse_transform(trans_diagonal_initial) + diag_params_initial = np.sqrt(diag_params_constraint.lower_bound * diag_params_constraint.upper_bound) + raw_diag_params_initial = diag_params_constraint.inverse_transform(diag_params_initial) # self.register_parameter( # name="raw_output_scale", parameter=torch.nn.Parameter(torch.ones(*self.batch_shape, 1).double()) # ) self.register_parameter( - name="raw_trans_diagonal", + name="raw_diag_params", parameter=torch.nn.Parameter( - raw_trans_diagonal_initial * torch.ones(*self.batch_shape, self.n_dim).double() + raw_diag_params_initial * torch.ones(*self.batch_shape, self.n_dim).double() ), ) # self.register_constraint("raw_output_scale", output_scale_constraint) - self.register_constraint("raw_trans_diagonal", trans_diagonal_constraint) + self.register_constraint("raw_diag_params", diag_params_constraint) if diagonal_prior: self.register_prior( - name="trans_diagonal_prior", - prior=gpytorch.priors.GammaPrior(concentration=1, rate=0.2), - param_or_closure=lambda m: m.trans_diagonal, - setting_closure=lambda m, v: m._set_trans_diagonal(v), + name="diag_params_prior", + prior=gpytorch.priors.GammaPrior(concentration=0.5, rate=0.2), + param_or_closure=lambda m: m.diag_params, + setting_closure=lambda m, v: m._set_diag_params(v), ) if self.off_diag: self.register_parameter( - name="raw_trans_off_diag", + name="raw_skew_params", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, self.n_off_diag).double()), ) - self.register_constraint("raw_trans_off_diag", trans_off_diag_constraint) + self.register_constraint("raw_skew_params", skew_params_constraint) # @property # def output_scale(self): # return self.raw_output_scale_constraint.transform(self.raw_output_scale) @property - def trans_diagonal(self): - return self.raw_trans_diagonal_constraint.transform(self.raw_trans_diagonal) + def diag_params(self): + return self.raw_diag_params_constraint.transform(self.raw_diag_params) @property - def trans_off_diag(self): - return self.raw_trans_off_diag_constraint.transform(self.raw_trans_off_diag) + def skew_params(self): + return self.raw_skew_params_constraint.transform(self.raw_skew_params) # @output_scale.setter # def output_scale(self, value): # self._set_output_scale(value) - @trans_diagonal.setter - def trans_diagonal(self, value): - self._set_trans_diagonal(value) + @diag_params.setter + def diag_params(self, value): + self._set_diag_params(value) - @trans_off_diag.setter - def trans_off_diag(self, value): - self._set_trans_off_diag(value) + @skew_params.setter + def skew_params(self, value): + self._set_skew_params(value) - def _set_trans_off_diag(self, value): + def _set_skew_params(self, value): if not torch.is_tensor(value): - value = torch.as_tensor(value).to(self.raw_trans_off_diag) - self.initialize(raw_trans_off_diag=self.raw_trans_off_diag_constraint.inverse_transform(value)) + value = torch.as_tensor(value).to(self.raw_skew_params) + self.initialize(raw_skew_params=self.raw_skew_params_constraint.inverse_transform(value)) # def _set_output_scale(self, value): # if not torch.is_tensor(value): # value = torch.as_tensor(value).to(self.raw_output_scale) # self.initialize(raw_output_scale=self.raw_output_scale_constraint.inverse_transform(value)) - def _set_trans_diagonal(self, value): + def _set_diag_params(self, value): if not torch.is_tensor(value): - value = torch.as_tensor(value).to(self.raw_trans_diagonal) - self.initialize(raw_trans_diagonal=self.raw_trans_diagonal_constraint.inverse_transform(value)) + value = torch.as_tensor(value).to(self.raw_diag_params) + self.initialize(raw_diag_params=self.raw_diag_params_constraint.inverse_transform(value)) @property def trans_matrix(self): @@ -102,11 +226,11 @@ def trans_matrix(self): # construct an orthogonal matrix. fun fact: exp(skew(N)) is the generator of SO(N) else: A = torch.zeros((self.n_dim, self.n_dim)).double() - A[np.triu_indices(self.n_dim, k=1)] = self.trans_off_diag + A[np.triu_indices(self.n_dim, k=1)] = self.skew_params A += -A.T T = torch.linalg.matrix_exp(A) - T = torch.matmul(torch.diag(self.trans_diagonal), T) + T = torch.matmul(torch.diag(self.diag_params), T) return T diff --git a/bloptools/bayesian/models.py b/bloptools/bayesian/models.py index d6ea2ea..91d2e7d 100644 --- a/bloptools/bayesian/models.py +++ b/bloptools/bayesian/models.py @@ -7,9 +7,67 @@ from . import kernels -class BoTorchSingleTaskGP(ExactGP, GPyTorchModel): +class BoTorchDirichletClassifier(botorch.models.gp_regression.SingleTaskGP): + # _num_outputs = 1 # to inform GPyTorchModel API + + def __init__(self, train_inputs, train_targets, *args, **kwargs): + super().__init__(train_inputs, train_targets, *args, **kwargs) + # self.mean_module = gpytorch.means.ConstantMean(batch_shape=len(train_targets.unique())) + # self.covar_module = gpytorch.kernels.ScaleKernel( + # kernels.LatentMaternKernel(n_dim=train_inputs.shape[-1], off_diag=False, diagonal_prior=False) + # ) + + # print(train_inputs.shape, train_targets.shape) + + self.mean_module = gpytorch.means.ConstantMean() # batch_shape=torch.Size((len(train_targets.unique()),))) + self.covar_module = gpytorch.kernels.ScaleKernel( + kernels.MultiOutputLatentKernel( + num_inputs=train_inputs.shape[-1], + num_outputs=train_targets.shape[-1], + off_diag=True, + diag_prior=True, + **kwargs + ) + ) + + def forward(self, x): + mean_x = self.mean_module(x) + covar_x = self.covar_module(x) + return gpytorch.distributions.MultivariateNormal(mean_x, covar_x) + + # def posterior(self, X): + # return botorch.posteriors.GPyTorchPosterior(gpytorch.distributions.MultivariateNormal()) + + def log_prob(self, x, n_samples=256): + *input_shape, n_dim = x.shape + samples = self.posterior(x.reshape(-1, n_dim)).sample(torch.Size((n_samples,))).exp() + return torch.log((samples / samples.sum(-1, keepdim=True)).mean(0)[:, 1]).reshape(*input_shape, 1) + + +class BoTorchSingleTaskGP(botorch.models.gp_regression.SingleTaskGP): + def __init__(self, train_inputs, train_targets, *args, **kwargs): + super().__init__(train_inputs, train_targets, *args, **kwargs) + + self.mean_module = gpytorch.means.ConstantMean() + self.covar_module = gpytorch.kernels.ScaleKernel( + kernels.MultiOutputLatentKernel( + num_inputs=train_inputs.shape[-1], + num_outputs=train_targets.shape[-1], + off_diag=True, + diag_prior=True, + **kwargs + ) + ) + + def forward(self, x): + mean_x = self.mean_module(x) + covar_x = self.covar_module(x) + return gpytorch.distributions.MultivariateNormal(mean_x, covar_x) + + +class OldBoTorchSingleTaskGP(ExactGP, GPyTorchModel): def __init__(self, train_inputs, train_targets, likelihood): - super(BoTorchSingleTaskGP, self).__init__(train_inputs, train_targets, likelihood) + super(OldBoTorchSingleTaskGP, self).__init__(train_inputs, train_targets, likelihood) self.mean_module = gpytorch.means.ConstantMean() self.covar_module = gpytorch.kernels.ScaleKernel( kernels.LatentMaternKernel(n_dim=train_inputs.shape[-1], off_diag=True, diagonal_prior=True) @@ -42,11 +100,11 @@ def forward(self, x): return gpytorch.distributions.MultitaskMultivariateNormal(mean_x, covar_x) -class BoTorchDirichletClassifier(gpytorch.models.ExactGP, botorch.models.gpytorch.GPyTorchModel): +class OldBoTorchDirichletClassifier(gpytorch.models.ExactGP, botorch.models.gpytorch.GPyTorchModel): _num_outputs = 1 # to inform GPyTorchModel API def __init__(self, train_inputs, train_targets, likelihood): - super(BoTorchDirichletClassifier, self).__init__(train_inputs, train_targets, likelihood) + super(OldBoTorchDirichletClassifier, self).__init__(train_inputs, train_targets, likelihood) self.mean_module = gpytorch.means.ConstantMean(batch_shape=len(train_targets.unique())) self.covar_module = gpytorch.kernels.ScaleKernel( kernels.LatentMaternKernel(n_dim=train_inputs.shape[-1], off_diag=False, diagonal_prior=False) diff --git a/bloptools/devices.py b/bloptools/devices.py new file mode 100644 index 0000000..6e56d2a --- /dev/null +++ b/bloptools/devices.py @@ -0,0 +1,29 @@ +import time as ttime + +from ophyd import Component as Cpt +from ophyd import Device, Signal, SignalRO + + +def dummy_dofs(n=2): + return [Signal(name=f"x{i+1}", value=0) for i in range(n)] + + +def get_dummy_device(name="dofs", n=2): + components = {} + + for i in range(n): + components[f"x{i+1}"] = Cpt(Signal, value=i + 1) + + cls = type("DOF", (Device,), components) + + device = cls(name=name) + + return [getattr(device, attr) for attr in device.read_attrs] + + +class TimeReadback(SignalRO): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def get(self): + return ttime.time() diff --git a/bloptools/experiments/sirepo/tes.py b/bloptools/experiments/sirepo/tes.py index 240f1ef..9fc26bb 100644 --- a/bloptools/experiments/sirepo/tes.py +++ b/bloptools/experiments/sirepo/tes.py @@ -6,30 +6,13 @@ def digestion(db, uid, image_name="w9"): Simulating a misaligned Gaussian beam. The optimum is at (1, 1, 1, 1) """ - table = db[uid].table(fill=True) - - products_keys = [ - "image", - "vertical_extent", - "horizontal_extent", - "flux", - "x_pos", - "y_pos", - "x_width", - "y_width", - ] - - products = {key: [] for key in products_keys} - - for index, entry in table.iterrows(): + products = db[uid].table(fill=True) + + for index, entry in products.iterrows(): image = getattr(entry, f"{image_name}_image") horizontal_extent = getattr(entry, f"{image_name}_horizontal_extent") vertical_extent = getattr(entry, f"{image_name}_vertical_extent") - products["image"].append(image) - products["vertical_extent"].append(vertical_extent) - products["horizontal_extent"].append(horizontal_extent) - flux = image.sum() n_y, n_x = image.shape @@ -47,13 +30,15 @@ def digestion(db, uid, image_name="w9"): bad |= any(np.isnan([mean_x, mean_y, sigma_x, sigma_y])) if not bad: - products["flux"].append(flux) - products["x_pos"].append(mean_x) - products["y_pos"].append(mean_y) - products["x_width"].append(2 * sigma_x) - products["y_width"].append(2 * sigma_y) + products.loc[index, ["flux", "x_pos", "y_pos", "x_width", "y_width"]] = ( + flux, + mean_x, + mean_y, + 2 * sigma_x, + 2 * sigma_y, + ) else: - for key in ["flux", "x_pos", "y_pos", "x_width", "y_width"]: - products[key].append(np.nan) + for col in ["flux", "x_pos", "y_pos", "x_width", "y_width"]: + products.loc[index, col] = np.nan return products diff --git a/bloptools/experiments/tests/__init__.py b/bloptools/experiments/tests/__init__.py deleted file mode 100644 index 89ce8a8..0000000 --- a/bloptools/experiments/tests/__init__.py +++ /dev/null @@ -1,60 +0,0 @@ -import numpy as np -from ophyd import Component as Cpt -from ophyd import Device, Signal - - -def get_dummy_dofs(n=2): - return [Signal(name=f"x{i+1}", value=0) for i in range(n)] - - -def get_dummy_device(name="dofs", n=2): - components = {} - - for i in range(n): - components[f"x{i+1}"] = Cpt(Signal, value=i + 1) - - cls = type("DOF", (Device,), components) - - device = cls(name=name) - - return [getattr(device, attr) for attr in device.read_attrs] - - -def booth(x1, x2): - return (x1 + 2 * x2 - 7) ** 2 + (2 * x1 + x2 - 5) ** 2 - - -def himmelblau(x1, x2): - return (x1**2 + x2 - 11) ** 2 + (x1 + x2**2 - 7) ** 2 - - -def gaussian_beam_waist(x1, x2): - return np.sqrt(1 + 0.25 * (x1 - x2) ** 2 + 16 * (x1 + x2 - 2) ** 2) - - -def himmelblau_digestion(db, uid): - table = db[uid].table() - products = {"himmelblau": []} - - for index, entry in table.iterrows(): - products["himmelblau"].append(himmelblau(entry.x1, entry.x2)) - - return products - - -def mock_kbs_digestion(db, uid): - """ - Simulating a misaligned Gaussian beam. The optimum is at (1, 1, 1, 1) - """ - - table = db[uid].table() - products = {"x_width": [], "y_width": []} - - for index, entry in table.iterrows(): - sigma_x = gaussian_beam_waist(entry.x1, entry.x2) - sigma_y = gaussian_beam_waist(entry.x3, entry.x4) - - products["x_width"].append(2 * sigma_x) - products["y_width"].append(2 * sigma_y) - - return products diff --git a/bloptools/tasks/__init__.py b/bloptools/tasks.py similarity index 95% rename from bloptools/tasks/__init__.py rename to bloptools/tasks.py index fa787e6..5437e46 100644 --- a/bloptools/tasks/__init__.py +++ b/bloptools/tasks.py @@ -5,8 +5,9 @@ class Task: def __init__(self, key, kind="max", name=None, transform=None, **kwargs): self.key = key self.kind = kind - self.name = name if name is not None else f"{kind}_{key}" + self.name = name if name is not None else f"{kind}_{key}_fitness" self.transform = transform if transform is not None else lambda x: x + self.weight = 1.0 if kind.lower() in ["min", "minimum", "minimize"]: self.sign = -1 diff --git a/bloptools/test_functions.py b/bloptools/test_functions.py new file mode 100644 index 0000000..e8cd621 --- /dev/null +++ b/bloptools/test_functions.py @@ -0,0 +1,39 @@ +import numpy as np + + +def booth(x1, x2): + return (x1 + 2 * x2 - 7) ** 2 + (2 * x1 + x2 - 5) ** 2 + + +def himmelblau(x1, x2): + return (x1**2 + x2 - 11) ** 2 + (x1 + x2**2 - 7) ** 2 + + +def gaussian_beam_waist(x1, x2): + return np.sqrt(1 + 0.25 * (x1 - x2) ** 2 + 16 * (x1 + x2 - 2) ** 2) + + +def himmelblau_digestion(db, uid): + products = db[uid].table() + + for index, entry in products.iterrows(): + products.loc[index, "himmelblau"] = himmelblau(entry.x1, entry.x2) + + return products + + +def mock_kbs_digestion(db, uid): + """ + Simulating a misaligned Gaussian beam. The optimum is at (1, 1, 1, 1) + """ + + products = db[uid].table() + + for index, entry in products.iterrows(): + sigma_x = gaussian_beam_waist(entry.x1, entry.x2) + sigma_y = gaussian_beam_waist(entry.x3, entry.x4) + + products.loc[index, "x_width"] = 2 * sigma_x + products.loc[index, "y_width"] = 2 * sigma_y + + return products diff --git a/bloptools/tests/test_bayesian_shadow.py b/bloptools/tests/test_bayesian_shadow.py index da32024..9d4eded 100644 --- a/bloptools/tests/test_bayesian_shadow.py +++ b/bloptools/tests/test_bayesian_shadow.py @@ -2,7 +2,7 @@ import pytest from sirepo_bluesky.sirepo_ophyd import create_classes -from bloptools.bayesian import Agent +import bloptools from bloptools.experiments.sirepo import tes from bloptools.tasks import Task @@ -19,24 +19,23 @@ def test_bayesian_agent_tes_shadow(RE, db, shadow_tes_simulation): kb_dofs = [kbv.x_rot, kbv.offz] kb_bounds = np.array([[-0.10, +0.10], [-0.50, +0.50]]) - for dof in kb_dofs: - dof.kind = "hinted" - beam_flux_task = Task(key="flux", kind="max", transform=lambda x: np.log(x)) beam_width_task = Task(key="x_width", kind="min", transform=lambda x: np.log(x)) beam_height_task = Task(key="y_width", kind="min", transform=lambda x: np.log(x)) - boa = Agent( - dofs=kb_dofs, - bounds=kb_bounds, + agent = bloptools.bayesian.Agent( + active_dofs=kb_dofs, + passive_dofs=[], detectors=[w9], + active_dof_bounds=kb_bounds, tasks=[beam_flux_task, beam_width_task, beam_height_task], digestion=tes.digestion, db=db, ) - RE(boa.initialize(init_scheme="quasi-random", n_init=4)) + RE(agent.initialize(init_scheme="quasi-random", n_init=4)) - RE(boa.learn(strategy="esti", n_iter=2, n_per_iter=2)) + RE(agent.learn(strategy="ei", n_iter=2)) + RE(agent.learn(strategy="pi", n_iter=2)) - boa.plot_tasks() + agent.plot_tasks() diff --git a/bloptools/tests/test_bayesian_test_funcs.py b/bloptools/tests/test_bayesian_test_funcs.py index 31cff3d..2eff571 100644 --- a/bloptools/tests/test_bayesian_test_funcs.py +++ b/bloptools/tests/test_bayesian_test_funcs.py @@ -1,48 +1,50 @@ import pytest import bloptools -from bloptools.experiments.tests import himmelblau_digestion, mock_kbs_digestion from bloptools.tasks import Task +from bloptools.test_functions import himmelblau_digestion, mock_kbs_digestion @pytest.mark.test_func def test_bayesian_agent_himmelblau(RE, db): - dofs = bloptools.experiments.tests.get_dummy_dofs(n=2) # get a list of two DOFs + dofs = bloptools.devices.dummy_dofs(n=2) # get a list of two DOFs bounds = [(-5.0, +5.0), (-5.0, +5.0)] task = Task(key="himmelblau", kind="min") - boa = bloptools.bayesian.Agent( - dofs=dofs, - bounds=bounds, - tasks=task, + agent = bloptools.bayesian.Agent( + active_dofs=dofs, + passive_dofs=[], + active_dof_bounds=bounds, + tasks=[task], digestion=himmelblau_digestion, db=db, ) - RE(boa.initialize(init_scheme="quasi-random", n_init=16)) + RE(agent.initialize(init_scheme="quasi-random", n_init=16)) - RE(boa.learn(strategy="esti", n_iter=2, n_per_iter=3)) + RE(agent.learn(strategy="ei", n_iter=2)) - boa.plot_tasks() + agent.plot_tasks() @pytest.mark.test_func def test_bayesian_agent_mock_kbs(RE, db): - dofs = bloptools.experiments.tests.get_dummy_dofs(n=4) # get a list of two DOFs + dofs = bloptools.devices.dummy_dofs(n=4) # get a list of two DOFs bounds = [(-4.0, +4.0), (-4.0, +4.0), (-4.0, +4.0), (-4.0, +4.0)] tasks = [Task(key="x_width", kind="min"), Task(key="y_width", kind="min")] - boa = bloptools.bayesian.Agent( - dofs=dofs, - bounds=bounds, + agent = bloptools.bayesian.Agent( + active_dofs=dofs, + passive_dofs=[], + active_dof_bounds=bounds, tasks=tasks, digestion=mock_kbs_digestion, db=db, ) - RE(boa.initialize(init_scheme="quasi-random", n_init=16)) + RE(agent.initialize(init_scheme="quasi-random", n_init=16)) - RE(boa.learn(strategy="esti", n_iter=2, n_per_iter=3)) + RE(agent.learn(strategy="ei", n_iter=4)) - boa.plot_tasks() + agent.plot_tasks() diff --git a/docs/source/tutorials/custom-acquisition.ipynb b/docs/source/tutorials/custom-acquisition.ipynb new file mode 100644 index 0000000..61e4752 --- /dev/null +++ b/docs/source/tutorials/custom-acquisition.ipynb @@ -0,0 +1,114 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "id": "e7b5e13a-c059-441d-8d4f-fff080d52054", + "metadata": {}, + "source": [ + "# Custom acquisition plans\n", + "\n", + "The simplest acqusition plan for a beamline is to move some motor inputs to some positions, and then trigger some detectors. Often, though, we want to investigate behaviors more complex than this. Consider the example of aligning the spectrometer at the Inner-Shell Spectroscopy (ISS) beamline at NSLS-II, which operates by modulating the energy of a beam on a sample and watching the resulting flux rise, peak, and then fall. \n", + "\n", + "We can build a toy model of this spectrometer using custom acquisition and digestion functions. Let's pretend that we can vary two inputs $\\mathbf{x} = (x_1, x_2)$, and that the resolution of the spectrometer is equal to \n", + "\n", + "$\\nu_\\sigma = \\big (1 + x_1^2 + (x_2 - 1)^2 \\big)^{1/2}$\n", + "\n", + "which has a single minimum at $(0,1)$. We can't sample this value directly, rather we can sample how the resulting flux peaks as we vary the energy $\\nu$. Let's pretend that it looks like a Gaussian:\n", + "\n", + "$I(\\nu) = \\exp \\Big [ - 0.5 (\\nu - \\nu_0)^2 / \\nu_\\sigma^2 \\Big ]$\n", + "\n", + "To find the inputs that lead to the tightest spectrum, we need to vary $\\mathbf{x}$, scan over $\\nu$, and then estimate the resolution for the agent to optimize over. Let's write acquisition and digestion functions to do this: " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "92c202e2", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "def acquisition(dofs, inputs, dets):\n", + "\n", + " uid = yield from bp.list_scan\n", + "\n", + " for x in inputs:\n", + "\n", + " res = np.sqrt(1 + np.square(x).sum()) # our resolution is \n", + "\n", + " nu_sigma = np.sqrt(1 + x1 ** 2 + (x2 - 1) ** 2)\n", + "\n", + " flux = np.exp(-0.5*np.square((nu-nu_0)/nu_sigma))\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0ba0bb16", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "from matplotlib import pyplot as plt\n", + "\n", + "\n", + "nu_0 = 100\n", + "\n", + "nu = np.linspace(90, 110, 256)\n", + "\n", + "for x1, x2 in [(0,0), (-2, 2), (-1,0), (0, 1)]:\n", + "\n", + " nu_sigma = np.sqrt(1 + x1 ** 2 + (x2 - 1) ** 2)\n", + "\n", + " flux = np.exp(-0.5*np.square((nu-nu_0)/nu_sigma))\n", + "\n", + " plt.plot(nu, flux, label=f\"(x1, x2) = ({x1}, {x2})\")\n", + "\n", + "plt.legend()\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "1075cb6a", + "metadata": {}, + "source": [ + "To find the inputs that lead to the tightest spectrum, we need to vary $\\mathbf{x}$, scan over $\\nu$, and then estimate the resolution for the agent to optimize over. Let's write acquisition and digestion functions to do this: " + ] + } + ], + "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.9.16 (main, Mar 8 2023, 14:00:05) \n[GCC 11.2.0]" + }, + "vscode": { + "interpreter": { + "hash": "9aced674e98d511b4f654e147532c84d38dc986fe042b1e92785fb9d8df41f75" + } + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/source/tutorials/introduction copy.ipynb b/docs/source/tutorials/introduction copy 2.ipynb similarity index 100% rename from docs/source/tutorials/introduction copy.ipynb rename to docs/source/tutorials/introduction copy 2.ipynb diff --git a/docs/source/tutorials/multi-task-sirepo.ipynb b/docs/source/tutorials/multi-task-sirepo.ipynb index 8a3478a..3986224 100644 --- a/docs/source/tutorials/multi-task-sirepo.ipynb +++ b/docs/source/tutorials/multi-task-sirepo.ipynb @@ -26,10 +26,7 @@ "%run -i ../../../examples/prepare_tes_shadow.py\n", "\n", "kb_dofs = [kbv.x_rot, kbh.x_rot]\n", - "kb_bounds = np.array([[-0.10, +0.10], [-0.10, +0.10]]) \n", - "\n", - "for dof in kb_dofs:\n", - " dof.kind = \"hinted\"" + "kb_bounds = np.array([[-0.10, +0.10], [-0.10, +0.10]]) " ] }, { @@ -51,15 +48,15 @@ "beam_height_task = Task(key=\"y_width\", kind=\"min\", transform=lambda x: np.log(x))\n", "\n", "boa = Agent(\n", - " dofs=kb_dofs, \n", - " bounds=kb_bounds,\n", + " active_dofs=kb_dofs, \n", + " active_dof_bounds=kb_bounds,\n", " detectors=[w9],\n", " tasks=[beam_flux_task, beam_width_task, beam_height_task],\n", " digestion=digestion,\n", " db=db,\n", ")\n", "\n", - "RE(boa.initialize(init_scheme='quasi-random', n_init=16))" + "RE(boa.initialize(init_scheme='quasi-random', n_init=2))" ] }, { diff --git a/docs/source/tutorials/passive-dofs.ipynb b/docs/source/tutorials/passive-dofs.ipynb new file mode 100644 index 0000000..e92011f --- /dev/null +++ b/docs/source/tutorials/passive-dofs.ipynb @@ -0,0 +1,335 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "id": "e7b5e13a-c059-441d-8d4f-fff080d52054", + "metadata": {}, + "source": [ + "## Introduction\n", + "\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "c18ef717", + "metadata": {}, + "source": [ + "This tutorial is an introduction to the syntax used by the optimizer, as well as the principles of Bayesian optimization in general.\n", + "\n", + "We'll start by minimizing Himmelblau's function, which looks like this:" + ] + }, + { + "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", + "\n", + "import bloptools\n", + "\n", + "x1 = x2 = np.linspace(-5.0, 5.0, 256)\n", + "X1, X2 = np.meshgrid(x1, x2)\n", + "\n", + "plt.pcolormesh(x1, \n", + " x2, \n", + " bloptools.experiments.tests.himmelblau(X1, X2), \n", + " norm=mpl.colors.LogNorm(), \n", + " shading=\"auto\")\n", + "\n", + "plt.xlabel(\"x1\")\n", + "plt.ylabel(\"x2\")\n", + "plt.colorbar()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f46924af", + "metadata": {}, + "outputs": [], + "source": [ + "from bloptools.objects import TimeReadback\n", + "\n", + "tr = TimeReadback(name=\"timestamp\")\n", + "\n", + "tr.read()\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "ecef8da5", + "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": "4c870567", + "metadata": {}, + "outputs": [], + "source": [ + "import bloptools\n", + "\n", + "dofs = bloptools.objects.get_dummy_dofs(n=2) # get a list of two DOFs\n", + "bounds = [(-5.0, +5.0), (-5.0, +5.0)]\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "7a88c7bd", + "metadata": {}, + "source": [ + "The agent automatically samples at different inputs, but we often need some post-processing after data collection. In this case, we need to give the agent a way to compute Himmelblau's function. We accomplish this with a digestion function:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e6bfcf73", + "metadata": {}, + "outputs": [], + "source": [ + "import bloptools\n", + "\n", + "dofs = bloptools.objects.get_dummy_dofs(n=2) # get a list of two DOFs\n", + "bounds = [(-5.0, +5.0), (-5.0, +5.0)]\n", + "\n", + "from bloptools.objects import TimeReadback\n", + "\n", + "tr = TimeReadback(name=\"timestamp\")\n", + "\n", + "tr.read()\n", + "\n", + "\n", + "\n", + "def digestion(db, uid):\n", + "\n", + " table = db[uid].table()\n", + " products = pd.DataFrame()\n", + "\n", + " for index, entry in table.iterrows():\n", + "\n", + " products.loc[index, \"himmelblau\"] = bloptools.test_functions.himmelblau(entry.x1, entry.x2)\n", + "\n", + " return products\n", + "\n", + "from bloptools.tasks import Task\n", + "\n", + "task = Task(key=\"himmelblau\", kind=\"min\")" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "4532b087", + "metadata": {}, + "source": [ + "The next ingredient is a task, which gives the agent something to do. We want it to minimize Himmelblau's function, so we make a task that will try to minimize the output of the digestion function called \"himmelblau\". We also include a transform function, which will make it easier to regress over the outputs of the function." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4c14d162", + "metadata": {}, + "outputs": [], + "source": [ + "from bloptools.tasks import Task\n", + "\n", + "task = Task(key=\"himmelblau\", kind=\"min\")" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "0d3d91c3", + "metadata": {}, + "source": [ + "Combining all of these with a databroker instance, we can make an agent:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "071a829f-a390-40dc-9d5b-ae75702e119e", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "%run -i ../../../examples/prepare_bluesky.py # prepare the bluesky environment\n", + "\n", + "boa = bloptools.bayesian.Agent(\n", + " dofs=dofs,\n", + " bounds=bounds,\n", + " passive_dims=[tr],\n", + " tasks=task,\n", + " digestion=digestion,\n", + " db=db, \n", + " )\n", + "\n", + "RE(boa.initialize(init_scheme='quasi-random', n_init=32))" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "d8f2da43", + "metadata": {}, + "source": [ + "We initialized the GP with the \"quasi-random\" strategy, as it doesn't require any prior data. We can view the state of the optimizer's posterior of the tasks over the input parameters:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5818143a", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e93c75a1", + "metadata": {}, + "outputs": [], + "source": [ + "np.atleast_1d([]).size" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "996c3c01-f91d-4a25-9b8d-eba5fa964504", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "boa.plot_tasks()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "d2eb855c", + "metadata": {}, + "source": [ + "We want to learn a bit more, so we can ask the agent where to sample based off of some strategy. Here we use the \"esti\" strategy, which maximizes the expected sum of tasks improvement." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f0e74651", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "RE(boa.learn(strategy='esti', n_iter=4))\n", + "boa.plot_tasks()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "aeab7a9b", + "metadata": {}, + "source": [ + "The agent has updated its model of the tasks, including refitting the hyperparameters. Continuing:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d6b39b54", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "RE(boa.learn(strategy='esti', n_iter=16))\n", + "boa.plot_tasks()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "e955233f", + "metadata": {}, + "source": [ + "Eventually, we reach a point of saturation where no more improvement takes place:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d73e4fd5", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "RE(boa.learn(strategy='esti', n_iter=32))\n", + "boa.plot_tasks()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5ad4b1e7", + "metadata": {}, + "outputs": [], + "source": [ + "boa.tasks[0].regressor.covar_module.base_kernel.trans_matrix" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "96e9abac", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "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.9.16" + }, + "vscode": { + "interpreter": { + "hash": "9aced674e98d511b4f654e147532c84d38dc986fe042b1e92785fb9d8df41f75" + } + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From bbb8e68abf08e070acf90f91efa9b51dd590da41 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 29 Jun 2023 16:42:36 -0400 Subject: [PATCH 03/11] fixed docs --- bloptools/bayesian/__init__.py | 524 +++++++++--------- bloptools/test_functions.py | 5 + docs/source/tutorials.rst | 1 + docs/source/tutorials/himmelblau.ipynb | 218 ++++++++ .../tutorials/introduction copy 2.ipynb | 302 ---------- docs/source/tutorials/introduction.ipynb | 137 ++--- 6 files changed, 521 insertions(+), 666 deletions(-) create mode 100644 docs/source/tutorials/himmelblau.ipynb delete mode 100644 docs/source/tutorials/introduction copy 2.ipynb diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index 793d7fc..5170334 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -8,7 +8,6 @@ import pandas as pd import scipy as sp import torch -from botorch.acquisition.analytic import LogExpectedImprovement, LogProbabilityOfImprovement from matplotlib import pyplot as plt from .. import utils @@ -38,7 +37,7 @@ def default_digestion_plan(db, uid): class Agent: - MAX_TEST_POINTS = 2**10 + MAX_TEST_POINTS = 2**11 def __init__( self, @@ -275,125 +274,39 @@ def tell(self, new_table=None, append=True, **kwargs): self.task_model = botorch.models.model.ModelList(*[task.regressor for task in self.tasks], self.feas_model) - # def sample_acqf(self, acqf, n_test=256, optimize=False): - # """ - # Return the point that optimizes a given acqusition function. - # """ - - # def acq_loss(x, *args): - # return -acqf(x, *args) - - # acq_args = (self,) - - # test_inputs = self.unnormalize_active_inputs(self.sampler(n=n_test)) - # init_inputs = test_inputs[acq_loss(test_inputs, *acq_args).argmin()] - - # if optimize: - # res = sp.optimize.minimize( - # fun=acq_loss, - # args=acq_args, - # x0=init_inputs, - # bounds=self.active_dof_bounds, - # method="SLSQP", - # options={"maxiter": 256}, - # ) - # acq_inputs = res.x - # else: - # acq_inputs = init_inputs - - # return acq_inputs, acq_loss(acq_inputs, *acq_args) - - # def ask( - # self, - # tasks=None, - # classifier=None, - # strategy=None, - # greedy=True, - # n=1, - # disappointment=0, - # route=True, - # cost_model=None, - # n_test=1024, - # optimize=True, - # ): - # """ - # The next $n$ points to sample, recommended by the agent. - # """ - - # if route: - # unrouted_points = self.ask( - # tasks=tasks, - # classifier=classifier, - # strategy=strategy, - # greedy=greedy, - # n=n, - # disappointment=disappointment, - # route=False, - # cost_model=cost_model, - # n_test=n_test, - # ) - - # routing_index, _ = utils.get_routing(self.read_active_dofs, unrouted_points) - # return unrouted_points[routing_index] - - # if strategy.lower() == "quasi-random": - # return self.unnormalize_active_inputs(self.sampler(n=n)) - - # if not self._initialized: - # raise RuntimeError("The agent is not initialized!") - - # if tasks is None: - # tasks = self.tasks - - # if classifier is None: - # classifier = self.feas_model - - # if (not greedy) or (n == 1): - # acqf = None - - # if strategy.lower() == "ei": # maximize the expected improvement - - # acq_func = LogExpectedImprovement(agent.task_model, - # best_f=agent.best_sum_of_tasks, - # posterior_transform=agent.scalarization) - - # acqf = self.acquisition.log_expected_improvement - - # if strategy.lower() == "pi": # maximize the probability improvement - # acqf = self.acquisition.log_probability_of_improvement - - # if acqf is None: - # raise ValueError(f'Unrecognized strategy "{strategy}".') - - # inputs_to_sample, loss = self.sample_acqf(acqf, optimize=optimize) - - # return np.atleast_2d(inputs_to_sample) - - # if greedy and (n > 1): - # dummy_tasks = [copy.deepcopy(task) for task in self.tasks] - - # inputs_to_sample = np.zeros((0, self.n_dof)) - - # for i in range(n): - - # new_input = self.ask(strategy=strategy, tasks=dummy_tasks, classifier=classifier, greedy=True, n=1) - - # inputs_to_sample = np.r_[inputs_to_sample, new_input] - - # if not (i + 1 == n): # no point if we're on the last iteration - # fantasy_model = botorch.models.model.ModelList(*[task.regressor for task in self.tasks]) - # fantasy_posterior = fantasy_model.posterior( - # torch.tensor(self.normalize_inputs(new_input)).double() - # ) - - # fantasy_targets = fantasy_posterior.mean - disappointment * fantasy_posterior.variance.sqrt() - # new_targets = self.unnormalize_targets(fantasy_targets.detach().numpy()) - - # self.tell(new_table=new_table) + def get_acquisition_function(self, strategy="ei", return_metadata=False, acqf_args={}, **kwargs): + if strategy.lower() == "ei": + acqf = botorch.acquisition.analytic.LogExpectedImprovement( + self.task_model, + best_f=self.best_sum_of_tasks, + posterior_transform=self.task_scalarization, + **kwargs, + ) + acqf_meta = {"name": "Expected Improvement", "args": {}} + + elif strategy.lower() == "pi": + acqf = botorch.acquisition.analytic.LogProbabilityOfImprovement( + self.task_model, + best_f=self.best_sum_of_tasks, + posterior_transform=self.task_scalarization, + **kwargs, + ) + acqf_meta = {"name": "Probability of Improvement", "args": {}} + + elif strategy.lower() == "ucb": + beta = acqf_args.get("beta", 0.1) + acqf = botorch.acquisition.analytic.UpperConfidenceBound( + self.task_model, + beta=beta, + posterior_transform=self.task_scalarization, + **kwargs, + ) + acqf_meta = {"name": "Upper Confidence Bound", "args": {"beta": beta}} - # self.forget(index=self.table.index[-n:]) # forget what you saw here + else: + raise ValueError(f'Unrecognized acquisition strategy "{strategy}".') - # return np.atleast_2d(inputs_to_sample) + return (acqf, acqf_meta) if return_metadata else acqf def ask( self, @@ -434,20 +347,7 @@ def ask( if not self._initialized: raise RuntimeError("The agent is not initialized!") - self.acqf = None - - if strategy.lower() == "ei": # maximize the expected improvement - self.acqf = LogExpectedImprovement( - self.task_model, best_f=self.best_sum_of_tasks, posterior_transform=self.task_scalarization - ) - - if strategy.lower() == "pi": # maximize the expected improvement - self.acqf = LogProbabilityOfImprovement( - self.task_model, best_f=self.best_sum_of_tasks, posterior_transform=self.task_scalarization - ) - - if self.acqf is None: - raise ValueError(f'Unrecognized strategy "{strategy}".') + self.acqf = self.get_acquisition_function(strategy=strategy) BATCH_SIZE = 1 NUM_RESTARTS = 7 @@ -455,14 +355,14 @@ def ask( candidates, _ = botorch.optim.optimize_acqf( acq_function=self.acqf, - bounds=torch.tensor(self.active_dof_bounds).T, + bounds=torch.tensor(self.dof_bounds).T, q=BATCH_SIZE, num_restarts=NUM_RESTARTS, raw_samples=RAW_SAMPLES, # used for intialization heuristic options={"batch_limit": 5, "maxiter": 200}, ) - return candidates.detach().numpy() + return candidates.detach().numpy()[..., self.dof_is_active_mask] def acquire(self, dof_inputs): """ @@ -508,21 +408,180 @@ def learn( self.tell(new_table=new_table, reuse_hypers=reuse_hypers) + def normalize_active_inputs(self, inputs): + return (inputs - self.active_dof_bounds.min(axis=1)) / self.active_dof_bounds.ptp(axis=1) + + def unnormalize_active_inputs(self, X): + return X * self.active_dof_bounds.ptp(axis=1) + self.active_dof_bounds.min(axis=1) + + def normalize_inputs(self, inputs): + return (inputs - self.input_bounds.min(axis=1)) / self.input_bounds.ptp(axis=1) + + def unnormalize_inputs(self, X): + return X * self.input_bounds.ptp(axis=1) + self.input_bounds.min(axis=1) + + def normalize_targets(self, targets): + return (targets - self.targets_mean) / (1e-20 + self.targets_scale) + + def unnormalize_targets(self, targets): + return targets * self.targets_scale + self.targets_mean + + @property + def test_inputs(self): + test_passive_inputs = ( + self.passive_inputs.values[-1][None] * np.ones(len(self.test_active_inputs))[..., None] + ) + return np.concatenate([self.test_active_inputs, test_passive_inputs], axis=-1) + + @property + def test_inputs_grid(self): + test_passive_inputs_grid = self.passive_inputs.values[-1] * np.ones( + (*self.test_active_inputs_grid.shape[:-1], self.n_passive_dof) + ) + return np.concatenate([self.test_active_inputs_grid, test_passive_inputs_grid], axis=-1) + + @property + def inputs(self): + return self.table.loc[:, self.dof_names].astype(float) + + @property + def active_inputs(self): + return self.inputs.loc[:, self.active_dof_names] + + @property + def passive_inputs(self): + return self.inputs.loc[:, self.passive_dof_names] + + @property + def targets(self): + return self.table.loc[:, self.task_names].astype(float) + + @property + def feasible(self): + with pd.option_context("mode.use_inf_as_null", True): + feasible = ~self.targets.isna() + return feasible + + # @property + # def input_bounds(self): + # lower_bound = np.r_[ + # self.active_dof_bounds[:, 0], np.nanmin(self.passive_inputs.astype(float).values, axis=0) + # ] + # upper_bound = np.r_[ + # self.active_dof_bounds[:, 1], np.nanmax(self.passive_inputs.astype(float).values, axis=0) + # ] + # return np.c_[lower_bound, upper_bound] + + @property + def targets_mean(self): + return np.nanmean(self.targets, axis=0) + + @property + def targets_scale(self): + return np.nanstd(self.targets, axis=0) + + @property + def normalized_targets(self): + return self.normalize_targets(self.targets) + + @property + def latest_passive_dof_values(self): + passive_inputs = self.passive_inputs + return [passive_inputs.loc[passive_inputs.last_valid_index(), col] for col in passive_inputs.columns] + + @property + def passive_dof_bounds(self): + # food for thought: should this be the current values, or the latest recorded values? + # the former leads to weird extrapolation (especially for time), and the latter to some latency. + # let's go with the first way for now + return np.outer(self.latest_passive_dof_values, [1.0, 1.0]) + + @property + def dof_is_active_mask(self): + return np.r_[np.ones(self.n_active_dof), np.zeros(self.n_passive_dof)].astype(bool) + + @property + def dof_bounds(self): + return np.r_[self.active_dof_bounds, self.passive_dof_bounds] + + @property + def read_active_dofs(self): + return np.array([dof.read()[dof.name]["value"] for dof in self.active_dofs]) + + @property + def read_passive_dofs(self): + return np.array([dof.read()[dof.name]["value"] for dof in self.passive_dofs]) + + @property + def read_dofs(self): + return np.r_[self.read_active_dofs, self.read_passive_dofs] + + @property + def active_dof_names(self): + return [dof.name for dof in self.active_dofs] + + @property + def passive_dof_names(self): + return [dof.name for dof in self.passive_dofs] + + @property + def dof_names(self): + return [dof.name for dof in self.dofs] + + @property + def det_names(self): + return [det.name for det in self.dets] + + @property + def target_names(self): + return [task.name for task in self.tasks] + + @property + def task_names(self): + return [task.name for task in self.tasks] + + @property + def task_weights(self): + return np.array([task.weight for task in self.tasks]) + + @property + def best_sum_of_tasks(self): + return self.targets.fillna(-np.inf).sum(axis=1).max() + + @property + def best_sum_of_tasks_inputs(self): + return self.inputs[np.nanargmax(self.targets.sum(axis=1))] + + @property + def go_to(self, inputs): + yield from bps.mv(*[_ for items in zip(self.dofs, np.atleast_1d(inputs).T) for _ in items]) + + @property + def go_to_best_sum_of_tasks(self): + yield from self.go_to(self.best_sum_of_tasks_inputs) + def plot_tasks(self, **kwargs): - if self.n_dof == 1: + if self.n_active_dof == 1: self._plot_tasks_one_dof(**kwargs) else: self._plot_tasks_many_dofs(**kwargs) def plot_feasibility(self, **kwargs): - if self.n_dof == 1: - self._plot_feasibility_one_dof(**kwargs) + if self.n_active_dof == 1: + self._plot_feas_one_dof(**kwargs) else: - self._plot_feasibility_many_dofs(**kwargs) + self._plot_feas_many_dofs(**kwargs) - def _plot_feasibility_one_dof(self, size=32): + def plot_acquisition(self, **kwargs): + if self.n_active_dof == 1: + self._plot_acq_one_dof(**kwargs) + + else: + self._plot_acq_many_dofs(**kwargs) + + def _plot_feas_one_dof(self, size=32): self.class_fig, self.class_ax = plt.subplots(1, 1, figsize=(4, 4), sharex=True, constrained_layout=True) self.class_ax.scatter(self.inputs.values, self.all_targets_valid.astype(int), s=size) @@ -534,9 +593,7 @@ def _plot_feasibility_one_dof(self, size=32): self.class_ax.set_xlim(*self.active_dof_bounds[0]) - def _plot_feasibility_many_dofs( - self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, size=32, gridded=None - ): + def _plot_feas_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, size=32, gridded=None): if gridded is None: gridded = self.n_dof == 2 @@ -568,7 +625,7 @@ def _plot_feasibility_many_dofs( ) else: - x = torch.tensor(self.normalize_inputs(self.test_inputs)).double() + x = torch.tensor(self.test_inputs).double() log_prob = self.dirichlet_classifier.log_prob(x).detach().numpy() self.class_axes[1].scatter( @@ -597,11 +654,9 @@ def _plot_tasks_one_dof(self, size=32, lw=1e0): self.task_axes[itask].set_ylabel(task.name) - task_posterior = task.regressor.posterior( - torch.tensor(self.normalize_inputs(self.test_inputs_grid)).double() - ) - task_mean = task_posterior.mean.detach().numpy() * task.targets_scale + task.targets_mean - task_sigma = task_posterior.variance.sqrt().detach().numpy() * task.targets_scale + task_posterior = task.regressor.posterior(torch.tensor(self.test_inputs_grid).double()) + task_mean = task_posterior.mean.detach().numpy().ravel() + task_sigma = task_posterior.variance.sqrt().detach().numpy().ravel() self.task_axes[itask].scatter(self.inputs.values, task.targets, s=size, color=color) self.task_axes[itask].plot(self.test_active_inputs_grid.ravel(), task_mean, lw=lw, color=color) @@ -667,7 +722,7 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL ) sigma_ax = self.task_axes[itask, 2].pcolormesh( *np.swapaxes(self.test_inputs_grid, 0, -1), - task_sigma.reshape(self.normalize_inputs(self.test_inputs_grid).shape[:-1]).T, + task_sigma.reshape(self.test_inputs_grid.shape[:-1]).T, shading=shading, cmap=cmap, ) @@ -687,137 +742,82 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL ax.set_xlim(*self.active_dof_bounds[axes[0]]) ax.set_ylim(*self.active_dof_bounds[axes[1]]) - def normalize_active_inputs(self, inputs): - return (inputs - self.active_dof_bounds.min(axis=1)) / self.active_dof_bounds.ptp(axis=1) - - def unnormalize_active_inputs(self, X): - return X * self.active_dof_bounds.ptp(axis=1) + self.active_dof_bounds.min(axis=1) - - def normalize_inputs(self, inputs): - return (inputs - self.input_bounds.min(axis=1)) / self.input_bounds.ptp(axis=1) - - def unnormalize_inputs(self, X): - return X * self.input_bounds.ptp(axis=1) + self.input_bounds.min(axis=1) - - def normalize_targets(self, targets): - return (targets - self.targets_mean) / (1e-20 + self.targets_scale) - - def unnormalize_targets(self, targets): - return targets * self.targets_scale + self.targets_mean + def _plot_acq_one_dof(self, size=32, lw=1e0, **kwargs): + strategies = np.atleast_1d(kwargs.get("strategy", "ei")) - @property - def test_inputs(self): - test_passive_inputs = ( - self.passive_inputs.values[-1][None] * np.ones(len(self.test_active_inputs))[..., None] - ) - return np.concatenate([self.test_active_inputs, test_passive_inputs], axis=-1) - - @property - def test_inputs_grid(self): - test_passive_inputs_grid = self.passive_inputs.values[-1] * np.ones( - (*self.test_active_inputs_grid.shape[:-1], self.n_passive_dof) + self.acq_fig, self.acq_axes = plt.subplots( + 1, + len(strategies), + figsize=(6 * len(strategies), 6), + sharex=True, + constrained_layout=True, ) - return np.concatenate([self.test_active_inputs_grid, test_passive_inputs_grid], axis=-1) - - @property - def inputs(self): - return self.table.loc[:, self.dof_names].astype(float) - - @property - def active_inputs(self): - return self.inputs.loc[:, self.active_dof_names] - - @property - def passive_inputs(self): - return self.inputs.loc[:, self.passive_dof_names] - - @property - def targets(self): - return self.table.loc[:, self.task_names].astype(float) - - @property - def feasible(self): - with pd.option_context("mode.use_inf_as_null", True): - feasible = ~self.targets.isna() - return feasible - @property - def input_bounds(self): - lower_bound = np.r_[ - self.active_dof_bounds[:, 0], np.nanmin(self.passive_inputs.astype(float).values, axis=0) - ] - upper_bound = np.r_[ - self.active_dof_bounds[:, 1], np.nanmax(self.passive_inputs.astype(float).values, axis=0) - ] - return np.c_[lower_bound, upper_bound] + self.acq_axes = np.atleast_1d(self.acq_axes) - @property - def targets_mean(self): - return np.nanmean(self.targets, axis=0) + for istrat, strategy in enumerate(strategies): + color = DEFAULT_COLOR_LIST[0] - @property - def targets_scale(self): - return np.nanstd(self.targets, axis=0) + acqf, acqf_meta = self.get_acquisition_function(strategy, return_metadata=True) - @property - def normalized_targets(self): - return self.normalize_targets(self.targets) + *grid_shape, dim = self.test_inputs_grid.shape + x = torch.tensor(self.test_inputs_grid.reshape(-1, 1, dim)).double() + obj = acqf.forward(x) - @property - def read_active_dofs(self): - return np.array([dof.read()[dof.name]["value"] for dof in self.active_dofs]) + if strategy in ["ei", "pi"]: + obj = obj.exp() - @property - def read_passive_dofs(self): - return np.array([dof.read()[dof.name]["value"] for dof in self.passive_dofs]) - - @property - def read_dofs(self): - return np.r_[self.read_active_dofs, self.read_passive_dofs] + self.acq_axes[istrat].set_title(acqf_meta["name"]) + self.acq_axes[istrat].plot( + self.test_active_inputs_grid.ravel(), obj.detach().numpy().ravel(), lw=lw, color=color + ) - @property - def active_dof_names(self): - return [dof.name for dof in self.active_dofs] + self.acq_axes[istrat].set_xlim(*self.active_dof_bounds[0]) - @property - def passive_dof_names(self): - return [dof.name for dof in self.passive_dofs] + def _plot_acq_many_dofs( + self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=32, **kwargs + ): + strategies = np.atleast_1d(kwargs.get("strategy", "ei")) - @property - def dof_names(self): - return [dof.name for dof in self.dofs] + self.acq_fig, self.acq_axes = plt.subplots( + 1, + len(strategies), + figsize=(4 * len(strategies), 5), + sharex=True, + sharey=True, + constrained_layout=True, + ) - @property - def det_names(self): - return [det.name for det in self.dets] + if gridded is None: + gridded = self.n_active_dof == 2 - @property - def target_names(self): - return [task.name for task in self.tasks] + self.acq_axes = np.atleast_1d(self.acq_axes) + self.acq_fig.suptitle(f"(x,y)=({self.dofs[axes[0]].name},{self.dofs[axes[1]].name})") - @property - def task_names(self): - return [task.name for task in self.tasks] + for istrat, strategy in enumerate(strategies): + acqf, acqf_meta = self.get_acquisition_function(strategy, return_metadata=True) - @property - def task_weights(self): - return np.array([task.weight for task in self.tasks]) + if gridded: + *grid_shape, dim = self.test_inputs_grid.shape + x = torch.tensor(self.test_inputs_grid.reshape(-1, 1, dim)).double() + obj = acqf.forward(x) - @property - def best_sum_of_tasks(self): - return self.targets.fillna(-np.inf).sum(axis=1).max() + if strategy in ["ei", "pi"]: + obj = obj.exp() - @property - def best_sum_of_tasks_inputs(self): - return self.inputs[np.nanargmax(self.targets.sum(axis=1))] + self.acq_axes[istrat].set_title(acqf_meta["name"]) + obj_ax = self.acq_axes[istrat].pcolormesh( + *np.swapaxes(self.test_inputs_grid, 0, -1), + obj.detach().numpy().reshape(grid_shape).T, + shading=shading, + cmap=cmap, + ) - @property - def go_to(self, inputs): - yield from bps.mv(*[_ for items in zip(self.dofs, np.atleast_1d(inputs).T) for _ in items]) + self.acq_fig.colorbar(obj_ax, ax=self.acq_axes[istrat], location="bottom", aspect=32, shrink=0.8) - @property - def go_to_best_sum_of_tasks(self): - yield from self.go_to(self.best_sum_of_tasks_inputs) + for ax in self.acq_axes.ravel(): + ax.set_xlim(*self.active_dof_bounds[axes[0]]) + ax.set_ylim(*self.active_dof_bounds[axes[1]]) def inspect_beam(self, index, border=None): im = self.images[index] diff --git a/bloptools/test_functions.py b/bloptools/test_functions.py index e8cd621..195f7d5 100644 --- a/bloptools/test_functions.py +++ b/bloptools/test_functions.py @@ -9,6 +9,11 @@ def himmelblau(x1, x2): return (x1**2 + x2 - 11) ** 2 + (x1 + x2**2 - 7) ** 2 +def rastrigin(*x): + X = np.c_[x] + return 10 * X.shape[-1] + (X**2 - 10 * np.cos(2 * np.pi * X)).sum(axis=1) + + def gaussian_beam_waist(x1, x2): return np.sqrt(1 + 0.25 * (x1 - x2) ** 2 + 16 * (x1 + x2 - 2) ** 2) diff --git a/docs/source/tutorials.rst b/docs/source/tutorials.rst index 53df520..fb721bd 100644 --- a/docs/source/tutorials.rst +++ b/docs/source/tutorials.rst @@ -5,4 +5,5 @@ Tutorials :maxdepth: 2 tutorials/introduction.ipynb + tutorials/himmelblau.ipynb tutorials/multi-task-sirepo.ipynb diff --git a/docs/source/tutorials/himmelblau.ipynb b/docs/source/tutorials/himmelblau.ipynb new file mode 100644 index 0000000..0b2244e --- /dev/null +++ b/docs/source/tutorials/himmelblau.ipynb @@ -0,0 +1,218 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "id": "e7b5e13a-c059-441d-8d4f-fff080d52054", + "metadata": {}, + "source": [ + "# Feasibility modeling\n", + "\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "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 < 36$" + ] + }, + { + "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(-8, 8, 256)\n", + "X1, X2 = np.meshgrid(x1, x2)\n", + "\n", + "F = test_functions.himmelblau(X1, X2)\n", + "F[X1 ** 2 + X2 ** 2 > 50] = np.nan\n", + "\n", + "plt.pcolormesh(x1, x2, F, norm=mpl.colors.LogNorm(), shading=\"auto\")\n", + "plt.colorbar()\n", + "plt.xlabel(\"x1\")\n", + "plt.ylabel(\"x2\")\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "7a88c7bd", + "metadata": {}, + "source": [ + "In our digestion function, we return a `NaN` when we violate the constraint:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e6bfcf73", + "metadata": {}, + "outputs": [], + "source": [ + "def digestion(db, uid):\n", + "\n", + " products = db[uid].table()\n", + "\n", + " for index, entry in products.iterrows():\n", + "\n", + " if entry.x1 ** 2 + entry.x2 ** 2 < 50:\n", + " products.loc[index, \"himmelblau\"] = test_functions.himmelblau(entry.x1, entry.x2)\n", + " else:\n", + " products.loc[index, \"himmelblau\"] = np.nan\n", + "\n", + " return products" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "0d3d91c3", + "metadata": {}, + "source": [ + "and create the agent in the usual way:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "071a829f-a390-40dc-9d5b-ae75702e119e", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "%run -i ../../../examples/prepare_bluesky.py # prepare the bluesky environment\n", + "\n", + "agent = bloptools.bayesian.Agent(\n", + " active_dofs=dofs,\n", + " passive_dofs=[],\n", + " active_dof_bounds=bounds,\n", + " tasks=[task],\n", + " digestion=digestion,\n", + " db=db,\n", + " )\n", + "\n", + "RE(agent.initialize(init_scheme='quasi-random', n_init=16))\n", + "\n", + "agent.plot_tasks()\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "9ab3be01", + "metadata": {}, + "source": [ + "In addition to modeling the fitness of the task, the agent models the probability that an input will be feasible:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "996c3c01-f91d-4a25-9b8d-eba5fa964504", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "agent.plot_feasibility()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "a79b56ac", + "metadata": {}, + "source": [ + "The agent automatically tries to avoid infeasible points, but will end up naturally exploring the boundary of the constraint. Let's see where the agent is thinking of going:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "84a76eb9", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "agent.plot_acquisition(strategy=[\"ei\", \"pi\", \"ucb\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ebc65169", + "metadata": {}, + "outputs": [], + "source": [ + "RE(agent.learn(strategy='ei', n_iter=4))\n", + "agent.plot_tasks()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "a43130e4", + "metadata": {}, + "source": [ + "The agent will naturally explore the whole parameter space" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "79fad8eb", + "metadata": {}, + "outputs": [], + "source": [ + "RE(agent.learn(strategy='ei', n_iter=16))\n", + "agent.plot_tasks()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "40491b5b", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "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.9.16" + }, + "vscode": { + "interpreter": { + "hash": "9aced674e98d511b4f654e147532c84d38dc986fe042b1e92785fb9d8df41f75" + } + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/source/tutorials/introduction copy 2.ipynb b/docs/source/tutorials/introduction copy 2.ipynb deleted file mode 100644 index 90061dc..0000000 --- a/docs/source/tutorials/introduction copy 2.ipynb +++ /dev/null @@ -1,302 +0,0 @@ -{ - "cells": [ - { - "attachments": {}, - "cell_type": "markdown", - "id": "e7b5e13a-c059-441d-8d4f-fff080d52054", - "metadata": {}, - "source": [ - "# Kernels and hyperparameters\n", - "\n" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "c18ef717", - "metadata": {}, - "source": [ - "This tutorial is an introduction to the syntax used by the optimizer, as well as the principles of Bayesian optimization in general.\n", - "\n", - "We'll start by minimizing Booth's function, which looks like this:" - ] - }, - { - "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", - "\n", - "import bloptools\n", - "\n", - "x1 = x2 = np.linspace(-10, 10, 256)\n", - "X1, X2 = np.meshgrid(x1, x2)\n", - "\n", - "plt.pcolormesh(x1, x2, bloptools.experiments.tests.himmelblau(X1, X2), norm=mpl.colors.LogNorm(), shading=\"auto\")\n", - "plt.colorbar()\n", - "plt.xlabel(\"x1\")\n", - "plt.ylabel(\"x2\")\n" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "ecef8da5", - "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": "4c870567", - "metadata": {}, - "outputs": [], - "source": [ - "import bloptools\n", - "\n", - "dofs = bloptools.experiments.tests.get_dummy_dofs(2)\n", - "bounds = [(-10, 10), (-10, 10)]" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "7a88c7bd", - "metadata": {}, - "source": [ - "The agent automatically samples at different inputs, but we often need some post-processing after data collection. In this case, we need to give the agent a way to compute Himmelblau's function. We accomplish this with a digestion function:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "e6bfcf73", - "metadata": {}, - "outputs": [], - "source": [ - "def digestion(db, uid):\n", - "\n", - " table = db[uid].table()\n", - " products = {\"himmelblau\": []}\n", - "\n", - " for index, entry in table.iterrows():\n", - "\n", - " products[\"himmelblau\"].append(bloptools.experiments.tests.himmelblau(entry.x1, entry.x2))\n", - "\n", - " return products" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "dad64303", - "metadata": {}, - "source": [ - "The next ingredient is a task, which gives the agent something to do. We want it to minimize Himmelblau's function, so we make a task that will try to minimize the output of the digestion function called \"himmelblau\". We also include a transform function, which will make it easier to regress over the outputs of the function." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4c14d162", - "metadata": {}, - "outputs": [], - "source": [ - "from bloptools.tasks import Task\n", - "\n", - "task = Task(key=\"himmelblau\", kind=\"min\", transform=lambda x: np.log(1 + 1e-2 * x))" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "0d3d91c3", - "metadata": {}, - "source": [ - "Combining all of these with a databroker instance, we can make an agent:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "071a829f-a390-40dc-9d5b-ae75702e119e", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "%run -i ../../../examples/prepare_bluesky.py # prepare the bluesky environment\n", - "\n", - "boa = bloptools.bayesian.Agent(\n", - " dofs=dofs, # things which we move around\n", - " bounds=bounds, # where we're allowed to move them\n", - " tasks=task, # tasks for the optimizer\n", - " digestion=digestion, # how to process the acquisition into task data\n", - " db=db, # a databroker instance\n", - " )\n", - "\n", - "RE(boa.initialize(init_scheme='quasi-random', n_init=16))" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "d8f2da43", - "metadata": {}, - "source": [ - "We initialized the GP with the \"quasi-random\" strategy, as it doesn't require any prior data. We can view the state of the optimizer's posterior of the tasks over the input parameters:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "996c3c01-f91d-4a25-9b8d-eba5fa964504", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "boa.plot_tasks()" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "4274930c", - "metadata": {}, - "source": [ - "We can also the agent's posterior about the probability of goodness over the parameters:" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "d2eb855c", - "metadata": {}, - "source": [ - "We want to learn a bit more, so we can ask the agent where to sample based off of some strategy. Here we use the \"esti\" strategy, which maximizes the expected sum of tasks improvement.\n", - "\n", - "We can ask the agent to \"route\" them using ``ortools``, so that we can sample them more quickly if it requires us to e.g. move motors." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "334f3c66", - "metadata": {}, - "outputs": [], - "source": [ - "X_to_sample = boa.ask(strategy='esti', n=16, optimize=True, route=True)\n", - "plt.scatter(*X_to_sample.T)\n", - "plt.plot(*X_to_sample.T)" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "296d9fd2", - "metadata": {}, - "source": [ - "Let's tell the agent to learn a bit more (it will direct itself):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "f0e74651", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "RE(boa.learn(strategy='esti', n_iter=1, n_per_iter=4))\n", - "boa.plot_tasks()" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "aeab7a9b", - "metadata": {}, - "source": [ - "The agent has updated its model of the tasks, including refitting the hyperparameters. Continuing:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "d6b39b54", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "RE(boa.learn(strategy='esti', n_iter=4, n_per_iter=4))\n", - "boa.plot_tasks()" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "e955233f", - "metadata": {}, - "source": [ - "Eventually, we reach a point of saturation where no more improvement takes place:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "d73e4fd5", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "RE(boa.learn(strategy='esti', n_iter=4, n_per_iter=4))\n", - "boa.plot_tasks()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5ad4b1e7", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "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.9.16" - }, - "vscode": { - "interpreter": { - "hash": "9aced674e98d511b4f654e147532c84d38dc986fe042b1e92785fb9d8df41f75" - } - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/docs/source/tutorials/introduction.ipynb b/docs/source/tutorials/introduction.ipynb index 90061dc..e5ea4c1 100644 --- a/docs/source/tutorials/introduction.ipynb +++ b/docs/source/tutorials/introduction.ipynb @@ -6,8 +6,7 @@ "id": "e7b5e13a-c059-441d-8d4f-fff080d52054", "metadata": {}, "source": [ - "# Kernels and hyperparameters\n", - "\n" + "# Bayesian optimization" ] }, { @@ -18,7 +17,7 @@ "source": [ "This tutorial is an introduction to the syntax used by the optimizer, as well as the principles of Bayesian optimization in general.\n", "\n", - "We'll start by minimizing Booth's function, which looks like this:" + "We'll start by minimizing the Rastrigin function in one dimension, which looks like this:" ] }, { @@ -32,15 +31,11 @@ "import matplotlib as mpl\n", "from matplotlib import pyplot as plt\n", "\n", - "import bloptools\n", + "from bloptools import test_functions\n", "\n", - "x1 = x2 = np.linspace(-10, 10, 256)\n", - "X1, X2 = np.meshgrid(x1, x2)\n", + "x = np.linspace(-2, 2, 256)\n", "\n", - "plt.pcolormesh(x1, x2, bloptools.experiments.tests.himmelblau(X1, X2), norm=mpl.colors.LogNorm(), shading=\"auto\")\n", - "plt.colorbar()\n", - "plt.xlabel(\"x1\")\n", - "plt.ylabel(\"x2\")\n" + "plt.plot(x, test_functions.rastrigin(x), c=\"b\")\n" ] }, { @@ -61,8 +56,9 @@ "source": [ "import bloptools\n", "\n", - "dofs = bloptools.experiments.tests.get_dummy_dofs(2)\n", - "bounds = [(-10, 10), (-10, 10)]" + "dofs = bloptools.devices.dummy_dofs(n=1) # an ophyd device that we can read/set\n", + "\n", + "bounds = [(-2., 2.)] # one set of bounds per dof" ] }, { @@ -71,7 +67,7 @@ "id": "7a88c7bd", "metadata": {}, "source": [ - "The agent automatically samples at different inputs, but we often need some post-processing after data collection. In this case, we need to give the agent a way to compute Himmelblau's function. We accomplish this with a digestion function:" + "This degree of freedom will move around a variable called `x1`. The agent automatically samples at different inputs, but we often need some post-processing after data collection. In this case, we need to give the agent a way to compute the Rastrigin function. We accomplish this with a digestion function, which always takes `(db, uid)` as an input. For each entry, we compute the function:\n" ] }, { @@ -82,13 +78,10 @@ "outputs": [], "source": [ "def digestion(db, uid):\n", + " products = db[uid].table()\n", "\n", - " table = db[uid].table()\n", - " products = {\"himmelblau\": []}\n", - "\n", - " for index, entry in table.iterrows():\n", - "\n", - " products[\"himmelblau\"].append(bloptools.experiments.tests.himmelblau(entry.x1, entry.x2))\n", + " for index, entry in products.iterrows():\n", + " products.loc[index, \"rastrigin\"] = test_functions.rastrigin(entry.x1)\n", "\n", " return products" ] @@ -99,7 +92,7 @@ "id": "dad64303", "metadata": {}, "source": [ - "The next ingredient is a task, which gives the agent something to do. We want it to minimize Himmelblau's function, so we make a task that will try to minimize the output of the digestion function called \"himmelblau\". We also include a transform function, which will make it easier to regress over the outputs of the function." + "The next ingredient is a task, which gives the agent something to do. We want it to minimize the Rastrigin function, so we make a task that will try to minimize the output of the digestion function called \"rastrigin\"." ] }, { @@ -111,7 +104,7 @@ "source": [ "from bloptools.tasks import Task\n", "\n", - "task = Task(key=\"himmelblau\", kind=\"min\", transform=lambda x: np.log(1 + 1e-2 * x))" + "task = Task(key=\"rastrigin\", kind=\"min\")" ] }, { @@ -134,15 +127,16 @@ "source": [ "%run -i ../../../examples/prepare_bluesky.py # prepare the bluesky environment\n", "\n", - "boa = bloptools.bayesian.Agent(\n", - " dofs=dofs, # things which we move around\n", - " bounds=bounds, # where we're allowed to move them\n", - " tasks=task, # tasks for the optimizer\n", - " digestion=digestion, # how to process the acquisition into task data\n", - " db=db, # a databroker instance\n", + "agent = bloptools.bayesian.Agent(\n", + " active_dofs=dofs,\n", + " passive_dofs=[],\n", + " active_dof_bounds=bounds,\n", + " tasks=[task],\n", + " digestion=digestion,\n", + " db=db,\n", " )\n", "\n", - "RE(boa.initialize(init_scheme='quasi-random', n_init=16))" + "RE(agent.initialize(init_scheme='quasi-random', n_init=4))" ] }, { @@ -163,114 +157,53 @@ }, "outputs": [], "source": [ - "boa.plot_tasks()" + "agent.plot_tasks()" ] }, { "attachments": {}, "cell_type": "markdown", - "id": "4274930c", + "id": "fa141636", "metadata": {}, "source": [ - "We can also the agent's posterior about the probability of goodness over the parameters:" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "d2eb855c", - "metadata": {}, - "source": [ - "We want to learn a bit more, so we can ask the agent where to sample based off of some strategy. Here we use the \"esti\" strategy, which maximizes the expected sum of tasks improvement.\n", + "Note that the value of the fitness is the negative value of the function: we always want to maximize the fitness of the tasks.\n", "\n", - "We can ask the agent to \"route\" them using ``ortools``, so that we can sample them more quickly if it requires us to e.g. move motors." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "334f3c66", - "metadata": {}, - "outputs": [], - "source": [ - "X_to_sample = boa.ask(strategy='esti', n=16, optimize=True, route=True)\n", - "plt.scatter(*X_to_sample.T)\n", - "plt.plot(*X_to_sample.T)" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "296d9fd2", - "metadata": {}, - "source": [ - "Let's tell the agent to learn a bit more (it will direct itself):" + "An important concept in Bayesian optimization is the acquisition function, which is how the agent decides where to sample next. Under the hood, the agent will see what inputs maximize the acquisition function to make its decision.\n", + "\n", + "We can see what the agent is thinking by asking it to plot a few different acquisition functions in its current state." ] }, { "cell_type": "code", "execution_count": null, - "id": "f0e74651", + "id": "8554d7c2", "metadata": { "tags": [] }, "outputs": [], "source": [ - "RE(boa.learn(strategy='esti', n_iter=1, n_per_iter=4))\n", - "boa.plot_tasks()" + "agent.plot_acquisition(strategy=[\"ei\", \"pi\", \"ucb\"])" ] }, { "attachments": {}, "cell_type": "markdown", - "id": "aeab7a9b", + "id": "2529763b", "metadata": {}, "source": [ - "The agent has updated its model of the tasks, including refitting the hyperparameters. Continuing:" + "Let's tell the agent to learn a little bit more. We just have to tell it what acquisition function to use (by passing a `strategy`) and how many iterations we'd like it to perform (by passing `n_iter`)." ] }, { "cell_type": "code", "execution_count": null, - "id": "d6b39b54", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "RE(boa.learn(strategy='esti', n_iter=4, n_per_iter=4))\n", - "boa.plot_tasks()" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "e955233f", + "id": "ebc65169", "metadata": {}, - "source": [ - "Eventually, we reach a point of saturation where no more improvement takes place:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "d73e4fd5", - "metadata": { - "tags": [] - }, "outputs": [], "source": [ - "RE(boa.learn(strategy='esti', n_iter=4, n_per_iter=4))\n", - "boa.plot_tasks()" + "RE(agent.learn(strategy='ei', n_iter=4))\n", + "agent.plot_tasks()" ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5ad4b1e7", - "metadata": {}, - "outputs": [], - "source": [] } ], "metadata": { From 5b7cfd03d116292734241718aad7a09ad96c2160 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 29 Jun 2023 18:38:46 -0400 Subject: [PATCH 04/11] don't import experiments --- bloptools/__init__.py | 2 +- bloptools/experiments/__init__.py | 1 - docs/source/tutorials/multi-task-sirepo.ipynb | 10 +++++----- 3 files changed, 6 insertions(+), 7 deletions(-) diff --git a/bloptools/__init__.py b/bloptools/__init__.py index 5591be5..8ba74e4 100644 --- a/bloptools/__init__.py +++ b/bloptools/__init__.py @@ -1,4 +1,4 @@ -from . import bayesian, devices, experiments, tasks # noqa F401 +from . import bayesian, devices, tasks # noqa F401 from ._version import get_versions __version__ = get_versions()["version"] diff --git a/bloptools/experiments/__init__.py b/bloptools/experiments/__init__.py index 35e5d2b..e69de29 100644 --- a/bloptools/experiments/__init__.py +++ b/bloptools/experiments/__init__.py @@ -1 +0,0 @@ -from . import tests # noqa F401 diff --git a/docs/source/tutorials/multi-task-sirepo.ipynb b/docs/source/tutorials/multi-task-sirepo.ipynb index 3986224..d528669 100644 --- a/docs/source/tutorials/multi-task-sirepo.ipynb +++ b/docs/source/tutorials/multi-task-sirepo.ipynb @@ -47,7 +47,7 @@ "beam_width_task = Task(key=\"x_width\", kind=\"min\", transform=lambda x: np.log(x))\n", "beam_height_task = Task(key=\"y_width\", kind=\"min\", transform=lambda x: np.log(x))\n", "\n", - "boa = Agent(\n", + "agent = Agent(\n", " active_dofs=kb_dofs, \n", " active_dof_bounds=kb_bounds,\n", " detectors=[w9],\n", @@ -56,7 +56,7 @@ " db=db,\n", ")\n", "\n", - "RE(boa.initialize(init_scheme='quasi-random', n_init=2))" + "RE(agent.initialize(init_scheme='quasi-random', n_init=8))" ] }, { @@ -77,7 +77,7 @@ }, "outputs": [], "source": [ - "boa.plot_tasks()" + "agent.plot_tasks()" ] }, { @@ -98,8 +98,8 @@ }, "outputs": [], "source": [ - "RE(boa.learn(strategy='esti', n_iter=4))\n", - "boa.plot_tasks()" + "RE(agent.learn(strategy='ei', n_iter=4))\n", + "agent.plot_tasks()" ] }, { From c1fb601a52323084091d26ae9b6b329cd55db28a Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Fri, 30 Jun 2023 11:11:30 -0400 Subject: [PATCH 05/11] better weighting for feasibility model --- bloptools/bayesian/__init__.py | 4 +- docs/source/tutorials/himmelblau.ipynb | 52 +++++++++++++++++++++++--- 2 files changed, 50 insertions(+), 6 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index 5170334..881a407 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -243,8 +243,10 @@ def tell(self, new_table=None, append=True, **kwargs): ) botorch.fit.fit_gpytorch_mll(task.regressor_mll, **kwargs) + log_feas_prob_weight = np.sqrt(np.sum(np.nanvar(self.targets.values, axis=0) * self.task_weights**2)) + self.task_scalarization = botorch.acquisition.objective.ScalarizedPosteriorTransform( - weights=torch.tensor([*[task.weight for task in self.tasks], 1]).double(), + weights=torch.tensor([*[task.weight for task in self.tasks], log_feas_prob_weight]).double(), offset=0, ) diff --git a/docs/source/tutorials/himmelblau.ipynb b/docs/source/tutorials/himmelblau.ipynb index 0b2244e..ca4fe61 100644 --- a/docs/source/tutorials/himmelblau.ipynb +++ b/docs/source/tutorials/himmelblau.ipynb @@ -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 < 36$" + "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$" ] }, { @@ -33,7 +33,9 @@ "\n", "x1 = x2 = np.linspace(-8, 8, 256)\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", "F[X1 ** 2 + X2 ** 2 > 50] = np.nan\n", "\n", @@ -49,7 +51,7 @@ "id": "7a88c7bd", "metadata": {}, "source": [ - "In our digestion function, we return a `NaN` when we violate the constraint:" + "where everything outside our constraint is undefined. In our digestion function, we return a `NaN` when we violate the constraint:" ] }, { @@ -93,6 +95,14 @@ "source": [ "%run -i ../../../examples/prepare_bluesky.py # prepare the bluesky environment\n", "\n", + "import bloptools\n", + "from bloptools.tasks import Task\n", + "\n", + "dofs = bloptools.devices.dummy_dofs(n=2)\n", + "bounds = [(-8, 8), (-8, 8)]\n", + "\n", + "task = Task(key=\"himmelblau\", kind=\"min\")\n", + "\n", "agent = bloptools.bayesian.Agent(\n", " active_dofs=dofs,\n", " passive_dofs=[],\n", @@ -140,7 +150,29 @@ { "cell_type": "code", "execution_count": null, - "id": "84a76eb9", + "id": "5e8c6458", + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "plt.scatter(*agent.test_inputs.T, c=agent.feas_model(torch.tensor(agent.test_inputs)))\n", + "plt.colorbar()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b9f0134d", + "metadata": {}, + "outputs": [], + "source": [ + "mean = agent.task_model.posterior(torch.tensor(agent.test_inputs).unsqueeze(1)).mean.squeeze(1).detach().numpy()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bc53bf67", "metadata": { "tags": [] }, @@ -156,7 +188,7 @@ "metadata": {}, "outputs": [], "source": [ - "RE(agent.learn(strategy='ei', n_iter=4))\n", + "RE(agent.learn(strategy='pi', n_iter=1))\n", "agent.plot_tasks()" ] }, @@ -172,7 +204,7 @@ { "cell_type": "code", "execution_count": null, - "id": "79fad8eb", + "id": "50b6582d", "metadata": {}, "outputs": [], "source": [ @@ -186,6 +218,16 @@ "id": "40491b5b", "metadata": {}, "outputs": [], + "source": [ + "agent.plot_acquisition()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fac7adf3", + "metadata": {}, + "outputs": [], "source": [] } ], From 12249753fb07ac89a9eb6a433661f67f335b8484 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Fri, 30 Jun 2023 11:57:38 -0400 Subject: [PATCH 06/11] moved draft notebooks to wip --- docs/{source/tutorials => wip}/custom-acquisition.ipynb | 0 docs/{source/tutorials => wip}/passive-dofs.ipynb | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename docs/{source/tutorials => wip}/custom-acquisition.ipynb (100%) rename docs/{source/tutorials => wip}/passive-dofs.ipynb (100%) diff --git a/docs/source/tutorials/custom-acquisition.ipynb b/docs/wip/custom-acquisition.ipynb similarity index 100% rename from docs/source/tutorials/custom-acquisition.ipynb rename to docs/wip/custom-acquisition.ipynb diff --git a/docs/source/tutorials/passive-dofs.ipynb b/docs/wip/passive-dofs.ipynb similarity index 100% rename from docs/source/tutorials/passive-dofs.ipynb rename to docs/wip/passive-dofs.ipynb From bdc66adac4910b4542887f1811c3c39bec8a3fca Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Wed, 12 Jul 2023 14:07:08 -0400 Subject: [PATCH 07/11] tweaks --- bloptools/bayesian/__init__.py | 312 +++++++++++------- bloptools/bayesian/kernels.py | 287 +++++----------- bloptools/bayesian/models.py | 60 ++-- bloptools/experiments/sirepo/tes.py | 12 +- bloptools/test_functions.py | 4 + bloptools/tests/test_bayesian_shadow.py | 10 +- bloptools/tests/test_bayesian_test_funcs.py | 8 +- bloptools/utils.py | 10 + docs/source/tutorials.rst | 4 +- ...lau.ipynb => constrained-himmelblau.ipynb} | 61 ++-- docs/source/tutorials/hyperparameters.ipynb | 156 +++++++++ docs/source/tutorials/introduction.ipynb | 31 +- .../tutorials/latent-toroid-dimensions.ipynb | 117 +++++++ docs/source/tutorials/multi-task-sirepo.ipynb | 48 ++- examples/prepare_tes_shadow.py | 6 +- 15 files changed, 686 insertions(+), 440 deletions(-) rename docs/source/tutorials/{himmelblau.ipynb => constrained-himmelblau.ipynb} (85%) create mode 100644 docs/source/tutorials/hyperparameters.ipynb create mode 100644 docs/source/tutorials/latent-toroid-dimensions.ipynb diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index 881a407..c43e64c 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -1,3 +1,5 @@ +import logging + import bluesky.plan_stubs as bps import bluesky.plans as bp # noqa F401 import botorch @@ -15,8 +17,8 @@ mpl.rc("image", cmap="coolwarm") -DEFAULT_COLOR_LIST = ["dodgerblue", "tomato", "mediumseagreen"] -DEFAULT_COLORMAP = "inferno" +DEFAULT_COLOR_LIST = ["dodgerblue", "tomato", "mediumseagreen", "goldenrod"] +DEFAULT_COLORMAP = "turbo" def default_acquisition_plan(dofs, inputs, dets): @@ -35,10 +37,10 @@ def default_digestion_plan(db, uid): # targets: these are what our model tries to predict from the inputs # tasks: these are quantities that our agent will try to optimize over +MAX_TEST_INPUTS = 2**11 -class Agent: - MAX_TEST_POINTS = 2**11 +class Agent: def __init__( self, active_dofs, @@ -65,8 +67,11 @@ def __init__( db : A databroker instance. """ + for dof in active_dofs: + dof.kind = "hinted" + self.active_dofs = np.atleast_1d(active_dofs) - self.active_dof_bounds = np.atleast_2d(active_dof_bounds) + self.active_dof_bounds = np.atleast_2d(active_dof_bounds).astype(float) self.tasks = np.atleast_1d(tasks) self.db = db @@ -77,6 +82,8 @@ def __init__( self.acquisition_plan = kwargs.get("acquisition_plan", default_acquisition_plan) self.digestion = kwargs.get("digestion", default_digestion_plan) + self.tolerate_acquisition_errors = kwargs.get("tolerate_acquisition_errors", True) + self.acquisition = acquisition.Acquisition() self.dets = np.atleast_1d(kwargs.get("detectors", [])) @@ -96,28 +103,59 @@ def __init__( self.training_iter = kwargs.get("training_iter", 256) - # self.test_normalized_active_inputs = self.sampler(n=self.MAX_TEST_POINTS) - self.test_active_inputs = self.unnormalize_active_inputs(self.sampler(n=self.MAX_TEST_POINTS)) + # self.test_normalized_active_inputs = self.sampler(n=self.MAX_TEST_INPUTS) + + # make some test points for sampling - n_per_active_dim = int(np.power(self.MAX_TEST_POINTS, 1 / self.n_active_dof)) + self.normalized_test_active_inputs = utils.normalized_sobol_sampler(n=MAX_TEST_INPUTS, d=self.n_active_dof) - test_normalized_active_inputs_grid = np.swapaxes( + n_per_active_dim = int(np.power(MAX_TEST_INPUTS, 1 / self.n_active_dof)) + + self.normalized_test_active_inputs_grid = np.swapaxes( np.r_[np.meshgrid(*self.n_active_dof * [np.linspace(0, 1, n_per_active_dim)])], 0, -1 ) - self.test_active_inputs_grid = self.unnormalize_active_inputs(test_normalized_active_inputs_grid) - self.table = pd.DataFrame() self._initialized = False + def normalize_active_inputs(self, x): + return (x - self.active_dof_bounds.min(axis=1)) / self.active_dof_bounds.ptp(axis=1) + + def unnormalize_active_inputs(self, x): + return x * self.active_dof_bounds.ptp(axis=1) + self.active_dof_bounds.min(axis=1) + + def active_inputs_sampler(self, n=MAX_TEST_INPUTS): + """ + Returns $n$ quasi-randomly sampled inputs in the bounded parameter space + """ + return self.unnormalize_active_inputs(utils.normalized_sobol_sampler(n, self.n_active_dof)) + + @property + def test_active_inputs(self): + """ + A static, quasi-randomly sampled set of test active inputs. + """ + return self.unnormalize_active_inputs(self.normalized_test_active_inputs) + + @property + def test_active_inputs_grid(self): + """ + A static, gridded set of test active inputs. + """ + return self.unnormalize_active_inputs(self.normalized_test_active_inputs_grid) + + # @property + # def input_transform(self): + # coefficient = torch.tensor(self.inputs.values.ptp(axis=0)) + # offset = torch.tensor(self.inputs.values.min(axis=0)) + # return botorch.models.transforms.input.AffineInputTransform( + # d=self.n_dof, coefficient=coefficient, offset=offset + # ) + @property def input_transform(self): - coefficient = torch.tensor(self.inputs.values.ptp(axis=0)) - offset = torch.tensor(self.inputs.values.min(axis=0)) - return botorch.models.transforms.input.AffineInputTransform( - d=self.n_dof, coefficient=coefficient, offset=offset - ) + return botorch.models.transforms.input.Normalize(d=self.n_dof) def save(self, filepath="./agent_data.h5"): """ @@ -139,13 +177,10 @@ def sampler(self, n): subset = np.random.choice(min_power_of_two, size=n, replace=False) return sp.stats.qmc.Sobol(d=self.n_active_dof, scramble=True).random(n=min_power_of_two)[subset] - def active_dof_sampler(self, n, q=1): - return botorch.utils.sampling.draw_sobol_samples(torch.tensor(self.active_dof_bounds.T), n=n, q=q) - def initialize( self, filepath=None, - init_scheme=None, + acqf=None, n_init=4, ): """ @@ -154,37 +189,33 @@ def initialize( It should be passed to a Bluesky RunEngine. """ - init_table = None - - if filepath is not None: - init_table = pd.read_hdf(filepath, key="table") - # experiment-specific stuff if self.initialization is not None: yield from self.initialization() + if filepath is not None: + self.tell(new_table=pd.read_hdf(filepath, key="table")) + # now let's get bayesian - if init_scheme == "quasi-random": - init_dof_inputs = self.ask(strategy="quasi-random", n=n_init, route=True) - init_table = yield from self.acquire(dof_inputs=init_dof_inputs) + elif acqf in ["qr"]: + yield from self.learn(acqf=acqf, n_iter=1, n_per_iter=n_init, route=True) else: raise Exception( - "Could not initialize model! Either pass initial X and data, or specify one of:" - "['quasi-random']." + """Could not initialize model! Either load a table, or specify an acqf from: +['qr'].""" ) - if init_table is None: - raise RuntimeError("Unhandled initialization error.") + # if init_table is None: + # raise RuntimeError("Unhandled initialization error.") - no_good_samples_tasks = np.isnan(init_table.loc[:, self.target_names]).all(axis=0) + no_good_samples_tasks = np.isnan(self.table.loc[:, self.target_names]).all(axis=0) if no_good_samples_tasks.any(): raise ValueError( f"The tasks {[self.tasks[i].name for i in np.where(no_good_samples_tasks)[0]]} " f"don't have any good samples." ) - self.tell(new_table=init_table, verbose=self.verbose) self._initialized = True def tell(self, new_table=None, append=True, **kwargs): @@ -230,7 +261,7 @@ def tell(self, new_table=None, append=True, **kwargs): ), ).double() - task.regressor = models.BoTorchSingleTaskGP( + task.regressor = models.LatentGP( train_inputs=train_inputs, train_targets=train_targets, likelihood=likelihood, @@ -241,9 +272,10 @@ def tell(self, new_table=None, append=True, **kwargs): task.regressor_mll = gpytorch.mlls.ExactMarginalLogLikelihood( task.regressor.likelihood, task.regressor ) - botorch.fit.fit_gpytorch_mll(task.regressor_mll, **kwargs) - log_feas_prob_weight = np.sqrt(np.sum(np.nanvar(self.targets.values, axis=0) * self.task_weights**2)) + log_feas_prob_weight = np.sqrt( + np.sum(np.nanvar(self.targets.values, axis=0) * np.square(self.task_weights)) + ) self.task_scalarization = botorch.acquisition.objective.ScalarizedPosteriorTransform( weights=torch.tensor([*[task.weight for task in self.tasks], log_feas_prob_weight]).double(), @@ -254,7 +286,7 @@ def tell(self, new_table=None, append=True, **kwargs): torch.as_tensor(self.all_targets_valid).long(), learn_additional_noise=True ).double() - self.dirichlet_classifier = models.BoTorchDirichletClassifier( + self.dirichlet_classifier = models.LatentDirichletClassifier( train_inputs=torch.tensor(self.inputs.values).double(), train_targets=dirichlet_likelihood.transformed_targets.transpose(-1, -2).double(), likelihood=dirichlet_likelihood, @@ -264,38 +296,49 @@ def tell(self, new_table=None, append=True, **kwargs): self.dirichlet_classifier_mll = gpytorch.mlls.ExactMarginalLogLikelihood( self.dirichlet_classifier.likelihood, self.dirichlet_classifier ) - botorch.fit.fit_gpytorch_mll(self.dirichlet_classifier_mll, **kwargs) self.feas_model = botorch.models.deterministic.GenericDeterministicModel( - f=lambda X: self.dirichlet_classifier.log_prob(X) + f=lambda X: -self.dirichlet_classifier.log_prob(X).square() ) + self.fit_models() + self.targets_model = botorch.models.model_list_gp_regression.ModelListGP( *[task.regressor for task in self.tasks] ) self.task_model = botorch.models.model.ModelList(*[task.regressor for task in self.tasks], self.feas_model) - def get_acquisition_function(self, strategy="ei", return_metadata=False, acqf_args={}, **kwargs): - if strategy.lower() == "ei": + def fit_models(self, **kwargs): + for task in self.tasks: + botorch.fit.fit_gpytorch_mll(task.regressor_mll, **kwargs) + botorch.fit.fit_gpytorch_mll(self.dirichlet_classifier_mll, **kwargs) + + def get_acquisition_function(self, acqf_name="ei", return_metadata=False, acqf_args={}, **kwargs): + if not self._initialized: + raise RuntimeError( + f'Can\'t construct acquisition function "{acqf_name}" (the agent is not initialized!)' + ) + + if acqf_name.lower() == "ei": acqf = botorch.acquisition.analytic.LogExpectedImprovement( self.task_model, best_f=self.best_sum_of_tasks, posterior_transform=self.task_scalarization, **kwargs, ) - acqf_meta = {"name": "Expected Improvement", "args": {}} + acqf_meta = {"name": "expected improvement", "args": {}} - elif strategy.lower() == "pi": + elif acqf_name.lower() == "pi": acqf = botorch.acquisition.analytic.LogProbabilityOfImprovement( self.task_model, best_f=self.best_sum_of_tasks, posterior_transform=self.task_scalarization, **kwargs, ) - acqf_meta = {"name": "Probability of Improvement", "args": {}} + acqf_meta = {"name": "probability of improvement", "args": {}} - elif strategy.lower() == "ucb": + elif acqf_name.lower() == "ucb": beta = acqf_args.get("beta", 0.1) acqf = botorch.acquisition.analytic.UpperConfidenceBound( self.task_model, @@ -303,10 +346,10 @@ def get_acquisition_function(self, strategy="ei", return_metadata=False, acqf_ar posterior_transform=self.task_scalarization, **kwargs, ) - acqf_meta = {"name": "Upper Confidence Bound", "args": {"beta": beta}} + acqf_meta = {"name": "upper confidence bound", "args": {"beta": beta}} else: - raise ValueError(f'Unrecognized acquisition strategy "{strategy}".') + raise ValueError(f'Unrecognized acquisition acqf_name "{acqf_name}".') return (acqf, acqf_meta) if return_metadata else acqf @@ -314,59 +357,61 @@ def ask( self, tasks=None, classifier=None, - strategy="ei", + acqf_name="ei", greedy=True, n=1, disappointment=0, - route=True, + route=False, cost_model=None, n_test=1024, optimize=True, + return_metadata=False, ): """ The next $n$ points to sample, recommended by the agent. """ - if route: - unrouted_points = self.ask( - tasks=tasks, - classifier=classifier, - strategy=strategy, - greedy=greedy, - n=n, - disappointment=disappointment, - route=False, - cost_model=cost_model, - n_test=n_test, - ) - - routing_index, _ = utils.get_routing(self.read_active_dofs, unrouted_points) - return unrouted_points[routing_index] - - if strategy.lower() == "quasi-random": - return self.unnormalize_active_inputs(self.sampler(n=n)) + # if route: + # unrouted_points = self.ask( + # tasks=tasks, + # classifier=classifier, + # strategy=strategy, + # greedy=greedy, + # n=n, + # disappointment=disappointment, + # route=False, + # cost_model=cost_model, + # n_test=n_test, + # ) + + # routing_index, _ = utils.get_routing(self.read_active_dofs, unrouted_points) + # x = unrouted_points[routing_index] + + if acqf_name.lower() == "qr": + x = self.active_inputs_sampler(n=n) + acqf_meta = {"name": "quasi-random", "args": {}} - if not self._initialized: - raise RuntimeError("The agent is not initialized!") - - self.acqf = self.get_acquisition_function(strategy=strategy) - - BATCH_SIZE = 1 - NUM_RESTARTS = 7 - RAW_SAMPLES = 512 + else: + acqf, acqf_meta = self.get_acquisition_function(acqf_name=acqf_name, return_metadata=True) + + BATCH_SIZE = 1 + NUM_RESTARTS = 7 + RAW_SAMPLES = 512 + + candidates, _ = botorch.optim.optimize_acqf( + acq_function=acqf, + bounds=torch.tensor(self.dof_bounds).T, + q=BATCH_SIZE, + num_restarts=NUM_RESTARTS, + raw_samples=RAW_SAMPLES, # used for intialization heuristic + options={"batch_limit": 3, "maxiter": 1024}, + ) - candidates, _ = botorch.optim.optimize_acqf( - acq_function=self.acqf, - bounds=torch.tensor(self.dof_bounds).T, - q=BATCH_SIZE, - num_restarts=NUM_RESTARTS, - raw_samples=RAW_SAMPLES, # used for intialization heuristic - options={"batch_limit": 5, "maxiter": 200}, - ) + x = candidates.detach().numpy()[..., self.dof_is_active_mask] - return candidates.detach().numpy()[..., self.dof_is_active_mask] + return (x, acqf_meta) if return_metadata else x - def acquire(self, dof_inputs): + def acquire(self, active_inputs): """ Acquire and digest according to the agent's acquisition and digestion plans. @@ -374,47 +419,43 @@ def acquire(self, dof_inputs): """ try: uid = yield from self.acquisition_plan( - self.dofs, dof_inputs, [*self.dets, *self.dofs, *self.passive_dofs] + self.dofs, active_inputs, [*self.dets, *self.dofs, *self.passive_dofs] ) + products = self.digestion(self.db, uid) - if "rejected" not in products.keys(): - products["rejected"] = False # compute the fitness for each task for index, entry in products.iterrows(): for task in self.tasks: products.loc[index, task.name] = task.get_fitness(entry) - except Exception as err: - raise err + except Exception as error: + if not self.tolerate_acquisition_errors: + raise error + logging.warning(f"Error in acquisition/digestion: {repr(error)}") + products = pd.DataFrame(active_inputs, columns=self.active_dof_names) + for task in self.tasks: + products.loc[:, task.name] = np.nan - if not len(dof_inputs) == len(products): - raise ValueError("The resulting table must be the same length as the sampled inputs!") + 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, strategy, n_iter=1, n_per_iter=1, reuse_hypers=True, upsample=1, verbose=True, plots=[], **kwargs - ): + def learn(self, acqf, n_iter=1, n_per_iter=1, reuse_hypers=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. """ - print(f'learning with strategy "{strategy}" ...') - for i in range(n_iter): - inputs_to_sample = np.atleast_2d(self.ask(n=n_per_iter, strategy=strategy, **kwargs)) - - new_table = yield from self.acquire(inputs_to_sample) + x, acqf_meta = self.ask(n=n_per_iter, acqf_name=acqf, return_metadata=True, **kwargs) - self.tell(new_table=new_table, reuse_hypers=reuse_hypers) + new_table = yield from self.acquire(x) - def normalize_active_inputs(self, inputs): - return (inputs - self.active_dof_bounds.min(axis=1)) / self.active_dof_bounds.ptp(axis=1) + new_table.loc[:, "acqf"] = acqf_meta["name"] - def unnormalize_active_inputs(self, X): - return X * self.active_dof_bounds.ptp(axis=1) + self.active_dof_bounds.min(axis=1) + self.tell(new_table=new_table, reuse_hypers=reuse_hypers) def normalize_inputs(self, inputs): return (inputs - self.input_bounds.min(axis=1)) / self.input_bounds.ptp(axis=1) @@ -631,7 +672,7 @@ def _plot_feas_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLO log_prob = self.dirichlet_classifier.log_prob(x).detach().numpy() self.class_axes[1].scatter( - *self.test_inputs.T[axes], s=size, c=np.exp(log_prob), vmin=0, vmax=1, cmap=cmap + *x.detach().numpy().T[axes], s=size, c=np.exp(log_prob), vmin=0, vmax=1, cmap=cmap ) self.class_fig.colorbar(data_ax, ax=self.class_axes[:2], location="bottom", aspect=32, shrink=0.8) @@ -731,10 +772,10 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL else: self.task_axes[itask, 1].scatter( - *self.test_inputs.T[axes], s=size, c=task_mean, norm=task_norm, cmap=cmap + *x.detach().numpy().T[axes], s=size, c=task_mean, norm=task_norm, cmap=cmap ) sigma_ax = self.task_axes[itask, 2].scatter( - *self.test_inputs.T[axes], s=size, c=task_sigma, cmap=cmap + *x.detach().numpy().T[axes], s=size, c=task_sigma, cmap=cmap ) self.task_fig.colorbar(data_ax, ax=self.task_axes[itask, :2], location="bottom", aspect=32, shrink=0.8) @@ -745,46 +786,46 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL ax.set_ylim(*self.active_dof_bounds[axes[1]]) def _plot_acq_one_dof(self, size=32, lw=1e0, **kwargs): - strategies = np.atleast_1d(kwargs.get("strategy", "ei")) + acqf_names = np.atleast_1d(kwargs.get("acqf", "ei")) self.acq_fig, self.acq_axes = plt.subplots( 1, - len(strategies), - figsize=(6 * len(strategies), 6), + len(acqf_names), + figsize=(6 * len(acqf_names), 6), sharex=True, constrained_layout=True, ) self.acq_axes = np.atleast_1d(self.acq_axes) - for istrat, strategy in enumerate(strategies): + for iacqf, acqf_name in enumerate(acqf_names): color = DEFAULT_COLOR_LIST[0] - acqf, acqf_meta = self.get_acquisition_function(strategy, return_metadata=True) + acqf, acqf_meta = self.get_acquisition_function(acqf_name, return_metadata=True) *grid_shape, dim = self.test_inputs_grid.shape x = torch.tensor(self.test_inputs_grid.reshape(-1, 1, dim)).double() obj = acqf.forward(x) - if strategy in ["ei", "pi"]: + if acqf_name in ["ei", "pi"]: obj = obj.exp() - self.acq_axes[istrat].set_title(acqf_meta["name"]) - self.acq_axes[istrat].plot( + self.acq_axes[iacqf].set_title(acqf_meta["name"]) + self.acq_axes[iacqf].plot( self.test_active_inputs_grid.ravel(), obj.detach().numpy().ravel(), lw=lw, color=color ) - self.acq_axes[istrat].set_xlim(*self.active_dof_bounds[0]) + self.acq_axes[iacqf].set_xlim(*self.active_dof_bounds[0]) def _plot_acq_many_dofs( self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=32, **kwargs ): - strategies = np.atleast_1d(kwargs.get("strategy", "ei")) + acqf_names = np.atleast_1d(kwargs.get("acqf", "ei")) self.acq_fig, self.acq_axes = plt.subplots( 1, - len(strategies), - figsize=(4 * len(strategies), 5), + len(acqf_names), + figsize=(4 * len(acqf_names), 5), sharex=True, sharey=True, constrained_layout=True, @@ -796,26 +837,43 @@ def _plot_acq_many_dofs( self.acq_axes = np.atleast_1d(self.acq_axes) self.acq_fig.suptitle(f"(x,y)=({self.dofs[axes[0]].name},{self.dofs[axes[1]].name})") - for istrat, strategy in enumerate(strategies): - acqf, acqf_meta = self.get_acquisition_function(strategy, return_metadata=True) + for iacqf, acqf_name in enumerate(acqf_names): + acqf, acqf_meta = self.get_acquisition_function(acqf_name, return_metadata=True) if gridded: *grid_shape, dim = self.test_inputs_grid.shape x = torch.tensor(self.test_inputs_grid.reshape(-1, 1, dim)).double() obj = acqf.forward(x) - if strategy in ["ei", "pi"]: + if acqf_name in ["ei", "pi"]: obj = obj.exp() - self.acq_axes[istrat].set_title(acqf_meta["name"]) - obj_ax = self.acq_axes[istrat].pcolormesh( + self.acq_axes[iacqf].set_title(acqf_meta["name"]) + obj_ax = self.acq_axes[iacqf].pcolormesh( *np.swapaxes(self.test_inputs_grid, 0, -1), obj.detach().numpy().reshape(grid_shape).T, shading=shading, cmap=cmap, ) - self.acq_fig.colorbar(obj_ax, ax=self.acq_axes[istrat], location="bottom", aspect=32, shrink=0.8) + self.acq_fig.colorbar(obj_ax, ax=self.acq_axes[iacqf], location="bottom", aspect=32, shrink=0.8) + + else: + *inputs_shape, dim = self.test_inputs.shape + x = torch.tensor(self.test_inputs.reshape(-1, 1, dim)).double() + obj = acqf.forward(x) + + if acqf_name in ["ei", "pi"]: + obj = obj.exp() + + self.acq_axes[iacqf].set_title(acqf_meta["name"]) + obj_ax = self.acq_axes[iacqf].scatter( + x.detach().numpy()[..., axes[0]], + x.detach().numpy()[..., axes[1]], + c=obj.detach().numpy().reshape(inputs_shape), + ) + + self.acq_fig.colorbar(obj_ax, ax=self.acq_axes[iacqf], location="bottom", aspect=32, shrink=0.8) for ax in self.acq_axes.ravel(): ax.set_xlim(*self.active_dof_bounds[axes[0]]) diff --git a/bloptools/bayesian/kernels.py b/bloptools/bayesian/kernels.py index 96481a2..2fe6f0b 100644 --- a/bloptools/bayesian/kernels.py +++ b/bloptools/bayesian/kernels.py @@ -1,275 +1,148 @@ -import math - import gpytorch import numpy as np import torch -class MultiOutputLatentKernel(gpytorch.kernels.Kernel): +class LatentKernel(gpytorch.kernels.Kernel): is_stationary = True + num_outputs = 1 + def __init__( self, num_inputs=1, - num_outputs=1, - off_diag=False, - diag_prior=False, + off_diag=True, + diag_prior=True, + scale_kernel=True, **kwargs, ): - super(MultiOutputLatentKernel, self).__init__() + super(LatentKernel, self).__init__() self.num_inputs = num_inputs - self.num_outputs = num_outputs self.n_off_diag = int(num_inputs * (num_inputs - 1) / 2) + self.off_diag = off_diag + self.scale_kernel = scale_kernel self.nu = kwargs.get("nu", 1.5) - # self.batch_shape = torch.Size([num_outputs]) + # kernel_scale_constraint = gpytorch.constraints.Positive() + diag_entries_constraint = gpytorch.constraints.Positive() # gpytorch.constraints.Interval(5e-1, 1e2) + skew_entries_constraint = gpytorch.constraints.Interval(-1e0, 1e0) - # output_scale_constraint = gpytorch.constraints.Positive() - diag_params_constraint = gpytorch.constraints.Interval(1e0, 1e2) - skew_params_constraint = gpytorch.constraints.Interval(-1e0, 1e0) - - diag_params_initial = np.sqrt(diag_params_constraint.lower_bound * diag_params_constraint.upper_bound) - raw_diag_params_initial = diag_params_constraint.inverse_transform(diag_params_initial) + # diag_entries_initial = np.ones() + # np.sqrt(diag_entries_constraint.lower_bound * diag_entries_constraint.upper_bound) + raw_diag_entries_initial = diag_entries_constraint.inverse_transform(torch.tensor(2)) self.register_parameter( - name="raw_diag_params", + name="raw_diag_entries", parameter=torch.nn.Parameter( - raw_diag_params_initial * torch.ones(self.num_outputs, self.num_inputs).double() + raw_diag_entries_initial * torch.ones(self.num_outputs, self.num_inputs).double() ), ) - - # self.register_constraint("raw_output_scale", output_scale_constraint) - self.register_constraint("raw_diag_params", diag_params_constraint) + self.register_constraint("raw_diag_entries", constraint=diag_entries_constraint) if diag_prior: self.register_prior( - name="diag_params_prior", - prior=gpytorch.priors.GammaPrior(concentration=0.5, rate=0.2), - param_or_closure=lambda m: m.diag_params, - setting_closure=lambda m, v: m._set_diag_params(v), + name="diag_entries_prior", + prior=gpytorch.priors.GammaPrior(concentration=2, rate=1), + param_or_closure=lambda m: m.diag_entries, + setting_closure=lambda m, v: m._set_diag_entries(v), ) if self.off_diag: self.register_parameter( - name="raw_skew_params", + name="raw_skew_entries", parameter=torch.nn.Parameter(torch.zeros(self.num_outputs, self.n_off_diag).double()), ) - self.register_constraint("raw_skew_params", skew_params_constraint) - - @property - def diag_params(self): - return self.raw_diag_params_constraint.transform(self.raw_diag_params) - - @property - def skew_params(self): - return self.raw_skew_params_constraint.transform(self.raw_skew_params) - - @diag_params.setter - def diag_params(self, value): - self._set_diag_params(value) - - @skew_params.setter - def skew_params(self, value): - self._set_skew_params(value) - - def _set_skew_params(self, value): - if not torch.is_tensor(value): - value = torch.as_tensor(value).to(self.raw_skew_params) - self.initialize(raw_skew_params=self.raw_skew_params_constraint.inverse_transform(value)) - - def _set_diag_params(self, value): - if not torch.is_tensor(value): - value = torch.as_tensor(value).to(self.raw_diag_params) - self.initialize(raw_diag_params=self.raw_diag_params_constraint.inverse_transform(value)) - - @property - def dimension_transform(self): - # no rotations - if not self.off_diag: - T = torch.eye(self.num_inputs, dtype=torch.float64) - - # construct an orthogonal matrix. fun fact: exp(skew(N)) is the generator of SO(N) - else: - A = torch.zeros((self.num_outputs, self.num_inputs, self.num_inputs), dtype=torch.float64) - upper_indices = np.triu_indices(self.num_inputs, k=1) - for output_index in range(self.num_outputs): - A[(output_index, *upper_indices)] = self.skew_params[output_index] - A += -A.transpose(-1, -2) - T = torch.linalg.matrix_exp(A) - - diagonal_transform = torch.cat([torch.diag(_values).unsqueeze(0) for _values in self.diag_params], dim=0) - T = torch.matmul(diagonal_transform, T) - - return T - - def forward(self, x1, x2, diag=False, **params): - # adapted from the Matern kernel - mean = x1.reshape(-1, x1.size(-1)).mean(0)[(None,) * (x1.dim() - 1)] - - trans_x1 = torch.matmul(self.dimension_transform.unsqueeze(1), (x1 - mean).unsqueeze(-1)).squeeze(-1) - trans_x2 = torch.matmul(self.dimension_transform.unsqueeze(1), (x2 - mean).unsqueeze(-1)).squeeze(-1) - - distance = self.covar_dist(trans_x1, trans_x2, diag=diag, **params) - - # if distance.shape[0] == 1: - # distance = distance.squeeze(0) # this is extremely necessary - - exp_component = torch.exp(-math.sqrt(self.nu * 2) * distance) - - if self.nu == 0.5: - constant_component = 1 - elif self.nu == 1.5: - constant_component = (math.sqrt(3) * distance).add(1) - elif self.nu == 2.5: - constant_component = (math.sqrt(5) * distance).add(1).add(5.0 / 3.0 * distance**2) + self.register_constraint("raw_skew_entries", skew_entries_constraint) - return constant_component * exp_component + if self.scale_kernel: + kernel_scale_constraint = gpytorch.constraints.Positive() + kernel_scale_prior = gpytorch.priors.GammaPrior(concentration=2, rate=0.15) + self.register_parameter( + name="raw_kernel_scale", + parameter=torch.nn.Parameter(torch.ones(1)), + ) -class LatentMaternKernel(gpytorch.kernels.Kernel): - def __init__( - self, - n_dim, - off_diag=False, - diagonal_prior=False, - **kwargs, - ): - super(LatentMaternKernel, self).__init__() - - self.n_dim = n_dim - self.n_off_diag = int(n_dim * (n_dim - 1) / 2) - self.off_diag = off_diag - - # output_scale_constraint = gpytorch.constraints.Positive() - diag_params_constraint = gpytorch.constraints.Interval(1e-1, 1e2) - skew_params_constraint = gpytorch.constraints.Interval(-1e0, 1e0) - - diag_params_initial = np.sqrt(diag_params_constraint.lower_bound * diag_params_constraint.upper_bound) - raw_diag_params_initial = diag_params_constraint.inverse_transform(diag_params_initial) - - # self.register_parameter( - # name="raw_output_scale", parameter=torch.nn.Parameter(torch.ones(*self.batch_shape, 1).double()) - # ) - self.register_parameter( - name="raw_diag_params", - parameter=torch.nn.Parameter( - raw_diag_params_initial * torch.ones(*self.batch_shape, self.n_dim).double() - ), - ) - - # self.register_constraint("raw_output_scale", output_scale_constraint) - self.register_constraint("raw_diag_params", diag_params_constraint) + self.register_constraint("raw_kernel_scale", constraint=kernel_scale_constraint) - if diagonal_prior: self.register_prior( - name="diag_params_prior", - prior=gpytorch.priors.GammaPrior(concentration=0.5, rate=0.2), - param_or_closure=lambda m: m.diag_params, - setting_closure=lambda m, v: m._set_diag_params(v), - ) - - if self.off_diag: - self.register_parameter( - name="raw_skew_params", - parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, self.n_off_diag).double()), + name="kernel_scale_prior", + prior=kernel_scale_prior, + param_or_closure=lambda m: m.kernel_scale, + setting_closure=lambda m, v: m._set_kernel_scale(v), ) - self.register_constraint("raw_skew_params", skew_params_constraint) - # @property - # def output_scale(self): - # return self.raw_output_scale_constraint.transform(self.raw_output_scale) + @property + def diag_entries(self): + return self.raw_diag_entries_constraint.transform(self.raw_diag_entries) @property - def diag_params(self): - return self.raw_diag_params_constraint.transform(self.raw_diag_params) + def skew_entries(self): + return self.raw_skew_entries_constraint.transform(self.raw_skew_entries) @property - def skew_params(self): - return self.raw_skew_params_constraint.transform(self.raw_skew_params) + def kernel_scale(self): + return self.raw_kernel_scale_constraint.transform(self.raw_kernel_scale) - # @output_scale.setter - # def output_scale(self, value): - # self._set_output_scale(value) + @diag_entries.setter + def diag_entries(self, value): + self._set_diag_entries(value) - @diag_params.setter - def diag_params(self, value): - self._set_diag_params(value) + @skew_entries.setter + def skew_entries(self, value): + self._set_skew_entries(value) - @skew_params.setter - def skew_params(self, value): - self._set_skew_params(value) + @kernel_scale.setter + def kernel_scale(self, value): + self._set_kernel_scale(value) - def _set_skew_params(self, value): + def _set_diag_entries(self, value): if not torch.is_tensor(value): - value = torch.as_tensor(value).to(self.raw_skew_params) - self.initialize(raw_skew_params=self.raw_skew_params_constraint.inverse_transform(value)) + value = torch.as_tensor(value).to(self.raw_diag_entries) + self.initialize(raw_diag_entries=self.raw_diag_entries_constraint.inverse_transform(value)) - # def _set_output_scale(self, value): - # if not torch.is_tensor(value): - # value = torch.as_tensor(value).to(self.raw_output_scale) - # self.initialize(raw_output_scale=self.raw_output_scale_constraint.inverse_transform(value)) + def _set_skew_entries(self, value): + if not torch.is_tensor(value): + value = torch.as_tensor(value).to(self.raw_skew_entries) + self.initialize(raw_skew_entries=self.raw_skew_entries_constraint.inverse_transform(value)) - def _set_diag_params(self, value): + def _set_kernel_scale(self, value): if not torch.is_tensor(value): - value = torch.as_tensor(value).to(self.raw_diag_params) - self.initialize(raw_diag_params=self.raw_diag_params_constraint.inverse_transform(value)) + value = torch.as_tensor(value).to(self.raw_kernel_scale) + self.initialize(raw_kernel_scale=self.raw_kernel_scale_constraint.inverse_transform(value)) @property - def trans_matrix(self): + def output_scale(self): + return self.kernel_scale.sqrt() + + @property + def latent_dimensions(self): # no rotations if not self.off_diag: - T = torch.eye(self.n_dim).double() + T = torch.eye(self.num_inputs, dtype=torch.float64) # construct an orthogonal matrix. fun fact: exp(skew(N)) is the generator of SO(N) else: - A = torch.zeros((self.n_dim, self.n_dim)).double() - A[np.triu_indices(self.n_dim, k=1)] = self.skew_params - A += -A.T + A = torch.zeros((self.num_inputs, self.num_inputs)).double() + A[np.triu_indices(self.num_inputs, k=1)] = self.skew_entries + A += -A.transpose(-1, -2) T = torch.linalg.matrix_exp(A) - T = torch.matmul(torch.diag(self.diag_params), T) + diagonal_transform = torch.cat([torch.diag(entries).unsqueeze(0) for entries in self.diag_entries], dim=0) + T = torch.matmul(diagonal_transform, T) return T - def forward(self, x1, x2=None, diag=False, auto=False, last_dim_is_batch=False, **params): - # returns the homoskedastic diagonal - if diag: - # return torch.square(self.output_scale[0]) * torch.ones((*self.batch_shape, *x1.shape[:-1])) - return torch.ones((*self.batch_shape, *x1.shape[:-1])) - - # computes the autocovariance of the process at the parameters - if auto: - x2 = x1 - - # print(x1, x2) - - # x1 and x2 are arrays of shape (..., n_1, n_dim) and (..., n_2, n_dim) - _x1, _x2 = torch.as_tensor(x1).double(), torch.as_tensor(x2).double() - - # dx has shape (..., n_1, n_2, n_dim) - dx = _x1.unsqueeze(-2) - _x2.unsqueeze(-3) - - # transform coordinates with hyperparameters (this applies lengthscale and rotations) - trans_dx = torch.matmul(self.trans_matrix, dx.unsqueeze(-1)) - - # total transformed distance. D has shape (..., n_1, n_2) - d_eff = torch.sqrt(torch.matmul(trans_dx.transpose(-1, -2), trans_dx).sum((-1, -2)) + 1e-12) - - # Matern covariance of effective order nu=3/2. - # nu=3/2 is a special case and has a concise closed-form expression - # In general, this is something between an exponential (n=1/2) and a Gaussian (n=infinity) - # https://en.wikipedia.org/wiki/Matern_covariance_function - - # C = torch.exp(-d_eff) # Matern_0.5 (exponential) - C = (1 + d_eff) * torch.exp(-d_eff) # Matern_1.5 - # C = (1 + d_eff + 1 / 3 * torch.square(d_eff)) * torch.exp(-d_eff) # Matern_2.5 - # C = torch.exp(-0.5 * np.square(d_eff)) # Matern_infinity (RBF) + def forward(self, x1, x2, diag=False, **params): + # adapted from the Matern kernel + mean = x1.reshape(-1, x1.size(-1)).mean(0)[(None,) * (x1.dim() - 1)] - # C = torch.square(self.output_scale[0]) * torch.exp(-torch.square(d_eff)) + trans_x1 = torch.matmul(self.latent_dimensions.unsqueeze(1), (x1 - mean).unsqueeze(-1)).squeeze(-1) + trans_x2 = torch.matmul(self.latent_dimensions.unsqueeze(1), (x2 - mean).unsqueeze(-1)).squeeze(-1) - # print(f'{diag = } {C.shape = }') + distance = self.covar_dist(trans_x1, trans_x2, diag=diag, **params) - return C + return self.kernel_scale * (1 + distance) * torch.exp(-distance) diff --git a/bloptools/bayesian/models.py b/bloptools/bayesian/models.py index 91d2e7d..21ebc36 100644 --- a/bloptools/bayesian/models.py +++ b/bloptools/bayesian/models.py @@ -7,62 +7,40 @@ from . import kernels -class BoTorchDirichletClassifier(botorch.models.gp_regression.SingleTaskGP): - # _num_outputs = 1 # to inform GPyTorchModel API - +class LatentDirichletClassifier(botorch.models.gp_regression.SingleTaskGP): def __init__(self, train_inputs, train_targets, *args, **kwargs): super().__init__(train_inputs, train_targets, *args, **kwargs) - # self.mean_module = gpytorch.means.ConstantMean(batch_shape=len(train_targets.unique())) - # self.covar_module = gpytorch.kernels.ScaleKernel( - # kernels.LatentMaternKernel(n_dim=train_inputs.shape[-1], off_diag=False, diagonal_prior=False) - # ) - - # print(train_inputs.shape, train_targets.shape) - self.mean_module = gpytorch.means.ConstantMean() # batch_shape=torch.Size((len(train_targets.unique()),))) - self.covar_module = gpytorch.kernels.ScaleKernel( - kernels.MultiOutputLatentKernel( - num_inputs=train_inputs.shape[-1], - num_outputs=train_targets.shape[-1], - off_diag=True, - diag_prior=True, - **kwargs - ) + self.mean_module = gpytorch.means.ConstantMean() + self.covar_module = kernels.LatentKernel( + num_inputs=train_inputs.shape[-1], + num_outputs=train_targets.shape[-1], + off_diag=True, + diag_prior=True, + scale_output=True, + **kwargs ) - def forward(self, x): - mean_x = self.mean_module(x) - covar_x = self.covar_module(x) - return gpytorch.distributions.MultivariateNormal(mean_x, covar_x) - - # def posterior(self, X): - # return botorch.posteriors.GPyTorchPosterior(gpytorch.distributions.MultivariateNormal()) - def log_prob(self, x, n_samples=256): *input_shape, n_dim = x.shape samples = self.posterior(x.reshape(-1, n_dim)).sample(torch.Size((n_samples,))).exp() return torch.log((samples / samples.sum(-1, keepdim=True)).mean(0)[:, 1]).reshape(*input_shape, 1) -class BoTorchSingleTaskGP(botorch.models.gp_regression.SingleTaskGP): +class LatentGP(botorch.models.gp_regression.SingleTaskGP): def __init__(self, train_inputs, train_targets, *args, **kwargs): super().__init__(train_inputs, train_targets, *args, **kwargs) - self.mean_module = gpytorch.means.ConstantMean() - self.covar_module = gpytorch.kernels.ScaleKernel( - kernels.MultiOutputLatentKernel( - num_inputs=train_inputs.shape[-1], - num_outputs=train_targets.shape[-1], - off_diag=True, - diag_prior=True, - **kwargs - ) - ) + self.mean_module = gpytorch.means.ConstantMean(constant_prior=gpytorch.priors.NormalPrior(loc=0, scale=1)) - def forward(self, x): - mean_x = self.mean_module(x) - covar_x = self.covar_module(x) - return gpytorch.distributions.MultivariateNormal(mean_x, covar_x) + self.covar_module = kernels.LatentKernel( + num_inputs=train_inputs.shape[-1], + num_outputs=train_targets.shape[-1], + off_diag=True, + diag_prior=True, + scale_output=True, + **kwargs + ) class OldBoTorchSingleTaskGP(ExactGP, GPyTorchModel): diff --git a/bloptools/experiments/sirepo/tes.py b/bloptools/experiments/sirepo/tes.py index 9fc26bb..a1b1267 100644 --- a/bloptools/experiments/sirepo/tes.py +++ b/bloptools/experiments/sirepo/tes.py @@ -1,11 +1,15 @@ import numpy as np -def digestion(db, uid, image_name="w9"): - """ - Simulating a misaligned Gaussian beam. The optimum is at (1, 1, 1, 1) - """ +def w8_digestion(db, uid): + return image_digestion(db, uid, image_name="w8") + +def w9_digestion(db, uid): + return image_digestion(db, uid, image_name="w9") + + +def image_digestion(db, uid, image_name): products = db[uid].table(fill=True) for index, entry in products.iterrows(): diff --git a/bloptools/test_functions.py b/bloptools/test_functions.py index 195f7d5..25693a2 100644 --- a/bloptools/test_functions.py +++ b/bloptools/test_functions.py @@ -5,6 +5,10 @@ def booth(x1, x2): return (x1 + 2 * x2 - 7) ** 2 + (2 * x1 + x2 - 5) ** 2 +def matyas(x1, x2): + return 0.26 * (x1**2 + x2**2) - 0.48 * x1 * x2 + + def himmelblau(x1, x2): return (x1**2 + x2 - 11) ** 2 + (x1 + x2**2 - 7) ** 2 diff --git a/bloptools/tests/test_bayesian_shadow.py b/bloptools/tests/test_bayesian_shadow.py index 9d4eded..aa6b908 100644 --- a/bloptools/tests/test_bayesian_shadow.py +++ b/bloptools/tests/test_bayesian_shadow.py @@ -3,7 +3,7 @@ from sirepo_bluesky.sirepo_ophyd import create_classes import bloptools -from bloptools.experiments.sirepo import tes +from bloptools.experiments.sirepo.tes import w9_digestion from bloptools.tasks import Task @@ -29,13 +29,13 @@ def test_bayesian_agent_tes_shadow(RE, db, shadow_tes_simulation): detectors=[w9], active_dof_bounds=kb_bounds, tasks=[beam_flux_task, beam_width_task, beam_height_task], - digestion=tes.digestion, + digestion=w9_digestion, db=db, ) - RE(agent.initialize(init_scheme="quasi-random", n_init=4)) + RE(agent.initialize(acqf="qr", n_init=4)) - RE(agent.learn(strategy="ei", n_iter=2)) - RE(agent.learn(strategy="pi", n_iter=2)) + RE(agent.learn(acqf="ei", n_iter=2)) + RE(agent.learn(acqf="pi", n_iter=2)) agent.plot_tasks() diff --git a/bloptools/tests/test_bayesian_test_funcs.py b/bloptools/tests/test_bayesian_test_funcs.py index 2eff571..058b709 100644 --- a/bloptools/tests/test_bayesian_test_funcs.py +++ b/bloptools/tests/test_bayesian_test_funcs.py @@ -20,9 +20,9 @@ def test_bayesian_agent_himmelblau(RE, db): db=db, ) - RE(agent.initialize(init_scheme="quasi-random", n_init=16)) + RE(agent.initialize(acqf="qr", n_init=16)) - RE(agent.learn(strategy="ei", n_iter=2)) + RE(agent.learn(acqf="ei", n_iter=2)) agent.plot_tasks() @@ -43,8 +43,8 @@ def test_bayesian_agent_mock_kbs(RE, db): db=db, ) - RE(agent.initialize(init_scheme="quasi-random", n_init=16)) + RE(agent.initialize(acqf="qr", n_init=16)) - RE(agent.learn(strategy="ei", n_iter=4)) + RE(agent.learn(acqf="ei", n_iter=4)) agent.plot_tasks() diff --git a/bloptools/utils.py b/bloptools/utils.py index 74e6d98..380ca69 100644 --- a/bloptools/utils.py +++ b/bloptools/utils.py @@ -1,8 +1,18 @@ +import botorch import numpy as np import scipy as sp +import torch from ortools.constraint_solver import pywrapcp, routing_enums_pb2 +def normalized_sobol_sampler(n, d): + """ + Returns $n$ quasi-randomly sampled points in the [0,1]^d hypercube + """ + x = botorch.utils.sampling.draw_sobol_samples(torch.outer(torch.tensor([0, 1]), torch.ones(d)), n=n, q=1) + return x.squeeze(1).detach().numpy() + + def estimate_root_indices(x): # or, indices_before_sign_changes i_whole = np.where(np.sign(x[1:]) != np.sign(x[:-1]))[0] diff --git a/docs/source/tutorials.rst b/docs/source/tutorials.rst index fb721bd..48218c9 100644 --- a/docs/source/tutorials.rst +++ b/docs/source/tutorials.rst @@ -5,5 +5,7 @@ Tutorials :maxdepth: 2 tutorials/introduction.ipynb - tutorials/himmelblau.ipynb + tutorials/hyperparameters.ipynb + tutorials/constrained-himmelblau.ipynb + tutorials/latent-toroid-dimensions.ipynb tutorials/multi-task-sirepo.ipynb diff --git a/docs/source/tutorials/himmelblau.ipynb b/docs/source/tutorials/constrained-himmelblau.ipynb similarity index 85% rename from docs/source/tutorials/himmelblau.ipynb rename to docs/source/tutorials/constrained-himmelblau.ipynb index ca4fe61..d078283 100644 --- a/docs/source/tutorials/himmelblau.ipynb +++ b/docs/source/tutorials/constrained-himmelblau.ipynb @@ -112,7 +112,7 @@ " db=db,\n", " )\n", "\n", - "RE(agent.initialize(init_scheme='quasi-random', n_init=16))\n", + "RE(agent.initialize(acqf='qr', n_init=32))\n", "\n", "agent.plot_tasks()\n" ] @@ -141,32 +141,41 @@ { "attachments": {}, "cell_type": "markdown", - "id": "a79b56ac", + "id": "ab608930", "metadata": {}, "source": [ - "The agent automatically tries to avoid infeasible points, but will end up naturally exploring the boundary of the constraint. Let's see where the agent is thinking of going:" + "It combines the estimate of the objective and the estimate of the feasibility in deciding where to go:" ] }, { "cell_type": "code", "execution_count": null, - "id": "5e8c6458", - "metadata": {}, + "id": "6c48549c", + "metadata": { + "tags": [] + }, "outputs": [], "source": [ - "import torch\n", - "plt.scatter(*agent.test_inputs.T, c=agent.feas_model(torch.tensor(agent.test_inputs)))\n", - "plt.colorbar()" + "agent.plot_acquisition(acqf=[\"ei\", \"pi\", \"ucb\"])" ] }, { "cell_type": "code", "execution_count": null, - "id": "b9f0134d", + "id": "ff1c5f1c", "metadata": {}, "outputs": [], "source": [ - "mean = agent.task_model.posterior(torch.tensor(agent.test_inputs).unsqueeze(1)).mean.squeeze(1).detach().numpy()" + "RE(agent.learn(\"ei\", n_iter=4))" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "a79b56ac", + "metadata": {}, + "source": [ + "The agent automatically tries to avoid infeasible points, but will end up naturally exploring the boundary of the constraint. Let's see where the agent is thinking of going:" ] }, { @@ -178,20 +187,10 @@ }, "outputs": [], "source": [ + "agent.plot_tasks()\n", "agent.plot_acquisition(strategy=[\"ei\", \"pi\", \"ucb\"])" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "ebc65169", - "metadata": {}, - "outputs": [], - "source": [ - "RE(agent.learn(strategy='pi', n_iter=1))\n", - "agent.plot_tasks()" - ] - }, { "attachments": {}, "cell_type": "markdown", @@ -208,27 +207,9 @@ "metadata": {}, "outputs": [], "source": [ - "RE(agent.learn(strategy='ei', n_iter=16))\n", + "RE(agent.learn('ei', n_iter=16))\n", "agent.plot_tasks()" ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "40491b5b", - "metadata": {}, - "outputs": [], - "source": [ - "agent.plot_acquisition()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "fac7adf3", - "metadata": {}, - "outputs": [], - "source": [] } ], "metadata": { diff --git a/docs/source/tutorials/hyperparameters.ipynb b/docs/source/tutorials/hyperparameters.ipynb new file mode 100644 index 0000000..7507193 --- /dev/null +++ b/docs/source/tutorials/hyperparameters.ipynb @@ -0,0 +1,156 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "id": "e7b5e13a-c059-441d-8d4f-fff080d52054", + "metadata": {}, + "source": [ + "# Hyperparameters\n", + "\n", + "Consider the test function of two inputs $f(x_1, x_2) = e^{- (x_1/10)^2 - x_2^2}$. We can see that this function varies more quickly in the $x_1$ direction than the $x_2$ direction, especially if we look at the plot:" + ] + }, + { + "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", + "\n", + "x1 = x2 = np.linspace(-10, 10, 256)\n", + "X1, X2 = np.meshgrid(x1, x2)\n", + "\n", + "F = - np.exp(-(X1/8)**2) * np.exp(-X2**2) \n", + "\n", + "plt.pcolormesh(x1, x2, F, shading=\"auto\")\n", + "plt.colorbar()\n", + "plt.xlabel(\"x1\")\n", + "plt.ylabel(\"x2\")\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "7a88c7bd", + "metadata": {}, + "source": [ + "The optimization goes faster if our model understands how the function changes as we change the inputs in different ways. The way it picks up on this is by starting from a general model that could describe a lot of functions, and making it specific to this one by choosing the right hyperparameters. Our Bayesian agent is very good at this, and only needs a few samples to figure out what the function looks like:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "071a829f-a390-40dc-9d5b-ae75702e119e", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "%run -i ../../../examples/prepare_bluesky.py # prepare the bluesky environment\n", + "\n", + "import bloptools\n", + "from bloptools.tasks import Task\n", + "\n", + "dofs = bloptools.devices.dummy_dofs(n=2)\n", + "bounds = [(-8, 8), (-8, 8)]\n", + "\n", + "task = Task(key=\"loss\", kind=\"min\")\n", + "\n", + "def digestion(db, uid):\n", + "\n", + " products = db[uid].table()\n", + "\n", + " for index, entry in products.iterrows():\n", + " products.loc[index, \"loss\"] = - np.exp(-1e-2 * entry.x1**2) * np.exp(-1e-1 * entry.x2**2) \n", + "\n", + " return products\n", + "\n", + "agent = bloptools.bayesian.Agent(\n", + " active_dofs=dofs,\n", + " passive_dofs=[],\n", + " active_dof_bounds=bounds,\n", + " tasks=[task],\n", + " digestion=digestion,\n", + " db=db,\n", + " )\n", + "\n", + "RE(agent.initialize(init_scheme='quasi-random', n_init=16))\n", + "\n", + "agent.plot_tasks()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8fba0f0c", + "metadata": {}, + "outputs": [], + "source": [ + "agent.tasks[0].regressor.covar_module.dimension_transform" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "9ab3be01", + "metadata": {}, + "source": [ + "In addition to modeling the fitness of the task, the agent models the probability that an input will be feasible:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bc53bf67", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "agent.plot_acquisition(strategy=[\"ei\", \"pi\", \"ucb\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ebc65169", + "metadata": {}, + "outputs": [], + "source": [ + "RE(agent.learn(strategy='ei', n_iter=8))\n", + "agent.plot_tasks()" + ] + } + ], + "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.9.16" + }, + "vscode": { + "interpreter": { + "hash": "9aced674e98d511b4f654e147532c84d38dc986fe042b1e92785fb9d8df41f75" + } + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/source/tutorials/introduction.ipynb b/docs/source/tutorials/introduction.ipynb index e5ea4c1..a2f67a9 100644 --- a/docs/source/tutorials/introduction.ipynb +++ b/docs/source/tutorials/introduction.ipynb @@ -33,7 +33,7 @@ "\n", "from bloptools import test_functions\n", "\n", - "x = np.linspace(-2, 2, 256)\n", + "x = np.linspace(-4, 4, 256)\n", "\n", "plt.plot(x, test_functions.rastrigin(x), c=\"b\")\n" ] @@ -58,7 +58,7 @@ "\n", "dofs = bloptools.devices.dummy_dofs(n=1) # an ophyd device that we can read/set\n", "\n", - "bounds = [(-2., 2.)] # one set of bounds per dof" + "bounds = [(-4., 4.)] # one set of bounds per dof" ] }, { @@ -136,7 +136,18 @@ " db=db,\n", " )\n", "\n", - "RE(agent.initialize(init_scheme='quasi-random', n_init=4))" + "RE(agent.initialize(init_scheme='quasi-random', n_init=8))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "22e77570", + "metadata": {}, + "outputs": [], + "source": [ + "import bluesky\n", + "bluesky.__version__" ] }, { @@ -157,6 +168,8 @@ }, "outputs": [], "source": [ + "# what are the points?\n", + "\n", "agent.plot_tasks()" ] }, @@ -182,6 +195,8 @@ }, "outputs": [], "source": [ + "# helper function to list acquisition functions\n", + "\n", "agent.plot_acquisition(strategy=[\"ei\", \"pi\", \"ucb\"])" ] }, @@ -204,6 +219,14 @@ "RE(agent.learn(strategy='ei', n_iter=4))\n", "agent.plot_tasks()" ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3e4e0fe6", + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { @@ -222,7 +245,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.16" + "version": "3.9.16 (main, Mar 8 2023, 14:00:05) \n[GCC 11.2.0]" }, "vscode": { "interpreter": { diff --git a/docs/source/tutorials/latent-toroid-dimensions.ipynb b/docs/source/tutorials/latent-toroid-dimensions.ipynb new file mode 100644 index 0000000..4f4a4f6 --- /dev/null +++ b/docs/source/tutorials/latent-toroid-dimensions.ipynb @@ -0,0 +1,117 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "id": "e7b5e13a-c059-441d-8d4f-fff080d52054", + "metadata": {}, + "source": [ + "# Finding latent dimensions for the toroidal mirror \n", + "\n", + "It is common that beamline inputs are highly coupled, and so the effect of an input on the beam cannot be understood except in concert with the others. In this example, we show how our agent figures out latent dimensions, as well as the benefit of doing so. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fa8a6989", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "%run -i ~/repos/bloptools/examples/prepare_bluesky.py\n", + "%run -i ~/repos/bloptools/examples/prepare_tes_shadow.py\n", + "\n", + "toroid_dofs = np.array([toroid.x_rot, toroid.offz])\n", + "toroid_bounds = np.array([[-1e-3, 1e-3], [-4e-1, +4e-1]])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "071a829f-a390-40dc-9d5b-ae75702e119e", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "\n", + "from bloptools.bayesian import Agent\n", + "from bloptools.experiments.sirepo.tes import w8_digestion\n", + "from bloptools.tasks import Task\n", + "\n", + "beam_flux_task = Task(key=\"flux\", kind=\"max\", transform=lambda x: np.log(x))\n", + "\n", + "agent = Agent(\n", + " active_dofs=toroid_dofs, \n", + " active_dof_bounds=toroid_bounds,\n", + " detectors=[w8],\n", + " tasks=[beam_flux_task],\n", + " digestion=w8_digestion,\n", + " db=db,)\n", + "\n", + "RE(agent.initialize(acqf=\"qr\", n_init=24))" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "a6259a4f", + "metadata": {}, + "source": [ + "We can see that the beam is only not cut off (i.e. it has a non-zero flux) in a diagonal strip, and that in fact this is really just a one-dimensional optimization problem in some diagonal dimension. Our agent has figured this out, with a transformation matrix that has a long coherence length in one dimension and a short coherence length orthogonal to it:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6e17e666", + "metadata": {}, + "outputs": [], + "source": [ + "agent.tasks[0].regressor.covar_module.latent_dimensions" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "996c3c01-f91d-4a25-9b8d-eba5fa964504", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "agent.plot_tasks()\n", + "agent.plot_feasibility()\n", + "agent.plot_acquisition(strategy=[\"ei\", \"pi\", \"ucb\"])" + ] + } + ], + "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.9.16" + }, + "vscode": { + "interpreter": { + "hash": "9aced674e98d511b4f654e147532c84d38dc986fe042b1e92785fb9d8df41f75" + } + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/source/tutorials/multi-task-sirepo.ipynb b/docs/source/tutorials/multi-task-sirepo.ipynb index d528669..5c02fe9 100644 --- a/docs/source/tutorials/multi-task-sirepo.ipynb +++ b/docs/source/tutorials/multi-task-sirepo.ipynb @@ -29,6 +29,16 @@ "kb_bounds = np.array([[-0.10, +0.10], [-0.10, +0.10]]) " ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "fe879ea5", + "metadata": {}, + "outputs": [], + "source": [ + "kb_bounds" + ] + }, { "cell_type": "code", "execution_count": null, @@ -40,7 +50,7 @@ "source": [ "\n", "from bloptools.bayesian import Agent\n", - "from bloptools.experiments.sirepo.tes import digestion\n", + "from bloptools.experiments.sirepo.tes import w9_digestion\n", "from bloptools.tasks import Task\n", "\n", "beam_flux_task = Task(key=\"flux\", kind=\"max\", transform=lambda x: np.log(x))\n", @@ -52,11 +62,11 @@ " active_dof_bounds=kb_bounds,\n", " detectors=[w9],\n", " tasks=[beam_flux_task, beam_width_task, beam_height_task],\n", - " digestion=digestion,\n", + " digestion=w9_digestion,\n", " db=db,\n", ")\n", "\n", - "RE(agent.initialize(init_scheme='quasi-random', n_init=8))" + "RE(agent.initialize(init_scheme='qr', n_init=16))" ] }, { @@ -89,6 +99,16 @@ "Let's optimize:" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "4bafc3cb", + "metadata": {}, + "outputs": [], + "source": [ + "agent.plot_acquisition(strategy=[\"ei\", \"pi\", \"ucb\"])" + ] + }, { "cell_type": "code", "execution_count": null, @@ -98,10 +118,30 @@ }, "outputs": [], "source": [ - "RE(agent.learn(strategy='ei', n_iter=4))\n", + "RE(agent.learn(acqf='ei', n_iter=2))\n", "agent.plot_tasks()" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "b8cb9315", + "metadata": {}, + "outputs": [], + "source": [ + "agent.tasks[2].regressor.covar_module.latent_dimensions" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "163f64c5", + "metadata": {}, + "outputs": [], + "source": [ + "plt.imshow(agent.table.w9_image.values[-1])" + ] + }, { "attachments": {}, "cell_type": "markdown", diff --git a/examples/prepare_tes_shadow.py b/examples/prepare_tes_shadow.py index 116ee76..ff5e8b8 100644 --- a/examples/prepare_tes_shadow.py +++ b/examples/prepare_tes_shadow.py @@ -12,13 +12,13 @@ classes, objects = create_classes(connection.data, connection=connection) globals().update(**objects) -data["models"]["simulation"]["npoint"] = 200000 -data["models"]["watchpointReport12"]["histogramBins"] = 32 +data["models"]["simulation"]["npoint"] = 50000 +data["models"]["watchpointReport12"]["histogramBins"] = 16 # w9.duration.kind = "hinted" bec.disable_baseline() bec.disable_heading() -bec.disable_table() +# bec.disable_table() # This should be done by installing the package with `pip install -e .` or something similar. # import sys From 8e2a98085f7ee6db8527360e676c4f05bafb85a6 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 13 Jul 2023 09:14:23 -0400 Subject: [PATCH 08/11] added black formatting to notebooks --- bloptools/tasks.py | 2 +- .../tutorials/constrained-himmelblau.ipynb | 30 ++++----- docs/source/tutorials/hyperparameters.ipynb | 39 ++++++----- docs/source/tutorials/introduction.ipynb | 45 ++++--------- .../tutorials/latent-toroid-dimensions.ipynb | 14 ++-- docs/source/tutorials/multi-task-sirepo.ipynb | 66 ++++--------------- 6 files changed, 72 insertions(+), 124 deletions(-) diff --git a/bloptools/tasks.py b/bloptools/tasks.py index 5437e46..61ce47b 100644 --- a/bloptools/tasks.py +++ b/bloptools/tasks.py @@ -1,5 +1,5 @@ class Task: - MIN_NOISE_LEVEL = 1e-4 + MIN_NOISE_LEVEL = 1e-6 MAX_NOISE_LEVEL = 1e-2 def __init__(self, key, kind="max", name=None, transform=None, **kwargs): diff --git a/docs/source/tutorials/constrained-himmelblau.ipynb b/docs/source/tutorials/constrained-himmelblau.ipynb index d078283..8a8f20e 100644 --- a/docs/source/tutorials/constrained-himmelblau.ipynb +++ b/docs/source/tutorials/constrained-himmelblau.ipynb @@ -37,12 +37,12 @@ "\n", "task = Task(key=\"himmelblau\", kind=\"min\")\n", "F = test_functions.himmelblau(X1, X2)\n", - "F[X1 ** 2 + X2 ** 2 > 50] = np.nan\n", + "F[X1**2 + X2**2 > 50] = np.nan\n", "\n", "plt.pcolormesh(x1, x2, F, norm=mpl.colors.LogNorm(), shading=\"auto\")\n", "plt.colorbar()\n", "plt.xlabel(\"x1\")\n", - "plt.ylabel(\"x2\")\n" + "plt.ylabel(\"x2\")" ] }, { @@ -62,12 +62,10 @@ "outputs": [], "source": [ "def digestion(db, uid):\n", - "\n", " products = db[uid].table()\n", "\n", " for index, entry in products.iterrows():\n", - "\n", - " if entry.x1 ** 2 + entry.x2 ** 2 < 50:\n", + " if entry.x1**2 + entry.x2**2 < 50:\n", " products.loc[index, \"himmelblau\"] = test_functions.himmelblau(entry.x1, entry.x2)\n", " else:\n", " products.loc[index, \"himmelblau\"] = np.nan\n", @@ -104,17 +102,17 @@ "task = Task(key=\"himmelblau\", kind=\"min\")\n", "\n", "agent = bloptools.bayesian.Agent(\n", - " active_dofs=dofs,\n", - " passive_dofs=[],\n", - " active_dof_bounds=bounds,\n", - " tasks=[task],\n", - " digestion=digestion,\n", - " db=db,\n", - " )\n", + " active_dofs=dofs,\n", + " passive_dofs=[],\n", + " active_dof_bounds=bounds,\n", + " tasks=[task],\n", + " digestion=digestion,\n", + " db=db,\n", + ")\n", "\n", - "RE(agent.initialize(acqf='qr', n_init=32))\n", + "RE(agent.initialize(acqf=\"qr\", n_init=32))\n", "\n", - "agent.plot_tasks()\n" + "agent.plot_tasks()" ] }, { @@ -150,7 +148,7 @@ { "cell_type": "code", "execution_count": null, - "id": "6c48549c", + "id": "28c5c0df", "metadata": { "tags": [] }, @@ -207,7 +205,7 @@ "metadata": {}, "outputs": [], "source": [ - "RE(agent.learn('ei', n_iter=16))\n", + "RE(agent.learn(\"ei\", n_iter=16))\n", "agent.plot_tasks()" ] } diff --git a/docs/source/tutorials/hyperparameters.ipynb b/docs/source/tutorials/hyperparameters.ipynb index 7507193..500282e 100644 --- a/docs/source/tutorials/hyperparameters.ipynb +++ b/docs/source/tutorials/hyperparameters.ipynb @@ -25,12 +25,12 @@ "x1 = x2 = np.linspace(-10, 10, 256)\n", "X1, X2 = np.meshgrid(x1, x2)\n", "\n", - "F = - np.exp(-(X1/8)**2) * np.exp(-X2**2) \n", + "F = -np.exp(-((X1 / 8) ** 2)) * np.exp(-(X2**2))\n", "\n", "plt.pcolormesh(x1, x2, F, shading=\"auto\")\n", "plt.colorbar()\n", "plt.xlabel(\"x1\")\n", - "plt.ylabel(\"x2\")\n" + "plt.ylabel(\"x2\")" ] }, { @@ -61,27 +61,28 @@ "\n", "task = Task(key=\"loss\", kind=\"min\")\n", "\n", - "def digestion(db, uid):\n", "\n", + "def digestion(db, uid):\n", " products = db[uid].table()\n", "\n", " for index, entry in products.iterrows():\n", - " products.loc[index, \"loss\"] = - np.exp(-1e-2 * entry.x1**2) * np.exp(-1e-1 * entry.x2**2) \n", + " products.loc[index, \"loss\"] = -np.exp(-1e-2 * entry.x1**2) * np.exp(-1e-1 * entry.x2**2)\n", "\n", " return products\n", "\n", + "\n", "agent = bloptools.bayesian.Agent(\n", - " active_dofs=dofs,\n", - " passive_dofs=[],\n", - " active_dof_bounds=bounds,\n", - " tasks=[task],\n", - " digestion=digestion,\n", - " db=db,\n", - " )\n", + " active_dofs=dofs,\n", + " passive_dofs=[],\n", + " active_dof_bounds=bounds,\n", + " tasks=[task],\n", + " digestion=digestion,\n", + " db=db,\n", + ")\n", "\n", - "RE(agent.initialize(init_scheme='quasi-random', n_init=16))\n", + "RE(agent.initialize(init_scheme=\"quasi-random\", n_init=16))\n", "\n", - "agent.plot_tasks()\n" + "agent.plot_tasks()" ] }, { @@ -122,9 +123,17 @@ "metadata": {}, "outputs": [], "source": [ - "RE(agent.learn(strategy='ei', n_iter=8))\n", + "RE(agent.learn(strategy=\"ei\", n_iter=8))\n", "agent.plot_tasks()" ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9ca1fa6c", + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { @@ -143,7 +152,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.16" + "version": "3.9.16 (main, Mar 8 2023, 14:00:05) \n[GCC 11.2.0]" }, "vscode": { "interpreter": { diff --git a/docs/source/tutorials/introduction.ipynb b/docs/source/tutorials/introduction.ipynb index a2f67a9..4aabba2 100644 --- a/docs/source/tutorials/introduction.ipynb +++ b/docs/source/tutorials/introduction.ipynb @@ -35,7 +35,7 @@ "\n", "x = np.linspace(-4, 4, 256)\n", "\n", - "plt.plot(x, test_functions.rastrigin(x), c=\"b\")\n" + "plt.plot(x, test_functions.rastrigin(x), c=\"b\")" ] }, { @@ -56,9 +56,9 @@ "source": [ "import bloptools\n", "\n", - "dofs = bloptools.devices.dummy_dofs(n=1) # an ophyd device that we can read/set\n", + "dofs = bloptools.devices.dummy_dofs(n=1) # an ophyd device that we can read/set\n", "\n", - "bounds = [(-4., 4.)] # one set of bounds per dof" + "bounds = [(-4.0, 4.0)] # one set of bounds per dof" ] }, { @@ -128,26 +128,15 @@ "%run -i ../../../examples/prepare_bluesky.py # prepare the bluesky environment\n", "\n", "agent = bloptools.bayesian.Agent(\n", - " active_dofs=dofs,\n", - " passive_dofs=[],\n", - " active_dof_bounds=bounds,\n", - " tasks=[task],\n", - " digestion=digestion,\n", - " db=db,\n", - " )\n", + " active_dofs=dofs,\n", + " passive_dofs=[],\n", + " active_dof_bounds=bounds,\n", + " tasks=[task],\n", + " digestion=digestion,\n", + " db=db,\n", + ")\n", "\n", - "RE(agent.initialize(init_scheme='quasi-random', n_init=8))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "22e77570", - "metadata": {}, - "outputs": [], - "source": [ - "import bluesky\n", - "bluesky.__version__" + "RE(agent.initialize(acqf=\"qr\", n_init=12))" ] }, { @@ -216,17 +205,9 @@ "metadata": {}, "outputs": [], "source": [ - "RE(agent.learn(strategy='ei', n_iter=4))\n", + "RE(agent.learn(acqf=\"ei\", n_iter=8))\n", "agent.plot_tasks()" ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "3e4e0fe6", - "metadata": {}, - "outputs": [], - "source": [] } ], "metadata": { @@ -245,7 +226,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.16 (main, Mar 8 2023, 14:00:05) \n[GCC 11.2.0]" + "version": "3.9.16" }, "vscode": { "interpreter": { diff --git a/docs/source/tutorials/latent-toroid-dimensions.ipynb b/docs/source/tutorials/latent-toroid-dimensions.ipynb index 4f4a4f6..ac54bad 100644 --- a/docs/source/tutorials/latent-toroid-dimensions.ipynb +++ b/docs/source/tutorials/latent-toroid-dimensions.ipynb @@ -36,7 +36,6 @@ }, "outputs": [], "source": [ - "\n", "from bloptools.bayesian import Agent\n", "from bloptools.experiments.sirepo.tes import w8_digestion\n", "from bloptools.tasks import Task\n", @@ -44,12 +43,13 @@ "beam_flux_task = Task(key=\"flux\", kind=\"max\", transform=lambda x: np.log(x))\n", "\n", "agent = Agent(\n", - " active_dofs=toroid_dofs, \n", - " active_dof_bounds=toroid_bounds,\n", - " detectors=[w8],\n", - " tasks=[beam_flux_task],\n", - " digestion=w8_digestion,\n", - " db=db,)\n", + " active_dofs=toroid_dofs,\n", + " active_dof_bounds=toroid_bounds,\n", + " detectors=[w8],\n", + " tasks=[beam_flux_task],\n", + " digestion=w8_digestion,\n", + " db=db,\n", + ")\n", "\n", "RE(agent.initialize(acqf=\"qr\", n_init=24))" ] diff --git a/docs/source/tutorials/multi-task-sirepo.ipynb b/docs/source/tutorials/multi-task-sirepo.ipynb index 5c02fe9..d3fefd0 100644 --- a/docs/source/tutorials/multi-task-sirepo.ipynb +++ b/docs/source/tutorials/multi-task-sirepo.ipynb @@ -26,17 +26,7 @@ "%run -i ../../../examples/prepare_tes_shadow.py\n", "\n", "kb_dofs = [kbv.x_rot, kbh.x_rot]\n", - "kb_bounds = np.array([[-0.10, +0.10], [-0.10, +0.10]]) " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "fe879ea5", - "metadata": {}, - "outputs": [], - "source": [ - "kb_bounds" + "kb_bounds = np.array([[-0.10, +0.10], [-0.10, +0.10]])" ] }, { @@ -48,7 +38,6 @@ }, "outputs": [], "source": [ - "\n", "from bloptools.bayesian import Agent\n", "from bloptools.experiments.sirepo.tes import w9_digestion\n", "from bloptools.tasks import Task\n", @@ -58,15 +47,15 @@ "beam_height_task = Task(key=\"y_width\", kind=\"min\", transform=lambda x: np.log(x))\n", "\n", "agent = Agent(\n", - " active_dofs=kb_dofs, \n", - " active_dof_bounds=kb_bounds,\n", - " detectors=[w9],\n", - " tasks=[beam_flux_task, beam_width_task, beam_height_task],\n", - " digestion=w9_digestion,\n", - " db=db,\n", + " active_dofs=kb_dofs,\n", + " active_dof_bounds=kb_bounds,\n", + " detectors=[w9],\n", + " tasks=[beam_flux_task, beam_width_task, beam_height_task],\n", + " digestion=w9_digestion,\n", + " db=db,\n", ")\n", "\n", - "RE(agent.initialize(init_scheme='qr', n_init=16))" + "RE(agent.initialize(acqf=\"qr\", n_init=16))" ] }, { @@ -87,7 +76,8 @@ }, "outputs": [], "source": [ - "agent.plot_tasks()" + "agent.plot_tasks()\n", + "agent.plot_acquisition(strategy=[\"ei\", \"pi\", \"ucb\"])" ] }, { @@ -96,17 +86,7 @@ "id": "296d9fd2", "metadata": {}, "source": [ - "Let's optimize:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4bafc3cb", - "metadata": {}, - "outputs": [], - "source": [ - "agent.plot_acquisition(strategy=[\"ei\", \"pi\", \"ucb\"])" + "We should find our optimum (or something close to it) on the very next iteration:" ] }, { @@ -118,30 +98,10 @@ }, "outputs": [], "source": [ - "RE(agent.learn(acqf='ei', n_iter=2))\n", + "RE(agent.learn(acqf=\"ei\", n_iter=2))\n", "agent.plot_tasks()" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "b8cb9315", - "metadata": {}, - "outputs": [], - "source": [ - "agent.tasks[2].regressor.covar_module.latent_dimensions" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "163f64c5", - "metadata": {}, - "outputs": [], - "source": [ - "plt.imshow(agent.table.w9_image.values[-1])" - ] - }, { "attachments": {}, "cell_type": "markdown", @@ -168,7 +128,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.16" + "version": "3.9.16 (main, Mar 8 2023, 14:00:05) \n[GCC 11.2.0]" }, "vscode": { "interpreter": { From b5d88673682149e331e063187ee7a9cdb81684b4 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 13 Jul 2023 14:44:35 -0400 Subject: [PATCH 09/11] fixed syntax error in docs --- docs/source/tutorials/hyperparameters.ipynb | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/docs/source/tutorials/hyperparameters.ipynb b/docs/source/tutorials/hyperparameters.ipynb index 500282e..0fe5e2a 100644 --- a/docs/source/tutorials/hyperparameters.ipynb +++ b/docs/source/tutorials/hyperparameters.ipynb @@ -80,7 +80,7 @@ " db=db,\n", ")\n", "\n", - "RE(agent.initialize(init_scheme=\"quasi-random\", n_init=16))\n", + "RE(agent.initialize(acqf=\"qr\", n_init=16))\n", "\n", "agent.plot_tasks()" ] @@ -92,7 +92,7 @@ "metadata": {}, "outputs": [], "source": [ - "agent.tasks[0].regressor.covar_module.dimension_transform" + "agent.tasks[0].regressor.covar_module.latent_dimensions" ] }, { @@ -123,17 +123,9 @@ "metadata": {}, "outputs": [], "source": [ - "RE(agent.learn(strategy=\"ei\", n_iter=8))\n", + "RE(agent.learn(\"ei\", n_iter=4))\n", "agent.plot_tasks()" ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "9ca1fa6c", - "metadata": {}, - "outputs": [], - "source": [] } ], "metadata": { @@ -152,7 +144,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.16 (main, Mar 8 2023, 14:00:05) \n[GCC 11.2.0]" + "version": "3.9.16" }, "vscode": { "interpreter": { From d03b8fa79935fc7f580cceb81d22b41e6015fa89 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 13 Jul 2023 14:57:26 -0400 Subject: [PATCH 10/11] last bug I swear --- docs/source/tutorials/latent-toroid-dimensions.ipynb | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/docs/source/tutorials/latent-toroid-dimensions.ipynb b/docs/source/tutorials/latent-toroid-dimensions.ipynb index ac54bad..d6ffe16 100644 --- a/docs/source/tutorials/latent-toroid-dimensions.ipynb +++ b/docs/source/tutorials/latent-toroid-dimensions.ipynb @@ -20,8 +20,8 @@ }, "outputs": [], "source": [ - "%run -i ~/repos/bloptools/examples/prepare_bluesky.py\n", - "%run -i ~/repos/bloptools/examples/prepare_tes_shadow.py\n", + "%run -i ../../../examples/prepare_bluesky.py\n", + "%run -i ../../../examples/prepare_tes_shadow.py\n", "\n", "toroid_dofs = np.array([toroid.x_rot, toroid.offz])\n", "toroid_bounds = np.array([[-1e-3, 1e-3], [-4e-1, +4e-1]])" @@ -86,6 +86,14 @@ "agent.plot_feasibility()\n", "agent.plot_acquisition(strategy=[\"ei\", \"pi\", \"ucb\"])" ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9093dd73", + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { From 21305cb2ff0a63daa1c7400dc7c42d6cdca11dc1 Mon Sep 17 00:00:00 2001 From: Max Rakitin Date: Sun, 16 Jul 2023 13:31:00 -0400 Subject: [PATCH 11/11] STY: apply max-line-length=125 to all files with black/isort/flake8 --- .isort.cfg | 2 +- .pre-commit-config.yaml | 3 ++ bloptools/bayesian/__init__.py | 58 ++++++++---------------------- bloptools/bayesian/kernels.py | 4 +-- bloptools/de/de_opt_utils.py | 8 ++--- bloptools/de/de_optimization.py | 24 ++++--------- bloptools/de/hardware_flyer.py | 4 +-- bloptools/experiments/nsls2/tes.py | 4 +-- bloptools/utils.py | 4 +-- docs/wip/custom-acquisition.ipynb | 22 ++++-------- docs/wip/passive-dofs.ipynb | 42 ++++++++++------------ examples/benchmark.py | 4 +-- pyproject.toml | 2 +- 13 files changed, 57 insertions(+), 124 deletions(-) diff --git a/.isort.cfg b/.isort.cfg index e0926f4..43f177b 100644 --- a/.isort.cfg +++ b/.isort.cfg @@ -1,4 +1,4 @@ [settings] -line_length = 115 +line_length = 125 multi_line_output = 3 include_trailing_comma = True diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 4336326..6de8684 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -11,6 +11,9 @@ repos: rev: 23.1.0 hooks: - id: black + language_version: python3.10 + - id: black-jupyter + language_version: python3.10 - repo: https://github.com/pycqa/flake8 rev: 6.0.0 hooks: diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index c43e64c..13c3297 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -245,9 +245,7 @@ def tell(self, new_table=None, append=True, **kwargs): task.feasibility = self.all_targets_valid.astype(int) train_inputs = torch.tensor(self.inputs.loc[task.feasibility == 1].values).double().unsqueeze(0) - train_targets = ( - torch.tensor(task.targets.loc[task.feasibility == 1].values).double().unsqueeze(0).unsqueeze(-1) - ) + train_targets = torch.tensor(task.targets.loc[task.feasibility == 1].values).double().unsqueeze(0).unsqueeze(-1) if train_inputs.ndim == 1: train_inputs = train_inputs.unsqueeze(-1) @@ -269,13 +267,9 @@ def tell(self, new_table=None, append=True, **kwargs): outcome_transform=botorch.models.transforms.outcome.Standardize(m=1, batch_shape=torch.Size((1,))), ).double() - task.regressor_mll = gpytorch.mlls.ExactMarginalLogLikelihood( - task.regressor.likelihood, task.regressor - ) + task.regressor_mll = gpytorch.mlls.ExactMarginalLogLikelihood(task.regressor.likelihood, task.regressor) - log_feas_prob_weight = np.sqrt( - np.sum(np.nanvar(self.targets.values, axis=0) * np.square(self.task_weights)) - ) + log_feas_prob_weight = np.sqrt(np.sum(np.nanvar(self.targets.values, axis=0) * np.square(self.task_weights))) self.task_scalarization = botorch.acquisition.objective.ScalarizedPosteriorTransform( weights=torch.tensor([*[task.weight for task in self.tasks], log_feas_prob_weight]).double(), @@ -303,9 +297,7 @@ def tell(self, new_table=None, append=True, **kwargs): self.fit_models() - self.targets_model = botorch.models.model_list_gp_regression.ModelListGP( - *[task.regressor for task in self.tasks] - ) + self.targets_model = botorch.models.model_list_gp_regression.ModelListGP(*[task.regressor for task in self.tasks]) self.task_model = botorch.models.model.ModelList(*[task.regressor for task in self.tasks], self.feas_model) @@ -316,9 +308,7 @@ def fit_models(self, **kwargs): def get_acquisition_function(self, acqf_name="ei", return_metadata=False, acqf_args={}, **kwargs): if not self._initialized: - raise RuntimeError( - f'Can\'t construct acquisition function "{acqf_name}" (the agent is not initialized!)' - ) + raise RuntimeError(f'Can\'t construct acquisition function "{acqf_name}" (the agent is not initialized!)') if acqf_name.lower() == "ei": acqf = botorch.acquisition.analytic.LogExpectedImprovement( @@ -418,9 +408,7 @@ def acquire(self, active_inputs): This should yield a table of sampled tasks with the same length as the sampled inputs. """ try: - uid = yield from self.acquisition_plan( - self.dofs, active_inputs, [*self.dets, *self.dofs, *self.passive_dofs] - ) + uid = yield from self.acquisition_plan(self.dofs, active_inputs, [*self.dets, *self.dofs, *self.passive_dofs]) products = self.digestion(self.db, uid) @@ -471,9 +459,7 @@ def unnormalize_targets(self, targets): @property def test_inputs(self): - test_passive_inputs = ( - self.passive_inputs.values[-1][None] * np.ones(len(self.test_active_inputs))[..., None] - ) + test_passive_inputs = self.passive_inputs.values[-1][None] * np.ones(len(self.test_active_inputs))[..., None] return np.concatenate([self.test_active_inputs, test_passive_inputs], axis=-1) @property @@ -654,9 +640,7 @@ def _plot_feas_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLO if gridded: x = torch.tensor(self.test_inputs_grid.reshape(-1, self.n_dof)).double() - log_prob = ( - self.dirichlet_classifier.log_prob(x).detach().numpy().reshape(self.test_inputs_grid.shape[:-1]) - ) + log_prob = self.dirichlet_classifier.log_prob(x).detach().numpy().reshape(self.test_inputs_grid.shape[:-1]) self.class_axes[1].pcolormesh( *np.swapaxes(self.test_inputs_grid, 0, -1), @@ -671,9 +655,7 @@ def _plot_feas_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLO x = torch.tensor(self.test_inputs).double() log_prob = self.dirichlet_classifier.log_prob(x).detach().numpy() - self.class_axes[1].scatter( - *x.detach().numpy().T[axes], s=size, c=np.exp(log_prob), vmin=0, vmax=1, cmap=cmap - ) + self.class_axes[1].scatter(*x.detach().numpy().T[axes], s=size, c=np.exp(log_prob), vmin=0, vmax=1, cmap=cmap) self.class_fig.colorbar(data_ax, ax=self.class_axes[:2], location="bottom", aspect=32, shrink=0.8) @@ -745,11 +727,7 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL *self.inputs.values.T[axes], s=size, c=task.targets, norm=task_norm, cmap=cmap ) - x = ( - torch.tensor(self.test_inputs_grid).double() - if gridded - else torch.tensor(self.test_inputs).double() - ) + x = torch.tensor(self.test_inputs_grid).double() if gridded else torch.tensor(self.test_inputs).double() task_posterior = task.regressor.posterior(x) task_mean = task_posterior.mean.detach().numpy() # * task.targets_scale + task.targets_mean @@ -771,12 +749,8 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL ) else: - self.task_axes[itask, 1].scatter( - *x.detach().numpy().T[axes], s=size, c=task_mean, norm=task_norm, cmap=cmap - ) - sigma_ax = self.task_axes[itask, 2].scatter( - *x.detach().numpy().T[axes], s=size, c=task_sigma, cmap=cmap - ) + self.task_axes[itask, 1].scatter(*x.detach().numpy().T[axes], s=size, c=task_mean, norm=task_norm, cmap=cmap) + sigma_ax = self.task_axes[itask, 2].scatter(*x.detach().numpy().T[axes], s=size, c=task_sigma, 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) @@ -811,15 +785,11 @@ def _plot_acq_one_dof(self, size=32, lw=1e0, **kwargs): obj = obj.exp() self.acq_axes[iacqf].set_title(acqf_meta["name"]) - self.acq_axes[iacqf].plot( - self.test_active_inputs_grid.ravel(), obj.detach().numpy().ravel(), lw=lw, color=color - ) + self.acq_axes[iacqf].plot(self.test_active_inputs_grid.ravel(), obj.detach().numpy().ravel(), lw=lw, color=color) self.acq_axes[iacqf].set_xlim(*self.active_dof_bounds[0]) - def _plot_acq_many_dofs( - self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=32, **kwargs - ): + def _plot_acq_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=32, **kwargs): acqf_names = np.atleast_1d(kwargs.get("acqf", "ei")) self.acq_fig, self.acq_axes = plt.subplots( diff --git a/bloptools/bayesian/kernels.py b/bloptools/bayesian/kernels.py index 2fe6f0b..0bb1430 100644 --- a/bloptools/bayesian/kernels.py +++ b/bloptools/bayesian/kernels.py @@ -36,9 +36,7 @@ def __init__( self.register_parameter( name="raw_diag_entries", - parameter=torch.nn.Parameter( - raw_diag_entries_initial * torch.ones(self.num_outputs, self.num_inputs).double() - ), + parameter=torch.nn.Parameter(raw_diag_entries_initial * torch.ones(self.num_outputs, self.num_inputs).double()), ) self.register_constraint("raw_diag_entries", constraint=diag_entries_constraint) diff --git a/bloptools/de/de_opt_utils.py b/bloptools/de/de_opt_utils.py index 033d68c..095326e 100644 --- a/bloptools/de/de_opt_utils.py +++ b/bloptools/de/de_opt_utils.py @@ -126,9 +126,7 @@ def _run_flyers(flyers): return uid_list -def run_hardware_fly( - motors, detector, population, max_velocity, min_velocity, start_det, read_det, stop_det, watch_func -): +def run_hardware_fly(motors, detector, population, max_velocity, min_velocity, start_det, read_det, stop_det, watch_func): flyers = generate_hardware_flyers( motors=motors, detector=detector, @@ -143,9 +141,7 @@ def run_hardware_fly( return _run_flyers(flyers) -def run_fly_sim( - population, num_interm_vals, num_scans_at_once, sim_id, server_name, root_dir, watch_name, run_parallel -): +def run_fly_sim(population, num_interm_vals, num_scans_at_once, sim_id, server_name, root_dir, watch_name, run_parallel): flyers = generate_sim_flyers( population=population, num_between_vals=num_interm_vals, diff --git a/bloptools/de/de_optimization.py b/bloptools/de/de_optimization.py index eb35748..54ed426 100644 --- a/bloptools/de/de_optimization.py +++ b/bloptools/de/de_optimization.py @@ -7,9 +7,7 @@ from .de_opt_utils import check_opt_bounds, move_to_optimized_positions -def omea_evaluation( - motors, bounds, popsize, num_interm_vals, num_scans_at_once, uids, flyer_name, intensity_name, db -): +def omea_evaluation(motors, bounds, popsize, num_interm_vals, num_scans_at_once, uids, flyer_name, intensity_name, db): if motors is not None: # hardware # get the data from databroker @@ -149,9 +147,7 @@ def rand_1(pop, popsize, target_indx, mut, bounds): for elem, param in x_1.items(): v_donor[elem] = {} for param_name in param.keys(): - v_donor[elem][param_name] = x_1[elem][param_name] + mut * ( - x_2[elem][param_name] - x_3[elem][param_name] - ) + v_donor[elem][param_name] = x_1[elem][param_name] + mut * (x_2[elem][param_name] - x_3[elem][param_name]) v_donor = ensure_bounds(vec=v_donor, bounds=bounds) return v_donor @@ -169,9 +165,7 @@ def best_1(pop, popsize, target_indx, mut, bounds, ind_sol): for elem, param in x_best.items(): v_donor[elem] = {} for param_name in param.items(): - v_donor[elem][param_name] = x_best[elem][param_name] + mut * ( - x_1[elem][param_name] - x_2[elem][param_name] - ) + v_donor[elem][param_name] = x_best[elem][param_name] + mut * (x_1[elem][param_name] - x_2[elem][param_name]) v_donor = ensure_bounds(vec=v_donor, bounds=bounds) return v_donor @@ -182,9 +176,7 @@ def mutate(population, strategy, mut, bounds, ind_sol): if strategy == "rand/1": v_donor = rand_1(pop=population, popsize=len(population), target_indx=i, mut=mut, bounds=bounds) elif strategy == "best/1": - v_donor = best_1( - pop=population, popsize=len(population), target_indx=i, mut=mut, bounds=bounds, ind_sol=ind_sol - ) + v_donor = best_1(pop=population, popsize=len(population), target_indx=i, mut=mut, bounds=bounds, ind_sol=ind_sol) # elif strategy == 'current-to-best/1': # v_donor = current_to_best_1(population, len(population), i, mut, bounds, ind_sol) # elif strategy == 'best/2': @@ -544,9 +536,7 @@ def optimization_plan( cross_trial_pop = crossover(population=pop_positions, mutated_indv=mutated_trial_pop, crosspb=crosspb) # select if opt_type == "hardware": - select_positions = yield from create_selection_params( - motors=motors, population=None, cross_indv=cross_trial_pop - ) + select_positions = yield from create_selection_params(motors=motors, population=None, cross_indv=cross_trial_pop) uid_list = yield from fly_plan( motors=motors, detector=detector, @@ -690,6 +680,4 @@ def optimization_plan( print(f"Convergence list: {best_fitness}") - yield from bp.count( - [], md={"best_fitness": best_fitness, "optimized_positions": optimized_positions, "uids": all_uids} - ) + yield from bp.count([], md={"best_fitness": best_fitness, "optimized_positions": optimized_positions, "uids": all_uids}) diff --git a/bloptools/de/hardware_flyer.py b/bloptools/de/hardware_flyer.py index 98ee9fa..5406b4f 100644 --- a/bloptools/de/hardware_flyer.py +++ b/bloptools/de/hardware_flyer.py @@ -121,9 +121,7 @@ def collect(self): motor_dict.update( { f"{self.name}_{motor_name}_velocity": self.velocities[motor_name], - f"{self.name}_{motor_name}_position": self.watch_positions[motor_name][field_name][ - ind - ], + f"{self.name}_{motor_name}_position": self.watch_positions[motor_name][field_name][ind], } ) data = {f"{self.name}_intensity": self.watch_intensities[ind]} diff --git a/bloptools/experiments/nsls2/tes.py b/bloptools/experiments/nsls2/tes.py index 33d3d89..c0a495a 100644 --- a/bloptools/experiments/nsls2/tes.py +++ b/bloptools/experiments/nsls2/tes.py @@ -51,9 +51,7 @@ def fitness(entry, args): background = args["background"] - x_min, x_max, y_min, y_max, separability = utils.get_beam_stats( - image - background, beam_prop=args["beam_prop"] - ) + x_min, x_max, y_min, y_max, separability = utils.get_beam_stats(image - background, beam_prop=args["beam_prop"]) n_y, n_x = image.shape diff --git a/bloptools/utils.py b/bloptools/utils.py index 380ca69..82953af 100644 --- a/bloptools/utils.py +++ b/bloptools/utils.py @@ -53,9 +53,7 @@ def get_routing(origin, points): delay_matrix = np.sqrt(np.square(rel_points[:, None, :] - rel_points[None, :, :]).sum(axis=-1)) delay_matrix = (1e3 * delay_matrix).astype(int) # it likes integers idk - manager = pywrapcp.RoutingIndexManager( - len(_points), 1, 0 - ) # number of depots, number of salesmen, starting index + manager = pywrapcp.RoutingIndexManager(len(_points), 1, 0) # number of depots, number of salesmen, starting index routing = pywrapcp.RoutingModel(manager) def delay_callback(from_index, to_index): diff --git a/docs/wip/custom-acquisition.ipynb b/docs/wip/custom-acquisition.ipynb index 61e4752..dbdae7f 100644 --- a/docs/wip/custom-acquisition.ipynb +++ b/docs/wip/custom-acquisition.ipynb @@ -31,22 +31,15 @@ "import numpy as np\n", "\n", "\n", - "\n", - "\n", - "\n", - "\n", "def acquisition(dofs, inputs, dets):\n", - "\n", " uid = yield from bp.list_scan\n", "\n", " for x in inputs:\n", + " res = np.sqrt(1 + np.square(x).sum()) # our resolution is\n", "\n", - " res = np.sqrt(1 + np.square(x).sum()) # our resolution is \n", - "\n", - " nu_sigma = np.sqrt(1 + x1 ** 2 + (x2 - 1) ** 2)\n", + " nu_sigma = np.sqrt(1 + x1**2 + (x2 - 1) ** 2)\n", "\n", - " flux = np.exp(-0.5*np.square((nu-nu_0)/nu_sigma))\n", - "\n" + " flux = np.exp(-0.5 * np.square((nu - nu_0) / nu_sigma))" ] }, { @@ -64,15 +57,14 @@ "\n", "nu = np.linspace(90, 110, 256)\n", "\n", - "for x1, x2 in [(0,0), (-2, 2), (-1,0), (0, 1)]:\n", - "\n", - " nu_sigma = np.sqrt(1 + x1 ** 2 + (x2 - 1) ** 2)\n", + "for x1, x2 in [(0, 0), (-2, 2), (-1, 0), (0, 1)]:\n", + " nu_sigma = np.sqrt(1 + x1**2 + (x2 - 1) ** 2)\n", "\n", - " flux = np.exp(-0.5*np.square((nu-nu_0)/nu_sigma))\n", + " flux = np.exp(-0.5 * np.square((nu - nu_0) / nu_sigma))\n", "\n", " plt.plot(nu, flux, label=f\"(x1, x2) = ({x1}, {x2})\")\n", "\n", - "plt.legend()\n" + "plt.legend()" ] }, { diff --git a/docs/wip/passive-dofs.ipynb b/docs/wip/passive-dofs.ipynb index e92011f..89d601e 100644 --- a/docs/wip/passive-dofs.ipynb +++ b/docs/wip/passive-dofs.ipynb @@ -37,15 +37,11 @@ "x1 = x2 = np.linspace(-5.0, 5.0, 256)\n", "X1, X2 = np.meshgrid(x1, x2)\n", "\n", - "plt.pcolormesh(x1, \n", - " x2, \n", - " bloptools.experiments.tests.himmelblau(X1, X2), \n", - " norm=mpl.colors.LogNorm(), \n", - " shading=\"auto\")\n", + "plt.pcolormesh(x1, x2, bloptools.experiments.tests.himmelblau(X1, X2), norm=mpl.colors.LogNorm(), shading=\"auto\")\n", "\n", "plt.xlabel(\"x1\")\n", "plt.ylabel(\"x2\")\n", - "plt.colorbar()\n" + "plt.colorbar()" ] }, { @@ -59,7 +55,7 @@ "\n", "tr = TimeReadback(name=\"timestamp\")\n", "\n", - "tr.read()\n" + "tr.read()" ] }, { @@ -80,8 +76,8 @@ "source": [ "import bloptools\n", "\n", - "dofs = bloptools.objects.get_dummy_dofs(n=2) # get a list of two DOFs\n", - "bounds = [(-5.0, +5.0), (-5.0, +5.0)]\n" + "dofs = bloptools.objects.get_dummy_dofs(n=2) # get a list of two DOFs\n", + "bounds = [(-5.0, +5.0), (-5.0, +5.0)]" ] }, { @@ -102,7 +98,7 @@ "source": [ "import bloptools\n", "\n", - "dofs = bloptools.objects.get_dummy_dofs(n=2) # get a list of two DOFs\n", + "dofs = bloptools.objects.get_dummy_dofs(n=2) # get a list of two DOFs\n", "bounds = [(-5.0, +5.0), (-5.0, +5.0)]\n", "\n", "from bloptools.objects import TimeReadback\n", @@ -112,18 +108,16 @@ "tr.read()\n", "\n", "\n", - "\n", "def digestion(db, uid):\n", - "\n", " table = db[uid].table()\n", " products = pd.DataFrame()\n", "\n", " for index, entry in table.iterrows():\n", - "\n", " products.loc[index, \"himmelblau\"] = bloptools.test_functions.himmelblau(entry.x1, entry.x2)\n", "\n", " return products\n", "\n", + "\n", "from bloptools.tasks import Task\n", "\n", "task = Task(key=\"himmelblau\", kind=\"min\")" @@ -171,15 +165,15 @@ "%run -i ../../../examples/prepare_bluesky.py # prepare the bluesky environment\n", "\n", "boa = bloptools.bayesian.Agent(\n", - " dofs=dofs,\n", - " bounds=bounds,\n", - " passive_dims=[tr],\n", - " tasks=task,\n", - " digestion=digestion,\n", - " db=db, \n", - " )\n", + " dofs=dofs,\n", + " bounds=bounds,\n", + " passive_dims=[tr],\n", + " tasks=task,\n", + " digestion=digestion,\n", + " db=db,\n", + ")\n", "\n", - "RE(boa.initialize(init_scheme='quasi-random', n_init=32))" + "RE(boa.initialize(init_scheme=\"quasi-random\", n_init=32))" ] }, { @@ -239,7 +233,7 @@ }, "outputs": [], "source": [ - "RE(boa.learn(strategy='esti', n_iter=4))\n", + "RE(boa.learn(strategy=\"esti\", n_iter=4))\n", "boa.plot_tasks()" ] }, @@ -261,7 +255,7 @@ }, "outputs": [], "source": [ - "RE(boa.learn(strategy='esti', n_iter=16))\n", + "RE(boa.learn(strategy=\"esti\", n_iter=16))\n", "boa.plot_tasks()" ] }, @@ -283,7 +277,7 @@ }, "outputs": [], "source": [ - "RE(boa.learn(strategy='esti', n_iter=32))\n", + "RE(boa.learn(strategy=\"esti\", n_iter=32))\n", "boa.plot_tasks()" ] }, diff --git a/examples/benchmark.py b/examples/benchmark.py index f0c07ae..1c7db37 100644 --- a/examples/benchmark.py +++ b/examples/benchmark.py @@ -38,9 +38,7 @@ plt.plot( timestamps - timestamps[0], [ - np.nanmax(bo.data.fitness.values[: i + 1]) - if not all(np.isnan(bo.data.fitness.values[: i + 1])) - else np.nan + np.nanmax(bo.data.fitness.values[: i + 1]) if not all(np.isnan(bo.data.fitness.values[: i + 1])) else np.nan for i in range(len(bo.data.fitness.values)) ], ) diff --git a/pyproject.toml b/pyproject.toml index 3239179..d852d68 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [tool.black] -line-length = 115 +line-length = 125 include = '\.pyi?$' exclude = ''' /(