add pareto functionality

src/blop/agent.py

@@ -17,15 +17,19 @@
 import pandas as pd
 import scipy as sp
 import torch
+
+# from botorch.utils.transforms import normalize
+from botorch.acquisition.objective import ScalarizedPosteriorTransform
 from botorch.models.deterministic import GenericDeterministicModel
 from botorch.models.model_list_gp_regression import ModelListGP
 from botorch.models.transforms.input import AffineInputTransform, ChainedInputTransform, Log10, Normalize
 from databroker import Broker
 from ophyd import Signal
 
-from . import utils
-from .bayesian import acquisition, models, plotting
-from .bayesian.transforms import TargetingPosteriorTransform
+from . import plotting, utils
+from .bayesian import acquisition, models
+
+# from .bayesian.transforms import TargetingPosteriorTransform
 from .digestion import default_digestion_function
 from .dofs import DOF, DOFList
 from .objectives import Objective, ObjectiveList
@@ -532,7 +536,7 @@ def benchmark(
 
     @property
     def model(self):
-        """A model encompassing all the objectives. A single GP in the single-objective case, or a model list."""
+        """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
         )
@@ -542,32 +546,86 @@ def posterior(self, x):
         return self.model.posterior(torch.tensor(x))
 
     @property
-    def targeting_transform(self):
-        return TargetingPosteriorTransform(
-            weights=torch.tensor(self.active_objs.weights, dtype=torch.double), targets=self.active_objs.targets
+    def fitness_model(self):
+        return (
+            ModelListGP(*[obj.model for obj in self.active_objs if obj.type == "fitness"])
+            if len(self.active_objs) > 1
+            else self.active_objs[0].model
         )
 
     @property
-    def scalarized_objectives(self):
-        """Returns a (n_obs,) array of scalarized objectives"""
-        return self.targeting_transform.evaluate(self.train_targets(active=True)).sum(axis=-1)
+    def fitness_weights(self):
+        return torch.tensor([obj.weight for obj in self.objectives if obj.type == "fitness"], dtype=torch.double)
 
     @property
-    def max_scalarized_objective(self):
-        """Returns the value of the best scalarized objective seen so far."""
-        f = self.scalarized_objectives
-        return np.max(np.where(np.isnan(f), -np.inf, f))
+    def fitness_scalarization(self):
+        return ScalarizedPosteriorTransform(weights=self.fitness_weights)
 
     @property
-    def argmax_scalarized_objective(self):
-        """Returns the index of the best scalarized objective seen so far."""
-        f = self.scalarized_objectives
-        return np.argmax(np.where(np.isnan(f), -np.inf, f))
+    def evaluated_fitnesses(self):
+        return self.train_targets()[:, self.objectives.type == "fitness"]
+
+    @property
+    def scalarized_fitnesses(self):
+        return self.fitness_scalarization.evaluate(self.evaluated_fitnesses)
+
+    @property
+    def evaluated_constraints(self):
+        if sum(self.objectives.type == "constraint"):
+            y = self.train_targets()
+            return torch.cat(
+                [
+                    ((y[:, i] >= obj.target[0]) & (y[:, i] <= obj.target[1])).unsqueeze(0)
+                    for i, obj in enumerate(self.objectives)
+                    if obj.type == "constraint"
+                ],
+                dim=0,
+            ).T
+        else:
+            return torch.ones(size=(len(self.table), 0), dtype=torch.bool)
+
+    @property
+    def argmax_best_f(self):
+        f = self.scalarized_fitnesses
+        c = self.evaluated_constraints.all(axis=-1)
+        mask = (~f.isnan()) & c
+
+        if not mask.sum():
+            raise ValueError("There are no valid points that satisfy the constraints!")
+
+        return torch.where(mask)[0][f[mask].argmax()]
+
+    @property
+    def best_f(self):
+        return self.scalarized_fitnesses[self.argmax_best_f]
+
+    @property
+    def pareto_front_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()[:, 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
+
+    @property
+    def pareto_front(self):
+        return self.table.loc[self.pareto_front_mask.numpy()]
+
+    @property
+    def min_ref_point(self):
+        y = self.train_targets()[:, self.objectives.type == "fitness"]
+        return y[y.argmax(axis=0)].min(axis=0).values
+
+    @property
+    def random_ref_point(self):
+        pareto_front_index = torch.where(self.pareto_front_mask)[0]
+        return self.evaluated_fitnesses[np.random.choice(pareto_front_index)]
 
     @property
     def all_objectives_valid(self):
         """A mask of whether all objectives are valid for each data point."""
-        return ~torch.isnan(self.scalarized_objectives)
+        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`."""
@@ -621,7 +679,7 @@ def _construct_model(self, obj, skew_dims=None):
                 trusted.long(), learn_additional_noise=True
             )
 
