diff --git a/blop/__init__.py b/blop/__init__.py index 80edaf0..c273c18 100644 --- a/blop/__init__.py +++ b/blop/__init__.py @@ -1,4 +1,8 @@ +from . import utils # noqa F401 from ._version import get_versions +from .agent import Agent # noqa F401 +from .dofs import DOF # noqa F401 +from .objectives import Objective # noqa F401 __version__ = get_versions()["version"] del get_versions diff --git a/blop/bayesian/agent.py b/blop/agent.py similarity index 58% rename from blop/bayesian/agent.py rename to blop/agent.py index a9d10bb..660731e 100644 --- a/blop/bayesian/agent.py +++ b/blop/agent.py @@ -11,29 +11,47 @@ import gpytorch import h5py import matplotlib as mpl +import napari import numpy as np import pandas as pd import scipy as sp import torch -from botorch.acquisition.objective import ScalarizedPosteriorTransform from botorch.models.deterministic import GenericDeterministicModel from botorch.models.model_list_gp_regression import ModelListGP +from botorch.models.transforms.input import AffineInputTransform, ChainedInputTransform, Log10, Normalize from databroker import Broker from ophyd import Signal -from .. import utils -from . import acquisition, models, plotting +from . import utils +from .bayesian import acquisition, models, plotting +from .bayesian.transforms import TargetingPosteriorTransform from .digestion import default_digestion_function from .dofs import DOF, DOFList from .objectives import Objective, ObjectiveList from .plans import default_acquisition_plan -from .transforms import TargetingPosteriorTransform warnings.filterwarnings("ignore", category=botorch.exceptions.warnings.InputDataWarning) mpl.rc("image", cmap="coolwarm") -MAX_TEST_INPUTS = 2**11 +DEFAULT_MAX_SAMPLES = 2**11 + + +def _validate_dofs_and_objs(dofs: DOFList, objs: ObjectiveList): + if len(dofs) == 0: + raise ValueError("You must supply at least one DOF.") + + if len(objs) == 0: + raise ValueError("You must supply at least one objective.") + + for obj in objs: + for latent_group in obj.latent_groups: + for dof_name in latent_group: + if dof_name not in dofs.names: + warnings.warn( + f"DOF name '{dof_name}' in latent group for objective '{obj.name}' does not exist." + "it will be ignored." + ) class Agent: @@ -98,6 +116,9 @@ def __init__( self.dofs = DOFList(list(np.atleast_1d(dofs))) self.objectives = ObjectiveList(list(np.atleast_1d(objectives))) + + _validate_dofs_and_objs(self.dofs, self.objectives) + self.db = db self.dets = dets @@ -119,196 +140,54 @@ def __init__( self.n_last_trained = 0 - def tell( - self, - data: Optional[Mapping] = {}, - x: Optional[Mapping] = {}, - y: Optional[Mapping] = {}, - metadata: Optional[Mapping] = {}, - append=True, - train=True, - hypers=None, - ): - """ - Inform the agent about new inputs and targets for the model. - - If run with no arguments, it will just reconstruct all the models. - - Parameters - ---------- - x : dict - A dict keyed by the name of each DOF, with a list of values for each DOF. - y : dict - A dict keyed by the name of each objective, with a list of values for each objective. - append: bool - If `True`, will append new data to old data. If `False`, will replace old data with new data. - train_models: bool - Whether to train the models on construction. - hypers: - A dict of hyperparameters for the model to assume a priori, instead of training. - """ - - if not data: - if not x and y: - raise ValueError("Must supply either x and y, or data.") - data = {**x, **y, **metadata} - - data = {k: list(np.atleast_1d(v)) for k, v in data.items()} - unique_field_lengths = {len(v) for v in data.values()} - - if len(unique_field_lengths) > 1: - raise ValueError("All supplies values must be the same length!") - - new_table = pd.DataFrame(data) - self.table = pd.concat([self.table, new_table]) if append else new_table - self.table.index = np.arange(len(self.table)) - - for obj in self.objectives: - t0 = ttime.monotonic() - - cached_hypers = obj.model.state_dict() if hasattr(obj, "model") else None + @property + def active_dofs(self): + return self.dofs.subset(active=True) - obj.model = self.construct_model(obj) + @property + def active_objs(self): + return self.objectives.subset(active=True) - if len(obj.model.train_targets) >= 2: - t0 = ttime.monotonic() - self.train_model(obj.model, hypers=(None if train else cached_hypers)) - if self.verbose: - print(f"trained model '{obj.name}' in {1e3*(ttime.monotonic() - t0):.00f} ms") + def __iter__(self): + for index in range(len(self)): + yield self.dofs[index] - # TODO: should this be per objective? - self.construct_classifier() + def __getattr__(self, attr): + acq_func_name = acquisition.parse_acq_func_identifier(attr) + if acq_func_name is not None: + return self._get_acquisition_function(identifier=acq_func_name) - def train_model(self, model, hypers=None, **kwargs): - """Fit all of the agent's models. All kwargs are passed to `botorch.fit.fit_gpytorch_mll`.""" - if hypers is not None: - model.load_state_dict(hypers) - else: - botorch.fit.fit_gpytorch_mll(gpytorch.mlls.ExactMarginalLogLikelihood(model.likelihood, model), **kwargs) - model.trained = True + raise AttributeError(f"No attribute named '{attr}'.") - def construct_model(self, obj, skew_dims=None): + def sample(self, n: int = DEFAULT_MAX_SAMPLES, method: str = "quasi-random") -> torch.Tensor: """ - Construct an untrained model for an objective. - """ - - skew_dims = skew_dims if skew_dims is not None else self.latent_dim_tuples - - likelihood = gpytorch.likelihoods.GaussianLikelihood( - noise_constraint=gpytorch.constraints.Interval( - torch.tensor(1e-4).square(), - torch.tensor(1 / obj.min_snr).square(), - ), - # noise_prior=gpytorch.priors.torch_priors.LogNormalPrior(loc=loc, scale=scale), - ) + Returns a (..., 1, n_active_dofs) tensor of points sampled within the parameter space. - outcome_transform = botorch.models.transforms.outcome.Standardize(m=1) # , batch_shape=torch.Size((1,))) - - train_inputs = self.train_inputs(active=True) - train_targets = self.train_targets(obj.name) - - safe = ~(torch.isnan(train_inputs).any(axis=1) | torch.isnan(train_targets).any(axis=1)) - - model = models.LatentGP( - train_inputs=train_inputs[safe], - train_targets=train_targets[safe], - likelihood=likelihood, - skew_dims=skew_dims, - input_transform=self.input_transform, - outcome_transform=outcome_transform, - ) - - model.trained = False + Parameters + ---------- + n : int + How many points to sample. + method : str + How to sample the points. Must be one of 'quasi-random', 'random', or 'grid'. + """ - return model + if method == "quasi-random": + X = utils.normalized_sobol_sampler(n, d=len(self.active_dofs)) - def construct_classifier(self, skew_dims=None): - skew_dims = skew_dims if skew_dims is not None else self.latent_dim_tuples + elif method == "random": + X = torch.rand(size=(n, 1, len(self.active_dofs))) - dirichlet_likelihood = gpytorch.likelihoods.DirichletClassificationLikelihood( - self.all_objectives_valid.long(), learn_additional_noise=True - ) + elif method == "grid": + n_side_if_settable = int(np.power(n, 1 / np.sum(~self.active_dofs.read_only))) + sides = [ + torch.linspace(0, 1, n_side_if_settable) if not dof.read_only else torch.zeros(1) for dof in self.active_dofs + ] + X = torch.cat([x.unsqueeze(-1) for x in torch.meshgrid(sides, indexing="ij")], dim=-1).unsqueeze(-2).double() - self.classifier = models.LatentDirichletClassifier( - train_inputs=self.train_inputs(active=True), - train_targets=dirichlet_likelihood.transformed_targets.transpose(-1, -2).double(), - skew_dims=skew_dims, - likelihood=dirichlet_likelihood, - input_transform=self.input_transform, - ) + else: + raise ValueError("'method' argument must be one of ['quasi-random', 'random', 'grid'].") - self.train_model(self.classifier) - self.constraint = GenericDeterministicModel(f=lambda x: self.classifier.probabilities(x)[..., -1]) - - # def construct_model(self, obj, skew_dims=None, a_priori_hypers=None): - # ''' - # Construct an untrained model for an objective. - # ''' - # skew_dims = skew_dims if skew_dims is not None else self.latent_dim_tuples - - # inputs = self.table.loc[:, self.dofs.subset(active=True).names].values.astype(float) - - # for i, obj in enumerate(self.objectives): - # values = self.train_targets(i) - # values = np.where(self.all_objectives_valid, values, np.nan) - - # train_index = ~np.isnan(values) - - # if not train_index.sum() >= 2: - # raise ValueError("There must be at least two valid data points per objective!") - - # train_inputs = torch.tensor(inputs[train_index], dtype=torch.double) - # train_values = torch.tensor(values[train_index], dtype=torch.double).unsqueeze(-1) # .unsqueeze(0) - - # # for constructing the log normal noise prior - # # target_snr = 2e2 - # # scale = 2e0 - # # loc = np.log(1 / target_snr**2) + scale**2 - - # likelihood = gpytorch.likelihoods.GaussianLikelihood( - # noise_constraint=gpytorch.constraints.Interval( - # torch.tensor(1e-4).square(), - # torch.tensor(1 / obj.min_snr).square(), - # ), - # # noise_prior=gpytorch.priors.torch_priors.LogNormalPrior(loc=loc, scale=scale), - # ) - - # outcome_transform = botorch.models.transforms.outcome.Standardize(m=1) # , batch_shape=torch.Size((1,))) - - # obj.model = models.LatentGP( - # train_inputs=train_inputs, - # train_targets=self.t, - # likelihood=likelihood, - # skew_dims=skew_dims, - # input_transform=self.input_transform, - # outcome_transform=outcome_transform, - # ) - - # dirichlet_likelihood = gpytorch.likelihoods.DirichletClassificationLikelihood( - # self.all_objectives_valid.long(), learn_additional_noise=True - # ) - - # self.classifier = models.LatentDirichletClassifier( - # train_inputs=torch.tensor(inputs).double(), - # train_targets=dirichlet_likelihood.transformed_targets.transpose(-1, -2).double(), - # skew_dims=skew_dims, - # likelihood=dirichlet_likelihood, - # input_transform=self._subset_input_transform(active=True), - # ) - - # if a_priori_hypers is not None: - # self._set_hypers(a_priori_hypers) - # else: - # self._train_models() - # # try: - - # # except botorch.exceptions.errors.ModelFittingError: - # # if self.initialized: - # # self._set_hypers(cached_hypers) - # # else: - # # raise RuntimeError('Could not fit model on initialization!') - - # self.constraint = GenericDeterministicModel(f=lambda x: self.classifier.probabilities(x)[..., -1]) + return self._sample_input_transform.untransform(X) def ask(self, acq_func_identifier="qei", n=1, route=True, sequential=True, upsample=1, **acq_func_kwargs): """Ask the agent for the best point to sample, given an acquisition function. @@ -331,29 +210,37 @@ def ask(self, acq_func_identifier="qei", n=1, route=True, sequential=True, upsam start_time = ttime.monotonic() - if self.verbose: - print(f"finding points with acquisition function '{acq_func_name}' ...") + if acq_func_name in ["quasi-random", "random", "grid"]: + candidates = self.sample(n=n, method=acq_func_name).squeeze(1).numpy() - if acq_func_type in ["analytic", "monte_carlo"]: + # define dummy acqf objective + acqf_obj = torch.zeros(len(candidates)) + + elif acq_func_type in ["analytic", "monte_carlo"]: if not all(hasattr(obj, "model") for obj in self.objectives): raise RuntimeError( f"Can't construct non-trivial acquisition function '{acq_func_identifier}'" f" (the agent is not initialized!)" ) + for obj in self.active_objs: + if obj.model_dofs != set(self.active_dofs.names): + self._construct_model(obj) + self._train_model(obj.model) + if acq_func_type == "analytic" and n > 1: raise ValueError("Can't generate multiple design points for analytic acquisition functions.") - acq_func, acq_func_meta = self.get_acquisition_function( + acq_func, _ = self._get_acquisition_function( identifier=acq_func_identifier, return_metadata=True, **acq_func_kwargs ) NUM_RESTARTS = 8 - RAW_SAMPLES = 1024 + RAW_SAMPLES = 256 candidates, acqf_obj = botorch.optim.optimize_acqf( acq_function=acq_func, - bounds=self.acquisition_function_bounds, + bounds=self._sample_bounds, q=n, sequential=sequential, num_restarts=NUM_RESTARTS, @@ -363,38 +250,10 @@ def ask(self, acq_func_identifier="qei", n=1, route=True, sequential=True, upsam # this includes both RO and non-RO DOFs candidates = candidates.numpy() - active_dofs_are_read_only = np.array([dof.read_only for dof in self.dofs.subset(active=True)]) - - acq_points = candidates[..., ~active_dofs_are_read_only] - read_only_values = candidates[..., active_dofs_are_read_only] - acq_func_meta["read_only_values"] = read_only_values - - else: - acqf_obj = None - - if acq_func_name == "random": - acq_points = torch.rand() - acq_func_meta = {"name": "random", "args": {}} - - if acq_func_name == "quasi-random": - acq_points = self._subset_inputs_sampler(n=n, active=True, read_only=False).squeeze(1).numpy() - acq_func_meta = {"name": "quasi-random", "args": {}} - - elif acq_func_name == "grid": - n_active_dims = len(self.dofs.subset(active=True, read_only=False)) - acq_points = self.test_inputs_grid(max_inputs=n).reshape(-1, n_active_dims).numpy() - acq_func_meta = {"name": "grid", "args": {}} - - else: - raise ValueError() - - # define dummy acqf objective - acqf_obj = 0 - - acq_func_meta["duration"] = duration = ttime.monotonic() - start_time + acq_points = candidates[..., ~self.active_dofs.read_only] + read_only_values = candidates[..., self.active_dofs.read_only] - if self.verbose: - print(f"found points {acq_points} in {1e3*duration:.01f} ms (obj = {acqf_obj})") + duration = 1e3 * (ttime.monotonic() - start_time) if route and n > 1: routing_index = utils.route(self.dofs.subset(active=True, read_only=False).readback, acq_points) @@ -405,74 +264,87 @@ def ask(self, acq_func_identifier="qei", n=1, route=True, sequential=True, upsam upsampled_idx = np.linspace(0, len(idx) - 1, upsample * len(idx) - 1) acq_points = sp.interpolate.interp1d(idx, acq_points, axis=0)(upsampled_idx) + p = self.posterior(candidates) if hasattr(self, "model") else None + res = { "points": acq_points, - "acq_func": acq_func_meta["name"], + "acq_func": acq_func_name, "acq_func_kwargs": acq_func_kwargs, - "duration": acq_func_meta["duration"], + "acq_func_obj": np.atleast_1d(acqf_obj.numpy()), + "duration_ms": duration, "sequential": sequential, "upsample": upsample, - "read_only_values": acq_func_meta.get("read_only_values"), + "read_only_values": read_only_values, + "posterior": p, } return res - def acquire(self, acquisition_inputs): - """Acquire and digest according to the self's acquisition and digestion plans. + def tell( + self, + data: Optional[Mapping] = {}, + x: Optional[Mapping] = {}, + y: Optional[Mapping] = {}, + metadata: Optional[Mapping] = {}, + append: bool = True, + update_models: bool = True, + train: bool = None, + ): + """ + Inform the agent about new inputs and targets for the model. + + If run with no arguments, it will just reconstruct all the models. Parameters ---------- - acquisition_inputs : - A 2D numpy array comprising inputs for the active and non-read-only DOFs to sample. + x : dict + A dict keyed by the name of each DOF, with a list of values for each DOF. + y : dict + A dict keyed by the name of each objective, with a list of values for each objective. + append: bool + If `True`, will append new data to old data. If `False`, will replace old data with new data. + train: bool + Whether to train the models on construction. + hypers: + A dict of hyperparameters for the model to assume a priori, instead of training. """ - if self.db is None: - raise ValueError("Cannot run acquistion without databroker instance!") - - try: - acquisition_devices = self.dofs.subset(active=True, read_only=False).devices - # read_only_devices = self.dofs.subset(active=True, read_only=True).devices - - # the acquisition plan always takes as arguments: - # (things to move, where to move them, things to trigger once you get there) - # with some other kwargs + if not data: + if not x and y: + raise ValueError("Must supply either x and y, or data.") + data = {**x, **y, **metadata} - uid = yield from self.acquisition_plan( - acquisition_devices, - acquisition_inputs.astype(float), - [*self.dets, *self.dofs.devices], - delay=self.trigger_delay, - ) + data = {k: list(np.atleast_1d(v)) for k, v in data.items()} + unique_field_lengths = {len(v) for v in data.values()} - products = self.digestion(self.db, uid) + if len(unique_field_lengths) > 1: + raise ValueError("All supplies values must be the same length!") - # compute the fitness for each objective - # for index, entry in products.iterrows(): - # for obj in self.objectives: - # products.loc[index, objective['key']] = getattr(entry, objective['key']) + new_table = pd.DataFrame(data) + self.table = pd.concat([self.table, new_table]) if append else new_table + self.table.index = np.arange(len(self.table)) - except KeyboardInterrupt as interrupt: - raise interrupt + if update_models: + for obj in self.active_objs: + t0 = ttime.monotonic() - except Exception as error: - if not self.tolerate_acquisition_errors: - raise error - logging.warning(f"Error in acquisition/digestion: {repr(error)}") - products = pd.DataFrame(acquisition_inputs, columns=self.dofs.subset(active=True, read_only=False).names) - for obj in self.objectives: - products.loc[:, obj.name] = np.nan + cached_hypers = obj.model.state_dict() if hasattr(obj, "model") else None + n_before_tell = obj.n + self._construct_model(obj) + n_after_tell = obj.n - if not len(acquisition_inputs) == len(products): - raise ValueError("The table returned by the digestion function must be the same length as the sampled inputs!") + if train is None: + train = int(n_after_tell / self.train_every) > int(n_before_tell / self.train_every) - return products + if len(obj.model.train_targets) >= 4: + if train: + t0 = ttime.monotonic() + self._train_model(obj.model) + if self.verbose: + print(f"trained model '{obj.name}' in {1e3*(ttime.monotonic() - t0):.00f} ms") - def load_data(self, data_file, append=True, train=True): - new_table = pd.read_hdf(data_file, key="table") - x = {key: new_table.pop(key).tolist() for key in self.dofs.names} - y = {key: new_table.pop(key).tolist() for key in self.objectives.names} - metadata = new_table.to_dict(orient="list") - self.tell(x=x, y=y, metadata=metadata, append=append, train=train) + else: + self._train_model(obj.model, hypers=cached_hypers) def learn( self, @@ -480,9 +352,10 @@ def learn( n: int = 1, iterations: int = 1, upsample: int = 1, - train: bool = True, + train: bool = None, append: bool = True, - hypers_file=None, + hypers: str = None, + route: bool = True, ): """This returns a Bluesky plan which iterates the learning algorithm, looping over ask -> acquire -> tell. @@ -511,14 +384,14 @@ def learn( """ if self.sample_center_on_init and not self.initialized: - center_inputs = np.atleast_2d(self.dofs.subset(active=True, read_only=False).limits.mean(axis=1)) + center_inputs = np.atleast_2d(self.dofs.subset(active=True, read_only=False).search_bounds.mean(axis=1)) new_table = yield from self.acquire(center_inputs) new_table.loc[:, "acq_func"] = "sample_center_on_init" for i in range(iterations): print(f"running iteration {i + 1} / {iterations}") for single_acq_func in np.atleast_1d(acq_func): - res = self.ask(n=n, acq_func_identifier=single_acq_func, upsample=upsample) + res = self.ask(n=n, acq_func_identifier=single_acq_func, upsample=upsample, route=route) new_table = yield from self.acquire(res["points"]) new_table.loc[:, "acq_func"] = res["acq_func"] @@ -527,17 +400,96 @@ def learn( metadata = new_table.to_dict(orient="list") self.tell(x=x, y=y, metadata=metadata, append=append, train=train) - def get_acquisition_function(self, identifier, return_metadata=False): - """Returns a BoTorch acquisition function for a given identifier. Acquisition functions can be - found in `agent.all_acq_funcs`. + def view(self, item: str = "mean", cmap: str = "turbo", max_inputs: int = 2**16): """ - return acquisition.get_acquisition_function(self, identifier=identifier, return_metadata=return_metadata) + Use napari to see a high-dimensional array. + + Parameters + ---------- + item : str + The thing to be viewed. Either 'mean', 'error', or an acquisition function. + """ + + test_grid = self.sample(n=max_inputs, method="grid") + + self.viewer = napari.Viewer() + + if item in ["mean", "error"]: + for obj in self.active_objs: + p = obj.model.posterior(test_grid) + + if item == "mean": + mean = p.mean.detach().numpy()[..., 0, 0] + self.viewer.add_image(data=mean, name=f"{obj.name}_mean", colormap=cmap) + + if item == "error": + error = np.sqrt(p.variance.detach().numpy()[..., 0, 0]) + self.viewer.add_image(data=error, name=f"{obj.name}_error", colormap=cmap) + + else: + try: + acq_func_identifier = acquisition.parse_acq_func_identifier(identifier=item) + except Exception: + raise ValueError("'item' must be either 'mean', 'error', or a valid acq func.") + + acq_func, acq_func_meta = self._get_acquisition_function(identifier=acq_func_identifier, return_metadata=True) + a = acq_func(test_grid).detach().numpy() + + self.viewer.add_image(data=a, name=f"{acq_func_identifier}", colormap=cmap) + + self.viewer.dims.axis_labels = self.dofs.names + + def acquire(self, acquisition_inputs): + """Acquire and digest according to the self's acquisition and digestion plans. + + Parameters + ---------- + acquisition_inputs : + A 2D numpy array comprising inputs for the active and non-read-only DOFs to sample. + """ + + if self.db is None: + raise ValueError("Cannot run acquistion without databroker instance!") + + try: + acquisition_devices = self.dofs.subset(active=True, read_only=False).devices + uid = yield from self.acquisition_plan( + acquisition_devices, + acquisition_inputs.astype(float), + [*self.dets, *self.dofs.devices], + delay=self.trigger_delay, + ) + + products = self.digestion(self.db, uid) + + except KeyboardInterrupt as interrupt: + raise interrupt + + except Exception as error: + if not self.tolerate_acquisition_errors: + raise error + logging.warning(f"Error in acquisition/digestion: {repr(error)}") + products = pd.DataFrame(acquisition_inputs, columns=self.dofs.subset(active=True, read_only=False).names) + for obj in self.active_objs: + products.loc[:, obj.name] = np.nan + + if not len(acquisition_inputs) == len(products): + raise ValueError("The table returned by the digestion function must be the same length as the sampled inputs!") + + return products + + def load_data(self, data_file, append=True, train=True): + new_table = pd.read_hdf(data_file, key="table") + x = {key: new_table.pop(key).tolist() for key in self.dofs.names} + y = {key: new_table.pop(key).tolist() for key in self.objectives.names} + metadata = new_table.to_dict(orient="list") + self.tell(x=x, y=y, metadata=metadata, append=append, train=train) def reset(self): """Reset the agent.""" self.table = pd.DataFrame() - for obj in self.objectives: + for obj in self.active_objs: if hasattr(obj, "model"): del obj.model @@ -569,36 +521,18 @@ def benchmark( @property def model(self): """A model encompassing all the objectives. A single GP in the single-objective case, or a model list.""" - return ModelListGP(*[obj.model for obj in self.objectives]) if len(self.objectives) > 1 else self.objectives[0].model + return ( + ModelListGP(*[obj.model for obj in self.active_objs]) if len(self.active_objs) > 1 else self.active_objs[0].model + ) def posterior(self, x): """A model encompassing all the objectives. A single GP in the single-objective case, or a model list.""" - return self.model.posterior(x) - - @property - def objective_weights_torch(self): - return torch.tensor(self.objectives.weights, dtype=torch.double) - - @property - def scalarizing_transform(self): - return ScalarizedPosteriorTransform(weights=self.objective_weights_torch, offset=0) + return self.model.posterior(torch.tensor(x)) @property def targeting_transform(self): - return TargetingPosteriorTransform(weights=self.objective_weights_torch, targets=self.objectives.targets) - - @property - def pseudo_targets(self): - """Targets for the posterior transform""" - return torch.tensor( - [ - self.train_targets(active=True)[..., i].max() - if t == "max" - else self.train_targets(active=True)[..., i].min() - if t == "min" - else t - for i, t in enumerate(self.objectives.targets) - ] + return TargetingPosteriorTransform( + weights=torch.tensor(self.active_objs.weights, dtype=torch.double), targets=self.active_objs.targets ) @property @@ -623,74 +557,151 @@ def all_objectives_valid(self): """A mask of whether all objectives are valid for each data point.""" return ~torch.isnan(self.scalarized_objectives) - def test_inputs_grid(self, max_inputs=MAX_TEST_INPUTS): - """Returns a (`n_side`, ..., `n_side`, 1, `n_active_dofs`) grid of test_inputs; `n_side` is 1 if a dof is read-only. - The value of `n_side` is the largest value such that the entire grid has less than `max_inputs` inputs. + def _train_model(self, model, hypers=None, **kwargs): + """Fit all of the agent's models. All kwargs are passed to `botorch.fit.fit_gpytorch_mll`.""" + if hypers is not None: + model.load_state_dict(hypers) + else: + botorch.fit.fit_gpytorch_mll(gpytorch.mlls.ExactMarginalLogLikelihood(model.likelihood, model), **kwargs) + model.trained = True + + def _construct_model(self, obj, skew_dims=None): + """ + Construct an untrained model for an objective. """ - n_settable_acq_func_dofs = len(self.dofs.subset(active=True, read_only=False)) - n_side_settable = int(np.power(max_inputs, n_settable_acq_func_dofs**-1)) - n_sides = [1 if dof.read_only else n_side_settable for dof in self.dofs.subset(active=True)] - return torch.cat( - [ - tensor.unsqueeze(-1) - for tensor in torch.meshgrid( - *[ - torch.linspace(lower_limit, upper_limit, n_side) - for (lower_limit, upper_limit), n_side in zip(self.dofs.subset(active=True).limits, n_sides) - ], - indexing="ij", - ) - ], - dim=-1, - ).unsqueeze(-2) - def test_inputs(self, n=MAX_TEST_INPUTS): - """Returns a (n, 1, n_active_dof) grid of test_inputs""" - return utils.sobol_sampler(self.acquisition_function_bounds, n=n) + skew_dims = skew_dims if skew_dims is not None else self._latent_dim_tuples(obj.name) - @property - def acquisition_function_bounds(self): - """Returns a (2, n_active_dof) array of bounds for the acquisition function""" - active_dofs = self.dofs.subset(active=True) + likelihood = gpytorch.likelihoods.GaussianLikelihood( + noise_constraint=gpytorch.constraints.Interval( + torch.tensor(obj.min_noise), + torch.tensor(obj.max_noise), + ), + ) - acq_func_lower_bounds = [dof.lower_limit if not dof.read_only else dof.readback for dof in active_dofs] - acq_func_upper_bounds = [dof.upper_limit if not dof.read_only else dof.readback for dof in active_dofs] + outcome_transform = botorch.models.transforms.outcome.Standardize(m=1) # , batch_shape=torch.Size((1,))) - return torch.tensor(np.vstack([acq_func_lower_bounds, acq_func_upper_bounds]), dtype=torch.double) + train_inputs = self.train_inputs(active=True) + train_targets = self.train_targets(obj.name) - @property - def latent_dim_tuples(self): + inputs_are_trusted = ~torch.isnan(train_inputs).any(axis=1) + targets_are_trusted = ~torch.isnan(train_targets).any(axis=1) + + trusted = inputs_are_trusted & targets_are_trusted + + obj.model = models.LatentGP( + train_inputs=train_inputs[trusted], + train_targets=train_targets[trusted], + likelihood=likelihood, + skew_dims=skew_dims, + input_transform=self._model_input_transform, + outcome_transform=outcome_transform, + ) + + obj.model_dofs = set(self.active_dofs.names) # if these change, retrain the model on self.ask() + + if trusted.all(): + obj.validity_conjugate_model = None + obj.validity_constraint = GenericDeterministicModel(f=lambda x: torch.ones(size=x.size())[..., -1]) + + else: + dirichlet_likelihood = gpytorch.likelihoods.DirichletClassificationLikelihood( + trusted.long(), learn_additional_noise=True + ) + + obj.validity_conjugate_model = models.LatentDirichletModel( + train_inputs=train_inputs[inputs_are_trusted], + train_targets=dirichlet_likelihood.transformed_targets.transpose(-1, -2)[inputs_are_trusted].double(), + skew_dims=skew_dims, + likelihood=dirichlet_likelihood, + input_transform=self._model_input_transform, + ) + + obj.validity_constraint = GenericDeterministicModel( + f=lambda x: obj.validity_conjugate_model.probabilities(x)[..., -1] + ) + + def _construct_all_models(self): + """Construct a model for each objective.""" + for obj in self.active_objs: + self._construct_model(obj) + + def _train_all_models(self, **kwargs): + """Fit all of the agent's models. All kwargs are passed to `botorch.fit.fit_gpytorch_mll`.""" + t0 = ttime.monotonic() + for obj in self.active_objs: + self._train_model(obj.model) + if obj.validity_conjugate_model is not None: + self._train_model(obj.validity_conjugate_model) + + if self.verbose: + print(f"trained models in {ttime.monotonic() - t0:.01f} seconds") + + self.n_last_trained = len(self.table) + + def _get_acquisition_function(self, identifier, return_metadata=False): + """Returns a BoTorch acquisition function for a given identifier. Acquisition functions can be + found in `agent.all_acq_funcs`. """ - Returns a list of tuples, where each tuple represent a group of dimension to find a latent representation of. + return acquisition.get_acquisition_function(self, identifier=identifier, return_metadata=return_metadata) + + def _latent_dim_tuples(self, obj_index=None): + """ + For the objective indexed by 'obj_index', return a list of tuples, where each tuple represents + a group of DOFs to fit a latent representation to. """ - latent_dim_labels = [dof.latent_group for dof in self.dofs.subset(active=True)] - u, uinv = np.unique(latent_dim_labels, return_inverse=True) + if obj_index is None: + return {obj.name: self._latent_dim_tuples(obj_index=obj.name) for obj in self.objectives} + + obj = self.objectives[obj_index] + + latent_group_index = {} + for dof in self.active_dofs: + latent_group_index[dof.name] = dof.name + for group_index, latent_group in enumerate(obj.latent_groups): + if dof.name in latent_group: + latent_group_index[dof.name] = group_index + u, uinv = np.unique(list(latent_group_index.values()), return_inverse=True) return [tuple(np.where(uinv == i)[0]) for i in range(len(u))] @property - def input_transform(self): - """ - A bounding transform for all the active DOFs. This is used for model fitting. - """ - return self._subset_input_transform(active=True) - - def _subset_input_transform(self, active=None, read_only=None, tags=[]): - # torch likes limits to be (2, n_dof) and not (n_dof, 2) - torch_limits = torch.tensor(self.dofs.subset(active, read_only, tags).limits.T, dtype=torch.double) - offset = torch_limits.min(dim=0).values - coefficient = torch_limits.max(dim=0).values - offset - return botorch.models.transforms.input.AffineInputTransform( - d=torch_limits.shape[-1], coefficient=coefficient, offset=offset - ) + def _sample_bounds(self): + return torch.tensor(self.active_dofs.search_bounds, dtype=torch.double).T + + @property + def _sample_input_transform(self): + tf1 = Log10(indices=list(np.where(self.active_dofs.log)[0])) - def _subset_inputs_sampler(self, active=None, read_only=None, tags=[], n=MAX_TEST_INPUTS): + transformed_sample_bounds = tf1.transform(self._sample_bounds) + + offset = transformed_sample_bounds.min(dim=0).values + coefficient = (transformed_sample_bounds.max(dim=0).values - offset).clamp(min=1e-16) + + tf2 = AffineInputTransform(d=len(offset), coefficient=coefficient, offset=offset) + + return ChainedInputTransform(tf1=tf1, tf2=tf2) + + @property + def _model_input_transform(self): """ - Returns $n$ quasi-randomly sampled inputs in the bounded parameter space + Suitably transforms model inputs to the unit hypercube. + + For modeling: + + Always normalize between min and max values. This is always inside the trust bounds, sometimes smaller. + + For sampling: + + Settable: normalize between search bounds + Read-only: constrain to the readback value """ - transform = self._subset_input_transform(active, read_only, tags) - return transform.untransform(utils.normalized_sobol_sampler(n, d=len(self.dofs.subset(active, read_only, tags)))) + + tf1 = Log10(indices=list(np.where(self.active_dofs.log)[0])) + tf2 = Normalize(d=len(self.active_dofs)) + + return ChainedInputTransform(tf1=tf1, tf2=tf2) def save_data(self, filepath="./self_data.h5"): """ @@ -700,44 +711,61 @@ def save_data(self, filepath="./self_data.h5"): self.table.to_hdf(filepath, key="table") - def forget(self, index, train=True): + def forget(self, last=None, index=None, train=True): """ - Make the agent forget some index of the data table. - """ - self.table.drop(index=index, inplace=True) - self._construct_models(train=train) + Make the agent forget some data. - def forget_last_n(self, n, train=True): - """ - Make the agent forget the last `n` data points taken. + Parameters + ---------- + index : + An index of samples to forget about. + last : int + Forget the last n=last points. """ - if n > len(self.table): - raise ValueError(f"Cannot forget {n} data points (only {len(self.table)} have been taken).") - self.forget(self.table.index.iloc[-n:], train=train) - def sampler(self, n, d): - """ - Returns $n$ quasi-randomly sampled points on the [0,1] ^ d 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=d, scramble=True).random(n=min_power_of_two)[subset] + if last is not None: + if last > len(self.table): + raise ValueError(f"Cannot forget last {last} data points (only {len(self.table)} samples have been taken).") + self.forget(index=self.table.index.values[-last:], train=train) + + elif index is not None: + self.table.drop(index=index, inplace=True) + self._construct_all_models() + if train: + self._train_all_models() + + else: + raise ValueError("Must supply either 'last' or 'index'.") def _set_hypers(self, hypers): - for obj in self.objectives: + for obj in self.active_objs: obj.model.load_state_dict(hypers[obj.name]) - self.classifier.load_state_dict(hypers["classifier"]) + self.validity_constraint.load_state_dict(hypers["validity_constraint"]) + + def constraint(self, x): + p = torch.ones(x.shape[:-1]) + for obj in self.active_objs: + # if the targeting constraint is non-trivial + if obj.use_as_constraint: + p *= obj.targeting_constraint(x) + # if the validity constaint is non-trivial + if obj.validity_conjugate_model is not None: + p *= obj.validity_constraint(x) + return p @property - def hypers(self): - """Returns a dict of all the hyperparameters in all the agent's models.""" - hypers = {"classifier": {}} - for key, value in self.classifier.state_dict().items(): - hypers["classifier"][key] = value + def hypers(self) -> dict: + """Returns a dict of all the hyperparameters for each model in each objective.""" + hypers = {} for obj in self.objectives: - hypers[obj.name] = {} + hypers[obj.name] = {"model": {}, "validity_conjugate_model": {}} + for key, value in obj.model.state_dict().items(): - hypers[obj.name][key] = value + hypers[obj.name]["model"][key] = value + + if obj.validity_conjugate_model is not None: + for key, value in obj.validity_conjugate_model.state_dict().items(): + hypers[obj.name]["validity_conjugate_model"][key] = value return hypers @@ -745,35 +773,32 @@ def save_hypers(self, filepath): """Save the agent's fitted hyperparameters to a given 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) + for obj_name in hypers.keys(): + f.create_group(obj_name) + f[obj_name].create_group("model") + f[obj_name].create_group("validity_conjugate_model") + + for key, value in hypers[obj_name]["model"].items(): + f[obj_name]["model"].create_dataset(key, data=value) + + for key, value in hypers[obj_name]["validity_conjugate_model"].items(): + f[obj_name]["validity_conjugate_model"].create_dataset(key, data=value) @staticmethod - def load_hypers(filepath): + def load_hypers(filepath) -> dict: """Load hyperparameters from a file.""" 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 + for obj_name in f.keys(): + hypers[obj_name] = {"model": OrderedDict(), "validity_conjugate_model": OrderedDict()} - def _train_models(self, **kwargs): - """Fit all of the agent's models. All kwargs are passed to `botorch.fit.fit_gpytorch_mll`.""" - t0 = ttime.monotonic() - for obj in self.objectives: - model = obj.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:.01f} seconds") + for key, value in f[obj_name]["model"].items(): + hypers[obj_name]["model"][key] = torch.tensor(np.atleast_1d(value[()])) - self.n_last_trained = len(self.table) + for key, value in f[obj_name]["validity_conjugate_model"].items(): + hypers[obj_name]["validity_conjugate_model"][key] = torch.tensor(np.atleast_1d(value[()])) + + return hypers @property def all_acq_funcs(self): @@ -787,41 +812,32 @@ def all_acq_funcs(self): print("\n\n".join(entries)) - @property - def inputs(self): - """A DataFrame of all DOF values.""" - return self.table.loc[:, self.dofs.names].astype(float) - - def train_inputs(self, dof_name=None, **subset_kwargs): + def train_inputs(self, index=None, **subset_kwargs): """A two-dimensional tensor of all DOF values.""" - if dof_name is None: + if index is None: return torch.cat([self.train_inputs(dof.name) for dof in self.dofs.subset(**subset_kwargs)], dim=-1) - dof = self.dofs[dof_name] + dof = self.dofs[index] inputs = self.table.loc[:, dof.name].values.copy() # check that inputs values are inside acceptable values - valid = (inputs >= dof.limits[0]) & (inputs <= dof.limits[1]) + valid = (inputs >= dof._trust_bounds[0]) & (inputs <= dof._trust_bounds[1]) inputs = np.where(valid, inputs, np.nan) - # transform if needed - if dof.log: - inputs = np.where(inputs > 0, np.log(inputs), np.nan) - return torch.tensor(inputs, dtype=torch.double).unsqueeze(-1) - def train_targets(self, obj_name=None, **subset_kwargs): + def train_targets(self, index=None, **subset_kwargs): """Returns the values associated with an objective name.""" - if obj_name is None: + if index is None: return torch.cat([self.train_targets(obj.name) for obj in self.objectives], dim=-1) - obj = self.objectives[obj_name] + obj = self.objectives[index] targets = self.table.loc[:, obj.name].values.copy() # check that targets values are inside acceptable values - valid = (targets >= obj.limits[0]) & (targets <= obj.limits[1]) + valid = (targets >= obj._trust_bounds[0]) & (targets <= obj._trust_bounds[1]) targets = np.where(valid, targets, np.nan) # transform if needed @@ -830,16 +846,6 @@ def train_targets(self, obj_name=None, **subset_kwargs): return torch.tensor(targets, dtype=torch.double).unsqueeze(-1) - @property - def active_inputs(self): - """A two-dimensional array of all inputs for model fitting.""" - return self.table.loc[:, self.dofs.subset(active=True).names].astype(float) - - @property - def acquisition_inputs(self): - """A two-dimensional array of all inputs for computing acquisition functions.""" - return self.table.loc[:, self.dofs.subset(active=True, read_only=False).names].astype(float) - @property def best(self): """Returns all data for the best point.""" @@ -900,7 +906,7 @@ def plot_acquisition(self, acq_func="ei", axes: Tuple = (0, 1), **kwargs): else: plotting._plot_acqf_many_dofs(self, acq_funcs=np.atleast_1d(acq_func), axes=axes, **kwargs) - def plot_constraint(self, axes: Tuple = (0, 1), **kwargs): + def plot_validity(self, axes: Tuple = (0, 1), **kwargs): """Plot the modeled constraint over test inputs sampling the limits of the parameter space. Parameters @@ -917,3 +923,7 @@ def plot_constraint(self, axes: Tuple = (0, 1), **kwargs): def plot_history(self, **kwargs): """Plot the improvement of the agent over time.""" plotting._plot_history(self, **kwargs) + + @property + def latent_transforms(self): + return {obj.name: obj.model.covar_module.latent_transform for obj in self.active_objs} diff --git a/blop/bayesian/__init__.py b/blop/bayesian/__init__.py index e49a76f..e69de29 100644 --- a/blop/bayesian/__init__.py +++ b/blop/bayesian/__init__.py @@ -1,3 +0,0 @@ -from .agent import * # noqa F401 -from .dofs import * # noqa F401 -from .objectives import * # noqa F401 diff --git a/blop/bayesian/acquisition/__init__.py b/blop/bayesian/acquisition/__init__.py index a66afbb..24bdacf 100644 --- a/blop/bayesian/acquisition/__init__.py +++ b/blop/bayesian/acquisition/__init__.py @@ -18,15 +18,10 @@ def parse_acq_func_identifier(identifier): - acq_func_name = None - for _acq_func_name in config.keys(): - if identifier.lower() in config[_acq_func_name]["identifiers"]: - acq_func_name = _acq_func_name - - if acq_func_name is None: - raise ValueError(f'Unrecognized acquisition function identifier "{identifier}".') - - return acq_func_name + for acq_func_name in config.keys(): + if identifier.lower() in config[acq_func_name]["identifiers"]: + return acq_func_name + return None def get_acquisition_function(agent, identifier="qei", return_metadata=True, verbose=False, **acq_func_kwargs): @@ -35,6 +30,9 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb """ acq_func_name = parse_acq_func_identifier(identifier) + if acq_func_name is None: + raise ValueError(f'Unrecognized acquisition function identifier "{identifier}".') + acq_func_config = config["upper_confidence_bound"] if config[acq_func_name]["multitask_only"] and (len(agent.objectives) == 1): diff --git a/blop/bayesian/dofs.py b/blop/bayesian/dofs.py deleted file mode 100644 index ad3f39f..0000000 --- a/blop/bayesian/dofs.py +++ /dev/null @@ -1,271 +0,0 @@ -import time as ttime -import uuid -from collections.abc import Sequence -from dataclasses import dataclass, field -from typing import Tuple, Union - -import numpy as np -import pandas as pd -from ophyd import Signal, SignalRO - -DEFAULT_BOUNDS = (-5.0, +5.0) -DOF_FIELDS = ["description", "readback", "lower_limit", "upper_limit", "units", "active", "read_only", "tags"] - -numeric = Union[float, int] - - -class ReadOnlyError(Exception): - ... - - -def _validate_dofs(dofs): - dof_names = [dof.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) - - -@dataclass -class DOF: - """A degree of freedom (DOF), to be used by an agent. - - Parameters - ---------- - name: str - The name of the DOF. This is used as a key. - description: str - A longer name for the DOF. - device: Signal, optional - An ophyd device. If None, a dummy ophyd device is generated. - limits: tuple, optional - A tuple of the lower and upper limit of the DOF. If the DOF is not read-only, the agent - will not explore outside the limits. If the DOF is read-only, the agent will reject all - sampled data where the DOF is outside the limits. - read_only: bool - If True, the agent will not try to set the DOF. Must be set to True if the supplied ophyd - device is read-only. - active: bool - If True, the agent will try to use the DOF in its optimization. If False, the agent will - still read the DOF but not include it any model or acquisition function. - units: str - The units of the DOF (e.g. mm or deg). This is only for plotting and general housekeeping. - tags: list - A list of tags. These make it easier to subset large groups of dofs. - latent_group: optional - An agent will fit latent dimensions to all DOFs with the same latent_group. If None, the - DOF will be modeled independently. - """ - - device: Signal = None - description: str = None - name: str = None - limits: Tuple[float, float] = (-10.0, 10.0) - units: str = "" - read_only: bool = False - active: bool = True - tags: list = field(default_factory=list) - log: bool = False - latent_group: str = None - - # Some post-processing. This is specific to dataclasses - def __post_init__(self): - self.uuid = str(uuid.uuid4()) - - if self.name is None: - self.name = self.device.name if hasattr(self.device, "name") else self.uuid - - if self.device is None: - self.device = Signal(name=self.name) - - if not self.read_only: - # check that the device has a put method - if isinstance(self.device, SignalRO): - raise ValueError("Must specify read_only=True for a read-only device!") - - if self.latent_group is None: - self.latent_group = str(uuid.uuid4()) - - # all dof degrees of freedom are hinted - self.device.kind = "hinted" - - @property - def lower_limit(self): - return float(self.limits[0]) - - @property - def upper_limit(self): - return float(self.limits[1]) - - @property - def readback(self): - return self.device.read()[self.device.name]["value"] - - @property - def summary(self) -> pd.Series: - series = pd.Series(index=DOF_FIELDS) - for attr in series.index: - series[attr] = getattr(self, attr) - return series - - @property - def label(self) -> str: - return f"{self.name}{f' [{self.units}]' if len(self.units) > 0 else ''}" - - @property - def has_model(self): - return hasattr(self, "model") - - -class DOFList(Sequence): - def __init__(self, dofs: list = []): - _validate_dofs(dofs) - self.dofs = dofs - - def __getitem__(self, i): - if type(i) is int: - return self.dofs[i] - elif type(i) is str: - return self.dofs[self.names.index(i)] - else: - raise ValueError(f"Invalid index {i}. A DOFList must be indexed by either an integer or a string.") - - def __len__(self): - return len(self.dofs) - - def __repr__(self): - return self.summary.__repr__() - - # def _repr_html_(self): - # return self.summary._repr_html_() - - @property - def summary(self) -> pd.DataFrame: - table = pd.DataFrame(columns=DOF_FIELDS) - for dof in self.dofs: - for attr in table.columns: - table.loc[dof.name, attr] = getattr(dof, attr) - - # convert dtypes - for attr in ["readback", "lower_limit", "upper_limit"]: - table[attr] = table[attr].astype(float) - - for attr in ["read_only", "active"]: - table[attr] = table[attr].astype(bool) - - return table - - @property - def names(self) -> list: - return [dof.name for dof in self.dofs] - - @property - def devices(self) -> list: - return [dof.device for dof in self.dofs] - - @property - def device_names(self) -> list: - return [dof.device.name for dof in self.dofs] - - @property - def lower_limits(self) -> np.array: - return np.array([dof.lower_limit for dof in self.dofs]) - - @property - def upper_limits(self) -> np.array: - return np.array([dof.upper_limit for dof in self.dofs]) - - @property - def limits(self) -> np.array: - """ - Returns a (n_dof, 2) array of bounds. - """ - return np.c_[self.lower_limits, self.upper_limits] - - @property - def readback(self) -> np.array: - return np.array([dof.readback for dof in self.dofs]) - - def add(self, dof): - _validate_dofs([*self.dofs, dof]) - self.dofs.append(dof) - - def _dof_read_only_mask(self, read_only=None): - return [dof.read_only == read_only if read_only is not None else True for dof in self.dofs] - - def _dof_active_mask(self, active=None): - return [dof.active == active if active is not None else True for dof in self.dofs] - - def _dof_tags_mask(self, tags=[]): - return [np.isin(dof["tags"], tags).any() if tags else True for dof in self.dofs] - - def _dof_mask(self, active=None, read_only=None, tags=[]): - return [ - (k and m and t) - for k, m, t in zip(self._dof_read_only_mask(read_only), self._dof_active_mask(active), self._dof_tags_mask(tags)) - ] - - def subset(self, active=None, read_only=None, tags=[]): - return DOFList([dof for dof, m in zip(self.dofs, self._dof_mask(active, read_only, tags)) if m]) - - def activate(self, read_only=None, active=None, tags=[]): - for dof in self._subset_dofs(read_only, active, tags): - dof.active = True - - def deactivate(self, read_only=None, active=None, tags=[]): - for dof in self._subset_dofs(read_only, active, tags): - dof.active = False - - -class BrownianMotion(SignalRO): - """ - Read-only degree of freedom simulating brownian motion - """ - - def __init__(self, name=None, theta=0.95, *args, **kwargs): - name = name if name is not None else str(uuid.uuid4()) - - super().__init__(name=name, *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() - - self.old_t = new_t - self.old_y = new_y - return new_y - - -class TimeReadback(SignalRO): - """ - Returns the current timestamp. - """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - def get(self): - return ttime.time() - - -class ConstantReadback(SignalRO): - """ - Returns a constant every time you read it (more useful than you'd think). - """ - - def __init__(self, constant=1, *args, **kwargs): - super().__init__(*args, **kwargs) - - self.constant = constant - - def get(self): - return self.constant diff --git a/blop/bayesian/kernels.py b/blop/bayesian/kernels.py index dc9f642..ecc7d5c 100644 --- a/blop/bayesian/kernels.py +++ b/blop/bayesian/kernels.py @@ -64,26 +64,23 @@ def __init__( self.n_skew_entries = len(self.skew_matrix_indices[0]) lengthscale_constraint = gpytorch.constraints.Positive() - raw_lengthscale_entries_initial = ( - lengthscale_constraint.inverse_transform(torch.tensor(1e-1)) - * torch.ones(self.num_outputs, self.num_inputs).double() + raw_lengthscales_initial = lengthscale_constraint.inverse_transform(torch.tensor(1e-1)) * torch.ones( + self.num_outputs, self.num_inputs, dtype=torch.double ) - self.register_parameter( - name="raw_lengthscale_entries", parameter=torch.nn.Parameter(raw_lengthscale_entries_initial) - ) - self.register_constraint(param_name="raw_lengthscale_entries", constraint=lengthscale_constraint) + self.register_parameter(name="raw_lengthscales", parameter=torch.nn.Parameter(raw_lengthscales_initial)) + self.register_constraint(param_name="raw_lengthscales", constraint=lengthscale_constraint) if priors: self.register_prior( - name="lengthscale_entries_prior", - prior=gpytorch.priors.GammaPrior(concentration=2, rate=1), - param_or_closure=lambda m: m.lengthscale_entries, - setting_closure=lambda m, v: m._set_lengthscale_entries(v), + name="lengthscales_prior", + prior=gpytorch.priors.GammaPrior(concentration=3, rate=6), + param_or_closure=lambda m: m.lengthscales, + setting_closure=lambda m, v: m._set_lengthscales(v), ) if self.n_skew_entries > 0: - skew_entries_constraint = gpytorch.constraints.Interval(-1e0, 1e0) + skew_entries_constraint = gpytorch.constraints.Interval(-2 * np.pi, 2 * np.pi) skew_entries_initial = torch.zeros((self.num_outputs, self.n_skew_entries), dtype=torch.float64) self.register_parameter(name="raw_skew_entries", parameter=torch.nn.Parameter(skew_entries_initial)) self.register_constraint(param_name="raw_skew_entries", constraint=skew_entries_constraint) @@ -94,7 +91,7 @@ def __init__( self.register_parameter( name="raw_outputscale", - parameter=torch.nn.Parameter(torch.ones(1)), + parameter=torch.nn.Parameter(torch.ones(1, dtype=torch.double)), ) self.register_constraint("raw_outputscale", constraint=outputscale_constraint) @@ -107,8 +104,8 @@ def __init__( ) @property - def lengthscale_entries(self): - return self.raw_lengthscale_entries_constraint.transform(self.raw_lengthscale_entries) + def lengthscales(self): + return self.raw_lengthscales_constraint.transform(self.raw_lengthscales) @property def skew_entries(self): @@ -118,9 +115,9 @@ def skew_entries(self): def outputscale(self): return self.raw_outputscale_constraint.transform(self.raw_outputscale) - @lengthscale_entries.setter - def lengthscale_entries(self, value): - self._set_lengthscale_entries(value) + @lengthscales.setter + def lengthscales(self, value): + self._set_lengthscales(value) @skew_entries.setter def skew_entries(self, value): @@ -130,10 +127,10 @@ def skew_entries(self, value): def outputscale(self, value): self._set_outputscale(value) - def _set_lengthscale_entries(self, value): + def _set_lengthscales(self, value): if not torch.is_tensor(value): - value = torch.as_tensor(value).to(self.raw_lengthscale_entries) - self.initialize(raw_lengthscale_entries=self.raw_lengthscale_entries_constraint.inverse_transform(value)) + value = torch.as_tensor(value).to(self.raw_lengthscales) + self.initialize(raw_lengthscales=self.raw_lengthscales_constraint.inverse_transform(value)) def _set_skew_entries(self, value): if not torch.is_tensor(value): @@ -157,7 +154,7 @@ def skew_matrix(self): @property def diag_matrix(self): D = torch.zeros((self.num_outputs, self.num_inputs, self.num_inputs), dtype=torch.float64) - D[self.diag_matrix_indices] = self.lengthscale_entries.ravel() ** -1 + D[self.diag_matrix_indices] = self.lengthscales.ravel() ** -1 return D @property diff --git a/blop/bayesian/models.py b/blop/bayesian/models.py index a50e84e..c0d165b 100644 --- a/blop/bayesian/models.py +++ b/blop/bayesian/models.py @@ -20,27 +20,15 @@ def __init__(self, train_inputs, train_targets, skew_dims=True, *args, **kwargs) **kwargs ) + self.trained = False -class TargetingGP(botorch.models.gp_regression.SingleTaskGP): - def __init__(self, train_inputs, train_targets, skew_dims=True, *args, **kwargs): - super().__init__(train_inputs, train_targets, *args, **kwargs) - - self.mean_module = gpytorch.means.ConstantMean(constant_prior=gpytorch.priors.NormalPrior(loc=0, scale=1)) - self.covar_module = kernels.LatentKernel( - num_inputs=train_inputs.shape[-1], - num_outputs=train_targets.shape[-1], - skew_dims=skew_dims, - priors=True, - scale=True, - **kwargs - ) - - -class LatentDirichletClassifier(LatentGP): +class LatentDirichletModel(LatentGP): def __init__(self, train_inputs, train_targets, skew_dims=True, *args, **kwargs): super().__init__(train_inputs, train_targets, skew_dims, *args, **kwargs) + self.trained = False + def probabilities(self, x, n_samples=1024): """ Takes in a (..., m) dimension tensor and returns a (..., n_classes) tensor diff --git a/blop/bayesian/objectives.py b/blop/bayesian/objectives.py deleted file mode 100644 index 61c8e06..0000000 --- a/blop/bayesian/objectives.py +++ /dev/null @@ -1,170 +0,0 @@ -from collections.abc import Sequence -from dataclasses import dataclass -from typing import Tuple, Union - -import numpy as np -import pandas as pd - -numeric = Union[float, int] - -DEFAULT_MINIMUM_SNR = 1e2 -OBJ_FIELDS = ["description", "target", "active", "limits", "weight", "log", "n", "snr", "min_snr"] - - -class DuplicateNameError(ValueError): - ... - - -def _validate_objectives(objectives): - names = [obj.name for obj in objectives] - unique_names, counts = np.unique(names, return_counts=True) - duplicate_names = unique_names[counts > 1] - if len(duplicate_names) > 0: - raise DuplicateNameError(f"Duplicate name(s) in supplied objectives: {duplicate_names}") - - -@dataclass -class Objective: - """An objective to be used by an agent. - - Parameters - ---------- - name: str - The name of the objective. This is used as a key. - description: str - A longer description for the objective. - target: float or str - One of 'min' or 'max', or a number. The agent will respectively minimize or maximize the - objective, or target the supplied number. - log: bool - Whether to apply a log to the objective, to make it more Gaussian. - weight: float - The relative importance of this objective, used when scalarizing in multi-objective optimization. - active: bool - If True, the agent will care about this objective during optimization. - limits: tuple of floats - The range of reliable measurements for the obejctive. Outside of this, data points will be rejected. - min_snr: float - The minimum signal-to-noise ratio of the objective, used when fitting the model. - units: str - A label representing the units of the objective. - """ - - name: str - description: str = None - target: Union[float, str] = "max" - log: bool = False - weight: numeric = 1.0 - active: bool = True - limits: Tuple[numeric, numeric] = None - min_snr: numeric = DEFAULT_MINIMUM_SNR - units: str = None - - def __post_init__(self): - if self.limits is None: - if self.log: - self.limits = (0, np.inf) - else: - self.limits = (-np.inf, np.inf) - - if type(self.target) is str: - if self.target not in ["min", "max"]: - raise ValueError("'target' must be either 'min', 'max', or a number.") - - @property - def label(self): - return f"{'log ' if self.log else ''}{self.description}" - - @property - def summary(self): - series = pd.Series() - for col in OBJ_FIELDS: - series[col] = getattr(self, col) - return series - - def __repr__(self): - return self.summary.__repr__() - - @property - def noise(self): - return self.model.likelihood.noise.item() if hasattr(self, "model") else None - - @property - def snr(self): - return np.round(1 / self.model.likelihood.noise.sqrt().item(), 1) if hasattr(self, "model") else None - - @property - def n(self): - return self.model.train_targets.shape[0] if hasattr(self, "model") else 0 - - -class ObjectiveList(Sequence): - def __init__(self, objectives: list = []): - _validate_objectives(objectives) - self.objectives = objectives - - def __getitem__(self, i): - if type(i) is int: - return self.objectives[i] - elif type(i) is str: - return self.objectives[self.names.index(i)] - else: - raise ValueError(f"Invalid index {i}. An ObjectiveList must be indexed by either an integer or a string.") - - def __len__(self): - return len(self.objectives) - - @property - def summary(self): - summary = pd.DataFrame(columns=OBJ_FIELDS) - for obj in self.objectives: - for col in summary.columns: - summary.loc[obj.name, col] = getattr(obj, col) - - # convert dtypes - for attr in ["log"]: - summary[attr] = summary[attr].astype(bool) - - return summary - - def __repr__(self): - return self.summary.__repr__() - - @property - def descriptions(self) -> list: - """ - Returns an array of the objective names. - """ - return [obj.description for obj in self.objectives] - - @property - def names(self) -> list: - """ - Returns an array of the objective names. - """ - return [obj.name for obj in self.objectives] - - @property - def targets(self) -> np.array: - """ - Returns an array of the objective targets. - """ - return [obj.target for obj in self.objectives] - - @property - def weights(self) -> np.array: - """ - Returns an array of the objective weights. - """ - return np.array([obj.weight for obj in self.objectives]) - - @property - def signed_weights(self) -> np.array: - """ - Returns a signed array of the objective weights. - """ - return np.array([(1 if obj.target == "max" else -1) * obj.weight for obj in self.objectives]) - - def add(self, objective): - _validate_objectives([*self.objectives, objective]) - self.objectives.append(objective) diff --git a/blop/bayesian/plotting.py b/blop/bayesian/plotting.py index fe5fb9f..a1c8917 100644 --- a/blop/bayesian/plotting.py +++ b/blop/bayesian/plotting.py @@ -31,7 +31,7 @@ def _plot_objs_one_dof(agent, size=16, lw=1e0): color = DEFAULT_COLOR_LIST[obj_index] - test_inputs = agent.test_inputs_grid() + test_inputs = agent.sample(method="grid") test_x = test_inputs[..., 0].detach().numpy() test_posterior = obj.model.posterior(test_inputs) @@ -50,7 +50,7 @@ def _plot_objs_one_dof(agent, size=16, lw=1e0): alpha=0.5**z, ) - agent.obj_axes[obj_index].set_xlim(*x_dof.limits) + agent.obj_axes[obj_index].set_xlim(*x_dof.search_bounds) agent.obj_axes[obj_index].set_xlabel(x_dof.label) agent.obj_axes[obj_index].set_ylabel(obj.label) @@ -67,10 +67,10 @@ def _plot_objs_many_dofs(agent, axes=(0, 1), shading="nearest", cmap=DEFAULT_COL agent.obj_fig, agent.obj_axes = plt.subplots( len(agent.objectives), - 3, - figsize=(10, 4 * len(agent.objectives)), + 4, + figsize=(12, 4 * len(agent.objectives)), constrained_layout=True, - dpi=256, + dpi=160, ) agent.obj_axes = np.atleast_2d(agent.obj_axes) @@ -82,7 +82,7 @@ def _plot_objs_many_dofs(agent, axes=(0, 1), shading="nearest", cmap=DEFAULT_COL # test_inputs has shape (*input_shape, 1, n_active_dofs) # test_x and test_y should be squeezeable - test_inputs = agent.test_inputs_grid() if gridded else agent.test_inputs(n=1024) + test_inputs = agent.sample(method="grid") if gridded else agent.sample(n=1024) test_x = test_inputs[..., 0, axes[0]].detach().squeeze().numpy() test_y = test_inputs[..., 0, axes[1]].detach().squeeze().numpy() @@ -144,7 +144,7 @@ def _plot_objs_many_dofs(agent, axes=(0, 1), shading="nearest", cmap=DEFAULT_COL obj_cbar.set_label(obj.label) err_cbar.set_label(f"{obj.label} error") - col_names = ["samples", "posterior mean", "posterior std. dev."] + col_names = ["samples", "posterior mean", "posterior std. dev.", "fitness"] pad = 5 @@ -179,8 +179,12 @@ def _plot_objs_many_dofs(agent, axes=(0, 1), shading="nearest", cmap=DEFAULT_COL for ax in agent.obj_axes.ravel(): ax.set_xlabel(x_dof.label) ax.set_ylabel(y_dof.label) - ax.set_xlim(*x_dof.limits) - ax.set_ylim(*y_dof.limits) + ax.set_xlim(*x_dof.search_bounds) + ax.set_ylim(*y_dof.search_bounds) + if x_dof.log: + ax.set_xscale("log") + if y_dof.log: + ax.set_yscale("log") def _plot_acqf_one_dof(agent, acq_funcs, lw=1e0, **kwargs): @@ -195,7 +199,7 @@ def _plot_acqf_one_dof(agent, acq_funcs, lw=1e0, **kwargs): agent.acq_axes = np.atleast_1d(agent.acq_axes) x_dof = agent.dofs.subset(active=True)[0] - test_inputs = agent.test_inputs_grid() + test_inputs = agent.sample(method="grid") for iacq_func, acq_func_identifier in enumerate(acq_funcs): color = DEFAULT_COLOR_LIST[iacq_func] @@ -205,7 +209,7 @@ def _plot_acqf_one_dof(agent, acq_funcs, lw=1e0, **kwargs): agent.acq_axes[iacq_func].plot(test_inputs.squeeze(-2), test_acqf, lw=lw, color=color) - agent.acq_axes[iacq_func].set_xlim(*x_dof.limits) + agent.acq_axes[iacq_func].set_xlim(*x_dof.search_bounds) agent.acq_axes[iacq_func].set_xlabel(x_dof.label) agent.acq_axes[iacq_func].set_ylabel(acq_func_meta["name"]) @@ -232,7 +236,7 @@ def _plot_acqf_many_dofs( x_dof, y_dof = plottable_dofs[axes[0]], plottable_dofs[axes[1]] # test_inputs has shape (..., 1, n_active_dofs) - test_inputs = agent.test_inputs_grid() if gridded else agent.test_inputs(n=1024) + test_inputs = agent.sample(n=1024, method="grid") if gridded else agent.sample(n=1024) *test_dim, input_dim = test_inputs.shape test_x = test_inputs[..., 0, axes[0]].detach().squeeze().numpy() test_y = test_inputs[..., 0, axes[1]].detach().squeeze().numpy() @@ -267,8 +271,12 @@ def _plot_acqf_many_dofs( for ax in agent.acq_axes.ravel(): ax.set_xlabel(x_dof.label) ax.set_ylabel(y_dof.label) - ax.set_xlim(*x_dof.limits) - ax.set_ylim(*y_dof.limits) + ax.set_xlim(*x_dof.search_bounds) + ax.set_ylim(*y_dof.search_bounds) + if x_dof.log: + ax.set_xscale("log") + if y_dof.log: + ax.set_yscale("log") def _plot_valid_one_dof(agent, size=16, lw=1e0): @@ -277,16 +285,16 @@ def _plot_valid_one_dof(agent, size=16, lw=1e0): x_dof = agent.dofs.subset(active=True)[0] x_values = agent.table.loc[:, x_dof.device.name].values - test_inputs = agent.test_inputs_grid() + test_inputs = agent.sample(method="grid") constraint = agent.constraint(test_inputs)[..., 0] agent.valid_ax.scatter(x_values, agent.all_objectives_valid, s=size) agent.valid_ax.plot(test_inputs.squeeze(-2), constraint, lw=lw) - agent.valid_ax.set_xlim(*x_dof.limits) + agent.valid_ax.set_xlim(*x_dof.search_bounds) def _plot_valid_many_dofs(agent, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, size=16, gridded=None): - agent.valid_fig, agent.valid_axes = plt.subplots(1, 2, figsize=(8, 4 * len(agent.objectives)), constrained_layout=True) + agent.valid_fig, agent.valid_axes = plt.subplots(1, 2, figsize=(8, 4), constrained_layout=True) plottable_dofs = agent.dofs.subset(active=True, read_only=False) @@ -296,7 +304,7 @@ def _plot_valid_many_dofs(agent, axes=[0, 1], shading="nearest", cmap=DEFAULT_CO x_dof, y_dof = plottable_dofs[axes[0]], plottable_dofs[axes[1]] # test_inputs has shape (..., 1, n_active_dofs) - test_inputs = agent.test_inputs_grid() if gridded else agent.test_inputs(n=1024) + test_inputs = agent.sample(method="grid") if gridded else agent.sample(n=1024) test_x = test_inputs[..., 0, axes[0]].detach().squeeze().numpy() test_y = test_inputs[..., 0, axes[1]].detach().squeeze().numpy() @@ -324,11 +332,15 @@ def _plot_valid_many_dofs(agent, axes=[0, 1], shading="nearest", cmap=DEFAULT_CO vmax=0, ) - for ax in agent.acq_axes.ravel(): + for ax in agent.valid_axes.ravel(): ax.set_xlabel(x_dof.label) ax.set_ylabel(y_dof.label) - ax.set_xlim(*x_dof.limits) - ax.set_ylim(*y_dof.limits) + ax.set_xlim(*x_dof.search_bounds) + ax.set_ylim(*y_dof.search_bounds) + if x_dof.log: + ax.set_xscale("log") + if y_dof.log: + ax.set_yscale("log") def _plot_history(agent, x_key="index", show_all_objs=False): diff --git a/blop/bayesian/transforms.py b/blop/bayesian/transforms.py index 1b28c6d..82bdb9d 100644 --- a/blop/bayesian/transforms.py +++ b/blop/bayesian/transforms.py @@ -7,6 +7,17 @@ from torch import Tensor +def targeting_transform(y, target): + if target == "max": + return y + if target == "min": + return -y + elif not isinstance(target, tuple): + return -(y - target).abs() + else: + return -((y - 0.5 * (target[1] + target[0])).abs() - 0.5 * (target[1] - target[0])).clamp(min=0) + + class TargetingPosteriorTransform(PosteriorTransform): r"""An affine posterior transform for scalarizing multi-output posteriors.""" @@ -23,22 +34,14 @@ def __init__(self, weights: Tensor, targets: Tensor) -> None: self.targets = targets self.register_buffer("weights", weights) - def sampled_transform(self, y): + def sample_transform(self, y): for i, target in enumerate(self.targets): - if target == "min": - y[..., i] = -y[..., i] - elif target != "max": - y[..., i] = -(y[..., i] - target).abs() - + y[..., i] = targeting_transform(y[..., i], target) return y @ self.weights.unsqueeze(-1) def mean_transform(self, mean, var): for i, target in enumerate(self.targets): - if target == "min": - mean[..., i] = -mean[..., i] - elif target != "max": - mean[..., i] = -(mean[..., i] - target).abs() - + mean[..., i] = targeting_transform(mean[..., i], target) return mean @ self.weights.unsqueeze(-1) def variance_transform(self, mean, var): @@ -53,7 +56,7 @@ def evaluate(self, Y: Tensor) -> Tensor: Returns: A `batch_shape x q`-dim tensor of transformed outcomes. """ - return self.sampled_transform(Y) + return self.sample_transform(Y) def forward(self, posterior: Union[GPyTorchPosterior, PosteriorList]) -> GPyTorchPosterior: r"""Compute the posterior of the affine transformation. @@ -68,7 +71,7 @@ def forward(self, posterior: Union[GPyTorchPosterior, PosteriorList]) -> GPyTorc return botorch.posteriors.transformed.TransformedPosterior( posterior, - sample_transform=self.sampled_transform, + sample_transform=self.sample_transform, mean_transform=self.mean_transform, variance_transform=self.variance_transform, ) diff --git a/blop/bayesian/digestion.py b/blop/digestion.py similarity index 100% rename from blop/bayesian/digestion.py rename to blop/digestion.py diff --git a/blop/dofs.py b/blop/dofs.py new file mode 100644 index 0000000..7e2ae6f --- /dev/null +++ b/blop/dofs.py @@ -0,0 +1,302 @@ +import time as ttime +import uuid +from collections.abc import Iterable, Sequence +from dataclasses import dataclass, field +from typing import Tuple + +import numpy as np +import pandas as pd +from ophyd import Signal, SignalRO + +DOF_FIELD_TYPES = { + "description": "str", + "readback": "float", + "search_bounds": "object", + "trust_bounds": "object", + "units": "str", + "active": "bool", + "read_only": "bool", + "log": "bool", + "tags": "object", +} + + +class ReadOnlyError(Exception): + ... + + +def _validate_dofs(dofs): + dof_names = [dof.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) + + +@dataclass +class DOF: + """A degree of freedom (DOF), to be used by an agent. + + Parameters + ---------- + name: str + The name of the DOF. This is used as a key to index observed data. + description: str, optional + A longer name for the DOF. + units: str + The units of the DOF (e.g. mm or deg). This is just for plotting and general sanity checking. + search_bounds: tuple + A tuple of the lower and upper limit of the DOF for the agent to search. + trust_bounds: tuple, optional + The agent will reject all data where the DOF value is outside the trust bounds. Must be larger than search bounds. + read_only: bool + If True, the agent will not try to set the DOF. Must be set to True if the supplied ophyd + device is read-only. + active: bool + If True, the agent will try to use the DOF in its optimization. If False, the agent will + still read the DOF but not include it any model or acquisition function. + log: bool + Whether to apply a log to the objective, i.e. to make the process outputs more Gaussian. + tags: list + A list of tags. These make it easier to subset large groups of dofs. + device: Signal, optional + An ophyd device. If not supplied, a dummy ophyd device will be generated. + """ + + name: str = None + description: str = "" + search_bounds: Tuple[float, float] = None + trust_bounds: Tuple[float, float] = None + units: str = "" + read_only: bool = False + active: bool = True + log: bool = False + tags: list = field(default_factory=list) + device: Signal = None + + # Some post-processing. This is specific to dataclasses + def __post_init__(self): + if self.search_bounds is None: + if not self.read_only: + raise ValueError("You must specify search_bounds if the device is not read-only.") + else: + self.search_bounds = tuple(self.search_bounds) + if len(self.search_bounds) != 2: + raise ValueError("'search_bounds' must be a 2-tuple of floats.") + if self.search_bounds[0] > self.search_bounds[1]: + raise ValueError("The lower search bound must be less than the upper search bound.") + + if self.trust_bounds is not None: + self.trust_bounds = tuple(self.trust_bounds) + if not self.read_only: + if (self.search_bounds[0] < self.trust_bounds[0]) or (self.search_bounds[1] > self.trust_bounds[1]): + raise ValueError("Trust bounds must be larger than search bounds.") + + if (self.name is None) ^ (self.device is None): + if self.name is None: + self.name = self.device.name + if self.device is None: + self.device = Signal(name=self.name) + else: + raise ValueError("DOF() accepts exactly one of either a name or an ophyd device.") + + if not self.read_only: + # check that the device has a put method + if isinstance(self.device, SignalRO): + raise ValueError("You must specify read_only=True for a read-only device.") + + if self.log: + if not self.search_bounds[0] > 0: + raise ValueError("Search bounds must be strictly positive if log=True.") + + # all dof degrees of freedom are hinted + self.device.kind = "hinted" + + @property + def _search_bounds(self): + if self.read_only: + _readback = self.readback + return (_readback, _readback) + return self.search_bounds + + @property + def _trust_bounds(self): + if self.trust_bounds is None: + return (0, np.inf) if self.log else (-np.inf, np.inf) + return self.trust_bounds + + @property + def readback(self): + return self.device.read()[self.device.name]["value"] + + @property + def summary(self) -> pd.Series: + series = pd.Series(index=list(DOF_FIELD_TYPES.keys()), dtype="object") + for attr in series.index: + series[attr] = getattr(self, attr) + return series + + @property + def label(self) -> str: + return f"{self.description}{f' [{self.units}]' if len(self.units) > 0 else ''}" + + @property + def has_model(self): + return hasattr(self, "model") + + +class DOFList(Sequence): + def __init__(self, dofs: list = []): + _validate_dofs(dofs) + self.dofs = dofs + + def __getattr__(self, attr): + # This is called if we can't find the attribute in the normal way. + if attr in DOF_FIELD_TYPES.keys(): + return np.array([getattr(dof, attr) for dof in self.dofs]) + if attr in self.names: + return self.__getitem__(attr) + + raise AttributeError(f"DOFList object has no attribute named '{attr}'.") + + def __getitem__(self, key): + if isinstance(key, str): + if key not in self.names: + raise ValueError(f"DOFList has no DOF named {key}.") + return self.dofs[self.names.index(key)] + elif isinstance(key, Iterable): + return [self.dofs[_key] for _key in key] + elif isinstance(key, slice): + return [self.dofs[i] for i in range(*key.indices(len(self)))] + elif isinstance(key, int): + return self.dofs[key] + else: + raise ValueError(f"Invalid index {key}.") + + def __len__(self): + return len(self.dofs) + + def __repr__(self): + return self.summary.__repr__() + + def __repr_html__(self): + return self.summary.__repr_html__() + + @property + def summary(self) -> pd.DataFrame: + table = pd.DataFrame(columns=list(DOF_FIELD_TYPES.keys()), index=self.names) + + for dof in self.dofs: + for attr, value in dof.summary.items(): + table.at[dof.name, attr] = value + + for attr, dtype in DOF_FIELD_TYPES.items(): + table[attr] = table[attr].astype(dtype) + + return table + + @property + def names(self) -> list: + return [dof.name for dof in self.dofs] + + @property + def devices(self) -> list: + return [dof.device for dof in self.dofs] + + @property + def search_bounds(self) -> np.array: + """ + Returns a (n_dof, 2) array of bounds. + """ + return np.array([dof._search_bounds for dof in self.dofs]) + + @property + def trust_bounds(self) -> np.array: + """ + Returns a (n_dof, 2) array of bounds. + """ + return np.array([dof._trust_bounds for dof in self.dofs]) + + def add(self, dof): + _validate_dofs([*self.dofs, dof]) + self.dofs.append(dof) + + @staticmethod + def _test_dof(dof, active=None, read_only=None, tag=None): + if active is not None: + if dof.active != active: + return False + if read_only is not None: + if dof.read_only != read_only: + return False + if tag is not None: + if not np.isin(np.atleast_1d(tag), dof.tags).any(): + return False + return True + + def subset(self, active=None, read_only=None, tag=None): + return DOFList([dof for dof in self.dofs if self._test_dof(dof, active=active, read_only=read_only, tag=tag)]) + + def activate(self, active=None, read_only=None, tag=None): + for dof in self.dofs: + if self._test_dof(dof, active=active, read_only=read_only, tag=tag): + dof.active = True + + def deactivate(self, active=None, read_only=None, tag=None): + for dof in self.dofs: + if self._test_dof(dof, active=active, read_only=read_only, tag=tag): + dof.active = False + + +class BrownianMotion(SignalRO): + """ + Read-only degree of freedom simulating brownian motion + """ + + def __init__(self, name=None, theta=0.95, *args, **kwargs): + name = name if name is not None else str(uuid.uuid4()) + + super().__init__(name=name, *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() + + self.old_t = new_t + self.old_y = new_y + return new_y + + +class TimeReadback(SignalRO): + """ + Returns the current timestamp. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def get(self): + return ttime.time() + + +class ConstantReadback(SignalRO): + """ + Returns a constant every time you read it (more useful than you'd think). + """ + + def __init__(self, constant=1, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.constant = constant + + def get(self): + return self.constant diff --git a/blop/objectives.py b/blop/objectives.py new file mode 100644 index 0000000..8c29d73 --- /dev/null +++ b/blop/objectives.py @@ -0,0 +1,251 @@ +from collections.abc import Iterable, Sequence +from dataclasses import dataclass, field +from typing import List, Tuple, Union + +import numpy as np +import pandas as pd +import torch +from torch.special import erf + +DEFAULT_MIN_NOISE_LEVEL = 1e-6 +DEFAULT_MAX_NOISE_LEVEL = 1e0 + +OBJ_FIELD_TYPES = { + "description": "object", + "target": "object", + "active": "bool", + "trust_bounds": "object", + "active": "bool", + "weight": "bool", + "units": "object", + "log": "bool", + "min_noise": "float", + "max_noise": "float", + "noise": "float", + "n": "int", + "latent_groups": "object", +} + + +class DuplicateNameError(ValueError): + ... + + +def _validate_objectives(objectives): + names = [obj.name for obj in objectives] + unique_names, counts = np.unique(names, return_counts=True) + duplicate_names = unique_names[counts > 1] + if len(duplicate_names) > 0: + raise DuplicateNameError(f"Duplicate name(s) in supplied objectives: {duplicate_names}") + + +@dataclass +class Objective: + """An objective to be used by an agent. + + Parameters + ---------- + name: str + The name of the objective. This is used as a key to index observed data. + description: str + A longer description for the objective. + target: str or float or tuple + One of "min", "max" , a float, or a tuple of floats. The agent will respectively minimize or maximize the + objective, target the supplied number, or target the interval of the tuple of numbers. + log: bool + Whether to apply a log to the objective, i.e. to make the process more stationary. + weight: float + The relative importance of this objective, to be used when scalarizing in multi-objective optimization. + active: bool + If True, the agent will care about this objective during optimization. + limits: tuple of floats + The range of reliable measurements for the objective. Outside of this, data points will be ignored. + min_noise: float + The minimum noise level of the fitted model. + max_noise: float + The maximum noise level of the fitted model. + units: str + A label representing the units of the objective. + latent_groups: list of tuples of strs, optional + An agent will fit latent dimensions to all DOFs with the same latent_group. All other + DOFs will be modeled independently. + """ + + name: str + description: str = "" + target: Union[Tuple[float, float], float, str] = "max" + log: bool = False + weight: float = 1.0 + active: bool = True + trust_bounds: Tuple[float, float] or None = None + min_noise: float = DEFAULT_MIN_NOISE_LEVEL + max_noise: float = DEFAULT_MAX_NOISE_LEVEL + units: str = None + latent_groups: List[Tuple[str, ...]] = field(default_factory=list) + + def __post_init__(self): + if isinstance(self.target, str): + if self.target not in ["min", "max"]: + raise ValueError("'target' must be either 'min', 'max', a number, or a tuple of numbers.") + + if isinstance(self.target, float): + if self.log and not self.target > 0: + return ValueError("'target' must strictly positive if log=True.") + + self.use_as_constraint = True if isinstance(self.target, tuple) else False + + @property + def _trust_bounds(self): + if self.trust_bounds is None: + return (0, np.inf) if self.log else (-np.inf, np.inf) + return self.trust_bounds + + @property + def label(self) -> str: + return f"{'log ' if self.log else ''}{self.description}" + + @property + def summary(self) -> pd.Series: + series = pd.Series(index=list(OBJ_FIELD_TYPES.keys()), dtype="object") + for attr in series.index: + value = getattr(self, attr) + if attr == "trust_bounds": + if value is None: + value = (0, np.inf) if self.log else (-np.inf, np.inf) + series[attr] = value + return series + + @property + def trust_lower_bound(self): + if self.trust_bounds is None: + return 0 if self.log else -np.inf + return float(self.trust_bounds[0]) + + @property + def trust_upper_bound(self): + if self.trust_bounds is None: + return np.inf + return float(self.trust_bounds[1]) + + @property + def noise(self) -> float: + return self.model.likelihood.noise.item() if hasattr(self, "model") else None + + @property + def snr(self) -> float: + return np.round(1 / self.model.likelihood.noise.sqrt().item(), 3) if hasattr(self, "model") else None + + @property + def n(self) -> int: + return self.model.train_targets.shape[0] if hasattr(self, "model") else 0 + + def targeting_constraint(self, x: torch.Tensor) -> torch.Tensor: + if not isinstance(self.target, tuple): + return None + + a, b = self.target + p = self.model.posterior(x) + m = p.mean + s = p.variance.sqrt() + + return 0.5 * (erf((b - m) / (np.sqrt(2) * s)) - erf((a - m) / (np.sqrt(2) * s)))[..., -1] + + +class ObjectiveList(Sequence): + def __init__(self, objectives: list = []): + _validate_objectives(objectives) + self.objectives = objectives + + def __getattr__(self, attr): + # This is called if we can't find the attribute in the normal way. + if attr in OBJ_FIELD_TYPES.keys(): + return np.array([getattr(obj, attr) for obj in self.objectives]) + if attr in self.names: + return self.__getitem__(attr) + + raise AttributeError(f"ObjectiveList object has no attribute named '{attr}'.") + + def __getitem__(self, key): + if isinstance(key, str): + if key not in self.names: + raise ValueError(f"ObjectiveList has no objective named {key}.") + return self.objectives[self.names.index(key)] + elif isinstance(key, Iterable): + return [self.objectives[_key] for _key in key] + elif isinstance(key, slice): + return [self.objectives[i] for i in range(*key.indices(len(self)))] + elif isinstance(key, int): + return self.objectives[key] + else: + raise ValueError(f"Invalid index {key}.") + + def __len__(self): + return len(self.objectives) + + @property + def summary(self) -> pd.DataFrame: + table = pd.DataFrame(columns=list(OBJ_FIELD_TYPES.keys()), index=self.names) + + for obj in self.objectives: + for attr, value in obj.summary.items(): + table.at[obj.name, attr] = value + + for attr, dtype in OBJ_FIELD_TYPES.items(): + table[attr] = table[attr].astype(dtype) + + return table + + def __repr__(self): + return self.summary.__repr__() + + def __repr_html__(self): + return self.summary.__repr_html__() + + @property + def descriptions(self) -> list: + """ + Returns an array of the objective names. + """ + return [obj.description for obj in self.objectives] + + @property + def names(self) -> list: + """ + Returns an array of the objective names. + """ + return [obj.name for obj in self.objectives] + + @property + def targets(self) -> list: + """ + Returns an array of the objective targets. + """ + return [obj.target for obj in self.objectives] + + @property + def weights(self) -> np.array: + """ + Returns an array of the objective weights. + """ + return np.array([obj.weight for obj in self.objectives]) + + @property + def signed_weights(self) -> np.array: + """ + Returns a signed array of the objective weights. + """ + return np.array([(1 if obj.target == "max" else -1) * obj.weight for obj in self.objectives]) + + def add(self, objective): + _validate_objectives([*self.objectives, objective]) + self.objectives.append(objective) + + @staticmethod + def _test_obj(obj, active=None): + if active is not None: + if obj.active != active: + return False + return True + + def subset(self, active=None): + return ObjectiveList([obj for obj in self.objectives if self._test_obj(obj, active=active)]) diff --git a/blop/bayesian/plans.py b/blop/plans.py similarity index 100% rename from blop/bayesian/plans.py rename to blop/plans.py diff --git a/blop/tests/conftest.py b/blop/tests/conftest.py index 3531243..6d3e781 100644 --- a/blop/tests/conftest.py +++ b/blop/tests/conftest.py @@ -8,8 +8,8 @@ from databroker import Broker from ophyd.utils import make_dir_tree -from blop.bayesian import DOF, Agent, Objective -from blop.bayesian.dofs import BrownianMotion +from blop import DOF, Agent, Objective +from blop.dofs import BrownianMotion from blop.utils import functions @@ -51,8 +51,8 @@ def agent(db): """ dofs = [ - DOF(name="x1", limits=(-8.0, 8.0)), - DOF(name="x2", limits=(-8.0, 8.0)), + DOF(name="x1", search_bounds=(-8.0, 8.0)), + DOF(name="x2", search_bounds=(-8.0, 8.0)), ] objectives = [Objective(name="himmelblau", target="min")] @@ -70,7 +70,7 @@ def agent(db): @pytest.fixture(scope="function") -def multi_agent(db): +def agent_with_multiple_objs(db): """ A simple agent minimizing two Himmelblau's functions """ @@ -85,8 +85,8 @@ def digestion(db, uid): return products dofs = [ - DOF(name="x1", limits=(-5.0, 5.0)), - DOF(name="x2", limits=(-5.0, 5.0)), + DOF(name="x1", search_bounds=(-5.0, 5.0)), + DOF(name="x2", search_bounds=(-5.0, 5.0)), ] objectives = [Objective(name="obj1", target="min"), Objective(name="obj2", target="min")] @@ -110,11 +110,11 @@ def agent_with_passive_dofs(db): """ dofs = [ - DOF(name="x1", limits=(-5.0, 5.0)), - DOF(name="x2", limits=(-5.0, 5.0)), - DOF(name="x3", limits=(-5.0, 5.0), active=False), - DOF(BrownianMotion(name="brownian1"), read_only=True), - DOF(BrownianMotion(name="brownian2"), read_only=True, active=False), + DOF(name="x1", search_bounds=(-5.0, 5.0)), + DOF(name="x2", search_bounds=(-5.0, 5.0)), + DOF(name="x3", search_bounds=(-5.0, 5.0), active=False), + DOF(device=BrownianMotion(name="brownian1"), read_only=True), + DOF(device=BrownianMotion(name="brownian2"), read_only=True, active=False), ] objectives = [ diff --git a/blop/tests/test_acq_funcs.py b/blop/tests/test_acq_funcs.py index fbfd5ff..be376b7 100644 --- a/blop/tests/test_acq_funcs.py +++ b/blop/tests/test_acq_funcs.py @@ -3,23 +3,23 @@ @pytest.mark.parametrize("acq_func", ["ei", "pi", "em", "ucb"]) def test_analytic_acq_funcs_single_objective(agent, RE, acq_func): - RE(agent.learn("qr", n=16)) + RE(agent.learn("qr", n=4)) RE(agent.learn(acq_func, n=1)) @pytest.mark.parametrize("acq_func", ["qei", "qpi", "qem", "qucb"]) def test_monte_carlo_acq_funcs_single_objective(agent, RE, acq_func): - RE(agent.learn("qr", n=16)) + RE(agent.learn("qr", n=4)) RE(agent.learn(acq_func, n=4)) @pytest.mark.parametrize("acq_func", ["ei", "pi", "em", "ucb"]) -def test_analytic_acq_funcs_multi_objective(multi_agent, RE, acq_func): - RE(multi_agent.learn("qr", n=16)) - RE(multi_agent.learn(acq_func, n=1)) +def test_analytic_acq_funcs_multi_objective(agent_with_multiple_objs, RE, acq_func): + RE(agent_with_multiple_objs.learn("qr", n=16)) + RE(agent_with_multiple_objs.learn(acq_func, n=1)) @pytest.mark.parametrize("acq_func", ["qei", "qpi", "qem", "qucb"]) -def test_monte_carlo_acq_funcs_multi_objective(multi_agent, RE, acq_func): - RE(multi_agent.learn("qr", n=16)) - RE(multi_agent.learn(acq_func, n=4)) +def test_monte_carlo_acq_funcs_multi_objective(agent_with_multiple_objs, RE, acq_func): + RE(agent_with_multiple_objs.learn("qr", n=16)) + RE(agent_with_multiple_objs.learn(acq_func, n=4)) diff --git a/blop/tests/test_agent.py b/blop/tests/test_agent.py new file mode 100644 index 0000000..22446c3 --- /dev/null +++ b/blop/tests/test_agent.py @@ -0,0 +1,10 @@ +import pytest # noqa F401 + + +def test_agent(agent, RE): + RE(agent.learn("qr", n=4)) + + +def test_forget(agent, RE): + RE(agent.learn("qr", n=4)) + agent.forget(last=2) diff --git a/blop/tests/test_napari.py b/blop/tests/test_napari.py new file mode 100644 index 0000000..231f280 --- /dev/null +++ b/blop/tests/test_napari.py @@ -0,0 +1,7 @@ +# import pytest # noqa F401 + + +# @pytest.mark.parametrize("item", ["mean", "error", "qei"]) +# def test_napari_viewer(agent, RE, item): +# RE(agent.learn("qr", n=4)) +# agent.view(item) diff --git a/blop/tests/test_passive_dofs.py b/blop/tests/test_passive_dofs.py index 797edde..b2034a9 100644 --- a/blop/tests/test_passive_dofs.py +++ b/blop/tests/test_passive_dofs.py @@ -4,7 +4,3 @@ @pytest.mark.test_func def test_passive_dofs(agent_with_passive_dofs, RE): RE(agent_with_passive_dofs.learn("qr", n=32)) - - agent_with_passive_dofs.plot_objectives() - agent_with_passive_dofs.plot_acquisition() - agent_with_passive_dofs.plot_constraint() diff --git a/blop/tests/test_plots.py b/blop/tests/test_plots.py index 706e042..f28758a 100644 --- a/blop/tests/test_plots.py +++ b/blop/tests/test_plots.py @@ -1,11 +1,28 @@ -import pytest +import pytest # noqa F401 -@pytest.mark.test_func def test_plots(RE, agent): - RE(agent.learn("qr", n=32)) + RE(agent.learn("qr", n=16)) agent.plot_objectives() agent.plot_acquisition() - agent.plot_constraint() + agent.plot_validity() agent.plot_history() + + +def test_plots_multiple_objs(RE, agent_with_multiple_objs): + RE(agent_with_multiple_objs.learn("qr", n=16)) + + agent_with_multiple_objs.plot_objectives() + agent_with_multiple_objs.plot_acquisition() + agent_with_multiple_objs.plot_validity() + agent_with_multiple_objs.plot_history() + + +def test_plots_passive_dofs(RE, agent_with_passive_dofs): + RE(agent_with_passive_dofs.learn("qr", n=16)) + + agent_with_passive_dofs.plot_objectives() + agent_with_passive_dofs.plot_acquisition() + agent_with_passive_dofs.plot_validity() + agent_with_passive_dofs.plot_history() diff --git a/blop/tests/test_read_write.py b/blop/tests/test_read_write.py index 9c225d9..db4b052 100644 --- a/blop/tests/test_read_write.py +++ b/blop/tests/test_read_write.py @@ -1,20 +1,16 @@ import pytest # noqa F401 -def test_agent(agent, RE): - RE(agent.learn("qr", n=16)) - - def test_agent_save_load_data(agent, RE): - RE(agent.learn("qr", n=16)) + RE(agent.learn("qr", n=4)) agent.save_data("/tmp/test_save_data.h5") agent.reset() - agent.load_data(data_file="/tmp/test_save_data.h5") - RE(agent.learn("qr", n=16)) + agent.load_data("/tmp/test_save_data.h5") + RE(agent.learn("qr", n=4)) def test_agent_save_load_hypers(agent, RE): - RE(agent.learn("qr", n=16)) + RE(agent.learn("qr", n=4)) agent.save_hypers("/tmp/test_save_hypers.h5") agent.reset() - RE(agent.learn("qr", n=16, hypers_file="/tmp/test_save_hypers.h5")) + RE(agent.learn("qr", n=16, hypers="/tmp/test_save_hypers.h5")) diff --git a/blop/utils/__init__.py b/blop/utils/__init__.py index 6f3bd85..0debc41 100644 --- a/blop/utils/__init__.py +++ b/blop/utils/__init__.py @@ -5,6 +5,10 @@ from ortools.constraint_solver import pywrapcp, routing_enums_pb2 +def cummax(x): + return [np.nanmax(x[: i + 1]) for i in range(len(np.atleast_1d(x)))] + + def sobol_sampler(bounds, n, q=1): """ Returns $n$ quasi-randomly sampled points within the bounds (a 2 by d tensor) diff --git a/blop/utils/functions.py b/blop/utils/functions.py index 810a517..0a2e949 100644 --- a/blop/utils/functions.py +++ b/blop/utils/functions.py @@ -183,6 +183,18 @@ def constrained_himmelblau_digestion(db, uid): return products +def himmelblau_digestion(db, uid): + """ + Digests Himmelblau's function into the feedback. + """ + products = db[uid].table() + + for index, entry in products.iterrows(): + products.loc[index, "himmelblau"] = himmelblau(entry.x1, entry.x2) + + return products + + def mock_kbs_digestion(db, uid): """ Digests a beam waist and height into the feedback. diff --git a/docs/source/agent.rst b/docs/source/agent.rst new file mode 100644 index 0000000..27ec35d --- /dev/null +++ b/docs/source/agent.rst @@ -0,0 +1,35 @@ +Agent ++++++ + +The blop ``Agent`` takes care of the entire optimization loop, from data acquisition to model fitting. + +.. code-block:: python + + from blop import DOF, Objective, Agent + + dofs = [ + DOF(name="x1", description="the first DOF", search_bounds=(-10, 10)) + DOF(name="x2", description="another DOF", search_bounds=(-5, 5)) + DOF(name="x3", description="ayet nother DOF", search_bounds=(0, 1)) + ] + + objective = [ + Objective(name="y1", description="something to minimize", target="min") + Objective(name="y2", description="something to maximize", target="max") + ] + + dets = [ + my_detector, # an ophyd device with a .trigger() method that determines "y1" + my_other_detector # a detector that measures "y2" + ] + + agent = Agent(dofs=dofs, objectives=objectives, dets=dets) + + +This creates an ``Agent`` with no data about the world, and thus no way to model it. +We have to start with asking the ``Agent`` to learn by randomly sampling the parameter space. +The ``Agent`` learns with Bluesky plans emitted by the ``agent.learn()`` method, which can be passed to a ``RunEngine``: + +.. code-block:: python + + RE(agent.learn("qr", n=16)) # the agent chooses 16 quasi-random points, samples them, and fits models to them diff --git a/docs/source/dofs.rst b/docs/source/dofs.rst new file mode 100644 index 0000000..42ec90a --- /dev/null +++ b/docs/source/dofs.rst @@ -0,0 +1,33 @@ +Degrees of freedom (DOFs) ++++++++++++++++++++++++++ + +A degree of freedom is a variable that affects our optimization objective. We can define a simple DOF as + +.. code-block:: python + + from blop import DOF + + dof = DOF(name="x1", description="my first DOF", search_bounds=(lower, upper)) + +This will instantiate a bunch of stuff under the hood, so that our agent knows how to move things and where to search. +Typically, this will correspond to a real, physical device available in Python. In that case, we can pass the DOF an ophyd device in place of a name + +.. code-block:: python + + from blop import DOF + + dof = DOF(device=my_ophyd_device, description="a real piece of hardware", search_bounds=(lower, upper)) + +In this case, the agent will control the device as it sees fit, moving it between the search bounds. + +Sometimes, a DOF may be something we can't directly control (e.g. a changing synchrotron current or a changing sample temperature) but want our agent to be aware of. +In this case, we can define a read-only DOF as + +.. code-block:: python + + from blop import DOF + + dof = DOF(device=a_read_only_ophyd_device, description="a thermometer or something", read_only=True, trust_bounds=(lower, upper)) + +and the agent will use the received values to model its objective, but won't try to move it. +We can also pass a set of ``trust_bounds``, so that our agent will ignore experiments where the DOF value jumps outside of the interval. diff --git a/docs/source/objectives.rst b/docs/source/objectives.rst new file mode 100644 index 0000000..69e489e --- /dev/null +++ b/docs/source/objectives.rst @@ -0,0 +1,34 @@ +Objectives +++++++++++ + +We can describe an optimization problem as a list of objectives to. A simple objective is + +.. code-block:: python + + from blop import Objective + + objective = Objective(name="y1", target="max") + +Given some data, the ``Objective`` object will try to model the quantity "y1" and find the corresponding inputs that maximize it. +The objective will expect that this quantity will be spit out by the experimentation loop, so we will check later that it is set up correctly. +There are many ways to specify an objective's behavior, which is done by changing the objective's target: + +.. code-block:: python + + from blop import Objective + + objective = Objective(name="y1", target="min") # minimize the quantity "y1" + + objective = Objective(name="y1", target=2) # get the quantity "y1" as close to 2 as possible + + objective = Objective(name="y1", target=(1, 3)) # find any input that puts "y1" between 1 and 3. + + +Often, the objective being modeled is some strictly positive quantity (e.g. the size of a beam being aligned). +In this case, a smart thing to do is to set ``log=True``, which will model and subsequently optimize the logarithm of the objective: + +.. code-block:: python + + from blop import Objective + + objective = Objective(name="some_strictly_positive_quantity", target="max", log=True) diff --git a/docs/source/release-history.rst b/docs/source/release-history.rst index 974375c..7f61bdb 100644 --- a/docs/source/release-history.rst +++ b/docs/source/release-history.rst @@ -2,6 +2,11 @@ Release History =============== +v0.6.0 (2024-01-30) +------------------- +- More sophisticated targeting capabilities for different objectives. +- More user-friendly agent controls. + v0.5.0 (2023-11-09) ------------------- - Added hypervolume acquisition and constraints. diff --git a/docs/source/tutorials/himmelblau.ipynb b/docs/source/tutorials/himmelblau.ipynb index b3e15d4..8b1a435 100644 --- a/docs/source/tutorials/himmelblau.ipynb +++ b/docs/source/tutorials/himmelblau.ipynb @@ -69,11 +69,11 @@ "metadata": {}, "outputs": [], "source": [ - "from blop.bayesian import DOF\n", + "from blop import DOF\n", "\n", "dofs = [\n", - " DOF(name=\"x1\", limits=(-6, 6)),\n", - " DOF(name=\"x2\", limits=(-6, 6)),\n", + " DOF(name=\"x1\", search_bounds=(-6, 6)),\n", + " DOF(name=\"x2\", search_bounds=(-6, 6)),\n", "]" ] }, @@ -92,7 +92,7 @@ "metadata": {}, "outputs": [], "source": [ - "from blop.bayesian import Objective\n", + "from blop import Objective\n", "\n", "objectives = [Objective(name=\"himmelblau\", description=\"Himmeblau's function\", target=\"min\")]" ] @@ -140,8 +140,7 @@ }, "outputs": [], "source": [ - "from blop.bayesian import Agent\n", - "\n", + "from blop import Agent\n", "\n", "agent = Agent(\n", " dofs=dofs,\n", @@ -295,7 +294,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "Python 3.11.5 64-bit", "language": "python", "name": "python3" }, @@ -309,11 +308,11 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.0" + "version": "3.11.5" }, "vscode": { "interpreter": { - "hash": "857d19a2fd370900ed798add63a0e418d98c0c9c9169a1442a8e3b86b5805755" + "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e" } } }, diff --git a/docs/source/tutorials/hyperparameters.ipynb b/docs/source/tutorials/hyperparameters.ipynb index ba4df90..f5f0ae2 100644 --- a/docs/source/tutorials/hyperparameters.ipynb +++ b/docs/source/tutorials/hyperparameters.ipynb @@ -72,11 +72,11 @@ "\n", "%run -i $prepare_re_env.__file__ --db-type=temp\n", "\n", - "from blop.bayesian import DOF, Objective, Agent\n", + "from blop import DOF, Objective, Agent\n", "\n", "dofs = [\n", - " DOF(name=\"x1\", limits=(-6, 6)),\n", - " DOF(name=\"x2\", limits=(-6, 6)),\n", + " DOF(name=\"x1\", search_bounds=(-6, 6)),\n", + " DOF(name=\"x2\", search_bounds=(-6, 6)),\n", "]\n", "\n", "objectives = [\n", @@ -127,6 +127,16 @@ "RE(agent.learn(\"qei\", n=4, iterations=4))\n", "agent.plot_objectives()" ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b70eaf9b", + "metadata": {}, + "outputs": [], + "source": [ + "agent.best" + ] } ], "metadata": { @@ -145,7 +155,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.5" + "version": "3.10.0" }, "vscode": { "interpreter": { diff --git a/docs/source/tutorials/passive-dofs.ipynb b/docs/source/tutorials/passive-dofs.ipynb index 17e8f3d..42a0948 100644 --- a/docs/source/tutorials/passive-dofs.ipynb +++ b/docs/source/tutorials/passive-dofs.ipynb @@ -39,13 +39,14 @@ "outputs": [], "source": [ "from blop.utils import functions\n", - "from blop.bayesian import DOF, Agent, BrownianMotion, Objective\n", + "from blop import DOF, Agent, Objective\n", + "from blop.dofs import BrownianMotion\n", "\n", "\n", "dofs = [\n", - " DOF(name=\"x1\", limits=(-5.0, 5.0)),\n", - " DOF(name=\"x2\", limits=(-5.0, 5.0)),\n", - " DOF(name=\"x3\", limits=(-5.0, 5.0), active=False),\n", + " DOF(name=\"x1\", search_bounds=(-5.0, 5.0)),\n", + " DOF(name=\"x2\", search_bounds=(-5.0, 5.0)),\n", + " DOF(name=\"x3\", search_bounds=(-5.0, 5.0), active=False),\n", " DOF(device=BrownianMotion(name=\"brownian1\"), read_only=True),\n", " DOF(device=BrownianMotion(name=\"brownian2\"), read_only=True, active=False),\n", "]\n", diff --git a/docs/source/usage.rst b/docs/source/usage.rst index 82e9d10..fa19d1a 100644 --- a/docs/source/usage.rst +++ b/docs/source/usage.rst @@ -1,109 +1,9 @@ -===== Usage ===== -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 blop - - my_agent = blop.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. +.. toctree:: + :maxdepth: 2 -#. 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. + objectives.rst + dofs.rst + agent.rst diff --git a/requirements.txt b/requirements.txt index e48f1be..36e3314 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,6 +4,7 @@ databroker gpytorch h5py matplotlib +napari numpy ophyd ortools