From 1f451195f7f93c9f163c3ec1ae9fd47c6d987719 Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Mon, 8 Jan 2024 11:19:32 -0500
Subject: [PATCH 01/20] added napari viewer to agent

---
 bloptools/bayesian/agent.py    | 274 +++++++++++++++------------------
 bloptools/bayesian/kernels.py  |  39 +++--
 bloptools/tests/test_napari.py |   7 +
 requirements.txt               |   1 +
 4 files changed, 150 insertions(+), 171 deletions(-)
 create mode 100644 bloptools/tests/test_napari.py

diff --git a/bloptools/bayesian/agent.py b/bloptools/bayesian/agent.py
index 77d1d06..487b312 100644
--- a/bloptools/bayesian/agent.py
+++ b/bloptools/bayesian/agent.py
@@ -11,6 +11,7 @@
 import gpytorch
 import h5py
 import matplotlib as mpl
+import napari
 import numpy as np
 import pandas as pd
 import scipy as sp
@@ -119,6 +120,45 @@ def __init__(
 
         self.n_last_trained = 0
 
+    def view(self, item: str = "mean", cmap: str = "turbo", max_inputs: int = MAX_TEST_INPUTS):
+        """
+        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.test_inputs_grid(max_inputs=max_inputs)
+
+        self.viewer = napari.Viewer()
+
+        if item in ["mean", "error"]:
+            for obj in self.objectives:
+                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 tell(
         self,
         data: Optional[Mapping] = {},
@@ -142,7 +182,7 @@ def tell(
             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
+        _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.
@@ -168,147 +208,16 @@ def tell(
 
             cached_hypers = obj.model.state_dict() if hasattr(obj, "model") else None
 
-            obj.model = self.construct_model(obj)
+            obj.model = self._construct_model(obj)
 
             if len(obj.model.train_targets) >= 2:
                 t0 = ttime.monotonic()
-                self.train_model(obj.model, hypers=(None if train else cached_hypers))
+                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")
 
         # TODO: should this be per objective?
-        self.construct_classifier()
-
-    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.
-        """
-
-        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),
-        )
-
-        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
-
-        return model
-
-    def construct_classifier(self, skew_dims=None):
-        skew_dims = skew_dims if skew_dims is not None else self.latent_dim_tuples
-
-        dirichlet_likelihood = gpytorch.likelihoods.DirichletClassificationLikelihood(
-            self.all_objectives_valid.long(), learn_additional_noise=True
-        )
-
-        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,
-        )
-
-        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])
+        self._construct_classifier()
 
     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.
@@ -527,6 +436,67 @@ def learn(
                 metadata = new_table.to_dict(orient="list")
                 self.tell(x=x, y=y, metadata=metadata, append=append, train=train)
 
+    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.
+        """
+
+        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),
+        )
+
+        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
+
+        return model
+
+    def _construct_classifier(self, skew_dims=None):
+        skew_dims = skew_dims if skew_dims is not None else self.latent_dim_tuples
+
+        dirichlet_likelihood = gpytorch.likelihoods.DirichletClassificationLikelihood(
+            self.all_objectives_valid.long(), learn_additional_noise=True
+        )
+
+        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,
+        )
+
+        self._train_model(self.classifier)
+        self.constraint = GenericDeterministicModel(f=lambda x: self.classifier.probabilities(x)[..., -1])
+
     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`.
@@ -630,19 +600,23 @@ def test_inputs_grid(self, max_inputs=MAX_TEST_INPUTS):
         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)
+        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)
+            .double()
+        )
 
     def test_inputs(self, n=MAX_TEST_INPUTS):
         """Returns a (n, 1, n_active_dof) grid of test_inputs"""
@@ -705,7 +679,7 @@ def forget(self, index, train=True):
         Make the agent forget some index of the data table.
         """
         self.table.drop(index=index, inplace=True)
-        self._construct_models(train=train)
+        self.__construct_models(train=train)
 
     def forget_last_n(self, n, train=True):
         """
@@ -761,7 +735,7 @@ def load_hypers(filepath):
                     hypers[model_key][param_key] = torch.tensor(np.atleast_1d(param_value[()]))
         return hypers
 
-    def _train_models(self, **kwargs):
+    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:
diff --git a/bloptools/bayesian/kernels.py b/bloptools/bayesian/kernels.py
index dc9f642..78bffcf 100644
--- a/bloptools/bayesian/kernels.py
+++ b/bloptools/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",
+                name="lengthscales_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),
+                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(-np.pi, 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/bloptools/tests/test_napari.py b/bloptools/tests/test_napari.py
new file mode 100644
index 0000000..75b8a42
--- /dev/null
+++ b/bloptools/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=16))
+    agent.view(item)
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

From a182e4352c4bc78796e7cab6245fd6a2df8a84ad Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Mon, 8 Jan 2024 16:54:05 -0500
Subject: [PATCH 02/20] improve test coverage

---
 .cookiecutter-bloptools.yaml                |  8 ++--
 .coveragerc                                 |  2 +-
 .flake8                                     |  4 +-
 .gitattributes                              |  2 +-
 CONTRIBUTING.rst                            | 22 +++++------
 MANIFEST.in                                 |  4 +-
 README.rst                                  | 12 +++---
 bloptools/_version.py                       |  2 +-
 bloptools/bayesian/agent.py                 | 42 ++++++++++++++-------
 bloptools/experiments/nsls2/iss.py          |  2 +-
 bloptools/tests/conftest.py                 |  6 +--
 bloptools/tests/test_acq_funcs.py           |  4 +-
 bloptools/tests/test_agent.py               | 10 +++++
 bloptools/tests/test_napari.py              |  2 +-
 bloptools/tests/test_read_write.py          | 10 ++---
 docs/source/conf.py                         | 26 ++++++-------
 docs/source/index.rst                       |  2 +-
 docs/source/installation.rst                |  4 +-
 docs/source/tutorials/himmelblau.ipynb      | 12 +++---
 docs/source/tutorials/hyperparameters.ipynb |  6 +--
 docs/source/tutorials/passive-dofs.ipynb    |  6 +--
 docs/source/usage.rst                       |  4 +-
 docs/wip/constrained-himmelblau copy.ipynb  | 16 ++++----
 docs/wip/introduction.ipynb                 |  8 ++--
 docs/wip/latent-toroid-dimensions.ipynb     |  8 ++--
 docs/wip/multi-task-sirepo.ipynb            |  6 +--
 examples/benchmark.py                       |  2 +-
 setup.cfg                                   |  4 +-
 setup.py                                    |  8 ++--
 29 files changed, 133 insertions(+), 111 deletions(-)
 create mode 100644 bloptools/tests/test_agent.py

diff --git a/.cookiecutter-bloptools.yaml b/.cookiecutter-bloptools.yaml
index 9427db7..291f6b7 100644
--- a/.cookiecutter-bloptools.yaml
+++ b/.cookiecutter-bloptools.yaml
@@ -2,10 +2,10 @@ default_context:
   full_name: "Brookhaven National Laboratory"
   email: "mrakitin@bnl.gov"
   github_username: "NSLS-II"
-  project_name: "bloptools"
-  package_dist_name: "bloptools"
-  package_dir_name: "bloptools"
-  repo_name: "bloptools"
+  project_name: "blop"
+  package_dist_name: "blop"
+  package_dir_name: "blop"
+  repo_name: "blop"
   project_short_description: "Beamline Optimization Tools"
   minimum_supported_python_version: "3.8"
 replay_dir: "~/tmp/new-repo"
diff --git a/.coveragerc b/.coveragerc
index cd5e110..8adc4d5 100644
--- a/.coveragerc
+++ b/.coveragerc
@@ -1,6 +1,6 @@
 [run]
 source =
-    bloptools
+    blop
 [report]
 omit =
     */python?.?/*
diff --git a/.flake8 b/.flake8
index 83b1308..446f349 100644
--- a/.flake8
+++ b/.flake8
@@ -2,12 +2,12 @@
 exclude =
     .git,
     __pycache__,
-    bloptools/_version.py,
+    blop/_version.py,
     build,
     dist,
     docs/source/conf.py
     examples/*.py,
-    bloptools/tests/test_bayesian_shadow.py,
+    blop/tests/test_bayesian_shadow.py,
     versioneer.py,
 max-line-length = 125
 # Ignore some style 'errors' produced while formatting by 'black'
diff --git a/.gitattributes b/.gitattributes
index e6c2d50..46a2f76 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -1 +1 @@
-bloptools/_version.py export-subst
+blop/_version.py export-subst
diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst
index 4d6b7b3..5542d59 100644
--- a/CONTRIBUTING.rst
+++ b/CONTRIBUTING.rst
@@ -13,7 +13,7 @@ Types of Contributions
 Report Bugs
 ~~~~~~~~~~~
 
-Report bugs at https://github.com/NSLS-II/bloptools/issues.
+Report bugs at https://github.com/NSLS-II/blop/issues.
 
 If you are reporting a bug, please include:
 
@@ -35,14 +35,14 @@ is open to whoever wants to implement it.
 Write Documentation
 ~~~~~~~~~~~~~~~~~~~
 
-bloptools could always use more documentation, whether
-as part of the official bloptools docs, in docstrings,
+blop could always use more documentation, whether
+as part of the official blop docs, in docstrings,
 or even on the web in blog posts, articles, and such.
 
 Submit Feedback
 ~~~~~~~~~~~~~~~
 
-The best way to send feedback is to file an issue at https://github.com/NSLS-II/bloptools/issues.
+The best way to send feedback is to file an issue at https://github.com/NSLS-II/blop/issues.
 
 If you are proposing a feature:
 
@@ -54,17 +54,17 @@ If you are proposing a feature:
 Get Started!
 ------------
 
-Ready to contribute? Here's how to set up `bloptools` for local development.
+Ready to contribute? Here's how to set up `blop` for local development.
 
-1. Fork the `bloptools` repo on GitHub.
+1. Fork the `blop` repo on GitHub.
 2. Clone your fork locally::
 
-    $ git clone git@github.com:your_name_here/bloptools.git
+    $ git clone git@github.com:your_name_here/blop.git
 
 3. Install your local copy into a virtualenv. Assuming you have virtualenvwrapper installed, this is how you set up your fork for local development::
 
-    $ mkvirtualenv bloptools
-    $ cd bloptools/
+    $ mkvirtualenv blop
+    $ cd blop/
     $ python setup.py develop
 
 4. Create a branch for local development::
@@ -75,7 +75,7 @@ Ready to contribute? Here's how to set up `bloptools` for local development.
 
 5. When you're done making changes, check that your changes pass flake8 and the tests, including testing other Python versions with tox::
 
-    $ flake8 bloptools tests
+    $ flake8 blop tests
     $ python setup.py test
     $ tox
 
@@ -99,5 +99,5 @@ Before you submit a pull request, check that it meets these guidelines:
    your new functionality into a function with a docstring, and add the
    feature to the list in README.rst.
 3. The pull request should work for Python 2.7, 3.3, 3.4, 3.5 and for PyPy. Check
-   https://travis-ci.org/NSLS-II/bloptools/pull_requests
+   https://travis-ci.org/NSLS-II/blop/pull_requests
    and make sure that the tests pass for all supported Python versions.
diff --git a/MANIFEST.in b/MANIFEST.in
index e0ceedc..82ce24f 100644
--- a/MANIFEST.in
+++ b/MANIFEST.in
@@ -10,8 +10,8 @@ recursive-exclude * *.py[co]
 recursive-include docs *.rst conf.py Makefile make.bat
 
 include versioneer.py
-include bloptools/_version.py
-include bloptools/bayesian/acquisition/config.yml
+include blop/_version.py
+include blop/bayesian/acquisition/config.yml
 
 # If including data files in the package, add them like:
 # include path/to/data_file
diff --git a/README.rst b/README.rst
index 1468901..4a1244c 100644
--- a/README.rst
+++ b/README.rst
@@ -1,16 +1,16 @@
 =========
-bloptools
+blop
 =========
 
-.. image:: https://github.com/NSLS-II/bloptools/actions/workflows/testing.yml/badge.svg
-   :target: https://github.com/NSLS-II/bloptools/actions/workflows/testing.yml
+.. image:: https://github.com/NSLS-II/blop/actions/workflows/testing.yml/badge.svg
+   :target: https://github.com/NSLS-II/blop/actions/workflows/testing.yml
 
 
-.. image:: https://img.shields.io/pypi/v/bloptools.svg
-        :target: https://pypi.python.org/pypi/bloptools
+.. image:: https://img.shields.io/pypi/v/blop.svg
+        :target: https://pypi.python.org/pypi/blop
 
 
 Beamline Optimization Tools
 
 * Free software: 3-clause BSD license
-* Documentation: https://NSLS-II.github.io/bloptools.
+* Documentation: https://NSLS-II.github.io/blop.
diff --git a/bloptools/_version.py b/bloptools/_version.py
index 4958fb9..46d2441 100644
--- a/bloptools/_version.py
+++ b/bloptools/_version.py
@@ -42,7 +42,7 @@ def get_config():
     cfg.style = "pep440-post"
     cfg.tag_prefix = "v"
     cfg.parentdir_prefix = "None"
-    cfg.versionfile_source = "bloptools/_version.py"
+    cfg.versionfile_source = "blop/_version.py"
     cfg.verbose = False
     return cfg
 
diff --git a/bloptools/bayesian/agent.py b/bloptools/bayesian/agent.py
index 487b312..1e355b0 100644
--- a/bloptools/bayesian/agent.py
+++ b/bloptools/bayesian/agent.py
@@ -182,7 +182,7 @@ def tell(
             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
+        train: bool
             Whether to train the models on construction.
         hypers:
             A dict of hyperparameters for the model to assume a priori, instead of training.
@@ -674,20 +674,31 @@ def save_data(self, filepath="./self_data.h5"):
 
         self.table.to_hdf(filepath, key="table")
 
-    def forget(self, index, train=True):
-        """
-        Make the agent forget some index of the data table.
+    def forget(self, last=None, index=None, train=True):
         """
-        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)
+
+        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 sampler(self, n, d):
         """
@@ -735,7 +746,12 @@ def load_hypers(filepath):
                     hypers[model_key][param_key] = torch.tensor(np.atleast_1d(param_value[()]))
         return hypers
 
-    def __train_models(self, **kwargs):
+    def _construct_all_models(self):
+        """Construct a model for each objective."""
+        for obj in self.objectives:
+            obj.model = 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.objectives:
diff --git a/bloptools/experiments/nsls2/iss.py b/bloptools/experiments/nsls2/iss.py
index 21f3efe..6a76a98 100644
--- a/bloptools/experiments/nsls2/iss.py
+++ b/bloptools/experiments/nsls2/iss.py
@@ -2,7 +2,7 @@
 import bluesky.plans as bp
 import numpy as np
 
-from bloptools.tasks import Task
+from blop.tasks import Task
 
 
 def initialize():
diff --git a/bloptools/tests/conftest.py b/bloptools/tests/conftest.py
index c86ff9f..3531243 100644
--- a/bloptools/tests/conftest.py
+++ b/bloptools/tests/conftest.py
@@ -8,9 +8,9 @@
 from databroker import Broker
 from ophyd.utils import make_dir_tree
 
-from bloptools.bayesian import DOF, Agent, Objective
-from bloptools.bayesian.dofs import BrownianMotion
-from bloptools.utils import functions
+from blop.bayesian import DOF, Agent, Objective
+from blop.bayesian.dofs import BrownianMotion
+from blop.utils import functions
 
 
 @pytest.fixture(scope="function")
diff --git a/bloptools/tests/test_acq_funcs.py b/bloptools/tests/test_acq_funcs.py
index fbfd5ff..a041dc9 100644
--- a/bloptools/tests/test_acq_funcs.py
+++ b/bloptools/tests/test_acq_funcs.py
@@ -3,13 +3,13 @@
 
 @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))
 
 
diff --git a/bloptools/tests/test_agent.py b/bloptools/tests/test_agent.py
new file mode 100644
index 0000000..22446c3
--- /dev/null
+++ b/bloptools/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/bloptools/tests/test_napari.py b/bloptools/tests/test_napari.py
index 75b8a42..bf6ac95 100644
--- a/bloptools/tests/test_napari.py
+++ b/bloptools/tests/test_napari.py
@@ -3,5 +3,5 @@
 
 @pytest.mark.parametrize("item", ["mean", "error", "qei"])
 def test_napari_viewer(agent, RE, item):
-    RE(agent.learn("qr", n=16))
+    RE(agent.learn("qr", n=4))
     agent.view(item)
diff --git a/bloptools/tests/test_read_write.py b/bloptools/tests/test_read_write.py
index 9c225d9..d6cf26b 100644
--- a/bloptools/tests/test_read_write.py
+++ b/bloptools/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))
+    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"))
diff --git a/docs/source/conf.py b/docs/source/conf.py
index c6993bb..9d4b0b5 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 #
-# bloptools documentation build configuration file, created by
+# blop documentation build configuration file, created by
 # sphinx-quickstart on Thu Jun 28 12:35:56 2018.
 #
 # This file is execfile()d with the current directory set to its
@@ -68,7 +68,7 @@
 master_doc = "index"
 
 # General information about the project.
-project = "bloptools"
+project = "blop"
 copyright = "2020, Brookhaven National Laboratory"
 author = "Brookhaven National Laboratory"
 
@@ -76,12 +76,12 @@
 # |version| and |release|, also used in various other places throughout the
 # built documents.
 #
-import bloptools
+import blop
 
 # The short X.Y version.
-version = bloptools.__version__
+version = blop.__version__
 # The full version, including alpha/beta/rc tags.
-release = bloptools.__version__
+release = blop.__version__
 
 # The language for content autogenerated by Sphinx. Refer to documentation
 # for a list of supported languages.
@@ -127,7 +127,7 @@
 # -- Options for HTMLHelp output ------------------------------------------
 
 # Output file base name for HTML help builder.
-htmlhelp_basename = "bloptools"
+htmlhelp_basename = "blop"
 
 
 # -- Options for LaTeX output ---------------------------------------------
