From fa6095bcee98ee88f60b3a043e5158d84f302d4f Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Sun, 14 Apr 2024 22:51:01 -0400
Subject: [PATCH 01/11] better control of DOFs and objectives

---
 src/blop/dofs.py       |  12 ++++
 src/blop/objectives.py | 124 ++++++++++++++++++++++-------------------
 2 files changed, 78 insertions(+), 58 deletions(-)

diff --git a/src/blop/dofs.py b/src/blop/dofs.py
index 5d3b90c..42c37fc 100644
--- a/src/blop/dofs.py
+++ b/src/blop/dofs.py
@@ -221,6 +221,18 @@ def __post_init__(self):
         # all dof degrees of freedom are hinted
         self.device.kind = "hinted"
 
+    @property
+    def domain(self):
+        """
+        The total domain of the DOF.
+        """
+        if self.transform is None:
+            if self.type == "continuous":
+                return (-np.inf, np.inf)
+            else:
+                return self.search_domain
+        return SUPPORTED_DOF_TRANSFORMS[self.transform]
+
     @property
     def _search_domain(self):
         """
diff --git a/src/blop/objectives.py b/src/blop/objectives.py
index f99e22b..8c710bd 100644
--- a/src/blop/objectives.py
+++ b/src/blop/objectives.py
@@ -212,6 +212,7 @@ def summary(self) -> pd.Series:
             series[attr] = value if value is not None else ""
         return series
 
+>>>>>>> 319737d (better control of DOFs and objectives)
     @property
     def noise(self) -> float:
         return self.model.likelihood.noise.item() if hasattr(self, "model") else np.nan
@@ -247,35 +248,35 @@ def targeting_constraint(self, x: torch.Tensor) -> torch.Tensor:
     #         f = -f
     #     return f
 
-    # def fitness_inverse(self, f):
-    #     y = f
-    #     if self.target == "min":
-    #         y = -y
-    #     if self.log:
-    #         y = np.exp(y)
-    #     return y
+    def fitness_inverse(self, f):
+        y = f
+        if self.target == "min":
+            y = -y
+        if self.log:
+            y = np.exp(y)
+        return y
 
-    # @property
-    # def is_fitness(self):
-    #     return self.target in ["min", "max"]
+    @property
+    def is_fitness(self):
+        return self.target in ["min", "max"]
 
-    # def value_prediction(self, X):
-    #     p = self.model.posterior(X)
+    def value_prediction(self, X):
+        p = self.model.posterior(X)
 
-    #     if self.is_fitness:
-    #         return self.fitness_inverse(p.mean)
+        if self.is_fitness:
+            return self.fitness_inverse(p.mean)
 
-    #     if isinstance(self.target, tuple):
-    #         return p.mean
+        if isinstance(self.target, tuple):
+            return p.mean
 
-    # def fitness_prediction(self, X):
-    #     p = self.model.posterior(X)
+    def fitness_prediction(self, X):
+        p = self.model.posterior(X)
 
-    #     if self.is_fitness:
-    #         return self.fitness_inverse(p.mean)
+        if self.is_fitness:
+            return self.fitness_inverse(p.mean)
 
-    #     if isinstance(self.target, tuple):
-    #         return self.targeting_constraint(X).log().clamp(min=-16)
+        if isinstance(self.target, tuple):
+            return self.targeting_constraint(X).log().clamp(min=-16)
 
 
 class ObjectiveList(Sequence):
@@ -330,42 +331,49 @@ def __repr__(self):
         return self.summary.__repr__()
 
     def _repr_html_(self):
-        return self.summary.T._repr_html_()
-
-    # @property
-    # def descriptions(self) -> list:
-    #     """
-    #     Returns an array of the objective names.
-    #     """
-    #     return [obj.description for obj in self.objectives]
-
-    # @property
-    # def names(self) -> list:
-    #     """
-    #     Returns an array of the objective names.
-    #     """
-    #     return [obj.name for obj in self.objectives]
-
-    # @property
-    # def targets(self) -> list:
-    #     """
-    #     Returns an array of the objective targets.
-    #     """
-    #     return [obj.target for obj in self.objectives]
-
-    # @property
-    # def weights(self) -> np.array:
-    #     """
-    #     Returns an array of the objective weights.
-    #     """
-    #     return np.array([obj.weight for obj in self.objectives])
-
-    # @property
-    # def signed_weights(self) -> np.array:
-    #     """
-    #     Returns a signed array of the objective weights.
-    #     """
-    #     return np.array([(1 if obj.target == "max" else -1) * obj.weight for obj in self.objectives])
+        return self.summary._repr_html_()
+
+    @property
+    def descriptions(self) -> list:
+        """
+        Returns an array of the objective names.
+        """
+        return [obj.description for obj in self.objectives]
+
+    @property
+    def names(self) -> list:
+        """
+        Returns an array of the objective names.
+        """
+        return [obj.name for obj in self.objectives]
+
+    @property
+    def targets(self) -> list:
+        """
+        Returns an array of the objective targets.
+        """
+        return [obj.target for obj in self.objectives]
+
+    @property
+    def weights(self) -> np.array:
+        """
+        Returns an array of the objective weights.
+        """
+        return np.array([obj.weight for obj in self.objectives])
+
+    @property
+    def is_fitness(self) -> np.array:
+        """
+        Returns an array of the objective weights.
+        """
+        return np.array([obj.target in ["min", "max"] for obj in self.objectives])
+
+    @property
+    def signed_weights(self) -> np.array:
+        """
+        Returns a signed array of the objective weights.
+        """
+        return np.array([(1 if obj.target == "max" else -1) * obj.weight for obj in self.objectives])
 
     def add(self, objective):
         _validate_objs([*self.objectives, objective])

From 6e65ec9040aa8ba9d4bd6599e45d92ad2bdb3b24 Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Thu, 18 Apr 2024 11:38:32 -0700
Subject: [PATCH 02/11] saving

---
 src/blop/agent.py                         | 201 +++++++++++-----------
 src/blop/bayesian/acquisition/__init__.py | 125 +++++++-------
 src/blop/bayesian/acquisition/config.yml  |  55 ++----
 src/blop/dofs.py                          |  19 +-
 src/blop/objectives.py                    |   5 +-
 src/blop/plotting.py                      |  61 +++----
 src/blop/tests/conftest.py                |  36 +++-
 src/blop/tests/test_acq_funcs.py          |  44 +++--
 src/blop/tests/test_agent.py              |  22 +--
 src/blop/tests/test_pareto.py             |  37 ++--
 src/blop/tests/test_passive_dofs.py       |   7 +-
 src/blop/tests/test_plots.py              |  24 ++-
 12 files changed, 345 insertions(+), 291 deletions(-)

diff --git a/src/blop/agent.py b/src/blop/agent.py
index 0a170cd..4fe4d93 100644
--- a/src/blop/agent.py
+++ b/src/blop/agent.py
@@ -28,6 +28,7 @@
 
 from . import plotting, utils
 from .bayesian import acquisition, models
+from .bayesian.acquisition import parse_acqf_identifier, _construct_acqf
 
 # from .bayesian.transforms import TargetingPosteriorTransform
 from .digestion import default_digestion_function
@@ -151,23 +152,15 @@ def unpack_run(self):
     def measurement_plan(self):
         return
 
-    @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]
 
     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)
-
+        acqf_config = acquisition.parse_acqf_identifier(attr)
+        if acqf_config is not None:
+            acqf, _ = _construct_acqf(acqf_name=acqf_config["name"])
+            return acqf
         raise AttributeError(f"No attribute named '{attr}'.")
 
     def sample(self, n: int = DEFAULT_MAX_SAMPLES, method: str = "quasi-random") -> torch.Tensor:
@@ -182,16 +175,18 @@ 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(self.active_dofs))
+            X = utils.normalized_sobol_sampler(n, d=len(active_dofs))
 
         elif method == "random":
-            X = torch.rand(size=(n, 1, len(self.active_dofs)))
+            X = torch.rand(size=(n, 1, len(active_dofs)))
 
         elif method == "grid":
-            n_side_if_settable = int(np.power(n, 1 / np.sum(~self.active_dofs.read_only)))
+            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 self.active_dofs
+                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()
 
@@ -200,13 +195,13 @@ def sample(self, n: int = DEFAULT_MAX_SAMPLES, method: str = "quasi-random") ->
 
         return self.dofs.subset(active=True).untransform(X)
 
-    def ask(self, acq_func_identifier="qei", n=1, route=True, sequential=True, upsample=1, **acq_func_kwargs):
+    def ask(self, acqf="qei", n=1, route=True, sequential=True, upsample=1, **acqf_kwargs):
         """Ask the agent for the best point to sample, given an acquisition function.
 
         Parameters
         ----------
-        acq_func_identifier :
-            Which acquisition function to use. Supported values can be found in `agent.all_acq_funcs`
+        acqf_identifier :
+            Which acquisition function to use. Supported values can be found in `agent.all_acqfs`
         n : int
             How many points you want
         route : bool
@@ -216,42 +211,48 @@ def ask(self, acq_func_identifier="qei", n=1, route=True, sequential=True, upsam
             finding one points and constructing a fantasy posterior about its value to generate the next point.
         """
 
-        acq_func_name = acquisition.parse_acq_func_identifier(acq_func_identifier)
-        acq_func_type = acquisition.config[acq_func_name]["type"]
+        acqf_config = parse_acqf_identifier(acqf)
+        if acqf_config is None:
+            raise ValueError(f"'{acqf}' is an invalid acquisition function.")
 
         start_time = ttime.monotonic()
 
-        if acq_func_name in ["quasi-random", "random", "grid"]:
-            candidates = self.sample(n=n, method=acq_func_name).squeeze(1).numpy()
+        active_dofs = self.dofs.subset(active=True)
+        active_objs = self.objectives.subset(active=True)
+
+        # these are the fake acquisiton functions that we don't need to construct
+        if acqf_config["name"] in ["quasi-random", "random", "grid"]:
+            candidates = self.sample(n=n, method=acqf_config["name"]).squeeze(1).numpy()
 
-            # define dummy acqf objective
-            acqf_obj = torch.zeros(len(candidates))
+            # define dummy acqf kwargs and objective
+            acqf_kwargs, acqf_obj = {}, torch.zeros(len(candidates))
 
-        elif acq_func_type in ["analytic", "monte_carlo"]:
-            if not all(hasattr(obj, "model") for obj in self.objectives):
+        else:
+            # check that all the objectives have models 
+            if not all(hasattr(obj, "model") for obj in active_objs):
                 raise RuntimeError(
-                    f"Can't construct non-trivial acquisition function '{acq_func_identifier}'"
+                    f"Can't construct non-trivial acquisition function '{acqf}'"
                     f" (the agent is not initialized!)"
                 )
 
-            for obj in self.active_objs:
-                if obj.model_dofs != set(self.active_dofs.names):
+            # if the model for any active objective mismatches the active dofs, reconstrut and train it
+            for obj in active_objs:
+                if obj.model_dofs != set(active_dofs.names):
                     self._construct_model(obj)
                     self._train_model(obj.model)
 
-            if acq_func_type == "analytic" and n > 1:
+            if acqf_config["type"] == "analytic" and n > 1:
                 raise ValueError("Can't generate multiple design points for analytic acquisition functions.")
 
