diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 183bfad..83ac879 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -12,7 +12,6 @@ jobs: strategy: matrix: host-os: ["ubuntu-latest"] - # host-os: ["ubuntu-latest", "macos-latest", "windows-latest"] python-version: ["3.8", "3.9", "3.10"] fail-fast: false diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index c15e1d0..d1d643a 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -2,6 +2,7 @@ import time as ttime import warnings from collections import OrderedDict +from collections.abc import Mapping import bluesky.plan_stubs as bps import bluesky.plans as bp # noqa F401 @@ -17,7 +18,9 @@ from matplotlib.patches import Patch from .. import utils -from . import acquisition, models +from . import models +from .acquisition import default_acquisition_plan +from .digestion import default_digestion_function warnings.filterwarnings("ignore", category=botorch.exceptions.warnings.InputDataWarning) @@ -27,46 +30,100 @@ DEFAULT_COLORMAP = "turbo" -def default_acquisition_plan(dofs, inputs, dets): - uid = yield from bp.list_scan(dets, *[_ for items in zip(dofs, np.atleast_2d(inputs).T) for _ in items]) - return uid - - -def default_digestion_plan(db, uid): - return db[uid].table() - +MAX_TEST_INPUTS = 2**11 -# 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 self will try to optimize over -MAX_TEST_INPUTS = 2**11 +TASK_CONFIG = {} -AVAILABLE_ACQFS = { +ACQ_FUNC_CONFIG = { + "quasi-random": { + "identifiers": ["qr", "quasi-random"], + "pretty_name": "Quasi-random", + "description": "Sobol-sampled quasi-random points.", + }, "expected_mean": { "identifiers": ["em", "expected_mean"], + "pretty_name": "Expected mean", + "description": "The expected value at each input.", }, "expected_improvement": { "identifiers": ["ei", "expected_improvement"], + "pretty_name": "Expected improvement", + "description": r"The expected value of max(f(x) - \nu, 0), where \nu is the current maximum.", }, "probability_of_improvement": { "identifiers": ["pi", "probability_of_improvement"], + "pretty_name": "Probability of improvement", + "description": "The probability that this input improves on the current maximum.", }, "upper_confidence_bound": { "identifiers": ["ucb", "upper_confidence_bound"], - "default_args": {"beta": 4}, + "default_args": {"z": 2}, + "pretty_name": "Upper confidence bound", + "description": r"The expected value, plus some multiple of the uncertainty (typically \mu + 2\sigma).", }, } +TASK_TRANSFORMS = {"log": lambda x: np.log(x)} + + +def _validate_and_prepare_dofs(dofs): + for dof in dofs: + if not isinstance(dof, Mapping): + raise ValueError("Supplied dofs must be an iterable of mappings (e.g. a dict)!") + if "device" not in dof.keys(): + raise ValueError("Each DOF must have a device!") + + dof["device"].kind = "hinted" + + if "limits" not in dof.keys(): + dof["limits"] = (-np.inf, np.inf) + dof["limits"] = tuple(np.array(dof["limits"], dtype=float)) + + # dofs are passive by default + dof["kind"] = dof.get("kind", "passive") + if dof["kind"] not in ["active", "passive"]: + raise ValueError('DOF kinds must be one of "active" or "passive"') + + dof["mode"] = dof.get("mode", "on" if dof["kind"] == "active" else "off") + if dof["mode"] not in ["on", "off"]: + raise ValueError('DOF modes must be one of "on" or "off"') + + dof_names = [dof["device"].name for dof in dofs] + + # check that dof names are unique + unique_dof_names, counts = np.unique(dof_names, return_counts=True) + duplicate_dof_names = unique_dof_names[counts > 1] + if len(duplicate_dof_names) > 0: + raise ValueError(f'Duplicate name(s) in supplied dofs: "{duplicate_dof_names}"') + + return list(dofs) + + +def _validate_and_prepare_tasks(tasks): + for task in tasks: + if not isinstance(task, Mapping): + raise ValueError("Supplied tasks must be an iterable of mappings (e.g. a dict)!") + if task["kind"] not in ["minimize", "maximize"]: + raise ValueError('"mode" must be specified as either "minimize" or "maximize"') + if "weight" not in task.keys(): + task["weight"] = 1 + if "limits" not in task.keys(): + task["limits"] = (-np.inf, np.inf) + + task_keys = [task["key"] for task in tasks] + unique_task_keys, counts = np.unique(task_keys, return_counts=True) + duplicate_task_keys = unique_task_keys[counts > 1] + if len(duplicate_task_keys) > 0: + raise ValueError(f'Duplicate key(s) in supplied tasks: "{duplicate_task_keys}"') + + return list(tasks) + class Agent: def __init__( self, - active_dofs, - active_dof_bounds, + dofs, tasks, db, **kwargs, @@ -77,11 +134,11 @@ def __init__( Parameters ---------- dofs : iterable of ophyd objects - The degrees of freedom that the self can control, which determine the output of the model. + The degrees of freedom that the agent can control, which determine the output of the model. bounds : iterable of lower and upper bounds The bounds on each degree of freedom. This should be an array of shape (n_dofs, 2). tasks : iterable of tasks - The tasks which the self will try to optimize. + The tasks which the agent will try to optimize. acquisition : Bluesky plan generator that takes arguments (dofs, inputs, dets) A plan that samples the beamline for some given inputs. digestion : function that takes arguments (db, uid) @@ -89,171 +146,60 @@ def __init__( db : A databroker instance. """ - self.active_dofs = list(np.atleast_1d(active_dofs)) - self.passive_dofs = list(np.atleast_1d(kwargs.get("passive_dofs", []))) - - for dof in self.dofs: - dof.kind = "hinted" - - self.active_dof_bounds = np.atleast_2d(active_dof_bounds).astype(float) - self.tasks = np.atleast_1d(tasks) + # DOFs are parametrized by kind ("active" or "passive") and mode ("on" or "off") + # + # below are the behaviors of DOFs of each kind and mode: + # + # "read": the agent will read the input on every acquisition (all dofs are always read) + # "move": the agent will try to set and optimize over these (there must be at least one of these) + # "input" means that the agent will use the value to make its posterior + # + # + # active passive + # +---------------------+---------------+ + # on | read, input, move | read, input | + # +---------------------+---------------+ + # off | read | read | + # +---------------------+---------------+ + # + # + + self.dofs = _validate_and_prepare_dofs(np.atleast_1d(dofs)) + self.tasks = _validate_and_prepare_tasks(np.atleast_1d(tasks)) self.db = db self.verbose = kwargs.get("verbose", False) - self.ignore_acquisition_errors = kwargs.get("ignore_acquisition_errors", False) - + self.allow_acquisition_errors = kwargs.get("allow_acquisition_errors", True) self.initialization = kwargs.get("initialization", None) self.acquisition_plan = kwargs.get("acquisition_plan", default_acquisition_plan) - self.digestion = kwargs.get("digestion", default_digestion_plan) + self.digestion = kwargs.get("digestion", default_digestion_function) + self.dets = list(np.atleast_1d(kwargs.get("dets", []))) - self.decoherence = kwargs.get("decoherence", False) - - self.tolerate_acquisition_errors = kwargs.get("tolerate_acquisition_errors", True) - - self.acquisition = acquisition.Acquisition() - - self.dets = np.atleast_1d(kwargs.get("detectors", [])) - - for i, task in enumerate(self.tasks): - task.index = i - - self.n_tasks = len(self.tasks) - - self.training_iter = kwargs.get("training_iter", 256) - - # make some test points for sampling - - self.normalized_test_active_inputs = utils.normalized_sobol_sampler(n=MAX_TEST_INPUTS, d=self.n_active_dofs) - - n_per_active_dim = int(np.power(MAX_TEST_INPUTS, 1 / self.n_active_dofs)) - - self.normalized_test_active_inputs_grid = np.swapaxes( - np.r_[np.meshgrid(*self.n_active_dofs * [np.linspace(0, 1, n_per_active_dim)])], 0, -1 - ) + self.acq_func_config = kwargs.get("acq_func_config", ACQ_FUNC_CONFIG) self.table = pd.DataFrame() self._initialized = False self._train_models = True - self.a_priori_hypers = None - 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): + def _subset_inputs_sampler(self, kind=None, mode=None, 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_dofs)) - - @property - def dofs(self): - return np.append(self.active_dofs, self.passive_dofs) - - @property - def n_active_dofs(self): - return len(self.active_dofs) - - @property - def n_passive_dofs(self): - return len(self.passive_dofs) - - @property - def n_dofs(self): - return self.n_active_dofs + self.n_passive_dofs - - @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): - # return botorch.models.transforms.input.Normalize(d=self.n_dofs) - - @property - def input_transform(self): - coefficient = torch.tensor(self.dof_bounds.ptp(axis=1)).unsqueeze(0) - offset = torch.tensor(self.dof_bounds.min(axis=1)).unsqueeze(0) - return botorch.models.transforms.input.AffineInputTransform(d=self.n_dofs, coefficient=coefficient, offset=offset) - - def save_data(self, filepath="./self_data.h5"): - """ - Save the sampled inputs and targets of the self to a file, which can be used - to initialize a future self. - """ - - self.table.to_hdf(filepath, key="table") - - 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_active_dof hypercube using Sobol sampling. - """ - 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_dofs, scramble=True).random(n=min_power_of_two)[subset] - - def _set_hypers(self, hypers): - for task in self.tasks: - task.regressor.load_state_dict(hypers[task.name]) - self.classifier.load_state_dict(hypers["classifier"]) - - @property - def hypers(self): - hypers = {"classifier": {}} - for key, value in self.classifier.state_dict().items(): - hypers["classifier"][key] = value - for task in self.tasks: - hypers[task.name] = {} - for key, value in task.regressor.state_dict().items(): - hypers[task.name][key] = value - - return hypers - - def save_hypers(self, filepath): - hypers = self.hypers - with h5py.File(filepath, "w") as f: - for model_key in hypers.keys(): - f.create_group(model_key) - for param_key, param_value in hypers[model_key].items(): - f[model_key].create_dataset(param_key, data=param_value) - - @staticmethod - def load_hypers(filepath): - hypers = {} - with h5py.File(filepath, "r") as f: - for model_key in f.keys(): - hypers[model_key] = OrderedDict() - for param_key, param_value in f[model_key].items(): - hypers[model_key][param_key] = torch.tensor(np.atleast_1d(param_value[()])) - return hypers + transform = self._subset_input_transform(kind=kind, mode=mode) + return transform.untransform(utils.normalized_sobol_sampler(n, d=self._len_subset_dofs(kind=kind, mode=mode))) def initialize( self, - acqf=None, + acq_func=None, n_init=4, data=None, hypers=None, ): """ An initialization plan for the self. - This must be run before the self can learn. + This must be run before the agent can learn. It should be passed to a Bluesky RunEngine. """ @@ -271,12 +217,12 @@ def initialize( self.tell(new_table=data) # now let's get bayesian - elif acqf in ["qr"]: + elif acq_func in ["qr"]: yield from self.learn("qr", n_iter=1, n_per_iter=n_init, route=True) else: raise Exception( - """Could not initialize model! Either load a table, or specify an acqf from: + """Could not initialize model! Either load a table, or specify an acq_func from: ['qr'].""" ) @@ -284,64 +230,62 @@ def initialize( def tell(self, new_table=None, append=True, train=True, **kwargs): """ - Inform the self about new inputs and targets for the model. + Inform the agent about new inputs and targets for the model. """ new_table = pd.DataFrame() if new_table is None else new_table - 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)) - skew_dims = [tuple(np.arange(self.n_active_dofs))] + fitnesses = self.task_fitnesses # computes from self.table - if not train: - hypers = self.hypers + # update fitness estimates + self.table.loc[:, fitnesses.columns] = fitnesses.values + self.table.loc[:, "total_fitness"] = fitnesses.values.sum(axis=1) - for task in self.tasks: - task.targets = self.targets.loc[:, task.name] + skew_dims = [tuple(np.arange(self._len_subset_dofs(mode="on")))] - task.feasibility = self.feasible_for_all_tasks + if self._initialized: + cached_hypers = self.hypers - if not task.feasibility.sum() >= 2: - raise ValueError("There must be at least two feasible data points per task!") + feasibility = ~fitnesses.isna().any(axis=1) - train_inputs = torch.tensor(self.inputs.loc[task.feasibility].values).double().unsqueeze(0) - train_targets = torch.tensor(task.targets.loc[task.feasibility].values).double().unsqueeze(0).unsqueeze(-1) + if not feasibility.sum() >= 2: + raise ValueError("There must be at least two feasible data points per task!") - if train_inputs.ndim == 1: - train_inputs = train_inputs.unsqueeze(-1) - if train_targets.ndim == 1: - train_targets = train_targets.unsqueeze(-1) + inputs = self.inputs.loc[feasibility, self._subset_dof_names(mode="on")].values + train_inputs = torch.tensor(inputs).double() # .unsqueeze(0) + + for task in self.tasks: + targets = self.table.loc[feasibility, f'{task["key"]}_fitness'].values + train_targets = torch.tensor(targets).double().unsqueeze(-1) # .unsqueeze(0) likelihood = gpytorch.likelihoods.GaussianLikelihood( noise_constraint=gpytorch.constraints.Interval( - torch.tensor(task.MIN_NOISE_LEVEL).square(), - torch.tensor(task.MAX_NOISE_LEVEL).square(), + torch.tensor(1e-6).square(), + torch.tensor(1e-2).square(), ), ).double() - task.regressor = models.LatentGP( + outcome_transform = botorch.models.transforms.outcome.Standardize(m=1) # , batch_shape=torch.Size((1,))) + + task["model"] = models.LatentGP( train_inputs=train_inputs, train_targets=train_targets, likelihood=likelihood, skew_dims=skew_dims, - input_transform=self.input_transform, - outcome_transform=botorch.models.transforms.outcome.Standardize(m=1, batch_shape=torch.Size((1,))), + input_transform=self._subset_input_transform(mode="on"), + outcome_transform=outcome_transform, ).double() - 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))) - + # this ensures that we have equal weight between task fitness and feasibility fitness self.task_scalarization = botorch.acquisition.objective.ScalarizedPosteriorTransform( - weights=torch.tensor([*[task.weight for task in self.tasks], log_feas_prob_weight]).double(), + weights=torch.tensor([*torch.ones(self.n_tasks), self.fitness_variance.sum().sqrt()]).double(), offset=0, ) dirichlet_likelihood = gpytorch.likelihoods.DirichletClassificationLikelihood( - torch.as_tensor(self.feasible_for_all_tasks.values).long(), learn_additional_noise=True + torch.tensor(feasibility).long(), learn_additional_noise=True ).double() self.classifier = models.LatentDirichletClassifier( @@ -349,139 +293,340 @@ def tell(self, new_table=None, append=True, train=True, **kwargs): train_targets=dirichlet_likelihood.transformed_targets.transpose(-1, -2).double(), skew_dims=skew_dims, likelihood=dirichlet_likelihood, - input_transform=self.input_transform, + input_transform=self._subset_input_transform(mode="on"), ).double() - self.classifier_mll = gpytorch.mlls.ExactMarginalLogLikelihood(self.classifier.likelihood, self.classifier) - - self.feas_model = botorch.models.deterministic.GenericDeterministicModel( - f=lambda X: -self.classifier.log_prob(X).square() - ) - if self.a_priori_hypers is not None: self._set_hypers(self.a_priori_hypers) elif not train: - self._set_hypers(hypers) + self._set_hypers(cached_hypers) else: - self.train_models() + try: + self.train_models() + except botorch.exceptions.errors.ModelFittingError: + if self._initialized: + self._set_hypers(cached_hypers) + else: + raise RuntimeError("Could not fit model on initialization!") + + feasibility_fitness_model = botorch.models.deterministic.GenericDeterministicModel( + f=lambda X: -self.classifier.log_prob(X).square() + ) + + self.model_list = botorch.models.model.ModelList(*[task["model"] for task in self.tasks], feasibility_fitness_model) + + @property + def task_fitnesses(self): + df = pd.DataFrame(index=self.table.index) + for task in self.tasks: + name = f'{task["key"]}_fitness' + + df.loc[:, name] = task["weight"] * self.table.loc[:, task["key"]] + + # check that task values are inside acceptable values + valid = (df.loc[:, name] > task["limits"][0]) & (df.loc[:, name] < task["limits"][1]) + + # transform if needed + if "transform" in task.keys(): + if task["transform"] == "log": + valid &= df.loc[:, name] > 0 + df.loc[valid, name] = np.log(df.loc[valid, name]) + df.loc[~valid, name] = np.nan + + if task["kind"] == "minimize": + df.loc[valid, name] *= -1 + return df + + def _dof_kind_mask(self, kind=None): + return [dof["kind"] == kind if kind is not None else True for dof in self.dofs] + + def _dof_mode_mask(self, mode=None): + return [dof["mode"] == mode if mode is not None else True for dof in self.dofs] + + def _dof_mask(self, kind=None, mode=None): + return [(k and m) for k, m in zip(self._dof_kind_mask(kind), self._dof_mode_mask(mode))] + + def _subset_dofs(self, kind=None, mode=None): + return [dof for dof, m in zip(self.dofs, self._dof_mask(kind, mode)) if m] + + def _len_subset_dofs(self, kind=None, mode=None): + return len(self._subset_dofs(kind, mode)) + + def _subset_devices(self, kind=None, mode=None): + return [dof["device"] for dof in self._subset_dofs(kind, mode)] + + def _read_subset_devices(self, kind=None, mode=None): + return [device.read()[device.name]["value"] for device in self._subset_devices(kind, mode)] + + def _subset_dof_names(self, kind=None, mode=None): + return [device.name for device in self._subset_devices(kind, mode)] + + def _subset_dof_limits(self, kind=None, mode=None): + dofs_subset = self._subset_dofs(kind, mode) + if len(dofs_subset) > 0: + return torch.tensor([dof["limits"] for dof in dofs_subset], dtype=torch.float64).T + return torch.empty((2, 0)) + + def test_inputs(self, n=MAX_TEST_INPUTS): + return utils.sobol_sampler(self._acq_func_bounds, n=n) + + @property + def test_inputs_grid(self): + n_side = int(MAX_TEST_INPUTS ** (1 / self._len_subset_dofs(kind="active", mode="on"))) + return torch.tensor( + np.r_[ + np.meshgrid( + *[ + np.linspace(*dof["limits"], n_side) + if dof["kind"] == "active" + else dof["device"].read()[dof["device"].name]["value"] + for dof in self._subset_dofs(mode="on") + ] + ) + ] + ).swapaxes(0, -1) + + @property + def _acq_func_bounds(self): + return torch.tensor( + [ + dof["limits"] if dof["kind"] == "active" else tuple(2 * [dof["device"].read()[dof["device"].name]["value"]]) + for dof in self.dofs + if dof["mode"] == "on" + ] + ).T + + @property + def n_tasks(self): + return len(self.tasks) + + @property + def det_names(self): + return [det.name for det in self.dets] + + @property + def task_keys(self): + return [task["key"] for task in self.tasks] + + @property + def task_models(self): + return [task["model"] for task in self.tasks] + + @property + def task_weights(self): + return torch.tensor([task["weight"] for task in self.tasks], dtype=torch.float64) + + @property + def task_signs(self): + return torch.tensor([(1 if task["kind"] == "maximize" else -1) for task in self.tasks], dtype=torch.long) + + def _subset_input_transform(self, kind=None, mode=None): + limits = self._subset_dof_limits(kind, mode) + offset = limits.min(dim=0).values + coefficient = limits.max(dim=0).values - offset + return botorch.models.transforms.input.AffineInputTransform( + d=limits.shape[-1], coefficient=coefficient, offset=offset + ) + + def save_data(self, filepath="./self_data.h5"): + """ + Save the sampled inputs and targets of the agent to a file, which can be used + to initialize a future self. + """ + + self.table.to_hdf(filepath, key="table") + + def forget(self, index): + self.tell(new_table=self.table.drop(index=index), append=False, train=False) + + def sampler(self, n): + """ + Returns $n$ quasi-randomly sampled points on the [0,1] ^ n_active_dof hypercube using Sobol sampling. + """ + 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._len_subset_dofs(kind="active", mode="on"), scramble=True).random( + n=min_power_of_two + )[subset] + + def _set_hypers(self, hypers): + for task in self.tasks: + task["model"].load_state_dict(hypers[task["key"]]) + self.classifier.load_state_dict(hypers["classifier"]) + + @property + def hypers(self): + hypers = {"classifier": {}} + for key, value in self.classifier.state_dict().items(): + hypers["classifier"][key] = value + for task in self.tasks: + hypers[task["key"]] = {} + for key, value in task["model"].state_dict().items(): + hypers[task["key"]][key] = value + + return hypers - self.task_model = botorch.models.model.ModelList(*[task.regressor for task in self.tasks], self.feas_model) + def save_hypers(self, filepath): + hypers = self.hypers + with h5py.File(filepath, "w") as f: + for model_key in hypers.keys(): + f.create_group(model_key) + for param_key, param_value in hypers[model_key].items(): + f[model_key].create_dataset(param_key, data=param_value) + + @staticmethod + def load_hypers(filepath): + hypers = {} + with h5py.File(filepath, "r") as f: + for model_key in f.keys(): + hypers[model_key] = OrderedDict() + for param_key, param_value in f[model_key].items(): + hypers[model_key][param_key] = torch.tensor(np.atleast_1d(param_value[()])) + return hypers def train_models(self, **kwargs): t0 = ttime.monotonic() for task in self.tasks: - botorch.fit.fit_gpytorch_mll(task.regressor_mll, **kwargs) - botorch.fit.fit_gpytorch_mll(self.classifier_mll, **kwargs) + model = task["model"] + botorch.fit.fit_gpytorch_mll(gpytorch.mlls.ExactMarginalLogLikelihood(model.likelihood, model), **kwargs) + botorch.fit.fit_gpytorch_mll( + gpytorch.mlls.ExactMarginalLogLikelihood(self.classifier.likelihood, self.classifier), **kwargs + ) if self.verbose: print(f"trained models in {ttime.monotonic() - t0:.02f} seconds") - def get_acquisition_function(self, acqf_identifier="ei", return_metadata=False, acqf_args={}, **kwargs): + @property + def acq_func_info(self): + entries = [] + for k, d in self.acq_func_config.items(): + ret = "" + ret += f'{d["pretty_name"].upper()} (identifiers: {d["identifiers"]})\n' + ret += f'-> {d["description"]}' + entries.append(ret) + + print("\n\n".join(entries)) + + def get_acquisition_function(self, acq_func_identifier="ei", return_metadata=False, acq_func_args={}, **kwargs): if not self._initialized: - raise RuntimeError(f'Can\'t construct acquisition function "{acqf_identifier}" (the self is not initialized!)') + raise RuntimeError( + f'Can\'t construct acquisition function "{acq_func_identifier}" (the agent is not initialized!)' + ) - if acqf_identifier.lower() in AVAILABLE_ACQFS["expected_improvement"]["identifiers"]: - acqf = botorch.acquisition.analytic.LogExpectedImprovement( - self.task_model, - best_f=self.best_sum_of_tasks, + if acq_func_identifier.lower() in ACQ_FUNC_CONFIG["expected_improvement"]["identifiers"]: + acq_func = botorch.acquisition.analytic.LogExpectedImprovement( + self.model_list, + best_f=self.scalarized_fitness.max(), posterior_transform=self.task_scalarization, **kwargs, ) - acqf_meta = {"name": "expected improvement", "args": {}} + acq_func_meta = {"name": "expected improvement", "args": {}} - elif acqf_identifier.lower() in AVAILABLE_ACQFS["probability_of_improvement"]["identifiers"]: - acqf = botorch.acquisition.analytic.LogProbabilityOfImprovement( - self.task_model, - best_f=self.best_sum_of_tasks, + elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["probability_of_improvement"]["identifiers"]: + acq_func = botorch.acquisition.analytic.LogProbabilityOfImprovement( + self.model_list, + best_f=self.scalarized_fitness.max(), posterior_transform=self.task_scalarization, **kwargs, ) - acqf_meta = {"name": "probability of improvement", "args": {}} + acq_func_meta = {"name": "probability of improvement", "args": {}} - elif acqf_identifier.lower() in AVAILABLE_ACQFS["expected_mean"]["identifiers"]: - acqf = botorch.acquisition.analytic.UpperConfidenceBound( - self.task_model, + elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["expected_mean"]["identifiers"]: + acq_func = botorch.acquisition.analytic.UpperConfidenceBound( + self.model_list, beta=0, posterior_transform=self.task_scalarization, **kwargs, ) - acqf_meta = {"name": "expected mean"} + acq_func_meta = {"name": "expected mean"} - elif acqf_identifier.lower() in AVAILABLE_ACQFS["upper_confidence_bound"]["identifiers"]: - beta = AVAILABLE_ACQFS["upper_confidence_bound"]["default_args"]["beta"] - acqf = botorch.acquisition.analytic.UpperConfidenceBound( - self.task_model, + elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["upper_confidence_bound"]["identifiers"]: + beta = ACQ_FUNC_CONFIG["upper_confidence_bound"]["default_args"]["z"] ** 2 + acq_func = botorch.acquisition.analytic.UpperConfidenceBound( + self.model_list, beta=beta, posterior_transform=self.task_scalarization, **kwargs, ) - acqf_meta = {"name": "upper confidence bound", "args": {"beta": beta}} + acq_func_meta = {"name": "upper confidence bound", "args": {"beta": beta}} else: - raise ValueError(f'Unrecognized acquisition acqf_identifier "{acqf_identifier}".') + raise ValueError(f'Unrecognized acquisition acq_func_identifier "{acq_func_identifier}".') - return (acqf, acqf_meta) if return_metadata else acqf + return (acq_func, acq_func_meta) if return_metadata else acq_func - def ask(self, acqf_identifier="ei", n=1, route=True, return_metadata=False): - if acqf_identifier.lower() == "qr": - x = self.active_inputs_sampler(n=n) - acqf_meta = {"name": "quasi-random", "args": {}} + def ask(self, acq_func_identifier="ei", n=1, route=True, return_metadata=False): + if acq_func_identifier.lower() == "qr": + active_X = self._subset_inputs_sampler(n=n, kind="active", mode="on").squeeze(1).numpy() + acq_func_meta = {"name": "quasi-random", "args": {}} elif n == 1: - x, acqf_meta = self.ask_single(acqf_identifier, return_metadata=True) - return (x, acqf_meta) if return_metadata else x + active_X, acq_func_meta = self.ask_single(acq_func_identifier, return_metadata=True) elif n > 1: + active_x_list = [] for i in range(n): - x, acqf_meta = self.ask_single(acqf_identifier, return_metadata=True) + active_x, acq_func_meta = self.ask_single(acq_func_identifier, return_metadata=True) + active_x_list.append(active_x) if i < (n - 1): - task_samples = [task.regressor.posterior(torch.tensor(x)).sample().item() for task in self.tasks] + x = np.c_[active_x, acq_func_meta["passive_values"]] + task_samples = [task["model"].posterior(torch.tensor(x)).sample().item() for task in self.tasks] fantasy_table = pd.DataFrame( - np.append(x, task_samples)[None], columns=[*self.dof_names, *self.task_names] + np.c_[active_x, acq_func_meta["passive_values"], np.atleast_2d(task_samples)], + columns=[ + *self._subset_dof_names(kind="active", mode="on"), + *self._subset_dof_names(kind="passive", mode="on"), + *self.task_keys, + ], ) self.tell(fantasy_table, train=False) - x = self.active_inputs.iloc[-n:].values + active_X = np.concatenate(active_x_list, axis=0) + self.forget(self.table.index[-(n - 1) :]) - if n > 1: - self.forget(self.table.index[-(n - 1) :]) + if route: + active_X = active_X[utils.route(self._read_subset_devices(kind="active", mode="on"), active_X)] - if route: - x = x[utils.route(self.read_active_dofs, x)] - - return (x, acqf_meta) if return_metadata else x + return (active_X, acq_func_meta) if return_metadata else active_X def ask_single( self, - acqf_identifier="ei", + acq_func_identifier="ei", return_metadata=False, ): """ - The next $n$ points to sample, recommended by the self. + The next $n$ points to sample, recommended by the self. Returns """ t0 = ttime.monotonic() - acqf, acqf_meta = self.get_acquisition_function(acqf_identifier=acqf_identifier, return_metadata=True) + acq_func, acq_func_meta = self.get_acquisition_function( + acq_func_identifier=acq_func_identifier, return_metadata=True + ) BATCH_SIZE = 1 NUM_RESTARTS = 8 RAW_SAMPLES = 256 candidates, _ = botorch.optim.optimize_acqf( - acq_function=acqf, - bounds=torch.tensor(self.dof_bounds).T, + acq_function=acq_func, + bounds=self._acq_func_bounds, q=BATCH_SIZE, num_restarts=NUM_RESTARTS, raw_samples=RAW_SAMPLES, # used for intialization heuristic ) - x = candidates.detach().numpy()[..., self.dof_is_active_mask] + x = candidates.numpy().astype(float) + + active_x = x[..., [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")]] + passive_x = x[..., [dof["kind"] != "active" for dof in self._subset_dofs(mode="on")]] + + acq_func_meta["passive_values"] = passive_x if self.verbose: print(f"found point {x} in {ttime.monotonic() - t0:.02f} seconds") - return (x, acqf_meta) if return_metadata else x + return (active_x, acq_func_meta) if return_metadata else active_x def acquire(self, active_inputs): """ @@ -490,22 +635,27 @@ 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]) + active_devices = self._subset_devices(kind="active", mode="on") + passive_devices = [*self._subset_devices(kind="passive"), *self._subset_devices(kind="active", mode="off")] + + uid = yield from self.acquisition_plan( + active_devices, active_inputs.astype(float), [*self.dets, *passive_devices] + ) products = self.digestion(self.db, uid) # compute the fitness for each task - for index, entry in products.iterrows(): - for task in self.tasks: - products.loc[index, task.name] = task.get_fitness(entry) + # for index, entry in products.iterrows(): + # for task in self.tasks: + # products.loc[index, task["key"]] = getattr(entry, task["key"]) except Exception as error: - if not self.tolerate_acquisition_errors: + if not self.allow_acquisition_errors: raise error logging.warning(f"Error in acquisition/digestion: {repr(error)}") - products = pd.DataFrame(active_inputs, columns=self.active_dof_names) + products = pd.DataFrame(active_inputs, columns=self._subset_dof_names(kind="active", mode="on")) for task in self.tasks: - products.loc[:, task.name] = np.nan + products.loc[:, task["key"]] = np.nan if not len(active_inputs) == len(products): raise ValueError("The table returned by the digestion must be the same length as the sampled inputs!") @@ -514,7 +664,7 @@ def acquire(self, active_inputs): def learn( self, - acqf_identifier, + acq_func_identifier, n_iter=1, n_per_iter=1, reuse_hypers=True, @@ -529,247 +679,47 @@ def learn( """ for iteration in range(n_iter): - x, acqf_meta = self.ask(n=n_per_iter, acqf_identifier=acqf_identifier, return_metadata=True, **kwargs) + x, acq_func_meta = self.ask( + n=n_per_iter, acq_func_identifier=acq_func_identifier, return_metadata=True, **kwargs + ) new_table = yield from self.acquire(x) - new_table.loc[:, "acqf"] = acqf_meta["name"] + new_table.loc[:, "acq_func"] = acq_func_meta["name"] 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) - - 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 batch_dimension(self): - return self.dof_names.index("training_batch") if "training_batch" in self.dof_names else None - - @property - def test_inputs(self): - test_passive_inputs = self.read_passive_dofs[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.read_passive_dofs * np.ones( - (*self.test_active_inputs_grid.shape[:-1], self.n_passive_dofs) - ) - 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] + return self.table.loc[:, self._subset_dof_names(mode="on")].astype(float) @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 + def fitness_variance(self): + return torch.tensor(np.nanvar(self.task_fitnesses.values, axis=0)) @property - def feasible_for_all_tasks(self): - # TODO: make this more robust - # with pd.option_context("mode.use_inf_as_null", True): - feasible = ~self.targets.isna().any(axis=1) - for task in self.tasks: - if task.min is not None: - feasible &= self.targets.loc[:, task.name].values > task.transform(task.min) - return feasible + def scalarized_fitness(self): + return self.task_fitnesses.sum(axis=1) # @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 second way for now - return np.outer(self.read_passive_dofs, [1.0, 1.0]) - - @property - def dof_is_active_mask(self): - return np.r_[np.ones(self.n_active_dofs), np.zeros(self.n_passive_dofs)].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))] + # def best_sum_of_tasks_inputs(self): + # return self.inputs[np.nanargmax(self.task_fitnesses.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]) + yield from bps.mv(*[_ for items in zip(self._subset_dofs(kind="active"), 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) + # @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_active_dofs == 1: + if self._len_subset_dofs(kind="active", mode="on") == 1: self._plot_tasks_one_dof(**kwargs) - else: self._plot_tasks_many_dofs(**kwargs) - def plot_feasibility(self, **kwargs): - if self.n_active_dofs == 1: - self._plot_feas_one_dof(**kwargs) - - else: - self._plot_feas_many_dofs(**kwargs) - - def plot_acquisition(self, **kwargs): - if self.n_active_dofs == 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.feasible_for_all_tasks.astype(int), s=size) - - x = torch.tensor(self.test_inputs_grid.reshape(-1, self.n_dofs)).double() - log_prob = self.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.active_dof_bounds[0]) - - 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_dofs == 2 - - self.class_fig, self.class_axes = plt.subplots( - 1, 2, figsize=(8, 4), sharex=True, sharey=True, constrained_layout=True - ) - - for ax in self.class_axes.ravel(): - ax.set_xlabel(self.dofs[axes[0]].name) - ax.set_ylabel(self.dofs[axes[1]].name) - - data_ax = self.class_axes[0].scatter( - *self.inputs.values.T[:2], s=size, c=self.feasible_for_all_tasks.astype(int), vmin=0, vmax=1, cmap=cmap - ) - - if gridded: - x = torch.tensor(self.test_inputs_grid.reshape(-1, self.n_dofs)).double() - log_prob = self.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), - np.exp(log_prob).T, - shading=shading, - cmap=cmap, - vmin=0, - vmax=1, - ) - - else: - x = torch.tensor(self.test_inputs).double() - log_prob = self.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_fig.colorbar(data_ax, ax=self.class_axes[:2], location="bottom", aspect=32, shrink=0.8) - - 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): + def _plot_tasks_one_dof(self, size=16, lw=1e0): self.task_fig, self.task_axes = plt.subplots( self.n_tasks, 1, @@ -783,30 +733,37 @@ def _plot_tasks_one_dof(self, size=32, lw=1e0): for itask, task in enumerate(self.tasks): color = DEFAULT_COLOR_LIST[itask] - self.task_axes[itask].set_ylabel(task.name) + self.task_axes[itask].set_ylabel(task["key"]) - 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() + x = self.test_inputs_grid + task_posterior = task["model"].posterior(x) + task_mean = task_posterior.mean.detach().numpy() + task_sigma = task_posterior.variance.sqrt().detach().numpy() - 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) + self.task_axes[itask].scatter( + self.inputs.loc[:, self._subset_dof_names(kind="active", mode="on")], + self.table.loc[:, f'{task["key"]}_fitness'], + s=size, + color=color, + ) + + on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] - for z in [1, 2]: + for z in [0, 1, 2]: self.task_axes[itask].fill_between( - self.test_inputs_grid.ravel(), - (task_mean - z * task_sigma).ravel(), - (task_mean + z * task_sigma).ravel(), + x[..., on_dofs_are_active_mask].squeeze(), + (task_mean - z * task_sigma).squeeze(), + (task_mean + z * task_sigma).squeeze(), lw=lw, color=color, alpha=0.5**z, ) - self.task_axes[itask].set_xlim(*self.active_dof_bounds[0]) + self.task_axes[itask].set_xlim(self._subset_dofs(kind="active", mode="on")[0]["limits"]) - def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=32): + def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=16): if gridded is None: - gridded = self.n_dofs == 2 + gridded = self._len_subset_dofs(kind="active", mode="on") == 2 self.task_fig, self.task_axes = plt.subplots( self.n_tasks, @@ -818,142 +775,238 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL ) self.task_axes = np.atleast_2d(self.task_axes) - self.task_fig.suptitle(f"(x,y)=({self.dofs[axes[0]].name},{self.dofs[axes[1]].name})") + # self.task_fig.suptitle(f"(x,y)=({self.dofs[axes[0]].name},{self.dofs[axes[1]].name})") + + feasible = ~self.task_fitnesses.isna().any(axis=1) for itask, task in enumerate(self.tasks): - task_norm = mpl.colors.Normalize(*np.nanpercentile(task.targets, q=[1, 99])) + sampled_fitness = np.where(feasible, self.table.loc[:, f'{task["key"]}_fitness'].values, np.nan) + task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) + task_norm = mpl.colors.Normalize(task_vmin, task_vmax) - self.task_axes[itask, 0].set_ylabel(task.name) + # if task["transform"] == "log": + # task_norm = mpl.colors.LogNorm(task_vmin, task_vmax) + # else: + + self.task_axes[itask, 0].set_ylabel(f'{task["key"]}_fitness') self.task_axes[itask, 0].set_title("samples") self.task_axes[itask, 1].set_title("posterior mean") self.task_axes[itask, 2].set_title("posterior std. dev.") data_ax = self.task_axes[itask, 0].scatter( - *self.inputs.values.T[axes], s=size, c=task.targets, norm=task_norm, cmap=cmap + *self.inputs.values.T[axes], s=size, c=sampled_fitness, norm=task_norm, cmap=cmap ) - x = torch.tensor(self.test_inputs_grid).double() if gridded else torch.tensor(self.test_inputs).double() + x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) - 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_posterior = task["model"].posterior(x) + task_mean = task_posterior.mean + task_sigma = task_posterior.variance.sqrt() if gridded: + if not x.ndim == 3: + raise ValueError() self.task_axes[itask, 1].pcolormesh( - *np.swapaxes(self.test_inputs_grid, 0, -1), - task_mean.reshape(self.test_active_inputs_grid.shape[:-1]).T, + x[..., 0], + x[..., 1], + task_mean[..., 0].detach().numpy(), shading=shading, cmap=cmap, norm=task_norm, ) sigma_ax = self.task_axes[itask, 2].pcolormesh( - *np.swapaxes(self.test_inputs_grid, 0, -1), - task_sigma.reshape(self.test_inputs_grid.shape[:-1]).T, + x[..., 0], + x[..., 1], + task_sigma[..., 0].detach().numpy(), shading=shading, cmap=cmap, ) 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()[..., axes[0]], + x.detach().numpy()[..., axes[1]], + c=task_mean[..., 0].detach().numpy(), + s=size, + norm=task_norm, + cmap=cmap, + ) + sigma_ax = self.task_axes[itask, 2].scatter( + x.detach().numpy()[..., axes[0]], + x.detach().numpy()[..., axes[1]], + c=task_sigma[..., 0].detach().numpy(), + s=size, + 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) for ax in self.task_axes.ravel(): - ax.set_xlim(*self.active_dof_bounds[axes[0]]) - ax.set_ylim(*self.active_dof_bounds[axes[1]]) + ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) + ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) + + def plot_acquisition(self, acq_funcs=["ei"], **kwargs): + if self._len_subset_dofs(kind="active", mode="on") == 1: + self._plot_acq_one_dof(acq_funcs=acq_funcs, **kwargs) - def _plot_acq_one_dof(self, size=32, lw=1e0, **kwargs): - acqf_names = np.atleast_1d(kwargs.get("acqf", "ei")) + else: + self._plot_acq_many_dofs(acq_funcs=acq_funcs, **kwargs) + def _plot_acq_one_dof(self, acq_funcs, lw=1e0, **kwargs): self.acq_fig, self.acq_axes = plt.subplots( 1, - len(acqf_names), - figsize=(6 * len(acqf_names), 6), + len(acq_funcs), + figsize=(4 * len(acq_funcs), 4), sharex=True, constrained_layout=True, ) self.acq_axes = np.atleast_1d(self.acq_axes) - for iacqf, acqf_name in enumerate(acqf_names): - color = DEFAULT_COLOR_LIST[0] + for iacq_func, acq_func_identifier in enumerate(acq_funcs): + color = DEFAULT_COLOR_LIST[iacq_func] - acqf, acqf_meta = self.get_acquisition_function(acqf_name, return_metadata=True) + acq_func, acq_func_meta = self.get_acquisition_function(acq_func_identifier, 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) + x = self.test_inputs_grid + *input_shape, input_dim = x.shape + obj = acq_func.forward(x.reshape(-1, 1, input_dim)).reshape(input_shape) - if acqf_name in ["ei", "pi"]: + if acq_func_identifier in ["ei", "pi"]: 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[iacq_func].set_title(acq_func_meta["name"]) - self.acq_axes[iacqf].set_xlim(*self.active_dof_bounds[0]) + on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] + self.acq_axes[iacq_func].plot( + x[..., on_dofs_are_active_mask].squeeze(), obj.detach().numpy(), lw=lw, color=color + ) - 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_axes[iacq_func].set_xlim(self._subset_dofs(kind="active", mode="on")[0]["limits"]) + def _plot_acq_many_dofs( + self, acq_funcs, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=16, **kwargs + ): self.acq_fig, self.acq_axes = plt.subplots( 1, - len(acqf_names), - figsize=(4 * len(acqf_names), 5), + len(acq_funcs), + figsize=(4 * len(acq_funcs), 4), sharex=True, sharey=True, constrained_layout=True, ) if gridded is None: - gridded = self.n_active_dofs == 2 + gridded = self._len_subset_dofs(kind="active", mode="on") == 2 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})") + # self.acq_fig.suptitle(f"(x,y)=({self.dofs[axes[0]].name},{self.dofs[axes[1]].name})") - for iacqf, acqf_name in enumerate(acqf_names): - acqf, acqf_meta = self.get_acquisition_function(acqf_name, return_metadata=True) + x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) + *input_shape, input_dim = x.shape - 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) + for iacq_func, acq_func_identifier in enumerate(acq_funcs): + acq_func, acq_func_meta = self.get_acquisition_function(acq_func_identifier, return_metadata=True) - if acqf_name in ["ei", "pi"]: - obj = obj.exp() + obj = acq_func.forward(x.reshape(-1, 1, input_dim)).reshape(input_shape) + if acq_func_identifier in ["ei", "pi"]: + obj = obj.exp() - self.acq_axes[iacqf].set_title(acqf_meta["name"]) - obj_ax = self.acq_axes[iacqf].pcolormesh( - *np.swapaxes(self.test_inputs_grid, 0, -1)[axes], - obj.detach().numpy().reshape(grid_shape).T, + if gridded: + self.acq_axes[iacq_func].set_title(acq_func_meta["name"]) + obj_ax = self.acq_axes[iacq_func].pcolormesh( + x[..., 0], + x[..., 1], + obj.detach().numpy(), shading=shading, cmap=cmap, ) - self.acq_fig.colorbar(obj_ax, ax=self.acq_axes[iacqf], location="bottom", aspect=32, shrink=0.8) + self.acq_fig.colorbar(obj_ax, ax=self.acq_axes[iacq_func], 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( + self.acq_axes[iacq_func].set_title(acq_func_meta["name"]) + obj_ax = self.acq_axes[iacq_func].scatter( x.detach().numpy()[..., axes[0]], x.detach().numpy()[..., axes[1]], - c=obj.detach().numpy().reshape(inputs_shape), + c=obj.detach().numpy(), ) - self.acq_fig.colorbar(obj_ax, ax=self.acq_axes[iacqf], location="bottom", aspect=32, shrink=0.8) + self.acq_fig.colorbar(obj_ax, ax=self.acq_axes[iacq_func], location="bottom", aspect=32, shrink=0.8) 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]]) + ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) + ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) + + def plot_feasibility(self, **kwargs): + if self._len_subset_dofs(kind="active", mode="on") == 1: + self._plot_feas_one_dof(**kwargs) + + else: + self._plot_feas_many_dofs(**kwargs) + + def _plot_feas_one_dof(self, size=16, lw=1e0): + self.feas_fig, self.feas_ax = plt.subplots(1, 1, figsize=(4, 4), sharex=True, constrained_layout=True) + + x = self.test_inputs_grid + *input_shape, input_dim = x.shape + log_prob = self.classifier.log_prob(x.reshape(-1, 1, input_dim)).reshape(input_shape) + + self.feas_ax.scatter(self.inputs.values, ~self.task_fitnesses.isna().any(axis=1), s=size) + + on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] + + self.feas_ax.plot(x[..., on_dofs_are_active_mask].squeeze(), log_prob.exp().detach().numpy(), lw=lw) + + self.feas_ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[0]["limits"]) + + def _plot_feas_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, size=16, gridded=None): + self.feas_fig, self.feas_axes = plt.subplots(1, 2, figsize=(8, 4), sharex=True, sharey=True, constrained_layout=True) + + if gridded is None: + gridded = self._len_subset_dofs(kind="active", mode="on") == 2 + + data_ax = self.feas_axes[0].scatter( + *self.inputs.values.T[:2], + c=~self.task_fitnesses.isna().any(axis=1), + s=size, + vmin=0, + vmax=1, + cmap=cmap, + ) + + x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) + *input_shape, input_dim = x.shape + log_prob = self.classifier.log_prob(x.reshape(-1, 1, input_dim)).reshape(input_shape) + + if gridded: + self.feas_axes[1].pcolormesh( + x[..., 0], + x[..., 1], + log_prob.exp().detach().numpy(), + shading=shading, + cmap=cmap, + vmin=0, + vmax=1, + ) + + # self.acq_fig.colorbar(obj_ax, ax=self.feas_axes[iacq_func], location="bottom", aspect=32, shrink=0.8) + + else: + # self.feas_axes.set_title(acq_func_meta["name"]) + self.feas_axes[1].scatter( + x.detach().numpy()[..., axes[0]], + x.detach().numpy()[..., axes[1]], + c=log_prob.exp().detach().numpy(), + ) + + self.feas_fig.colorbar(data_ax, ax=self.feas_axes[:2], location="bottom", aspect=32, shrink=0.8) + + for ax in self.feas_axes.ravel(): + ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) + ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) def inspect_beam(self, index, border=None): im = self.images[index] @@ -987,16 +1040,18 @@ def plot_history(self, x_key="index", show_all_tasks=False): ) hist_axes = np.atleast_1d(hist_axes) - unique_strategies, acqf_index, acqf_inverse = np.unique(self.table.acqf, return_index=True, return_inverse=True) + unique_strategies, acq_func_index, acq_func_inverse = np.unique( + self.table.acq_func, return_index=True, return_inverse=True + ) - sample_colors = np.array(DEFAULT_COLOR_LIST)[acqf_inverse] + sample_colors = np.array(DEFAULT_COLOR_LIST)[acq_func_inverse] if show_all_tasks: for itask, task in enumerate(self.tasks): - y = task.targets.values + y = self.table.loc[:, f'{task["key"]}_fitness'].values hist_axes[itask].scatter(x, y, c=sample_colors) hist_axes[itask].plot(x, y, lw=5e-1, c="k") - hist_axes[itask].set_ylabel(task.name) + hist_axes[itask].set_ylabel(task["key"]) y = self.table.total_fitness @@ -1011,9 +1066,9 @@ def plot_history(self, x_key="index", show_all_tasks=False): hist_axes[-1].set_xlabel(x_key) handles = [] - for i_acqf, acqf in enumerate(unique_strategies): - # i_acqf = np.argsort(acqf_index)[i_handle] - handles.append(Patch(color=DEFAULT_COLOR_LIST[i_acqf], label=acqf)) + for i_acq_func, acq_func in enumerate(unique_strategies): + # i_acq_func = np.argsort(acq_func_index)[i_handle] + handles.append(Patch(color=DEFAULT_COLOR_LIST[i_acq_func], label=acq_func)) legend = hist_axes[0].legend(handles=handles, fontsize=8) legend.set_title("acquisition function") diff --git a/bloptools/bayesian/acquisition.py b/bloptools/bayesian/acquisition.py index 0a041dc..3354d19 100644 --- a/bloptools/bayesian/acquisition.py +++ b/bloptools/bayesian/acquisition.py @@ -1,37 +1,7 @@ -import torch -from botorch.acquisition.analytic import LogExpectedImprovement, LogProbabilityOfImprovement +import bluesky.plans as bp +import numpy as np -class Acquisition: - def __init__(self, *args, **kwargs): - ... - - @staticmethod - def log_expected_improvement(candidates, agent): - *input_shape, n_dim = candidates.shape - - x = torch.as_tensor(candidates.reshape(-1, 1, n_dim)).double() - - 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() +def default_acquisition_plan(dofs, inputs, dets): + uid = yield from bp.list_scan(dets, *[_ for items in zip(dofs, np.atleast_2d(inputs).T) for _ in items]) + return uid diff --git a/bloptools/bayesian/digestion.py b/bloptools/bayesian/digestion.py new file mode 100644 index 0000000..147d2d3 --- /dev/null +++ b/bloptools/bayesian/digestion.py @@ -0,0 +1,2 @@ +def default_digestion_function(db, uid): + return db[uid].table(fill=True) diff --git a/bloptools/bayesian/kernels.py b/bloptools/bayesian/kernels.py index db2d99f..dc46438 100644 --- a/bloptools/bayesian/kernels.py +++ b/bloptools/bayesian/kernels.py @@ -21,8 +21,7 @@ def __init__( self.num_inputs = num_inputs self.scale_output = scale_output - self.nu = kwargs.get("nu", 1.5) - self.batch_dimension = kwargs.get("batch_dimension", None) + self.nu = kwargs.get("nu", 2.5) if type(skew_dims) is bool: if skew_dims: @@ -172,4 +171,15 @@ def forward(self, x1, x2, diag=False, **params): distance = self.covar_dist(trans_x1, trans_x2, diag=diag, **params) - return (self.outputscale if self.scale_output else 1.0) * (1 + distance) * torch.exp(-distance) + if self.num_outputs == 1: + distance = distance.squeeze(0) + + outputscale = self.outputscale if self.scale_output else 1.0 + + # special cases of the Matern function + if self.nu == 0.5: + return outputscale * torch.exp(-distance) + if self.nu == 1.5: + return outputscale * (1 + distance) * torch.exp(-distance) + if self.nu == 2.5: + return outputscale * (1 + distance + distance**2 / 3) * torch.exp(-distance) diff --git a/bloptools/devices.py b/bloptools/devices.py index 5840fbd..0afc20e 100644 --- a/bloptools/devices.py +++ b/bloptools/devices.py @@ -1,28 +1,55 @@ import time as ttime -from ophyd import Component as Cpt -from ophyd import Device, Signal, SignalRO +import numpy as np +from ophyd import Signal, SignalRO +DEFAULT_BOUNDS = (-5.0, +5.0) -def dummy_dof(name): - return Signal(name=name, value=0.0) +class ReadOnlyError(Exception): + ... -def dummy_dofs(n=2): - return [dummy_dof(name=f"x{i+1}") for i in range(n)] +class DOF(Signal): + """ + Degree of freedom + """ + + ... + + +class DOFRO(DOF): + """ + Read-only degree of freedom + """ -def get_dummy_device(name="dofs", n=2): - components = {} + def put(self, value, *, timestamp=None, force=False): + raise ReadOnlyError(f'Cannot put, DOF "{self.name}" is read-only!') - for i in range(n): - components[f"x{i+1}"] = Cpt(Signal, value=i + 1) + def set(self, value, *, timestamp=None, force=False): + raise ReadOnlyError(f'Cannot set, DOF "{self.name}" is read-only!') - cls = type("DOF", (Device,), components) - device = cls(name=name) +class BrownianMotion(DOFRO): + """ + Read-only degree of freedom simulating brownian motion + """ + + def __init__(self, theta=0.95, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.theta = theta + self.old_t = ttime.monotonic() + self.old_y = 0.0 + + def get(self): + new_t = ttime.monotonic() + alpha = self.theta ** (new_t - self.old_t) + new_y = alpha * self.old_y + np.sqrt(1 - alpha**2) * np.random.standard_normal() - return [getattr(device, attr) for attr in device.read_attrs] + self.old_t = new_t + self.old_y = new_y + return new_y class TimeReadback(SignalRO): diff --git a/bloptools/test_functions.py b/bloptools/test_functions.py index 9b1983e..d61b7f8 100644 --- a/bloptools/test_functions.py +++ b/bloptools/test_functions.py @@ -2,24 +2,37 @@ def booth(x1, x2): + """ + The Booth function (https://en.wikipedia.org/wiki/Test_functions_for_optimization) + """ return (x1 + 2 * x2 - 7) ** 2 + (2 * x1 + x2 - 5) ** 2 def matyas(x1, x2): + """ + The Matyas function (https://en.wikipedia.org/wiki/Test_functions_for_optimization) + """ return 0.26 * (x1**2 + x2**2) - 0.48 * x1 * x2 def himmelblau(x1, x2): + """ + Himmelblau's function (https://en.wikipedia.org/wiki/Himmelblau%27s_function) + """ return (x1**2 + x2 - 11) ** 2 + (x1 + x2**2 - 7) ** 2 def constrained_himmelblau(x1, x2): - if x1**2 + x2**2 > 50: - return np.nan - return himmelblau(x1, x2) + """ + Himmelblau's function, returns NaN outside the constraint + """ + return np.where(x1**2 + x2**2 < 50, himmelblau(x1, x2), np.nan) def skewed_himmelblau(x1, x2): + """ + Himmelblau's function, with skewed coordinates + """ _x1 = 2 * x1 + x2 _x2 = 0.5 * (x1 - 2 * x2) @@ -27,20 +40,32 @@ def skewed_himmelblau(x1, x2): def bukin(x1, x2): + """ + Bukin function N.6 (https://en.wikipedia.org/wiki/Test_functions_for_optimization) + """ return 100 * np.sqrt(np.abs(x2 - 1e-2 * x1**2)) + 0.01 * np.abs(x1) def rastrigin(*x): + """ + The Rastrigin function in arbitrary dimensions (https://en.wikipedia.org/wiki/Rastrigin_function) + """ X = np.c_[x] return 10 * X.shape[-1] + (X**2 - 10 * np.cos(2 * np.pi * X)).sum(axis=1) def styblinski_tang(*x): + """ + Styblinski-Tang function in arbitrary dimensions (https://en.wikipedia.org/wiki/Test_functions_for_optimization) + """ X = np.c_[x] return 0.5 * (X**4 - 16 * X**2 + 5 * X).sum(axis=1) def ackley(*x): + """ + The Ackley function in arbitrary dimensions (https://en.wikipedia.org/wiki/Ackley_function) + """ X = np.c_[x] return ( -20 * np.exp(-0.2 * np.sqrt(0.5 * (X**2).sum(axis=1))) @@ -51,10 +76,16 @@ def ackley(*x): def gaussian_beam_waist(x1, x2): + """ + Simulating a misaligned Gaussian beam. The optimum is at (1, 1, 1, 1) + """ return np.sqrt(1 + 0.25 * (x1 - x2) ** 2 + 16 * (x1 + x2 - 2) ** 2) def himmelblau_digestion(db, uid): + """ + Digests Himmelblau's function into the feedback. + """ products = db[uid].table() for index, entry in products.iterrows(): @@ -65,7 +96,7 @@ def himmelblau_digestion(db, uid): def mock_kbs_digestion(db, uid): """ - Simulating a misaligned Gaussian beam. The optimum is at (1, 1, 1, 1) + Digests a beam waist and height into the feedback. """ products = db[uid].table() diff --git a/bloptools/tests/conftest.py b/bloptools/tests/conftest.py index 9e23e8c..a02f9ca 100644 --- a/bloptools/tests/conftest.py +++ b/bloptools/tests/conftest.py @@ -13,7 +13,6 @@ from sirepo_bluesky.srw_handler import SRWFileHandler from bloptools.bayesian import Agent -from bloptools.tasks import Task from .. import devices, test_functions @@ -57,27 +56,57 @@ def RE(db): @pytest.fixture(scope="function") def agent(db): """ - A simple agent maximizing two Styblinski-Tang functions + A simple agent minimizing Himmelblau's function """ - dofs = devices.dummy_dofs(n=2) - bounds = [(-5, 5), (-5, 5)] + dofs = [ + {"device": devices.DOF(name="x1"), "limits": (-5, 5), "kind": "active"}, + {"device": devices.DOF(name="x2"), "limits": (-5, 5), "kind": "active"}, + ] + + tasks = [ + {"key": "himmelblau", "kind": "minimize"}, + ] + + agent = Agent( + dofs=dofs, + tasks=tasks, + digestion=test_functions.himmelblau_digestion, + db=db, + verbose=True, + tolerate_acquisition_errors=False, + ) + + return agent + + +@pytest.fixture(scope="function") +def multitask_agent(db): + """ + A simple agent minimizing two Styblinski-Tang functions + """ def digestion(db, uid): products = db[uid].table() for index, entry in products.iterrows(): - products.loc[index, "loss1"] = test_functions.styblinski_tang(entry.x1, entry.x2) - products.loc[index, "loss2"] = test_functions.styblinski_tang(entry.x1, -entry.x2) + products.loc[index, "ST1"] = test_functions.styblinski_tang(entry.x1, entry.x2) + products.loc[index, "ST2"] = test_functions.styblinski_tang(entry.x1, -entry.x2) return products - tasks = [Task(key="loss1", kind="min"), Task(key="loss2", kind="min")] + dofs = [ + {"device": devices.DOF(name="x1"), "limits": (-5, 5), "kind": "active"}, + {"device": devices.DOF(name="x2"), "limits": (-5, 5), "kind": "active"}, + ] + + tasks = [ + {"key": "ST1", "kind": "minimize"}, + {"key": "ST2", "kind": "minimize"}, + ] agent = Agent( - active_dofs=dofs, - passive_dofs=[], - active_dof_bounds=bounds, + dofs=dofs, tasks=tasks, digestion=digestion, db=db, diff --git a/bloptools/tests/test_agent.py b/bloptools/tests/test_agent.py index 2d660b3..66657ce 100644 --- a/bloptools/tests/test_agent.py +++ b/bloptools/tests/test_agent.py @@ -3,7 +3,7 @@ import pytest -@pytest.mark.agent +@pytest.mark.test_func def test_writing_hypers(RE, agent): RE(agent.initialize("qr", n_init=32)) @@ -12,3 +12,14 @@ def test_writing_hypers(RE, agent): RE(agent.initialize("qr", n_init=8, hypers="hypers.h5")) os.remove("hypers.h5") + + +@pytest.mark.test_func +def test_writing_hypers_multitask(RE, multitask_agent): + RE(multitask_agent.initialize("qr", n_init=32)) + + multitask_agent.save_hypers("hypers.h5") + + RE(multitask_agent.initialize("qr", n_init=8, hypers="hypers.h5")) + + os.remove("hypers.h5") diff --git a/bloptools/tests/test_bayesian_shadow.py b/bloptools/tests/test_bayesian_shadow.py index d281582..001f6e8 100644 --- a/bloptools/tests/test_bayesian_shadow.py +++ b/bloptools/tests/test_bayesian_shadow.py @@ -1,34 +1,34 @@ -import numpy as np import pytest from sirepo_bluesky.sirepo_ophyd import create_classes import bloptools from bloptools.experiments.sirepo.tes import w9_digestion -from bloptools.tasks import Task @pytest.mark.shadow def test_bayesian_agent_tes_shadow(RE, db, shadow_tes_simulation): data, schema = shadow_tes_simulation.auth("shadow", "00000002") - classes, objects = create_classes(shadow_tes_simulation.data, connection=shadow_tes_simulation) + classes, objects = create_classes(connection=shadow_tes_simulation) globals().update(**objects) data["models"]["simulation"]["npoint"] = 100000 data["models"]["watchpointReport12"]["histogramBins"] = 32 - kb_dofs = [kbv.x_rot, kbv.offz] - kb_bounds = np.array([[-0.10, +0.10], [-0.50, +0.50]]) + dofs = [ + {"device": kbv.x_rot, "limits": (-0.1, 0.1), "kind": "active"}, + {"device": kbv.offz, "limits": (-0.5, 0.5), "kind": "active"}, + ] - 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)) + tasks = [ + {"key": "flux", "kind": "maximize"}, + {"key": "w9_fwhm_x", "kind": "minimize"}, + {"key": "w9_fwhm_y", "kind": "minimize"}, + ] 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], + dofs=dofs, + tasks=tasks, + dets=[w9], digestion=w9_digestion, db=db, ) diff --git a/bloptools/tests/test_bayesian_test_funcs.py b/bloptools/tests/test_bayesian_test_funcs.py deleted file mode 100644 index 04eeb55..0000000 --- a/bloptools/tests/test_bayesian_test_funcs.py +++ /dev/null @@ -1,50 +0,0 @@ -import pytest - -import bloptools -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.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") - - agent = bloptools.bayesian.Agent( - active_dofs=dofs, - passive_dofs=[], - active_dof_bounds=bounds, - tasks=[task], - digestion=himmelblau_digestion, - db=db, - ) - - RE(agent.initialize("qr", n_init=16)) - - RE(agent.learn("ei", n_iter=2)) - - agent.plot_tasks() - - -@pytest.mark.test_func -def test_bayesian_agent_mock_kbs(RE, db): - 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")] - - agent = bloptools.bayesian.Agent( - active_dofs=dofs, - passive_dofs=[], - active_dof_bounds=bounds, - tasks=tasks, - digestion=mock_kbs_digestion, - db=db, - ) - - RE(agent.initialize("qr", n_init=16)) - - RE(agent.learn("ei", n_iter=4)) - - agent.plot_tasks() diff --git a/bloptools/tests/test_passive_dofs.py b/bloptools/tests/test_passive_dofs.py new file mode 100644 index 0000000..f170a79 --- /dev/null +++ b/bloptools/tests/test_passive_dofs.py @@ -0,0 +1,33 @@ +import pytest + +from bloptools import devices, test_functions +from bloptools.bayesian import Agent + + +@pytest.mark.test_func +def test_passive_dofs(RE, db): + dofs = [ + {"device": devices.DOF(name="x1"), "limits": (-5, 5), "kind": "active"}, + {"device": devices.DOF(name="x2"), "limits": (-5, 5), "kind": "active"}, + {"device": devices.BrownianMotion(name="brownian1"), "limits": (-2, 2), "kind": "passive"}, + {"device": devices.BrownianMotion(name="brownian2"), "limits": (-2, 2), "kind": "passive"}, + ] + + tasks = [ + {"key": "himmelblau", "kind": "minimize"}, + ] + + agent = Agent( + dofs=dofs, + tasks=tasks, + digestion=test_functions.himmelblau_digestion, + db=db, + verbose=True, + tolerate_acquisition_errors=False, + ) + + RE(agent.initialize("qr", n_init=32)) + + agent.plot_tasks() + agent.plot_acquisition() + agent.plot_feasibility() diff --git a/bloptools/tests/test_plots.py b/bloptools/tests/test_plots.py new file mode 100644 index 0000000..4675b01 --- /dev/null +++ b/bloptools/tests/test_plots.py @@ -0,0 +1,10 @@ +import pytest + + +@pytest.mark.test_func +def test_plots(RE, agent): + RE(agent.initialize("qr", n_init=32)) + + agent.plot_tasks() + agent.plot_acquisition() + agent.plot_feasibility() diff --git a/bloptools/utils.py b/bloptools/utils.py index 492ef5c..67e1ec1 100644 --- a/bloptools/utils.py +++ b/bloptools/utils.py @@ -5,12 +5,20 @@ from ortools.constraint_solver import pywrapcp, routing_enums_pb2 +def sobol_sampler(bounds, n, q=1): + """ + Returns $n$ quasi-randomly sampled points within the bounds (a 2 by d tensor) + and batch size $q$ + """ + return botorch.utils.sampling.draw_sobol_samples(bounds, n=n, q=q) + + 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() + normalized_bounds = torch.outer(torch.tensor([0, 1]), torch.ones(d)) + return sobol_sampler(normalized_bounds, n=n, q=1) def estimate_root_indices(x): diff --git a/docs/source/conf.py b/docs/source/conf.py index 1575f81..c6993bb 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -107,10 +107,7 @@ # The theme to use for HTML and HTML Help pages. See the documentation for # a list of builtin themes. # -html_theme = "sphinx_rtd_theme" -import sphinx_rtd_theme - -html_theme_path = [sphinx_rtd_theme.get_html_theme_path()] +html_theme = "furo" # Theme options are theme-specific and customize the look and feel of a theme # further. For a list of options available for each theme, see the @@ -125,15 +122,6 @@ # Custom sidebar templates, must be a dictionary that maps document names # to template names. -# -# This is required for the alabaster theme -# refs: http://alabaster.readthedocs.io/en/latest/installation.html#sidebars -html_sidebars = { - "**": [ - "relations.html", # needs 'show_related': True theme option to display - "searchbox.html", - ] -} # -- Options for HTMLHelp output ------------------------------------------ diff --git a/docs/source/tutorials.rst b/docs/source/tutorials.rst index 48218c9..3a57df0 100644 --- a/docs/source/tutorials.rst +++ b/docs/source/tutorials.rst @@ -5,7 +5,8 @@ Tutorials :maxdepth: 2 tutorials/introduction.ipynb - tutorials/hyperparameters.ipynb tutorials/constrained-himmelblau.ipynb + tutorials/hyperparameters.ipynb + tutorials/passive-dofs.ipynb tutorials/latent-toroid-dimensions.ipynb tutorials/multi-task-sirepo.ipynb diff --git a/docs/source/tutorials/constrained-himmelblau.ipynb b/docs/source/tutorials/constrained-himmelblau.ipynb index 0a36267..01c06d8 100644 --- a/docs/source/tutorials/constrained-himmelblau.ipynb +++ b/docs/source/tutorials/constrained-himmelblau.ipynb @@ -36,8 +36,7 @@ "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", + "F = test_functions.constrained_himmelblau(X1, X2)\n", "\n", "plt.pcolormesh(x1, x2, F, norm=mpl.colors.LogNorm(), shading=\"auto\")\n", "plt.colorbar()\n", @@ -90,24 +89,26 @@ "source": [ "%run -i ../../../examples/prepare_bluesky.py # prepare the bluesky environment\n", "\n", - "import bloptools\n", - "from bloptools.tasks import Task\n", + "from bloptools import devices\n", + "from bloptools.bayesian import Agent\n", "\n", - "dofs = bloptools.devices.dummy_dofs(n=2)\n", - "bounds = [(-10, 10), (-10, 10)]\n", + "dofs = [\n", + " {\"device\": devices.DOF(name=\"x1\"), \"limits\": (-8, 8), \"kind\": \"active\"},\n", + " {\"device\": devices.DOF(name=\"x2\"), \"limits\": (-8, 8), \"kind\": \"active\"},\n", + "]\n", "\n", - "task = Task(key=\"himmelblau\", kind=\"min\")\n", + "tasks = [\n", + " {\"key\": \"himmelblau\", \"kind\": \"minimize\"},\n", + "]\n", "\n", - "agent = bloptools.bayesian.Agent(\n", - " active_dofs=dofs,\n", - " passive_dofs=[],\n", - " active_dof_bounds=bounds,\n", - " tasks=[task],\n", + "agent = Agent(\n", + " dofs=dofs,\n", + " tasks=tasks,\n", " digestion=digestion,\n", " db=db,\n", ")\n", "\n", - "RE(agent.initialize(\"qr\", n_init=32))\n", + "RE(agent.initialize(\"qr\", n_init=64))\n", "\n", "agent.plot_tasks()" ] @@ -151,7 +152,7 @@ }, "outputs": [], "source": [ - "agent.plot_acquisition(acqf=[\"ei\", \"pi\", \"ucb\"])" + "agent.plot_acquisition(acq_func=[\"ei\", \"pi\", \"ucb\"])" ] }, { diff --git a/docs/source/tutorials/hyperparameters.ipynb b/docs/source/tutorials/hyperparameters.ipynb index 124295d..d37b5c1 100644 --- a/docs/source/tutorials/hyperparameters.ipynb +++ b/docs/source/tutorials/hyperparameters.ipynb @@ -43,6 +43,22 @@ "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": "7e9c949e", + "metadata": {}, + "outputs": [], + "source": [ + "def digestion(db, uid):\n", + " products = db[uid].table()\n", + "\n", + " for index, entry in products.iterrows():\n", + " products.loc[index, \"booth\"] = test_functions.booth(entry.x1, entry.x2)\n", + "\n", + " return products" + ] + }, { "cell_type": "code", "execution_count": null, @@ -54,48 +70,30 @@ "source": [ "%run -i ../../../examples/prepare_bluesky.py # prepare the bluesky environment\n", "\n", - "import bloptools\n", - "from bloptools.tasks import Task\n", + "from bloptools import devices\n", + "from bloptools.bayesian import Agent\n", "\n", - "dofs = bloptools.devices.dummy_dofs(n=2)\n", - "bounds = [(-8, 8), (-8, 8)]\n", + "dofs = [\n", + " {\"device\": devices.DOF(name=\"x1\"), \"limits\": (-5, 5), \"kind\": \"active\"},\n", + " {\"device\": devices.DOF(name=\"x2\"), \"limits\": (-5, 5), \"kind\": \"active\"},\n", + "]\n", "\n", - "task = Task(key=\"booth\", kind=\"min\")\n", + "tasks = [\n", + " {\"key\": \"booth\", \"kind\": \"minimize\"},\n", + "]\n", "\n", - "\n", - "def digestion(db, uid):\n", - " products = db[uid].table()\n", - "\n", - " for index, entry in products.iterrows():\n", - " products.loc[index, \"booth\"] = test_functions.booth(entry.x1, entry.x2)\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", + "agent = Agent(\n", + " dofs=dofs,\n", + " tasks=tasks,\n", " digestion=digestion,\n", " db=db,\n", ")\n", "\n", - "RE(agent.initialize(acqf=\"qr\", n_init=16))\n", + "RE(agent.initialize(acq_func=\"qr\", n_init=16))\n", "\n", "agent.plot_tasks()" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "8fba0f0c", - "metadata": {}, - "outputs": [], - "source": [ - "agent.tasks[0].regressor.covar_module.latent_transform" - ] - }, { "attachments": {}, "cell_type": "markdown", diff --git a/docs/source/tutorials/introduction.ipynb b/docs/source/tutorials/introduction.ipynb index c960697..1d4eef1 100644 --- a/docs/source/tutorials/introduction.ipynb +++ b/docs/source/tutorials/introduction.ipynb @@ -17,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 the Rastrigin function in one dimension, which looks like this:" + "We'll start by minimizing the Styblinski-Tang function in one dimension, which looks like this:" ] }, { @@ -35,7 +35,8 @@ "\n", "x = np.linspace(-5, 5, 256)\n", "\n", - "plt.plot(x, test_functions.ackley(x), c=\"b\")" + "plt.plot(x, test_functions.styblinski_tang(x), c=\"b\")\n", + "plt.xlim(-5, 5)" ] }, { @@ -54,66 +55,58 @@ "metadata": {}, "outputs": [], "source": [ - "import bloptools\n", + "from bloptools import devices\n", "\n", - "dofs = bloptools.devices.dummy_dofs(n=1) # an ophyd device that we can read/set\n", - "\n", - "bounds = [(-4.0, 4.0)] # one set of bounds per dof" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "7a88c7bd", - "metadata": {}, - "source": [ - "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" + "dofs = [\n", + " {\"device\": devices.DOF(name=\"x\"), \"limits\": (-5, 5), \"kind\": \"active\"},\n", + "]" ] }, { "cell_type": "code", "execution_count": null, - "id": "e6bfcf73", + "id": "c8556bc9", "metadata": {}, "outputs": [], "source": [ - "def digestion(db, uid):\n", - " products = db[uid].table()\n", - "\n", - " for index, entry in products.iterrows():\n", - " products.loc[index, \"ackley\"] = test_functions.ackley(entry.x1)\n", - "\n", - " return products" + "tasks = [\n", + " {\"key\": \"styblinski-tang\", \"kind\": \"minimize\"},\n", + "]" ] }, { "attachments": {}, "cell_type": "markdown", - "id": "dad64303", + "id": "7a88c7bd", "metadata": {}, "source": [ - "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\"." + "\n", + "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 Styblinski-Tang function. We accomplish this with a digestion function, which always takes `(db, uid)` as an input. For each entry, we compute the function:\n" ] }, { "cell_type": "code", "execution_count": null, - "id": "4c14d162", + "id": "e6bfcf73", "metadata": {}, "outputs": [], "source": [ - "from bloptools.tasks import Task\n", + "def digestion(db, uid):\n", + " products = db[uid].table()\n", + "\n", + " for index, entry in products.iterrows():\n", + " products.loc[index, \"styblinski-tang\"] = test_functions.styblinski_tang(entry.x)\n", "\n", - "task = Task(key=\"ackley\", kind=\"min\")" + " return products" ] }, { "attachments": {}, "cell_type": "markdown", - "id": "0d3d91c3", + "id": "dad64303", "metadata": {}, "source": [ - "Combining all of these with a databroker instance, we can make an agent:" + "The next ingredient is a task, which gives the agent something to do. We want it to minimize the Styblinski-Tang function, so we make a task that will try to minimize the output of the digestion function called \"styblinski-tang\"." ] }, { @@ -127,11 +120,11 @@ "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", + "from bloptools.bayesian import Agent\n", + "\n", + "agent = Agent(\n", + " dofs=dofs,\n", + " tasks=tasks,\n", " digestion=digestion,\n", " db=db,\n", ")\n", @@ -175,6 +168,16 @@ "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": "589a263b", + "metadata": {}, + "outputs": [], + "source": [ + "agent.acq_func_info" + ] + }, { "cell_type": "code", "execution_count": null, @@ -184,9 +187,7 @@ }, "outputs": [], "source": [ - "# helper function to list acquisition functions\n", - "\n", - "agent.plot_acquisition(strategy=[\"ei\", \"pi\", \"ucb\"])" + "agent.plot_acquisition(acq_funcs=[\"ei\", \"pi\", \"ucb\"])" ] }, { @@ -205,17 +206,9 @@ "metadata": {}, "outputs": [], "source": [ - "RE(agent.learn(\"ei\", n_per_iter=4))\n", + "RE(agent.learn(\"ei\", n_iter=4))\n", "agent.plot_tasks()" ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "87908041", - "metadata": {}, - "outputs": [], - "source": [] } ], "metadata": { diff --git a/docs/source/tutorials/latent-toroid-dimensions.ipynb b/docs/source/tutorials/latent-toroid-dimensions.ipynb index 0bd3aa4..5ad5fbe 100644 --- a/docs/source/tutorials/latent-toroid-dimensions.ipynb +++ b/docs/source/tutorials/latent-toroid-dimensions.ipynb @@ -21,10 +21,7 @@ "outputs": [], "source": [ "%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]])" + "%run -i ../../../examples/prepare_tes_shadow.py" ] }, { @@ -36,17 +33,20 @@ }, "outputs": [], "source": [ - "from bloptools.bayesian import Agent\n", + "import bloptools\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", + "dofs = [\n", + " {\"device\": toroid.x_rot, \"limits\": (-0.001, 0.001), \"kind\": \"active\"},\n", + " {\"device\": toroid.offz, \"limits\": (-0.5, 0.5), \"kind\": \"active\"},\n", + "]\n", + "\n", + "tasks = [{\"key\": \"flux\", \"kind\": \"maximize\", \"transform\": \"log\"}]\n", "\n", - "agent = Agent(\n", - " active_dofs=toroid_dofs,\n", - " active_dof_bounds=toroid_bounds,\n", - " detectors=[w8],\n", - " tasks=[beam_flux_task],\n", + "agent = bloptools.bayesian.Agent(\n", + " dofs=dofs,\n", + " tasks=tasks,\n", + " dets=[w8],\n", " digestion=w8_digestion,\n", " db=db,\n", ")\n", @@ -70,7 +70,7 @@ "metadata": {}, "outputs": [], "source": [ - "agent.tasks[0].regressor.covar_module.latent_transform" + "agent.tasks[0][\"model\"].covar_module.latent_transform" ] }, { @@ -86,6 +86,14 @@ "agent.plot_feasibility()\n", "agent.plot_acquisition(strategy=[\"ei\", \"pi\", \"ucb\"])" ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "efd5ab4f", + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { diff --git a/docs/source/tutorials/multi-task-sirepo.ipynb b/docs/source/tutorials/multi-task-sirepo.ipynb index 1d318be..b965327 100644 --- a/docs/source/tutorials/multi-task-sirepo.ipynb +++ b/docs/source/tutorials/multi-task-sirepo.ipynb @@ -23,10 +23,17 @@ "outputs": [], "source": [ "%run -i ../../../examples/prepare_bluesky.py\n", - "%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]])" + "%run -i ../../../examples/prepare_tes_shadow.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8672bf9e", + "metadata": {}, + "outputs": [], + "source": [ + "toroid.offz" ] }, { @@ -40,17 +47,22 @@ "source": [ "from bloptools.bayesian import Agent\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", - "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", + "dofs = [\n", + " {\"device\": kbv.x_rot, \"limits\": (-0.1, 0.1), \"kind\": \"active\"},\n", + " {\"device\": kbh.x_rot, \"limits\": (-0.1, 0.1), \"kind\": \"active\"},\n", + "]\n", + "\n", + "tasks = [\n", + " {\"key\": \"flux\", \"kind\": \"maximize\", \"transform\": \"log\"},\n", + " {\"key\": \"w9_fwhm_x\", \"kind\": \"minimize\", \"transform\": \"log\"},\n", + " {\"key\": \"w9_fwhm_y\", \"kind\": \"minimize\", \"transform\": \"log\"},\n", + "]\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", + " dofs=dofs,\n", + " tasks=tasks,\n", + " dets=[w9],\n", " digestion=w9_digestion,\n", " db=db,\n", ")\n", diff --git a/docs/source/tutorials/passive-dofs.ipynb b/docs/source/tutorials/passive-dofs.ipynb new file mode 100644 index 0000000..6f5adea --- /dev/null +++ b/docs/source/tutorials/passive-dofs.ipynb @@ -0,0 +1,92 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "id": "e7b5e13a-c059-441d-8d4f-fff080d52054", + "metadata": {}, + "source": [ + "# Passive degrees of freedom\n", + "\n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "c18ef717", + "metadata": {}, + "source": [ + "Passive dofs!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e6bfcf73", + "metadata": {}, + "outputs": [], + "source": [ + "%run -i ../../../examples/prepare_bluesky.py # prepare the bluesky environment\n", + "\n", + "from bloptools import devices, test_functions\n", + "from bloptools.bayesian import Agent\n", + "\n", + "\n", + "def digestion(db, uid):\n", + " products = db[uid].table()\n", + "\n", + " for index, entry in products.iterrows():\n", + " products.loc[index, \"styblinksi-tang\"] = test_functions.styblinski_tang(entry.x - 1e-1 * entry.brownian)\n", + "\n", + " return products\n", + "\n", + "\n", + "dofs = [\n", + " {\"device\": devices.DOF(name=\"x\"), \"limits\": (-5, 5), \"kind\": \"active\"},\n", + " {\"device\": devices.BrownianMotion(name=\"brownian\"), \"limits\": (-2, 2), \"kind\": \"passive\"},\n", + "]\n", + "\n", + "tasks = [\n", + " {\"key\": \"styblinksi-tang\", \"kind\": \"minimize\"},\n", + "]\n", + "\n", + "agent = Agent(\n", + " dofs=dofs,\n", + " tasks=tasks,\n", + " digestion=digestion,\n", + " db=db,\n", + ")\n", + "\n", + "RE(agent.initialize(\"qr\", n_init=32))\n", + "\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/usage.rst b/docs/source/usage.rst index 0958322..905c0e6 100644 --- a/docs/source/usage.rst +++ b/docs/source/usage.rst @@ -2,8 +2,108 @@ Usage ===== -Start by importing bloptools. +Working in the Bluesky environment, we need to pass four ingredients to the Bayesian agent: + +* ``dofs``: A list of degrees of freedom for the agent to optimize over. +* ``tasks``: A list of tasks for the agent to optimize. +* ``digestion``: A function that processes the output of the acquisition into the task values. +* ``dets``: (Optional) A list of detectors to be triggered during acquisition. +* ``acquisition``: (Optional) A Bluesky plan to run for each set of inputs. + + +Degrees of freedom +++++++++++++++++++ + +Degrees of freedom (DOFs) are passed as an iterable of dicts, each containing at least the device and set of limits. + +.. code-block:: python + + my_dofs = [ + {"device": some_motor, "limits": (lower_limit, upper_limit)}, + {"device": another_motor, "limits": (lower_limit, upper_limit)}, + ] + +Here ``some_motor`` and ``another_motor`` are ``ophyd`` objects. + + + +Tasks ++++++ + +Tasks are what we want our agent to try to optimize (either maximize or minimize). We can pass as many as we'd like: + +.. code-block:: python + + my_tasks = [ + {"key": "something_to_maximize", "kind": "maximize"} + {"key": "something_to_minimize", "kind": "minimize"} + ] + + + +Digestion ++++++++++ + +The digestion function is how we go from what is spit out by the acquisition to the actual values of the tasks. + +.. code-block:: python + + def my_digestion_function(db, uid): + + products = db[uid].table(fill=True) # a pandas DataFrame + + # for each entry, do some + for index, entry in products.iterrows(): + + raw_output_1 = entry.raw_output_1 + raw_output_2 = entry.raw_output_2 + + entry.loc[index, "thing_to_maximize"] = some_fitness_function(raw_output_1, raw_output_2) + + return products + + +Detectors ++++++++++ + +Detectors are triggered for each input. + +.. code-block:: python + + my_dets = [some_detector, some_other_detector] + + + +Acquisition ++++++++++++ + +We run this plan for each set of DOF inputs. By default, this just moves the active DOFs to the desired points and triggers the supplied detectors. + + + + +Building the agent +++++++++++++++++++ + +Combining these with a databroker instance will construct an agent. .. code-block:: python import bloptools + + my_agent = bloptools.bayesian.Agent( + dofs=my_dofs, + dets=my_dets, + tasks=my_tasks, + digestion=my_digestion_function, + db=db, # a databroker instance + ) + + RE(agent.initialize("qr", n_init=24)) + + +In the example below, the agent will loop over the following steps in each iteration of learning. + +#. Find the most interesting point (or points) to sample, and move the degrees of freedom there. +#. For each point, run an acquisition plan (e.g., trigger and read the detectors). +#. Digest the results of the acquisition to find the value of the task. diff --git a/docs/wip/passive-dofs.ipynb b/docs/wip/passive-dofs.ipynb deleted file mode 100644 index 89d601e..0000000 --- a/docs/wip/passive-dofs.ipynb +++ /dev/null @@ -1,329 +0,0 @@ -{ - "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, 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()" - ] - }, - { - "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()" - ] - }, - { - "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)]" - ] - }, - { - "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", - "def digestion(db, uid):\n", - " table = db[uid].table()\n", - " products = pd.DataFrame()\n", - "\n", - " for index, entry in table.iterrows():\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\")" - ] - }, - { - "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 -} diff --git a/examples/prepare_chx_shadow.py b/examples/prepare_chx_shadow.py index 252476f..faacf5f 100644 --- a/examples/prepare_chx_shadow.py +++ b/examples/prepare_chx_shadow.py @@ -10,13 +10,9 @@ connection = SirepoBluesky("http://localhost:8000") data, schema = connection.auth("shadow", "I1Flcbdw") -classes, objects = create_classes(connection.data, connection=connection) +classes, objects = create_classes(connection=connection) globals().update(**objects) -# data["models"]["simulation"]["npoint"] = 100000 -# data["models"]["watchpointReport12"]["histogramBins"] = 32 -# w9.duration.kind = "hinted" - bec.disable_baseline() bec.disable_heading() bec.disable_table() diff --git a/examples/prepare_tes_shadow.py b/examples/prepare_tes_shadow.py index e590741..708eaaa 100644 --- a/examples/prepare_tes_shadow.py +++ b/examples/prepare_tes_shadow.py @@ -16,7 +16,7 @@ connection = SirepoBluesky("http://localhost:8000") data, schema = connection.auth("shadow", "00000002") -classes, objects = create_classes(connection.data, connection=connection) +classes, objects = create_classes(connection=connection) globals().update(**objects) data["models"]["simulation"]["npoint"] = 100000 @@ -30,9 +30,3 @@ import warnings warnings.filterwarnings("ignore", module="sirepo_bluesky") - -# This should be done by installing the package with `pip install -e .` or something similar. -# import sys -# sys.path.insert(0, "../") - -kb_dofs = [kbv.x_rot, kbv.offz, kbh.x_rot, kbh.offz] # noqa F821 diff --git a/examples/prepare_tes_srw.py b/examples/prepare_tes_srw.py index a8de346..8c25c1f 100644 --- a/examples/prepare_tes_srw.py +++ b/examples/prepare_tes_srw.py @@ -9,7 +9,7 @@ connection = SirepoBluesky("http://localhost:8000") data, schema = connection.auth("srw", "00000002") -classes, objects = create_classes(connection.data, connection=connection) +classes, objects = create_classes(connection=connection) globals().update(**objects) # w9.duration.kind = "hinted" diff --git a/requirements-dev.txt b/requirements-dev.txt index 4169c3c..6782e7c 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -4,6 +4,7 @@ black pytest-codecov coverage flake8 +furo isort nbstripout pre-commit diff --git a/requirements.txt b/requirements.txt index 65eb816..0977478 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,5 +7,6 @@ numpy ophyd ortools scipy -sirepo-bluesky +sirepo-bluesky>=0.7.0 +tables torch