Simplify and extend meta recommender logic

CHANGELOG.md

@@ -19,6 +19,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `arrays_to_dataframes` decorator to create lookups from array-based callables
 - `DiscreteConstraint.get_valid` to conveniently access valid candidates
 - Functionality for persisting benchmarking results on S3 from a manual pipeline run
+- `remain_switched` option to `TwoPhaseMetaRecommender`
+- `BatchSizeAdaptiveMetaRecommender` for selecting recommenders based on batch size
+- `is_stateful` class variable to `MetaRecommenders`
 
 ### Changed
 - `SubstanceParameter` encodings are now computed exclusively with the
@@ -29,6 +32,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `CustomDiscreteParameter` does not allow duplicated rows in `data` anymore
 - De-/activating Polars via `BAYBE_DEACTIVATE_POLARS` now requires passing values
   compatible with `strtobool`
+- All arguments to `MetaRecommender.select_recommender` are now optional
+- `MetaRecommender`s can now be composed of other `MetaRecommender`s
 
 ### Fixed
 - Rare bug arising from degenerate `SubstanceParameter.comp_df` rows that caused
@@ -58,6 +63,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `SubstanceEncoding` value `RDKIT`. As a replacement, `RDKIT2DDESCRIPTORS` can be used.
 - The `metadata` attribute of `SubspaceDiscrete` no longer exists. Metadata is now
   exclusively handled by the `Campaign` class.
+- `get_current_recommender` and `get_next_recommender` of `MetaRecommender` have become
+  obsolete and calling them is no longer possible
 
 ## [0.11.3] - 2024-11-06
 ### Fixed

baybe/campaign.py

@@ -403,9 +403,6 @@ def recommend(
     def posterior(self, candidates: pd.DataFrame) -> Posterior:
         """Get the posterior predictive distribution for the given candidates.
 
-        The predictive distribution is based on the surrogate model of the last used
-        recommender.
-
         Args:
             candidates: The candidate points in experimental recommendations.
                 For details, see :meth:`baybe.surrogates.base.Surrogate.posterior`.
@@ -430,9 +427,19 @@ def posterior(self, candidates: pd.DataFrame) -> Posterior:
         with torch.no_grad():
             return surrogate.posterior(candidates)
 
-    def get_surrogate(self) -> SurrogateProtocol:
+    def get_surrogate(
+        self,
+        batch_size: int | None = None,
+        pending_experiments: pd.DataFrame | None = None,
+    ) -> SurrogateProtocol:
         """Get the current surrogate model.
 
+        Args:
+            batch_size: See :meth:`recommend`.
+                Only required when using meta recommenders that demand it.
+            pending_experiments: See :meth:`recommend`.
+                Only required when using meta recommenders that demand it.
+
         Raises:
             RuntimeError: If the current recommender does not provide a surrogate model.
 
@@ -453,7 +460,13 @@ def get_surrogate(self) -> SurrogateProtocol:
 
         pure_recommender: RecommenderProtocol
         if isinstance(self.recommender, MetaRecommender):
-            pure_recommender = self.recommender.get_current_recommender()
+            pure_recommender = self.recommender.select_recommender(
+                batch_size,
+                self.searchspace,
+                self.objective,
+                self.measurements,
+                pending_experiments,
+            )
         else:
             pure_recommender = self.recommender
 

baybe/recommenders/meta/adaptive.py

@@ -0,0 +1,48 @@
+"""Meta recommenders that adaptively select recommenders based on the context."""
+
+import pandas as pd
+from attrs import define, field
+from attrs.validators import deep_iterable, instance_of
+from typing_extensions import override
+
+from baybe.objectives.base import Objective
+from baybe.recommenders.base import RecommenderProtocol
+from baybe.recommenders.meta.base import MetaRecommender
+from baybe.searchspace.core import SearchSpace
+from baybe.utils.interval import Partition
+
+
+@define
+class BatchSizeAdaptiveMetaRecommender(MetaRecommender):
+    """A meta recommender that selects recommenders according to the batch size."""
+
+    recommenders: list[RecommenderProtocol] = field(
+        converter=list, validator=deep_iterable(instance_of(RecommenderProtocol))
+    )
+    """The recommenders for the individual batch size intervals."""
+
+    partition: Partition = field(
+        converter=lambda x: Partition(x) if not isinstance(x, Partition) else x
+    )
+    """The partition mapping batch size intervals to recommenders. """
+
+    @partition.validator
+    def _validate_partitioning(self, _, value):
+        if (lr := len(self.recommenders)) != (lp := len(value)):
+            raise ValueError(
+                f"The number of recommenders (given: {lr}) must be equal to the number "
+                f"of intervals defined by the partition (given: {lp})."
+            )
+
+    @override
+    def select_recommender(
+        self,
+        batch_size: int | None = None,
+        searchspace: SearchSpace | None = None,
+        objective: Objective | None = None,
+        measurements: pd.DataFrame | None = None,
+        pending_experiments: pd.DataFrame | None = None,
+    ) -> RecommenderProtocol:
+        if batch_size is None:
+            raise ValueError("A batch size is required.")
+        return self.recommenders[self.partition.get_interval_index(batch_size)]

baybe/recommenders/meta/base.py

@@ -2,13 +2,14 @@
 
 import gc
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import Any, ClassVar
 
 import cattrs
 import pandas as pd
-from attrs import define, field
+from attrs import define
 from typing_extensions import override
 
+from baybe.exceptions import DeprecationError
 from baybe.objectives.base import Objective
 from baybe.recommenders.base import RecommenderProtocol
 from baybe.recommenders.pure.base import PureRecommender
@@ -22,76 +23,39 @@
 class MetaRecommender(SerialMixin, RecommenderProtocol, ABC):
     """Abstract base class for all meta recommenders."""
 
-    _current_recommender: PureRecommender | None = field(default=None, init=False)
-    """The current recommender."""
-
-    _used_recommender_ids: set[int] = field(factory=set, init=False)
-    """Set of ids from recommenders that were used by this meta recommender."""
+    is_stateful: ClassVar[bool] = False
+    """Boolean flag indicating if the meta recommender is stateful."""
 
     @abstractmethod
     def select_recommender(
         self,
-        batch_size: int,
-        searchspace: SearchSpace,
+        batch_size: int | None = None,
+        searchspace: SearchSpace | None = None,
         objective: Objective | None = None,
         measurements: pd.DataFrame | None = None,
         pending_experiments: pd.DataFrame | None = None,
-    ) -> PureRecommender:
+    ) -> RecommenderProtocol:
         """Select a pure recommender for the given experimentation context.
 
         See :meth:`baybe.recommenders.base.RecommenderProtocol.recommend` for details
         on the method arguments.
         """
 
     def get_current_recommender(self) -> PureRecommender:
-        """Get the current recommender, if available."""
-        if self._current_recommender is None:
-            raise RuntimeError(
-                f"No recommendation has been requested from the "
-                f"'{self.__class__.__name__}' yet. Because the recommender is a "
-                f"'{MetaRecommender.__name__}', this means no actual recommender has "
-                f"been selected so far. The recommender will be available after the "
-                f"next '{self.recommend.__name__}' call."
-            )
-        return self._current_recommender
-
-    def get_next_recommender(
-        self,
-        batch_size: int,
-        searchspace: SearchSpace,
-        objective: Objective | None = None,
-        measurements: pd.DataFrame | None = None,
-        pending_experiments: pd.DataFrame | None = None,
-    ) -> PureRecommender:
-        """Get the recommender for the next recommendation.
-
-        Returns the next recommender in row that has not yet been used for generating
-        recommendations. In case of multiple consecutive calls, this means that
-        the same recommender instance is returned until its :meth:`recommend` method
-        is called.
+        """Deprecated! Use :meth:`select_recommender` instead."""  # noqa: D401
+        raise DeprecationError(
+            f"'{MetaRecommender.__name__}.get_current_recommender' has become "
+            f"obsolete. Use "
+            f"'{MetaRecommender.__name__}.{self.select_recommender.__name__}' instead."
+        )
 
-        See :meth:`baybe.recommenders.base.RecommenderProtocol.recommend` for details
-        on the method arguments.
-        """
-        # Check if the stored recommender instance can be returned
-        if (
-            self._current_recommender is not None
-            and id(self._current_recommender) not in self._used_recommender_ids
-        ):
-            recommender = self._current_recommender
-
-        # Otherwise, fetch the next recommender waiting in row
-        else:
-            recommender = self.select_recommender(
-                batch_size=batch_size,
-                searchspace=searchspace,
-                objective=objective,
-                measurements=measurements,
-                pending_experiments=pending_experiments,
-            )
-            self._current_recommender = recommender
-
-        return recommender
+    def get_next_recommender(self) -> PureRecommender:
+        """Deprecated! Use :meth:`select_recommender` instead."""  # noqa: D401
+        raise DeprecationError(
+            f"'{MetaRecommender.__name__}.get_next_recommender' has become "
+            f"obsolete. Use "
+            f"'{MetaRecommender.__name__}.{self.select_recommender.__name__}' instead."
+        )
 
     @override
     def recommend(
@@ -103,7 +67,7 @@ def recommend(
         pending_experiments: pd.DataFrame | None = None,
     ) -> pd.DataFrame:
         """See :meth:`baybe.recommenders.base.RecommenderProtocol.recommend`."""
-        recommender = self.get_next_recommender(
+        recommender = self.select_recommender(
             batch_size=batch_size,
             searchspace=searchspace,
             objective=objective,
@@ -123,15 +87,12 @@ def recommend(
             }
         )
 
-        recommendations = recommender.recommend(
+        return recommender.recommend(
             batch_size=batch_size,
             searchspace=searchspace,
             pending_experiments=pending_experiments,
             **optional_args,
         )
-        self._used_recommender_ids.add(id(recommender))
-
-        return recommendations
 
 
 # Register (un-)structure hooks