Lazify get_candidates

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   anywhere. In particular, recommendation dataframes no longer reflect discrete subspace
   locations via their index. Similarly, `_recommend_discrete` and kin now return a
   dataframe-based subselection instead of a `pd.Index`.
+- `Objective.handle_missing_values` is now private (`_handle_missing_values`)
 - All optional arguments of `SubspaceDiscrete.from_simplex` after `simplex_parameters` 
   are now keyword-only
 - `Campaign.measurements` no longer contains `FitNr` or `BatchNr` metadata columns

baybe/acquisition/_builder.py

@@ -4,9 +4,9 @@
 from functools import cached_property
 from inspect import signature
 from types import MappingProxyType
-from typing import Any
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
-import pandas as pd
+import narwhals.stable.v2 as nw
 import torch
 from attrs import Attribute, asdict, define, field, fields
 from attrs.validators import instance_of, optional
@@ -42,7 +42,12 @@
 from baybe.targets.numerical import NumericalTarget
 from baybe.transformations import IdentityTransformation
 from baybe.utils.basic import match_attributes
-from baybe.utils.dataframe import handle_missing_values, to_tensor
+from baybe.utils.dataframe import _df_with_backend, handle_missing_values, to_tensor
+
+if TYPE_CHECKING:
+    from narwhals.stable.v2.typing import IntoDataFrameT
+else:
+    IntoDataFrameT = TypeVar("IntoDataFrameT")
 
 _OPT_FIELD: None = object()  # type: ignore[assignment]
 """Sentinel value indicating optional acquisition function attributes."""
@@ -96,7 +101,7 @@ def collect(self) -> dict[str, Any]:
 
 
 @define
-class BotorchAcquisitionFunctionBuilder:
+class BotorchAcquisitionFunctionBuilder(Generic[IntoDataFrameT]):
     """A class for building BoTorch acquisition functions from BayBE objects."""
 
     # The BayBE acquisition function to be translated
@@ -106,8 +111,8 @@ class BotorchAcquisitionFunctionBuilder:
     surrogate: SurrogateProtocol = field()
     searchspace: SearchSpace = field()
     objective: Objective = field()
-    measurements: pd.DataFrame = field()
-    pending_experiments: pd.DataFrame | None = field(default=None)
+    measurements: IntoDataFrameT = field()
+    pending_experiments: IntoDataFrameT | None = field(default=None)
 
     # Context shared across building methods
     _args: BotorchAcquisitionArgs = field(init=False)
@@ -134,7 +139,7 @@ def _botorch_surrogate(self) -> Model:
         return self.surrogate.to_botorch()
 
     @cached_property
-    def _train_x(self) -> pd.DataFrame:
+    def _train_x(self) -> IntoDataFrameT:
         """The training parameter values."""
         return self.searchspace.transform(self.measurements, allow_extra=True)
 
@@ -161,7 +166,7 @@ def _posterior_mean_comp(self) -> Tensor:
         return mean.squeeze(-2)
 
     @cached_property
-    def _target_configurations(self) -> pd.DataFrame:
+    def _target_configurations(self) -> IntoDataFrameT:
         """The target configurations used for reference point calculation.
 
         Only completely measured points are considered.
@@ -172,9 +177,11 @@ def _target_configurations(self) -> pd.DataFrame:
         Raises:
             ValueError: If no complete measurement exists.
         """
+        target_names = [t.name for t in self.objective.targets]
+        measurements_nw = nw.from_native(self.measurements, eager_only=True)
         configurations = handle_missing_values(
-            self.measurements[[t.name for t in self.objective.targets]],
-            [t.name for t in self.objective.targets],
+            measurements_nw.select(target_names).to_pandas(),
+            target_names,
             drop=True,
         )
 
@@ -190,7 +197,10 @@ def _target_configurations(self) -> pd.DataFrame:
                 f"argument of '{self.acqf.__class__.__name__}' explicitly."
             )
 
-        return configurations
+        return _df_with_backend(
+            nw.from_native(configurations, eager_only=True),
+            measurements_nw.implementation,
+        ).to_native()
 
     def build(self) -> BoAcquisitionFunction:
         """Build the BoTorch acquisition function object."""

baybe/acquisition/acqfs.py

@@ -5,8 +5,9 @@
 import gc
 import math
 from abc import ABC
-from typing import ClassVar
+from typing import TYPE_CHECKING, ClassVar
 
+import narwhals.stable.v2 as nw
 import numpy as np
 import numpy.typing as npt
 import pandas as pd
@@ -16,9 +17,14 @@
 from attrs.validators import gt, instance_of, le
 from typing_extensions import override
 
+if TYPE_CHECKING:
+    from narwhals.stable.v2.typing import IntoDataFrame
+
 from baybe.acquisition.base import AcquisitionFunction
 from baybe.searchspace import SearchSpace