-            acq_func, _ = self._get_acquisition_function(
-                identifier=acq_func_identifier, return_metadata=True, **acq_func_kwargs
-            )
+            # we may pick up some more kwargs 
+            acqf, acqf_kwargs = _construct_acqf(self, acqf_name=acqf_config["name"], **acqf_kwargs)
 
             NUM_RESTARTS = 16
             RAW_SAMPLES = 1024
 
             candidates, acqf_obj = botorch.optim.optimize_acqf(
-                acq_function=acq_func,
-                bounds=self._sample_domain,
+                acq_function=acqf,
+                bounds=self.sample_domain,
                 q=n,
                 sequential=sequential,
                 num_restarts=NUM_RESTARTS,
@@ -264,25 +265,29 @@ def ask(self, acq_func_identifier="qei", n=1, route=True, sequential=True, upsam
 
         p = self.posterior(candidates) if hasattr(self, "model") else None
 
-        acq_points = candidates[..., ~self.active_dofs.read_only]
-        read_only_values = candidates[..., self.active_dofs.read_only]
+        active_dofs = self.dofs.subset(active=True)
+
+        points = candidates[..., ~active_dofs.read_only]
+        read_only_values = candidates[..., active_dofs.read_only]
 
         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)
-            acq_points = acq_points[routing_index]
+            routing_index = utils.route(self.dofs.subset(active=True, read_only=False).readback, points)
+            points = points[routing_index]
 
         if upsample > 1:
-            idx = np.arange(len(acq_points))
+            if n == 1:
+                raise ValueError("Cannot upsample points unless n > 1.")
+            idx = np.arange(len(points))
             upsampled_idx = np.linspace(0, len(idx) - 1, upsample * len(idx) - 1)
-            acq_points = sp.interpolate.interp1d(idx, acq_points, axis=0)(upsampled_idx)
+            points = sp.interpolate.interp1d(idx, points, axis=0)(upsampled_idx)
 
         res = {
-            "points": acq_points,
-            "acq_func": acq_func_name,
-            "acq_func_kwargs": acq_func_kwargs,
-            "acq_func_obj": np.atleast_1d(acqf_obj.numpy()),
+            "points": points,
+            "acqf_name": acqf_config["name"],
+            "acqf_obj": np.atleast_1d(acqf_obj.numpy()),
+            "acqf_kwargs": acqf_kwargs,
             "duration_ms": duration,
             "sequential": sequential,
             "upsample": upsample,
@@ -337,7 +342,7 @@ def tell(
         self.table.index = np.arange(len(self.table))
 
         if update_models:
-            for obj in self.active_objs:
+            for obj in self.objectives.subset(active=True):
                 t0 = ttime.monotonic()
 
                 cached_hypers = obj.model.state_dict() if hasattr(obj, "model") else None
@@ -360,7 +365,7 @@ def tell(
 
     def learn(
         self,
-        acq_func: str = "qei",
+        acqf: str = "qei",
         n: int = 1,
         iterations: int = 1,
         upsample: int = 1,
@@ -368,7 +373,7 @@ def learn(
         append: bool = True,
         hypers: str = None,
         route: bool = True,
-        **acq_func_kwargs,
+        **acqf_kwargs,
     ):
         """This returns a Bluesky plan which iterates the learning algorithm, looping over ask -> acquire -> tell.
 
@@ -379,7 +384,7 @@ def learn(
 
         Parameters
         ----------
-        acq_func : str
+        acqf : str
             A valid identifier for an implemented acquisition function.
         n : int
             How many points to sample on each iteration.
@@ -399,7 +404,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).search_domain.mean(axis=1))
             new_table = yield from self.acquire(center_inputs)
-            new_table.loc[:, "acq_func"] = "sample_center_on_init"
+            new_table.loc[:, "acqf"] = "sample_center_on_init"
 
         for i in range(iterations):
             if self.verbose:
@@ -407,7 +412,7 @@ def learn(
             for single_acq_func in np.atleast_1d(acq_func):
                 res = self.ask(n=n, acq_func_identifier=single_acq_func, upsample=upsample, route=route, **acq_func_kwargs)
                 new_table = yield from self.acquire(res["points"])
-                new_table.loc[:, "acq_func"] = res["acq_func"]
+                new_table.loc[:, "acqf"] = res["acqf_name"]
 
                 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}
@@ -431,7 +436,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.active_objs:
+            for obj in self.objectives.subset(active=True):
                 p = obj.model.posterior(test_grid)
 
                 if item == "mean":
@@ -444,14 +449,14 @@ def view(self, item: str = "mean", cmap: str = "turbo", max_inputs: int = 2**16)
 
         else:
             try:
-                acq_func_identifier = acquisition.parse_acq_func_identifier(identifier=item)
+                acqf_identifier = acquisition.parse_acqf_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()
+            acqf, acqf_meta = self._get_acquisition_function(identifier=acqf_identifier, return_metadata=True)
+            a = acqf(test_grid).detach().numpy()
 
-            self.viewer.add_image(data=a, name=f"{acq_func_identifier}", colormap=cmap)
+            self.viewer.add_image(data=a, name=f"{acqf_identifier}", colormap=cmap)
 
         self.viewer.dims.axis_labels = self.dofs.names
 
@@ -486,7 +491,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.active_objs:
+            for obj in self.objectives.subset(active=True):
                 products.loc[:, obj.name] = np.nan
 
         if not len(acquisition_inputs) == len(products):
@@ -505,7 +510,7 @@ def reset(self):
         """Reset the agent."""
         self.table = pd.DataFrame()
 
-        for obj in self.active_objs:
+        for obj in self.objectives.subset(active=True):
             if hasattr(obj, "model"):
                 del obj.model
 
@@ -515,7 +520,7 @@ def benchmark(
         self,
         output_dir="./",
         iterations=16,
-        per_iter_learn_kwargs_list=[{"acq_func": "qr", "n": 32}, {"acq_func": "qei", "n": 4, "iterations": 4}],
+        per_iter_learn_kwargs_list=[{"acqf": "qr", "n": 32}, {"acqf": "qei", "n": 1, "iterations": 4}],
     ):
         """Iterate over having the agent learn from scratch, and save the results to an output directory.
 
@@ -540,9 +545,8 @@ def benchmark(
     @property
     def model(self):
         """A model encompassing all the fitnesses and constraints."""
-        return (
-            ModelListGP(*[obj.model for obj in self.active_objs]) if len(self.active_objs) > 1 else self.active_objs[0].model
-        )
+        active_objs = self.objectives.subset(active=True)
+        return ModelListGP(*[obj.model for obj in active_objs]) if len(active_objs) > 1 else 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."""
@@ -550,11 +554,12 @@ def posterior(self, x):
 
     @property
     def fitness_model(self):
-        return (
-            ModelListGP(*[obj.model for obj in self.objectives.subset(active=True, kind="fitness")])
-            if len(self.active_objs) > 1
-            else self.active_objs[0].model
-        )
+        active_fitness_models = self.objectives.subset(active=True, kind="fitness")
+        if len(active_fitness_models) == 0:
+            raise ValueError("Having no fitness objectives is unhandled.")
+        if len(active_fitness_models) == 1:
+            return active_fitness_models[0].model
+        return ModelListGP(*[obj.model for obj in active_fitness_models])
 
     @property
     def evaluated_constraints(self):
@@ -662,11 +667,11 @@ def _construct_model(self, obj, skew_dims=None):
             train_targets=train_targets[trusted],
             likelihood=likelihood,
             skew_dims=skew_dims,
-            input_transform=self._model_input_transform,
+            input_transform=self.input_normalization,
             outcome_transform=outcome_transform,
         )
 
-        obj.model_dofs = set(self.active_dofs.names)  # if these change, retrain the model on self.ask()
+        obj.model_dofs = set(self.dofs.subset(active=True).names)  # if these change, retrain the model on self.ask()
 
         if trusted.all():
             obj.validity_conjugate_model = None
@@ -682,7 +687,7 @@ def _construct_model(self, obj, skew_dims=None):
                 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,
+                input_transform=self.input_normalization,
             )
 
             obj.validity_constraint = GenericDeterministicModel(
@@ -691,13 +696,13 @@ def _construct_model(self, obj, skew_dims=None):
 
     def _construct_all_models(self):
         """Construct a model for each objective."""
-        for obj in self.active_objs:
+        for obj in self.objectives.subset(active=True):
             self._construct_model(obj)
 
     def _train_all_models(self, **kwargs):
         """Fit all of the agent's models. All kwargs are passed to `botorch.fit.fit_gpytorch_mll`."""
         t0 = ttime.monotonic()
-        for obj in self.active_objs:
+        for obj in self.objectives.subset(active=True):
             self._train_model(obj.model)
             if obj.validity_conjugate_model is not None:
                 self._train_model(obj.validity_conjugate_model)
@@ -709,9 +714,12 @@ def _train_all_models(self, **kwargs):
 
     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`.
+        found in `agent.all_acqfs`.
         """
-        return acquisition.get_acquisition_function(self, identifier=identifier, return_metadata=return_metadata)
+
+        acquisition._construct_acqf(self, identifier=identifier, return_metadata=return_metadata)
+
+        return 
 
     def _latent_dim_tuples(self, obj_index=None):
         """
@@ -725,7 +733,7 @@ def _latent_dim_tuples(self, obj_index=None):
         obj = self.objectives[obj_index]
 
         latent_group_index = {}
-        for dof in self.active_dofs:
+        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:
@@ -744,7 +752,7 @@ def _sample_domain(self):
         return self.dofs.subset(active=True).transform(self.dofs.subset(active=True).search_domain.T)
 
     @property
-    def _model_input_transform(self):
+    def input_normalization(self):
         """
         Suitably transforms model inputs to the unit hypercube.
 
@@ -758,7 +766,7 @@ def _model_input_transform(self):
         Read-only: constrain to the readback value
         """
 
-        return Normalize(d=len(self.active_dofs))
+        return Normalize(d=self.dofs.active.sum())
 
     def save_data(self, path="./data.h5"):
         """
@@ -797,7 +805,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.active_objs:
+        for obj in self.objectives.subset(active=True):
             obj.model.load_state_dict(hypers[obj.name])
         self.validity_constraint.load_state_dict(hypers["validity_constraint"])
 
@@ -805,14 +813,14 @@ def constraint(self, x):
         x = self.dofs.subset(active=True).transform(x)
 
         p = torch.ones(x.shape[:-1])
-        for obj in self.active_objs:
+        for obj in self.objectives.subset(active=True):
             # if the targeting constraint is non-trivial
             # if obj.kind == "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  # + 1e-6 * normalize(x, self._sample_domain).square().sum(axis=-1)
+        return p  # + 1e-6 * normalize(x, self.sample_domain).square().sum(axis=-1)
 
     @property
     def hypers(self) -> dict:
@@ -862,16 +870,11 @@ def load_hypers(filepath) -> dict:
         return hypers
 
     @property
-    def all_acq_funcs(self):
-        """Description and identifiers for all supported acquisition functions."""
-        entries = []
-        for k, d in acquisition.config.items():
-            ret = ""
-            ret += f"{d['pretty_name'].upper()} (identifiers: {d['identifiers']})\n"
-            ret += f"-> {d['description']}"
-            entries.append(ret)
-
-        print("\n\n".join(entries))
+    def all_acqfs(self):
+        """
+        Description and identifiers for all supported acquisition functions.
+        """
+        return acquisition.all_acqfs()
 
     def raw_inputs(self, index=None, **subset_kwargs):
         """
@@ -958,26 +961,30 @@ def plot_objectives(self, axes: Tuple = (0, 1), **kwargs):
         axes :
             A tuple specifying which DOFs to plot as a function of. Can be either an int or the name of DOFs.
         """
+
         if len(self.dofs.subset(active=True, read_only=False)) == 1:
-            plotting._plot_objs_one_dof(self, **kwargs)
+            if len(self.objectives.subset(active=True, kind="fitness")) > 0:
+                plotting._plot_fitness_objs_one_dof(self, **kwargs)
+            if len(self.objectives.subset(active=True, kind="constraint")) > 0:
+                plotting._plot_constraint_objs_one_dof(self, **kwargs)
         else:
             plotting._plot_objs_many_dofs(self, axes=axes, **kwargs)
 
-    def plot_acquisition(self, acq_func="ei", axes: Tuple = (0, 1), **kwargs):
+    def plot_acquisition(self, acqf="ei", axes: Tuple = (0, 1), **kwargs):
         """Plot an acquisition function over test inputs sampling the limits of the parameter space.
 
         Parameters
         ----------
-        acq_func :
+        acqf :
             Which acquisition function to plot. Can also take a list of acquisition functions.
         axes :
             A tuple specifying which DOFs to plot as a function of. Can be either an int or the name of DOFs.
         """
         if len(self.dofs.subset(active=True, read_only=False)) == 1:
-            plotting._plot_acqf_one_dof(self, acq_funcs=np.atleast_1d(acq_func), **kwargs)
+            plotting._plot_acqf_one_dof(self, acqfs=np.atleast_1d(acqf), **kwargs)
 
         else:
-            plotting._plot_acqf_many_dofs(self, acq_funcs=np.atleast_1d(acq_func), axes=axes, **kwargs)
+            plotting._plot_acqf_many_dofs(self, acqfs=np.atleast_1d(acqf), axes=axes, **kwargs)
 
     def plot_validity(self, axes: Tuple = (0, 1), **kwargs):
         """Plot the modeled constraint over test inputs sampling the limits of the parameter space.
@@ -999,7 +1006,7 @@ def plot_history(self, **kwargs):
 
     @property
     def latent_transforms(self):
-        return {obj.name: obj.model.covar_module.latent_transform for obj in self.active_objs}
+        return {obj.name: obj.model.covar_module.latent_transform for obj in self.objectives.subset(active=True)}
 
     def plot_pareto_front(self, **kwargs):
         """Plot the improvement of the agent over time."""
diff --git a/src/blop/bayesian/acquisition/__init__.py b/src/blop/bayesian/acquisition/__init__.py
index 329c933..65b9c41 100644
--- a/src/blop/bayesian/acquisition/__init__.py
+++ b/src/blop/bayesian/acquisition/__init__.py
@@ -1,6 +1,7 @@
 import os
 
 import yaml
+import pandas as pd
 from botorch.utils.transforms import normalize
 
 from . import analytic, monte_carlo
@@ -12,116 +13,116 @@
 with open(f"{here}/config.yml", "r") as f:
     config = yaml.safe_load(f)
 
-
-# supplying the full name is also a valid identifier
-for acq_func_name in config.keys():
-    config[acq_func_name]["identifiers"].append(acq_func_name)
-
-
-def parse_acq_func_identifier(identifier):
-    for acq_func_name in config.keys():
-        if identifier.lower() in config[acq_func_name]["identifiers"]:
-            return acq_func_name
+def all_acqfs(columns=["identifier", "type", "multitask_only", "description"]):
+    acqfs = pd.DataFrame(config).T[columns]
+    acqfs.index.name = "name"
+    return acqfs.sort_values(["type", "name"])
+
+def parse_acqf_identifier(identifier, strict=True):
+    for acqf_name in config.keys():
+        if identifier.lower() in [acqf_name, config[acqf_name]["identifier"]]:
+            return {"name": acqf_name, **config[acqf_name]}
+    if strict:
+        raise ValueError(f"'{identifier}' is not a valid acquisition function identifier.")
     return None
 
-
-def get_acquisition_function(agent, identifier="qei", return_metadata=True, verbose=False, **acq_func_kwargs):
+def _construct_acqf(agent, acqf_name, **acqf_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`.
+    their identifiers can be found at `agent.all_acqfs`.
     """
 
-    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"]
+    acqf_config = config["upper_confidence_bound"]
 
-    if config[acq_func_name]["multitask_only"] and (len(agent.objectives) == 1):
-        raise ValueError(f'Acquisition function "{acq_func_name}" is only for multi-task optimization problems!')
+    if config[acqf_name]["multitask_only"] and (len(agent.objectives) == 1):
+        raise ValueError(f'Acquisition function "{acqf_name}" is only for multi-task optimization problems!')
 
     # there is probably a better way to structure this
-    if acq_func_name == "expected_improvement":
-        acq_func = analytic.ConstrainedLogExpectedImprovement(
+    if acqf_name == "expected_improvement":
+
+        acqf_kwargs["best_f"] = agent.best_f(weights="default")
+
+        acqf = analytic.ConstrainedLogExpectedImprovement(
             constraint=agent.constraint,
             model=agent.fitness_model,
-            best_f=agent.best_f(),
-            posterior_transform=agent.fitness_scalarization(),
+            posterior_transform=agent.fitness_scalarization(weights="default"),
+            **acqf_kwargs,
         )
-        acq_func_meta = {"name": acq_func_name, "args": {}}
 
-    elif acq_func_name == "monte_carlo_expected_improvement":
-        acq_func = monte_carlo.qConstrainedExpectedImprovement(
+    elif acqf_name == "monte_carlo_expected_improvement":
+
+        acqf_kwargs["best_f"] = agent.best_f(weights="default")
+
+        acqf = monte_carlo.qConstrainedExpectedImprovement(
             constraint=agent.constraint,
             model=agent.fitness_model,
-            best_f=agent.best_f(),
-            posterior_transform=agent.fitness_scalarization(),
+            posterior_transform=agent.fitness_scalarization(weights="default"),
+            **acqf_kwargs,
         )
-        acq_func_meta = {"name": acq_func_name, "args": {}}
 
-    elif acq_func_name == "probability_of_improvement":
-        acq_func = analytic.ConstrainedLogProbabilityOfImprovement(
+    elif acqf_name == "probability_of_improvement":
+
+        acqf_kwargs["best_f"] = agent.best_f(weights="default")
+
+        acqf = analytic.ConstrainedLogProbabilityOfImprovement(
             constraint=agent.constraint,
             model=agent.fitness_model,
-            best_f=agent.best_f(),
             posterior_transform=agent.fitness_scalarization(),
+            **acqf_kwargs,
         )
-        acq_func_meta = {"name": acq_func_name, "args": {}}
 
-    elif acq_func_name == "monte_carlo_probability_of_improvement":
-        acq_func = monte_carlo.qConstrainedProbabilityOfImprovement(
+    elif acqf_name == "monte_carlo_probability_of_improvement":
+        acqf = monte_carlo.qConstrainedProbabilityOfImprovement(
             constraint=agent.constraint,
             model=agent.fitness_model,
             best_f=agent.best_f(),
             posterior_transform=agent.fitness_scalarization(),
         )
-        acq_func_meta = {"name": acq_func_name, "args": {}}
 
-    elif acq_func_name == "lower_bound_max_value_entropy":
-        acq_func = monte_carlo.qConstrainedLowerBoundMaxValueEntropy(
+    elif acqf_name == "lower_bound_max_value_entropy":
+        acqf = monte_carlo.qConstrainedLowerBoundMaxValueEntropy(
             constraint=agent.constraint,
             model=agent.fitness_model,
             candidate_set=agent.test_inputs(n=1024).squeeze(1),
         )
-        acq_func_meta = {"name": acq_func_name, "args": {}}
 
-    elif acq_func_name == "monte_carlo_noisy_expected_hypervolume_improvement":
-        acq_func = monte_carlo.qConstrainedNoisyExpectedHypervolumeImprovement(
+    elif acqf_name == "monte_carlo_noisy_expected_hypervolume_improvement":
+
+        acqf_kwargs["ref_point"] = acqf_kwargs.get("ref_point", agent.random_ref_point)
+
+        acqf = monte_carlo.qConstrainedNoisyExpectedHypervolumeImprovement(
             constraint=agent.constraint,
             model=agent.fitness_model,
-            ref_point=agent.random_ref_point,
-            X_baseline=normalize(agent.train_inputs(), agent._sample_domain),
+            X_baseline=agent.input_normalization.forward(agent.train_inputs()),
             prune_baseline=True,
+            **acqf_kwargs
         )
-        acq_func_meta = {"name": acq_func_name, "args": {}}
 
-    elif acq_func_name == "upper_confidence_bound":
-        beta = acq_func_kwargs.get("beta", acq_func_config["default_args"]["beta"])
+    elif acqf_name == "upper_confidence_bound":
+        acqf_kwargs["beta"] = acqf_kwargs.get("beta", acqf_config["default_kwargs"]["beta"])
 
-        acq_func = analytic.ConstrainedUpperConfidenceBound(
+        acqf = analytic.ConstrainedUpperConfidenceBound(
             constraint=agent.constraint,
             model=agent.fitness_model,
-            beta=beta,
             posterior_transform=agent.fitness_scalarization(),
+            **acqf_kwargs,
         )
-        acq_func_meta = {"name": acq_func_name, "args": {"beta": beta}}
 
-    elif acq_func_name == "monte_carlo_upper_confidence_bound":
-        beta = acq_func_kwargs.get("beta", acq_func_config["default_args"]["beta"])
+    elif acqf_name == "monte_carlo_upper_confidence_bound":
+        acqf_kwargs["beta"] = acqf_kwargs.get("beta", acqf_config["default_kwargs"]["beta"])
 
-        acq_func = monte_carlo.qConstrainedUpperConfidenceBound(
+        acqf = monte_carlo.qConstrainedUpperConfidenceBound(
             constraint=agent.constraint,
             model=agent.fitness_model,
-            beta=beta,
             posterior_transform=agent.fitness_scalarization(),
+            **acqf_kwargs,
         )
-        acq_func_meta = {"name": acq_func_name, "args": {"beta": beta}}
 
-    elif acq_func_name == "expected_mean":
-        acq_func = get_acquisition_function(agent, identifier="ucb", beta=0, return_metadata=False)
-        acq_func_meta = {"name": acq_func_name, "args": {}}
+    elif acqf_name == "expected_mean":
+        acqf, _ = _construct_acqf(agent, acqf_name="upper_confidence_bound", beta=0)
+        acqf_kwargs = {}
 
-    elif acq_func_name == "monte_carlo_expected_mean":
-        acq_func = get_acquisition_function(agent, identifier="qucb", beta=0, return_metadata=False)
-        acq_func_meta = {"name": acq_func_name, "args": {}}
+    elif acqf_name == "monte_carlo_expected_mean":
+        acqf, _ = _construct_acqf(agent, acqf_name="monte_carlo_upper_confidence_bound", beta=0)
+        acqf_kwargs = {}
 
-    return (acq_func, acq_func_meta) if return_metadata else acq_func
+    return acqf, acqf_kwargs
diff --git a/src/blop/bayesian/acquisition/config.yml b/src/blop/bayesian/acquisition/config.yml
index 512470a..91a187d 100644
--- a/src/blop/bayesian/acquisition/config.yml
+++ b/src/blop/bayesian/acquisition/config.yml
@@ -1,118 +1,101 @@
 expected_improvement:
   pretty_name: Expected improvement
   description: The expected value of max(f(x) - \nu, 0), where \nu is the current maximum.
-  identifiers:
-  - ei
+  identifier: ei
   multitask_only: false
   type: analytic
 
 monte_carlo_expected_improvement:
   description: The expected value of max(f(x) - \nu, 0), where \nu is the current maximum.
-  identifiers:
-  - qei
+  identifier: qei
   multitask_only: false
   pretty_name: Monte Carlo Expected improvement
   type: monte_carlo
 
 expected_mean:
   description: The expected value at each input.
-  identifiers:
-  - em
+  identifier: em
   multitask_only: false
   pretty_name: Expected mean
   type: analytic
 
 monte_carlo_expected_mean:
   description: The expected value at each input.
-  identifiers:
-  - qem
+  identifier: qem
   multitask_only: false
   pretty_name: Monte Carlo expected mean
   type: monte_carlo
 
 lower_bound_max_value_entropy:
   description: Max entropy search, basically
-  identifiers:
-  - lbmve
-  - lbmes
-  - gibbon
+  identifier: lbmve
   multitask_only: false
   pretty_name: Lower bound max value entropy
   type: monte_carlo
 
 noisy_expected_hypervolume_improvement:
   description: It's like a big box. How big is the box?
-  identifiers:
-  - nehvi
+  identifier: nehvi
   multitask_only: true
   pretty_name: Noisy expected hypervolume improvement
   type: analytic
 
 monte_carlo_noisy_expected_hypervolume_improvement:
   description: It's like a big box. How big is the box?
-  identifiers:
-  - qnehvi
+  identifier: qnehvi
   multitask_only: true
   pretty_name: Noisy expected hypervolume improvement
   type: monte_carlo
 
 probability_of_improvement:
   description: The probability that this input improves on the current maximum.
-  identifiers:
-  - pi
+  identifier: pi
   multitask_only: false
   pretty_name: Probability of improvement
   type: analytic
 
 monte_carlo_probability_of_improvement:
   description: The probability that this input improves on the current maximum.
-  identifiers:
-  - qpi
+  identifier: qpi
   multitask_only: false
   pretty_name: Monte Carlo probability of improvement
   type: monte_carlo
 
 random:
   description: Uniformly-sampled random points.
-  identifiers:
-  - r
+  identifier: r
   multitask_only: false
   pretty_name: Random
-  type: none
+  type: random
 
 quasi-random:
   description: Sobol-sampled quasi-random points.
-  identifiers:
-  - qr
+  identifier: qr
   multitask_only: false
   pretty_name: Quasi-random
-  type: none
+  type: random
 
 grid:
   description: A grid scan over the parameters.
-  identifiers:
-  - g
-  - gs
+  identifier: g
   multitask_only: false
   pretty_name: Grid scan
-  type: none
+  type: random
 
 upper_confidence_bound:
-  default_args:
+  default_kwargs:
     beta: 4
   description: The expected value, plus some multiple of the uncertainty (typically \mu + 2\sigma).
-  identifiers:
-  - ucb
+  identifier: ucb
   multitask_only: false
   pretty_name: Upper confidence bound
   type: analytic
 
 monte_carlo_upper_confidence_bound:
-  default_args:
+  default_kwargs:
     beta: 4
   description: The expected value, plus some multiple of the uncertainty (typically \mu + 2\sigma).
-  identifiers:
-  - qucb
+  identifier: qucb
   multitask_only: false
   pretty_name: Monte Carlo upper confidence bound
   type: monte_carlo
diff --git a/src/blop/dofs.py b/src/blop/dofs.py
index 42c37fc..ecbf484 100644
--- a/src/blop/dofs.py
+++ b/src/blop/dofs.py
@@ -351,6 +351,17 @@ def __init__(self, dofs: list = []):
         _validate_dofs(dofs)
         self.dofs = dofs
 
+    @property
+    def names(self) -> list:
+        return [dof.name for dof in self.dofs]
+
+    @property
+    def devices(self) -> list:
+        return [dof.device for dof in self.dofs]
+
+    def __call__(self, *args, **kwargs):
+        return self.subset(*args, **kwargs)
+
     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():
@@ -429,14 +440,6 @@ def summary(self) -> pd.DataFrame:
 
         return table
 
-    @property
-    def names(self) -> list:
-        return [dof.name for dof in self.dofs]
-
-    @property
-    def devices(self) -> list:
-        return [dof.device for dof in self.dofs]
-
     @property
     def search_domain(self) -> np.array:
         """
diff --git a/src/blop/objectives.py b/src/blop/objectives.py
index 8c710bd..6aacba7 100644
--- a/src/blop/objectives.py
+++ b/src/blop/objectives.py
@@ -284,11 +284,14 @@ def __init__(self, objectives: list = []):
         _validate_objs(objectives)
         self.objectives = objectives
 
+    def __call__(self, *args, **kwargs):
+        return self.subset(*args, **kwargs)
+
     @property
     def names(self):
         return [obj.name for obj in self.objectives]
 
-    def __getattr__(self, attr):
+    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(), "kind"]:
             return np.array([getattr(obj, attr) for obj in self.objectives])
diff --git a/src/blop/plotting.py b/src/blop/plotting.py
index 2f69d9d..06b3da2 100644
--- a/src/blop/plotting.py
+++ b/src/blop/plotting.py
@@ -3,7 +3,7 @@
 from matplotlib import pyplot as plt
 from matplotlib.patches import Patch
 
-from .bayesian import acquisition
+from .bayesian.acquisition import parse_acqf_identifier, _construct_acqf
 
 DEFAULT_COLOR_LIST = ["dodgerblue", "tomato", "mediumseagreen", "goldenrod"]
 # Note: the values near 1 are hard to see on a white background. Turbo goes from red to blue and isn't white in the middle.
@@ -314,11 +314,11 @@ def _plot_objs_many_dofs(agent, axes=(0, 1), shading="nearest", cmap=DEFAULT_COL
             ax.set_yscale("log")
 
 
-def _plot_acqf_one_dof(agent, acq_funcs, lw=1e0, **kwargs):
+def _plot_acqf_one_dof(agent, acqfs, lw=1e0, **kwargs):
     agent.acq_fig, agent.acq_axes = plt.subplots(
         1,
-        len(acq_funcs),
-        figsize=(4 * len(acq_funcs), 4),
+        len(acqfs),
+        figsize=(4 * len(acqfs), 4),
         sharex=True,
         constrained_layout=True,
     )
@@ -328,26 +328,26 @@ def _plot_acqf_one_dof(agent, acq_funcs, lw=1e0, **kwargs):
 
     test_inputs = agent.sample(method="grid")
 
-    for iacq_func, acq_func_identifier in enumerate(acq_funcs):
-        color = DEFAULT_COLOR_LIST[iacq_func]
+    for iacqf, acqf_identifier in enumerate(acqfs):
+        color = DEFAULT_COLOR_LIST[iacqf]
 
-        acq_func, acq_func_meta = acquisition.get_acquisition_function(agent, acq_func_identifier)
-        test_acqf = acq_func(test_inputs).detach().numpy()
+        acqf, acqf_meta = acquisition.get_acquisition_function(agent, acqf_identifier)
+        test_acqf = acqf(test_inputs).detach().numpy()
 
-        agent.acq_axes[iacq_func].plot(test_inputs.squeeze(-2), test_acqf, lw=lw, color=color)
+        agent.acq_axes[iacqf].plot(test_inputs.squeeze(-2), test_acqf, lw=lw, color=color)
 
-        agent.acq_axes[iacq_func].set_xlim(*x_dof.search_domain)
-        agent.acq_axes[iacq_func].set_xlabel(x_dof.label_with_units)
-        agent.acq_axes[iacq_func].set_ylabel(acq_func_meta["name"])
+        agent.acq_axes[iacqf].set_xlim(*x_dof.search_domain)
+        agent.acq_axes[iacqf].set_xlabel(x_dof.label_with_units)
+        agent.acq_axes[iacqf].set_ylabel(acqf_meta["name"])
 
 
 def _plot_acqf_many_dofs(
-    agent, acq_funcs, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=16, **kwargs
+    agent, acqfs, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=16, **kwargs
 ):
     agent.acq_fig, agent.acq_axes = plt.subplots(
         1,
-        len(acq_funcs),
-        figsize=(4 * len(acq_funcs), 4),
+        len(acqfs),
+        figsize=(4 * len(acqfs), 4),
         sharex=True,
         sharey=True,
         constrained_layout=True,
@@ -368,14 +368,17 @@ def _plot_acqf_many_dofs(
     test_x = test_inputs[..., 0, axes[0]].detach().squeeze().numpy()
     test_y = test_inputs[..., 0, axes[1]].detach().squeeze().numpy()
 
-    for iacq_func, acq_func_identifier in enumerate(acq_funcs):
-        acq_func, acq_func_meta = acquisition.get_acquisition_function(agent, acq_func_identifier)
+    for iacqf, acqf_identifier in enumerate(acqfs):
 
-        test_acqf = acq_func(test_inputs.reshape(-1, 1, input_dim)).detach().reshape(test_dim).squeeze().numpy()
+        acqf_config = par
+
+        acqf, acqf_meta = acquisition.get_acquisition_function(agent, acqf_identifier)
+
+        test_acqf = acqf(test_inputs.reshape(-1, 1, input_dim)).detach().reshape(test_dim).squeeze().numpy()
 
         if gridded:
-            agent.acq_axes[iacq_func].set_title(acq_func_meta["name"])
-            obj_ax = agent.acq_axes[iacq_func].pcolormesh(
+            agent.acq_axes[iacqf].set_title(acqf_meta["name"])
+            obj_ax = agent.acq_axes[iacqf].pcolormesh(
                 test_x,
                 test_y,
                 test_acqf,
@@ -383,17 +386,17 @@ def _plot_acqf_many_dofs(
                 cmap=cmap,
             )
 
-            agent.acq_fig.colorbar(obj_ax, ax=agent.acq_axes[iacq_func], location="bottom", aspect=32, shrink=0.8)
+            agent.acq_fig.colorbar(obj_ax, ax=agent.acq_axes[iacqf], location="bottom", aspect=32, shrink=0.8)
 
         else:
-            agent.acq_axes[iacq_func].set_title(acq_func_meta["name"])
-            obj_ax = agent.acq_axes[iacq_func].scatter(
+            agent.acq_axes[iacqf].set_title(acqf_meta["name"])
+            obj_ax = agent.acq_axes[iacqf].scatter(
                 test_x,
                 test_y,
                 c=test_acqf,
             )
 
-            agent.acq_fig.colorbar(obj_ax, ax=agent.acq_axes[iacq_func], location="bottom", aspect=32, shrink=0.8)
+            agent.acq_fig.colorbar(obj_ax, ax=agent.acq_axes[iacqf], location="bottom", aspect=32, shrink=0.8)
 
     for ax in agent.acq_axes.ravel():
         ax.set_xlabel(x_dof.label_with_units)
@@ -484,11 +487,11 @@ def _plot_history(agent, x_key="index", show_all_objs=False):
     )
     hist_axes = np.atleast_1d(hist_axes)
 
-    unique_strategies, acq_func_index, acq_func_inverse = np.unique(
-        agent.table.acq_func, return_index=True, return_inverse=True
+    unique_strategies, acqf_index, acqf_inverse = np.unique(
+        agent.table.acqf, return_index=True, return_inverse=True
     )
 
-    sample_colors = np.array(DEFAULT_COLOR_LIST)[acq_func_inverse]
+    sample_colors = np.array(DEFAULT_COLOR_LIST)[acqf_inverse]
 
     if show_all_objs:
         for obj_index, obj in enumerate(agent.objectives):
@@ -510,8 +513,8 @@ def _plot_history(agent, x_key="index", show_all_objs=False):
     hist_axes[-1].set_xlabel(x_key)
 
     handles = []
-    for i_acq_func, acq_func in enumerate(unique_strategies):
-        handles.append(Patch(color=DEFAULT_COLOR_LIST[i_acq_func], label=acq_func))
+    for i_acqf, acqf in enumerate(unique_strategies):
+        handles.append(Patch(color=DEFAULT_COLOR_LIST[i_acqf], label=acqf))
     legend = hist_axes[0].legend(handles=handles, fontsize=8)
     legend.set_title("acquisition function")
 
diff --git a/src/blop/tests/conftest.py b/src/blop/tests/conftest.py
index 9c81da8..864efcf 100644
--- a/src/blop/tests/conftest.py
+++ b/src/blop/tests/conftest.py
@@ -46,7 +46,37 @@ def RE(db):
 
 
 @pytest.fixture(scope="function")
-def agent(db):
+def agent_1dof_1fit(db):
+    """
+    A one dimensional agent.
+    """
+
+    def digestion(db, uid):
+        products = db[uid].table()
+
+        for index, entry in products.iterrows():
+            products.loc[index, "f1"]  = functions.himmelblau(entry.x1, 3)
+
+        return products
+
+    dofs = DOF(description="The first DOF", name="x1", search_domain=(-5.0, 5.0))
+    objectives = Objective(description="f1", name="f1", target="min")
+
+
+    agent = Agent(
+        dofs=dofs,
+        objectives=objectives,
+        digestion=digestion,
+        db=db,
+    )
+
+
+    return agent
+
+
+
+@pytest.fixture(scope="function")
+def agent_2dof_1fit(db):
     """
     A simple agent minimizing Himmelblau's function
     """
@@ -71,7 +101,7 @@ def agent(db):
 
 
 @pytest.fixture(scope="function")
-def mo_agent(db):
+def agent_2dof_2fit(db):
     """
     An agent minimizing two Himmelblau's functions
     """
@@ -105,7 +135,7 @@ def digestion(db, uid):
 
 
 @pytest.fixture(scope="function")
-def constrained_agent(db):
+def agent_2dof_2fit_2con(db):
     """
     Chankong and Haimes function from https://en.wikipedia.org/wiki/Test_functions_for_optimization
     """
diff --git a/src/blop/tests/test_acq_funcs.py b/src/blop/tests/test_acq_funcs.py
index 2a9d89d..e785baa 100644
--- a/src/blop/tests/test_acq_funcs.py
+++ b/src/blop/tests/test_acq_funcs.py
@@ -1,25 +1,37 @@
 import pytest
 
 
-@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=4))
-    RE(agent.learn(acq_func, n=1))
+@pytest.mark.parametrize("acqf", ["ei", "pi", "em", "ucb"])
+def test_analytic_acqfs_one_dimensional(agent_1dof_1fit, RE, acqf):
+    RE(agent_1dof_1fit.learn("qr", n=16))
+    RE(agent_1dof_1fit.learn(acqf, 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=4))
-    RE(agent.learn(acq_func, n=4))
+@pytest.mark.parametrize("acqf", ["qei", "qpi", "qem", "qucb"])
+def test_monte_carlo_acqfs_one_dimensional(agent_1dof_1fit, RE, acqf):
+    RE(agent_1dof_1fit.learn("qr", n=16))
+    RE(agent_1dof_1fit.learn(acqf, n=2))
 
 
-@pytest.mark.parametrize("acq_func", ["ei", "pi", "em", "ucb"])
-def test_analytic_acq_funcs_multi_objective(mo_agent, RE, acq_func):
-    RE(mo_agent.learn("qr", n=16))
-    RE(mo_agent.learn(acq_func, n=1))
+@pytest.mark.parametrize("acqf", ["ei", "pi", "em", "ucb"])
+def test_analytic_acqfs_single_objective(agent_2dof_1fit, RE, acqf):
+    RE(agent_2dof_1fit.learn("qr", n=16))
+    RE(agent_2dof_1fit.learn(acqf, n=1))
 
 
-@pytest.mark.parametrize("acq_func", ["qei", "qpi", "qem", "qucb"])
-def test_monte_carlo_acq_funcs_multi_objective(mo_agent, RE, acq_func):
-    RE(mo_agent.learn("qr", n=16))
-    RE(mo_agent.learn(acq_func, n=4))
+@pytest.mark.parametrize("acqf", ["qei", "qpi", "qem", "qucb"])
+def test_monte_carlo_acqfs_single_objective(agent_2dof_1fit, RE, acqf):
+    RE(agent_2dof_1fit.learn("qr", n=16))
+    RE(agent_2dof_1fit.learn(acqf, n=2))
+
+
+@pytest.mark.parametrize("acqf", ["ei", "pi", "em", "ucb"])
+def test_analytic_acqfs_multi_objective(agent_2dof_2fit, RE, acqf):
+    RE(agent_2dof_2fit.learn("qr", n=16))
+    RE(agent_2dof_2fit.learn(acqf, n=1))
+
+
+@pytest.mark.parametrize("acqf", ["qei", "qpi", "qem", "qucb"])
+def test_monte_carlo_acqfs_multi_objective(agent_2dof_2fit, RE, acqf):
+    RE(agent_2dof_2fit.learn("qr", n=16))
+    RE(agent_2dof_2fit.learn(acqf, n=2))
diff --git a/src/blop/tests/test_agent.py b/src/blop/tests/test_agent.py
index a04b432..a62ce1f 100644
--- a/src/blop/tests/test_agent.py
+++ b/src/blop/tests/test_agent.py
@@ -1,18 +1,20 @@
 import pytest  # noqa F401
 
 
-def test_agent(agent, RE):
-    RE(agent.learn("qr", n=4))
+def test_agent(agent_2dof_1fit, RE):
+    RE(agent_2dof_1fit.learn("qr", n=4))
 
-    assert [dof.name in agent.best for dof in agent.dofs]
-    assert [obj.name in agent.best for obj in agent.objectives]
+    agent_2dof_1fit.dofs
+    agent_2dof_1fit.objectives
+    agent_2dof_1fit.best
+    agent_2dof_1fit.best_inputs
 
 
-def test_forget(agent, RE):
-    RE(agent.learn("qr", n=4))
-    agent.forget(last=2)
+def test_forget(agent_2dof_1fit, RE):
+    RE(agent_2dof_1fit.learn("qr", n=4))
+    agent_2dof_1fit.forget(last=2)
 
 
-def test_benchmark(agent, RE):
-    per_iter_learn_kwargs_list = [{"acq_func": "qr", "n": 64}, {"acq_func": "qei", "n": 2, "iterations": 2}]
-    RE(agent.benchmark(output_dir="/tmp/blop", iterations=1, per_iter_learn_kwargs_list=per_iter_learn_kwargs_list))
+def test_benchmark(agent_2dof_1fit, RE):
+    per_iter_learn_kwargs_list = [{"acqf": "qr", "n": 32}, {"acqf": "qei", "n": 2, "iterations": 2}]
+    RE(agent_2dof_1fit.benchmark(output_dir="/tmp/blop", iterations=1, per_iter_learn_kwargs_list=per_iter_learn_kwargs_list))
diff --git a/src/blop/tests/test_pareto.py b/src/blop/tests/test_pareto.py
index 75f1fb6..3b5e1ed 100644
--- a/src/blop/tests/test_pareto.py
+++ b/src/blop/tests/test_pareto.py
@@ -2,24 +2,25 @@
 
 
 @pytest.mark.test_func
-def test_pareto(mo_agent, RE):
-    RE(mo_agent.learn("qr", n=16))
-    mo_agent.plot_pareto_front()
-
-
-@pytest.mark.parametrize("acq_func", ["qnehvi"])
-def test_monte_carlo_pareto_acq_funcs(mo_agent, RE, acq_func):
-    RE(mo_agent.learn("qr", n=16))
-    RE(mo_agent.learn(acq_func, n=2))
+def test_pareto(agent_2dof_2fit, RE):
+    agent = agent_2dof_2fit
+    RE(agent.learn("qr", n=16))
+    agent.plot_pareto_front()
 
+@pytest.mark.parametrize("acqf", ["qnehvi"])
+def test_monte_carlo_pareto_acqfs(agent_2dof_2fit, RE, acqf):
+    agent = agent_2dof_2fit
+    RE(agent.learn("qr", n=16))
+    RE(agent.learn(acqf, n=2))
 
 @pytest.mark.test_func
-def test_constrained_pareto(constrained_agent, RE):
-    RE(constrained_agent.learn("qr", n=16))
-    constrained_agent.plot_pareto_front()
-
-
-@pytest.mark.parametrize("acq_func", ["qnehvi"])
-def test_constrained_monte_carlo_pareto_acq_funcs(constrained_agent, RE, acq_func):
-    RE(constrained_agent.learn("qr", n=16))
-    RE(constrained_agent.learn(acq_func, n=2))
+def test_constrained_pareto(agent_2dof_2fit_2con, RE):
+    agent = agent_2dof_2fit_2con
+    RE(agent.learn("qr", n=16))
+    agent.plot_pareto_front()
+
+@pytest.mark.parametrize("acqf", ["qnehvi"])
+def test_constrained_monte_carlo_pareto_acqfs(agent_2dof_2fit_2con, RE, acqf):
+    agent = agent_2dof_2fit_2con
+    RE(agent.learn("qr", n=16))
+    RE(agent.learn(acqf, n=2))
diff --git a/src/blop/tests/test_passive_dofs.py b/src/blop/tests/test_passive_dofs.py
index fb8624a..0208bb6 100644
--- a/src/blop/tests/test_passive_dofs.py
+++ b/src/blop/tests/test_passive_dofs.py
@@ -2,6 +2,7 @@
 
 
 @pytest.mark.test_func
-def test_read_only_dofs(agent_with_read_only_dofs, RE):
-    RE(agent_with_read_only_dofs.learn("qr", n=32))
-    RE(agent_with_read_only_dofs.learn("qei", n=2))
+def test_passive_dofs(agent_with_passive_dofs, RE):
+    agent = agent_with_passive_dofs
+    RE(agent.learn("qr", n=32))
+    RE(agent.learn("qei", n=2))
diff --git a/src/blop/tests/test_plots.py b/src/blop/tests/test_plots.py
index 4f2c535..9dec6f7 100644
--- a/src/blop/tests/test_plots.py
+++ b/src/blop/tests/test_plots.py
@@ -1,22 +1,30 @@
 import pytest  # noqa F401
 
 
-def test_plots(RE, agent):
+def test_plots_one_dimensional(RE, agent_1dof_1fit):
+    agent = agent_1dof_1fit
     RE(agent.learn("qr", n=16))
-
     agent.plot_objectives()
     agent.plot_acquisition()
     agent.plot_validity()
     agent.plot_history()
 
+def test_plots_two_dimensional(RE, agent_1dof_1fit):
+    agent = agent_1dof_1fit
+    RE(agent.learn("qr", n=16))
+    agent.plot_objectives()
+    agent.plot_acquisition()
+    agent.plot_validity()
+    agent.plot_history()
 
-def test_plots_multiple_objs(RE, mo_agent):
-    RE(mo_agent.learn("qr", n=16))
 
-    mo_agent.plot_objectives()
-    mo_agent.plot_acquisition()
-    mo_agent.plot_validity()
-    mo_agent.plot_history()
+def test_plots_multiple_objs(RE, agent_2dof_2fit):
+    agent = agent_2dof_2fit
+    RE(agent.learn("qr", n=16))
+    agent.plot_objectives()
+    agent.plot_acquisition()
+    agent.plot_validity()
+    agent.plot_history()
 
 
 def test_plots_read_only_dofs(RE, agent_with_read_only_dofs):

From 07046c22067d84282ab3ab75d38d88787296fa41 Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Thu, 18 Apr 2024 12:59:37 -0700
Subject: [PATCH 03/11] fix acqf plots

---
 src/blop/agent.py                 | 15 ++++++-------
 src/blop/plotting.py              | 32 +++++++++++++--------------
 src/blop/tests/conftest.py        | 13 +++++------
 src/blop/tests/test_acq_funcs.py  | 36 +++++++++++++++----------------
 src/blop/tests/test_agent.py      | 22 +++++++++----------
 src/blop/tests/test_pareto.py     | 19 +++++++++-------
 src/blop/tests/test_plots.py      | 13 +++++------
 src/blop/tests/test_read_write.py | 22 +++++++++----------
 8 files changed, 85 insertions(+), 87 deletions(-)

diff --git a/src/blop/agent.py b/src/blop/agent.py
index 4fe4d93..f892c72 100644
--- a/src/blop/agent.py
+++ b/src/blop/agent.py
@@ -28,7 +28,7 @@
 
 from . import plotting, utils
 from .bayesian import acquisition, models
-from .bayesian.acquisition import parse_acqf_identifier, _construct_acqf
+from .bayesian.acquisition import _construct_acqf, parse_acqf_identifier
 
 # from .bayesian.transforms import TargetingPosteriorTransform
 from .digestion import default_digestion_function
@@ -157,7 +157,7 @@ def __iter__(self):
             yield self.dofs[index]
 
     def __getattr__(self, attr):
-        acqf_config = acquisition.parse_acqf_identifier(attr)
+        acqf_config = acquisition.parse_acqf_identifier(attr, strict=False)
         if acqf_config is not None:
             acqf, _ = _construct_acqf(acqf_name=acqf_config["name"])
             return acqf
@@ -228,11 +228,10 @@ def ask(self, acqf="qei", n=1, route=True, sequential=True, upsample=1, **acqf_k
             acqf_kwargs, acqf_obj = {}, torch.zeros(len(candidates))
 
         else:
-            # check that all the objectives have models 
+            # check that all the objectives have models
             if not all(hasattr(obj, "model") for obj in active_objs):
                 raise RuntimeError(
-                    f"Can't construct non-trivial acquisition function '{acqf}'"
-                    f" (the agent is not initialized!)"
+                    f"Can't construct non-trivial acquisition function '{acqf}'" f" (the agent is not initialized!)"
                 )
 
             # if the model for any active objective mismatches the active dofs, reconstrut and train it
@@ -244,7 +243,7 @@ def ask(self, acqf="qei", n=1, route=True, sequential=True, upsample=1, **acqf_k
             if acqf_config["type"] == "analytic" and n > 1:
                 raise ValueError("Can't generate multiple design points for analytic acquisition functions.")
 
-            # we may pick up some more kwargs 
+            # we may pick up some more kwargs
             acqf, acqf_kwargs = _construct_acqf(self, acqf_name=acqf_config["name"], **acqf_kwargs)
 
             NUM_RESTARTS = 16
@@ -628,7 +627,7 @@ def random_ref_point(self):
     @property
     def all_objectives_valid(self):
         """A mask of whether all objectives are valid for each data point."""
-        return ~torch.isnan(self.scalarized_fitnesses)
+        return ~torch.isnan(self.scalarized_fitnesses())
 
     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`."""
@@ -719,7 +718,7 @@ def _get_acquisition_function(self, identifier, return_metadata=False):
 
         acquisition._construct_acqf(self, identifier=identifier, return_metadata=return_metadata)
 
-        return 
+        return
 
     def _latent_dim_tuples(self, obj_index=None):
         """
diff --git a/src/blop/plotting.py b/src/blop/plotting.py
index 06b3da2..f18d293 100644
--- a/src/blop/plotting.py
+++ b/src/blop/plotting.py
@@ -3,7 +3,7 @@
 from matplotlib import pyplot as plt
 from matplotlib.patches import Patch
 
-from .bayesian.acquisition import parse_acqf_identifier, _construct_acqf
+from .bayesian.acquisition import _construct_acqf, parse_acqf_identifier
 
 DEFAULT_COLOR_LIST = ["dodgerblue", "tomato", "mediumseagreen", "goldenrod"]
 # Note: the values near 1 are hard to see on a white background. Turbo goes from red to blue and isn't white in the middle.
@@ -331,14 +331,16 @@ def _plot_acqf_one_dof(agent, acqfs, lw=1e0, **kwargs):
     for iacqf, acqf_identifier in enumerate(acqfs):
         color = DEFAULT_COLOR_LIST[iacqf]
 
-        acqf, acqf_meta = acquisition.get_acquisition_function(agent, acqf_identifier)
-        test_acqf = acqf(test_inputs).detach().numpy()
+        acqf_config = parse_acqf_identifier(acqf_identifier)
+        acqf, _ = _construct_acqf(agent, acqf_config["name"])
 
-        agent.acq_axes[iacqf].plot(test_inputs.squeeze(-2), test_acqf, lw=lw, color=color)
+        test_acqf_value = acqf(test_inputs).detach().numpy()
+
+        agent.acq_axes[iacqf].plot(test_inputs.squeeze(-2), test_acqf_value, lw=lw, color=color)
 
         agent.acq_axes[iacqf].set_xlim(*x_dof.search_domain)
         agent.acq_axes[iacqf].set_xlabel(x_dof.label_with_units)
-        agent.acq_axes[iacqf].set_ylabel(acqf_meta["name"])
+        agent.acq_axes[iacqf].set_ylabel(acqf_config["name"])
 
 
 def _plot_acqf_many_dofs(
@@ -369,19 +371,17 @@ def _plot_acqf_many_dofs(
     test_y = test_inputs[..., 0, axes[1]].detach().squeeze().numpy()
 
     for iacqf, acqf_identifier in enumerate(acqfs):
+        acqf_config = parse_acqf_identifier(acqf_identifier)
+        acqf, _ = _construct_acqf(agent, acqf_config["name"])
 
-        acqf_config = par
-
-        acqf, acqf_meta = acquisition.get_acquisition_function(agent, acqf_identifier)
-
-        test_acqf = acqf(test_inputs.reshape(-1, 1, input_dim)).detach().reshape(test_dim).squeeze().numpy()
+        test_acqf_value = acqf(test_inputs.reshape(-1, 1, input_dim)).detach().reshape(test_dim).squeeze().numpy()
 
         if gridded:
-            agent.acq_axes[iacqf].set_title(acqf_meta["name"])
+            agent.acq_axes[iacqf].set_title(acqf_config["name"])
             obj_ax = agent.acq_axes[iacqf].pcolormesh(
                 test_x,
                 test_y,
-                test_acqf,
+                test_acqf_value,
                 shading=shading,
                 cmap=cmap,
             )
@@ -389,11 +389,11 @@ def _plot_acqf_many_dofs(
             agent.acq_fig.colorbar(obj_ax, ax=agent.acq_axes[iacqf], location="bottom", aspect=32, shrink=0.8)
 
         else:
-            agent.acq_axes[iacqf].set_title(acqf_meta["name"])
+            agent.acq_axes[iacqf].set_title(acqf_config["name"])
             obj_ax = agent.acq_axes[iacqf].scatter(
                 test_x,
                 test_y,
-                c=test_acqf,
+                c=test_acqf_value,
             )
 
             agent.acq_fig.colorbar(obj_ax, ax=agent.acq_axes[iacqf], location="bottom", aspect=32, shrink=0.8)
@@ -487,9 +487,7 @@ def _plot_history(agent, x_key="index", show_all_objs=False):
     )
     hist_axes = np.atleast_1d(hist_axes)
 
-    unique_strategies, acqf_index, acqf_inverse = np.unique(
-        agent.table.acqf, return_index=True, return_inverse=True
-    )
+    unique_strategies, acqf_index, acqf_inverse = np.unique(agent.table.acqf, return_index=True, return_inverse=True)
 
     sample_colors = np.array(DEFAULT_COLOR_LIST)[acqf_inverse]
 
diff --git a/src/blop/tests/conftest.py b/src/blop/tests/conftest.py
index 864efcf..12ef1e4 100644
--- a/src/blop/tests/conftest.py
+++ b/src/blop/tests/conftest.py
@@ -46,7 +46,7 @@ def RE(db):
 
 
 @pytest.fixture(scope="function")
-def agent_1dof_1fit(db):
+def agent_1d_1f(db):
     """
     A one dimensional agent.
     """
@@ -55,14 +55,13 @@ def digestion(db, uid):
         products = db[uid].table()
 
         for index, entry in products.iterrows():
-            products.loc[index, "f1"]  = functions.himmelblau(entry.x1, 3)
+            products.loc[index, "f1"] = functions.himmelblau(entry.x1, 3)
 
         return products
 
     dofs = DOF(description="The first DOF", name="x1", search_domain=(-5.0, 5.0))
     objectives = Objective(description="f1", name="f1", target="min")
 
-
     agent = Agent(
         dofs=dofs,
         objectives=objectives,
@@ -70,13 +69,11 @@ def digestion(db, uid):
         db=db,
     )
 
-
     return agent
 
 
-
 @pytest.fixture(scope="function")
-def agent_2dof_1fit(db):
+def agent_2d_1f(db):
     """
     A simple agent minimizing Himmelblau's function
     """
@@ -101,7 +98,7 @@ def agent_2dof_1fit(db):
 
 
 @pytest.fixture(scope="function")
-def agent_2dof_2fit(db):
+def agent_2d_2f(db):
     """
     An agent minimizing two Himmelblau's functions
     """
@@ -135,7 +132,7 @@ def digestion(db, uid):
 
 
 @pytest.fixture(scope="function")
-def agent_2dof_2fit_2con(db):
+def agent_2d_2f_2c(db):
     """
     Chankong and Haimes function from https://en.wikipedia.org/wiki/Test_functions_for_optimization
     """
diff --git a/src/blop/tests/test_acq_funcs.py b/src/blop/tests/test_acq_funcs.py
index e785baa..5bf1562 100644
--- a/src/blop/tests/test_acq_funcs.py
+++ b/src/blop/tests/test_acq_funcs.py
@@ -2,36 +2,36 @@
 
 
 @pytest.mark.parametrize("acqf", ["ei", "pi", "em", "ucb"])
-def test_analytic_acqfs_one_dimensional(agent_1dof_1fit, RE, acqf):
-    RE(agent_1dof_1fit.learn("qr", n=16))
-    RE(agent_1dof_1fit.learn(acqf, n=1))
+def test_analytic_acqfs_one_dimensional(agent_1d_1f, RE, acqf):
+    RE(agent_1d_1f.learn("qr", n=16))
+    RE(agent_1d_1f.learn(acqf, n=1))
 
 
 @pytest.mark.parametrize("acqf", ["qei", "qpi", "qem", "qucb"])
-def test_monte_carlo_acqfs_one_dimensional(agent_1dof_1fit, RE, acqf):
-    RE(agent_1dof_1fit.learn("qr", n=16))
-    RE(agent_1dof_1fit.learn(acqf, n=2))
+def test_monte_carlo_acqfs_one_dimensional(agent_1d_1f, RE, acqf):
+    RE(agent_1d_1f.learn("qr", n=16))
+    RE(agent_1d_1f.learn(acqf, n=2))
 
 
 @pytest.mark.parametrize("acqf", ["ei", "pi", "em", "ucb"])
-def test_analytic_acqfs_single_objective(agent_2dof_1fit, RE, acqf):
-    RE(agent_2dof_1fit.learn("qr", n=16))
-    RE(agent_2dof_1fit.learn(acqf, n=1))
+def test_analytic_acqfs_single_objective(agent_2d_1f, RE, acqf):
+    RE(agent_2d_1f.learn("qr", n=16))
+    RE(agent_2d_1f.learn(acqf, n=1))
 
 
 @pytest.mark.parametrize("acqf", ["qei", "qpi", "qem", "qucb"])
-def test_monte_carlo_acqfs_single_objective(agent_2dof_1fit, RE, acqf):
-    RE(agent_2dof_1fit.learn("qr", n=16))
-    RE(agent_2dof_1fit.learn(acqf, n=2))
+def test_monte_carlo_acqfs_single_objective(agent_2d_1f, RE, acqf):
+    RE(agent_2d_1f.learn("qr", n=16))
+    RE(agent_2d_1f.learn(acqf, n=2))
 
 
 @pytest.mark.parametrize("acqf", ["ei", "pi", "em", "ucb"])
-def test_analytic_acqfs_multi_objective(agent_2dof_2fit, RE, acqf):
-    RE(agent_2dof_2fit.learn("qr", n=16))
-    RE(agent_2dof_2fit.learn(acqf, n=1))
+def test_analytic_acqfs_multi_objective(agent_2d_2f, RE, acqf):
+    RE(agent_2d_2f.learn("qr", n=16))
+    RE(agent_2d_2f.learn(acqf, n=1))
 
 
 @pytest.mark.parametrize("acqf", ["qei", "qpi", "qem", "qucb"])
-def test_monte_carlo_acqfs_multi_objective(agent_2dof_2fit, RE, acqf):
-    RE(agent_2dof_2fit.learn("qr", n=16))
-    RE(agent_2dof_2fit.learn(acqf, n=2))
+def test_monte_carlo_acqfs_multi_objective(agent_2d_2f, RE, acqf):
+    RE(agent_2d_2f.learn("qr", n=16))
+    RE(agent_2d_2f.learn(acqf, n=2))
diff --git a/src/blop/tests/test_agent.py b/src/blop/tests/test_agent.py
index a62ce1f..2f35102 100644
--- a/src/blop/tests/test_agent.py
+++ b/src/blop/tests/test_agent.py
@@ -1,20 +1,20 @@
 import pytest  # noqa F401
 
 
-def test_agent(agent_2dof_1fit, RE):
-    RE(agent_2dof_1fit.learn("qr", n=4))
+def test_agent(agent_2d_1f, RE):
+    RE(agent_2d_1f.learn("qr", n=4))
 
-    agent_2dof_1fit.dofs
-    agent_2dof_1fit.objectives
-    agent_2dof_1fit.best
-    agent_2dof_1fit.best_inputs
+    agent_2d_1f.dofs
+    agent_2d_1f.objectives
+    agent_2d_1f.best
+    agent_2d_1f.best_inputs
 
 
-def test_forget(agent_2dof_1fit, RE):
-    RE(agent_2dof_1fit.learn("qr", n=4))
-    agent_2dof_1fit.forget(last=2)
+def test_forget(agent_2d_1f, RE):
+    RE(agent_2d_1f.learn("qr", n=4))
+    agent_2d_1f.forget(last=2)
 
 
-def test_benchmark(agent_2dof_1fit, RE):
+def test_benchmark(agent_2d_1f, RE):
     per_iter_learn_kwargs_list = [{"acqf": "qr", "n": 32}, {"acqf": "qei", "n": 2, "iterations": 2}]
-    RE(agent_2dof_1fit.benchmark(output_dir="/tmp/blop", iterations=1, per_iter_learn_kwargs_list=per_iter_learn_kwargs_list))
+    RE(agent_2d_1f.benchmark(output_dir="/tmp/blop", iterations=1, per_iter_learn_kwargs_list=per_iter_learn_kwargs_list))
diff --git a/src/blop/tests/test_pareto.py b/src/blop/tests/test_pareto.py
index 3b5e1ed..5934300 100644
--- a/src/blop/tests/test_pareto.py
+++ b/src/blop/tests/test_pareto.py
@@ -2,25 +2,28 @@
 
 
 @pytest.mark.test_func
-def test_pareto(agent_2dof_2fit, RE):
-    agent = agent_2dof_2fit
+def test_pareto(agent_2d_2f, RE):
+    agent = agent_2d_2f
     RE(agent.learn("qr", n=16))
     agent.plot_pareto_front()
 
+
 @pytest.mark.parametrize("acqf", ["qnehvi"])
-def test_monte_carlo_pareto_acqfs(agent_2dof_2fit, RE, acqf):
-    agent = agent_2dof_2fit
+def test_monte_carlo_pareto_acqfs(agent_2d_2f, RE, acqf):
+    agent = agent_2d_2f
     RE(agent.learn("qr", n=16))
     RE(agent.learn(acqf, n=2))
 
+
 @pytest.mark.test_func
-def test_constrained_pareto(agent_2dof_2fit_2con, RE):
-    agent = agent_2dof_2fit_2con
+def test_constrained_pareto(agent_2d_2f_2c, RE):
+    agent = agent_2d_2f_2c
     RE(agent.learn("qr", n=16))
     agent.plot_pareto_front()
 
+
 @pytest.mark.parametrize("acqf", ["qnehvi"])
-def test_constrained_monte_carlo_pareto_acqfs(agent_2dof_2fit_2con, RE, acqf):
-    agent = agent_2dof_2fit_2con
+def test_constrained_monte_carlo_pareto_acqfs(agent_2d_2f_2c, RE, acqf):
+    agent = agent_2d_2f_2c
     RE(agent.learn("qr", n=16))
     RE(agent.learn(acqf, n=2))
diff --git a/src/blop/tests/test_plots.py b/src/blop/tests/test_plots.py
index 9dec6f7..77c20bf 100644
--- a/src/blop/tests/test_plots.py
+++ b/src/blop/tests/test_plots.py
@@ -1,16 +1,17 @@
 import pytest  # noqa F401
 
 
-def test_plots_one_dimensional(RE, agent_1dof_1fit):
-    agent = agent_1dof_1fit
+def test_plots_one_dimensional(RE, agent_1d_1f):
+    agent = agent_1d_1f
     RE(agent.learn("qr", n=16))
     agent.plot_objectives()
     agent.plot_acquisition()
     agent.plot_validity()
     agent.plot_history()
 
-def test_plots_two_dimensional(RE, agent_1dof_1fit):
-    agent = agent_1dof_1fit
+
+def test_plots_two_dimensional(RE, agent_1d_1f):
+    agent = agent_1d_1f
     RE(agent.learn("qr", n=16))
     agent.plot_objectives()
     agent.plot_acquisition()
@@ -18,8 +19,8 @@ def test_plots_two_dimensional(RE, agent_1dof_1fit):
     agent.plot_history()
 
 
-def test_plots_multiple_objs(RE, agent_2dof_2fit):
-    agent = agent_2dof_2fit
+def test_plots_multiple_objs(RE, agent_2d_2f):
+    agent = agent_2d_2f
     RE(agent.learn("qr", n=16))
     agent.plot_objectives()
     agent.plot_acquisition()
diff --git a/src/blop/tests/test_read_write.py b/src/blop/tests/test_read_write.py
index db4b052..cb37deb 100644
--- a/src/blop/tests/test_read_write.py
+++ b/src/blop/tests/test_read_write.py
@@ -1,16 +1,16 @@
 import pytest  # noqa F401
 
 
-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("/tmp/test_save_data.h5")
-    RE(agent.learn("qr", n=4))
+def test_agent_2d_1f_save_load_data(agent_2d_1f, RE):
+    RE(agent_2d_1f.learn("qr", n=4))
+    agent_2d_1f.save_data("/tmp/test_save_data.h5")
+    agent_2d_1f.reset()
+    agent_2d_1f.load_data("/tmp/test_save_data.h5")
+    RE(agent_2d_1f.learn("qr", n=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="/tmp/test_save_hypers.h5"))
+def test_agent_2d_1f_save_load_hypers(agent_2d_1f, RE):
+    RE(agent_2d_1f.learn("qr", n=4))
+    agent_2d_1f.save_hypers("/tmp/test_save_hypers.h5")
+    agent_2d_1f.reset()
+    RE(agent_2d_1f.learn("qr", n=16, hypers="/tmp/test_save_hypers.h5"))

From 6a6c34a655c814f5245b1b46f8b2130bafa15e30 Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Mon, 22 Apr 2024 11:03:13 -0700
Subject: [PATCH 04/11] fix bugs after rebasing

---
 src/blop/agent.py                   | 32 ++++++++++----------
 src/blop/objectives.py              | 47 ++---------------------------
 src/blop/plotting.py                |  6 ++--
 src/blop/tests/test_agent.py        | 11 ++++---
 src/blop/tests/test_passive_dofs.py |  4 +--
 5 files changed, 30 insertions(+), 70 deletions(-)

diff --git a/src/blop/agent.py b/src/blop/agent.py
index f892c72..cc5f222 100644
--- a/src/blop/agent.py
+++ b/src/blop/agent.py
@@ -345,9 +345,9 @@ def tell(
                 t0 = ttime.monotonic()
 
                 cached_hypers = obj.model.state_dict() if hasattr(obj, "model") else None
-                n_before_tell = obj.n
+                n_before_tell = obj.n_valid
                 self._construct_model(obj)
-                n_after_tell = obj.n
+                n_after_tell = obj.n_valid
 
                 if train is None:
                     train = int(n_after_tell / self.train_every) > int(n_before_tell / self.train_every)
@@ -408,8 +408,8 @@ def learn(
         for i in range(iterations):
             if self.verbose:
                 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, route=route, **acq_func_kwargs)
+            for single_acqf in np.atleast_1d(acqf):
+                res = self.ask(n=n, acqf=single_acqf, upsample=upsample, route=route, **acqf_kwargs)
                 new_table = yield from self.acquire(res["points"])
                 new_table.loc[:, "acqf"] = res["acqf_name"]
 
@@ -596,24 +596,24 @@ def best_f(self, weights="default"):
         return float(self.scalarized_fitnesses(weights=weights, constrained=True).max())
 
     @property
-    def pareto_front_mask(self):
-        """
-        A mask for all data points that is true when the point is on the Pareto front.
-        A point is on the Pareto front if it is Pareto dominant
-        A point is Pareto dominant if there is no other point that is better at every objective
-        Points that violate any constraint are excluded.
-        """
-        y = self.train_targets()[:, self.objectives.type == "fitness"]
-        in_pareto_front = ~(y.unsqueeze(1) > y.unsqueeze(0)).all(axis=-1).any(axis=0)
-        all_constraints_satisfied = self.evaluated_constraints.all(axis=-1)
-        return in_pareto_front & all_constraints_satisfied
+    def pareto_mask(self):
+        # a point is on the Pareto front if it is Pareto dominant
+        # a point is Pareto dominant if it is there is no other point that is better at every objective
+        Y = self.train_targets(active=True, kind="fitness")
+
+        # nuke the bad points
+        Y[~self.evaluated_constraints.all(axis=-1)] = -np.inf
+        if Y.shape[-1] < 2:
+            raise ValueError("Computing the Pareto front requires at least 2 fitness objectives.")
+        in_pareto_front = ~(Y.unsqueeze(1) > Y.unsqueeze(0)).all(axis=-1).any(axis=0)
+        return in_pareto_front & self.evaluated_constraints.all(axis=-1)
 
     @property
     def pareto_front(self):
         """
         A subset of the data table containing only points on the Pareto front.
         """
-        return self.table.loc[self.pareto_front_mask.numpy()]
+        return self.table.loc[self.pareto_mask.numpy()]
 
     @property
     def min_ref_point(self):
diff --git a/src/blop/objectives.py b/src/blop/objectives.py
index 6aacba7..368411d 100644
--- a/src/blop/objectives.py
+++ b/src/blop/objectives.py
@@ -212,7 +212,6 @@ def summary(self) -> pd.Series:
             series[attr] = value if value is not None else ""
         return series
 
->>>>>>> 319737d (better control of DOFs and objectives)
     @property
     def noise(self) -> float:
         return self.model.likelihood.noise.item() if hasattr(self, "model") else np.nan
@@ -222,7 +221,7 @@ def snr(self) -> float:
         return np.round(1 / self.model.likelihood.noise.sqrt().item(), 3) if hasattr(self, "model") else None
 
     @property
-    def n(self) -> int:
+    def n_valid(self) -> int:
         return int((~self.model.train_targets.isnan()).sum()) if hasattr(self, "model") else 0
 
     def targeting_constraint(self, x: torch.Tensor) -> torch.Tensor:
@@ -291,7 +290,7 @@ def __call__(self, *args, **kwargs):
     def names(self):
         return [obj.name for obj in self.objectives]
 
-    def __getattr__(self, attr):    
+    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(), "kind"]:
             return np.array([getattr(obj, attr) for obj in self.objectives])
@@ -336,48 +335,6 @@ def __repr__(self):
     def _repr_html_(self):
         return self.summary._repr_html_()
 
-    @property
-    def descriptions(self) -> list:
-        """
-        Returns an array of the objective names.
-        """
-        return [obj.description for obj in self.objectives]
-
-    @property
-    def names(self) -> list:
-        """
-        Returns an array of the objective names.
-        """
-        return [obj.name for obj in self.objectives]
-
-    @property
-    def targets(self) -> list:
-        """
-        Returns an array of the objective targets.
-        """
-        return [obj.target for obj in self.objectives]
-
-    @property
-    def weights(self) -> np.array:
-        """
-        Returns an array of the objective weights.
-        """
-        return np.array([obj.weight for obj in self.objectives])
-
-    @property
-    def is_fitness(self) -> np.array:
-        """
-        Returns an array of the objective weights.
-        """
-        return np.array([obj.target in ["min", "max"] for obj in self.objectives])
-
-    @property
-    def signed_weights(self) -> np.array:
-        """
-        Returns a signed array of the objective weights.
-        """
-        return np.array([(1 if obj.target == "max" else -1) * obj.weight for obj in self.objectives])
-
     def add(self, objective):
         _validate_objs([*self.objectives, objective])
         self.objectives.append(objective)
diff --git a/src/blop/plotting.py b/src/blop/plotting.py
index f18d293..00547ee 100644
--- a/src/blop/plotting.py
+++ b/src/blop/plotting.py
@@ -547,12 +547,12 @@ def _plot_pareto_front(agent, obj_indices=(0, 1)):
 
     y = agent.train_targets()[:, agent.objectives.kind == "fitness"]
 
-    pareto_front_mask = agent.pareto_front_mask
+    pareto_mask = agent.pareto_mask
     constraint = agent.evaluated_constraints.all(axis=-1)
 
-    ax.scatter(*y[(~pareto_front_mask) & constraint].T[[i, j]], c="k")
+    ax.scatter(*y[(~pareto_mask) & constraint].T[[i, j]], c="k")
     ax.scatter(*y[~constraint].T[[i, j]], c="r", marker="x", label="invalid")
-    ax.scatter(*y[pareto_front_mask].T[[i, j]], c="b", label="Pareto front")
+    ax.scatter(*y[pareto_mask].T[[i, j]], c="b", label="Pareto front")
 
     ax.set_xlabel(f"{f_objs[i].name} fitness")
     ax.set_ylabel(f"{f_objs[j].name} fitness")
diff --git a/src/blop/tests/test_agent.py b/src/blop/tests/test_agent.py
index 2f35102..028a7da 100644
--- a/src/blop/tests/test_agent.py
+++ b/src/blop/tests/test_agent.py
@@ -4,10 +4,13 @@
 def test_agent(agent_2d_1f, RE):
     RE(agent_2d_1f.learn("qr", n=4))
 
-    agent_2d_1f.dofs
-    agent_2d_1f.objectives
-    agent_2d_1f.best
-    agent_2d_1f.best_inputs
+    best = agent_2d_1f.best
+
+    assert [dof.name in best for dof in agent_2d_1f.dofs]
+    assert [obj.name in best for obj in agent_2d_1f.objectives]
+
+    print(agent_2d_1f.dofs)
+    print(agent_2d_1f.objectives)
 
 
 def test_forget(agent_2d_1f, RE):
diff --git a/src/blop/tests/test_passive_dofs.py b/src/blop/tests/test_passive_dofs.py
index 0208bb6..f20e2d7 100644
--- a/src/blop/tests/test_passive_dofs.py
+++ b/src/blop/tests/test_passive_dofs.py
@@ -2,7 +2,7 @@
 
 
 @pytest.mark.test_func
-def test_passive_dofs(agent_with_passive_dofs, RE):
-    agent = agent_with_passive_dofs
+def test_read_only_dofs(agent_with_read_only_dofs, RE):
+    agent = agent_with_read_only_dofs
     RE(agent.learn("qr", n=32))
     RE(agent.learn("qei", n=2))

From d28c620598ab2d78b1952cf6047a2b6639dd26d4 Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Tue, 23 Apr 2024 13:06:37 -0700
Subject: [PATCH 05/11] reconfigure digestion syntax

---
 docs/source/tutorials/hyperparameters.ipynb | 14 +++---
 docs/source/tutorials/introduction.ipynb    | 16 +++----
 docs/source/tutorials/pareto-fronts.ipynb   | 16 +++----
 scripts/gui.py                              | 18 +++----
 src/blop/agent.py                           | 24 +++++++---
 src/blop/bayesian/acquisition/__init__.py   | 18 +++----
 src/blop/digestion.py                       |  8 ++--
 src/blop/dofs.py                            | 27 ++++-------
 src/blop/objectives.py                      | 11 +++--
 src/blop/plotting.py                        |  4 +-
 src/blop/tests/conftest.py                  | 52 +++++++++++----------
 src/blop/tests/test_acq_funcs.py            | 36 +++++++++-----
 src/blop/tests/test_agent.py                | 33 +++++++++----
 src/blop/tests/test_pareto.py               |  4 ++
 src/blop/utils/__init__.py                  |  6 +--
 src/blop/utils/functions.py                 | 42 ++++++++---------
 16 files changed, 182 insertions(+), 147 deletions(-)

diff --git a/docs/source/tutorials/hyperparameters.ipynb b/docs/source/tutorials/hyperparameters.ipynb
index 9170ffb..e2c03b5 100644
--- a/docs/source/tutorials/hyperparameters.ipynb
+++ b/docs/source/tutorials/hyperparameters.ipynb
@@ -50,13 +50,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def digestion(db, uid):\n",
-    "    products = db[uid].table()\n",
+    "def digestion(df):\n",
+    "    for index, entry in df.iterrows():\n",
+    "        df.loc[index, \"booth\"] = functions.booth(entry.x1, entry.x2)\n",
     "\n",
-    "    for index, entry in products.iterrows():\n",
-    "        products.loc[index, \"booth\"] = functions.booth(entry.x1, entry.x2)\n",
-    "\n",
-    "    return products"
+    "    return df"
    ]
   },
   {
@@ -91,7 +89,7 @@
     "    db=db,\n",
     ")\n",
     "\n",
-    "RE(agent.learn(acq_func=\"qr\", n=16))\n",
+    "RE(agent.learn(acqf=\"qr\", n=16))\n",
     "\n",
     "agent.plot_objectives()"
    ]
@@ -114,7 +112,7 @@
    },
    "outputs": [],
    "source": [
-    "agent.plot_acquisition(acq_func=\"qei\")"
+    "agent.plot_acquisition(acqf=\"qei\")"
    ]
   },
   {
diff --git a/docs/source/tutorials/introduction.ipynb b/docs/source/tutorials/introduction.ipynb
index 143fb6a..214cfff 100644
--- a/docs/source/tutorials/introduction.ipynb
+++ b/docs/source/tutorials/introduction.ipynb
@@ -113,13 +113,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def digestion(db, uid):\n",
-    "    products = db[uid].table()\n",
+    "def digestion(df):\n",
+    "    for index, entry in df.iterrows():\n",
+    "        df.loc[index, \"himmelblau\"] = functions.himmelblau(entry.x1, entry.x2)\n",
     "\n",
-    "    for index, entry in products.iterrows():\n",
-    "        products.loc[index, \"himmelblau\"] = functions.himmelblau(entry.x1, entry.x2)\n",
-    "\n",
-    "    return products"
+    "    return df"
    ]
   },
   {
@@ -184,7 +182,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "agent.all_acq_funcs"
+    "agent.all_acqfs"
    ]
   },
   {
@@ -203,7 +201,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "agent.plot_acquisition(acq_func=\"qei\")"
+    "agent.plot_acquisition(acqf=\"qei\")"
    ]
   },
   {
@@ -246,7 +244,7 @@
    "outputs": [],
    "source": [
     "res = agent.ask(\"qei\", n=8, route=True)\n",
-    "agent.plot_acquisition(acq_func=\"qei\")\n",
+    "agent.plot_acquisition(acqf=\"qei\")\n",
     "plt.scatter(*res[\"points\"].T, marker=\"d\", facecolor=\"w\", edgecolor=\"k\")\n",
     "plt.plot(\n",
     "    *res[\"points\"].T,\n",
diff --git a/docs/source/tutorials/pareto-fronts.ipynb b/docs/source/tutorials/pareto-fronts.ipynb
index 4b8a13c..33793e9 100644
--- a/docs/source/tutorials/pareto-fronts.ipynb
+++ b/docs/source/tutorials/pareto-fronts.ipynb
@@ -37,18 +37,16 @@
     "import numpy as np\n",
     "\n",
     "\n",
-    "def digestion(db, uid):\n",
-    "    products = db[uid].table()\n",
-    "\n",
-    "    for index, entry in products.iterrows():\n",
+    "def digestion(df):\n",
+    "    for index, entry in df.iterrows():\n",
     "        x1, x2 = entry.x1, entry.x2\n",
     "\n",
-    "        products.loc[index, \"f1\"] = (x1 - 2) ** 2 + (x2 - 1) + 2\n",
-    "        products.loc[index, \"f2\"] = 9 * x1 - (x2 - 1) + 2\n",
-    "        products.loc[index, \"c1\"] = x1**2 + x2**2\n",
-    "        products.loc[index, \"c2\"] = x1 - 3 * x2 + 10\n",
+    "        df.loc[index, \"f1\"] = (x1 - 2) ** 2 + (x2 - 1) + 2\n",
+    "        df.loc[index, \"f2\"] = 9 * x1 - (x2 - 1) + 2\n",
+    "        df.loc[index, \"c1\"] = x1**2 + x2**2\n",
+    "        df.loc[index, \"c2\"] = x1 - 3 * x2 + 10\n",
     "\n",
-    "    return products\n",
+    "    return df\n",
     "\n",
     "\n",
     "dofs = [\n",
diff --git a/scripts/gui.py b/scripts/gui.py
index 9955180..9529d20 100644
--- a/scripts/gui.py
+++ b/scripts/gui.py
@@ -72,8 +72,8 @@
     err_cbar = obj_plt.fig.colorbar(mappable=im3, ax=[ax3], location="bottom", aspect=16)
 
     for ax in [ax1, ax2, ax3]:
-        ax.set_xlabel(agent.dofs[0].label)
-        ax.set_ylabel(agent.dofs[1].label)
+        ax.set_xlabel(agent.dofs[0].label_with_units)
+        ax.set_ylabel(agent.dofs[1].label_with_units)
 
 
 acqf_configs = {
@@ -103,8 +103,8 @@
         acqf_plt_objs[acqf]["hist"] = ax.scatter([], [])
         acqf_plt_objs[acqf]["best"] = ax.scatter([], [])
 
-        ax.set_xlabel(agent.dofs[0].label)
-        ax.set_ylabel(agent.dofs[1].label)
+        ax.set_xlabel(agent.dofs[0].label_with_units)
+        ax.set_ylabel(agent.dofs[1].label_with_units)
 
 
 acqf_button_options = {index: config["name"] for index, config in acqf_configs.items()}
@@ -135,11 +135,12 @@ def learn():
     with obj_plt:
         obj = agent.objectives[0]
 
-        x_samples = agent.train_inputs().detach().numpy()
-        y_samples = agent.train_targets(obj.name).detach().numpy()[..., 0]
+        x_samples = agent.raw_inputs().detach().numpy()
+        y_samples = agent.raw_targets(obj.name).detach().numpy()[..., 0]
 
         x = agent.sample(method="grid", n=20000)  # (n, n, 1, d)
-        p = obj.model.posterior(x)
+        model_x = agent.dofs.transform(x)
+        p = obj.model.posterior(model_x)
 
         m = p.mean.squeeze(-1, -2).detach().numpy()
         e = p.variance.sqrt().squeeze(-1, -2).detach().numpy()
@@ -164,12 +165,13 @@ def learn():
 
     with acq_plt:
         x = agent.sample(method="grid", n=20000)  # (n, n, 1, d)
+        model_x = agent.dofs.transform(x)
         x_samples = agent.train_inputs().detach().numpy()
 
         for acqf in acqf_plt_objs.keys():
             ax = acqf_plt_objs[acqf]["ax"]
 
-            acqf_obj = getattr(agent, acqf)(x).detach().numpy()
+            acqf_obj = getattr(agent, acqf)(model_x).detach().numpy()
 
             acqf_norm = mpl.colors.Normalize(vmin=np.nanmin(acqf_obj), vmax=np.nanmax(acqf_obj))
             acqf_plt_objs[acqf]["im"].set_data(acqf_obj.T[::-1])
diff --git a/src/blop/agent.py b/src/blop/agent.py
index cc5f222..2f9aa69 100644
--- a/src/blop/agent.py
+++ b/src/blop/agent.py
@@ -69,6 +69,7 @@ def __init__(
         dets: Sequence[Signal] = [],
         acquistion_plan=default_acquisition_plan,
         digestion: Callable = default_digestion_function,
+        digestion_kwargs: dict = {},
         verbose: bool = False,
         tolerate_acquisition_errors=False,
         sample_center_on_init=False,
@@ -89,7 +90,9 @@ def __init__(
         acquisition_plan : optional
             A plan that samples the beamline for some given inputs.
         digestion :
-            A function to digest the output of the acquisition, taking arguments (db, uid).
+            A function to digest the output of the acquisition, taking a DataFrame as an argument.
+        digestion_kwargs :
+            Some kwargs for the digestion function.
         db : optional
             A databroker instance.
         verbose : bool
@@ -130,6 +133,7 @@ def __init__(
         self.dets = dets
         self.acquisition_plan = acquistion_plan
         self.digestion = digestion
+        self.digestion_kwargs = digestion_kwargs
 
         self.verbose = verbose
 
@@ -159,10 +163,17 @@ def __iter__(self):
     def __getattr__(self, attr):
         acqf_config = acquisition.parse_acqf_identifier(attr, strict=False)
         if acqf_config is not None:
-            acqf, _ = _construct_acqf(acqf_name=acqf_config["name"])
+            acqf, _ = _construct_acqf(agent=self, acqf_name=acqf_config["name"])
             return acqf
         raise AttributeError(f"No attribute named '{attr}'.")
 
+    def refresh(self):
+        self._construct_all_models()
+        self._train_all_models()
+
+    def redigest(self):
+        self.table = self.digestion(self.table, **self.digestion_kwargs)
+
     def sample(self, n: int = DEFAULT_MAX_SAMPLES, method: str = "quasi-random") -> torch.Tensor:
         """
         Returns a (..., 1, n_active_dofs) tensor of points sampled within the parameter space.
@@ -272,7 +283,9 @@ def ask(self, acqf="qei", n=1, route=True, sequential=True, upsample=1, **acqf_k
         duration = 1e3 * (ttime.monotonic() - start_time)
 
         if route and n > 1:
-            routing_index = utils.route(self.dofs.subset(active=True, read_only=False).readback, points)
+            current_points = np.array([dof.readback for dof in active_dofs if not dof.read_only])
+            travel_expenses = np.array([dof.travel_expense for dof in active_dofs if not dof.read_only])
+            routing_index = utils.route(current_points, points, dim_weights=travel_expenses)
             points = points[routing_index]
 
         if upsample > 1:
@@ -479,8 +492,7 @@ def acquire(self, acquisition_inputs):
                 [*self.dets, *self.dofs.devices],
                 delay=self.trigger_delay,
             )
-
-            products = self.digestion(self.db, uid)
+            products = self.digestion(self.db[uid].table(), **self.digestion_kwargs)
 
         except KeyboardInterrupt as interrupt:
             raise interrupt
@@ -549,7 +561,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(torch.tensor(x))
+        return self.model.posterior(self.dofs.transform(torch.tensor(x)))
 
     @property
     def fitness_model(self):
diff --git a/src/blop/bayesian/acquisition/__init__.py b/src/blop/bayesian/acquisition/__init__.py
index 65b9c41..8cefc2a 100644
--- a/src/blop/bayesian/acquisition/__init__.py
+++ b/src/blop/bayesian/acquisition/__init__.py
@@ -1,23 +1,27 @@
 import os
 
-import yaml
 import pandas as pd
-from botorch.utils.transforms import normalize
+import yaml
 
 from . import analytic, monte_carlo
 from .analytic import *  # noqa F401
 from .monte_carlo import *  # noqa F401
 
+# from botorch.utils.transforms import normalize
+
+
 here, this_filename = os.path.split(__file__)
 
 with open(f"{here}/config.yml", "r") as f:
     config = yaml.safe_load(f)
 
+
 def all_acqfs(columns=["identifier", "type", "multitask_only", "description"]):
     acqfs = pd.DataFrame(config).T[columns]
     acqfs.index.name = "name"
     return acqfs.sort_values(["type", "name"])
 
+
 def parse_acqf_identifier(identifier, strict=True):
     for acqf_name in config.keys():
         if identifier.lower() in [acqf_name, config[acqf_name]["identifier"]]:
@@ -26,6 +30,7 @@ def parse_acqf_identifier(identifier, strict=True):
         raise ValueError(f"'{identifier}' is not a valid acquisition function identifier.")
     return None
 
+
 def _construct_acqf(agent, acqf_name, **acqf_kwargs):
     """Generates an acquisition function from a supplied identifier. A list of acquisition functions and
     their identifiers can be found at `agent.all_acqfs`.
@@ -38,7 +43,6 @@ def _construct_acqf(agent, acqf_name, **acqf_kwargs):
 
     # there is probably a better way to structure this
     if acqf_name == "expected_improvement":
-
         acqf_kwargs["best_f"] = agent.best_f(weights="default")
 
         acqf = analytic.ConstrainedLogExpectedImprovement(
@@ -49,7 +53,6 @@ def _construct_acqf(agent, acqf_name, **acqf_kwargs):
         )
 
     elif acqf_name == "monte_carlo_expected_improvement":
-
         acqf_kwargs["best_f"] = agent.best_f(weights="default")
 
         acqf = monte_carlo.qConstrainedExpectedImprovement(
@@ -60,7 +63,6 @@ def _construct_acqf(agent, acqf_name, **acqf_kwargs):
         )
 
     elif acqf_name == "probability_of_improvement":
-
         acqf_kwargs["best_f"] = agent.best_f(weights="default")
 
         acqf = analytic.ConstrainedLogProbabilityOfImprovement(
@@ -86,15 +88,15 @@ def _construct_acqf(agent, acqf_name, **acqf_kwargs):
         )
 
     elif acqf_name == "monte_carlo_noisy_expected_hypervolume_improvement":
-
         acqf_kwargs["ref_point"] = acqf_kwargs.get("ref_point", agent.random_ref_point)
 
         acqf = monte_carlo.qConstrainedNoisyExpectedHypervolumeImprovement(
             constraint=agent.constraint,
             model=agent.fitness_model,
-            X_baseline=agent.input_normalization.forward(agent.train_inputs()),
+            # X_baseline=agent.input_normalization.forward(agent.train_inputs())[],
+            X_baseline=agent.dofs.transform(agent.train_inputs(active=True)),
             prune_baseline=True,
-            **acqf_kwargs
+            **acqf_kwargs,
         )
 
     elif acqf_name == "upper_confidence_bound":
diff --git a/src/blop/digestion.py b/src/blop/digestion.py
index 05805a0..bdfd1b3 100644
--- a/src/blop/digestion.py
+++ b/src/blop/digestion.py
@@ -1,3 +1,5 @@
-def default_digestion_function(db, uid):
-    products = db[uid].table(fill=True)
-    return products
+import pandas as pd
+
+
+def default_digestion_function(df: pd.DataFrame) -> pd.DataFrame:
+    return df
diff --git a/src/blop/dofs.py b/src/blop/dofs.py
index ecbf484..2f6c185 100644
--- a/src/blop/dofs.py
+++ b/src/blop/dofs.py
@@ -15,14 +15,14 @@
     "description": "str",
     "readback": "object",
     "type": "str",
+    "units": "str",
+    "tags": "object",
     "transform": "str",
     "search_domain": "object",
     "trust_domain": "object",
     "domain": "object",
     "active": "bool",
     "read_only": "bool",
-    "units": "str",
-    "tags": "object",
 }
 
 DOF_TYPES = ["continuous", "binary", "ordinal", "categorical"]
@@ -126,14 +126,15 @@ class DOF:
     name: str = None
     description: str = ""
     type: str = None
+    transform: str = None
     search_domain: Union[Tuple[float, float], Sequence] = None
     trust_domain: Union[Tuple[float, float], Sequence] = None
     units: str = None
-    read_only: bool = False
     active: bool = True
-    transform: str = None
+    read_only: bool = False
     tags: list = field(default_factory=list)
     device: Signal = None
+    travel_expense: float = 1
 
     def __repr__(self):
         nodef_f_vals = ((f.name, attrgetter(f.name)(self)) for f in fields(self))
@@ -221,18 +222,6 @@ def __post_init__(self):
         # all dof degrees of freedom are hinted
         self.device.kind = "hinted"
 
-    @property
-    def domain(self):
-        """
-        The total domain of the DOF.
-        """
-        if self.transform is None:
-            if self.type == "continuous":
-                return (-np.inf, np.inf)
-            else:
-                return self.search_domain
-        return SUPPORTED_DOF_TRANSFORMS[self.transform]
-
     @property
     def _search_domain(self):
         """
@@ -364,8 +353,10 @@ def __call__(self, *args, **kwargs):
 
     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 all([hasattr(dof, attr) for dof in self.dofs]):
+            if DOF_FIELD_TYPES.get(attr) in ["float", "int", "bool"]:
+                return np.array([getattr(dof, attr) for dof in self.dofs])
+            return [getattr(dof, attr) for dof in self.dofs]
         if attr in self.names:
             return self.__getitem__(attr)
 
diff --git a/src/blop/objectives.py b/src/blop/objectives.py
index 368411d..9d35d7e 100644
--- a/src/blop/objectives.py
+++ b/src/blop/objectives.py
@@ -13,7 +13,6 @@
 
 OBJ_FIELD_TYPES = {
     "description": "object",
-    # "kind": "str",
     "type": "str",
     "target": "object",
     "transform": "str",
@@ -292,8 +291,10 @@ def names(self):
 
     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(), "kind"]:
-            return np.array([getattr(obj, attr) for obj in self.objectives])
+        if all([hasattr(obj, attr) for obj in self.objectives]):
+            if OBJ_FIELD_TYPES.get(attr) in ["float", "int", "bool"]:
+                return np.array([getattr(obj, attr) for obj in self.objectives])
+            return [getattr(obj, attr) for obj in self.objectives]
         if attr in self.names:
             return self.__getitem__(attr)
 
@@ -330,10 +331,10 @@ def summary(self) -> pd.DataFrame:
         return table
 
     def __repr__(self):
-        return self.summary.__repr__()
+        return self.summary.T.__repr__()
 
     def _repr_html_(self):
-        return self.summary._repr_html_()
+        return self.summary.T._repr_html_()
 
     def add(self, objective):
         _validate_objs([*self.objectives, objective])
diff --git a/src/blop/plotting.py b/src/blop/plotting.py
index 00547ee..9878007 100644
--- a/src/blop/plotting.py
+++ b/src/blop/plotting.py
@@ -537,7 +537,7 @@ def inspect_beam(agent, index, border=None):
 
 
 def _plot_pareto_front(agent, obj_indices=(0, 1)):
-    f_objs = agent.objectives[agent.objectives.kind == "fitness"]
+    f_objs = agent.objectives.subset(kind="fitness")
     (i, j) = obj_indices
 
     if len(f_objs) < 2:
@@ -545,7 +545,7 @@ def _plot_pareto_front(agent, obj_indices=(0, 1)):
 
     fig, ax = plt.subplots(1, 1, figsize=(6, 6))
 
-    y = agent.train_targets()[:, agent.objectives.kind == "fitness"]
+    y = agent.train_targets(kind="fitness")
 
     pareto_mask = agent.pareto_mask
     constraint = agent.evaluated_constraints.all(axis=-1)
diff --git a/src/blop/tests/conftest.py b/src/blop/tests/conftest.py
index 12ef1e4..227bd2a 100644
--- a/src/blop/tests/conftest.py
+++ b/src/blop/tests/conftest.py
@@ -51,13 +51,11 @@ def agent_1d_1f(db):
     A one dimensional agent.
     """
 
-    def digestion(db, uid):
-        products = db[uid].table()
+    def digestion(df):
+        for index, entry in df.iterrows():
+            df.loc[index, "f1"] = functions.himmelblau(entry.x1, 3)
 
-        for index, entry in products.iterrows():
-            products.loc[index, "f1"] = functions.himmelblau(entry.x1, 3)
-
-        return products
+        return df
 
     dofs = DOF(description="The first DOF", name="x1", search_domain=(-5.0, 5.0))
     objectives = Objective(description="f1", name="f1", target="min")
@@ -69,6 +67,8 @@ def digestion(db, uid):
         db=db,
     )
 
+    agent.dofs.add(DOF(name="dummy", search_domain=(0, 1), read_only=False))
+
     return agent
 
 
@@ -94,6 +94,8 @@ def agent_2d_1f(db):
         tolerate_acquisition_errors=False,
     )
 
+    agent.dofs.add(DOF(name="dummy", search_domain=(0, 1), read_only=False))
+
     return agent
 
 
@@ -103,14 +105,12 @@ def agent_2d_2f(db):
     An agent minimizing two Himmelblau's functions
     """
 
-    def digestion(db, uid):
-        products = db[uid].table()
-
-        for index, entry in products.iterrows():
-            products.loc[index, "f1"] = functions.himmelblau(entry.x1, entry.x2)
-            products.loc[index, "f2"] = functions.himmelblau(entry.x2, entry.x1)
+    def digestion(df):
+        for index, entry in df.iterrows():
+            df.loc[index, "f1"] = functions.himmelblau(entry.x1, entry.x2)
+            df.loc[index, "f2"] = functions.himmelblau(entry.x2, entry.x1)
 
-        return products
+        return df
 
     dofs = [
         DOF(name="x1", search_domain=(-5.0, 5.0)),
@@ -128,6 +128,8 @@ def digestion(db, uid):
         tolerate_acquisition_errors=False,
     )
 
+    agent.dofs.add(DOF(name="dummy", search_domain=(0, 1), read_only=False))
+
     return agent
 
 
@@ -137,20 +139,18 @@ def agent_2d_2f_2c(db):
     Chankong and Haimes function from https://en.wikipedia.org/wiki/Test_functions_for_optimization
     """
 
-    def digestion(db, uid):
-        products = db[uid].table()
-
-        for index, entry in products.iterrows():
-            products.loc[index, "f1"] = (entry.x1 - 2) ** 2 + (entry.x2 - 1) + 2
-            products.loc[index, "f2"] = 9 * entry.x1 - (entry.x2 - 1) + 2
-            products.loc[index, "c1"] = entry.x1**2 + entry.x2**2
-            products.loc[index, "c2"] = entry.x1 - 3 * entry.x2 + 10
+    def digestion(df):
+        for index, entry in df.iterrows():
+            df.loc[index, "f1"] = (entry.x1 - 2) ** 2 + (entry.x2 - 1) + 2
+            df.loc[index, "f2"] = 9 * entry.x1 - (entry.x2 - 1) + 2
+            df.loc[index, "c1"] = entry.x1**2 + entry.x2**2
+            df.loc[index, "c2"] = entry.x1 - 3 * entry.x2 + 10
 
-        return products
+        return df
 
     dofs = [
-        DOF(description="The first DOF", name="x1", search_domain=(-20, 20)),
-        DOF(description="The second DOF", name="x2", search_domain=(-20, 20)),
+        DOF(description="The first DOF", name="x1", search_domain=(-20, 20), travel_expense=1.0),
+        DOF(description="The second DOF", name="x2", search_domain=(-20, 20), travel_expense=2.0),
     ]
 
     objectives = [
@@ -169,6 +169,8 @@ def digestion(db, uid):
         tolerate_acquisition_errors=False,
     )
 
+    agent.dofs.add(DOF(name="dummy", search_domain=(0, 1), read_only=False))
+
     return agent
 
 
@@ -199,6 +201,8 @@ def agent_with_read_only_dofs(db):
         tolerate_acquisition_errors=False,
     )
 
+    agent.dofs.add(DOF(name="dummy", search_domain=(0, 1), read_only=False))
+
     return agent
 
 
diff --git a/src/blop/tests/test_acq_funcs.py b/src/blop/tests/test_acq_funcs.py
index 5bf1562..a2e0c31 100644
--- a/src/blop/tests/test_acq_funcs.py
+++ b/src/blop/tests/test_acq_funcs.py
@@ -3,35 +3,47 @@
 
 @pytest.mark.parametrize("acqf", ["ei", "pi", "em", "ucb"])
 def test_analytic_acqfs_one_dimensional(agent_1d_1f, RE, acqf):
-    RE(agent_1d_1f.learn("qr", n=16))
-    RE(agent_1d_1f.learn(acqf, n=1))
+    a = agent_1d_1f
+    RE(a.learn("qr", n=16))
+    RE(a.learn(acqf, n=1))
+    getattr(a, acqf)
 
 
 @pytest.mark.parametrize("acqf", ["qei", "qpi", "qem", "qucb"])
 def test_monte_carlo_acqfs_one_dimensional(agent_1d_1f, RE, acqf):
-    RE(agent_1d_1f.learn("qr", n=16))
-    RE(agent_1d_1f.learn(acqf, n=2))
+    a = agent_1d_1f
+    RE(a.learn("qr", n=16))
+    RE(a.learn(acqf, n=2))
+    getattr(a, acqf)
 
 
 @pytest.mark.parametrize("acqf", ["ei", "pi", "em", "ucb"])
 def test_analytic_acqfs_single_objective(agent_2d_1f, RE, acqf):
-    RE(agent_2d_1f.learn("qr", n=16))
-    RE(agent_2d_1f.learn(acqf, n=1))
+    a = agent_2d_1f
+    RE(a.learn("qr", n=4))
+    RE(a.learn(acqf, n=1))
+    getattr(a, acqf)
 
 
 @pytest.mark.parametrize("acqf", ["qei", "qpi", "qem", "qucb"])
 def test_monte_carlo_acqfs_single_objective(agent_2d_1f, RE, acqf):
-    RE(agent_2d_1f.learn("qr", n=16))
-    RE(agent_2d_1f.learn(acqf, n=2))
+    a = agent_2d_1f
+    RE(a.learn("qr", n=4))
+    RE(a.learn(acqf, n=2))
+    getattr(a, acqf)
 
 
 @pytest.mark.parametrize("acqf", ["ei", "pi", "em", "ucb"])
 def test_analytic_acqfs_multi_objective(agent_2d_2f, RE, acqf):
-    RE(agent_2d_2f.learn("qr", n=16))
-    RE(agent_2d_2f.learn(acqf, n=1))
+    a = agent_2d_2f
+    RE(a.learn("qr", n=4))
+    RE(a.learn(acqf, n=1))
+    getattr(a, acqf)
 
 
 @pytest.mark.parametrize("acqf", ["qei", "qpi", "qem", "qucb"])
 def test_monte_carlo_acqfs_multi_objective(agent_2d_2f, RE, acqf):
-    RE(agent_2d_2f.learn("qr", n=16))
-    RE(agent_2d_2f.learn(acqf, n=2))
+    a = agent_2d_2f
+    RE(a.learn("qr", n=4))
+    RE(a.learn(acqf, n=2))
+    getattr(a, acqf)
diff --git a/src/blop/tests/test_agent.py b/src/blop/tests/test_agent.py
index 028a7da..78b0b1d 100644
--- a/src/blop/tests/test_agent.py
+++ b/src/blop/tests/test_agent.py
@@ -2,22 +2,37 @@
 
 
 def test_agent(agent_2d_1f, RE):
-    RE(agent_2d_1f.learn("qr", n=4))
+    a = agent_2d_1f
+    RE(a.learn("qr", n=4))
 
-    best = agent_2d_1f.best
+    best = a.best
+    assert [dof.name in best for dof in a.dofs]
+    assert [obj.name in best for obj in a.objectives]
+    assert a.dofs.x1 is a.dofs[0]
 
-    assert [dof.name in best for dof in agent_2d_1f.dofs]
-    assert [obj.name in best for obj in agent_2d_1f.objectives]
+    print(a.dofs)
+    print(a.objectives)
 
-    print(agent_2d_1f.dofs)
-    print(agent_2d_1f.objectives)
+
+def test_refresh(agent_2d_1f, RE):
+    """
+    Test that the agent can determine when the number of DOFs has changed, and adjust.
+    """
+    a = agent_2d_1f
+    RE(a.learn("qr", n=4))
+
+    RE(a.learn("qei", n=1))
+    a.dofs[2].activate()
+    RE(a.learn("qei", n=1))
 
 
 def test_forget(agent_2d_1f, RE):
-    RE(agent_2d_1f.learn("qr", n=4))
-    agent_2d_1f.forget(last=2)
+    a = agent_2d_1f
+    RE(a.learn("qr", n=4))
+    a.forget(last=2)
 
 
 def test_benchmark(agent_2d_1f, RE):
+    a = agent_2d_1f
     per_iter_learn_kwargs_list = [{"acqf": "qr", "n": 32}, {"acqf": "qei", "n": 2, "iterations": 2}]
-    RE(agent_2d_1f.benchmark(output_dir="/tmp/blop", iterations=1, per_iter_learn_kwargs_list=per_iter_learn_kwargs_list))
+    RE(a.benchmark(output_dir="/tmp/blop", iterations=1, per_iter_learn_kwargs_list=per_iter_learn_kwargs_list))
diff --git a/src/blop/tests/test_pareto.py b/src/blop/tests/test_pareto.py
index 5934300..f242168 100644
--- a/src/blop/tests/test_pareto.py
+++ b/src/blop/tests/test_pareto.py
@@ -13,6 +13,8 @@ def test_monte_carlo_pareto_acqfs(agent_2d_2f, RE, acqf):
     agent = agent_2d_2f
     RE(agent.learn("qr", n=16))
     RE(agent.learn(acqf, n=2))
+    agent.dofs[0].deactivate()
+    RE(agent.learn(acqf, n=2))
 
 
 @pytest.mark.test_func
@@ -27,3 +29,5 @@ def test_constrained_monte_carlo_pareto_acqfs(agent_2d_2f_2c, RE, acqf):
     agent = agent_2d_2f_2c
     RE(agent.learn("qr", n=16))
     RE(agent.learn(acqf, n=2))
+    agent.dofs[0].deactivate()
+    RE(agent.learn(acqf, n=2))
diff --git a/src/blop/utils/__init__.py b/src/blop/utils/__init__.py
index 0debc41..a25bf4f 100644
--- a/src/blop/utils/__init__.py
+++ b/src/blop/utils/__init__.py
@@ -48,7 +48,7 @@ def mprod(*M):
     return res
 
 
-def route(start_point, points):
+def route(start_point, points, dim_weights=1):
     """
     Returns the indices of the most efficient way to visit `points`, starting from `start_point`.
     """
@@ -60,9 +60,9 @@ def route(start_point, points):
     if dim_mask.sum() == 0:
         return np.arange(len(points))
 
-    normalized_points = (total_points - total_points.min(axis=0))[:, dim_mask] / points_scale[dim_mask]
+    scaled_points = (total_points - total_points.min(axis=0)) * (dim_weights / np.where(points_scale > 0, points_scale, 1))
 
-    delay_matrix = np.sqrt(np.square(normalized_points[:, None, :] - normalized_points[None, :, :]).sum(axis=-1))
+    delay_matrix = np.sqrt(np.square(scaled_points[:, None, :] - scaled_points[None, :, :]).sum(axis=-1))
     delay_matrix = (1e4 * delay_matrix).astype(int)  # it likes integers idk
 
     manager = pywrapcp.RoutingIndexManager(len(total_points), 1, 0)  # number of depots, number of salesmen, starting index
diff --git a/src/blop/utils/functions.py b/src/blop/utils/functions.py
index 553d79d..1323fa8 100644
--- a/src/blop/utils/functions.py
+++ b/src/blop/utils/functions.py
@@ -1,4 +1,5 @@
 import numpy as np
+import pandas as pd
 import torch
 
 
@@ -60,18 +61,17 @@ def binh_korn(x1, x2):
     return np.where(c, f1, np.nan), np.where(c, f2, np.nan)
 
 
-def binh_korn_digestion(db, uid):
+def binh_korn_digestion(df: pd.DataFrame) -> pd.DataFrame:
     """
     Digests Himmelblau's function into the feedback.
     """
-    products = db[uid].table()
 
-    for index, entry in products.iterrows():
+    for index, entry in df.iterrows():
         f1, f2 = binh_korn(entry.x1, entry.x2)
-        products.loc[index, "f1"] = f1
-        products.loc[index, "f2"] = f2
+        df.loc[index, "f1"] = f1
+        df.loc[index, "f2"] = f2
 
-    return products
+    return df
 
 
 def skewed_himmelblau(x1, x2):
@@ -184,42 +184,38 @@ def kb_tradeoff_4d(x1, x2, x3, x4):
     return x_width, y_width, flux
 
 
-def constrained_himmelblau_digestion(db, uid):
+def constrained_himmelblau_digestion(df: pd.DataFrame) -> pd.DataFrame:
     """
     Digests Himmelblau's function into the feedback.
     """
-    products = db[uid].table()
 
-    for index, entry in products.iterrows():
-        products.loc[index, "himmelblau"] = constrained_himmelblau(entry.x1, entry.x2)
+    for index, entry in df.iterrows():
+        df.loc[index, "himmelblau"] = constrained_himmelblau(entry.x1, entry.x2)
 
-    return products
+    return df
 
 
-def himmelblau_digestion(db, uid):
+def himmelblau_digestion(df: pd.DataFrame) -> pd.DataFrame:
     """
     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)
+    for index, entry in df.iterrows():
+        df.loc[index, "himmelblau"] = himmelblau(entry.x1, entry.x2)
 
-    return products
+    return df
 
 
-def mock_kbs_digestion(db, uid):
+def mock_kbs_digestion(df: pd.DataFrame) -> pd.DataFrame:
     """
     Digests a beam waist and height into the feedback.
     """
 
-    products = db[uid].table()
-
-    for index, entry in products.iterrows():
+    for index, entry in df.iterrows():
         sigma_x = gaussian_beam_waist(entry.x1, entry.x2)
         sigma_y = gaussian_beam_waist(entry.x3, entry.x4)
 
-        products.loc[index, "x_width"] = 2 * sigma_x
-        products.loc[index, "y_width"] = 2 * sigma_y
+        df.loc[index, "x_width"] = 2 * sigma_x
+        df.loc[index, "y_width"] = 2 * sigma_y
 
-    return products
+    return df

From e0f312c7be1eab0c0731740ae1e6df8e08edfc53 Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Tue, 23 Apr 2024 14:34:14 -0700
Subject: [PATCH 06/11] rebased on main

---
 src/blop/agent.py                         | 74 +++++++++++------------
 src/blop/bayesian/acquisition/__init__.py |  2 +-
 src/blop/plotting.py                      | 26 ++++----
 3 files changed, 51 insertions(+), 51 deletions(-)

diff --git a/src/blop/agent.py b/src/blop/agent.py
index 2f9aa69..b38a13f 100644
--- a/src/blop/agent.py
+++ b/src/blop/agent.py
@@ -186,7 +186,7 @@ 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)
+        active_dofs = self.dofs(active=True)
 
         if method == "quasi-random":
             X = utils.normalized_sobol_sampler(n, d=len(active_dofs))
@@ -204,7 +204,7 @@ def sample(self, n: int = DEFAULT_MAX_SAMPLES, method: str = "quasi-random") ->
         else:
             raise ValueError("'method' argument must be one of ['quasi-random', 'random', 'grid'].")
 
-        return self.dofs.subset(active=True).untransform(X)
+        return self.dofs(active=True).untransform(X)
 
     def ask(self, acqf="qei", n=1, route=True, sequential=True, upsample=1, **acqf_kwargs):
         """Ask the agent for the best point to sample, given an acquisition function.
@@ -228,8 +228,8 @@ def ask(self, acqf="qei", n=1, route=True, sequential=True, upsample=1, **acqf_k
 
         start_time = ttime.monotonic()
 
-        active_dofs = self.dofs.subset(active=True)
-        active_objs = self.objectives.subset(active=True)
+        active_dofs = self.dofs(active=True)
+        active_objs = self.objectives(active=True)
 
         # these are the fake acquisiton functions that we don't need to construct
         if acqf_config["name"] in ["quasi-random", "random", "grid"]:
@@ -271,11 +271,11 @@ def ask(self, acqf="qei", n=1, route=True, sequential=True, upsample=1, **acqf_k
 
             # this includes both RO and non-RO DOFs.
             # and is in the transformed model space
-            candidates = self.dofs.subset(active=True).untransform(candidates).numpy()
+            candidates = self.dofs(active=True).untransform(candidates).numpy()
 
         p = self.posterior(candidates) if hasattr(self, "model") else None
 
-        active_dofs = self.dofs.subset(active=True)
+        active_dofs = self.dofs(active=True)
 
         points = candidates[..., ~active_dofs.read_only]
         read_only_values = candidates[..., active_dofs.read_only]
@@ -354,7 +354,7 @@ def tell(
         self.table.index = np.arange(len(self.table))
 
         if update_models:
-            for obj in self.objectives.subset(active=True):
+            for obj in self.objectives(active=True):
                 t0 = ttime.monotonic()
 
                 cached_hypers = obj.model.state_dict() if hasattr(obj, "model") else None
@@ -414,7 +414,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).search_domain.mean(axis=1))
+            center_inputs = np.atleast_2d(self.dofs(active=True, read_only=False).search_domain.mean(axis=1))
             new_table = yield from self.acquire(center_inputs)
             new_table.loc[:, "acqf"] = "sample_center_on_init"
 
@@ -448,7 +448,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.subset(active=True):
+            for obj in self.objectives(active=True):
                 p = obj.model.posterior(test_grid)
 
                 if item == "mean":
@@ -485,7 +485,7 @@ def acquire(self, acquisition_inputs):
             raise ValueError("Cannot run acquistion without databroker instance!")
 
         try:
-            acquisition_devices = self.dofs.subset(active=True, read_only=False).devices
+            acquisition_devices = self.dofs(active=True, read_only=False).devices
             uid = yield from self.acquisition_plan(
                 acquisition_devices,
                 acquisition_inputs.astype(float),
@@ -501,8 +501,8 @@ def acquire(self, acquisition_inputs):
             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.subset(active=True):
+            products = pd.DataFrame(acquisition_inputs, columns=self.dofs(active=True, read_only=False).names)
+            for obj in self.objectives(active=True):
                 products.loc[:, obj.name] = np.nan
 
         if not len(acquisition_inputs) == len(products):
@@ -521,7 +521,7 @@ def reset(self):
         """Reset the agent."""
         self.table = pd.DataFrame()
 
-        for obj in self.objectives.subset(active=True):
+        for obj in self.objectives(active=True):
             if hasattr(obj, "model"):
                 del obj.model
 
@@ -556,16 +556,16 @@ def benchmark(
     @property
     def model(self):
         """A model encompassing all the fitnesses and constraints."""
-        active_objs = self.objectives.subset(active=True)
+        active_objs = self.objectives(active=True)
         return ModelListGP(*[obj.model for obj in active_objs]) if len(active_objs) > 1 else 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(self.dofs.transform(torch.tensor(x)))
+        return self.model.posterior(self.dofs(active=True).transform(torch.tensor(x)))
 
     @property
     def fitness_model(self):
-        active_fitness_models = self.objectives.subset(active=True, kind="fitness")
+        active_fitness_models = self.objectives(active=True, kind="fitness")
         if len(active_fitness_models) == 0:
             raise ValueError("Having no fitness objectives is unhandled.")
         if len(active_fitness_models) == 1:
@@ -574,14 +574,14 @@ def fitness_model(self):
 
     @property
     def evaluated_constraints(self):
-        constraint_objectives = self.objectives.subset(kind="constraint")
+        constraint_objectives = self.objectives(kind="constraint")
         if len(constraint_objectives):
             return torch.cat([obj.constrain(self.raw_targets(obj.name)) for obj in constraint_objectives], dim=-1)
         else:
             return torch.ones(size=(len(self.table), 0), dtype=torch.bool)
 
     def fitness_scalarization(self, weights="default"):
-        fitness_objectives = self.objectives.subset(active=True, kind="fitness")
+        fitness_objectives = self.objectives(active=True, kind="fitness")
         if weights == "default":
             weights = torch.tensor([obj.weight for obj in fitness_objectives], dtype=torch.double)
         elif weights == "equal":
@@ -682,7 +682,7 @@ def _construct_model(self, obj, skew_dims=None):
             outcome_transform=outcome_transform,
         )
 
-        obj.model_dofs = set(self.dofs.subset(active=True).names)  # if these change, retrain the model on self.ask()
+        obj.model_dofs = set(self.dofs(active=True).names)  # if these change, retrain the model on self.ask()
 
         if trusted.all():
             obj.validity_conjugate_model = None
@@ -707,13 +707,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.subset(active=True):
+        for obj in self.objectives(active=True):
             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.subset(active=True):
+        for obj in self.objectives(active=True):
             self._train_model(obj.model)
             if obj.validity_conjugate_model is not None:
                 self._train_model(obj.validity_conjugate_model)
@@ -744,7 +744,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.dofs(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:
@@ -754,13 +754,13 @@ def _latent_dim_tuples(self, obj_index=None):
         return [tuple(np.where(uinv == i)[0]) for i in range(len(u))]
 
     @property
-    def _sample_domain(self):
+    def sample_domain(self):
         """
         Returns a (2, n_active_dof) array of lower and upper bounds for dofs.
         Read-only DOFs are set to exactly their last known value.
         Discrete DOFs are relaxed to some continuous domain.
         """
-        return self.dofs.subset(active=True).transform(self.dofs.subset(active=True).search_domain.T)
+        return self.dofs(active=True).transform(self.dofs(active=True).search_domain.T)
 
     @property
     def input_normalization(self):
@@ -816,15 +816,15 @@ 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.subset(active=True):
+        for obj in self.objectives(active=True):
             obj.model.load_state_dict(hypers[obj.name])
         self.validity_constraint.load_state_dict(hypers["validity_constraint"])
 
     def constraint(self, x):
-        x = self.dofs.subset(active=True).transform(x)
+        x = self.dofs(active=True).transform(x)
 
         p = torch.ones(x.shape[:-1])
-        for obj in self.objectives.subset(active=True):
+        for obj in self.objectives(active=True):
             # if the targeting constraint is non-trivial
             # if obj.kind == "constraint":
             #     p *= obj.targeting_constraint(x)
@@ -892,14 +892,14 @@ def raw_inputs(self, index=None, **subset_kwargs):
         Get the raw, untransformed inputs for a DOF (or for a subset).
         """
         if index is None:
-            return torch.cat([self.raw_inputs(dof.name) for dof in self.dofs.subset(**subset_kwargs)], dim=-1)
+            return torch.cat([self.raw_inputs(dof.name) for dof in self.dofs(**subset_kwargs)], dim=-1)
         return torch.tensor(self.table.loc[:, self.dofs[index].name].values, dtype=torch.double).unsqueeze(-1)
 
     def train_inputs(self, index=None, **subset_kwargs):
         """A two-dimensional tensor of all DOF values."""
 
         if index is None:
-            return torch.cat([self.train_inputs(index=dof.name) for dof in self.dofs.subset(**subset_kwargs)], dim=-1)
+            return torch.cat([self.train_inputs(index=dof.name) for dof in self.dofs(**subset_kwargs)], dim=-1)
 
         dof = self.dofs[index]
         raw_inputs = self.raw_inputs(index=index, **subset_kwargs)
@@ -915,14 +915,14 @@ def raw_targets(self, index=None, **subset_kwargs):
         Get the raw, untransformed inputs for an objective (or for a subset).
         """
         if index is None:
-            return torch.cat([self.raw_targets(index=obj.name) for obj in self.objectives.subset(**subset_kwargs)], dim=-1)
+            return torch.cat([self.raw_targets(index=obj.name) for obj in self.objectives(**subset_kwargs)], dim=-1)
         return torch.tensor(self.table.loc[:, self.objectives[index].name].values, dtype=torch.double).unsqueeze(-1)
 
     def train_targets(self, index=None, **subset_kwargs):
         """Returns the values associated with an objective name."""
 
         if index is None:
-            return torch.cat([self.train_targets(obj.name) for obj in self.objectives.subset(**subset_kwargs)], dim=-1)
+            return torch.cat([self.train_targets(obj.name) for obj in self.objectives(**subset_kwargs)], dim=-1)
 
         obj = self.objectives[index]
         raw_targets = self.raw_targets(index=index, **subset_kwargs)
@@ -973,10 +973,10 @@ def plot_objectives(self, axes: Tuple = (0, 1), **kwargs):
             A tuple specifying which DOFs to plot as a function of. Can be either an int or the name of DOFs.
         """
 
-        if len(self.dofs.subset(active=True, read_only=False)) == 1:
-            if len(self.objectives.subset(active=True, kind="fitness")) > 0:
+        if len(self.dofs(active=True, read_only=False)) == 1:
+            if len(self.objectives(active=True, kind="fitness")) > 0:
                 plotting._plot_fitness_objs_one_dof(self, **kwargs)
-            if len(self.objectives.subset(active=True, kind="constraint")) > 0:
+            if len(self.objectives(active=True, kind="constraint")) > 0:
                 plotting._plot_constraint_objs_one_dof(self, **kwargs)
         else:
             plotting._plot_objs_many_dofs(self, axes=axes, **kwargs)
@@ -991,7 +991,7 @@ def plot_acquisition(self, acqf="ei", axes: Tuple = (0, 1), **kwargs):
         axes :
             A tuple specifying which DOFs to plot as a function of. Can be either an int or the name of DOFs.
         """
-        if len(self.dofs.subset(active=True, read_only=False)) == 1:
+        if len(self.dofs(active=True, read_only=False)) == 1:
             plotting._plot_acqf_one_dof(self, acqfs=np.atleast_1d(acqf), **kwargs)
 
         else:
@@ -1005,7 +1005,7 @@ def plot_validity(self, axes: Tuple = (0, 1), **kwargs):
         axes :
             A tuple specifying which DOFs to plot as a function of. Can be either an int or the name of DOFs.
         """
-        if len(self.dofs.subset(active=True, read_only=False)) == 1:
+        if len(self.dofs(active=True, read_only=False)) == 1:
             plotting._plot_valid_one_dof(self, **kwargs)
 
         else:
@@ -1017,7 +1017,7 @@ def plot_history(self, **kwargs):
 
     @property
     def latent_transforms(self):
-        return {obj.name: obj.model.covar_module.latent_transform for obj in self.objectives.subset(active=True)}
+        return {obj.name: obj.model.covar_module.latent_transform for obj in self.objectives(active=True)}
 
     def plot_pareto_front(self, **kwargs):
         """Plot the improvement of the agent over time."""
diff --git a/src/blop/bayesian/acquisition/__init__.py b/src/blop/bayesian/acquisition/__init__.py
index 8cefc2a..e8e4657 100644
--- a/src/blop/bayesian/acquisition/__init__.py
+++ b/src/blop/bayesian/acquisition/__init__.py
@@ -94,7 +94,7 @@ def _construct_acqf(agent, acqf_name, **acqf_kwargs):
             constraint=agent.constraint,
             model=agent.fitness_model,
             # X_baseline=agent.input_normalization.forward(agent.train_inputs())[],
-            X_baseline=agent.dofs.transform(agent.train_inputs(active=True)),
+            X_baseline=agent.dofs(active=True).transform(agent.train_inputs(active=True)),
             prune_baseline=True,
             **acqf_kwargs,
         )
diff --git a/src/blop/plotting.py b/src/blop/plotting.py
index 9878007..bbeaa19 100644
--- a/src/blop/plotting.py
+++ b/src/blop/plotting.py
@@ -14,7 +14,7 @@
 
 
 def _plot_fitness_objs_one_dof(agent, size=16, lw=1e0):
-    fitness_objs = agent.objectives.subset(kind="fitness")
+    fitness_objs = agent.objectives(kind="fitness")
 
     agent.obj_fig, agent.obj_axes = plt.subplots(
         len(fitness_objs),
@@ -26,11 +26,11 @@ def _plot_fitness_objs_one_dof(agent, size=16, lw=1e0):
 
     agent.obj_axes = np.atleast_1d(agent.obj_axes)
 
-    x_dof = agent.dofs.subset(active=True)[0]
+    x_dof = agent.dofs(active=True)[0]
     x_values = agent.table.loc[:, x_dof.device.name].values
 
     test_inputs = agent.sample(method="grid")
-    test_model_inputs = agent.dofs.transform(test_inputs)
+    test_model_inputs = agent.dofs(active=True).transform(test_inputs)
 
     for obj_index, obj in enumerate(fitness_objs):
         obj_values = agent.train_targets(obj.name).squeeze(-1).numpy()
@@ -62,7 +62,7 @@ def _plot_fitness_objs_one_dof(agent, size=16, lw=1e0):
 
 
 def _plot_constraint_objs_one_dof(agent, size=16, lw=1e0):
-    constraint_objs = agent.objectives.subset(kind="constraint")
+    constraint_objs = agent.objectives(kind="constraint")
 
     agent.obj_fig, agent.obj_axes = plt.subplots(
         len(constraint_objs),
@@ -74,11 +74,11 @@ def _plot_constraint_objs_one_dof(agent, size=16, lw=1e0):
 
     agent.obj_axes = np.atleast_2d(agent.obj_axes)
 
-    x_dof = agent.dofs.subset(active=True)[0]
+    x_dof = agent.dofs(active=True)[0]
     x_values = agent.table.loc[:, x_dof.device.name].values
 
     test_inputs = agent.sample(method="grid")
-    test_model_inputs = agent.dofs.transform(test_inputs)
+    test_model_inputs = agent.dofs(active=True).transform(test_inputs)
 
     for obj_index, obj in enumerate(constraint_objs):
         val_ax = agent.obj_axes[obj_index, 0]
@@ -129,7 +129,7 @@ def _plot_objs_many_dofs(agent, axes=(0, 1), shading="nearest", cmap=DEFAULT_COL
     Axes represents which active, non-read-only axes to plot with
     """
 
-    plottable_dofs = agent.dofs.subset(active=True, read_only=False)
+    plottable_dofs = agent.dofs(active=True, read_only=False)
 
     if gridded is None:
         gridded = len(plottable_dofs) == 2
@@ -155,7 +155,7 @@ def _plot_objs_many_dofs(agent, axes=(0, 1), shading="nearest", cmap=DEFAULT_COL
     test_x = test_inputs[..., 0, axes[0]].detach().squeeze().numpy()
     test_y = test_inputs[..., 0, axes[1]].detach().squeeze().numpy()
 
-    model_inputs = agent.dofs.subset(active=True).transform(test_inputs)
+    model_inputs = agent.dofs(active=True).transform(test_inputs)
 
     for obj_index, obj in enumerate(agent.objectives):
         targets = agent.train_targets(obj.name).squeeze(-1).numpy()
@@ -324,7 +324,7 @@ def _plot_acqf_one_dof(agent, acqfs, lw=1e0, **kwargs):
     )
 
     agent.acq_axes = np.atleast_1d(agent.acq_axes)
-    x_dof = agent.dofs.subset(active=True)[0]
+    x_dof = agent.dofs(active=True)[0]
 
     test_inputs = agent.sample(method="grid")
 
@@ -355,7 +355,7 @@ def _plot_acqf_many_dofs(
         constrained_layout=True,
     )
 
-    plottable_dofs = agent.dofs.subset(active=True, read_only=False)
+    plottable_dofs = agent.dofs(active=True, read_only=False)
 
     if gridded is None:
         gridded = len(plottable_dofs) == 2
@@ -412,7 +412,7 @@ def _plot_acqf_many_dofs(
 def _plot_valid_one_dof(agent, size=16, lw=1e0):
     agent.valid_fig, agent.valid_ax = plt.subplots(1, 1, figsize=(6, 4 * len(agent.objectives)), constrained_layout=True)
 
-    x_dof = agent.dofs.subset(active=True)[0]
+    x_dof = agent.dofs(active=True)[0]
     x_values = agent.table.loc[:, x_dof.device.name].values
 
     test_inputs = agent.sample(method="grid")
@@ -426,7 +426,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), constrained_layout=True)
 
-    plottable_dofs = agent.dofs.subset(active=True, read_only=False)
+    plottable_dofs = agent.dofs(active=True, read_only=False)
 
     if gridded is None:
         gridded = len(plottable_dofs) == 2
@@ -537,7 +537,7 @@ def inspect_beam(agent, index, border=None):
 
 
 def _plot_pareto_front(agent, obj_indices=(0, 1)):
-    f_objs = agent.objectives.subset(kind="fitness")
+    f_objs = agent.objectives(kind="fitness")
     (i, j) = obj_indices
 
     if len(f_objs) < 2:

From a43f7f36807980283394b03b5aa5ccafc7402760 Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Wed, 24 Apr 2024 14:07:04 -0700
Subject: [PATCH 07/11] better test coverage

---
 src/blop/agent.py                   |  45 +++--
 src/blop/objectives.py              |  33 +---
 src/blop/tests/conftest.py          | 250 ++++++++++------------------
 src/blop/tests/test_acq_funcs.py    |  49 ------
 src/blop/tests/test_acqfs.py        |  21 +++
 src/blop/tests/test_agent.py        |  38 -----
 src/blop/tests/test_agents.py       |  54 ++++++
 src/blop/tests/test_napari.py       |   7 -
 src/blop/tests/test_pareto.py       |  30 +---
 src/blop/tests/test_passive_dofs.py |   8 -
 src/blop/tests/test_plots.py        |  37 ----
 src/blop/tests/test_read_write.py   |  16 --
 src/blop/utils/functions.py         |  33 +++-
 13 files changed, 236 insertions(+), 385 deletions(-)
 delete mode 100644 src/blop/tests/test_acq_funcs.py
 create mode 100644 src/blop/tests/test_acqfs.py
 delete mode 100644 src/blop/tests/test_agent.py
 create mode 100644 src/blop/tests/test_agents.py
 delete mode 100644 src/blop/tests/test_napari.py
 delete mode 100644 src/blop/tests/test_passive_dofs.py
 delete mode 100644 src/blop/tests/test_plots.py
 delete mode 100644 src/blop/tests/test_read_write.py

diff --git a/src/blop/agent.py b/src/blop/agent.py
index b38a13f..18a81d8 100644
--- a/src/blop/agent.py
+++ b/src/blop/agent.py
@@ -273,7 +273,7 @@ def ask(self, acqf="qei", n=1, route=True, sequential=True, upsample=1, **acqf_k
             # and is in the transformed model space
             candidates = self.dofs(active=True).untransform(candidates).numpy()
 
-        p = self.posterior(candidates) if hasattr(self, "model") else None
+        # p = self.posterior(candidates) if hasattr(self, "model") else None
 
         active_dofs = self.dofs(active=True)
 
@@ -304,7 +304,7 @@ def ask(self, acqf="qei", n=1, route=True, sequential=True, upsample=1, **acqf_k
             "sequential": sequential,
             "upsample": upsample,
             "read_only_values": read_only_values,
-            "posterior": p,
+            # "posterior": p,
         }
 
         return res
@@ -426,9 +426,11 @@ def learn(
                 new_table = yield from self.acquire(res["points"])
                 new_table.loc[:, "acqf"] = res["acqf_name"]
 
-                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")
+                x = {key: new_table.loc[:, key].tolist() for key in self.dofs.names}
+                y = {key: new_table.loc[:, key].tolist() for key in self.objectives.names}
+                metadata = {
+                    key: new_table.loc[:, key].tolist() for key in new_table.columns if (key not in x) and (key not in y)
+                }
                 self.tell(x=x, y=y, metadata=metadata, append=append, train=train)
 
     def view(self, item: str = "mean", cmap: str = "turbo", max_inputs: int = 2**16):
@@ -510,12 +512,10 @@ def acquire(self, acquisition_inputs):
 
         return products
 
-    def load_data(self, data_file, append=True, train=True):
+    def load_data(self, data_file, append=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)
+        self.table = pd.concat([self.table, new_table]) if append else new_table
+        self.refresh()
 
     def reset(self):
         """Reset the agent."""
@@ -557,7 +557,9 @@ def benchmark(
     def model(self):
         """A model encompassing all the fitnesses and constraints."""
         active_objs = self.objectives(active=True)
-        return ModelListGP(*[obj.model for obj in active_objs]) if len(active_objs) > 1 else active_objs[0].model
+        if all(hasattr(obj, "model") for obj in active_objs):
+            return ModelListGP(*[obj.model for obj in active_objs]) if len(active_objs) > 1 else active_objs[0].model
+        raise ValueError("Not all active objectives have models.")
 
     def posterior(self, x):
         """A model encompassing all the objectives. A single GP in the single-objective case, or a model list."""
@@ -567,7 +569,7 @@ def posterior(self, x):
     def fitness_model(self):
         active_fitness_models = self.objectives(active=True, kind="fitness")
         if len(active_fitness_models) == 0:
-            raise ValueError("Having no fitness objectives is unhandled.")
+            return GenericDeterministicModel(f=lambda x: torch.ones(x.shape[:-1]).unsqueeze(-1))
         if len(active_fitness_models) == 1:
             return active_fitness_models[0].model
         return ModelListGP(*[obj.model for obj in active_fitness_models])
@@ -594,12 +596,21 @@ def fitness_scalarization(self, weights="default"):
         return ScalarizedPosteriorTransform(weights=weights)
 
     def scalarized_fitnesses(self, weights="default", constrained=True):
-        f = self.fitness_scalarization(weights=weights).evaluate(self.train_targets(active=True, kind="fitness"))
+        """
+        Return the scalar fitness for each sample, scalarized by the weighting scheme.
+
+        If constrained=True, the points that satisfy the most constraints are automatically better than the others.
+        """
+        fitness_objs = self.objectives(kind="fitness")
+        if len(fitness_objs) >= 1:
+            f = self.fitness_scalarization(weights=weights).evaluate(self.train_targets(active=True, kind="fitness"))
+        else:
+            f = torch.zeros(len(self.table), dtype=torch.double)
         if constrained:
-            c = self.evaluated_constraints.all(axis=-1)
-            if not c.sum():
-                raise ValueError("There are no valid points that satisfy the constraints!")
-        return torch.where(c, f, -np.inf)
+            # how many constraints are satisfied?
+            c = self.evaluated_constraints.sum(axis=-1)
+            f = torch.where(c < c.max(), -np.inf, f)
+        return f
 
     def argmax_best_f(self, weights="default"):
         return int(self.scalarized_fitnesses(weights=weights, constrained=True).argmax())
diff --git a/src/blop/objectives.py b/src/blop/objectives.py
index 9d35d7e..5f7ae71 100644
--- a/src/blop/objectives.py
+++ b/src/blop/objectives.py
@@ -12,6 +12,7 @@
 DEFAULT_MAX_NOISE_LEVEL = 1e0
 
 OBJ_FIELD_TYPES = {
+    "name": "str",
     "description": "object",
     "type": "str",
     "target": "object",
@@ -35,14 +36,6 @@ class DuplicateNameError(ValueError):
     ...
 
 
-def _validate_objs(objs):
-    names = [obj.name for obj in objs]
-    unique_names, counts = np.unique(names, return_counts=True)
-    duplicate_names = unique_names[counts > 1]
-    if len(duplicate_names) > 0:
-        raise DuplicateNameError(f"Duplicate name(s) in supplied objectives: {duplicate_names}")
-
-
 domains = {"log"}
 
 
@@ -238,22 +231,6 @@ def targeting_constraint(self, x: torch.Tensor) -> torch.Tensor:
             0.5 * (approximate_erf((b - m) / (np.sqrt(2) * sish)) - approximate_erf((a - m) / (np.sqrt(2) * sish)))[..., -1]
         )
 
-    # def fitness_forward(self, y):
-    #     f = y
-    #     if self.log:
-    #         f = np.log(f)
-    #     if self.target == "min":
-    #         f = -f
-    #     return f
-
-    def fitness_inverse(self, f):
-        y = f
-        if self.target == "min":
-            y = -y
-        if self.log:
-            y = np.exp(y)
-        return y
-
     @property
     def is_fitness(self):
         return self.target in ["min", "max"]
@@ -279,7 +256,6 @@ def fitness_prediction(self, X):
 
 class ObjectiveList(Sequence):
     def __init__(self, objectives: list = []):
-        _validate_objs(objectives)
         self.objectives = objectives
 
     def __call__(self, *args, **kwargs):
@@ -319,11 +295,11 @@ def __len__(self):
 
     @property
     def summary(self) -> pd.DataFrame:
-        table = pd.DataFrame(columns=list(OBJ_FIELD_TYPES.keys()), index=self.names)
+        table = pd.DataFrame(columns=list(OBJ_FIELD_TYPES.keys()), index=np.arange(len(self)))
 
-        for obj in self.objectives:
+        for index, obj in enumerate(self.objectives):
             for attr, value in obj.summary.items():
-                table.at[obj.name, attr] = value
+                table.at[index, attr] = value
 
         for attr, dtype in OBJ_FIELD_TYPES.items():
             table[attr] = table[attr].astype(dtype)
@@ -337,7 +313,6 @@ def _repr_html_(self):
         return self.summary.T._repr_html_()
 
     def add(self, objective):
-        _validate_objs([*self.objectives, objective])
         self.objectives.append(objective)
 
     @staticmethod
diff --git a/src/blop/tests/conftest.py b/src/blop/tests/conftest.py
index 227bd2a..6bba490 100644
--- a/src/blop/tests/conftest.py
+++ b/src/blop/tests/conftest.py
@@ -1,5 +1,5 @@
+# content of conftest.py
 import asyncio
-import datetime
 
 import databroker
 import numpy as np
@@ -7,7 +7,6 @@
 from bluesky.callbacks import best_effort
 from bluesky.run_engine import RunEngine
 from databroker import Broker
-from ophyd.utils import make_dir_tree
 
 from blop import DOF, Agent, Objective
 from blop.dofs import BrownianMotion
@@ -45,168 +44,103 @@ def RE(db):
     return RE
 
 
-@pytest.fixture(scope="function")
-def agent_1d_1f(db):
-    """
-    A one dimensional agent.
-    """
-
-    def digestion(df):
-        for index, entry in df.iterrows():
-            df.loc[index, "f1"] = functions.himmelblau(entry.x1, 3)
-
-        return df
-
-    dofs = DOF(description="The first DOF", name="x1", search_domain=(-5.0, 5.0))
-    objectives = Objective(description="f1", name="f1", target="min")
-
-    agent = Agent(
-        dofs=dofs,
-        objectives=objectives,
-        digestion=digestion,
-        db=db,
-    )
-
-    agent.dofs.add(DOF(name="dummy", search_domain=(0, 1), read_only=False))
-
-    return agent
-
-
-@pytest.fixture(scope="function")
-def agent_2d_1f(db):
-    """
-    A simple agent minimizing Himmelblau's function
-    """
-
-    dofs = [
-        DOF(name="x1", search_domain=(-5.0, 5.0)),
-        DOF(name="x2", search_domain=(-5.0, 5.0)),
-    ]
-
-    objectives = [Objective(name="himmelblau", target="min")]
-
-    agent = Agent(
-        dofs=dofs,
-        objectives=objectives,
-        digestion=functions.himmelblau_digestion,
-        db=db,
-        verbose=True,
-        tolerate_acquisition_errors=False,
-    )
-
-    agent.dofs.add(DOF(name="dummy", search_domain=(0, 1), read_only=False))
-
-    return agent
-
+single_task_agents = [
+    "1d_1f",
+    # "1d_1c",
+    "2d_1f",
+    "2d_1f_1c",
+    "2d_2f_2c",
+    "3d_2r_2f_1c",
+]
 
-@pytest.fixture(scope="function")
-def agent_2d_2f(db):
-    """
-    An agent minimizing two Himmelblau's functions
-    """
-
-    def digestion(df):
-        for index, entry in df.iterrows():
-            df.loc[index, "f1"] = functions.himmelblau(entry.x1, entry.x2)
-            df.loc[index, "f2"] = functions.himmelblau(entry.x2, entry.x1)
-
-        return df
-
-    dofs = [
-        DOF(name="x1", search_domain=(-5.0, 5.0)),
-        DOF(name="x2", search_domain=(-5.0, 5.0)),
-    ]
-
-    objectives = [Objective(name="f1", target="min"), Objective(name="f2", target="min")]
-
-    agent = Agent(
-        dofs=dofs,
-        objectives=objectives,
-        digestion=digestion,
-        db=db,
-        verbose=True,
-        tolerate_acquisition_errors=False,
-    )
-
-    agent.dofs.add(DOF(name="dummy", search_domain=(0, 1), read_only=False))
-
-    return agent
+pareto_agents = ["2d_2f_2c", "3d_2r_2f_1c"]
 
+all_agents = [*single_task_agents, *pareto_agents]
 
-@pytest.fixture(scope="function")
-def agent_2d_2f_2c(db):
-    """
-    Chankong and Haimes function from https://en.wikipedia.org/wiki/Test_functions_for_optimization
-    """
-
-    def digestion(df):
-        for index, entry in df.iterrows():
-            df.loc[index, "f1"] = (entry.x1 - 2) ** 2 + (entry.x2 - 1) + 2
-            df.loc[index, "f2"] = 9 * entry.x1 - (entry.x2 - 1) + 2
-            df.loc[index, "c1"] = entry.x1**2 + entry.x2**2
-            df.loc[index, "c2"] = entry.x1 - 3 * entry.x2 + 10
-
-        return df
-
-    dofs = [
-        DOF(description="The first DOF", name="x1", search_domain=(-20, 20), travel_expense=1.0),
-        DOF(description="The second DOF", name="x2", search_domain=(-20, 20), travel_expense=2.0),
-    ]
-
-    objectives = [
-        Objective(description="f1", name="f1", target="min"),
-        Objective(description="f2", name="f2", target="min"),
-        Objective(description="c1", name="c1", target=(-np.inf, 225)),
-        Objective(description="c2", name="c2", target=(-np.inf, 0)),
-    ]
-
-    agent = Agent(
-        dofs=dofs,
-        objectives=objectives,
-        digestion=digestion,
-        db=db,
-        verbose=True,
-        tolerate_acquisition_errors=False,
-    )
-
-    agent.dofs.add(DOF(name="dummy", search_domain=(0, 1), read_only=False))
-
-    return agent
 
-
-@pytest.fixture(scope="function")
-def agent_with_read_only_dofs(db):
+def get_agent(param):
     """
-    A simple agent minimizing two Himmelblau's functions
+    Generate a bunch of different agents
     """
 
-    dofs = [
-        DOF(name="x1", search_domain=(-5.0, 5.0)),
-        DOF(name="x2", search_domain=(-5.0, 5.0)),
-        DOF(name="x3", search_domain=(-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 = [
-        Objective(name="himmelblau", target="min"),
-    ]
-
-    agent = Agent(
-        dofs=dofs,
-        objectives=objectives,
-        digestion=functions.himmelblau_digestion,
-        db=db,
-        verbose=True,
-        tolerate_acquisition_errors=False,
-    )
-
-    agent.dofs.add(DOF(name="dummy", search_domain=(0, 1), read_only=False))
+    if param == "1d_1f":
+        return Agent(
+            dofs=DOF(description="The first DOF", name="x1", search_domain=(-5.0, 5.0)),
+            objectives=Objective(description="Himmelblau’s function", name="himmelblau", target="min"),
+            digestion=functions.himmelblau_digestion,
+        )
+
+    elif param == "1d_1c":
+        return Agent(
+            dofs=DOF(description="The first DOF", name="x1", search_domain=(-5.0, 5.0)),
+            objectives=Objective(description="Himmelblau’s function", name="himmelblau", target=(95, 105)),
+            digestion=functions.himmelblau_digestion,
+        )
+
+    elif param == "2d_1f":
+        return Agent(
+            dofs=[
+                DOF(description="The first DOF", name="x1", search_domain=(-5.0, 5.0)),
+                DOF(description="The first DOF", name="x2", search_domain=(-5.0, 5.0)),
+            ],
+            objectives=Objective(description="Himmelblau’s function", name="himmelblau", target="min"),
+            digestion=functions.himmelblau_digestion,
+        )
+
+    elif param == "2d_1f_1c":
+        return Agent(
+            dofs=[
+                DOF(description="The first DOF", name="x1", search_domain=(-5.0, 5.0)),
+                DOF(description="The first DOF", name="x2", search_domain=(-5.0, 5.0)),
+            ],
+            objectives=[
+                Objective(description="Himmelblau’s function", name="himmelblau", target="min"),
+                Objective(description="Himmelblau’s function", name="himmelblau", target=(95, 105)),
+            ],
+            digestion=functions.himmelblau_digestion,
+        )
+
+    elif param == "2d_2f_2c":
+        return Agent(
+            dofs=[
+                DOF(description="The first DOF", name="x1", search_domain=(-5.0, 5.0)),
+                DOF(description="The first DOF", name="x2", search_domain=(-5.0, 5.0)),
+            ],
+            objectives=[
+                Objective(description="f1", name="f1", target="min"),
+                Objective(description="f2", name="f2", target="min"),
+                Objective(description="c1", name="c1", target=(-np.inf, 225)),
+                Objective(description="c2", name="c2", target=(-np.inf, 0)),
+            ],
+            digestion=functions.chankong_and_haimes_digestion,
+        )
+
+    elif param == "3d_2r_2f_1c":
+        return Agent(
+            dofs=[
+                DOF(name="x1", search_domain=(-5.0, 5.0)),
+                DOF(name="x2", search_domain=(-5.0, 5.0)),
+                DOF(name="x3", search_domain=(-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=[
+                Objective(name="himmelblau", target="min"),
+                Objective(name="himmelblau_transpose", target="min"),
+                Objective(description="Himmelblau’s function", name="himmelblau", target=(95, 105)),
+            ],
+            digestion=functions.himmelblau_digestion,
+        )
+
+    else:
+        raise ValueError(f"Invalid agent parameter '{param}'.")
+
+
+@pytest.fixture
+def agent(request):
+    agent = get_agent(request.param)
+
+    # add a useless DOF to try and break things
+    agent.dofs.add(DOF(name="dummy", search_domain=(0, 1), active=False))
 
     return agent
-
-
-@pytest.fixture(scope="function")
-def make_dirs():
-    root_dir = "/tmp/sirepo-bluesky-data"
-    _ = make_dir_tree(datetime.datetime.now().year, base_path=root_dir)
diff --git a/src/blop/tests/test_acq_funcs.py b/src/blop/tests/test_acq_funcs.py
deleted file mode 100644
index a2e0c31..0000000
--- a/src/blop/tests/test_acq_funcs.py
+++ /dev/null
@@ -1,49 +0,0 @@
-import pytest
-
-
-@pytest.mark.parametrize("acqf", ["ei", "pi", "em", "ucb"])
-def test_analytic_acqfs_one_dimensional(agent_1d_1f, RE, acqf):
-    a = agent_1d_1f
-    RE(a.learn("qr", n=16))
-    RE(a.learn(acqf, n=1))
-    getattr(a, acqf)
-
-
-@pytest.mark.parametrize("acqf", ["qei", "qpi", "qem", "qucb"])
-def test_monte_carlo_acqfs_one_dimensional(agent_1d_1f, RE, acqf):
-    a = agent_1d_1f
-    RE(a.learn("qr", n=16))
-    RE(a.learn(acqf, n=2))
-    getattr(a, acqf)
-
-
-@pytest.mark.parametrize("acqf", ["ei", "pi", "em", "ucb"])
-def test_analytic_acqfs_single_objective(agent_2d_1f, RE, acqf):
-    a = agent_2d_1f
-    RE(a.learn("qr", n=4))
-    RE(a.learn(acqf, n=1))
-    getattr(a, acqf)
-
-
-@pytest.mark.parametrize("acqf", ["qei", "qpi", "qem", "qucb"])
-def test_monte_carlo_acqfs_single_objective(agent_2d_1f, RE, acqf):
-    a = agent_2d_1f
-    RE(a.learn("qr", n=4))
-    RE(a.learn(acqf, n=2))
-    getattr(a, acqf)
-
-
-@pytest.mark.parametrize("acqf", ["ei", "pi", "em", "ucb"])
-def test_analytic_acqfs_multi_objective(agent_2d_2f, RE, acqf):
-    a = agent_2d_2f
-    RE(a.learn("qr", n=4))
-    RE(a.learn(acqf, n=1))
-    getattr(a, acqf)
-
-
-@pytest.mark.parametrize("acqf", ["qei", "qpi", "qem", "qucb"])
-def test_monte_carlo_acqfs_multi_objective(agent_2d_2f, RE, acqf):
-    a = agent_2d_2f
-    RE(a.learn("qr", n=4))
-    RE(a.learn(acqf, n=2))
-    getattr(a, acqf)
diff --git a/src/blop/tests/test_acqfs.py b/src/blop/tests/test_acqfs.py
new file mode 100644
index 0000000..3548d18
--- /dev/null
+++ b/src/blop/tests/test_acqfs.py
@@ -0,0 +1,21 @@
+import pytest
+
+from .conftest import all_agents
+
+
+@pytest.mark.parametrize("acqf", ["ei", "pi", "em", "ucb"])
+@pytest.mark.parametrize("agent", all_agents, indirect=True)
+def test_analytic_acqfs(agent, RE, db, acqf):
+    agent.db = db
+    RE(agent.learn("qr", n=16))
+    RE(agent.learn(acqf, n=1))
+    getattr(agent, acqf)
+
+
+@pytest.mark.parametrize("acqf", ["qei", "qpi", "qem", "qucb"])
+@pytest.mark.parametrize("agent", all_agents, indirect=True)
+def test_monte_carlo_acqfs(agent, RE, db, acqf):
+    agent.db = db
+    RE(agent.learn("qr", n=4))
+    RE(agent.learn(acqf, n=1))
+    getattr(agent, acqf)
diff --git a/src/blop/tests/test_agent.py b/src/blop/tests/test_agent.py
deleted file mode 100644
index 78b0b1d..0000000
--- a/src/blop/tests/test_agent.py
+++ /dev/null
@@ -1,38 +0,0 @@
-import pytest  # noqa F401
-
-
-def test_agent(agent_2d_1f, RE):
-    a = agent_2d_1f
-    RE(a.learn("qr", n=4))
-
-    best = a.best
-    assert [dof.name in best for dof in a.dofs]
-    assert [obj.name in best for obj in a.objectives]
-    assert a.dofs.x1 is a.dofs[0]
-
-    print(a.dofs)
-    print(a.objectives)
-
-
-def test_refresh(agent_2d_1f, RE):
-    """
-    Test that the agent can determine when the number of DOFs has changed, and adjust.
-    """
-    a = agent_2d_1f
-    RE(a.learn("qr", n=4))
-
-    RE(a.learn("qei", n=1))
-    a.dofs[2].activate()
-    RE(a.learn("qei", n=1))
-
-
-def test_forget(agent_2d_1f, RE):
-    a = agent_2d_1f
-    RE(a.learn("qr", n=4))
-    a.forget(last=2)
-
-
-def test_benchmark(agent_2d_1f, RE):
-    a = agent_2d_1f
-    per_iter_learn_kwargs_list = [{"acqf": "qr", "n": 32}, {"acqf": "qei", "n": 2, "iterations": 2}]
-    RE(a.benchmark(output_dir="/tmp/blop", iterations=1, per_iter_learn_kwargs_list=per_iter_learn_kwargs_list))
diff --git a/src/blop/tests/test_agents.py b/src/blop/tests/test_agents.py
new file mode 100644
index 0000000..8a8dc52
--- /dev/null
+++ b/src/blop/tests/test_agents.py
@@ -0,0 +1,54 @@
+import pytest  # noqa F401
+
+from .conftest import all_agents
+
+
+@pytest.mark.parametrize("agent", all_agents, indirect=True)
+def test_agent(agent, RE, db):
+    """
+    All agents should be able to do these things.
+    """
+
+    agent.db = db
+    RE(agent.learn("qr", n=4))
+
+    best = agent.best
+    assert [dof.name in best for dof in agent.dofs]
+    assert [obj.name in best for obj in agent.objectives]
+    assert agent.dofs.x1 is agent.dofs[0]
+
+    print(agent.dofs)
+    print(agent.objectives)
+
+    # test refreshing
+    RE(agent.learn("qei", n=1))
+    agent.dofs.activate()
+    RE(agent.learn("qei", n=1))
+
+    # test forgetting
+    RE(agent.learn("qr", n=4))
+    agent.forget(last=2)
+
+    # test some functions
+    agent.refresh()
+    agent.redigest()
+
+    # save the data, reset the agent, and get the data back
+    agent.save_data("/tmp/test_save_data.h5")
+    agent.reset()
+    agent.load_data("/tmp/test_save_data.h5")
+
+    RE(agent.learn("qei", n=2))
+
+    # test plots
+    agent.plot_objectives()
+    agent.plot_acquisition()
+    agent.plot_validity()
+    agent.plot_history()
+
+
+@pytest.mark.parametrize("agent", all_agents, indirect=True)
+def test_benchmark(agent, RE, db):
+    agent.db = db
+    per_iter_learn_kwargs_list = [{"acqf": "qr", "n": 32}, {"acqf": "qei", "n": 2, "iterations": 2}]
+    RE(agent.benchmark(output_dir="/tmp/blop", iterations=1, per_iter_learn_kwargs_list=per_iter_learn_kwargs_list))
diff --git a/src/blop/tests/test_napari.py b/src/blop/tests/test_napari.py
deleted file mode 100644
index 231f280..0000000
--- a/src/blop/tests/test_napari.py
+++ /dev/null
@@ -1,7 +0,0 @@
-# import pytest  # noqa F401
-
-
-# @pytest.mark.parametrize("item", ["mean", "error", "qei"])
-# def test_napari_viewer(agent, RE, item):
-#     RE(agent.learn("qr", n=4))
-#     agent.view(item)
diff --git a/src/blop/tests/test_pareto.py b/src/blop/tests/test_pareto.py
index f242168..d2d9a09 100644
--- a/src/blop/tests/test_pareto.py
+++ b/src/blop/tests/test_pareto.py
@@ -1,33 +1,21 @@
 import pytest
 
-
-@pytest.mark.test_func
-def test_pareto(agent_2d_2f, RE):
-    agent = agent_2d_2f
-    RE(agent.learn("qr", n=16))
-    agent.plot_pareto_front()
+from .conftest import pareto_agents
 
 
-@pytest.mark.parametrize("acqf", ["qnehvi"])
-def test_monte_carlo_pareto_acqfs(agent_2d_2f, RE, acqf):
-    agent = agent_2d_2f
-    RE(agent.learn("qr", n=16))
-    RE(agent.learn(acqf, n=2))
-    agent.dofs[0].deactivate()
-    RE(agent.learn(acqf, n=2))
-
-
-@pytest.mark.test_func
-def test_constrained_pareto(agent_2d_2f_2c, RE):
-    agent = agent_2d_2f_2c
+@pytest.mark.parametrize("agent", pareto_agents, indirect=True)
+def test_pareto(agent, RE, db):
+    agent.db = db
     RE(agent.learn("qr", n=16))
     agent.plot_pareto_front()
 
 
 @pytest.mark.parametrize("acqf", ["qnehvi"])
-def test_constrained_monte_carlo_pareto_acqfs(agent_2d_2f_2c, RE, acqf):
-    agent = agent_2d_2f_2c
-    RE(agent.learn("qr", n=16))
+@pytest.mark.parametrize("agent", pareto_agents, indirect=True)
+def test_monte_carlo_pareto_acqfs(agent, RE, db, acqf):
+    agent.db = db
+    RE(agent.learn("qr", n=4))
     RE(agent.learn(acqf, n=2))
     agent.dofs[0].deactivate()
     RE(agent.learn(acqf, n=2))
+    getattr(agent, acqf)
diff --git a/src/blop/tests/test_passive_dofs.py b/src/blop/tests/test_passive_dofs.py
deleted file mode 100644
index f20e2d7..0000000
--- a/src/blop/tests/test_passive_dofs.py
+++ /dev/null
@@ -1,8 +0,0 @@
-import pytest
-
-
-@pytest.mark.test_func
-def test_read_only_dofs(agent_with_read_only_dofs, RE):
-    agent = agent_with_read_only_dofs
-    RE(agent.learn("qr", n=32))
-    RE(agent.learn("qei", n=2))
diff --git a/src/blop/tests/test_plots.py b/src/blop/tests/test_plots.py
deleted file mode 100644
index 77c20bf..0000000
--- a/src/blop/tests/test_plots.py
+++ /dev/null
@@ -1,37 +0,0 @@
-import pytest  # noqa F401
-
-
-def test_plots_one_dimensional(RE, agent_1d_1f):
-    agent = agent_1d_1f
-    RE(agent.learn("qr", n=16))
-    agent.plot_objectives()
-    agent.plot_acquisition()
-    agent.plot_validity()
-    agent.plot_history()
-
-
-def test_plots_two_dimensional(RE, agent_1d_1f):
-    agent = agent_1d_1f
-    RE(agent.learn("qr", n=16))
-    agent.plot_objectives()
-    agent.plot_acquisition()
-    agent.plot_validity()
-    agent.plot_history()
-
-
-def test_plots_multiple_objs(RE, agent_2d_2f):
-    agent = agent_2d_2f
-    RE(agent.learn("qr", n=16))
-    agent.plot_objectives()
-    agent.plot_acquisition()
-    agent.plot_validity()
-    agent.plot_history()
-
-
-def test_plots_read_only_dofs(RE, agent_with_read_only_dofs):
-    RE(agent_with_read_only_dofs.learn("qr", n=16))
-
-    agent_with_read_only_dofs.plot_objectives()
-    agent_with_read_only_dofs.plot_acquisition()
-    agent_with_read_only_dofs.plot_validity()
-    agent_with_read_only_dofs.plot_history()
diff --git a/src/blop/tests/test_read_write.py b/src/blop/tests/test_read_write.py
deleted file mode 100644
index cb37deb..0000000
--- a/src/blop/tests/test_read_write.py
+++ /dev/null
@@ -1,16 +0,0 @@
-import pytest  # noqa F401
-
-
-def test_agent_2d_1f_save_load_data(agent_2d_1f, RE):
-    RE(agent_2d_1f.learn("qr", n=4))
-    agent_2d_1f.save_data("/tmp/test_save_data.h5")
-    agent_2d_1f.reset()
-    agent_2d_1f.load_data("/tmp/test_save_data.h5")
-    RE(agent_2d_1f.learn("qr", n=4))
-
-
-def test_agent_2d_1f_save_load_hypers(agent_2d_1f, RE):
-    RE(agent_2d_1f.learn("qr", n=4))
-    agent_2d_1f.save_hypers("/tmp/test_save_hypers.h5")
-    agent_2d_1f.reset()
-    RE(agent_2d_1f.learn("qr", n=16, hypers="/tmp/test_save_hypers.h5"))
diff --git a/src/blop/utils/functions.py b/src/blop/utils/functions.py
index 1323fa8..e802441 100644
--- a/src/blop/utils/functions.py
+++ b/src/blop/utils/functions.py
@@ -184,24 +184,47 @@ def kb_tradeoff_4d(x1, x2, x3, x4):
     return x_width, y_width, flux
 
 
-def constrained_himmelblau_digestion(df: pd.DataFrame) -> pd.DataFrame:
+def himmelblau_digestion(df: pd.DataFrame) -> pd.DataFrame:
     """
     Digests Himmelblau's function into the feedback.
     """
-
     for index, entry in df.iterrows():
-        df.loc[index, "himmelblau"] = constrained_himmelblau(entry.x1, entry.x2)
+        if not hasattr(entry, "x1"):
+            df.loc[index, "x1"] = x1 = 0
+        else:
+            x1 = entry.x1
+        if not hasattr(entry, "x2"):
+            df.loc[index, "x2"] = x2 = 0
+        else:
+            x2 = entry.x2
+        df.loc[index, "himmelblau"] = himmelblau(x1=x1, x2=x2)
+        df.loc[index, "himmelblau_transpose"] = himmelblau(x1=x2, x2=x1)
 
     return df
 
 
-def himmelblau_digestion(df: pd.DataFrame) -> pd.DataFrame:
+def constrained_himmelblau_digestion(df: pd.DataFrame) -> pd.DataFrame:
     """
     Digests Himmelblau's function into the feedback.
     """
 
+    df = himmelblau_digestion(df)
+    df.loc[:, "himmelblau"] = np.where(df.x1.values**2 + df.x1.values**2 < 36, df.himmelblau.values, np.nan)
+
+    return df
+
+
+"""
+Chankong and Haimes function from https://en.wikipedia.org/wiki/Test_functions_for_optimization
+"""
+
+
+def chankong_and_haimes_digestion(df):
     for index, entry in df.iterrows():
-        df.loc[index, "himmelblau"] = himmelblau(entry.x1, entry.x2)
+        df.loc[index, "f1"] = (entry.x1 - 2) ** 2 + (entry.x2 - 1) + 2
+        df.loc[index, "f2"] = 9 * entry.x1 - (entry.x2 - 1) + 2
+        df.loc[index, "c1"] = entry.x1**2 + entry.x2**2
+        df.loc[index, "c2"] = entry.x1 - 3 * entry.x2 + 10
 
     return df
 

From c1d776ac663df2d53901630c4f06f225faee405e Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Fri, 26 Apr 2024 10:12:46 -0700
Subject: [PATCH 08/11] make data types more robust

---
 src/blop/agent.py                             | 10 +-
 src/blop/bayesian/kernels.py                  |  6 +-
 .../{digestion.py => digestion/__init__.py}   |  0
 src/blop/digestion/tests.py                   | 89 +++++++++++++++++
 src/blop/experiments/__init__.py              |  0
 src/blop/experiments/atf/__init__.py          |  0
 src/blop/experiments/atf/atf.py               |  0
 src/blop/experiments/nsls2/__init__.py        |  0
 src/blop/experiments/nsls2/iss.py             | 46 ---------
 src/blop/experiments/nsls2/tes.py             | 96 -------------------
 src/blop/experiments/sirepo/__init__.py       |  0
 src/blop/experiments/sirepo/tes.py            | 49 ----------
 src/blop/tests/conftest.py                    | 16 ++--
 src/blop/utils/functions.py                   | 74 --------------
 14 files changed, 107 insertions(+), 279 deletions(-)
 rename src/blop/{digestion.py => digestion/__init__.py} (100%)
 create mode 100644 src/blop/digestion/tests.py
 delete mode 100644 src/blop/experiments/__init__.py
 delete mode 100644 src/blop/experiments/atf/__init__.py
 delete mode 100644 src/blop/experiments/atf/atf.py
 delete mode 100644 src/blop/experiments/nsls2/__init__.py
 delete mode 100644 src/blop/experiments/nsls2/iss.py
 delete mode 100644 src/blop/experiments/nsls2/tes.py
 delete mode 100644 src/blop/experiments/sirepo/__init__.py
 delete mode 100644 src/blop/experiments/sirepo/tes.py

diff --git a/src/blop/agent.py b/src/blop/agent.py
index 18a81d8..49c69a0 100644
--- a/src/blop/agent.py
+++ b/src/blop/agent.py
@@ -204,7 +204,7 @@ def sample(self, n: int = DEFAULT_MAX_SAMPLES, method: str = "quasi-random") ->
         else:
             raise ValueError("'method' argument must be one of ['quasi-random', 'random', 'grid'].")
 
-        return self.dofs(active=True).untransform(X)
+        return self.dofs(active=True).untransform(X).double()
 
     def ask(self, acqf="qei", n=1, route=True, sequential=True, upsample=1, **acqf_kwargs):
         """Ask the agent for the best point to sample, given an acquisition function.
@@ -257,8 +257,8 @@ def ask(self, acqf="qei", n=1, route=True, sequential=True, upsample=1, **acqf_k
             # we may pick up some more kwargs
             acqf, acqf_kwargs = _construct_acqf(self, acqf_name=acqf_config["name"], **acqf_kwargs)
 
-            NUM_RESTARTS = 16
-            RAW_SAMPLES = 1024
+            NUM_RESTARTS = 8
+            RAW_SAMPLES = 256
 
             candidates, acqf_obj = botorch.optim.optimize_acqf(
                 acq_function=acqf,
@@ -267,6 +267,7 @@ def ask(self, acqf="qei", n=1, route=True, sequential=True, upsample=1, **acqf_k
                 sequential=sequential,
                 num_restarts=NUM_RESTARTS,
                 raw_samples=RAW_SAMPLES,  # used for intialization heuristic
+                fixed_features={i: dof._transform(dof.readback) for i, dof in enumerate(active_dofs) if dof.read_only},
             )
 
             # this includes both RO and non-RO DOFs.
@@ -604,8 +605,9 @@ def scalarized_fitnesses(self, weights="default", constrained=True):
         fitness_objs = self.objectives(kind="fitness")
         if len(fitness_objs) >= 1:
             f = self.fitness_scalarization(weights=weights).evaluate(self.train_targets(active=True, kind="fitness"))
+            f = torch.where(f.isnan(), -np.inf, f)  # remove all nans
         else:
-            f = torch.zeros(len(self.table), dtype=torch.double)
+            f = torch.zeros(len(self.table), dtype=torch.double)  # if there are no fitnesses, use a constant dummy fitness
         if constrained:
             # how many constraints are satisfied?
             c = self.evaluated_constraints.sum(axis=-1)
diff --git a/src/blop/bayesian/kernels.py b/src/blop/bayesian/kernels.py
index ecc7d5c..b847469 100644
--- a/src/blop/bayesian/kernels.py
+++ b/src/blop/bayesian/kernels.py
@@ -165,8 +165,10 @@ def forward(self, x1, x2, diag=False, **params):
         # adapted from the Matern kernel
         mean = x1.reshape(-1, x1.size(-1)).mean(0)[(None,) * (x1.dim() - 1)]
 
-        trans_x1 = torch.matmul(self.latent_transform.unsqueeze(1), (x1 - mean).unsqueeze(-1)).squeeze(-1)
-        trans_x2 = torch.matmul(self.latent_transform.unsqueeze(1), (x2 - mean).unsqueeze(-1)).squeeze(-1)
+        transform = self.latent_transform.unsqueeze(1)
+
+        trans_x1 = torch.matmul(transform, (x1 - mean).unsqueeze(-1)).squeeze(-1)
+        trans_x2 = torch.matmul(transform, (x2 - mean).unsqueeze(-1)).squeeze(-1)
 
         distance = self.covar_dist(trans_x1, trans_x2, diag=diag, **params)
 
diff --git a/src/blop/digestion.py b/src/blop/digestion/__init__.py
similarity index 100%
rename from src/blop/digestion.py
rename to src/blop/digestion/__init__.py
diff --git a/src/blop/digestion/tests.py b/src/blop/digestion/tests.py
new file mode 100644
index 0000000..f194769
--- /dev/null
+++ b/src/blop/digestion/tests.py
@@ -0,0 +1,89 @@
+import numpy as np
+import pandas as pd
+
+from ..utils import functions
+
+
+def himmelblau_digestion(df: pd.DataFrame) -> pd.DataFrame:
+    """
+    Digests Himmelblau's function into the feedback.
+    """
+    for index, entry in df.iterrows():
+        if not hasattr(entry, "x1"):
+            df.loc[index, "x1"] = x1 = 0
+        else:
+            x1 = entry.x1
+        if not hasattr(entry, "x2"):
+            df.loc[index, "x2"] = x2 = 0
+        else:
+            x2 = entry.x2
+        df.loc[index, "himmelblau"] = functions.himmelblau(x1=x1, x2=x2)
+        df.loc[index, "himmelblau_transpose"] = functions.himmelblau(x1=x2, x2=x1)
+
+    return df
+
+
+def constrained_himmelblau_digestion(df: pd.DataFrame) -> pd.DataFrame:
+    """
+    Digests Himmelblau's function into the feedback, constrained with NaN for a distance of more than 6 from the origin.
+    """
+
+    df = himmelblau_digestion(df)
+    df.loc[:, "himmelblau"] = np.where(df.x1.values**2 + df.x1.values**2 < 36, df.himmelblau.values, np.nan)
+
+    return df
+
+
+def sketchy_himmelblau_digestion(df: pd.DataFrame, p=0.1) -> pd.DataFrame:
+    """
+    Evaluates the constrained Himmelblau, where every point is bad with probability p.
+    """
+
+    df = constrained_himmelblau_digestion(df)
+    bad = np.random.choice(a=[True, False], size=len(df), p=[p, 1 - p])
+    df.loc[:, "himmelblau"] = np.where(bad, np.nan, df.himmelblau.values)
+
+    return df
+
+
+"""
+Chankong and Haimes function from https://en.wikipedia.org/wiki/Test_functions_for_optimization
+"""
+
+
+def chankong_and_haimes_digestion(df):
+    for index, entry in df.iterrows():
+        df.loc[index, "f1"] = (entry.x1 - 2) ** 2 + (entry.x2 - 1) + 2
+        df.loc[index, "f2"] = 9 * entry.x1 - (entry.x2 - 1) + 2
+        df.loc[index, "c1"] = entry.x1**2 + entry.x2**2
+        df.loc[index, "c2"] = entry.x1 - 3 * entry.x2 + 10
+
+    return df
+
+
+def mock_kbs_digestion(df: pd.DataFrame) -> pd.DataFrame:
+    """
+    Digests a beam waist and height into the feedback.
+    """
+
+    for index, entry in df.iterrows():
+        sigma_x = functions.gaussian_beam_waist(entry.x1, entry.x2)
+        sigma_y = functions.gaussian_beam_waist(entry.x3, entry.x4)
+
+        df.loc[index, "x_width"] = 2 * sigma_x
+        df.loc[index, "y_width"] = 2 * sigma_y
+
+    return df
+
+
+def binh_korn_digestion(df: pd.DataFrame) -> pd.DataFrame:
+    """
+    Digests Himmelblau's function into the feedback.
+    """
+
+    for index, entry in df.iterrows():
+        f1, f2 = functions.binh_korn(entry.x1, entry.x2)
+        df.loc[index, "f1"] = f1
+        df.loc[index, "f2"] = f2
+
+    return df
diff --git a/src/blop/experiments/__init__.py b/src/blop/experiments/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/src/blop/experiments/atf/__init__.py b/src/blop/experiments/atf/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/src/blop/experiments/atf/atf.py b/src/blop/experiments/atf/atf.py
deleted file mode 100644
index e69de29..0000000
diff --git a/src/blop/experiments/nsls2/__init__.py b/src/blop/experiments/nsls2/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/src/blop/experiments/nsls2/iss.py b/src/blop/experiments/nsls2/iss.py
deleted file mode 100644
index 6a76a98..0000000
--- a/src/blop/experiments/nsls2/iss.py
+++ /dev/null
@@ -1,46 +0,0 @@
-import bluesky.plan_stubs as bps
-import bluesky.plans as bp
-import numpy as np
-
-from blop.tasks import Task
-
-
-def initialize():
-    yield from bps.null()  # do nothing
-
-
-class MaxBeamFlux(Task):
-    name = "max_flux"
-
-    def get_fitness(processed_entry):
-        return getattr(processed_entry, "flux")
-
-
-def acquisition(dofs, inputs, dets):
-    uid = yield from bp.list_scan(dets, *[_ for items in zip(dofs, np.atleast_2d(inputs).T) for _ in items])
-    return uid
-
-
-def flux_digestion(db, uid):
-    """
-    This method eats the output of a Bluesky scan, and returns a dict with inputs for the tasks
-    """
-
-    table = db[uid].table(fill=True)
-
-    products_keys = [
-        "image",
-        "vertical_extent",
-        "horizontal_extent",
-        "flux",
-        "x_pos",
-        "y_pos",
-        "x_width",
-        "y_width",
-    ]
-    products = {key: [] for key in products_keys}
-
-    for index, entry in table.iterrows():
-        products["apb_ch4"].append(getattr(entry, "apb_ch4"))
-
-    return products
diff --git a/src/blop/experiments/nsls2/tes.py b/src/blop/experiments/nsls2/tes.py
deleted file mode 100644
index c0a495a..0000000
--- a/src/blop/experiments/nsls2/tes.py
+++ /dev/null
@@ -1,96 +0,0 @@
-import time as ttime
-
-import bluesky.plan_stubs as bps
-import bluesky.plans as bp  # noqa F401
-import numpy as np
-
-from ... import utils
-
-TARGET_BEAM_WIDTH_X = 0
-TARGET_BEAM_WIDTH_Y = 0
-BEAM_PROP = 0.5
-MIN_SEPARABILITY = 0.1
-OOB_REL_BUFFER = 1 / 32
-
-
-def initialize(shutter, detectors):
-    timeout = 10
-
-    yield from bps.mv(shutter.close_cmd, 1)
-    yield from bps.sleep(2.0)
-
-    start_time = ttime.monotonic()
-    while (ttime.monotonic() - start_time < timeout) and (shutter.status.get(as_string=True) == "Open"):
-        print(f"Shutter not closed, retrying ... (closed_status = {shutter.status.get(as_string=True)})")
-        yield from bps.sleep(1.0)
-        yield from bps.mv(shutter.close_cmd, 1)
-
-    uid = yield from bp.count(detectors)
-
-    yield from bps.mv(shutter.open_cmd, 1)
-    yield from bps.sleep(2.0)
-
-    timeout = 10
-    start_time = ttime.monotonic()
-    while (ttime.monotonic() - start_time < timeout) and (shutter.status.get(as_string=True) != "Open"):
-        print(f"Shutter not open, retrying ... (closed_status = {shutter.status.get(as_string=True)})")
-        yield from bps.sleep(1.0)
-        yield from bps.mv(shutter.open_cmd, 1)
-
-    yield from bps.sleep(10.0)
-
-    return uid
-
-
-def fitness(entry, args):
-    image = getattr(entry, args["image"])
-    # if (args['flux'] not in entry.index) or (args['flux'] is not None):
-    #    flux = 1e0
-    # else:
-    flux = getattr(entry, args["flux"])
-
-    background = args["background"]
-
-    x_min, x_max, y_min, y_max, separability = utils.get_beam_stats(image - background, beam_prop=args["beam_prop"])
-
-    n_y, n_x = image.shape
-
-    # u, s, v = np.linalg.svd(image - background)
-
-    # separability = np.square(s[0]) / np.square(s).sum()
-
-    # ymode, xmode = u[:,0], v[0]
-
-    # x_roots = utils.estimate_root_indices(np.abs(xmode) - args['beam_prop'] * np.abs(xmode).max())
-    # y_roots = utils.estimate_root_indices(np.abs(ymode) - args['beam_prop'] * np.abs(ymode).max())
-
-    # x_min, x_max = x_roots.min(), x_roots.max()
-    # y_min, y_max = y_roots.min(), y_roots.max()
-
-    width_x = x_max - x_min
-    width_y = y_max - y_min
-
-    bad = x_min < (n_x + 1) * args["OOB_rel_buffer"]
-    bad |= x_max > (n_x + 1) * (1 - args["OOB_rel_buffer"])
-    bad |= y_min < (n_y + 1) * args["OOB_rel_buffer"]
-    bad |= y_max > (n_y + 1) * (1 - args["OOB_rel_buffer"])
-    bad |= separability < args["min_separability"]
-    bad |= width_x < 2
-    bad |= width_y < 2
-
-    if bad:
-        fitness = np.nan
-
-    else:
-        fitness = np.log(flux * separability / (width_x**2 + width_y**2))
-
-    return ("fitness", "x_min", "x_max", "y_min", "y_max", "width_x", "width_y", "separability"), (
-        fitness,
-        x_min,
-        x_max,
-        y_min,
-        y_max,
-        width_x,
-        width_y,
-        separability,
-    )
diff --git a/src/blop/experiments/sirepo/__init__.py b/src/blop/experiments/sirepo/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/src/blop/experiments/sirepo/tes.py b/src/blop/experiments/sirepo/tes.py
deleted file mode 100644
index cbf0599..0000000
--- a/src/blop/experiments/sirepo/tes.py
+++ /dev/null
@@ -1,49 +0,0 @@
-import numpy as np
-
-
-def w8_digestion(db, uid):
-    return image_digestion(db, uid, image_name="w8")
-
-
-def w9_digestion(db, uid):
-    return image_digestion(db, uid, image_name="w9")
-
-
-def image_digestion(db, uid, image_name):
-    products = db[uid].table(fill=True)
-
-    for index, entry in products.iterrows():
-        image = getattr(entry, f"{image_name}_image")
-        horizontal_extent = getattr(entry, f"{image_name}_horizontal_extent")
-        vertical_extent = getattr(entry, f"{image_name}_vertical_extent")
-
-        flux = image.sum()
-        n_y, n_x = image.shape
-
-        # reject if there is no flux, or we can't estimate the position and size of the beam for some reason
-        bad = False
-        bad |= not (flux > 0)
-        if not bad:
-            X, Y = np.meshgrid(np.linspace(*horizontal_extent, n_x), np.linspace(*vertical_extent, n_y))
-
-            mean_x = np.sum(X * image) / np.sum(image)
-            mean_y = np.sum(Y * image) / np.sum(image)
-
-            sigma_x = np.sqrt(np.sum((X - mean_x) ** 2 * image) / np.sum(image))
-            sigma_y = np.sqrt(np.sum((Y - mean_y) ** 2 * image) / np.sum(image))
-
-            bad |= any(np.isnan([mean_x, mean_y, sigma_x, sigma_y]))
-
-        if not bad:
-            products.loc[index, ["flux", "x_pos", "y_pos", "x_width", "y_width"]] = (
-                flux,
-                mean_x,
-                mean_y,
-                2 * sigma_x,
-                2 * sigma_y,
-            )
-        else:
-            for col in ["flux", "x_pos", "y_pos", "x_width", "y_width"]:
-                products.loc[index, col] = np.nan
-
-    return products
diff --git a/src/blop/tests/conftest.py b/src/blop/tests/conftest.py
index 6bba490..9de6513 100644
--- a/src/blop/tests/conftest.py
+++ b/src/blop/tests/conftest.py
@@ -9,8 +9,8 @@
 from databroker import Broker
 
 from blop import DOF, Agent, Objective
+from blop.digestion.tests import chankong_and_haimes_digestion, sketchy_himmelblau_digestion
 from blop.dofs import BrownianMotion
-from blop.utils import functions
 
 
 @pytest.fixture(scope="function")
@@ -60,21 +60,21 @@ def RE(db):
 
 def get_agent(param):
     """
-    Generate a bunch of different agents
+    Generate a bunch of different agents.
     """
 
     if param == "1d_1f":
         return Agent(
             dofs=DOF(description="The first DOF", name="x1", search_domain=(-5.0, 5.0)),
             objectives=Objective(description="Himmelblau’s function", name="himmelblau", target="min"),
-            digestion=functions.himmelblau_digestion,
+            digestion=sketchy_himmelblau_digestion,
         )
 
     elif param == "1d_1c":
         return Agent(
             dofs=DOF(description="The first DOF", name="x1", search_domain=(-5.0, 5.0)),
             objectives=Objective(description="Himmelblau’s function", name="himmelblau", target=(95, 105)),
-            digestion=functions.himmelblau_digestion,
+            digestion=sketchy_himmelblau_digestion,
         )
 
     elif param == "2d_1f":
@@ -84,7 +84,7 @@ def get_agent(param):
                 DOF(description="The first DOF", name="x2", search_domain=(-5.0, 5.0)),
             ],
             objectives=Objective(description="Himmelblau’s function", name="himmelblau", target="min"),
-            digestion=functions.himmelblau_digestion,
+            digestion=sketchy_himmelblau_digestion,
         )
 
     elif param == "2d_1f_1c":
@@ -97,7 +97,7 @@ def get_agent(param):
                 Objective(description="Himmelblau’s function", name="himmelblau", target="min"),
                 Objective(description="Himmelblau’s function", name="himmelblau", target=(95, 105)),
             ],
-            digestion=functions.himmelblau_digestion,
+            digestion=sketchy_himmelblau_digestion,
         )
 
     elif param == "2d_2f_2c":
@@ -112,7 +112,7 @@ def get_agent(param):
                 Objective(description="c1", name="c1", target=(-np.inf, 225)),
                 Objective(description="c2", name="c2", target=(-np.inf, 0)),
             ],
-            digestion=functions.chankong_and_haimes_digestion,
+            digestion=chankong_and_haimes_digestion,
         )
 
     elif param == "3d_2r_2f_1c":
@@ -129,7 +129,7 @@ def get_agent(param):
                 Objective(name="himmelblau_transpose", target="min"),
                 Objective(description="Himmelblau’s function", name="himmelblau", target=(95, 105)),
             ],
-            digestion=functions.himmelblau_digestion,
+            digestion=sketchy_himmelblau_digestion,
         )
 
     else:
diff --git a/src/blop/utils/functions.py b/src/blop/utils/functions.py
index e802441..26b16b6 100644
--- a/src/blop/utils/functions.py
+++ b/src/blop/utils/functions.py
@@ -1,5 +1,4 @@
 import numpy as np
-import pandas as pd
 import torch
 
 
@@ -61,19 +60,6 @@ def binh_korn(x1, x2):
     return np.where(c, f1, np.nan), np.where(c, f2, np.nan)
 
 
-def binh_korn_digestion(df: pd.DataFrame) -> pd.DataFrame:
-    """
-    Digests Himmelblau's function into the feedback.
-    """
-
-    for index, entry in df.iterrows():
-        f1, f2 = binh_korn(entry.x1, entry.x2)
-        df.loc[index, "f1"] = f1
-        df.loc[index, "f2"] = f2
-
-    return df
-
-
 def skewed_himmelblau(x1, x2):
     """
     Himmelblau's function, with skewed coordinates
@@ -182,63 +168,3 @@ def kb_tradeoff_4d(x1, x2, x3, x4):
     flux = np.exp(-0.5 * np.where(d < 5, np.where(d > -5, 0, d + 5), d - 5) ** 2)
 
     return x_width, y_width, flux
-
-
-def himmelblau_digestion(df: pd.DataFrame) -> pd.DataFrame:
-    """
-    Digests Himmelblau's function into the feedback.
-    """
-    for index, entry in df.iterrows():
-        if not hasattr(entry, "x1"):
-            df.loc[index, "x1"] = x1 = 0
-        else:
-            x1 = entry.x1
-        if not hasattr(entry, "x2"):
-            df.loc[index, "x2"] = x2 = 0
-        else:
-            x2 = entry.x2
-        df.loc[index, "himmelblau"] = himmelblau(x1=x1, x2=x2)
-        df.loc[index, "himmelblau_transpose"] = himmelblau(x1=x2, x2=x1)
-
-    return df
-
-
-def constrained_himmelblau_digestion(df: pd.DataFrame) -> pd.DataFrame:
-    """
-    Digests Himmelblau's function into the feedback.
-    """
-
-    df = himmelblau_digestion(df)
-    df.loc[:, "himmelblau"] = np.where(df.x1.values**2 + df.x1.values**2 < 36, df.himmelblau.values, np.nan)
-
-    return df
-
-
-"""
-Chankong and Haimes function from https://en.wikipedia.org/wiki/Test_functions_for_optimization
-"""
-
-
-def chankong_and_haimes_digestion(df):
-    for index, entry in df.iterrows():
-        df.loc[index, "f1"] = (entry.x1 - 2) ** 2 + (entry.x2 - 1) + 2
-        df.loc[index, "f2"] = 9 * entry.x1 - (entry.x2 - 1) + 2
-        df.loc[index, "c1"] = entry.x1**2 + entry.x2**2
-        df.loc[index, "c2"] = entry.x1 - 3 * entry.x2 + 10
-
-    return df
-
-
-def mock_kbs_digestion(df: pd.DataFrame) -> pd.DataFrame:
-    """
-    Digests a beam waist and height into the feedback.
-    """
-
-    for index, entry in df.iterrows():
-        sigma_x = gaussian_beam_waist(entry.x1, entry.x2)
-        sigma_y = gaussian_beam_waist(entry.x3, entry.x4)
-
-        df.loc[index, "x_width"] = 2 * sigma_x
-        df.loc[index, "y_width"] = 2 * sigma_y
-
-    return df

From dbcba80a2ad26ffff18955f5fbf3b2842945bab7 Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Tue, 30 Apr 2024 13:16:35 -0400
Subject: [PATCH 09/11] fix docs bug

---
 docs/source/tutorials/passive-dofs.ipynb | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/source/tutorials/passive-dofs.ipynb b/docs/source/tutorials/passive-dofs.ipynb
index a79ae5b..5ff058e 100644
--- a/docs/source/tutorials/passive-dofs.ipynb
+++ b/docs/source/tutorials/passive-dofs.ipynb
@@ -38,7 +38,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from blop.utils import functions\n",
+    "from blop.digestion.tests import constrained_himmelblau_digestion\n",
     "from blop import DOF, Agent, Objective\n",
     "from blop.dofs import BrownianMotion\n",
     "\n",
@@ -58,7 +58,7 @@
     "agent = Agent(\n",
     "    dofs=dofs,\n",
     "    objectives=objectives,\n",
-    "    digestion=functions.constrained_himmelblau_digestion,\n",
+    "    digestion=constrained_himmelblau_digestion,\n",
     "    db=db,\n",
     "    verbose=True,\n",
     "    tolerate_acquisition_errors=False,\n",

From a708a64c3f83311275e7c8b6c782f577e15c2681 Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Thu, 2 May 2024 15:11:17 -0400
Subject: [PATCH 10/11] add labels to points returned by .ask()

---
 docs/source/tutorials/introduction.ipynb |  7 ++-----
 src/blop/agent.py                        | 22 ++++++++++++++--------
 src/blop/plans.py                        | 12 ++++++++----
 src/blop/plotting.py                     | 10 ++++++----
 4 files changed, 30 insertions(+), 21 deletions(-)

diff --git a/docs/source/tutorials/introduction.ipynb b/docs/source/tutorials/introduction.ipynb
index 214cfff..277a246 100644
--- a/docs/source/tutorials/introduction.ipynb
+++ b/docs/source/tutorials/introduction.ipynb
@@ -245,11 +245,8 @@
    "source": [
     "res = agent.ask(\"qei\", n=8, route=True)\n",
     "agent.plot_acquisition(acqf=\"qei\")\n",
-    "plt.scatter(*res[\"points\"].T, marker=\"d\", facecolor=\"w\", edgecolor=\"k\")\n",
-    "plt.plot(\n",
-    "    *res[\"points\"].T,\n",
-    "    color=\"r\",\n",
-    ")"
+    "plt.scatter(res[\"points\"][\"x1\"], res[\"points\"][\"x2\"], marker=\"d\", facecolor=\"w\", edgecolor=\"k\")\n",
+    "plt.plot(res[\"points\"][\"x1\"], res[\"points\"][\"x2\"], color=\"r\")"
    ]
   },
   {
diff --git a/src/blop/agent.py b/src/blop/agent.py
index 49c69a0..a99248e 100644
--- a/src/blop/agent.py
+++ b/src/blop/agent.py
@@ -297,9 +297,9 @@ def ask(self, acqf="qei", n=1, route=True, sequential=True, upsample=1, **acqf_k
             points = sp.interpolate.interp1d(idx, points, axis=0)(upsampled_idx)
 
         res = {
-            "points": points,
+            "points": {dof.name: list(points[..., i]) for i, dof in enumerate(active_dofs(read_only=False))},
             "acqf_name": acqf_config["name"],
-            "acqf_obj": np.atleast_1d(acqf_obj.numpy()),
+            "acqf_obj": list(np.atleast_1d(acqf_obj.numpy())),
             "acqf_kwargs": acqf_kwargs,
             "duration_ms": duration,
             "sequential": sequential,
@@ -475,7 +475,7 @@ def view(self, item: str = "mean", cmap: str = "turbo", max_inputs: int = 2**16)
 
         self.viewer.dims.axis_labels = self.dofs.names
 
-    def acquire(self, acquisition_inputs):
+    def acquire(self, points):
         """Acquire and digest according to the self's acquisition and digestion plans.
 
         Parameters
@@ -487,11 +487,17 @@ def acquire(self, acquisition_inputs):
         if self.db is None:
             raise ValueError("Cannot run acquistion without databroker instance!")
 
+        acquisition_dofs = self.dofs(active=True, read_only=False)
+        for dof in acquisition_dofs:
+            if dof.name not in points:
+                raise ValueError(f"Cannot acquire points; missing values for {dof.name}.")
+
+        n = len(points[dof.name])
+
         try:
-            acquisition_devices = self.dofs(active=True, read_only=False).devices
             uid = yield from self.acquisition_plan(
-                acquisition_devices,
-                acquisition_inputs.astype(float),
+                acquisition_dofs,
+                points,
                 [*self.dets, *self.dofs.devices],
                 delay=self.trigger_delay,
             )
@@ -504,11 +510,11 @@ def acquire(self, acquisition_inputs):
             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(active=True, read_only=False).names)
+            products = pd.DataFrame(points)
             for obj in self.objectives(active=True):
                 products.loc[:, obj.name] = np.nan
 
-        if not len(acquisition_inputs) == len(products):
+        if len(products) != n:
             raise ValueError("The table returned by the digestion function must be the same length as the sampled inputs!")
 
         return products
diff --git a/src/blop/plans.py b/src/blop/plans.py
index 26f0940..d84c957 100644
--- a/src/blop/plans.py
+++ b/src/blop/plans.py
@@ -1,6 +1,5 @@
 import bluesky.plan_stubs as bps
 import bluesky.plans as bp
-import numpy as np
 
 
 def list_scan_with_delay(*args, delay=0, **kwargs):
@@ -19,11 +18,16 @@ def one_nd_step_with_delay(detectors, step, pos_cache):
 
 
 def default_acquisition_plan(dofs, inputs, dets, **kwargs):
+    """
+    "dofs" is a list of dofs.
+    "inputs" is a dict of a list of inputs per dof, keyed by dof.name
+    "dets" is a list of detectors to trigger
+    """
     delay = kwargs.get("delay", 0)
     args = []
-    for dof, points in zip(dofs, np.atleast_2d(inputs).T):
-        args.append(dof)
-        args.append(list(points))
+    for dof in dofs:
+        args.append(dof.device)
+        args.append(inputs[dof.name])
 
     uid = yield from list_scan_with_delay(dets, *args, delay=delay)
     return uid
diff --git a/src/blop/plotting.py b/src/blop/plotting.py
index bbeaa19..5492a65 100644
--- a/src/blop/plotting.py
+++ b/src/blop/plotting.py
@@ -327,6 +327,7 @@ def _plot_acqf_one_dof(agent, acqfs, lw=1e0, **kwargs):
     x_dof = agent.dofs(active=True)[0]
 
     test_inputs = agent.sample(method="grid")
+    model_inputs = agent.dofs.transform(test_inputs)
 
     for iacqf, acqf_identifier in enumerate(acqfs):
         color = DEFAULT_COLOR_LIST[iacqf]
@@ -334,7 +335,7 @@ def _plot_acqf_one_dof(agent, acqfs, lw=1e0, **kwargs):
         acqf_config = parse_acqf_identifier(acqf_identifier)
         acqf, _ = _construct_acqf(agent, acqf_config["name"])
 
-        test_acqf_value = acqf(test_inputs).detach().numpy()
+        test_acqf_value = acqf(model_inputs).detach().numpy()
 
         agent.acq_axes[iacqf].plot(test_inputs.squeeze(-2), test_acqf_value, lw=lw, color=color)
 
@@ -366,6 +367,7 @@ def _plot_acqf_many_dofs(
 
     # test_inputs has shape (..., 1, n_active_dofs)
     test_inputs = agent.sample(n=1024, method="grid") if gridded else agent.sample(n=1024)
+    model_inputs = agent.dofs.transform(test_inputs)
     *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()
@@ -374,7 +376,7 @@ def _plot_acqf_many_dofs(
         acqf_config = parse_acqf_identifier(acqf_identifier)
         acqf, _ = _construct_acqf(agent, acqf_config["name"])
 
-        test_acqf_value = acqf(test_inputs.reshape(-1, 1, input_dim)).detach().reshape(test_dim).squeeze().numpy()
+        test_acqf_value = acqf(model_inputs.reshape(-1, 1, input_dim)).detach().reshape(test_dim).squeeze().numpy()
 
         if gridded:
             agent.acq_axes[iacqf].set_title(acqf_config["name"])
@@ -416,7 +418,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.constraint(test_inputs)[..., 0]
+    constraint = agent.constraint(agent.dofs.transform(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)
@@ -438,7 +440,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.constraint(test_inputs)[..., 0].squeeze().numpy()
+    constraint = agent.constraint(agent.dofs.transform(test_inputs))[..., 0].squeeze().numpy()
 
     if gridded:
         _ = agent.valid_axes[1].pcolormesh(

From 0dcbfb478f37e9bdeb79d6b670a6c7405b6e83b0 Mon Sep 17 00:00:00 2001
From: Thomas Morris <thomasmorris@princeton.edu>
Date: Fri, 10 May 2024 13:49:08 -0400
Subject: [PATCH 11/11] apply suggestions from PR review

---
 src/blop/agent.py             |  8 +++++---
 src/blop/plans.py             | 11 ++++++++---
 src/blop/tests/test_agents.py |  2 +-
 3 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/src/blop/agent.py b/src/blop/agent.py
index a99248e..80ef55f 100644
--- a/src/blop/agent.py
+++ b/src/blop/agent.py
@@ -242,7 +242,7 @@ def ask(self, acqf="qei", n=1, route=True, sequential=True, upsample=1, **acqf_k
             # check that all the objectives have models
             if not all(hasattr(obj, "model") for obj in active_objs):
                 raise RuntimeError(
-                    f"Can't construct non-trivial acquisition function '{acqf}'" f" (the agent is not initialized!)"
+                    f"Can't construct non-trivial acquisition function '{acqf}' as the agent is not initialized."
                 )
 
             # if the model for any active objective mismatches the active dofs, reconstrut and train it
@@ -628,8 +628,10 @@ def best_f(self, weights="default"):
 
     @property
     def pareto_mask(self):
-        # a point is on the Pareto front if it is Pareto dominant
-        # a point is Pareto dominant if it is there is no other point that is better at every objective
+        """
+        Returns a mask of all points that satisfy all constraints and are Pareto efficient.
+        A point is Pareto efficient if it is there is no other point that is better at every objective.
+        """
         Y = self.train_targets(active=True, kind="fitness")
 
         # nuke the bad points
diff --git a/src/blop/plans.py b/src/blop/plans.py
index d84c957..e26725d 100644
--- a/src/blop/plans.py
+++ b/src/blop/plans.py
@@ -19,9 +19,14 @@ def one_nd_step_with_delay(detectors, step, pos_cache):
 
 def default_acquisition_plan(dofs, inputs, dets, **kwargs):
     """
-    "dofs" is a list of dofs.
-    "inputs" is a dict of a list of inputs per dof, keyed by dof.name
-    "dets" is a list of detectors to trigger
+    Parameters
+    ----------
+    x : list of DOFs or DOFList
+        A list of DOFs
+    inputs: dict
+        A dict of a list of inputs per dof, keyed by dof.name
+    dets: list
+        A list of detectors to trigger
     """
     delay = kwargs.get("delay", 0)
     args = []
diff --git a/src/blop/tests/test_agents.py b/src/blop/tests/test_agents.py
index 8a8dc52..84f434c 100644
--- a/src/blop/tests/test_agents.py
+++ b/src/blop/tests/test_agents.py
@@ -10,7 +10,7 @@ def test_agent(agent, RE, db):
     """
 
     agent.db = db
-    RE(agent.learn("qr", n=4))
+    RE(agent.learn("qr", n=16))
 
     best = agent.best
     assert [dof.name in best for dof in agent.dofs]