Refactor recommendation mechanics

CHANGELOG.md

@@ -33,6 +33,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `SubspaceDiscrete.from_product` and `SubspaceDiscrete.from_simplex` now split
   their `constraints` argument into filtering constraints (applied during construction)
   and batch constraints (stored in `batch_constraints`)
+- Internal search space and recommender logic simplified by reducing indirection and
+  argument passing between methods
 
 ### Fixed
 - Deserialization with constructor selection now correctly respects converter settings

baybe/campaign.py

@@ -31,7 +31,6 @@
 from baybe.recommenders.meta.sequential import TwoPhaseMetaRecommender
 from baybe.recommenders.pure.bayesian.base import BayesianRecommender
 from baybe.recommenders.pure.nonpredictive.base import NonPredictiveRecommender
-from baybe.searchspace._filtered import FilteredSubspaceDiscrete
 from baybe.searchspace.core import (
     SearchSpace,
     SearchSpaceType,
@@ -613,8 +612,8 @@ def recommend(
                 )
             searchspace = evolve(
                 self.searchspace,
-                discrete=FilteredSubspaceDiscrete.from_subspace(
-                    self.searchspace.discrete, ~mask_todrop.to_numpy()
+                discrete=evolve(
+                    self.searchspace.discrete, exp_rep=exp_rep.loc[~mask_todrop]
                 ),
             )
         else:

baybe/recommenders/naive.py

@@ -92,9 +92,6 @@ def recommend(
         cont_part = searchspace.continuous.sample_uniform(1)
         cont_part_tensor = to_tensor(cont_part).unsqueeze(-2)
 
-        # Get discrete candidates
-        candidates_exp, _ = searchspace.discrete.get_candidates()
-
         # We now check whether the discrete recommender is bayesian.
         if isinstance(self.disc_recommender, BayesianRecommender):
             # Get access to the recommenders acquisition function
@@ -115,7 +112,6 @@ def recommend(
         # Call the private function of the discrete recommender and get the candidates
         disc_rec = self.disc_recommender._recommend_discrete(
             subspace_discrete=searchspace.discrete,
-            candidates_exp=candidates_exp,
             batch_size=batch_size,
         )
 

baybe/recommenders/pure/base.py

@@ -133,25 +133,18 @@ def recommend(
                 subspace_continuous=searchspace.continuous, batch_size=batch_size
             )
         else:
-            return self._recommend_with_discrete_parts(
-                searchspace,
-                batch_size,
-                pending_experiments=pending_experiments,
-            )
+            return self._recommend_with_discrete_parts(searchspace, batch_size)
 
     def _recommend_discrete(
         self,
         subspace_discrete: SubspaceDiscrete,
-        candidates_exp: pd.DataFrame,
         batch_size: int,
     ) -> pd.DataFrame:
         """Generate recommendations from a discrete search space.
 
         Args:
             subspace_discrete: The discrete subspace from which to generate
                 recommendations.
-            candidates_exp: The experimental representation of all discrete candidate
-                points to be considered.
             batch_size: The size of the recommendation batch.
 
         Raises:
@@ -166,7 +159,6 @@ def _recommend_discrete(
         try:
             return self._recommend_hybrid(
                 searchspace=SearchSpace(discrete=subspace_discrete),
-                candidates_exp=candidates_exp,
                 batch_size=batch_size,
             )
         except NotImplementedError as exc:
@@ -202,7 +194,6 @@ def _recommend_continuous(
         try:
             return self._recommend_hybrid(
                 searchspace=SearchSpace(continuous=subspace_continuous),
-                candidates_exp=pd.DataFrame(),
                 batch_size=batch_size,
             )
         except NotImplementedError as exc:
@@ -218,7 +209,6 @@ def _recommend_continuous(
     def _recommend_hybrid(
         self,
         searchspace: SearchSpace,
-        candidates_exp: pd.DataFrame,
         batch_size: int,
     ) -> pd.DataFrame:
         """Generate recommendations from a hybrid search space.
@@ -230,8 +220,6 @@ def _recommend_hybrid(
         Args:
             searchspace: The hybrid search space from which to generate
                 recommendations.
-            candidates_exp: The experimental representation of all discrete candidate
-                points to be considered.
             batch_size: The size of the recommendation batch.
 
         Raises:
@@ -246,7 +234,6 @@ def _recommend_with_discrete_parts(
         self,
         searchspace: SearchSpace,
         batch_size: int,
-        pending_experiments: pd.DataFrame | None,
     ) -> pd.DataFrame:
         """Obtain recommendations in search spaces with a discrete part.
 
@@ -256,7 +243,6 @@ def _recommend_with_discrete_parts(
         Args:
             searchspace: The search space from which to generate recommendations.
             batch_size: The size of the recommendation batch.
-            pending_experiments: Pending experiments in experimental representation.
 
         Returns:
             A dataframe containing the recommendations as individual rows.
@@ -281,26 +267,21 @@ def _recommend_with_discrete_parts(
                 f"{constraint_types}."
             )
 
-        # Get discrete candidates
-        candidates_exp, _ = searchspace.discrete.get_candidates()
-
         # TODO: Introduce new flag to recommend batches larger than the search space
 
         # Check if enough candidates are left
         # TODO [15917]: This check is not perfectly correct.
-        if (not is_hybrid_space) and (len(candidates_exp) < batch_size):
+        if (not is_hybrid_space) and (len(searchspace.discrete.exp_rep) < batch_size):
             raise NotEnoughPointsLeftError(
                 f"Using the current settings, there are fewer than {batch_size} "
                 f"possible data points left to recommend."
             )
 
         # Get recommendations
         if is_hybrid_space:
-            rec = self._recommend_hybrid(searchspace, candidates_exp, batch_size)
+            rec = self._recommend_hybrid(searchspace, batch_size)
         else:
-            rec = self._recommend_discrete(
-                searchspace.discrete, candidates_exp, batch_size
-            )
+            rec = self._recommend_discrete(searchspace.discrete, batch_size)
 
         # Return recommendations
         return rec

baybe/recommenders/pure/bayesian/botorch/core.py

@@ -156,7 +156,6 @@ def __str__(self) -> str:
     def _recommend_discrete(
         self,
         subspace_discrete: SubspaceDiscrete,
-        candidates_exp: pd.DataFrame,
         batch_size: int,
     ) -> pd.DataFrame:
         """Generate recommendations from a discrete search space.
@@ -167,21 +166,15 @@ def _recommend_discrete(
         Args:
             subspace_discrete: The discrete subspace from which to generate
                 recommendations.
-            candidates_exp: The experimental representation of all discrete candidate
-                points to be considered.
             batch_size: The size of the recommendation batch.
 
         Returns:
             A dataframe containing the recommendations as a subset of rows from the
             provided experimental representation.
         """
         if subspace_discrete.n_subsets > 0:
-            return recommend_discrete_with_subsets(
-                self, subspace_discrete, candidates_exp, batch_size
-            )
-        return recommend_discrete_without_subsets(
-            self, subspace_discrete, candidates_exp, batch_size
-        )
+            return recommend_discrete_with_subsets(self, subspace_discrete, batch_size)
+        return recommend_discrete_without_subsets(self, subspace_discrete, batch_size)
 
     @override
     def _recommend_continuous(
@@ -221,7 +214,6 @@ def _recommend_continuous(
     def _recommend_hybrid(
         self,
         searchspace: SearchSpace,
-        candidates_exp: pd.DataFrame,
         batch_size: int,
     ) -> pd.DataFrame:
         """Generate recommendations from a hybrid search space.
@@ -231,20 +223,14 @@ def _recommend_hybrid(
 
         Args:
             searchspace: The search space in which the recommendations should be made.
-            candidates_exp: The experimental representation of the candidates
-                of the discrete subspace.
             batch_size: The size of the calculated batch.
 
         Returns:
             The recommended points.
         """
         if searchspace.n_subsets > 0:
-            return recommend_hybrid_with_subsets(
-                self, searchspace, candidates_exp, batch_size
-            )
-        return recommend_hybrid_without_subsets(
-            self, searchspace, candidates_exp, batch_size
-        )
+            return recommend_hybrid_with_subsets(self, searchspace, batch_size)
+        return recommend_hybrid_without_subsets(self, searchspace, batch_size)
 
     def _optimize_over_subsets(
         self,

baybe/recommenders/pure/bayesian/botorch/discrete.py

@@ -8,6 +8,7 @@
 import numpy as np
 import numpy.typing as npt
 import pandas as pd
+from attrs import evolve
 
 from baybe.searchspace import SubspaceDiscrete
 from baybe.utils.dataframe import to_tensor
@@ -21,7 +22,6 @@
 def recommend_discrete_with_subsets(
     recommender: BotorchRecommender,
     subspace_discrete: SubspaceDiscrete,
-    candidates_exp: pd.DataFrame,
     batch_size: int,
 ) -> pd.DataFrame:
     """Recommend from a discrete space with subset-generating constraints.
@@ -35,7 +35,6 @@ def recommend_discrete_with_subsets(
         recommender: The recommender instance.
         subspace_discrete: The discrete subspace from which to generate
             recommendations.
-        candidates_exp: The experimental representation of candidates.
         batch_size: The size of the recommendation batch.
 
     Returns:
@@ -47,21 +46,25 @@ def recommend_discrete_with_subsets(
     masks: Iterable[npt.NDArray[np.bool_]]
     if subspace_discrete.n_subsets <= recommender.max_n_subsets:
         masks = subspace_discrete.subset_masks(
-            candidates_exp, min_candidates=batch_size
+            subspace_discrete.exp_rep, min_candidates=batch_size
         )
     else:
         masks = subspace_discrete.sample_subset_masks(
-            candidates_exp, recommender.max_n_subsets, min_candidates=batch_size
+            subspace_discrete.exp_rep,
+            recommender.max_n_subsets,
+            min_candidates=batch_size,
         )
 
     def make_callable(
         mask: np.ndarray,
     ) -> Callable[[], tuple[pd.DataFrame, Tensor]]:
         def optimize() -> tuple[pd.DataFrame, Tensor]:
-            subset = candidates_exp.loc[mask]
+            subset_subspace = evolve(
+                subspace_discrete, exp_rep=subspace_discrete.exp_rep.loc[mask]
+            )
 
             rec = recommend_discrete_without_subsets(
-                recommender, subspace_discrete, subset, batch_size
+                recommender, subset_subspace, batch_size
             )
 
             comp = subspace_discrete.transform(rec)
@@ -79,7 +82,6 @@ def optimize() -> tuple[pd.DataFrame, Tensor]:
 def recommend_discrete_without_subsets(
     recommender: BotorchRecommender,
     subspace_discrete: SubspaceDiscrete,
-    candidates_exp: pd.DataFrame,
     batch_size: int,
 ) -> pd.DataFrame:
     """Generate recommendations from a discrete search space.
@@ -88,8 +90,6 @@ def recommend_discrete_without_subsets(
         recommender: The recommender instance.
         subspace_discrete: The discrete subspace from which to generate
             recommendations.
-        candidates_exp: The experimental representation of all discrete candidate
-            points to be considered.
         batch_size: The size of the recommendation batch.
 
     Raises:
@@ -120,13 +120,13 @@ def recommend_discrete_without_subsets(
 
     from botorch.optim import optimize_acqf_discrete
 
-    # determine the next set of points to be tested
-    candidates_comp = subspace_discrete.transform(candidates_exp)
+    # Determine the next set of points to be tested
+    candidates_comp = subspace_discrete.comp_rep
     points, _ = optimize_acqf_discrete(
         recommender._botorch_acqf, batch_size, to_tensor(candidates_comp)
     )
 
-    # retrieve the rows from the input dataframe corresponding to the selected points
+    # Retrieve the rows from the subspace corresponding to the selected points
     # IMPROVE: The merging procedure is conceptually similar to what
     #   `SearchSpace._match_measurement_with_searchspace_indices` does, though using
     #   a simpler matching logic. When refactoring the SearchSpace class to
@@ -140,4 +140,4 @@ def recommend_discrete_without_subsets(
         )["index"]
     )
 
-    return candidates_exp.loc[idxs]
+    return subspace_discrete.exp_rep.loc[idxs]