-            obj.validity_conjugate_model = models.LatentDirichletModel(
+            obj.validity_conjugate_model = models.LatentDirichletClassifier(
                 train_inputs=train_inputs[inputs_are_trusted],
                 train_targets=dirichlet_likelihood.transformed_targets.transpose(-1, -2)[inputs_are_trusted].double(),
                 skew_dims=skew_dims,
@@ -760,12 +818,12 @@ def constraint(self, x):
         p = torch.ones(x.shape[:-1])
         for obj in self.active_objs:
             # if the targeting constraint is non-trivial
-            if obj.use_as_constraint:
+            if obj.type == "constraint":
                 p *= obj.targeting_constraint(x)
             # if the validity constaint is non-trivial
             if obj.validity_conjugate_model is not None:
                 p *= obj.validity_constraint(x)
-        return p
+        return p  # + 1e-6 * normalize(x, self._sample_domain).square().sum(axis=-1)
 
     @property
     def hypers(self) -> dict:
@@ -848,22 +906,20 @@ def train_targets(self, index=None, **subset_kwargs):
             return torch.cat([self.train_targets(obj.name) for obj in self.objectives], dim=-1)
 
         obj = self.objectives[index]
-        targets = self.table.loc[:, obj.name].values.copy()
+        values = self.table.loc[:, obj.name].values.copy()
 
         # check that targets values are inside acceptable values
-        valid = (targets >= obj._trust_domain[0]) & (targets <= obj._trust_domain[1])
-        targets = np.where(valid, targets, np.nan)
+        valid = (values >= obj._trust_domain[0]) & (values <= obj._trust_domain[1])
+        values = np.where(valid, values, np.nan)
 
-        # transform if needed
-        if obj.log:
-            targets = np.where(targets > 0, np.log(targets), np.nan)
+        targets = obj.fitness_forward(values)
 
         return torch.tensor(targets, dtype=torch.double).unsqueeze(-1)
 
     @property
     def best(self):
         """Returns all data for the best point."""
-        return self.table.loc[self.argmax_scalarized_objective]
+        return self.table.loc[self.argmax_best_f]
 
     @property
     def best_inputs(self):
@@ -941,3 +997,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}
+
+    def plot_pareto_front(self, **kwargs):
+        """Plot the improvement of the agent over time."""
+        plotting._plot_pareto_front(self, **kwargs)

src/blop/bayesian/acquisition/__init__.py

@@ -1,6 +1,7 @@
 import os
 
 import yaml
+from botorch.utils.transforms import normalize
 
 from . import analytic, monte_carlo
 from .analytic import *  # noqa F401
@@ -42,53 +43,53 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
     if acq_func_name == "expected_improvement":
         acq_func = analytic.ConstrainedLogExpectedImprovement(
             constraint=agent.constraint,
-            model=agent.model,
-            best_f=agent.max_scalarized_objective,
-            posterior_transform=agent.targeting_transform,
+            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 == "monte_carlo_expected_improvement":
         acq_func = monte_carlo.qConstrainedExpectedImprovement(
             constraint=agent.constraint,
-            model=agent.model,
-            best_f=agent.max_scalarized_objective,
-            posterior_transform=agent.targeting_transform,
+            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 == "probability_of_improvement":
         acq_func = analytic.ConstrainedLogProbabilityOfImprovement(
             constraint=agent.constraint,
-            model=agent.model,
-            best_f=agent.max_scalarized_objective,
-            posterior_transform=agent.targeting_transform,
+            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 == "monte_carlo_probability_of_improvement":
         acq_func = monte_carlo.qConstrainedProbabilityOfImprovement(
             constraint=agent.constraint,
-            model=agent.model,
-            best_f=agent.max_scalarized_objective,
-            posterior_transform=agent.targeting_transform,
+            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(
             constraint=agent.constraint,
-            model=agent.model,
+            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 == "noisy_expected_hypervolume_improvement":
+    elif acq_func_name == "monte_carlo_noisy_expected_hypervolume_improvement":
         acq_func = monte_carlo.qConstrainedNoisyExpectedHypervolumeImprovement(
             constraint=agent.constraint,
-            model=agent.model,
-            ref_point=agent.train_targets.min(dim=0).values,
-            X_baseline=agent.train_inputs,
+            model=agent.fitness_model,
+            ref_point=agent.random_ref_point,
+            X_baseline=normalize(agent.train_inputs(), agent._sample_domain),
             prune_baseline=True,
         )
         acq_func_meta = {"name": acq_func_name, "args": {}}
@@ -98,9 +99,9 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
 
         acq_func = analytic.ConstrainedUpperConfidenceBound(
             constraint=agent.constraint,
-            model=agent.model,
+            model=agent.fitness_model,
             beta=beta,
-            posterior_transform=agent.targeting_transform,
+            posterior_transform=agent.fitness_scalarization,
         )
         acq_func_meta = {"name": acq_func_name, "args": {"beta": beta}}
 
@@ -109,9 +110,9 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
 
         acq_func = monte_carlo.qConstrainedUpperConfidenceBound(
             constraint=agent.constraint,
-            model=agent.model,
+            model=agent.fitness_model,
             beta=beta,
-            posterior_transform=agent.targeting_transform,
+            posterior_transform=agent.fitness_scalarization,
         )
         acq_func_meta = {"name": acq_func_name, "args": {"beta": beta}}
 

src/blop/bayesian/acquisition/config.yml

@@ -48,6 +48,14 @@ noisy_expected_hypervolume_improvement:
   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
+  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:

src/blop/bayesian/models.py

@@ -1,11 +1,11 @@
-import botorch
 import gpytorch
 import torch
+from botorch.models.gp_regression import SingleTaskGP
 
 from . import kernels
 
 
-class LatentGP(botorch.models.gp_regression.SingleTaskGP):
+class LatentGP(SingleTaskGP):
     def __init__(self, train_inputs, train_targets, skew_dims=True, *args, **kwargs):
         super().__init__(train_inputs, train_targets, *args, **kwargs)
 
@@ -23,7 +23,22 @@ def __init__(self, train_inputs, train_targets, skew_dims=True, *args, **kwargs)
         self.trained = False
 
 
-class LatentDirichletModel(LatentGP):
+class LatentConstraintModel(LatentGP):
+    def __init__(self, train_inputs, train_targets, skew_dims=True, *args, **kwargs):
+        super().__init__(train_inputs, train_targets, skew_dims, *args, **kwargs)
+
+        self.trained = False
+
+    def fitness(self, x, n_samples=1024):
+        """
+        Takes in a (..., m) dimension tensor and returns a (..., n_classes) tensor
+        """
+        *input_shape, n_dim = x.shape
+        samples = self.posterior(x.reshape(-1, n_dim)).sample(torch.Size((n_samples,))).exp()
+        return (samples / samples.sum(-1, keepdim=True)).mean(0).reshape(*input_shape, -1)
+
+
+class LatentDirichletClassifier(LatentGP):
     def __init__(self, train_inputs, train_targets, skew_dims=True, *args, **kwargs):
         super().__init__(train_inputs, train_targets, skew_dims, *args, **kwargs)
 

src/blop/bayesian/transforms.py

@@ -6,6 +6,8 @@
 from botorch.posteriors.posterior_list import PosteriorList
 from torch import Tensor
 
+# This file is deprecated for now, but we may need it later
+
 
 def targeting_transform(y, target):
     if target == "max":

src/blop/dofs.py

@@ -105,14 +105,16 @@ def __post_init__(self):
                 if not self.read_only:
                     raise ValueError("You must specify search_domain if the device is not read-only.")
             else:
-                self.search_domain = tuple(self.search_domain)
                 if len(self.search_domain) != 2:
-                    raise ValueError("'search_domain' must be a 2-tuple of floats.")
+                    raise ValueError("'search_domain' must have length 2.")
+                self.search_domain = tuple([float(self.search_domain[0]), float(self.search_domain[1])])
                 if self.search_domain[0] > self.search_domain[1]:
                     raise ValueError("The lower search bound must be less than the upper search bound.")
 
             if self.trust_domain is not None:
-                self.trust_domain = tuple(self.trust_domain)
+                if len(self.trust_domain) != 2:
+                    raise ValueError("'trust_domain' must have length 2.")
+                self.trust_domain = tuple([float(self.trust_domain[0]), float(self.trust_domain[1])])
                 if not self.read_only:
                     if (self.search_domain[0] < self.trust_domain[0]) or (self.search_domain[1] > self.trust_domain[1]):
                         raise ValueError("Trust domain must be larger than search domain.")