@@ -152,8 +152,8 @@
 latex_documents = [
     (
         master_doc,
-        "bloptools.tex",
-        "bloptools Documentation",
+        "blop.tex",
+        "blop Documentation",
         "Contributors",
         "manual",
     ),
@@ -167,8 +167,8 @@
 man_pages = [
     (
         master_doc,
-        "bloptools",
-        "bloptools Documentation",
+        "blop",
+        "blop Documentation",
         [author],
         1,
     )
@@ -183,10 +183,10 @@
 texinfo_documents = [
     (
         master_doc,
-        "bloptools",
-        "bloptools Documentation",
+        "blop",
+        "blop Documentation",
         author,
-        "bloptools",
+        "blop",
         "Beamline Optimization Tools",
         "Miscellaneous",
     ),
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 4a5132a..49f2703 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -3,7 +3,7 @@
    You can adapt this file completely to your liking, but it should at least
    contain the root `toctree` directive.
 
-bloptools Documentation
+blop Documentation
 =======================
 
 .. toctree::
diff --git a/docs/source/installation.rst b/docs/source/installation.rst
index ba107a6..895c9e4 100644
--- a/docs/source/installation.rst
+++ b/docs/source/installation.rst
@@ -11,13 +11,13 @@ To install the package using the ``pip`` package manager, run the following comm
 
 .. code:: bash
 
-   $ python3 -m pip install bloptools
+   $ python3 -m pip install blop
 
 To install the package using the ``conda`` package manager, run the following command:
 
 .. code:: bash
 
-   $ conda install -c conda-forge bloptools
+   $ conda install -c conda-forge blop
 
 If you'd like to use the Sirepo backend and ``sirepo-bluesky`` ophyd objects, please
 follow the `Sirepo/Sirepo-Bluesky installation & configuration instructions
diff --git a/docs/source/tutorials/himmelblau.ipynb b/docs/source/tutorials/himmelblau.ipynb
index 6fb98ac..b3e15d4 100644
--- a/docs/source/tutorials/himmelblau.ipynb
+++ b/docs/source/tutorials/himmelblau.ipynb
@@ -16,7 +16,7 @@
    "id": "c18ef717",
    "metadata": {},
    "source": [
-    "Let's use ``bloptools`` to minimize Himmelblau's function, which has four global minima:"
+    "Let's use ``blop`` to minimize Himmelblau's function, which has four global minima:"
    ]
   },
   {
@@ -26,7 +26,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from bloptools.utils import prepare_re_env\n",
+    "from blop.utils import prepare_re_env\n",
     "\n",
     "%run -i $prepare_re_env.__file__ --db-type=temp"
    ]
@@ -41,7 +41,7 @@
     "import numpy as np\n",
     "import matplotlib as mpl\n",
     "from matplotlib import pyplot as plt\n",
-    "from bloptools.utils import functions\n",
+    "from blop.utils import functions\n",
     "\n",
     "x1 = x2 = np.linspace(-6, 6, 256)\n",
     "X1, X2 = np.meshgrid(x1, x2)\n",
@@ -69,7 +69,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from bloptools.bayesian import DOF\n",
+    "from blop.bayesian import DOF\n",
     "\n",
     "dofs = [\n",
     "    DOF(name=\"x1\", limits=(-6, 6)),\n",
@@ -92,7 +92,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from bloptools.bayesian import Objective\n",
+    "from blop.bayesian import Objective\n",
     "\n",
     "objectives = [Objective(name=\"himmelblau\", description=\"Himmeblau's function\", target=\"min\")]"
    ]
@@ -140,7 +140,7 @@
    },
    "outputs": [],
    "source": [
-    "from bloptools.bayesian import Agent\n",
+    "from blop.bayesian import Agent\n",
     "\n",
     "\n",
     "agent = Agent(\n",
diff --git a/docs/source/tutorials/hyperparameters.ipynb b/docs/source/tutorials/hyperparameters.ipynb
index 5ec7645..ba4df90 100644
--- a/docs/source/tutorials/hyperparameters.ipynb
+++ b/docs/source/tutorials/hyperparameters.ipynb
@@ -21,7 +21,7 @@
     "import numpy as np\n",
     "import matplotlib as mpl\n",
     "from matplotlib import pyplot as plt\n",
-    "from bloptools.utils import functions\n",
+    "from blop.utils import functions\n",
     "\n",
     "x1 = x2 = np.linspace(-10, 10, 256)\n",
     "X1, X2 = np.meshgrid(x1, x2)\n",
@@ -68,11 +68,11 @@
    },
    "outputs": [],
    "source": [
-    "from bloptools.utils import prepare_re_env\n",
+    "from blop.utils import prepare_re_env\n",
     "\n",
     "%run -i $prepare_re_env.__file__ --db-type=temp\n",
     "\n",
-    "from bloptools.bayesian import DOF, Objective, Agent\n",
+    "from blop.bayesian import DOF, Objective, Agent\n",
     "\n",
     "dofs = [\n",
     "    DOF(name=\"x1\", limits=(-6, 6)),\n",
diff --git a/docs/source/tutorials/passive-dofs.ipynb b/docs/source/tutorials/passive-dofs.ipynb
index 71330a6..17e8f3d 100644
--- a/docs/source/tutorials/passive-dofs.ipynb
+++ b/docs/source/tutorials/passive-dofs.ipynb
@@ -26,7 +26,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from bloptools.utils import prepare_re_env\n",
+    "from blop.utils import prepare_re_env\n",
     "\n",
     "%run -i $prepare_re_env.__file__ --db-type=temp"
    ]
@@ -38,8 +38,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from bloptools.utils import functions\n",
-    "from bloptools.bayesian import DOF, Agent, BrownianMotion, Objective\n",
+    "from blop.utils import functions\n",
+    "from blop.bayesian import DOF, Agent, BrownianMotion, Objective\n",
     "\n",
     "\n",
     "dofs = [\n",
diff --git a/docs/source/usage.rst b/docs/source/usage.rst
index 905c0e6..82e9d10 100644
--- a/docs/source/usage.rst
+++ b/docs/source/usage.rst
@@ -89,9 +89,9 @@ Combining these with a databroker instance will construct an agent.
 
 .. code-block:: python
 
-    import bloptools
+    import blop
 
-    my_agent = bloptools.bayesian.Agent(
+    my_agent = blop.bayesian.Agent(
         dofs=my_dofs,
         dets=my_dets,
         tasks=my_tasks,
diff --git a/docs/wip/constrained-himmelblau copy.ipynb b/docs/wip/constrained-himmelblau copy.ipynb
index 10b3a2c..aa0e99f 100644
--- a/docs/wip/constrained-himmelblau copy.ipynb	
+++ b/docs/wip/constrained-himmelblau copy.ipynb	
@@ -16,7 +16,7 @@
    "id": "c18ef717",
    "metadata": {},
    "source": [
-    "Let's use ``bloptools`` to minimize Himmelblau's function, subject to the constraint that $x_1^2 + x_2^2 < 50$. Our function looks like this:"
+    "Let's use ``blop`` to minimize Himmelblau's function, subject to the constraint that $x_1^2 + x_2^2 < 50$. Our function looks like this:"
    ]
   },
   {
@@ -29,11 +29,11 @@
     "import numpy as np\n",
     "import matplotlib as mpl\n",
     "from matplotlib import pyplot as plt\n",
-    "from bloptools.utils import functions\n",
+    "from blop.utils import functions\n",
     "\n",
     "x1 = x2 = np.linspace(-8, 8, 256)\n",
     "X1, X2 = np.meshgrid(x1, x2)\n",
-    "from bloptools.tasks import Task\n",
+    "from blop.tasks import Task\n",
     "\n",
     "task = Task(name=\"himmelblau\", kind=\"min\")\n",
     "F = functions.constrained_himmelblau(X1, X2)\n",
@@ -59,7 +59,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from bloptools import devices\n",
+    "from blop import devices\n",
     "\n",
     "dofs = [\n",
     "    {\"device\": devices.DOF(name=\"x1\"), \"limits\": (-8, 8), \"kind\": \"active\"},\n",
@@ -130,10 +130,10 @@
    },
    "outputs": [],
    "source": [
-    "from bloptools.utils import prepare_re_env\n",
+    "from blop.utils import prepare_re_env\n",
     "\n",
     "%run -i $prepare_re_env.__file__ --db-type=temp\n",
-    "from bloptools.bayesian import Agent\n",
+    "from blop.bayesian import Agent\n",
     "\n",
     "agent = Agent(\n",
     "    dofs=dofs,\n",
@@ -150,9 +150,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import bloptools\n",
+    "import blop\n",
     "\n",
-    "bloptools.bayesian.acquisition.parse_acq_func_identifier(acq_func_identifier=\"quasi-random\")"
+    "blop.bayesian.acquisition.parse_acq_func_identifier(acq_func_identifier=\"quasi-random\")"
    ]
   },
   {
diff --git a/docs/wip/introduction.ipynb b/docs/wip/introduction.ipynb
index 07db65e..cd124f2 100644
--- a/docs/wip/introduction.ipynb
+++ b/docs/wip/introduction.ipynb
@@ -31,7 +31,7 @@
     "import matplotlib as mpl\n",
     "from matplotlib import pyplot as plt\n",
     "\n",
-    "from bloptools.utils import functions\n",
+    "from blop.utils import functions\n",
     "\n",
     "x = np.linspace(-5, 5, 256)\n",
     "\n",
@@ -55,7 +55,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from bloptools import devices\n",
+    "from blop import devices\n",
     "\n",
     "dofs = [\n",
     "    {\"device\": devices.DOF(name=\"x\"), \"limits\": (-5, 5), \"kind\": \"active\"},\n",
@@ -118,11 +118,11 @@
    },
    "outputs": [],
    "source": [
-    "from bloptools.utils import prepare_re_env\n",
+    "from blop.utils import prepare_re_env\n",
     "\n",
     "%run -i $prepare_re_env.__file__ --db-type=temp\n",
     "\n",
-    "from bloptools.bayesian import Agent\n",
+    "from blop.bayesian import Agent\n",
     "\n",
     "agent = Agent(\n",
     "    dofs=dofs,\n",
diff --git a/docs/wip/latent-toroid-dimensions.ipynb b/docs/wip/latent-toroid-dimensions.ipynb
index 5148279..7d1ab91 100644
--- a/docs/wip/latent-toroid-dimensions.ipynb
+++ b/docs/wip/latent-toroid-dimensions.ipynb
@@ -20,7 +20,7 @@
    },
    "outputs": [],
    "source": [
-    "from bloptools.utils import prepare_re_env\n",
+    "from blop.utils import prepare_re_env\n",
     "\n",
     "%run -i $prepare_re_env.__file__ --db-type=temp\n",
     "%run -i ../../../examples/prepare_tes_shadow.py"
@@ -35,8 +35,8 @@
    },
    "outputs": [],
    "source": [
-    "import bloptools\n",
-    "from bloptools.experiments.sirepo.tes import w8_digestion\n",
+    "import blop\n",
+    "from blop.experiments.sirepo.tes import w8_digestion\n",
     "\n",
     "dofs = [\n",
     "    {\"device\": toroid.x_rot, \"limits\": (-0.001, 0.001), \"kind\": \"active\"},\n",
@@ -45,7 +45,7 @@
     "\n",
     "tasks = [{\"key\": \"flux\", \"kind\": \"maximize\", \"transform\": \"log\"}]\n",
     "\n",
-    "agent = bloptools.bayesian.Agent(\n",
+    "agent = blop.bayesian.Agent(\n",
     "    dofs=dofs,\n",
     "    tasks=tasks,\n",
     "    dets=[w8],\n",
diff --git a/docs/wip/multi-task-sirepo.ipynb b/docs/wip/multi-task-sirepo.ipynb
index d1bd14c..62a7982 100644
--- a/docs/wip/multi-task-sirepo.ipynb
+++ b/docs/wip/multi-task-sirepo.ipynb
@@ -22,7 +22,7 @@
    },
    "outputs": [],
    "source": [
-    "from bloptools.utils import prepare_re_env\n",
+    "from blop.utils import prepare_re_env\n",
     "\n",
     "%run -i $prepare_re_env.__file__ --db-type=temp\n",
     "%run -i ../../../examples/prepare_tes_shadow.py"
@@ -37,8 +37,8 @@
    },
    "outputs": [],
    "source": [
-    "from bloptools.bayesian import Agent\n",
-    "from bloptools.experiments.sirepo.tes import w9_digestion\n",
+    "from blop.bayesian import Agent\n",
+    "from blop.experiments.sirepo.tes import w9_digestion\n",
     "\n",
     "dofs = [\n",
     "    {\"device\": kbv.x_rot, \"limits\": (-0.1, 0.1), \"kind\": \"active\"},\n",
diff --git a/examples/benchmark.py b/examples/benchmark.py
index 1c7db37..24db2e2 100644
--- a/examples/benchmark.py
+++ b/examples/benchmark.py
@@ -9,7 +9,7 @@
 
 import numpy as np
 
-from bloptools import gp
+from blop import gp
 
 bo = gp.BayesianOptimizer(
     init_scheme="quasi-random",
diff --git a/setup.cfg b/setup.cfg
index 90adc44..3b4e323 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,6 +1,6 @@
 [versioneer]
 VCS = git
 style = pep440-post
-versionfile_source = bloptools/_version.py
-versionfile_build = bloptools/_version.py
+versionfile_source = blop/_version.py
+versionfile_build = blop/_version.py
 tag_prefix = v
diff --git a/setup.py b/setup.py
index f4cd55a..26b5822 100644
--- a/setup.py
+++ b/setup.py
@@ -14,7 +14,7 @@
 )
 if sys.version_info < min_version:
     error = """
-bloptools does not support Python {0}.{1}.
+blop does not support Python {0}.{1}.
 Python {2}.{3} and above is required. Check your Python version like so:
 
 python3 --version
@@ -39,14 +39,14 @@
 
 
 setup(
-    name="bloptools",
+    name="blop",
     version=versioneer.get_version(),
     cmdclass=versioneer.get_cmdclass(),
     description="Beamline Optimization Tools",
     long_description=readme,
     author="Brookhaven National Laboratory",
     author_email="mrakitin@bnl.gov",
-    url="https://github.com/NSLS-II/bloptools",
+    url="https://github.com/NSLS-II/blop",
     python_requires=">={}".format(".".join(str(n) for n in min_version)),
     packages=find_packages(exclude=["docs", "tests"]),
     entry_points={
@@ -56,7 +56,7 @@
     },
     include_package_data=True,
     package_data={
-        "bloptools": [
+        "blop": [
             # When adding files here, remember to update MANIFEST.in as well,
             # or else they will not be included in the distribution on PyPI!
             # 'path/to/data_file',

From dd0fe80e808267c42c15d1fbdc337287c5f08b16 Mon Sep 17 00:00:00 2001
From: Thomas Morris <tmorris@bnl.gov>
Date: Fri, 12 Jan 2024 12:36:43 -0500
Subject: [PATCH 03/20] sirepo-bluesky patch

---
 blop/bayesian/agent.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/blop/bayesian/agent.py b/blop/bayesian/agent.py
index 77d1d06..a9d10bb 100644
--- a/blop/bayesian/agent.py
+++ b/blop/bayesian/agent.py
@@ -153,7 +153,7 @@ def tell(
                 raise ValueError("Must supply either x and y, or data.")
             data = {**x, **y, **metadata}
 
-        data = {k: np.atleast_1d(v) for k, v in data.items()}
+        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:

From 616b2df9e744a2de88806928cd693e5c2eae2cf8 Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Fri, 12 Jan 2024 16:26:10 -0500
Subject: [PATCH 04/20] better controls for dofs and objs

---
 blop/__init__.py                  |   3 +
 blop/bayesian/__init__.py         |   2 -
 blop/bayesian/agent.py            |  39 ++++-------
 blop/bayesian/plotting.py         |  26 ++++----
 blop/bayesian/transforms.py       |  22 ++++---
 blop/{bayesian => }/dofs.py       | 106 +++++++++++++++++++-----------
 blop/{bayesian => }/objectives.py |  65 ++++++++++++------
 blop/utils/__init__.py            |   4 ++
 blop/utils/functions.py           |  12 ++++
 9 files changed, 168 insertions(+), 111 deletions(-)
 rename blop/{bayesian => }/dofs.py (72%)
 rename blop/{bayesian => }/objectives.py (73%)

diff --git a/blop/__init__.py b/blop/__init__.py
index 80edaf0..35a88e7 100644
--- a/blop/__init__.py
+++ b/blop/__init__.py
@@ -1,4 +1,7 @@
+from . import utils  # noqa F401
 from ._version import get_versions
+from .dofs import DOF  # noqa F401
+from .objectives import Objective  # noqa F401
 
 __version__ = get_versions()["version"]
 del get_versions
diff --git a/blop/bayesian/__init__.py b/blop/bayesian/__init__.py
index e49a76f..732303d 100644
--- a/blop/bayesian/__init__.py
+++ b/blop/bayesian/__init__.py
@@ -1,3 +1 @@
 from .agent import *  # noqa F401
-from .dofs import *  # noqa F401
-from .objectives import *  # noqa F401
diff --git a/blop/bayesian/agent.py b/blop/bayesian/agent.py
index 1e355b0..4a39889 100644
--- a/blop/bayesian/agent.py
+++ b/blop/bayesian/agent.py
@@ -23,10 +23,10 @@
 from ophyd import Signal
 
 from .. import utils
+from ..dofs import DOF, DOFList
+from ..objectives import Objective, ObjectiveList
 from . import acquisition, models, plotting
 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
 
@@ -420,7 +420,7 @@ 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"
 
@@ -453,8 +453,8 @@ def _construct_model(self, obj, skew_dims=None):
 
         likelihood = gpytorch.likelihoods.GaussianLikelihood(
             noise_constraint=gpytorch.constraints.Interval(
-                torch.tensor(1e-4).square(),
-                torch.tensor(1 / obj.min_snr).square(),
+                torch.tensor(obj.min_noise),
+                torch.tensor(obj.max_noise),
             ),
             # noise_prior=gpytorch.priors.torch_priors.LogNormalPrior(loc=loc, scale=scale),
         )
@@ -557,20 +557,6 @@ def scalarizing_transform(self):
     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)
-            ]
-        )
-
     @property
     def scalarized_objectives(self):
         """Returns a (n_obs,) array of scalarized objectives"""
@@ -607,7 +593,9 @@ def test_inputs_grid(self, max_inputs=MAX_TEST_INPUTS):
                     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)
+                            for (lower_limit, upper_limit), n_side in zip(
+                                self.dofs.subset(active=True).search_bounds, n_sides
+                            )
                         ],
                         indexing="ij",
                     )
@@ -625,10 +613,9 @@ def test_inputs(self, n=MAX_TEST_INPUTS):
     @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)
 
-        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]
+        acq_func_lower_bounds = np.where(self.dofs.read_only, self.dofs.readback, self.dofs.search_lower_bounds)
+        acq_func_upper_bounds = np.where(self.dofs.read_only, self.dofs.readback, self.dofs.search_upper_bounds)
 
         return torch.tensor(np.vstack([acq_func_lower_bounds, acq_func_upper_bounds]), dtype=torch.double)
 
@@ -652,7 +639,7 @@ def input_transform(self):
 
     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)
+        torch_limits = torch.tensor(self.dofs.subset(active, read_only, tags).search_bounds.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(
@@ -792,7 +779,7 @@ def train_inputs(self, dof_name=None, **subset_kwargs):
         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
@@ -811,7 +798,7 @@ def train_targets(self, obj_name=None, **subset_kwargs):
         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
diff --git a/blop/bayesian/plotting.py b/blop/bayesian/plotting.py
index fe5fb9f..4dee48b 100644
--- a/blop/bayesian/plotting.py
+++ b/blop/bayesian/plotting.py
@@ -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)
@@ -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,8 @@ 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)
 
 
 def _plot_acqf_one_dof(agent, acq_funcs, lw=1e0, **kwargs):
@@ -205,7 +205,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"])
 
@@ -267,8 +267,8 @@ 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)
 
 
 def _plot_valid_one_dof(agent, size=16, lw=1e0):
@@ -282,7 +282,7 @@ def _plot_valid_one_dof(agent, size=16, lw=1e0):
 
     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):
@@ -327,8 +327,8 @@ def _plot_valid_many_dofs(agent, axes=[0, 1], shading="nearest", cmap=DEFAULT_CO
     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)
 
 
 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..0533249 100644
--- a/blop/bayesian/transforms.py
+++ b/blop/bayesian/transforms.py
@@ -7,6 +7,16 @@
 from torch import Tensor
 
 
+def targeting_transform(y, target):
+    if target == "min":
+        y = -y
+    elif not isinstance(target, tuple):
+        y = -(y - target).abs()
+    else:
+        y = -((y - 0.5 * (target[1] + target[0])).abs() - 0.5 * (target[1] - target[0])).clamp(min=0)
+    return y
+
+
 class TargetingPosteriorTransform(PosteriorTransform):
     r"""An affine posterior transform for scalarizing multi-output posteriors."""
 
@@ -25,20 +35,12 @@ def __init__(self, weights: Tensor, targets: Tensor) -> None:
 
     def sampled_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):
diff --git a/blop/bayesian/dofs.py b/blop/dofs.py
similarity index 72%
rename from blop/bayesian/dofs.py
rename to blop/dofs.py
index ad3f39f..9b15122 100644
--- a/blop/bayesian/dofs.py
+++ b/blop/dofs.py
@@ -2,16 +2,23 @@
 import uuid
 from collections.abc import Sequence
 from dataclasses import dataclass, field
-from typing import Tuple, Union
+from typing import Tuple
 
 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]
+DOF_FIELD_TYPES = {
+    "description": "object",
+    "readback": "float",
+    "search_bounds": "object",
+    "trust_bounds": "object",
+    "units": "object",
+    "active": "bool",
+    "read_only": "bool",
+    "log": "bool",
+    "tags": "object",
+}
 
 
 class ReadOnlyError(Exception):
@@ -42,16 +49,16 @@ class DOF:
         A longer name for the DOF.
     device: Signal, optional
         An ophyd device. If None, a dummy ophyd device is generated.
-    limits: tuple, optional
+    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.
+    search_bounds: 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.
+        will not explore outside the search_bounds. If the DOF is read-only, the agent will reject all
+        sampled data where the DOF is outside the 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.
     units: str
         The units of the DOF (e.g. mm or deg). This is only for plotting and general housekeeping.
     tags: list
@@ -64,7 +71,8 @@ class DOF:
     device: Signal = None
     description: str = None
     name: str = None
-    limits: Tuple[float, float] = (-10.0, 10.0)
+    search_bounds: Tuple[float, float] = (-10.0, 10.0)
+    trust_bounds: Tuple[float, float] = (-np.inf, np.inf)
     units: str = ""
     read_only: bool = False
     active: bool = True
@@ -94,12 +102,12 @@ def __post_init__(self):
         self.device.kind = "hinted"
 
     @property
-    def lower_limit(self):
-        return float(self.limits[0])
+    def search_lower_bound(self):
+        return float(self.search_bounds[0])
 
     @property
-    def upper_limit(self):
-        return float(self.limits[1])
+    def search_upper_bound(self):
+        return float(self.search_bounds[1])
 
     @property
     def readback(self):
@@ -107,14 +115,14 @@ def readback(self):
 
     @property
     def summary(self) -> pd.Series:
-        series = pd.Series(index=DOF_FIELDS)
+        series = pd.Series(index=list(DOF_FIELD_TYPES.keys()))
         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 ''}"
+        return f"{self.description}{f' [{self.units}]' if len(self.units) > 0 else ''}"
 
     @property
     def has_model(self):
@@ -140,22 +148,17 @@ def __len__(self):
     def __repr__(self):
         return self.summary.__repr__()
 
-    # def _repr_html_(self):
-    #     return self.summary._repr_html_()
+    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)
+        table = pd.DataFrame(columns=list(DOF_FIELD_TYPES.keys()), index=self.names)
 
-        # 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)
+        for attr, dtype in DOF_FIELD_TYPES.items():
+            for dof in self.dofs:
+                table.at[dof.name, attr] = getattr(dof, attr)
+            table[attr] = table[attr].astype(dtype)
 
         return table
 
@@ -172,33 +175,56 @@ 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])
+    def search_lower_bounds(self) -> np.array:
+        return np.array([dof.search_lower_bound for dof in self.dofs])
 
     @property
-    def upper_limits(self) -> np.array:
-        return np.array([dof.upper_limit for dof in self.dofs])
+    def search_upper_bounds(self) -> np.array:
+        return np.array([dof.search_upper_bound for dof in self.dofs])
 
     @property
-    def limits(self) -> np.array:
+    def search_bounds(self) -> np.array:
         """
         Returns a (n_dof, 2) array of bounds.
         """
-        return np.c_[self.lower_limits, self.upper_limits]
+        return np.c_[self.search_lower_bounds, self.search_upper_bounds]
+
+    @property
+    def trust_lower_bounds(self) -> np.array:
+        return np.array([dof.trust_lower_bound for dof in self.dofs])
+
+    @property
+    def trust_upper_bounds(self) -> np.array:
+        return np.array([dof.trust_upper_bound for dof in self.dofs])
+
+    @property
+    def trust_bounds(self) -> np.array:
+        """
+        Returns a (n_dof, 2) array of bounds.
+        """
+        return np.c_[self.trust_lower_bounds, self.trust_upper_bounds]
 
     @property
     def readback(self) -> np.array:
         return np.array([dof.readback for dof in self.dofs])
 
+    @property
+    def active(self):
+        return np.array([dof.active for dof in self.dofs])
+
+    @property
+    def read_only(self):
+        return np.array([dof.read_only 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]
+        return [_active == active if active is not None else True for _active in self.active]
+
+    def _dof_read_only_mask(self, read_only=None):
+        return [_read_only == read_only if read_only is not None else True for _read_only in self.read_only]
 
     def _dof_tags_mask(self, tags=[]):
         return [np.isin(dof["tags"], tags).any() if tags else True for dof in self.dofs]
diff --git a/blop/bayesian/objectives.py b/blop/objectives.py
similarity index 73%
rename from blop/bayesian/objectives.py
rename to blop/objectives.py
index 61c8e06..47aa213 100644
--- a/blop/bayesian/objectives.py
+++ b/blop/objectives.py
@@ -7,8 +7,25 @@
 
 numeric = Union[float, int]
 
-DEFAULT_MINIMUM_SNR = 1e2
-OBJ_FIELDS = ["description", "target", "active", "limits", "weight", "log", "n", "snr", "min_snr"]
+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",
+}
+
+OBJ_SUMMARY_STYLE = {"min_noise": "{:.2E}", "max_noise": "{:.2E}"}
 
 
 class DuplicateNameError(ValueError):
@@ -34,7 +51,7 @@ class Objective:
     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
+        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.
@@ -52,20 +69,21 @@ class Objective:
 
     name: str
     description: str = None
-    target: Union[float, str] = "max"
+    target: Union[Tuple[numeric, numeric], float, str] = "max"
     log: bool = False
     weight: numeric = 1.0
     active: bool = True
-    limits: Tuple[numeric, numeric] = None
-    min_snr: numeric = DEFAULT_MINIMUM_SNR
+    trust_bounds: Tuple[numeric, numeric] = None
+    min_noise: numeric = DEFAULT_MIN_NOISE_LEVEL
+    max_noise: numeric = DEFAULT_MAX_NOISE_LEVEL
     units: str = None
 
     def __post_init__(self):
-        if self.limits is None:
+        if self.trust_bounds is None:
             if self.log:
-                self.limits = (0, np.inf)
+                self.trust_bounds = (0, np.inf)
             else:
-                self.limits = (-np.inf, np.inf)
+                self.trust_bounds = (-np.inf, np.inf)
 
         if type(self.target) is str:
             if self.target not in ["min", "max"]:
@@ -78,13 +96,18 @@ def label(self):
     @property
     def summary(self):
         series = pd.Series()
-        for col in OBJ_FIELDS:
-            series[col] = getattr(self, col)
+        for attr, dtype in OBJ_FIELD_TYPES.items():
+            series[attr] = getattr(self, attr)
+            series[attr] = series[attr].astype(dtype)
+
         return series
 
     def __repr__(self):
         return self.summary.__repr__()
 
+    def __repr_html__(self):
+        return self.summary.__repr_html__()
+
     @property
     def noise(self):
         return self.model.likelihood.noise.item() if hasattr(self, "model") else None
@@ -115,21 +138,23 @@ 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)
+    def summary(self) -> pd.DataFrame:
+        table = pd.DataFrame(columns=list(OBJ_FIELD_TYPES.keys()), index=self.names)
 
-        # convert dtypes
-        for attr in ["log"]:
-            summary[attr] = summary[attr].astype(bool)
+        for attr, dtype in OBJ_FIELD_TYPES.items():
+            for obj in self.objectives:
+                table.at[obj.name, attr] = getattr(obj, attr)
 
-        return summary
+            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:
         """
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.

From aa7ebeee7a92785f555c0eafbdef281f6e63a1c3 Mon Sep 17 00:00:00 2001
From: Thomas Morris <tmorris@bnl.gov>
Date: Tue, 16 Jan 2024 16:09:40 -0500
Subject: [PATCH 05/20] specify latent dimensions by objective

---
 blop/__init__.py                      |  1 +
 blop/{bayesian => }/agent.py          | 99 ++++++++++++++++++++-------
 blop/bayesian/__init__.py             |  1 -
 blop/bayesian/acquisition/__init__.py | 17 ++---
 blop/bayesian/kernels.py              |  4 +-
 blop/bayesian/plotting.py             |  1 +
 blop/bayesian/transforms.py           |  9 +--
 blop/{bayesian => }/digestion.py      |  0
 blop/dofs.py                          | 35 ++++++----
 blop/objectives.py                    | 34 +++++----
 blop/{bayesian => }/plans.py          |  0
 11 files changed, 133 insertions(+), 68 deletions(-)
 rename blop/{bayesian => }/agent.py (92%)
 rename blop/{bayesian => }/digestion.py (100%)
 rename blop/{bayesian => }/plans.py (100%)

diff --git a/blop/__init__.py b/blop/__init__.py
index 35a88e7..cc149b7 100644
--- a/blop/__init__.py
+++ b/blop/__init__.py
@@ -1,5 +1,6 @@
 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
 
diff --git a/blop/bayesian/agent.py b/blop/agent.py
similarity index 92%
rename from blop/bayesian/agent.py
rename to blop/agent.py
index d053b47..22fff5d 100644
--- a/blop/bayesian/agent.py
+++ b/blop/agent.py
@@ -22,13 +22,13 @@
 from databroker import Broker
 from ophyd import Signal
 
-from .. import utils
-from ..dofs import DOF, DOFList
-from ..objectives import Objective, ObjectiveList
-from . import acquisition, models, plotting
+from . import utils
+from .dofs import DOF, DOFList
+from .objectives import Objective, ObjectiveList
+from .bayesian import acquisition, models, plotting
 from .digestion import default_digestion_function
 from .plans import default_acquisition_plan
-from .transforms import TargetingPosteriorTransform
+from .bayesian.transforms import TargetingPosteriorTransform
 
 warnings.filterwarnings("ignore", category=botorch.exceptions.warnings.InputDataWarning)
 
@@ -37,6 +37,22 @@
 MAX_TEST_INPUTS = 2**11
 
 
+
+def _validate_dofs_and_objs(dofs: DOFList, objs: ObjectiveList):
+
+    if len(dofs) == 0:
+        raise ValueError(f"You must supply at least one DOF.")
+
+    if len(objs) == 0:
+        raise ValueError(f"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 not dof_name in dofs.names:
+                    raise ValueError(f"DOF name '{dof_name}' in latent group for objective '{obj.name}' does not exist.")
+
+
 class Agent:
     def __init__(
         self,
@@ -99,6 +115,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
@@ -120,6 +139,19 @@ def __init__(
 
         self.n_last_trained = 0
 
+    def __iter__(self):
+        for index in range(len(self)):
+            yield self.dofs[index]
+
+    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)
+        
+        raise AttributeError(f"DOFList object has no attribute named '{attr}'.")
+
+
     def view(self, item: str = "mean", cmap: str = "turbo", max_inputs: int = MAX_TEST_INPUTS):
         """
         Use napari to see a high-dimensional array.
@@ -391,7 +423,8 @@ def learn(
         upsample: int = 1,
         train: bool = True,
         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.
 
@@ -427,7 +460,7 @@ def learn(
         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"]
 
@@ -449,7 +482,7 @@ def _construct_model(self, obj, skew_dims=None):
         Construct an untrained model for an objective.
         """
 
-        skew_dims = skew_dims if skew_dims is not None else self.latent_dim_tuples
+        skew_dims = skew_dims if skew_dims is not None else self.latent_dim_tuples(obj.name)
 
         likelihood = gpytorch.likelihoods.GaussianLikelihood(
             noise_constraint=gpytorch.constraints.Interval(
@@ -464,11 +497,11 @@ def _construct_model(self, obj, skew_dims=None):
         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))
+        trusted = ~(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],
+            train_inputs=train_inputs[trusted],
+            train_targets=train_targets[trusted],
             likelihood=likelihood,
             skew_dims=skew_dims,
             input_transform=self.input_transform,
@@ -480,14 +513,17 @@ def _construct_model(self, obj, skew_dims=None):
         return model
 
     def _construct_classifier(self, skew_dims=None):
-        skew_dims = skew_dims if skew_dims is not None else self.latent_dim_tuples
+        skew_dims = [tuple([i]) for i in range(len(self.dofs))]
+
+        train_inputs = self.train_inputs(active=True)
+        trusted = ~torch.isnan(train_inputs).any(axis=1)
 
         dirichlet_likelihood = gpytorch.likelihoods.DirichletClassificationLikelihood(
-            self.all_objectives_valid.long(), learn_additional_noise=True
+            self.all_objectives_valid.long()[trusted], learn_additional_noise=True
         )
 
         self.classifier = models.LatentDirichletClassifier(
-            train_inputs=self.train_inputs(active=True),
+            train_inputs=train_inputs[trusted],
             train_targets=dirichlet_likelihood.transformed_targets.transpose(-1, -2).double(),
             skew_dims=skew_dims,
             likelihood=dirichlet_likelihood,
@@ -618,16 +654,27 @@ def acquisition_function_bounds(self):
         acq_func_upper_bounds = np.where(self.dofs.read_only, self.dofs.readback, self.dofs.search_upper_bounds)
 
         return torch.tensor(np.vstack([acq_func_lower_bounds, acq_func_upper_bounds]), dtype=torch.double)
+    
 
-    @property
-    def latent_dim_tuples(self):
+    def latent_dim_tuples(self, obj_index=None):
         """
-        Returns a list of tuples, where each tuple represent a group of dimension to find a latent representation of.
+        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.dofs.subset(active=True):
+            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
@@ -769,13 +816,13 @@ 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
@@ -788,13 +835,13 @@ def train_inputs(self, dof_name=None, **subset_kwargs):
 
         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
diff --git a/blop/bayesian/__init__.py b/blop/bayesian/__init__.py
index 732303d..e69de29 100644
--- a/blop/bayesian/__init__.py
+++ b/blop/bayesian/__init__.py
@@ -1 +0,0 @@
-from .agent import *  # noqa F401
diff --git a/blop/bayesian/acquisition/__init__.py b/blop/bayesian/acquisition/__init__.py
index a66afbb..f578bd5 100644
--- a/blop/bayesian/acquisition/__init__.py
+++ b/blop/bayesian/acquisition/__init__.py
@@ -18,16 +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):
     """Generates an acquisition function from a supplied identifier. A list of acquisition functions and
@@ -35,6 +29,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/kernels.py b/blop/bayesian/kernels.py
index 78bffcf..abfcdc4 100644
--- a/blop/bayesian/kernels.py
+++ b/blop/bayesian/kernels.py
@@ -74,13 +74,13 @@ def __init__(
         if priors:
             self.register_prior(
                 name="lengthscales_prior",
-                prior=gpytorch.priors.GammaPrior(concentration=2, rate=1),
+                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(-np.pi, np.pi)
+            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)
diff --git a/blop/bayesian/plotting.py b/blop/bayesian/plotting.py
index 4dee48b..1bbbad3 100644
--- a/blop/bayesian/plotting.py
+++ b/blop/bayesian/plotting.py
@@ -87,6 +87,7 @@ def _plot_objs_many_dofs(agent, axes=(0, 1), shading="nearest", cmap=DEFAULT_COL
     test_y = test_inputs[..., 0, axes[1]].detach().squeeze().numpy()
 
     for obj_index, obj in enumerate(agent.objectives):
+
         targets = agent.train_targets(obj.name).squeeze(-1).numpy()
 
         obj_vmin, obj_vmax = np.nanpercentile(targets, q=[1, 99])
diff --git a/blop/bayesian/transforms.py b/blop/bayesian/transforms.py
index 0533249..09b718f 100644
--- a/blop/bayesian/transforms.py
+++ b/blop/bayesian/transforms.py
@@ -8,13 +8,14 @@
 
 
 def targeting_transform(y, target):
+    if target == "max":
+        return y
     if target == "min":
-        y = -y
+        return -y
     elif not isinstance(target, tuple):
-        y = -(y - target).abs()
+        return -(y - target).abs()
     else:
-        y = -((y - 0.5 * (target[1] + target[0])).abs() - 0.5 * (target[1] - target[0])).clamp(min=0)
-    return y
+        return -((y - 0.5 * (target[1] + target[0])).abs() - 0.5 * (target[1] - target[0])).clamp(min=0)
 
 
 class TargetingPosteriorTransform(PosteriorTransform):
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
index 9b15122..751ad54 100644
--- a/blop/dofs.py
+++ b/blop/dofs.py
@@ -63,9 +63,6 @@ class DOF:
         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
@@ -78,7 +75,6 @@ class DOF:
     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):
@@ -95,9 +91,6 @@ def __post_init__(self):
             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"
 
@@ -108,6 +101,14 @@ def search_lower_bound(self):
     @property
     def search_upper_bound(self):
         return float(self.search_bounds[1])
+    
+    @property
+    def trust_lower_bound(self):
+        return float(self.trust_bounds[0])
+
+    @property
+    def trust_upper_bound(self):
+        return float(self.trust_bounds[1])
 
     @property
     def readback(self):
@@ -134,13 +135,21 @@ 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)]
+    def __getattr__(self, attr):
+        if attr in self.names:
+            return self.__getitem__(attr)
+
+        raise AttributeError(f"DOFList object has no attribute named '{attr}'.")
+
+    def __getitem__(self, index):
+        if type(index) is int:
+            return self.dofs[index]
+        elif type(index) is str:
+            if index not in self.names:
+                raise ValueError(f"DOFList has no DOF named {index}.")
+            return self.dofs[self.names.index(index)]
         else:
-            raise ValueError(f"Invalid index {i}. A DOFList must be indexed by either an integer or a string.")
+            raise ValueError(f"Invalid index {index}. A DOFList must be indexed by either an integer or a string.")
 
     def __len__(self):
         return len(self.dofs)
diff --git a/blop/objectives.py b/blop/objectives.py
index 47aa213..762d1c3 100644
--- a/blop/objectives.py
+++ b/blop/objectives.py
@@ -1,12 +1,10 @@
 from collections.abc import Sequence
-from dataclasses import dataclass
-from typing import Tuple, Union
+from dataclasses import dataclass, field
+from typing import Tuple, List, Union
 
 import numpy as np
 import pandas as pd
 
-numeric = Union[float, int]
-
 DEFAULT_MIN_NOISE_LEVEL = 1e-6
 DEFAULT_MAX_NOISE_LEVEL = 1e0
 
@@ -23,6 +21,7 @@
     "max_noise": "float",
     "noise": "float",
     "n": "int",
+    "latent_groups": "object",
 }
 
 OBJ_SUMMARY_STYLE = {"min_noise": "{:.2E}", "max_noise": "{:.2E}"}
@@ -61,22 +60,28 @@ class Objective:
         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.
+    min_noise: float
+        The minimum noise level of the fitted model.
+    max_noise: float
+        The minimum noise level of the fitted model.
     units: str
         A label representing the units of the objective.
+    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.
     """
 
     name: str
     description: str = None
-    target: Union[Tuple[numeric, numeric], float, str] = "max"
+    target: Union[Tuple[float, float], float, str] = "max"
     log: bool = False
-    weight: numeric = 1.0
+    weight: float = 1.0
     active: bool = True
-    trust_bounds: Tuple[numeric, numeric] = None
-    min_noise: numeric = DEFAULT_MIN_NOISE_LEVEL
-    max_noise: numeric = DEFAULT_MAX_NOISE_LEVEL
+    trust_bounds: Tuple[float, float] = 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 self.trust_bounds is None:
@@ -98,7 +103,6 @@ def summary(self):
         series = pd.Series()
         for attr, dtype in OBJ_FIELD_TYPES.items():
             series[attr] = getattr(self, attr)
-            series[attr] = series[attr].astype(dtype)
 
         return series
 
@@ -126,6 +130,12 @@ def __init__(self, objectives: list = []):
         _validate_objectives(objectives)
         self.objectives = objectives
 
+    def __getattr__(self, attr):
+        if attr in self.names:
+            return self.__getitem__(attr)
+        else:
+            raise AttributeError(f'No attribute named {attr}')
+
     def __getitem__(self, i):
         if type(i) is int:
             return self.objectives[i]
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

From ca18732a2ca0bc4494e42c696e6e3a835495b8ac Mon Sep 17 00:00:00 2001
From: Thomas Morris <tmorris@bnl.gov>
Date: Tue, 16 Jan 2024 16:18:54 -0500
Subject: [PATCH 06/20] fixed notebooks

---
 blop/__init__.py                            |  2 +-
 blop/agent.py                               | 27 +++++++++------------
 blop/bayesian/acquisition/__init__.py       |  3 ++-
 blop/bayesian/kernels.py                    |  2 +-
 blop/bayesian/plotting.py                   |  1 -
 blop/dofs.py                                |  2 +-
 blop/objectives.py                          |  4 +--
 docs/source/tutorials/himmelblau.ipynb      | 11 ++++-----
 docs/source/tutorials/hyperparameters.ipynb | 18 +++++++++++---
 docs/source/tutorials/passive-dofs.ipynb    | 21 ++++++++++++----
 10 files changed, 53 insertions(+), 38 deletions(-)

diff --git a/blop/__init__.py b/blop/__init__.py
index cc149b7..c273c18 100644
--- a/blop/__init__.py
+++ b/blop/__init__.py
@@ -1,6 +1,6 @@
 from . import utils  # noqa F401
 from ._version import get_versions
-from .agent import Agent # noqa F401
+from .agent import Agent  # noqa F401
 from .dofs import DOF  # noqa F401
 from .objectives import Objective  # noqa F401
 
diff --git a/blop/agent.py b/blop/agent.py
index 22fff5d..8dc3e0f 100644
--- a/blop/agent.py
+++ b/blop/agent.py
@@ -23,12 +23,12 @@
 from ophyd import Signal
 
 from . import utils
-from .dofs import DOF, DOFList
-from .objectives import Objective, ObjectiveList
 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 .bayesian.transforms import TargetingPosteriorTransform
 
 warnings.filterwarnings("ignore", category=botorch.exceptions.warnings.InputDataWarning)
 
@@ -37,19 +37,17 @@
 MAX_TEST_INPUTS = 2**11
 
 
-
 def _validate_dofs_and_objs(dofs: DOFList, objs: ObjectiveList):
-
     if len(dofs) == 0:
-        raise ValueError(f"You must supply at least one DOF.")
+        raise ValueError("You must supply at least one DOF.")
 
     if len(objs) == 0:
-        raise ValueError(f"You must supply at least one objective.")
+        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 not dof_name in dofs.names:
+                if dof_name not in dofs.names:
                     raise ValueError(f"DOF name '{dof_name}' in latent group for objective '{obj.name}' does not exist.")
 
 
@@ -144,13 +142,11 @@ def __iter__(self):
             yield self.dofs[index]
 
     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)
-        
-        raise AttributeError(f"DOFList object has no attribute named '{attr}'.")
 
+        raise AttributeError(f"DOFList object has no attribute named '{attr}'.")
 
     def view(self, item: str = "mean", cmap: str = "turbo", max_inputs: int = MAX_TEST_INPUTS):
         """
@@ -513,7 +509,7 @@ def _construct_model(self, obj, skew_dims=None):
         return model
 
     def _construct_classifier(self, skew_dims=None):
-        skew_dims = [tuple([i]) for i in range(len(self.dofs))]
+        skew_dims = [tuple([i]) for i in range(len(self.dofs.subset(active=True)))]
 
         train_inputs = self.train_inputs(active=True)
         trusted = ~torch.isnan(train_inputs).any(axis=1)
@@ -654,7 +650,6 @@ def acquisition_function_bounds(self):
         acq_func_upper_bounds = np.where(self.dofs.read_only, self.dofs.readback, self.dofs.search_upper_bounds)
 
         return torch.tensor(np.vstack([acq_func_lower_bounds, acq_func_upper_bounds]), dtype=torch.double)
-    
 
     def latent_dim_tuples(self, obj_index=None):
         """
@@ -663,8 +658,8 @@ def latent_dim_tuples(self, obj_index=None):
         """
 
         if obj_index is None:
-            return {obj.name:self.latent_dim_tuples(obj_index=obj.name) for obj in self.objectives}
-        
+            return {obj.name: self.latent_dim_tuples(obj_index=obj.name) for obj in self.objectives}
+
         obj = self.objectives[obj_index]
 
         latent_group_index = {}
@@ -673,7 +668,7 @@ def latent_dim_tuples(self, obj_index=None):
             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))]
 
diff --git a/blop/bayesian/acquisition/__init__.py b/blop/bayesian/acquisition/__init__.py
index f578bd5..24bdacf 100644
--- a/blop/bayesian/acquisition/__init__.py
+++ b/blop/bayesian/acquisition/__init__.py
@@ -23,6 +23,7 @@ def parse_acq_func_identifier(identifier):
             return acq_func_name
     return None
 
+
 def get_acquisition_function(agent, identifier="qei", return_metadata=True, verbose=False, **acq_func_kwargs):
     """Generates an acquisition function from a supplied identifier. A list of acquisition functions and
     their identifiers can be found at `agent.all_acq_funcs`.
@@ -31,7 +32,7 @@ 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/kernels.py b/blop/bayesian/kernels.py
index abfcdc4..ecc7d5c 100644
--- a/blop/bayesian/kernels.py
+++ b/blop/bayesian/kernels.py
@@ -80,7 +80,7 @@ def __init__(
             )
 
         if self.n_skew_entries > 0:
-            skew_entries_constraint = gpytorch.constraints.Interval(-2*np.pi, 2*np.pi)
+            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)
diff --git a/blop/bayesian/plotting.py b/blop/bayesian/plotting.py
index 1bbbad3..4dee48b 100644
--- a/blop/bayesian/plotting.py
+++ b/blop/bayesian/plotting.py
@@ -87,7 +87,6 @@ def _plot_objs_many_dofs(agent, axes=(0, 1), shading="nearest", cmap=DEFAULT_COL
     test_y = test_inputs[..., 0, axes[1]].detach().squeeze().numpy()
 
     for obj_index, obj in enumerate(agent.objectives):
-
         targets = agent.train_targets(obj.name).squeeze(-1).numpy()
 
         obj_vmin, obj_vmax = np.nanpercentile(targets, q=[1, 99])
diff --git a/blop/dofs.py b/blop/dofs.py
index 751ad54..d1969da 100644
--- a/blop/dofs.py
+++ b/blop/dofs.py
@@ -101,7 +101,7 @@ def search_lower_bound(self):
     @property
     def search_upper_bound(self):
         return float(self.search_bounds[1])
-    
+
     @property
     def trust_lower_bound(self):
         return float(self.trust_bounds[0])
diff --git a/blop/objectives.py b/blop/objectives.py
index 762d1c3..b4d0804 100644
--- a/blop/objectives.py
+++ b/blop/objectives.py
@@ -1,6 +1,6 @@
 from collections.abc import Sequence
 from dataclasses import dataclass, field
-from typing import Tuple, List, Union
+from typing import List, Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -134,7 +134,7 @@ def __getattr__(self, attr):
         if attr in self.names:
             return self.__getitem__(attr)
         else:
-            raise AttributeError(f'No attribute named {attr}')
+            raise AttributeError(f"No attribute named {attr}")
 
     def __getitem__(self, i):
         if type(i) is int:
diff --git a/docs/source/tutorials/himmelblau.ipynb b/docs/source/tutorials/himmelblau.ipynb
index b3e15d4..c871dc6 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",
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..bd96f7c 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",
@@ -75,6 +76,16 @@
    "source": [
     "agent.plot_objectives()"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "312610b1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "agent.latent_dim_tuples()"
+   ]
   }
  ],
  "metadata": {
@@ -93,7 +104,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.5"
+   "version": "3.10.0"
   },
   "vscode": {
    "interpreter": {

From a7b6439d40101172e5e327917fbc982c67915046 Mon Sep 17 00:00:00 2001
From: Thomas Morris <tmorris@bnl.gov>
Date: Thu, 18 Jan 2024 13:00:30 -0500
Subject: [PATCH 07/20] add docstrings

---
 blop/dofs.py       | 62 ++++++++++++++++++++++++++++++----------------
 blop/objectives.py | 24 +++++++++---------
 2 files changed, 52 insertions(+), 34 deletions(-)

diff --git a/blop/dofs.py b/blop/dofs.py
index d1969da..fa31130 100644
--- a/blop/dofs.py
+++ b/blop/dofs.py
@@ -18,6 +18,7 @@
     "read_only": "bool",
     "log": "bool",
     "tags": "object",
+    "device": "object",
 }
 
 
@@ -44,40 +45,51 @@ class DOF:
     Parameters
     ----------
     name: str
-        The name of the DOF. This is used as a key.
-    description: 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.
-    device: Signal, optional
-        An ophyd device. If None, a dummy ophyd device is generated.
-    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.
-    search_bounds: 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 search_bounds. If the DOF is read-only, the agent will reject all
-        sampled data where the DOF is outside the search_bounds.
+    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.
     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.
-    units: str
-        The units of the DOF (e.g. mm or deg). This is only for plotting and general housekeeping.
+    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.
     """
 
-    device: Signal = None
-    description: str = None
-    name: str = None
-    search_bounds: Tuple[float, float] = (-10.0, 10.0)
-    trust_bounds: Tuple[float, float] = (-np.inf, np.inf)
+    name: str
+    description: str = ""
     units: str = ""
+    search_bounds: Tuple[float, float]
+    trust_bounds: Tuple[float, float] = None
     read_only: bool = False
     active: bool = True
-    tags: list = field(default_factory=list)
     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.trust_bounds is None:
+            if self.log:
+                self.trust_bounds = (0, np.inf)
+            else:
+                self.trust_bounds = (-np.inf, np.inf)
+
+        self.search_bounds = tuple(self.search_bounds)
+        self.trust_bounds = tuple(self.trust_bounds)
+
         self.uuid = str(uuid.uuid4())
 
         if self.name is None:
@@ -118,7 +130,11 @@ def readback(self):
     def summary(self) -> pd.Series:
         series = pd.Series(index=list(DOF_FIELD_TYPES.keys()))
         for attr in series.index:
-            series[attr] = getattr(self, attr)
+            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
@@ -164,9 +180,11 @@ def __repr_html__(self):
     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():
-            for dof in self.dofs:
-                table.at[dof.name, attr] = getattr(dof, attr)
             table[attr] = table[attr].astype(dtype)
 
         return table
diff --git a/blop/objectives.py b/blop/objectives.py
index b4d0804..205cc8f 100644
--- a/blop/objectives.py
+++ b/blop/objectives.py
@@ -46,33 +46,33 @@ class Objective:
     Parameters
     ----------
     name: str
-        The name of the objective. This is used as a key.
+        The name of the objective. This is used as a key to index observed data.
     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.
+    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, to make it more Gaussian.
+        Whether to apply a log to the objective, i.e. to make the process more stationary.
     weight: float
-        The relative importance of this objective, used when scalarizing in multi-objective optimization.
+        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 obejctive. Outside of this, data points will be rejected.
+        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 minimum noise level of the fitted model.
+        The maximum noise level of the fitted model.
     units: str
         A label representing the units of the objective.
-    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.
+    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 = None
+    description: str = ""
     target: Union[Tuple[float, float], float, str] = "max"
     log: bool = False
     weight: float = 1.0

From 3c21a4b96c88d5ede97336aedea292729f357bac Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Thu, 18 Jan 2024 13:50:41 -0500
Subject: [PATCH 08/20] fixed read-only DOFs

---
 blop/agent.py                 |  6 ++---
 blop/dofs.py                  | 46 ++++++++++++++++++++---------------
 blop/objectives.py            | 32 ++++++++++++++++--------
 blop/tests/conftest.py        | 22 ++++++++---------
 blop/tests/test_read_write.py |  4 +--
 5 files changed, 64 insertions(+), 46 deletions(-)

diff --git a/blop/agent.py b/blop/agent.py
index 8dc3e0f..1a21b9f 100644
--- a/blop/agent.py
+++ b/blop/agent.py
@@ -148,7 +148,7 @@ def __getattr__(self, attr):
 
         raise AttributeError(f"DOFList object has no attribute named '{attr}'.")
 
-    def view(self, item: str = "mean", cmap: str = "turbo", max_inputs: int = MAX_TEST_INPUTS):
+    def view(self, item: str = "mean", cmap: str = "turbo", max_inputs: int = 2**16):
         """
         Use napari to see a high-dimensional array.
 
@@ -821,7 +821,7 @@ def train_inputs(self, index=None, **subset_kwargs):
         inputs = self.table.loc[:, dof.name].values.copy()
 
         # check that inputs values are inside acceptable values
-        valid = (inputs >= dof.trust_bounds[0]) & (inputs <= dof.trust_bounds[1])
+        valid = (inputs >= dof.trust_lower_bound) & (inputs <= dof.trust_upper_bound)
         inputs = np.where(valid, inputs, np.nan)
 
         # transform if needed
@@ -840,7 +840,7 @@ def train_targets(self, index=None, **subset_kwargs):
         targets = self.table.loc[:, obj.name].values.copy()
 
         # check that targets values are inside acceptable values
-        valid = (targets >= obj.trust_bounds[0]) & (targets <= obj.trust_bounds[1])
+        valid = (targets >= obj.trust_lower_bound) & (targets <= obj.trust_upper_bound)
         targets = np.where(valid, targets, np.nan)
 
         # transform if needed
diff --git a/blop/dofs.py b/blop/dofs.py
index fa31130..cf4b43c 100644
--- a/blop/dofs.py
+++ b/blop/dofs.py
@@ -18,7 +18,6 @@
     "read_only": "bool",
     "log": "bool",
     "tags": "object",
-    "device": "object",
 }
 
 
@@ -68,11 +67,11 @@ class DOF:
         An ophyd device. If not supplied, a dummy ophyd device will be generated.
     """
 
-    name: str
+    name: str = None
     description: str = ""
-    units: str = ""
-    search_bounds: Tuple[float, float]
+    search_bounds: Tuple[float, float] = None
     trust_bounds: Tuple[float, float] = None
+    units: str = ""
     read_only: bool = False
     active: bool = True
     log: bool = False
@@ -81,14 +80,13 @@ class DOF:
 
     # Some post-processing. This is specific to dataclasses
     def __post_init__(self):
-        if self.trust_bounds is None:
-            if self.log:
-                self.trust_bounds = (0, np.inf)
-            else:
-                self.trust_bounds = (-np.inf, np.inf)
-
-        self.search_bounds = tuple(self.search_bounds)
-        self.trust_bounds = tuple(self.trust_bounds)
+        if self.trust_bounds is not None:
+            self.trust_bounds = tuple(self.trust_bounds)
+        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)
 
         self.uuid = str(uuid.uuid4())
 
@@ -101,26 +99,34 @@ def __post_init__(self):
         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!")
+                raise ValueError("You must specify read_only=True for a read-only device.")
+
+        if self.log:
+            if not self.search_lower_bound > 0:
+                raise ValueError("Search bounds must be positive if log=True.")
 
         # all dof degrees of freedom are hinted
         self.device.kind = "hinted"
 
     @property
     def search_lower_bound(self):
-        return float(self.search_bounds[0])
+        if self.read_only:
+            raise ValueError("Read-only DOFs do not have search bounds.")
+        return float(self.summary.search_bounds[0])
 
     @property
     def search_upper_bound(self):
-        return float(self.search_bounds[1])
+        if self.read_only:
+            raise ValueError("Read-only DOFs do not have search bounds.")
+        return float(self.summary.search_bounds[1])
 
     @property
     def trust_lower_bound(self):
-        return float(self.trust_bounds[0])
+        return float(self.summary.trust_bounds[0])
 
     @property
     def trust_upper_bound(self):
-        return float(self.trust_bounds[1])
+        return float(self.summary.trust_bounds[1])
 
     @property
     def readback(self):
@@ -128,7 +134,7 @@ def readback(self):
 
     @property
     def summary(self) -> pd.Series:
-        series = pd.Series(index=list(DOF_FIELD_TYPES.keys()))
+        series = pd.Series(index=list(DOF_FIELD_TYPES.keys()), dtype="object")
         for attr in series.index:
             value = getattr(self, attr)
             if attr == "trust_bounds":
@@ -203,11 +209,11 @@ def device_names(self) -> list:
 
     @property
     def search_lower_bounds(self) -> np.array:
-        return np.array([dof.search_lower_bound for dof in self.dofs])
+        return np.array([dof.search_lower_bound if not dof.read_only else dof.readback for dof in self.dofs])
 
     @property
     def search_upper_bounds(self) -> np.array:
-        return np.array([dof.search_upper_bound for dof in self.dofs])
+        return np.array([dof.search_upper_bound if not dof.read_only else dof.readback for dof in self.dofs])
 
     @property
     def search_bounds(self) -> np.array:
diff --git a/blop/objectives.py b/blop/objectives.py
index 205cc8f..f7df627 100644
--- a/blop/objectives.py
+++ b/blop/objectives.py
@@ -99,11 +99,14 @@ def label(self):
         return f"{'log ' if self.log else ''}{self.description}"
 
     @property
-    def summary(self):
-        series = pd.Series()
-        for attr, dtype in OBJ_FIELD_TYPES.items():
-            series[attr] = getattr(self, attr)
-
+    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
 
     def __repr__(self):
@@ -112,13 +115,21 @@ def __repr__(self):
     def __repr_html__(self):
         return self.summary.__repr_html__()
 
+    @property
+    def trust_lower_bound(self):
+        return float(self.summary.trust_bounds[0])
+
+    @property
+    def trust_upper_bound(self):
+        return float(self.summary.trust_bounds[1])
+
     @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
+        return np.round(1 / self.model.likelihood.noise.sqrt().item(), 3) if hasattr(self, "model") else None
 
     @property
     def n(self):
@@ -151,10 +162,11 @@ def __len__(self):
     def summary(self) -> pd.DataFrame:
         table = pd.DataFrame(columns=list(OBJ_FIELD_TYPES.keys()), index=self.names)
 
-        for attr, dtype in OBJ_FIELD_TYPES.items():
-            for obj in self.objectives:
-                table.at[obj.name, attr] = getattr(obj, attr)
+        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
@@ -180,7 +192,7 @@ def names(self) -> list:
         return [obj.name for obj in self.objectives]
 
     @property
-    def targets(self) -> np.array:
+    def targets(self) -> list:
         """
         Returns an array of the objective targets.
         """
diff --git a/blop/tests/conftest.py b/blop/tests/conftest.py
index 3531243..f381e6a 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")]
@@ -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_read_write.py b/blop/tests/test_read_write.py
index d6cf26b..db4b052 100644
--- a/blop/tests/test_read_write.py
+++ b/blop/tests/test_read_write.py
@@ -5,7 +5,7 @@ def test_agent_save_load_data(agent, RE):
     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")
+    agent.load_data("/tmp/test_save_data.h5")
     RE(agent.learn("qr", n=4))
 
 
@@ -13,4 +13,4 @@ def test_agent_save_load_hypers(agent, RE):
     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"))

From 616af0081329dd8df80260bf061275d346387de0 Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Mon, 22 Jan 2024 22:10:44 -0500
Subject: [PATCH 09/20] add log transform for DOFs

---
 blop/agent.py                            | 122 +++++++++++++++--------
 blop/dofs.py                             |  32 ++++--
 blop/objectives.py                       |   8 +-
 blop/tests/test_napari.py                |  10 +-
 docs/source/tutorials/passive-dofs.ipynb |   2 +-
 5 files changed, 117 insertions(+), 57 deletions(-)

diff --git a/blop/agent.py b/blop/agent.py
index 1a21b9f..671ab2c 100644
--- a/blop/agent.py
+++ b/blop/agent.py
@@ -19,6 +19,7 @@
 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
 
@@ -48,7 +49,10 @@ def _validate_dofs_and_objs(dofs: DOFList, objs: ObjectiveList):
         for latent_group in obj.latent_groups:
             for dof_name in latent_group:
                 if dof_name not in dofs.names:
-                    raise ValueError(f"DOF name '{dof_name}' in latent group for objective '{obj.name}' does not exist.")
+                    warnings.warn(
+                        f"DOF name '{dof_name}' in latent group for objective '{obj.name}' does not exist."
+                        "it will be ignored."
+                    )
 
 
 class Agent:
@@ -300,27 +304,21 @@ 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 == "random":
+            #     candidates = self.sample(n=n, active=True, read_only=False).squeeze(1).numpy()
+            #     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()
+                candidates = self.sample(n=n).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": {}}
+            # elif acq_func_name == "grid":
+            #     n_active_dims = len(self.dofs.subset(active=True, read_only=False))
+            #     candidates = self.test_inputs_grid(max_inputs=n).reshape(-1, n_active_dims).numpy()
+            #     acq_func_meta = {"name": "grid", "args": {}}
 
             else:
                 raise ValueError()
@@ -328,6 +326,12 @@ def ask(self, acq_func_identifier="qei", n=1, route=True, sequential=True, upsam
             # define dummy acqf objective
             acqf_obj = 0
 
+        active_dofs_are_read_only = self.dofs.subset(active=True).read_only
+
+        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
+
         acq_func_meta["duration"] = duration = ttime.monotonic() - start_time
 
         if self.verbose:
@@ -346,6 +350,7 @@ def ask(self, acq_func_identifier="qei", n=1, route=True, sequential=True, upsam
             "points": acq_points,
             "acq_func": acq_func_meta["name"],
             "acq_func_kwargs": acq_func_kwargs,
+            "acq_func_obj": acqf_obj,
             "duration": acq_func_meta["duration"],
             "sequential": sequential,
             "upsample": upsample,
@@ -500,7 +505,7 @@ def _construct_model(self, obj, skew_dims=None):
             train_targets=train_targets[trusted],
             likelihood=likelihood,
             skew_dims=skew_dims,
-            input_transform=self.input_transform,
+            input_transform=self.model_input_transform,
             outcome_transform=outcome_transform,
         )
 
@@ -523,7 +528,7 @@ def _construct_classifier(self, skew_dims=None):
             train_targets=dirichlet_likelihood.transformed_targets.transpose(-1, -2).double(),
             skew_dims=skew_dims,
             likelihood=dirichlet_likelihood,
-            input_transform=self.input_transform,
+            input_transform=self.model_input_transform,
         )
 
         self._train_model(self.classifier)
@@ -673,27 +678,70 @@ def latent_dim_tuples(self, obj_index=None):
         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).search_bounds.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.dofs.subset(active=True).search_bounds, dtype=torch.double).T
 
-    def _subset_inputs_sampler(self, active=None, read_only=None, tags=[], n=MAX_TEST_INPUTS):
+    @property
+    def sample_input_transform(self):
+        tf1 = Log10(indices=list(np.where(self.dofs.subset(active=True).log)[0]))
+
+        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))))
+
+        active_dofs = self.dofs.subset(active=True)
+
+        tf1 = Log10(indices=list(np.where(active_dofs.log)[0]))
+        tf2 = Normalize(d=len(active_dofs))
+
+        return ChainedInputTransform(tf1=tf1, tf2=tf2)
+
+    def sample(self, n, method="quasi-random"):
+        active_dofs = self.dofs.subset(active=True)
+
+        if method == "quasi-random":
+            X = utils.normalized_sobol_sampler(n, d=len(active_dofs))
+
+        if method == "random":
+            X = torch.rand(size=(n, 1, len(active_dofs)))
+
+        return self.sample_input_transform.untransform(X)
+
+    # 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).search_bounds.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 _subset_inputs_sampler(self, active=None, read_only=None, tags=[], n=MAX_TEST_INPUTS):
+    #     """
+    #     Returns $n$ quasi-randomly sampled inputs in the bounded parameter space
+    #     """
+    #     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))))
 
     def save_data(self, filepath="./self_data.h5"):
         """
@@ -824,10 +872,6 @@ def train_inputs(self, index=None, **subset_kwargs):
         valid = (inputs >= dof.trust_lower_bound) & (inputs <= dof.trust_upper_bound)
         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, index=None, **subset_kwargs):
diff --git a/blop/dofs.py b/blop/dofs.py
index cf4b43c..d7b2323 100644
--- a/blop/dofs.py
+++ b/blop/dofs.py
@@ -52,7 +52,7 @@ class DOF:
     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.
+        The agent will reject all data where the DOF value is outside the trust bounds. Much 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.
@@ -80,14 +80,18 @@ class DOF:
 
     # Some post-processing. This is specific to dataclasses
     def __post_init__(self):
-        if self.trust_bounds is not None:
-            self.trust_bounds = tuple(self.trust_bounds)
         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 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.")
+
         self.uuid = str(uuid.uuid4())
 
         if self.name is None:
@@ -112,21 +116,25 @@ def __post_init__(self):
     def search_lower_bound(self):
         if self.read_only:
             raise ValueError("Read-only DOFs do not have search bounds.")
-        return float(self.summary.search_bounds[0])
+        return float(self.search_bounds[0])
 
     @property
     def search_upper_bound(self):
         if self.read_only:
             raise ValueError("Read-only DOFs do not have search bounds.")
-        return float(self.summary.search_bounds[1])
+        return float(self.search_bounds[1])
 
     @property
     def trust_lower_bound(self):
-        return float(self.summary.trust_bounds[0])
+        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):
-        return float(self.summary.trust_bounds[1])
+        if self.trust_bounds is None:
+            return np.inf
+        return float(self.trust_bounds[1])
 
     @property
     def readback(self):
@@ -137,9 +145,9 @@ def summary(self) -> pd.Series:
         series = pd.Series(index=list(DOF_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)
+            # 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
 
@@ -249,6 +257,10 @@ def active(self):
     def read_only(self):
         return np.array([dof.read_only for dof in self.dofs])
 
+    @property
+    def log(self):
+        return np.array([dof.log for dof in self.dofs])
+
     def add(self, dof):
         _validate_dofs([*self.dofs, dof])
         self.dofs.append(dof)
diff --git a/blop/objectives.py b/blop/objectives.py
index f7df627..2f29a9d 100644
--- a/blop/objectives.py
+++ b/blop/objectives.py
@@ -117,11 +117,15 @@ def __repr_html__(self):
 
     @property
     def trust_lower_bound(self):
-        return float(self.summary.trust_bounds[0])
+        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):
-        return float(self.summary.trust_bounds[1])
+        if self.trust_bounds is None:
+            return np.inf
+        return float(self.trust_bounds[1])
 
     @property
     def noise(self):
diff --git a/blop/tests/test_napari.py b/blop/tests/test_napari.py
index bf6ac95..231f280 100644
--- a/blop/tests/test_napari.py
+++ b/blop/tests/test_napari.py
@@ -1,7 +1,7 @@
-import pytest  # noqa F401
+# 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)
+# @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/docs/source/tutorials/passive-dofs.ipynb b/docs/source/tutorials/passive-dofs.ipynb
index bd96f7c..4c3f4fd 100644
--- a/docs/source/tutorials/passive-dofs.ipynb
+++ b/docs/source/tutorials/passive-dofs.ipynb
@@ -104,7 +104,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.0"
+   "version": "3.11.5"
   },
   "vscode": {
    "interpreter": {

From 74086d9851a8f4a52b3327d8d4c08d55535ba4ff Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Tue, 23 Jan 2024 15:53:32 -0500
Subject: [PATCH 10/20] finish log transformed dofs

---
 blop/agent.py                         | 580 +++++++++++---------------
 blop/bayesian/acquisition/__init__.py |  16 +-
 blop/bayesian/models.py               |  18 +-
 blop/bayesian/plotting.py             |  32 +-
 blop/dofs.py                          |  39 +-
 blop/tests/conftest.py                |   2 +-
 blop/tests/test_acq_funcs.py          |  12 +-
 blop/tests/test_passive_dofs.py       |   4 -
 blop/tests/test_plots.py              |  25 +-
 9 files changed, 328 insertions(+), 400 deletions(-)

diff --git a/blop/agent.py b/blop/agent.py
index 671ab2c..29f1dd4 100644
--- a/blop/agent.py
+++ b/blop/agent.py
@@ -35,7 +35,7 @@
 
 mpl.rc("image", cmap="coolwarm")
 
-MAX_TEST_INPUTS = 2**11
+DEFAULT_MAX_SAMPLES = 2**11
 
 
 def _validate_dofs_and_objs(dofs: DOFList, objs: ObjectiveList):
@@ -148,108 +148,41 @@ def __iter__(self):
     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)
+            return self._get_acquisition_function(identifier=acq_func_name)
 
-        raise AttributeError(f"DOFList object has no attribute named '{attr}'.")
+        raise AttributeError(f"No attribute named '{attr}'.")
 
-    def view(self, item: str = "mean", cmap: str = "turbo", max_inputs: int = 2**16):
+    def sample(self, n: int = DEFAULT_MAX_SAMPLES, method: str = "quasi-random") -> torch.Tensor:
         """
-        Use napari to see a high-dimensional array.
+        Returns a (..., 1, n_active_dofs) tensor of points sampled within the parameter space.
 
         Parameters
         ----------
-        item : str
-            The thing to be viewed. Either 'mean', 'error', or an acquisition function.
+        n : int
+            How many points to sample.
+        method : str
+            How to sample the points. Must be one of 'quasi-random', 'random', or 'grid'.
         """
 
-        test_grid = self.test_inputs_grid(max_inputs=max_inputs)
-
-        self.viewer = napari.Viewer()
+        active_dofs = self.dofs.subset(active=True)
 
-        if item in ["mean", "error"]:
-            for obj in self.objectives:
-                p = obj.model.posterior(test_grid)
+        if method == "quasi-random":
+            X = utils.normalized_sobol_sampler(n, d=len(active_dofs))
 
-                if item == "mean":
-                    mean = p.mean.detach().numpy()[..., 0, 0]
-                    self.viewer.add_image(data=mean, name=f"{obj.name}_mean", colormap=cmap)
+        elif method == "random":
+            X = torch.rand(size=(n, 1, len(active_dofs)))
 
-                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)
+        elif method == "grid":
+            n_side_if_settable = int(np.power(n, 1 / np.sum(~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 active_dofs
+            ]
+            X = torch.cat([x.unsqueeze(-1) for x in torch.meshgrid(sides, indexing="ij")], dim=-1).unsqueeze(-2).double()
 
         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 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: 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
-
-            obj.model = self._construct_model(obj)
+            raise ValueError("'method' argument must be one of ['quasi-random', 'random', 'grid'].")
 
-            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")
-
-        # TODO: should this be per objective?
-        self._construct_classifier()
+        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.
@@ -272,10 +205,13 @@ 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()
+
+            # define dummy acqf objective
+            acqf_obj = torch.zeros(len(candidates))
 
-        if acq_func_type in ["analytic", "monte_carlo"]:
+        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}'"
@@ -285,7 +221,7 @@ def ask(self, acq_func_identifier="qei", n=1, route=True, sequential=True, upsam
             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
             )
 
@@ -294,7 +230,7 @@ def ask(self, acq_func_identifier="qei", n=1, route=True, sequential=True, upsam
 
             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,
@@ -304,38 +240,12 @@ 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()
 
-        else:
-            acqf_obj = None
-
-            # if acq_func_name == "random":
-            #     candidates = self.sample(n=n, active=True, read_only=False).squeeze(1).numpy()
-            #     acq_func_meta = {"name": "random", "args": {}}
-
-            if acq_func_name == "quasi-random":
-                candidates = self.sample(n=n).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))
-            #     candidates = 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
-
         active_dofs_are_read_only = self.dofs.subset(active=True).read_only
 
         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
-
-        acq_func_meta["duration"] = duration = ttime.monotonic() - start_time
 
-        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)
@@ -346,75 +256,81 @@ 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,
-            "acq_func_obj": acqf_obj,
-            "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,
+        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
-
-            uid = yield from self.acquisition_plan(
-                acquisition_devices,
-                acquisition_inputs.astype(float),
-                [*self.dets, *self.dofs.devices],
-                delay=self.trigger_delay,
-            )
+        if not data:
+            if not x and y:
+                raise ValueError("Must supply either x and y, or data.")
+            data = {**x, **y, **metadata}
 
-            products = self.digestion(self.db, uid)
+        data = {k: list(np.atleast_1d(v)) for k, v in data.items()}
+        unique_field_lengths = {len(v) for v in data.values()}
 
-            # 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'])
+        if len(unique_field_lengths) > 1:
+            raise ValueError("All supplies values must be the same length!")
 
-        except KeyboardInterrupt as interrupt:
-            raise interrupt
+        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 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
+        for obj in self.objectives:
+            t0 = ttime.monotonic()
 
-        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!")
+            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
 
-        return products
+            if train is None:
+                train = int(n_after_tell / self.train_every) > int(n_before_tell / self.train_every)
 
-    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)
+            if (len(obj.model.train_targets) >= 2) and train:
+                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 learn(
         self,
@@ -422,7 +338,7 @@ def learn(
         n: int = 1,
         iterations: int = 1,
         upsample: int = 1,
-        train: bool = True,
+        train: bool = None,
         append: bool = True,
         hypers: str = None,
         route: bool = True,
@@ -470,75 +386,90 @@ def learn(
                 metadata = new_table.to_dict(orient="list")
                 self.tell(x=x, y=y, metadata=metadata, append=append, train=train)
 
-    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):
+    def view(self, item: str = "mean", cmap: str = "turbo", max_inputs: int = 2**16):
         """
-        Construct an untrained model for an objective.
+        Use napari to see a high-dimensional array.
+
+        Parameters
+        ----------
+        item : str
+            The thing to be viewed. Either 'mean', 'error', or an acquisition function.
         """
 
-        skew_dims = skew_dims if skew_dims is not None else self.latent_dim_tuples(obj.name)
+        test_grid = self.sample(n=max_inputs, method="grid")
 
-        likelihood = gpytorch.likelihoods.GaussianLikelihood(
-            noise_constraint=gpytorch.constraints.Interval(
-                torch.tensor(obj.min_noise),
-                torch.tensor(obj.max_noise),
-            ),
-            # noise_prior=gpytorch.priors.torch_priors.LogNormalPrior(loc=loc, scale=scale),
-        )
+        self.viewer = napari.Viewer()
 
-        outcome_transform = botorch.models.transforms.outcome.Standardize(m=1)  # , batch_shape=torch.Size((1,)))
+        if item in ["mean", "error"]:
+            for obj in self.objectives:
+                p = obj.model.posterior(test_grid)
 
-        train_inputs = self.train_inputs(active=True)
-        train_targets = self.train_targets(obj.name)
+                if item == "mean":
+                    mean = p.mean.detach().numpy()[..., 0, 0]
+                    self.viewer.add_image(data=mean, name=f"{obj.name}_mean", colormap=cmap)
 
-        trusted = ~(torch.isnan(train_inputs).any(axis=1) | torch.isnan(train_targets).any(axis=1))
+                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)
 
-        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,
-        )
+        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.")
 
-        model.trained = False
+            acq_func, acq_func_meta = self._get_acquisition_function(identifier=acq_func_identifier, return_metadata=True)
+            a = acq_func(test_grid).detach().numpy()
 
-        return model
+            self.viewer.add_image(data=a, name=f"{acq_func_identifier}", colormap=cmap)
 
-    def _construct_classifier(self, skew_dims=None):
-        skew_dims = [tuple([i]) for i in range(len(self.dofs.subset(active=True)))]
+        self.viewer.dims.axis_labels = self.dofs.names
 
-        train_inputs = self.train_inputs(active=True)
-        trusted = ~torch.isnan(train_inputs).any(axis=1)
+    def acquire(self, acquisition_inputs):
+        """Acquire and digest according to the self's acquisition and digestion plans.
 
-        dirichlet_likelihood = gpytorch.likelihoods.DirichletClassificationLikelihood(
-            self.all_objectives_valid.long()[trusted], learn_additional_noise=True
-        )
+        Parameters
+        ----------
+        acquisition_inputs :
+            A 2D numpy array comprising inputs for the active and non-read-only DOFs to sample.
+        """
 
-        self.classifier = models.LatentDirichletClassifier(
-            train_inputs=train_inputs[trusted],
-            train_targets=dirichlet_likelihood.transformed_targets.transpose(-1, -2).double(),
-            skew_dims=skew_dims,
-            likelihood=dirichlet_likelihood,
-            input_transform=self.model_input_transform,
-        )
+        if self.db is None:
+            raise ValueError("Cannot run acquistion without databroker instance!")
 
-        self._train_model(self.classifier)
-        self.constraint = GenericDeterministicModel(f=lambda x: self.classifier.probabilities(x)[..., -1])
+        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,
+            )
 
-    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`.
-        """
-        return acquisition.get_acquisition_function(self, identifier=identifier, return_metadata=return_metadata)
+            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.objectives:
+                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."""
@@ -580,7 +511,7 @@ def model(self):
 
     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)
+        return self.model.posterior(torch.tensor(x))
 
     @property
     def objective_weights_torch(self):
@@ -616,54 +547,95 @@ 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):
         """
-        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).search_bounds, n_sides
-                            )
-                        ],
-                        indexing="ij",
-                    )
-                ],
-                dim=-1,
-            )
-            .unsqueeze(-2)
-            .double()
+        Construct an untrained model for an objective.
+        """
+
+        skew_dims = skew_dims if skew_dims is not None else self._latent_dim_tuples(obj.name)
+
+        likelihood = gpytorch.likelihoods.GaussianLikelihood(
+            noise_constraint=gpytorch.constraints.Interval(
+                torch.tensor(obj.min_noise),
+                torch.tensor(obj.max_noise),
+            ),
         )
 
-    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)
+        outcome_transform = botorch.models.transforms.outcome.Standardize(m=1)  # , batch_shape=torch.Size((1,)))
 
