Narwhalify model layer

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
@@ -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
 
@@ -799,7 +800,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 +826,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 +846,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 +877,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 +911,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.
 

baybe/objectives/base.py

@@ -15,7 +15,12 @@
 from baybe.targets.base import Target
 from baybe.targets.numerical import NumericalTarget
 from baybe.utils.basic import is_all_instance
-from baybe.utils.dataframe import _copy_index, get_transform_objects, to_tensor
+from baybe.utils.dataframe import (
+    _copy_index,
+    _df_with_backend,
+    get_transform_objects,
+    to_tensor,
+)
 from baybe.utils.dataframe import (
     handle_missing_values as df_handle_missing_values,
 )
@@ -111,9 +116,9 @@ def _full_transformation(self) -> MCAcquisitionObjective:
         """The end-to-end transformation applied, from targets to objective values."""
         return self.to_botorch()
 
-    def handle_missing_values(
-        self, measurements: pd.DataFrame
-    ) -> dict[str, pd.DataFrame]:
+    def _handle_missing_values(
+        self, measurements: IntoDataFrameT
+    ) -> dict[str, IntoDataFrameT]:
         """Handle missing values in the given measurements for each modeled quantity.
 
         Args:
@@ -122,10 +127,16 @@ def handle_missing_values(
         Returns:
             A dictionary with one dataframe for each modeled quantity.
         """
-        cleaned: dict[str, pd.DataFrame] = {}
+        # TODO: The logic should be reworked. Currently, the returned dataframes contain
+        #    all original columns, even those not related to the modeled quantities.
+        measurements_nw = nw.from_native(measurements, eager_only=True)
+        measurements_pd = measurements_nw.to_pandas()
+        cleaned: dict[str, IntoDataFrameT] = {}
         for quantity, target_names in self._model_quantities_to_target_names.items():
-            data = df_handle_missing_values(measurements, target_names, drop=True)
-            cleaned[quantity] = data
+            data_pd = df_handle_missing_values(measurements_pd, target_names, drop=True)
+            cleaned[quantity] = _df_with_backend(
+                nw.from_native(data_pd, eager_only=True), measurements_nw.implementation
+            ).to_native()
 
         return cleaned
 

baybe/recommenders/naive.py

@@ -100,23 +100,11 @@ def recommend(
         cont_part = searchspace.continuous.sample_uniform(1)
         cont_part_tensor = to_tensor(cont_part).unsqueeze(-2)
 
-        # Convert to pandas for BoTorch internals (acqf setup requires pd.DataFrame)
-        measurements_pd = (
-            nw.from_native(measurements, eager_only=True).to_pandas()
-            if measurements is not None
-            else None
-        )
-        pending_experiments_pd = (
-            nw.from_native(pending_experiments, eager_only=True).to_pandas()
-            if pending_experiments is not None
-            else None
-        )
-
         # We now check whether the discrete recommender is bayesian.
         if isinstance(self.disc_recommender, BayesianRecommender):
             # Get access to the recommenders acquisition function
             self.disc_recommender._setup_botorch_acqf(
-                searchspace, objective, measurements_pd, pending_experiments_pd
+                searchspace, objective, measurements, pending_experiments
             )
 
             # Construct the partial acquisition function that attaches cont_part
@@ -143,7 +131,7 @@ def recommend(
 
             # Setup a fresh acquisition function for the continuous recommender
             self.cont_recommender._setup_botorch_acqf(
-                searchspace, objective, measurements_pd, pending_experiments_pd
+                searchspace, objective, measurements, pending_experiments
             )
 
             # Construct the continuous space as a standalone space