+from baybe.settings import active_settings
 from baybe.utils.basic import classproperty, convert_to_float
+from baybe.utils.dataframe import _df_with_backend
 from baybe.utils.sampling_algorithms import DiscreteSamplingMethod, sample_numerical_df
 from baybe.utils.validation import finite_float
 
@@ -81,7 +87,7 @@ def _non_botorch_attrs(cls: type[AttrsInstance]) -> tuple[str, ...]:
             flds.sampling_fraction.name,
         )
 
-    def get_integration_points(self, searchspace: SearchSpace) -> pd.DataFrame:
+    def get_integration_points(self, searchspace: SearchSpace) -> IntoDataFrame:
         """Sample points from a search space for integration purposes.
 
         Sampling of the discrete part can be controlled via 'sampling_method', but
@@ -105,7 +111,7 @@ def get_integration_points(self, searchspace: SearchSpace) -> pd.DataFrame:
         # Discrete part
         if not searchspace.discrete.is_empty:
             candidates_discrete = searchspace.discrete.transform(
-                searchspace.discrete.get_candidates()
+                searchspace.discrete._get_candidates().collect().to_pandas()
             )
             n_candidates = self.sampling_n_points or math.ceil(
                 self.sampling_fraction * len(candidates_discrete)  # type: ignore[operator]
@@ -137,7 +143,10 @@ def get_integration_points(self, searchspace: SearchSpace) -> pd.DataFrame:
         # Combine different search space parts
         result = pd.concat(sampled_parts, axis=1)
 
-        return result
+        return _df_with_backend(
+            nw.from_native(result, eager_only=True),
+            active_settings.default_dataframe_backend,
+        ).to_native()
 
 
 ########################################################################################

baybe/acquisition/base.py

@@ -6,7 +6,7 @@
 from abc import ABC
 from typing import TYPE_CHECKING, ClassVar, Literal, overload
 
-import pandas as pd
+import narwhals.stable.v2 as nw
 from attrs import define
 
 from baybe.exceptions import (
@@ -19,10 +19,11 @@
 from baybe.surrogates.base import SurrogateProtocol
 from baybe.utils.basic import classproperty
 from baybe.utils.boolean import is_abstract
-from baybe.utils.dataframe import to_tensor
+from baybe.utils.dataframe import _copy_index, to_tensor
 
 if TYPE_CHECKING:
     from botorch.acquisition import AcquisitionFunction as BotorchAcquisitionFunction
+    from narwhals.stable.v2.typing import IntoDataFrameT, IntoSeries
 
 
 @define(frozen=True)
@@ -63,8 +64,8 @@ def to_botorch(
         surrogate: SurrogateProtocol,
         searchspace: SearchSpace,
         objective: Objective,
-        measurements: pd.DataFrame,
-        pending_experiments: pd.DataFrame | None = None,
+        measurements: IntoDataFrameT,
+        pending_experiments: IntoDataFrameT | None = None,
     ) -> BotorchAcquisitionFunction:
         """Create the botorch-ready representation of the function.
 
@@ -86,40 +87,40 @@ def to_botorch(
     @overload
     def evaluate(
         self,
-        candidates: pd.DataFrame,
+        candidates: IntoDataFrameT,
         surrogate: SurrogateProtocol,
         searchspace: SearchSpace,
         objective: Objective,
-        measurements: pd.DataFrame,
-        pending_experiments: pd.DataFrame | None = None,
+        measurements: IntoDataFrameT,
+        pending_experiments: IntoDataFrameT | None = None,
         *,
         jointly: Literal[True],
     ) -> float: ...
 
     @overload
     def evaluate(
         self,
-        candidates: pd.DataFrame,
+        candidates: IntoDataFrameT,
         surrogate: SurrogateProtocol,
         searchspace: SearchSpace,
         objective: Objective,
-        measurements: pd.DataFrame,
-        pending_experiments: pd.DataFrame | None = None,
+        measurements: IntoDataFrameT,
+        pending_experiments: IntoDataFrameT | None = None,
         *,
         jointly: Literal[False] = False,
-    ) -> pd.Series: ...
+    ) -> IntoSeries: ...
 
     def evaluate(
         self,
-        candidates: pd.DataFrame,
+        candidates: IntoDataFrameT,
         surrogate: SurrogateProtocol,
         searchspace: SearchSpace,
         objective: Objective,
-        measurements: pd.DataFrame,
-        pending_experiments: pd.DataFrame | None = None,
+        measurements: IntoDataFrameT,
+        pending_experiments: IntoDataFrameT | None = None,
         *,
         jointly: bool = False,
-    ) -> pd.Series | float:
+    ) -> IntoSeries | float:
         """Get the acquisition values for the given candidates.
 
         Args:
@@ -144,6 +145,8 @@ def evaluate(
         """
         import torch
 