-    @property
-    def acquisition_function_bounds(self):
-        """Returns a (2, n_active_dof) array of bounds for the acquisition function"""
+        train_inputs = self.train_inputs(active=True)
+        train_targets = self.train_targets(obj.name)
+
+        trusted = ~(torch.isnan(train_inputs).any(axis=1) | torch.isnan(train_targets).any(axis=1))
+
+        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,
+        )
+
+        if trusted.all():
+            obj.classifier_conjugate_model = None
+            obj.classifier = GenericDeterministicModel(f=lambda x: torch.ones(size=x.size())[..., -1])
+
+        else:
+            dirichlet_likelihood = gpytorch.likelihoods.DirichletClassificationLikelihood(
+                trusted.long(), learn_additional_noise=True
+            )
+
+            obj.classifier_conjugate_model = models.LatentDirichletClassifier(
+                train_inputs=train_inputs,
+                train_targets=dirichlet_likelihood.transformed_targets.transpose(-1, -2).double(),
+                skew_dims=skew_dims,
+                likelihood=dirichlet_likelihood,
+                input_transform=self._model_input_transform,
+            )
+
+            obj.classifier = GenericDeterministicModel(f=lambda x: obj.classifier_conjugate_model.probabilities(x)[..., -1])
+
+    def _construct_all_models(self):
+        """Construct a model for each objective."""
+        for obj in self.objectives:
+            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.objectives:
+            self._train_model(obj.model)
+            if obj.classifier_conjugate_model is not None:
+                self._train_model(obj.classifier_conjugate_model)
 