+        candidates_nw = nw.from_native(candidates, eager_only=True)
+
         # Assemble the Botorch acquisition function and its input
         botorch_acqf = self.to_botorch(
             surrogate, searchspace, objective, measurements, pending_experiments
@@ -156,7 +159,11 @@ def evaluate(
             out = botorch_acqf(in_)
         if jointly:
             return out.item()
-        return pd.Series(out.numpy(), index=candidates.index)
+
+        backend = nw.get_native_namespace(candidates_nw)
+        return _copy_index(
+            nw.new_series("", out.numpy(), backend=backend), candidates_nw
+        ).to_native()
 
 
 def _get_botorch_acqf_class(

baybe/campaign.py

@@ -56,6 +56,7 @@
 if TYPE_CHECKING:
     from botorch.acquisition import AcquisitionFunction as BoAcquisitionFunction
     from botorch.posteriors import Posterior
+    from narwhals.stable.v2.typing import IntoDataFrame, IntoDataFrameT, IntoSeries
 
     from baybe.acquisition.base import AcquisitionFunction
 
@@ -430,7 +431,7 @@ def toggle_discrete_candidates(  # noqa: DOC501
         #  * Additional shortcuts might be possible.
         self.clear_cache()
 
-        df = self.searchspace.discrete.get_candidates()
+        df = self.searchspace.discrete._get_candidates().collect().to_pandas()
 
         if isinstance(constraints, pd.DataFrame):
             # Determine the candidate subset to be toggled
@@ -531,7 +532,9 @@ def recommend(
         if self.searchspace.type is SearchSpaceType.DISCRETE:
             # TODO: This implementation should at some point be hidden behind an
             #   appropriate public interface, like `SubspaceDiscrete.filter()`
-            candidates = self.searchspace.discrete.get_candidates()
+            candidates = (
+                self.searchspace.discrete._get_candidates().collect().to_pandas()
+            )
             mask_todrop = pd.Series(False, index=candidates.index)
             if not self._excluded_experiments.empty:
                 mask_todrop |= (
@@ -799,7 +802,7 @@ def get_surrogate(
     def _get_non_meta_recommender(
         self,
         batch_size: int | None = None,
-        pending_experiments: pd.DataFrame | None = None,
+        pending_experiments: IntoDataFrame | None = None,
     ) -> RecommenderProtocol:
         """Get the current recommender.
 
@@ -825,7 +828,7 @@ def _get_non_meta_recommender(
     def _get_bayesian_recommender(
         self,
         batch_size: int | None = None,
-        pending_experiments: pd.DataFrame | None = None,
+        pending_experiments: IntoDataFrame | None = None,
     ) -> BayesianRecommender:
         """Get the current Bayesian recommender (if available).
 
@@ -845,7 +848,7 @@ def _get_bayesian_recommender(
     def get_acquisition_function(
         self,
         batch_size: int | None = None,
-        pending_experiments: pd.DataFrame | None = None,
+        pending_experiments: IntoDataFrame | None = None,
     ) -> BoAcquisitionFunction:
         """Get the current BoTorch acquisition function.
 
@@ -876,12 +879,12 @@ def get_acquisition_function(
 
     def acquisition_values(
         self,
-        candidates: pd.DataFrame,
+        candidates: IntoDataFrameT,
         acquisition_function: AcquisitionFunction | None = None,
         *,
         batch_size: int | None = None,
-        pending_experiments: pd.DataFrame | None = None,
-    ) -> pd.Series:
+        pending_experiments: IntoDataFrameT | None = None,
+    ) -> IntoSeries:
         """Compute the acquisition values for the given candidates.
 
         Args:
@@ -910,11 +913,11 @@ def acquisition_values(
 
     def joint_acquisition_value(  # noqa: DOC101, DOC103
         self,
-        candidates: pd.DataFrame,
+        candidates: IntoDataFrameT,
         acquisition_function: AcquisitionFunction | None = None,
         *,
         batch_size: int | None = None,
-        pending_experiments: pd.DataFrame | None = None,
+        pending_experiments: IntoDataFrameT | None = None,
     ) -> float:
         """Compute the joint acquisition values for the given candidate batch.
 
@@ -1081,7 +1084,9 @@ def _structure_campaign(d: dict, cl: type) -> Campaign:
     # >>>>>>>>>> Deprecation
     # Post-structure reconstruction from legacy metadata indices
     if legacy_recommended_idxs is not None or legacy_excluded_idxs is not None:
-        candidates = campaign.searchspace.discrete.get_candidates()
+        candidates = (
+            campaign.searchspace.discrete._get_candidates().collect().to_pandas()
+        )
         if legacy_recommended_idxs is not None:
             campaign._recommended_experiments = candidates.loc[
                 legacy_recommended_idxs