-        acq_func_lower_bounds = np.where(self.dofs.read_only, self.dofs.readback, self.dofs.search_lower_bounds)
-        acq_func_upper_bounds = np.where(self.dofs.read_only, self.dofs.readback, self.dofs.search_upper_bounds)
+        if self.verbose:
+            print(f"trained models in {ttime.monotonic() - t0:.01f} seconds")
+
+        self.n_last_trained = len(self.table)
 
-        return torch.tensor(np.vstack([acq_func_lower_bounds, acq_func_upper_bounds]), dtype=torch.double)
+    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`.
+        """
+        return acquisition.get_acquisition_function(self, identifier=identifier, return_metadata=return_metadata)
 
-    def latent_dim_tuples(self, obj_index=None):
+    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.
         """
 
         if obj_index is None:
-            return {obj.name: self.latent_dim_tuples(obj_index=obj.name) for obj in self.objectives}
+            return {obj.name: self._latent_dim_tuples(obj_index=obj.name) for obj in self.objectives}
 
         obj = self.objectives[obj_index]
 
@@ -678,14 +650,14 @@ def latent_dim_tuples(self, obj_index=None):
         return [tuple(np.where(uinv == i)[0]) for i in range(len(u))]
 
     @property
-    def sample_bounds(self):
+    def _sample_bounds(self):
         return torch.tensor(self.dofs.subset(active=True).search_bounds, dtype=torch.double).T
 
     @property
-    def sample_input_transform(self):
+    def _sample_input_transform(self):
         tf1 = Log10(indices=list(np.where(self.dofs.subset(active=True).log)[0]))
 
-        transformed_sample_bounds = tf1.transform(self.sample_bounds)
+        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)
@@ -695,7 +667,7 @@ def sample_input_transform(self):
         return ChainedInputTransform(tf1=tf1, tf2=tf2)
 
     @property
-    def model_input_transform(self):
+    def _model_input_transform(self):
         """
         Suitably transforms model inputs to the unit hypercube.
 
@@ -716,33 +688,6 @@ def model_input_transform(self):
 
         return ChainedInputTransform(tf1=tf1, tf2=tf2)
 
-    def sample(self, n, method="quasi-random"):
-        active_dofs = self.dofs.subset(active=True)
-
-        if method == "quasi-random":
-            X = utils.normalized_sobol_sampler(n, d=len(active_dofs))
-
-        if method == "random":
-            X = torch.rand(size=(n, 1, len(active_dofs)))
-
-        return self.sample_input_transform.untransform(X)
-
-    # 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).search_bounds.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 _subset_inputs_sampler(self, active=None, read_only=None, tags=[], n=MAX_TEST_INPUTS):
-    #     """
-    #     Returns $n$ quasi-randomly sampled inputs in the bounded parameter space
-    #     """
-    #     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))))
-
     def save_data(self, filepath="./self_data.h5"):
         """
         Save the sampled inputs and targets of the agent to a file, which can be used
@@ -777,19 +722,22 @@ def forget(self, last=None, index=None, train=True):
         else:
             raise ValueError("Must supply either 'last' or 'index'.")
 
-    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]
-
     def _set_hypers(self, hypers):
         for obj in self.objectives:
             obj.model.load_state_dict(hypers[obj.name])
         self.classifier.load_state_dict(hypers["classifier"])
 
+    @property
+    def classifier(self):
+        def f(x):
+            p = torch.ones(x.shape[:-1])
+            for obj in self.objectives:
+                if obj.classifier_conjugate_model is not None:
+                    p *= obj.classifier(x)
+            return p
+
+        return GenericDeterministicModel(f=f)
+
     @property
     def hypers(self):
         """Returns a dict of all the hyperparameters in all the agent's models."""
@@ -823,25 +771,6 @@ def load_hypers(filepath):
                     hypers[model_key][param_key] = torch.tensor(np.atleast_1d(param_value[()]))
         return hypers
 
-    def _construct_all_models(self):
-        """Construct a model for each objective."""
-        for obj in self.objectives:
-            obj.model = 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.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")
-
-        self.n_last_trained = len(self.table)
-
     @property
     def all_acq_funcs(self):
         """Description and identifiers for all supported acquisition functions."""
@@ -854,11 +783,6 @@ 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, index=None, **subset_kwargs):
         """A two-dimensional tensor of all DOF values."""
 
@@ -893,16 +817,6 @@ def train_targets(self, index=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."""
@@ -963,7 +877,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
diff --git a/blop/bayesian/acquisition/__init__.py b/blop/bayesian/acquisition/__init__.py
index 24bdacf..6740e65 100644
--- a/blop/bayesian/acquisition/__init__.py
+++ b/blop/bayesian/acquisition/__init__.py
@@ -41,7 +41,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
     # there is probably a better way to structure this
     if acq_func_name == "expected_improvement":
         acq_func = analytic.ConstrainedLogExpectedImprovement(
-            constraint=agent.constraint,
+            constraint=agent.classifier,
             model=agent.model,
             best_f=agent.max_scalarized_objective,
             posterior_transform=agent.targeting_transform,
@@ -50,7 +50,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
 
     elif acq_func_name == "monte_carlo_expected_improvement":
         acq_func = monte_carlo.qConstrainedExpectedImprovement(
-            constraint=agent.constraint,
+            constraint=agent.classifier,
             model=agent.model,
             best_f=agent.max_scalarized_objective,
             posterior_transform=agent.targeting_transform,
@@ -59,7 +59,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
 
     elif acq_func_name == "probability_of_improvement":
         acq_func = analytic.ConstrainedLogProbabilityOfImprovement(
-            constraint=agent.constraint,
+            constraint=agent.classifier,
             model=agent.model,
             best_f=agent.max_scalarized_objective,
             posterior_transform=agent.targeting_transform,
@@ -68,7 +68,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
 
     elif acq_func_name == "monte_carlo_probability_of_improvement":
         acq_func = monte_carlo.qConstrainedProbabilityOfImprovement(
-            constraint=agent.constraint,
+            constraint=agent.classifier,
             model=agent.model,
             best_f=agent.max_scalarized_objective,
             posterior_transform=agent.targeting_transform,
@@ -77,7 +77,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
 
     elif acq_func_name == "lower_bound_max_value_entropy":
         acq_func = monte_carlo.qConstrainedLowerBoundMaxValueEntropy(
-            constraint=agent.constraint,
+            constraint=agent.classifier,
             model=agent.model,
             candidate_set=agent.test_inputs(n=1024).squeeze(1),
         )
@@ -85,7 +85,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
 
     elif acq_func_name == "noisy_expected_hypervolume_improvement":
         acq_func = monte_carlo.qConstrainedNoisyExpectedHypervolumeImprovement(
-            constraint=agent.constraint,
+            constraint=agent.classifier,
             model=agent.model,
             ref_point=agent.train_targets.min(dim=0).values,
             X_baseline=agent.train_inputs,
@@ -97,7 +97,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
         beta = acq_func_kwargs.get("beta", acq_func_config["default_args"]["beta"])
 
         acq_func = analytic.ConstrainedUpperConfidenceBound(
-            constraint=agent.constraint,
+            constraint=agent.classifier,
             model=agent.model,
             beta=beta,
             posterior_transform=agent.targeting_transform,
@@ -108,7 +108,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
         beta = acq_func_kwargs.get("beta", acq_func_config["default_args"]["beta"])
 
         acq_func = monte_carlo.qConstrainedUpperConfidenceBound(
-            constraint=agent.constraint,
+            constraint=agent.classifier,
             model=agent.model,
             beta=beta,
             posterior_transform=agent.targeting_transform,
diff --git a/blop/bayesian/models.py b/blop/bayesian/models.py
index a50e84e..9dd15cc 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
         )
 
-
-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
-        )
+        self.trained = False
 
 
 class LatentDirichletClassifier(LatentGP):
     def __init__(self, train_inputs, train_targets, skew_dims=True, *args, **kwargs):
         super().__init__(train_inputs, train_targets, skew_dims, *args, **kwargs)
 
+        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/plotting.py b/blop/bayesian/plotting.py
index 4dee48b..a9968f0 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)
@@ -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()
 
@@ -181,6 +181,10 @@ def _plot_objs_many_dofs(agent, axes=(0, 1), shading="nearest", cmap=DEFAULT_COL
         ax.set_ylabel(y_dof.label)
         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]
@@ -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()
@@ -269,6 +273,10 @@ def _plot_acqf_many_dofs(
         ax.set_ylabel(y_dof.label)
         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,8 +285,8 @@ 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()
-    constraint = agent.constraint(test_inputs)[..., 0]
+    test_inputs = agent.sample(method="grid")
+    constraint = agent.classifier(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)
@@ -286,7 +294,7 @@ def _plot_valid_one_dof(agent, size=16, lw=1e0):
 
 
 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,11 +304,11 @@ 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()
 
-    constraint = agent.constraint(test_inputs)[..., 0].squeeze().numpy()
+    constraint = agent.classifier(test_inputs)[..., 0].squeeze().numpy()
 
     if gridded:
         _ = agent.valid_axes[1].pcolormesh(
@@ -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.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/dofs.py b/blop/dofs.py
index d7b2323..cfb5e85 100644
--- a/blop/dofs.py
+++ b/blop/dofs.py
@@ -113,28 +113,33 @@ def __post_init__(self):
         self.device.kind = "hinted"
 
     @property
-    def search_lower_bound(self):
+    def _search_bounds(self):
         if self.read_only:
-            raise ValueError("Read-only DOFs do not have search bounds.")
-        return float(self.search_bounds[0])
+            _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 search_lower_bound(self):
+        return float(self._search_bounds[0])
 
     @property
     def search_upper_bound(self):
-        if self.read_only:
-            raise ValueError("Read-only DOFs do not have search bounds.")
-        return float(self.search_bounds[1])
+        return float(self._search_bounds[1])
 
     @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])
+        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])
+        return float(self._trust_bounds[1])
 
     @property
     def readback(self):
@@ -144,11 +149,7 @@ def readback(self):
     def summary(self) -> pd.Series:
         series = pd.Series(index=list(DOF_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
+            series[attr] = getattr(self, attr)
         return series
 
     @property
@@ -228,7 +229,7 @@ def search_bounds(self) -> np.array:
         """
         Returns a (n_dof, 2) array of bounds.
         """
-        return np.c_[self.search_lower_bounds, self.search_upper_bounds]
+        return np.array([dof._search_bounds for dof in self.dofs])
 
     @property
     def trust_lower_bounds(self) -> np.array:
@@ -243,7 +244,7 @@ def trust_bounds(self) -> np.array:
         """
         Returns a (n_dof, 2) array of bounds.
         """
-        return np.c_[self.trust_lower_bounds, self.trust_upper_bounds]
+        return np.array([dof._trust_bounds for dof in self.dofs])
 
     @property
     def readback(self) -> np.array:
diff --git a/blop/tests/conftest.py b/blop/tests/conftest.py
index f381e6a..6d3e781 100644
--- a/blop/tests/conftest.py
+++ b/blop/tests/conftest.py
@@ -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
     """
diff --git a/blop/tests/test_acq_funcs.py b/blop/tests/test_acq_funcs.py
index a041dc9..be376b7 100644
--- a/blop/tests/test_acq_funcs.py
+++ b/blop/tests/test_acq_funcs.py
@@ -14,12 +14,12 @@ def test_monte_carlo_acq_funcs_single_objective(agent, RE, acq_func):
 
 
 @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_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()

From 7f7ea3d3839f87cd063611e007c71923c0871aac Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Fri, 26 Jan 2024 13:02:58 -0500
Subject: [PATCH 11/20] better subsetter for dofs and objs

---
 blop/agent.py      |  75 ++++++++++++-----------
 blop/dofs.py       | 144 ++++++++++++++++-----------------------------
 blop/objectives.py |  26 +++++---
 3 files changed, 105 insertions(+), 140 deletions(-)

diff --git a/blop/agent.py b/blop/agent.py
index 29f1dd4..507bb94 100644
--- a/blop/agent.py
+++ b/blop/agent.py
@@ -16,7 +16,6 @@
 import pandas as pd
 import scipy as sp
 import torch
-from botorch.acquisition.objective import ScalarizedPosteriorTransform
 from botorch.models.deterministic import GenericDeterministicModel
 from botorch.models.model_list_gp_regression import ModelListGP
 from botorch.models.transforms.input import AffineInputTransform, ChainedInputTransform, Log10, Normalize
@@ -141,6 +140,14 @@ def __init__(
 
         self.n_last_trained = 0
 
+    @property
+    def active_dofs(self):
+        return self.dofs.subset(active=True)
+
+    @property
+    def active_objs(self):
+        return self.objectives.subset(active=True)
+
     def __iter__(self):
         for index in range(len(self)):
             yield self.dofs[index]
@@ -164,18 +171,16 @@ def sample(self, n: int = DEFAULT_MAX_SAMPLES, method: str = "quasi-random") ->
             How to sample the points. Must be one of 'quasi-random', 'random', or 'grid'.
         """
 
-        active_dofs = self.dofs.subset(active=True)
-
         if method == "quasi-random":
-            X = utils.normalized_sobol_sampler(n, d=len(active_dofs))
+            X = utils.normalized_sobol_sampler(n, d=len(self.active_dofs))
 
         elif method == "random":
-            X = torch.rand(size=(n, 1, len(active_dofs)))
+            X = torch.rand(size=(n, 1, len(self.active_dofs)))
 
         elif method == "grid":
-            n_side_if_settable = int(np.power(n, 1 / np.sum(~active_dofs.read_only)))
+            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 active_dofs
+                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()
 
@@ -240,10 +245,8 @@ 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 = self.dofs.subset(active=True).read_only
-
-        acq_points = candidates[..., ~active_dofs_are_read_only]
-        read_only_values = candidates[..., active_dofs_are_read_only]
+        acq_points = candidates[..., ~self.active_dofs.read_only]
+        read_only_values = candidates[..., self.active_dofs.read_only]
 
         duration = 1e3 * (ttime.monotonic() - start_time)
 
@@ -315,7 +318,7 @@ def tell(
         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:
+        for obj in self.active_objs:
             t0 = ttime.monotonic()
 
             cached_hypers = obj.model.state_dict() if hasattr(obj, "model") else None
@@ -401,7 +404,7 @@ def view(self, item: str = "mean", cmap: str = "turbo", max_inputs: int = 2**16)
         self.viewer = napari.Viewer()
 
         if item in ["mean", "error"]:
-            for obj in self.objectives:
+            for obj in self.active_objs:
                 p = obj.model.posterior(test_grid)
 
                 if item == "mean":
@@ -456,7 +459,7 @@ def acquire(self, acquisition_inputs):
                 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:
+            for obj in self.active_objs:
                 products.loc[:, obj.name] = np.nan
 
         if not len(acquisition_inputs) == len(products):
@@ -475,7 +478,7 @@ 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
 
@@ -507,23 +510,19 @@ 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(torch.tensor(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)
-
     @property
     def targeting_transform(self):
-        return TargetingPosteriorTransform(weights=self.objective_weights_torch, targets=self.objectives.targets)
+        return TargetingPosteriorTransform(
+            weights=torch.tensor(self.active_objs.weights, dtype=torch.double), targets=self.active_objs.targets
+        )
 
     @property
     def scalarized_objectives(self):
@@ -606,13 +605,13 @@ def _construct_model(self, obj, skew_dims=None):
 
     def _construct_all_models(self):
         """Construct a model for each objective."""
-        for obj in self.objectives:
+        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.objectives:
+        for obj in self.active_objs:
             self._train_model(obj.model)
             if obj.classifier_conjugate_model is not None:
                 self._train_model(obj.classifier_conjugate_model)
@@ -640,7 +639,7 @@ def _latent_dim_tuples(self, obj_index=None):
         obj = self.objectives[obj_index]
 
         latent_group_index = {}
-        for dof in self.dofs.subset(active=True):
+        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:
@@ -651,11 +650,11 @@ def _latent_dim_tuples(self, obj_index=None):
 
     @property
     def _sample_bounds(self):
-        return torch.tensor(self.dofs.subset(active=True).search_bounds, dtype=torch.double).T
+        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.dofs.subset(active=True).log)[0]))
+        tf1 = Log10(indices=list(np.where(self.active_dofs.log)[0]))
 
         transformed_sample_bounds = tf1.transform(self._sample_bounds)
 
@@ -681,10 +680,8 @@ def _model_input_transform(self):
         Read-only: constrain to the readback value
         """
 
-        active_dofs = self.dofs.subset(active=True)
-
-        tf1 = Log10(indices=list(np.where(active_dofs.log)[0]))
-        tf2 = Normalize(d=len(active_dofs))
+        tf1 = Log10(indices=list(np.where(self.active_dofs.log)[0]))
+        tf2 = Normalize(d=len(self.active_dofs))
 
         return ChainedInputTransform(tf1=tf1, tf2=tf2)
 
@@ -723,7 +720,7 @@ def forget(self, last=None, index=None, train=True):
             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"])
 
@@ -731,7 +728,7 @@ def _set_hypers(self, hypers):
     def classifier(self):
         def f(x):
             p = torch.ones(x.shape[:-1])
-            for obj in self.objectives:
+            for obj in self.active_objs:
                 if obj.classifier_conjugate_model is not None:
                     p *= obj.classifier(x)
             return p
@@ -744,7 +741,7 @@ def hypers(self):
         hypers = {"classifier": {}}
         for key, value in self.classifier.state_dict().items():
             hypers["classifier"][key] = value
-        for obj in self.objectives:
+        for obj in self.active_objs:
             hypers[obj.name] = {}
             for key, value in obj.model.state_dict().items():
                 hypers[obj.name][key] = value
@@ -793,7 +790,7 @@ def train_inputs(self, index=None, **subset_kwargs):
         inputs = self.table.loc[:, dof.name].values.copy()
 
         # check that inputs values are inside acceptable values
-        valid = (inputs >= dof.trust_lower_bound) & (inputs <= dof.trust_upper_bound)
+        valid = (inputs >= dof._trust_bounds[0]) & (inputs <= dof._trust_bounds[1])
         inputs = np.where(valid, inputs, np.nan)
 
         return torch.tensor(inputs, dtype=torch.double).unsqueeze(-1)
@@ -808,7 +805,7 @@ def train_targets(self, index=None, **subset_kwargs):
         targets = self.table.loc[:, obj.name].values.copy()
 
         # check that targets values are inside acceptable values
-        valid = (targets >= obj.trust_lower_bound) & (targets <= obj.trust_upper_bound)
+        valid = (targets >= obj._trust_bounds[0]) & (targets <= obj._trust_bounds[1])
         targets = np.where(valid, targets, np.nan)
 
         # transform if needed
diff --git a/blop/dofs.py b/blop/dofs.py
index cfb5e85..6a68d83 100644
--- a/blop/dofs.py
+++ b/blop/dofs.py
@@ -1,6 +1,6 @@
 import time as ttime
 import uuid
-from collections.abc import Sequence
+from collections.abc import Iterable, Sequence
 from dataclasses import dataclass, field
 from typing import Tuple
 
@@ -9,11 +9,11 @@
 from ophyd import Signal, SignalRO
 
 DOF_FIELD_TYPES = {
-    "description": "object",
+    "description": "str",
     "readback": "float",
     "search_bounds": "object",
     "trust_bounds": "object",
-    "units": "object",
+    "units": "str",
     "active": "bool",
     "read_only": "bool",
     "log": "bool",
@@ -92,13 +92,16 @@ def __post_init__(self):
                 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.")
 
-        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.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 self.device is None:
-            self.device = Signal(name=self.name)
+        if self.description is None:
+            self.description = self.name
 
         if not self.read_only:
             # check that the device has a put method
@@ -125,22 +128,6 @@ def _trust_bounds(self):
             return (0, np.inf) if self.log else (-np.inf, np.inf)
         return self.trust_bounds
 
-    @property
-    def search_lower_bound(self):
-        return float(self._search_bounds[0])
-
-    @property
-    def search_upper_bound(self):
-        return float(self._search_bounds[1])
-
-    @property
-    def trust_lower_bound(self):
-        return float(self._trust_bounds[0])
-
-    @property
-    def trust_upper_bound(self):
-        return float(self._trust_bounds[1])
-
     @property
     def readback(self):
         return self.device.read()[self.device.name]["value"]
@@ -167,20 +154,27 @@ def __init__(self, dofs: list = []):
         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, index):
-        if type(index) is int:
-            return self.dofs[index]
-        elif type(index) is str:
-            if index not in self.names:
-                raise ValueError(f"DOFList has no DOF named {index}.")
-            return self.dofs[self.names.index(index)]
+    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 {index}. A DOFList must be indexed by either an integer or a string.")
+            raise ValueError(f"Invalid index {key}.")
 
     def __len__(self):
         return len(self.dofs)
@@ -212,18 +206,6 @@ def names(self) -> list:
     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 search_lower_bounds(self) -> np.array:
-        return np.array([dof.search_lower_bound if not dof.read_only else dof.readback for dof in self.dofs])
-
-    @property
-    def search_upper_bounds(self) -> np.array:
-        return np.array([dof.search_upper_bound if not dof.read_only else dof.readback for dof in self.dofs])
-
     @property
     def search_bounds(self) -> np.array:
         """
@@ -231,14 +213,6 @@ def search_bounds(self) -> np.array:
         """
         return np.array([dof._search_bounds for dof in self.dofs])
 
-    @property
-    def trust_lower_bounds(self) -> np.array:
-        return np.array([dof.trust_lower_bound for dof in self.dofs])
-
-    @property
-    def trust_upper_bounds(self) -> np.array:
-        return np.array([dof.trust_upper_bound for dof in self.dofs])
-
     @property
     def trust_bounds(self) -> np.array:
         """
@@ -246,51 +220,35 @@ def trust_bounds(self) -> np.array:
         """
         return np.array([dof._trust_bounds for dof in self.dofs])
 
-    @property
-    def readback(self) -> np.array:
-        return np.array([dof.readback for dof in self.dofs])
-
-    @property
-    def active(self):
-        return np.array([dof.active for dof in self.dofs])
-
-    @property
-    def read_only(self):
-        return np.array([dof.read_only for dof in self.dofs])
-
-    @property
-    def log(self):
-        return np.array([dof.log for dof in self.dofs])
-
     def add(self, dof):
         _validate_dofs([*self.dofs, dof])
         self.dofs.append(dof)
 
-    def _dof_active_mask(self, active=None):
-        return [_active == active if active is not None else True for _active in self.active]
-
-    def _dof_read_only_mask(self, read_only=None):
-        return [_read_only == read_only if read_only is not None else True for _read_only in self.read_only]
-
-    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
+    @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, read_only=None, active=None, tags=[]):
-        for dof in self._subset_dofs(read_only, active, tags):
-            dof.active = False
+    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):
diff --git a/blop/objectives.py b/blop/objectives.py
index 2f29a9d..1be23a5 100644
--- a/blop/objectives.py
+++ b/blop/objectives.py
@@ -77,22 +77,22 @@ class Objective:
     log: bool = False
     weight: float = 1.0
     active: bool = True
-    trust_bounds: Tuple[float, float] = None
+    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 self.trust_bounds is None:
-            if self.log:
-                self.trust_bounds = (0, np.inf)
-            else:
-                self.trust_bounds = (-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.")
+                raise ValueError("'target' must be either 'min', 'max', a number, or a tuple of numbers.")
+
+    @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):
@@ -219,3 +219,13 @@ def signed_weights(self) -> np.array:
     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)])

From c29f2c01cc8ca22302db167fc013663b1ea947eb Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Fri, 26 Jan 2024 17:45:20 -0500
Subject: [PATCH 12/20] method for active dofs and objs

---
 blop/agent.py      | 25 ++++++++++++++++++++-----
 blop/dofs.py       |  5 +----
 blop/objectives.py | 27 ++++++++++++++++++---------
 3 files changed, 39 insertions(+), 18 deletions(-)

diff --git a/blop/agent.py b/blop/agent.py
index 507bb94..df97a71 100644
--- a/blop/agent.py
+++ b/blop/agent.py
@@ -223,6 +223,11 @@ def ask(self, acq_func_identifier="qei", n=1, route=True, sequential=True, upsam
                     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.")
 
@@ -329,11 +334,15 @@ def tell(
             if train is None:
                 train = int(n_after_tell / self.train_every) > int(n_before_tell / self.train_every)
 
-            if (len(obj.model.train_targets) >= 2) and train:
-                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")
+            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")
+
+                else:
+                    self._train_model(obj.model, hypers=cached_hypers)
 
     def learn(
         self,
@@ -584,6 +593,8 @@ def _construct_model(self, obj, skew_dims=None):
             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.classifier_conjugate_model = None
             obj.classifier = GenericDeterministicModel(f=lambda x: torch.ones(size=x.size())[..., -1])
@@ -891,3 +902,7 @@ def plot_validity(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_dims(self):
+        return {obj.name: obj.model.covar_module.latent_transform for obj in self.active_objs}
diff --git a/blop/dofs.py b/blop/dofs.py
index 6a68d83..363856c 100644
--- a/blop/dofs.py
+++ b/blop/dofs.py
@@ -100,9 +100,6 @@ def __post_init__(self):
         else:
             raise ValueError("DOF() accepts exactly one of either a name or an ophyd device.")
 
-        if self.description is None:
-            self.description = self.name
-
         if not self.read_only:
             # check that the device has a put method
             if isinstance(self.device, SignalRO):
@@ -110,7 +107,7 @@ def __post_init__(self):
 
         if self.log:
             if not self.search_lower_bound > 0:
-                raise ValueError("Search bounds must be positive if log=True.")
+                raise ValueError("Search bounds must be strictly positive if log=True.")
 
         # all dof degrees of freedom are hinted
         self.device.kind = "hinted"
diff --git a/blop/objectives.py b/blop/objectives.py
index 1be23a5..9361698 100644
--- a/blop/objectives.py
+++ b/blop/objectives.py
@@ -1,4 +1,4 @@
-from collections.abc import Sequence
+from collections.abc import Iterable, Sequence
 from dataclasses import dataclass, field
 from typing import List, Tuple, Union
 
@@ -146,18 +146,27 @@ def __init__(self, objectives: list = []):
         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)
-        else:
-            raise AttributeError(f"No attribute named {attr}")
 
-    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)]
+        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 {i}. An ObjectiveList must be indexed by either an integer or a string.")
+            raise ValueError(f"Invalid index {key}.")
 
     def __len__(self):
         return len(self.objectives)

From 5c3003c4803d5cc984d04c63b97df4acc629552f Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Tue, 30 Jan 2024 12:01:17 -0500
Subject: [PATCH 13/20] model outcome constraints for targeted objectives

---
 blop/agent.py                         |  30 ++++---
 blop/bayesian/acquisition/__init__.py |  16 ++--
 blop/bayesian/models.py               |   2 +-
 blop/bayesian/plotting.py             |   4 +-
 blop/bayesian/transforms.py           |  55 +++++++++----
 blop/dofs.py                          |   6 +-
 blop/objectives.py                    |  35 +++++---
 docs/source/dofs.rst                  | 111 ++++++++++++++++++++++++++
 docs/source/objectives.rst            | 109 +++++++++++++++++++++++++
 9 files changed, 318 insertions(+), 50 deletions(-)
 create mode 100644 docs/source/dofs.rst
 create mode 100644 docs/source/objectives.rst

diff --git a/blop/agent.py b/blop/agent.py
index df97a71..c9792d7 100644
--- a/blop/agent.py
+++ b/blop/agent.py
@@ -596,15 +596,15 @@ def _construct_model(self, obj, skew_dims=None):
         obj.model_dofs = set(self.active_dofs.names)  # if these change, retrain the model on self.ask()
 
         if trusted.all():
-            obj.classifier_conjugate_model = None
-            obj.classifier = GenericDeterministicModel(f=lambda x: torch.ones(size=x.size())[..., -1])
+            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.classifier_conjugate_model = models.LatentDirichletClassifier(
+            obj.validity_conjugate_model = models.LatentDirichletModel(
                 train_inputs=train_inputs,
                 train_targets=dirichlet_likelihood.transformed_targets.transpose(-1, -2).double(),
                 skew_dims=skew_dims,
@@ -612,7 +612,7 @@ def _construct_model(self, obj, skew_dims=None):
                 input_transform=self._model_input_transform,
             )
 
-            obj.classifier = GenericDeterministicModel(f=lambda x: obj.classifier_conjugate_model.probabilities(x)[..., -1])
+            obj.validity_ = GenericDeterministicModel(f=lambda x: obj.validity_conjugate_model.probabilities(x)[..., -1])
 
     def _construct_all_models(self):
         """Construct a model for each objective."""
@@ -624,8 +624,8 @@ def _train_all_models(self, **kwargs):
         t0 = ttime.monotonic()
         for obj in self.active_objs:
             self._train_model(obj.model)
-            if obj.classifier_conjugate_model is not None:
-                self._train_model(obj.classifier_conjugate_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")
@@ -733,15 +733,19 @@ def forget(self, last=None, index=None, train=True):
     def _set_hypers(self, hypers):
         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"])
 
     @property
-    def classifier(self):
+    def constraint(self):
         def f(x):
             p = torch.ones(x.shape[:-1])
             for obj in self.active_objs:
-                if obj.classifier_conjugate_model is not None:
-                    p *= obj.classifier(x)
+                # 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
 
         return GenericDeterministicModel(f=f)
@@ -749,9 +753,9 @@ def f(x):
     @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
+        hypers = {"validity_constraint": {}}
+        for key, value in self.validity_constraint.state_dict().items():
+            hypers["validity_constraint"][key] = value
         for obj in self.active_objs:
             hypers[obj.name] = {}
             for key, value in obj.model.state_dict().items():
diff --git a/blop/bayesian/acquisition/__init__.py b/blop/bayesian/acquisition/__init__.py
index 6740e65..24bdacf 100644
--- a/blop/bayesian/acquisition/__init__.py
+++ b/blop/bayesian/acquisition/__init__.py
@@ -41,7 +41,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
     # there is probably a better way to structure this
     if acq_func_name == "expected_improvement":
         acq_func = analytic.ConstrainedLogExpectedImprovement(
-            constraint=agent.classifier,
+            constraint=agent.constraint,
             model=agent.model,
             best_f=agent.max_scalarized_objective,
             posterior_transform=agent.targeting_transform,
@@ -50,7 +50,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
 
     elif acq_func_name == "monte_carlo_expected_improvement":
         acq_func = monte_carlo.qConstrainedExpectedImprovement(
-            constraint=agent.classifier,
+            constraint=agent.constraint,
             model=agent.model,
             best_f=agent.max_scalarized_objective,
             posterior_transform=agent.targeting_transform,
@@ -59,7 +59,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
 
     elif acq_func_name == "probability_of_improvement":
         acq_func = analytic.ConstrainedLogProbabilityOfImprovement(
-            constraint=agent.classifier,
+            constraint=agent.constraint,
             model=agent.model,
             best_f=agent.max_scalarized_objective,
             posterior_transform=agent.targeting_transform,
@@ -68,7 +68,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
 
     elif acq_func_name == "monte_carlo_probability_of_improvement":
         acq_func = monte_carlo.qConstrainedProbabilityOfImprovement(
-            constraint=agent.classifier,
+            constraint=agent.constraint,
             model=agent.model,
             best_f=agent.max_scalarized_objective,
             posterior_transform=agent.targeting_transform,
@@ -77,7 +77,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
 
     elif acq_func_name == "lower_bound_max_value_entropy":
         acq_func = monte_carlo.qConstrainedLowerBoundMaxValueEntropy(
-            constraint=agent.classifier,
+            constraint=agent.constraint,
             model=agent.model,
             candidate_set=agent.test_inputs(n=1024).squeeze(1),
         )
@@ -85,7 +85,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
 
     elif acq_func_name == "noisy_expected_hypervolume_improvement":
         acq_func = monte_carlo.qConstrainedNoisyExpectedHypervolumeImprovement(
-            constraint=agent.classifier,
+            constraint=agent.constraint,
             model=agent.model,
             ref_point=agent.train_targets.min(dim=0).values,
             X_baseline=agent.train_inputs,
@@ -97,7 +97,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
         beta = acq_func_kwargs.get("beta", acq_func_config["default_args"]["beta"])
 
         acq_func = analytic.ConstrainedUpperConfidenceBound(
-            constraint=agent.classifier,
+            constraint=agent.constraint,
             model=agent.model,
             beta=beta,
             posterior_transform=agent.targeting_transform,
@@ -108,7 +108,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
         beta = acq_func_kwargs.get("beta", acq_func_config["default_args"]["beta"])
 
         acq_func = monte_carlo.qConstrainedUpperConfidenceBound(
-            constraint=agent.classifier,
+            constraint=agent.constraint,
             model=agent.model,
             beta=beta,
             posterior_transform=agent.targeting_transform,
diff --git a/blop/bayesian/models.py b/blop/bayesian/models.py
index 9dd15cc..c0d165b 100644
--- a/blop/bayesian/models.py
+++ b/blop/bayesian/models.py
@@ -23,7 +23,7 @@ def __init__(self, train_inputs, train_targets, skew_dims=True, *args, **kwargs)
         self.trained = False
 
 
-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)
 
diff --git a/blop/bayesian/plotting.py b/blop/bayesian/plotting.py
index a9968f0..a1c8917 100644
--- a/blop/bayesian/plotting.py
+++ b/blop/bayesian/plotting.py
@@ -286,7 +286,7 @@ def _plot_valid_one_dof(agent, size=16, lw=1e0):
     x_values = agent.table.loc[:, x_dof.device.name].values
 
     test_inputs = agent.sample(method="grid")
-    constraint = agent.classifier(test_inputs)[..., 0]
+    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)
@@ -308,7 +308,7 @@ def _plot_valid_many_dofs(agent, axes=[0, 1], shading="nearest", cmap=DEFAULT_CO
     test_x = test_inputs[..., 0, axes[0]].detach().squeeze().numpy()
     test_y = test_inputs[..., 0, axes[1]].detach().squeeze().numpy()
 
-    constraint = agent.classifier(test_inputs)[..., 0].squeeze().numpy()
+    constraint = agent.constraint(test_inputs)[..., 0].squeeze().numpy()
 
     if gridded:
         _ = agent.valid_axes[1].pcolormesh(
diff --git a/blop/bayesian/transforms.py b/blop/bayesian/transforms.py
index 09b718f..9423cae 100644
--- a/blop/bayesian/transforms.py
+++ b/blop/bayesian/transforms.py
@@ -1,13 +1,17 @@
 from typing import Union
 
 import botorch
+import numpy as np
+import torch
 from botorch.acquisition.objective import PosteriorTransform
 from botorch.posteriors.gpytorch import GPyTorchPosterior
 from botorch.posteriors.posterior_list import PosteriorList
-from torch import Tensor
+from torch.special import erf
 
+sqrt2 = np.sqrt(2)
 
-def targeting_transform(y, target):
+
+def targeting_sample_transform(y: torch.Tensor, target) -> torch.Tensor:
     if target == "max":
         return y
     if target == "min":
@@ -15,7 +19,28 @@ def targeting_transform(y, target):
     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)
+        return y * 0  # torch.where((y > target[0]) & (y < target[1]), 0, -np.inf)
+
+
+def targeting_mean_transform(mean: torch.Tensor, variance: torch.Tensor, target) -> torch.Tensor:
+    if target == "max":
+        return mean
+    if target == "min":
+        return -mean
+    elif not isinstance(target, tuple):
+        return -(mean - target).abs()
+    else:
+        s = variance.sqrt()
+        return torch.log(0.5 * (erf((target[1] - mean) / (sqrt2 * s)) - erf((target[0] - mean) / (sqrt2 * s))))
+    # else:
+    #    return -((mean - 0.5 * (target[1] + target[0])).abs() - 0.5 * (target[1] - target[0])).clamp(min=0)
+
+
+def targeting_variance_transform(mean: torch.Tensor, variance: torch.Tensor, target) -> torch.Tensor:
+    if isinstance(target, tuple):
+        return 0
+    else:
+        return variance
 
 
 class TargetingPosteriorTransform(PosteriorTransform):
@@ -23,10 +48,10 @@ class TargetingPosteriorTransform(PosteriorTransform):
 
     scalarize: bool = True
 
-    def __init__(self, weights: Tensor, targets: Tensor) -> None:
+    def __init__(self, weights: torch.Tensor, targets: torch.Tensor) -> None:
         r"""
         Args:
-            weights: A one-dimensional tensor with `m` elements representing the
+            weights: A one-dimensional torch.Tensor with `m` elements representing the
                 linear weights on the outputs.
             offset: An offset to be added to posterior mean.
         """
@@ -34,29 +59,31 @@ 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):
-            y[..., i] = targeting_transform(y[..., i], target)
+            y[..., i] = targeting_sample_transform(y[..., i], target)
         return y @ self.weights.unsqueeze(-1)
 
     def mean_transform(self, mean, var):
         for i, target in enumerate(self.targets):
-            mean[..., i] = targeting_transform(mean[..., i], target)
+            mean[..., i] = targeting_mean_transform(mean[..., i], var[..., i], target)
         return mean @ self.weights.unsqueeze(-1)
 
     def variance_transform(self, mean, var):
-        return var @ self.weights.unsqueeze(-1)
+        for i, target in enumerate(self.targets):
+            mean[..., i] = targeting_variance_transform(mean[..., i], var[..., i], target)
+        return mean @ self.weights.unsqueeze(-1)
 
-    def evaluate(self, Y: Tensor) -> Tensor:
+    def evaluate(self, Y: torch.Tensor) -> torch.Tensor:
         r"""Evaluate the transform on a set of outcomes.
 
         Args:
-            Y: A `batch_shape x q x m`-dim tensor of outcomes.
+            Y: A `batch_shape x q x m`-dim torch.Tensor of outcomes.
 
         Returns:
-            A `batch_shape x q`-dim tensor of transformed outcomes.
+            A `batch_shape x q`-dim torch.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.
@@ -71,7 +98,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/dofs.py b/blop/dofs.py
index 363856c..9cca0fb 100644
--- a/blop/dofs.py
+++ b/blop/dofs.py
@@ -85,6 +85,10 @@ def __post_init__(self):
                 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)
@@ -106,7 +110,7 @@ def __post_init__(self):
                 raise ValueError("You must specify read_only=True for a read-only device.")
 
         if self.log:
-            if not self.search_lower_bound > 0:
+            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
diff --git a/blop/objectives.py b/blop/objectives.py
index 9361698..129aece 100644
--- a/blop/objectives.py
+++ b/blop/objectives.py
@@ -4,6 +4,8 @@
 
 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
@@ -84,10 +86,16 @@ class Objective:
     latent_groups: List[Tuple[str, ...]] = field(default_factory=list)
 
     def __post_init__(self):
-        if type(self.target) is str:
+        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:
@@ -95,7 +103,7 @@ def _trust_bounds(self):
         return self.trust_bounds
 
     @property
-    def label(self):
+    def label(self) -> str:
         return f"{'log ' if self.log else ''}{self.description}"
 
     @property
@@ -109,12 +117,6 @@ def summary(self) -> pd.Series:
             series[attr] = value
         return series
 
-    def __repr__(self):
-        return self.summary.__repr__()
-
-    def __repr_html__(self):
-        return self.summary.__repr_html__()
-
     @property
     def trust_lower_bound(self):
         if self.trust_bounds is None:
@@ -128,17 +130,28 @@ def trust_upper_bound(self):
         return float(self.trust_bounds[1])
 
     @property
-    def noise(self):
+    def noise(self) -> float:
         return self.model.likelihood.noise.item() if hasattr(self, "model") else None
 
     @property
-    def snr(self):
+    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):
+    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 = []):
diff --git a/docs/source/dofs.rst b/docs/source/dofs.rst
new file mode 100644
index 0000000..9baf23f
--- /dev/null
+++ b/docs/source/dofs.rst
@@ -0,0 +1,111 @@
+Degrees of freedom (DOFs)
++++++++++++++++++++++++++
+
+A degree of freedom is how
+
+.. code-block:: python
+
+    from blop import DOF
+
+    dof = DOF(name="x1", search_bounds=(lower_limit, upper_limit))
+
+
+
+
+
+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.
+
+#. Find the most interesting point (or points) to sample, and move the degrees of freedom there.
+#. For each point, run an acquisition plan (e.g., trigger and read the detectors).
+#. Digest the results of the acquisition to find the value of the task.
diff --git a/docs/source/objectives.rst b/docs/source/objectives.rst
new file mode 100644
index 0000000..82e9d10
--- /dev/null
+++ b/docs/source/objectives.rst
@@ -0,0 +1,109 @@
+=====
+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.
+
+#. 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.

From a70195527f65e53de46ae1ca000c74a19be151b1 Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Tue, 30 Jan 2024 12:47:24 -0500
Subject: [PATCH 14/20] fix writing hyperparameters

---
 blop/agent.py               | 52 ++++++++++++++++++++++++-------------
 blop/bayesian/transforms.py | 49 ++++++++--------------------------
 2 files changed, 45 insertions(+), 56 deletions(-)

diff --git a/blop/agent.py b/blop/agent.py
index c9792d7..f597102 100644
--- a/blop/agent.py
+++ b/blop/agent.py
@@ -612,7 +612,9 @@ def _construct_model(self, obj, skew_dims=None):
                 input_transform=self._model_input_transform,
             )
 
-            obj.validity_ = GenericDeterministicModel(f=lambda x: obj.validity_conjugate_model.probabilities(x)[..., -1])
+            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."""
@@ -751,15 +753,18 @@ def f(x):
         return GenericDeterministicModel(f=f)
 
     @property
-    def hypers(self):
-        """Returns a dict of all the hyperparameters in all the agent's models."""
-        hypers = {"validity_constraint": {}}
-        for key, value in self.validity_constraint.state_dict().items():
-            hypers["validity_constraint"][key] = value
-        for obj in self.active_objs:
-            hypers[obj.name] = {}
+    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] = {"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
 
@@ -767,20 +772,31 @@ 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[()]))
+            for obj_name in f.keys():
+                hypers[obj_name] = {"model": OrderedDict(), "validity_conjugate_model": OrderedDict()}
+
+                for key, value in f[obj_name]["model"].items():
+                    hypers[obj_name]["model"][key] = torch.tensor(np.atleast_1d(value[()]))
+
+                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
diff --git a/blop/bayesian/transforms.py b/blop/bayesian/transforms.py
index 9423cae..82bdb9d 100644
--- a/blop/bayesian/transforms.py
+++ b/blop/bayesian/transforms.py
@@ -1,17 +1,13 @@
 from typing import Union
 
 import botorch
-import numpy as np
-import torch
 from botorch.acquisition.objective import PosteriorTransform
 from botorch.posteriors.gpytorch import GPyTorchPosterior
 from botorch.posteriors.posterior_list import PosteriorList
-from torch.special import erf
+from torch import Tensor
 
-sqrt2 = np.sqrt(2)
 
-
-def targeting_sample_transform(y: torch.Tensor, target) -> torch.Tensor:
+def targeting_transform(y, target):
     if target == "max":
         return y
     if target == "min":
@@ -19,28 +15,7 @@ def targeting_sample_transform(y: torch.Tensor, target) -> torch.Tensor:
     elif not isinstance(target, tuple):
         return -(y - target).abs()
     else:
-        return y * 0  # torch.where((y > target[0]) & (y < target[1]), 0, -np.inf)
-
-
-def targeting_mean_transform(mean: torch.Tensor, variance: torch.Tensor, target) -> torch.Tensor:
-    if target == "max":
-        return mean
-    if target == "min":
-        return -mean
-    elif not isinstance(target, tuple):
-        return -(mean - target).abs()
-    else:
-        s = variance.sqrt()
-        return torch.log(0.5 * (erf((target[1] - mean) / (sqrt2 * s)) - erf((target[0] - mean) / (sqrt2 * s))))
-    # else:
-    #    return -((mean - 0.5 * (target[1] + target[0])).abs() - 0.5 * (target[1] - target[0])).clamp(min=0)
-
-
-def targeting_variance_transform(mean: torch.Tensor, variance: torch.Tensor, target) -> torch.Tensor:
-    if isinstance(target, tuple):
-        return 0
-    else:
-        return variance
+        return -((y - 0.5 * (target[1] + target[0])).abs() - 0.5 * (target[1] - target[0])).clamp(min=0)
 
 
 class TargetingPosteriorTransform(PosteriorTransform):
@@ -48,10 +23,10 @@ class TargetingPosteriorTransform(PosteriorTransform):
 
     scalarize: bool = True
 
-    def __init__(self, weights: torch.Tensor, targets: torch.Tensor) -> None:
+    def __init__(self, weights: Tensor, targets: Tensor) -> None:
         r"""
         Args:
-            weights: A one-dimensional torch.Tensor with `m` elements representing the
+            weights: A one-dimensional tensor with `m` elements representing the
                 linear weights on the outputs.
             offset: An offset to be added to posterior mean.
         """
@@ -61,27 +36,25 @@ def __init__(self, weights: torch.Tensor, targets: torch.Tensor) -> None:
 
     def sample_transform(self, y):
         for i, target in enumerate(self.targets):
-            y[..., i] = targeting_sample_transform(y[..., i], target)
+            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):
-            mean[..., i] = targeting_mean_transform(mean[..., i], var[..., i], target)
+            mean[..., i] = targeting_transform(mean[..., i], target)
         return mean @ self.weights.unsqueeze(-1)
 
     def variance_transform(self, mean, var):
-        for i, target in enumerate(self.targets):
-            mean[..., i] = targeting_variance_transform(mean[..., i], var[..., i], target)
-        return mean @ self.weights.unsqueeze(-1)
+        return var @ self.weights.unsqueeze(-1)
 
-    def evaluate(self, Y: torch.Tensor) -> torch.Tensor:
+    def evaluate(self, Y: Tensor) -> Tensor:
         r"""Evaluate the transform on a set of outcomes.
 
         Args:
-            Y: A `batch_shape x q x m`-dim torch.Tensor of outcomes.
+            Y: A `batch_shape x q x m`-dim tensor of outcomes.
 
         Returns:
-            A `batch_shape x q`-dim torch.Tensor of transformed outcomes.
+            A `batch_shape x q`-dim tensor of transformed outcomes.
         """
         return self.sample_transform(Y)
 

From 63d8e5e794fcdde1a609e35d3f1fee5d8d5c6a6a Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Tue, 30 Jan 2024 14:03:32 -0500
Subject: [PATCH 15/20] fix docs

---
 blop/agent.py                            |  28 +++---
 docs/source/agent.rst                    |  35 ++++++++
 docs/source/dofs.rst                     | 104 +++------------------
 docs/source/objectives.rst               | 107 ++++------------------
 docs/source/tutorials/himmelblau.ipynb   |   6 +-
 docs/source/tutorials/passive-dofs.ipynb |  10 ---
 docs/source/usage.rst                    | 110 ++---------------------
 7 files changed, 84 insertions(+), 316 deletions(-)
 create mode 100644 docs/source/agent.rst

diff --git a/blop/agent.py b/blop/agent.py
index f597102..973185c 100644
--- a/blop/agent.py
+++ b/blop/agent.py
@@ -236,7 +236,7 @@ def ask(self, acq_func_identifier="qei", n=1, route=True, sequential=True, upsam
             )
 
             NUM_RESTARTS = 8
-            RAW_SAMPLES = 1024
+            RAW_SAMPLES = 256
 
             candidates, acqf_obj = botorch.optim.optimize_acqf(
                 acq_function=acq_func,
@@ -737,20 +737,16 @@ def _set_hypers(self, hypers):
             obj.model.load_state_dict(hypers[obj.name])
         self.validity_constraint.load_state_dict(hypers["validity_constraint"])
 
-    @property
-    def constraint(self):
-        def f(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
-
-        return GenericDeterministicModel(f=f)
+    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) -> dict:
@@ -924,5 +920,5 @@ def plot_history(self, **kwargs):
         plotting._plot_history(self, **kwargs)
 
     @property
-    def latent_dims(self):
+    def latent_transforms(self):
         return {obj.name: obj.model.covar_module.latent_transform for obj in self.active_objs}
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
index 9baf23f..42ec90a 100644
--- a/docs/source/dofs.rst
+++ b/docs/source/dofs.rst
@@ -1,111 +1,33 @@
 Degrees of freedom (DOFs)
 +++++++++++++++++++++++++
 
-A degree of freedom is how
+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", search_bounds=(lower_limit, upper_limit))
+    dof = DOF(name="x1", description="my first DOF", search_bounds=(lower, upper))
 
-
-
-
-
-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.
+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
 
-    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.
-
-
+    from blop import DOF
 
+    dof = DOF(device=my_ophyd_device, description="a real piece of hardware", search_bounds=(lower, upper))
 
-Building the agent
-++++++++++++++++++
+In this case, the agent will control the device as it sees fit, moving it between the search bounds.
 
-Combining these with a databroker instance will construct an agent.
+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
 
-    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))
-
+    from blop import DOF
 
-In the example below, the agent will loop over the following steps in each iteration of learning.
+    dof = DOF(device=a_read_only_ophyd_device, description="a thermometer or something", read_only=True, trust_bounds=(lower, upper))
 
-#. 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.
+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
index 82e9d10..69e489e 100644
--- a/docs/source/objectives.rst
+++ b/docs/source/objectives.rst
@@ -1,109 +1,34 @@
-=====
-Usage
-=====
+Objectives
+++++++++++
 
-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.
+We can describe an optimization problem as a list of objectives to. A simple objective is
 
 .. 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.
+    from blop import Objective
 
+    objective = Objective(name="y1", target="max")
 
-
-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:
+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
 
-    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
+    from blop import Objective
 
-            entry.loc[index, "thing_to_maximize"] = some_fitness_function(raw_output_1, raw_output_2)
+    objective = Objective(name="y1", target="min") # minimize the quantity "y1"
 
-        return products
+    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.
 
-Detectors
-+++++++++
 
-Detectors are triggered for each input.
+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
 
-    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.
+    from blop import Objective
 
-#. 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.
+    objective = Objective(name="some_strictly_positive_quantity", target="max", log=True)
diff --git a/docs/source/tutorials/himmelblau.ipynb b/docs/source/tutorials/himmelblau.ipynb
index c871dc6..8b1a435 100644
--- a/docs/source/tutorials/himmelblau.ipynb
+++ b/docs/source/tutorials/himmelblau.ipynb
@@ -294,7 +294,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "Python 3.11.5 64-bit",
    "language": "python",
    "name": "python3"
   },
@@ -308,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/passive-dofs.ipynb b/docs/source/tutorials/passive-dofs.ipynb
index 4c3f4fd..42a0948 100644
--- a/docs/source/tutorials/passive-dofs.ipynb
+++ b/docs/source/tutorials/passive-dofs.ipynb
@@ -76,16 +76,6 @@
    "source": [
     "agent.plot_objectives()"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "312610b1",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "agent.latent_dim_tuples()"
-   ]
   }
  ],
  "metadata": {
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

From af7b501b478ca4effea7239abf9a58a706584811 Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Tue, 30 Jan 2024 14:14:10 -0500
Subject: [PATCH 16/20] update models on tell is optional

---
 blop/agent.py | 34 ++++++++++++++++++----------------
 1 file changed, 18 insertions(+), 16 deletions(-)

diff --git a/blop/agent.py b/blop/agent.py
index 973185c..9d19363 100644
--- a/blop/agent.py
+++ b/blop/agent.py
@@ -287,6 +287,7 @@ def tell(
         y: Optional[Mapping] = {},
         metadata: Optional[Mapping] = {},
         append: bool = True,
+        update_models: bool = True,
         train: bool = None,
     ):
         """
@@ -323,26 +324,27 @@ def tell(
         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.active_objs:
-            t0 = ttime.monotonic()
+        if update_models:
+            for obj in self.active_objs:
+                t0 = ttime.monotonic()
 
-            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
+                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 train is None:
-                train = int(n_after_tell / self.train_every) > int(n_before_tell / self.train_every)
+                if train is None:
+                    train = int(n_after_tell / self.train_every) > int(n_before_tell / self.train_every)
 
-            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")
+                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")
 
-                else:
-                    self._train_model(obj.model, hypers=cached_hypers)
+                    else:
+                        self._train_model(obj.model, hypers=cached_hypers)
 
     def learn(
         self,

From 848148eff10252d7c1fa2b18510e7ec8a090abc3 Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Tue, 30 Jan 2024 15:01:28 -0500
Subject: [PATCH 17/20] exclude invalid inputs from validity constraint

---
 blop/agent.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/blop/agent.py b/blop/agent.py
index 9d19363..660731e 100644
--- a/blop/agent.py
+++ b/blop/agent.py
@@ -584,7 +584,10 @@ def _construct_model(self, obj, skew_dims=None):
         train_inputs = self.train_inputs(active=True)
         train_targets = self.train_targets(obj.name)
 
-        trusted = ~(torch.isnan(train_inputs).any(axis=1) | torch.isnan(train_targets).any(axis=1))
+        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],
@@ -607,8 +610,8 @@ def _construct_model(self, obj, skew_dims=None):
             )
 
             obj.validity_conjugate_model = models.LatentDirichletModel(
-                train_inputs=train_inputs,
-                train_targets=dirichlet_likelihood.transformed_targets.transpose(-1, -2).double(),
+                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,

From aac25dfe51d9dff882dc34e86c2c9e6cbcce8531 Mon Sep 17 00:00:00 2001
From: Thomas Morris <41275226+thomaswmorris@users.noreply.github.com>
Date: Tue, 30 Jan 2024 18:23:08 -0500
Subject: [PATCH 18/20] Update blop/dofs.py

Co-authored-by: Max Rakitin <mrakitin@users.noreply.github.com>
---
 blop/dofs.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/blop/dofs.py b/blop/dofs.py
index 9cca0fb..7e2ae6f 100644
--- a/blop/dofs.py
+++ b/blop/dofs.py
@@ -52,7 +52,7 @@ class DOF:
     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. Much be larger than search bounds.
+        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.

From 535041f1bbfd8fdaccf68eef6f8fea42ad90560a Mon Sep 17 00:00:00 2001
From: Thomas Morris <41275226+thomaswmorris@users.noreply.github.com>
Date: Tue, 30 Jan 2024 18:23:18 -0500
Subject: [PATCH 19/20] Update blop/objectives.py

Co-authored-by: Max Rakitin <mrakitin@users.noreply.github.com>
---
 blop/objectives.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/blop/objectives.py b/blop/objectives.py
index 129aece..8c29d73 100644
--- a/blop/objectives.py
+++ b/blop/objectives.py
@@ -26,8 +26,6 @@
     "latent_groups": "object",
 }
 
-OBJ_SUMMARY_STYLE = {"min_noise": "{:.2E}", "max_noise": "{:.2E}"}
-
 
 class DuplicateNameError(ValueError):
     ...

From 0af49dc8712eadc9bc46242ccf03cc846e800620 Mon Sep 17 00:00:00 2001
From: Thomas Morris <41275226+thomaswmorris@users.noreply.github.com>
Date: Tue, 30 Jan 2024 18:27:25 -0500
Subject: [PATCH 20/20] Update release-history.rst

---
 docs/source/release-history.rst | 5 +++++
 1 file changed, 5 insertions(+)

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.